Implement multipart upload(AST-24859) (#1343)

* implemented multipart upload

* Added wrapper method in mock

* fix unit test cases and lint issue

* Add test cases

* fix unit test case

* debug unit test case

* debug unit test case

* unit test case

* unit test case

* Remove print msg

* handle  multipart_file_size value is empty

* rename the MultipartFileSizeKey

* fix lint issue

* Made code chnages according to review comments

* remove extra comments

---------

Co-authored-by: Anurag Dalke <[email protected]>

Author: cx-sumit-morchhale

internal/commands/scan.go

@@ -153,6 +153,7 @@ var (
 	aditionalParameters []string
 	kicsErrorCodes      = []string{"60", "50", "40", "30", "20"}
 	containerResolver   wrappers.ContainerResolverWrapper
+	MaxSizeBytes        int64 = 5 * 1024 * 1024 * 1024 // 5 GB in bytes
 )
 
 func NewScanCommand(
@@ -2082,7 +2083,20 @@ func uploadZip(uploadsWrapper wrappers.UploadsWrapper, zipFilePath string, unzip
 	var zipFilePathErr error
 	// Send a request to uploads service
 	var preSignedURL *string
-	preSignedURL, zipFilePathErr = uploadsWrapper.UploadFile(zipFilePath, featureFlagsWrapper)
+
+	fileInfo, err := os.Stat(zipFilePath)
+	if err != nil {
+		return "", zipFilePath, errors.Wrapf(err, "Failed to check the size - %s", zipFilePath)
+	}
+	logger.PrintIfVerbose(fmt.Sprintf("Zip size before upload:  %.2fMB\n", float64(fileInfo.Size())/mbBytes))
+
+	flagResponse, _ := featureFlagsWrapper.GetSpecificFlag(wrappers.IncreaseFileUploadLimit)
+	if flagResponse.Status && fileInfo.Size() > MaxSizeBytes {
+		logger.PrintIfVerbose("Uploading source code in multiple parts.")
+		preSignedURL, zipFilePathErr = uploadsWrapper.UploadFileInMultipart(zipFilePath, featureFlagsWrapper)
+	} else {
+		preSignedURL, zipFilePathErr = uploadsWrapper.UploadFile(zipFilePath, featureFlagsWrapper)
+	}
 	if zipFilePathErr != nil {
 		if unzip || !userProvidedZip {
 			return "", zipFilePath, errors.Wrapf(zipFilePathErr, "%s: Failed to upload sources file\n", failedCreating)

internal/commands/scan_test.go

@@ -2972,15 +2972,31 @@ func TestResubmitConfig_ProjectDoesNotExist_ReturnedEmptyConfig(t *testing.T) {
 func TestUploadZip_whenUserProvideZip_shouldReturnEmptyZipFilePathInSuccessCase(t *testing.T) {
 	uploadWrapper := mock.UploadsMockWrapper{}
 	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
-	_, zipPath, err := uploadZip(&uploadWrapper, "test.zip", false, true, featureFlagsWrapper)
+	_, zipPath, err := uploadZip(&uploadWrapper, "data/sources.zip", false, true, featureFlagsWrapper)
 	assert.NilError(t, err)
 	assert.Equal(t, zipPath, "")
 }
 
 func TestUploadZip_whenUserProvideZip_shouldReturnEmptyZipFilePathInFailureCase(t *testing.T) {
+	// Create a temporary zip file
+	dir := t.TempDir()
+	zipPathTemp := filepath.Join(dir, "failureCase.zip")
+
+	// Create the zip file
+	zipFile, err := os.Create(zipPathTemp)
+	if err != nil {
+		t.Fatalf("Failed to create zip file: %v", err)
+	}
+	defer func(zipFile *os.File) {
+		err := zipFile.Close()
+		if err != nil {
+			t.Fatalf("Failed to close zip file: %v", err)
+		}
+	}(zipFile)
+
 	uploadWrapper := mock.UploadsMockWrapper{}
 	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
-	_, zipPath, err := uploadZip(&uploadWrapper, "failureCase.zip", false, true, featureFlagsWrapper)
+	_, zipPath, err := uploadZip(&uploadWrapper, zipPathTemp, false, true, featureFlagsWrapper)
 	assert.Assert(t, err != nil)
 	assert.Assert(t, strings.Contains(err.Error(), "error from UploadFile"), err.Error())
 	assert.Equal(t, zipPath, "")
@@ -2989,18 +3005,33 @@ func TestUploadZip_whenUserProvideZip_shouldReturnEmptyZipFilePathInFailureCase(
 func TestUploadZip_whenUserNotProvideZip_shouldReturnZipFilePathInSuccessCase(t *testing.T) {
 	uploadWrapper := mock.UploadsMockWrapper{}
 	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
-	_, zipPath, err := uploadZip(&uploadWrapper, "test.zip", false, false, featureFlagsWrapper)
+	_, zipPath, err := uploadZip(&uploadWrapper, "data/sources.zip", false, false, featureFlagsWrapper)
 	assert.NilError(t, err)
-	assert.Equal(t, zipPath, "test.zip")
+	assert.Equal(t, zipPath, "data/sources.zip")
 }
 
 func TestUploadZip_whenUserNotProvideZip_shouldReturnZipFilePathInFailureCase(t *testing.T) {
+	// Create a temporary zip file
+	dir := t.TempDir()
+	zipPathTemp := filepath.Join(dir, "failureCase.zip")
+
+	// Create the zip file
+	zipFile, err := os.Create(zipPathTemp)
+	if err != nil {
+		t.Fatalf("Failed to create zip file: %v", err)
+	}
+	defer func(zipFile *os.File) {
+		err := zipFile.Close()
+		if err != nil {
+			t.Fatalf("Failed to close zip file: %v", err)
+		}
+	}(zipFile)
 	uploadWrapper := mock.UploadsMockWrapper{}
 	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
-	_, zipPath, err := uploadZip(&uploadWrapper, "failureCase.zip", false, false, featureFlagsWrapper)
+	_, zipPath, err := uploadZip(&uploadWrapper, zipPathTemp, false, false, featureFlagsWrapper)
 	assert.Assert(t, err != nil)
 	assert.Assert(t, strings.Contains(err.Error(), "error from UploadFile"), err.Error())
-	assert.Equal(t, zipPath, "failureCase.zip")
+	assert.Equal(t, zipPath, zipPathTemp)
 }
 
 func TestAddSastScan_ScanFlags(t *testing.T) {
@@ -4249,3 +4280,96 @@ func TestEnforceLocalResolutionForTarFiles_Integration(t *testing.T) {
 		})
 	}
 }
+
+func TestUploadZip_AsMultipartUpload_when_FF_Enable_ZIP_Exceeds_5GB(t *testing.T) {
+	// Simulate a file size > 5GB by setting MaxSizeBytes to less than actual size
+	MaxSizeBytes = 1
+	defer func() { MaxSizeBytes = 5 * 1024 * 1024 * 1024 }() // Reset after test
+	uploadWrapper := mock.UploadsMockWrapper{}
+	mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.IncreaseFileUploadLimit, Status: true}
+	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
+	url, zipPath, err := uploadZip(&uploadWrapper, "data/sources.zip", false, true, featureFlagsWrapper)
+	assert.NilError(t, err)
+	assert.Equal(t, zipPath, "")
+	assert.Equal(t, url, "multiPart/path/to/nowhere")
+}
+
+func TestUploadZip_AsMultipartUpload_when_FF_Disable_ZIP_Exceeds_5GB(t *testing.T) {
+	fileInfo, err := os.Stat("data/sources.zip")
+	if err != nil {
+		t.Fatalf("Failed to close zip file: %v", err)
+	}
+	// Simulate a file size > 5GB by setting MaxSizeBytes to less than actual size
+	MaxSizeBytes = fileInfo.Size() - 1
+	defer func() { MaxSizeBytes = 5 * 1024 * 1024 * 1024 }() // Reset after test
+
+	uploadWrapper := mock.UploadsMockWrapper{}
+	mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.IncreaseFileUploadLimit, Status: false}
+	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
+	url, zipPath, err := uploadZip(&uploadWrapper, "data/sources.zip", false, true, featureFlagsWrapper)
+	assert.NilError(t, err)
+	assert.Equal(t, zipPath, "")
+	assert.Equal(t, url, "singlePart/path/to/nowhere")
+}
+
+func TestUploadZip_AsMultipartUpload_when_FF_Enable_ZIP_LessThan_5GB(t *testing.T) {
+	uploadWrapper := mock.UploadsMockWrapper{}
+	mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.IncreaseFileUploadLimit, Status: true}
+	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
+	url, zipPath, err := uploadZip(&uploadWrapper, "data/sources.zip", false, true, featureFlagsWrapper)
+	assert.NilError(t, err)
+	assert.Equal(t, zipPath, "")
+	assert.Equal(t, url, "singlePart/path/to/nowhere")
+}
+
+func TestUploadZip_AsMultipartUpload_when_FF_Disable_ZIP_LessThan_5GB(t *testing.T) {
+	uploadWrapper := mock.UploadsMockWrapper{}
+	mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.IncreaseFileUploadLimit, Status: false}
+	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
+	url, zipPath, err := uploadZip(&uploadWrapper, "data/sources.zip", false, true, featureFlagsWrapper)
+	assert.NilError(t, err)
+	assert.Equal(t, zipPath, "")
+	assert.Equal(t, url, "singlePart/path/to/nowhere")
+}
+
+func TestUploadZip_AsMultipartUpload_when_FF_Enable_ZIP_Exceeds_5GB_Error(t *testing.T) {
+	// Create a temporary zip file
+	dir := t.TempDir()
+	zipPathTemp := filepath.Join(dir, "failureCaseLarge.zip")
+
+	// Create the zip file
+	zipFile, err := os.Create(zipPathTemp)
+	if err != nil {
+		t.Fatalf("Failed to create zip file: %v", err)
+	}
+	defer func(zipFile *os.File) {
+		err := zipFile.Close()
+		if err != nil {
+			t.Fatalf("Failed to close zip file: %v", err)
+		}
+	}(zipFile)
+
+	// Seek to 5KB + 1 byte
+	_, err = zipFile.Seek(5*1024+1, 0) // 5121 bytes
+	if err != nil {
+		panic("Failed to seek in zip file: " + err.Error())
+	}
+
+	// Write a single byte to allocate space
+	_, err = zipFile.Write([]byte{0})
+	if err != nil {
+		panic("Failed to write to zip file: " + err.Error())
+	}
+
+	// Simulate a file size > 5GB by setting MaxSizeBytes to less than actual size
+	MaxSizeBytes = 1
+	defer func() { MaxSizeBytes = 5 * 1024 * 1024 * 1024 }() //
+
+	uploadWrapper := mock.UploadsMockWrapper{}
+	mock.Flag = wrappers.FeatureFlagResponseModel{Name: wrappers.IncreaseFileUploadLimit, Status: true}
+	featureFlagsWrapper := &mock.FeatureFlagsMockWrapper{}
+	_, zipPath, err := uploadZip(&uploadWrapper, zipPathTemp, false, true, featureFlagsWrapper)
+	assert.Assert(t, err != nil)
+	assert.Assert(t, strings.Contains(err.Error(), "error from UploadFileInMultipart"), err.Error())
+	assert.Equal(t, zipPath, "")
+}

internal/params/binds.go

@@ -80,4 +80,8 @@ var EnvVarsBinds = []struct {
 	{RiskManagementPathKey, RiskManagementPathEnv, "api/risk-management/projects/%s/results?scanID=%s"},
 	{ConfigFilePathKey, ConfigFilePathEnv, ""},
 	{RealtimeScannerPathKey, RealtimeScannerPathEnv, "api/realtime-scanner"},
+	{StartMultiPartUploadPathKey, StartMultiPartUploadPathEnv, "api/uploads/start-multipart-upload"},
+	{MultipartPresignedPathKey, MultipartPresignedPathEnv, "api/uploads/multipart-presigned"},
+	{CompleteMultiPartUploadPathKey, CompleteMultipartUploadPathEnv, "api/uploads/complete-multipart-upload"},
+	{MultipartFileSizeKey, MultipartFileSizeEnv, "2"},
 }

internal/params/envs.go

@@ -80,4 +80,8 @@ const (
 	RiskManagementPathEnv               = "CX_RISK_MANAGEMENT_PATH"
 	ConfigFilePathEnv                   = "CX_CONFIG_FILE_PATH"
 	RealtimeScannerPathEnv              = "CX_REALTIME_SCANNER_PATH"
+	StartMultiPartUploadPathEnv         = "CX_START_MULTIPART_UPLOAD_PATH"
+	MultipartPresignedPathEnv           = "CX_MULTIPART_PRESIGNED_URL_PATH"
+	CompleteMultipartUploadPathEnv      = "CX_COMPLETE_MULTIPART_UPLOAD_PATH"
+	MultipartFileSizeEnv                = "MULTIPART_FILE_SIZE"
 )

internal/params/keys.go

@@ -79,4 +79,8 @@ var (
 	RiskManagementPathKey               = strings.ToLower(RiskManagementPathEnv)
 	ConfigFilePathKey                   = strings.ToLower(ConfigFilePathEnv)
 	RealtimeScannerPathKey              = strings.ToLower(RealtimeScannerPathEnv)
+	StartMultiPartUploadPathKey         = strings.ToLower(StartMultiPartUploadPathEnv)
+	MultipartPresignedPathKey           = strings.ToLower(MultipartPresignedPathEnv)
+	CompleteMultiPartUploadPathKey      = strings.ToLower(CompleteMultipartUploadPathEnv)
+	MultipartFileSizeKey                = strings.ToLower(MultipartFileSizeEnv)
 )

internal/wrappers/feature-flags.go

@@ -17,6 +17,7 @@ const OssRealtimeEnabled = "OSS_REALTIME_ENABLED"
 const ScsLicensingV2Enabled = "SSCS_NEW_LICENSING_ENABLED"
 const DirectAssociationEnabled = "DIRECT_APP_ASSOCIATION_ENABLED"
 const maxRetries = 3
+const IncreaseFileUploadLimit = "INCREASE_FILE_UPLOAD_LIMIT"
 
 var DefaultFFLoad bool = false
 

internal/wrappers/mock/uploads-mock.go

@@ -2,6 +2,7 @@ package mock
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/pkg/errors"
 
@@ -11,11 +12,20 @@ import (
 type UploadsMockWrapper struct {
 }
 
+func (u *UploadsMockWrapper) UploadFileInMultipart(filePath string, wrapper wrappers.FeatureFlagsWrapper) (*string, error) {
+	fmt.Println("UploadFileInMultipart called Create in UploadsMockWrapper")
+	if strings.Contains(filePath, "failureCaseLarge.zip") {
+		return nil, errors.New("error from UploadFileInMultipart")
+	}
+	url := "multiPart/path/to/nowhere"
+	return &url, nil
+}
+
 func (u *UploadsMockWrapper) UploadFile(filePath string, featureFlagsWrapper wrappers.FeatureFlagsWrapper) (*string, error) {
 	fmt.Println("Called Create in UploadsMockWrapper")
-	if filePath == "failureCase.zip" {
+	if strings.Contains(filePath, "failureCase.zip") {
 		return nil, errors.New("error from UploadFile")
 	}
-	url := "/path/to/nowhere"
+	url := "singlePart/path/to/nowhere"
 	return &url, nil
 }