Merge pull request #60 from snyk/fix/upload-revision-request-body

fix(uploadrevision): load request body into memory

Author: attriaayush

internal/fileupload/client_integration_test.go

@@ -50,6 +50,24 @@ func TestUploadDirectoryIntegration(t *testing.T) {
 	assert.NotEqual(t, uuid.Nil, fileuploadRevisionID)
 }
 
+func TestUploadLargeFileIntegration(t *testing.T) {
+	fileUploadClient := setupFileUploadClient(t)
+
+	content := generateFileOfSizeMegabytes(t, 30)
+	files := []uploadrevision.LoadedFile{
+		{Path: "src/main.go", Content: content},
+	}
+
+	dir, dirCleanup := createDirWithFiles(t, files)
+	defer dirCleanup()
+
+	fileuploadRevisionID, err := fileUploadClient.CreateRevisionFromFile(t.Context(), filepath.Join(dir.Name(), files[0].Path), fileupload.UploadOptions{})
+	if err != nil {
+		t.Errorf("failed to create fileupload revision: %s", err.Error())
+	}
+	assert.NotEqual(t, uuid.Nil, fileuploadRevisionID)
+}
+
 func setupFileUploadClient(t *testing.T) fileupload.Client {
 	t.Helper()
 
@@ -73,6 +91,12 @@ func setupFileUploadClient(t *testing.T) fileupload.Client {
 	)
 }
 
+func generateFileOfSizeMegabytes(t *testing.T, megabytes int) string {
+	t.Helper()
+	content := make([]byte, megabytes*1024*1024)
+	return string(content)
+}
+
 type testConfig struct {
 	BaseURL  string
 	OrgID    uuid.UUID

internal/fileupload/uploadrevision/client.go

@@ -56,9 +56,6 @@ func NewClient(cfg Config, opts ...Opt) *HTTPSealableClient {
 		opt(&c)
 	}
 
-	crt := NewCompressionRoundTripper(c.httpClient.Transport)
-	c.httpClient.Transport = crt
-
 	return &c
 }
 
@@ -127,13 +124,22 @@ func (c *HTTPSealableClient) UploadFiles(ctx context.Context, orgID OrgID, revis
 	mpartWriter := multipart.NewWriter(pipeWriter)
 
 	go streamFilesToPipe(pipeWriter, mpartWriter, files)
+	body := compressRequestBody(pipeReader)
+
+	// Load body bytes into memmory so go can determine the Content-Length
+	// and not send the request chunked
+	bts, err := io.ReadAll(body)
+	if err != nil {
+		return fmt.Errorf("failed to create upload files request: %w", err)
+	}
 
 	url := fmt.Sprintf("%s/hidden/orgs/%s/upload_revisions/%s/files?version=%s", c.cfg.BaseURL, orgID, revisionID, apiVersion)
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, pipeReader)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bts))
 	if err != nil {
 		return fmt.Errorf("failed to create upload files request: %w", err)
 	}
 	req.Header.Set(ContentType, mpartWriter.FormDataContentType())
+	req.Header.Set(ContentEncoding, "gzip")
 
 	res, err := c.httpClient.Do(req)
 	if err != nil {

internal/fileupload/uploadrevision/client_test.go

@@ -26,7 +26,6 @@ func TestClient_CreateRevision(t *testing.T) {
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		assert.Equal(t, http.MethodPost, r.Method)
 		assert.Equal(t, "application/vnd.api+json", r.Header.Get("Content-Type"))
-		assert.Equal(t, "gzip", r.Header.Get("Content-Encoding"))
 		assert.Equal(t, fmt.Sprintf("/hidden/orgs/%s/upload_revisions", orgID), r.URL.Path)
 		assert.Equal(t, "2024-10-15", r.URL.Query().Get("version"))
 
@@ -383,7 +382,6 @@ func TestClient_SealRevision(t *testing.T) {
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		assert.Equal(t, http.MethodPatch, r.Method)
 		assert.Equal(t, "application/vnd.api+json", r.Header.Get("Content-Type"))
-		assert.Equal(t, "gzip", r.Header.Get("Content-Encoding"))
 		assert.Equal(t, fmt.Sprintf("/hidden/orgs/%s/upload_revisions/%s", orgID, revID), r.URL.Path)
 		assert.Equal(t, "2024-10-15", r.URL.Query().Get("version"))