Merge pull request #45 from snyk/feat/file-upload-gzip-compression

feat: gzip compress file upload requests

Author: paulrosca-snyk

internal/fileupload/lowlevel/client.go

@@ -44,17 +44,20 @@ const FileSizeLimit = 50_000_000 // arbitrary number, chosen to support max size
 // FileCountLimit specifies the maximum number of files allowed in a single upload.
 const FileCountLimit = 100 // arbitrary number, will need to be re-evaluated
 
-// ContentType is the HTTP header name for content type.
-const ContentType = "Content-Type"
-
 // NewClient creates a new file upload client with the given configuration and options.
 func NewClient(cfg Config, opts ...Opt) *HTTPClient {
-	c := HTTPClient{cfg, http.DefaultClient}
+	httpClient := &http.Client{
+		Transport: http.DefaultTransport,
+	}
+	c := HTTPClient{cfg, httpClient}
 
 	for _, opt := range opts {
 		opt(&c)
 	}
 
+	crt := NewCompressionRoundTripper(c.httpClient.Transport)
+	c.httpClient.Transport = crt
+
 	return &c
 }
 
@@ -80,7 +83,7 @@ func (c *HTTPClient) CreateRevision(ctx context.Context, orgID OrgID) (*UploadRe
 	url := fmt.Sprintf("%s/hidden/orgs/%s/upload_revisions?version=%s", c.cfg.BaseURL, orgID, APIVersion)
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, buff)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create request body: %w", err)
+		return nil, fmt.Errorf("failed to create revision request: %w", err)
 	}
 	req.Header.Set(ContentType, "application/vnd.api+json")
 
@@ -91,7 +94,7 @@ func (c *HTTPClient) CreateRevision(ctx context.Context, orgID OrgID) (*UploadRe
 	defer res.Body.Close()
 
 	if res.StatusCode != http.StatusCreated {
-		return nil, c.handleUnexpectedStatusCodes(res.Body, res.StatusCode, res.Status, "create upload revision")
+		return nil, handleUnexpectedStatusCodes(res.Body, res.StatusCode, res.Status, "create upload revision")
 	}
 
 	var respBody UploadRevisionResponseBody
@@ -112,44 +115,23 @@ func (c *HTTPClient) UploadFiles(ctx context.Context, orgID OrgID, revisionID Re
 		return ErrEmptyRevisionID
 	}
 
-	if len(files) > FileCountLimit {
-		return NewFileCountLimitError(len(files), FileCountLimit)
+	if err := validateFiles(files); err != nil {
+		return err
 	}
 
-	if len(files) == 0 {
-		return ErrNoFilesProvided
-	}
-
-	for _, file := range files {
-		fileInfo, err := file.File.Stat()
-		if err != nil {
-			return NewFileAccessError(file.Path, err)
-		}
-
-		if fileInfo.IsDir() {
-			return NewDirectoryError(file.Path)
-		}
+	// Create pipe for multipart data
+	pipeReader, pipeWriter := io.Pipe()
+	defer pipeReader.Close()
 
-		if fileInfo.Size() > FileSizeLimit {
-			return NewFileSizeLimitError(file.Path, fileInfo.Size(), FileSizeLimit)
-		}
-	}
-
-	// Create pipe for streaming multipart data
-	pReader, pWriter := io.Pipe()
-	mpartWriter := multipart.NewWriter(pWriter)
+	mpartWriter := multipart.NewWriter(pipeWriter)
 
-	// Start goroutine to write multipart data
-	go c.streamFilesToPipe(pWriter, mpartWriter, files)
+	go streamFilesToPipe(pipeWriter, mpartWriter, files)
 
-	// Create HTTP request with streaming body
 	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, pReader)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, pipeReader)
 	if err != nil {
-		pReader.Close()
 		return fmt.Errorf("failed to create upload files request: %w", err)
 	}
-
 	req.Header.Set(ContentType, mpartWriter.FormDataContentType())
 
 	res, err := c.httpClient.Do(req)
@@ -159,18 +141,21 @@ func (c *HTTPClient) UploadFiles(ctx context.Context, orgID OrgID, revisionID Re
 	defer res.Body.Close()
 
 	if res.StatusCode != http.StatusNoContent {
-		return c.handleUnexpectedStatusCodes(res.Body, res.StatusCode, res.Status, "upload files")
+		return handleUnexpectedStatusCodes(res.Body, res.StatusCode, res.Status, "upload files")
 	}
 
 	return nil
 }
 
-func (c *HTTPClient) streamFilesToPipe(pWriter *io.PipeWriter, mpartWriter *multipart.Writer, files []UploadFile) {
+// streamFilesToPipe writes files to the multipart form.
+func streamFilesToPipe(pipeWriter *io.PipeWriter, mpartWriter *multipart.Writer, files []UploadFile) {
 	var streamError error
 	defer func() {
-		pWriter.CloseWithError(streamError)
+		if closeErr := mpartWriter.Close(); closeErr != nil && streamError == nil {
+			streamError = closeErr
+		}
+		pipeWriter.CloseWithError(streamError)
 	}()
-	defer mpartWriter.Close()
 
 	for _, file := range files {
 		// Create form file part
@@ -180,14 +165,41 @@ func (c *HTTPClient) streamFilesToPipe(pWriter *io.PipeWriter, mpartWriter *mult
 			return
 		}
 
-		_, err = io.Copy(part, file.File)
-		if err != nil {
+		if _, err := io.Copy(part, file.File); err != nil {
 			streamError = fmt.Errorf("failed to copy file content for %s: %w", file.Path, err)
 			return
 		}
 	}
 }
 
+// validateFiles validates the files before upload.
+func validateFiles(files []UploadFile) error {
+	if len(files) > FileCountLimit {
+		return NewFileCountLimitError(len(files), FileCountLimit)
+	}
+
+	if len(files) == 0 {
+		return ErrNoFilesProvided
+	}
+
+	for _, file := range files {
+		fileInfo, err := file.File.Stat()
+		if err != nil {
+			return NewFileAccessError(file.Path, err)
+		}
+
+		if fileInfo.IsDir() {
+			return NewDirectoryError(file.Path)
+		}
+
+		if fileInfo.Size() > FileSizeLimit {
+			return NewFileSizeLimitError(file.Path, fileInfo.Size(), FileSizeLimit)
+		}
+	}
+
+	return nil
+}
+
 // SealRevision seals the specified upload revision, marking it as complete.
 func (c *HTTPClient) SealRevision(ctx context.Context, orgID OrgID, revisionID RevisionID) (*SealUploadRevisionResponseBody, error) {
 	if orgID == uuid.Nil {
@@ -215,7 +227,7 @@ func (c *HTTPClient) SealRevision(ctx context.Context, orgID OrgID, revisionID R
 	url := fmt.Sprintf("%s/hidden/orgs/%s/upload_revisions/%s?version=%s", c.cfg.BaseURL, orgID, revisionID, APIVersion)
 	req, err := http.NewRequestWithContext(ctx, http.MethodPatch, url, buff)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create request body: %w", err)
+		return nil, fmt.Errorf("failed to create seal request: %w", err)
 	}
 	req.Header.Set(ContentType, "application/vnd.api+json")
 
@@ -226,7 +238,7 @@ func (c *HTTPClient) SealRevision(ctx context.Context, orgID OrgID, revisionID R
 	defer res.Body.Close()
 
 	if res.StatusCode != http.StatusOK {
-		return nil, c.handleUnexpectedStatusCodes(res.Body, res.StatusCode, res.Status, "seal upload revision")
+		return nil, handleUnexpectedStatusCodes(res.Body, res.StatusCode, res.Status, "seal upload revision")
 	}
 
 	var respBody SealUploadRevisionResponseBody
@@ -237,7 +249,7 @@ func (c *HTTPClient) SealRevision(ctx context.Context, orgID OrgID, revisionID R
 	return &respBody, nil
 }
 
-func (c *HTTPClient) handleUnexpectedStatusCodes(body io.ReadCloser, statusCode int, status, operation string) error {
+func handleUnexpectedStatusCodes(body io.ReadCloser, statusCode int, status, operation string) error {
 	bts, err := io.ReadAll(body)
 	if err != nil {
 		return fmt.Errorf("failed to read response body: %w", err)

internal/fileupload/lowlevel/client_test.go

@@ -1,6 +1,7 @@
 package lowlevel_fileupload_test
 
 import (
+	"compress/gzip"
 	"context"
 	"errors"
 	"fmt"
@@ -25,6 +26,7 @@ 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"))
 
@@ -88,6 +90,7 @@ func TestClient_UploadFiles(t *testing.T) {
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		assert.Equal(t, http.MethodPost, r.Method)
 		assert.Contains(t, r.Header.Get("Content-Type"), "multipart/form-data")
+		assert.Equal(t, "gzip", r.Header.Get("Content-Encoding"))
 		assert.Equal(t, fmt.Sprintf("/hidden/orgs/%s/upload_revisions/%s/files", orgID, revID), r.URL.Path)
 		assert.Equal(t, "2024-10-15", r.URL.Query().Get("version"))
 
@@ -98,7 +101,9 @@ func TestClient_UploadFiles(t *testing.T) {
 		boundary := params["boundary"]
 		require.NotEmpty(t, boundary, "multipart boundary should be present")
 
-		reader := multipart.NewReader(r.Body, boundary)
+		gzipReader, err := gzip.NewReader(r.Body)
+		require.NoError(t, err)
+		reader := multipart.NewReader(gzipReader, boundary)
 
 		// Read the first (and should be only) part
 		part, err := reader.NextPart()
@@ -152,6 +157,7 @@ func TestClient_UploadFiles_MultipleFiles(t *testing.T) {
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		assert.Equal(t, http.MethodPost, r.Method)
 		assert.Contains(t, r.Header.Get("Content-Type"), "multipart/form-data")
+		assert.Equal(t, "gzip", r.Header.Get("Content-Encoding"))
 		assert.Equal(t, fmt.Sprintf("/hidden/orgs/%s/upload_revisions/%s/files", orgID, revID), r.URL.Path)
 		assert.Equal(t, "2024-10-15", r.URL.Query().Get("version"))
 
@@ -162,7 +168,9 @@ func TestClient_UploadFiles_MultipleFiles(t *testing.T) {
 		boundary := params["boundary"]
 		require.NotEmpty(t, boundary, "multipart boundary should be present")
 
-		reader := multipart.NewReader(r.Body, boundary)
+		gzipReader, err := gzip.NewReader(r.Body)
+		require.NoError(t, err)
+		reader := multipart.NewReader(gzipReader, boundary)
 		filesReceived := make(map[string]string)
 
 		// Read all parts
@@ -287,7 +295,7 @@ func TestClient_UploadFiles_FileSizeLimit(t *testing.T) {
 	assert.Error(t, err)
 	var fileSizeErr *lowlevel_fileupload.FileSizeLimitError
 	assert.ErrorAs(t, err, &fileSizeErr)
-	assert.Equal(t, "large_file.txt", fileSizeErr.FileName)
+	assert.Equal(t, "large_file.txt", fileSizeErr.FilePath)
 	assert.Equal(t, int64(lowlevel_fileupload.FileSizeLimit+1), fileSizeErr.FileSize)
 	assert.Equal(t, int64(lowlevel_fileupload.FileSizeLimit), fileSizeErr.Limit)
 }
@@ -375,6 +383,7 @@ 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"))
 

internal/fileupload/lowlevel/compression.go

@@ -0,0 +1,67 @@
+package lowlevel_fileupload //nolint:revive // underscore naming is intentional for this internal package.
+
+import (
+	"compress/gzip"
+	"io"
+	"net/http"
+)
+
+// CompressionRoundTripper is an http.RoundTripper that automatically compresses
+// request bodies using gzip compression. It wraps another RoundTripper and adds
+// Content-Encoding: gzip header while removing Content-Length to allow proper
+// compression handling.
+type CompressionRoundTripper struct {
+	defaultRoundTripper http.RoundTripper
+}
+
+// NewCompressionRoundTripper creates a new CompressionRoundTripper that wraps
+// the provided RoundTripper. If drt is nil, http.DefaultTransport is used.
+// All HTTP requests with a body will be automatically compressed using gzip.
+func NewCompressionRoundTripper(drt http.RoundTripper) *CompressionRoundTripper {
+	rt := drt
+	if rt == nil {
+		rt = http.DefaultTransport
+	}
+	return &CompressionRoundTripper{rt}
+}
+
+// compressRequestBody wraps the given reader with gzip compression.
+func compressRequestBody(body io.Reader) io.ReadCloser {
+	pipeReader, pipeWriter := io.Pipe()
+
+	go func() {
+		var err error
+		gzWriter := gzip.NewWriter(pipeWriter)
+
+		_, err = io.Copy(gzWriter, body)
+
+		if closeErr := gzWriter.Close(); closeErr != nil && err == nil {
+			err = closeErr
+		}
+		pipeWriter.CloseWithError(err)
+	}()
+
+	return pipeReader
+}
+
+// RoundTrip implements the http.RoundTripper interface. It compresses the request
+// body using gzip if a body is present, sets the Content-Encoding header to "gzip",
+// and removes the Content-Length header to allow Go's HTTP client to calculate
+// the correct length after compression. Requests without a body are passed through
+// unchanged to the wrapped RoundTripper.
+func (crt *CompressionRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	if r.Body == nil || r.Body == http.NoBody {
+		//nolint:wrapcheck // No need to wrap the error here.
+		return crt.defaultRoundTripper.RoundTrip(r)
+	}
+
+	compressedBody := compressRequestBody(r.Body)
+
+	r.Body = compressedBody
+	r.Header.Set(ContentEncoding, "gzip")
+	r.Header.Del("Content-Length")
+	r.ContentLength = -1 // Let Go calculate the length
+
+	//nolint:wrapcheck // No need to wrap the error here.
+	return crt.defaultRoundTripper.RoundTrip(r)
+}