From 50200d7c3a9484ae70658de3db155bef0dbfe137 Mon Sep 17 00:00:00 2001 From: louisgls <11679637+louisgls@users.noreply.github.com> Date: Wed, 22 Apr 2026 15:43:56 +0200 Subject: [PATCH] fix: reduce object body memory copies --- s3proxy/internal/router/body.go | 27 +++++++++++++++++++ s3proxy/internal/router/handler.go | 2 +- s3proxy/internal/router/object.go | 9 +++++-- s3proxy/internal/router/router_test.go | 24 +++++++++++++++++ s3proxy/internal/s3/s3.go | 34 ++++++++++++++--------- s3proxy/internal/s3/s3_test.go | 37 ++++++++++++++++++++++++++ 6 files changed, 117 insertions(+), 16 deletions(-) create mode 100644 s3proxy/internal/router/body.go create mode 100644 s3proxy/internal/s3/s3_test.go diff --git a/s3proxy/internal/router/body.go b/s3proxy/internal/router/body.go new file mode 100644 index 0000000..d61eb7e --- /dev/null +++ b/s3proxy/internal/router/body.go @@ -0,0 +1,27 @@ +/* +Copyright (c) Intrinsec 2026 + +SPDX-License-Identifier: AGPL-3.0-only +*/ + +package router + +import ( + "fmt" + "io" +) + +func readBody(body io.Reader, contentLength int64) ([]byte, error) { + if contentLength <= 0 { + return io.ReadAll(body) + } + if contentLength > int64(int(^uint(0)>>1)) { + return nil, fmt.Errorf("content length %d exceeds maximum supported size", contentLength) + } + + bodyBytes := make([]byte, int(contentLength)) + if _, err := io.ReadFull(body, bodyBytes); err != nil { + return nil, err + } + return bodyBytes, nil +} diff --git a/s3proxy/internal/router/handler.go b/s3proxy/internal/router/handler.go index 682908b..2cafc84 100644 --- a/s3proxy/internal/router/handler.go +++ b/s3proxy/internal/router/handler.go @@ -78,7 +78,7 @@ func handlePutObject(client *s3.Client, key string, bucket string, log *logger.L return } - body, err := io.ReadAll(req.Body) + body, err := readBody(req.Body, req.ContentLength) if err != nil { log.WithField("error", err).Error("PutObject reading body") http.Error(w, "failed to read request body", http.StatusInternalServerError) diff --git a/s3proxy/internal/router/object.go b/s3proxy/internal/router/object.go index 47af0fc..2d0ab13 100644 --- a/s3proxy/internal/router/object.go +++ b/s3proxy/internal/router/object.go @@ -12,7 +12,6 @@ import ( "encoding/hex" "errors" "fmt" - "io" "net/http" "net/url" "regexp" @@ -81,7 +80,11 @@ func (o object) get(w http.ResponseWriter, r *http.Request) { } }() - body, err := io.ReadAll(output.Body) + contentLength := int64(-1) + if output.ContentLength != nil { + contentLength = *output.ContentLength + } + body, err := readBody(output.Body, contentLength) if err != nil { o.log.WithField("requestID", requestID).WithField("error", err).Error("GetObject reading S3 response") http.Error(w, fmt.Sprintf("failed to read response: %v", err), http.StatusInternalServerError) @@ -99,6 +102,7 @@ func (o object) get(w http.ResponseWriter, r *http.Request) { } plaintext, err = crypto.Decrypt(body, encryptedDEK, o.kek) + body = nil if err != nil { o.log.WithField("requestID", requestID).WithField("error", err).Error("GetObject decrypting response") http.Error(w, "failed to decrypt object", http.StatusInternalServerError) @@ -133,6 +137,7 @@ func (o object) put(w http.ResponseWriter, r *http.Request) { http.Error(w, err.Error(), http.StatusInternalServerError) return } + o.data = nil o.metadata[dekTag] = hex.EncodeToString(encryptedDEK) output, err := o.client.PutObject(context.WithoutCancel(r.Context()), o.bucket, o.key, o.tags, o.contentType, o.objectLockLegalHoldStatus, o.objectLockMode, o.sseCustomerAlgorithm, o.sseCustomerKey, o.sseCustomerKeyMD5, o.objectLockRetainUntilDate, o.metadata, ciphertext) diff --git a/s3proxy/internal/router/router_test.go b/s3proxy/internal/router/router_test.go index a690ce6..f5d150d 100644 --- a/s3proxy/internal/router/router_test.go +++ b/s3proxy/internal/router/router_test.go @@ -6,9 +6,12 @@ SPDX-License-Identifier: AGPL-3.0-only package router import ( + "io" + "strings" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestValidateContentMD5(t *testing.T) { @@ -85,3 +88,24 @@ func TestByteSliceToByteArray(t *testing.T) { }) } } + +func TestReadBodyUsesKnownContentLength(t *testing.T) { + body, err := readBody(strings.NewReader("hello"), 5) + + require.NoError(t, err) + assert.Equal(t, []byte("hello"), body) + assert.Equal(t, 5, cap(body)) +} + +func TestReadBodyFallsBackWhenContentLengthUnknown(t *testing.T) { + body, err := readBody(strings.NewReader("hello"), -1) + + require.NoError(t, err) + assert.Equal(t, []byte("hello"), body) +} + +func TestReadBodyReturnsErrorOnShortBody(t *testing.T) { + _, err := readBody(strings.NewReader("hi"), 5) + + assert.ErrorIs(t, err, io.ErrUnexpectedEOF) +} diff --git a/s3proxy/internal/s3/s3.go b/s3proxy/internal/s3/s3.go index f3fb35a..a5d23c2 100644 --- a/s3proxy/internal/s3/s3.go +++ b/s3proxy/internal/s3/s3.go @@ -18,6 +18,7 @@ import ( "encoding/base64" "fmt" "io" + "net/http" "time" "github.com/aws/aws-sdk-go-v2/aws" @@ -51,7 +52,22 @@ func (m *ErrorRawResponse) Error() string { return m.RawResponse } -// Middleware to capture the raw response in the Send phase by cloning and storing the response body +func readBody(body io.Reader, contentLength int64) ([]byte, error) { + if contentLength <= 0 { + return io.ReadAll(body) + } + if contentLength > int64(int(^uint(0)>>1)) { + return nil, fmt.Errorf("content length %d exceeds maximum supported size", contentLength) + } + + bodyBytes := make([]byte, int(contentLength)) + if _, err := io.ReadFull(body, bodyBytes); err != nil { + return nil, err + } + return bodyBytes, nil +} + +// Middleware to capture error response bodies without cloning successful object bodies. func addCaptureRawResponseDeserializeMiddleware(log *logger.Logger) func(*middleware.Stack) error { return func(stack *middleware.Stack) error { return stack.Deserialize.Add(middleware.DeserializeMiddlewareFunc("CaptureRawResponseDeserialize", func( @@ -61,25 +77,17 @@ func addCaptureRawResponseDeserializeMiddleware(log *logger.Logger) func(*middle ) { out, metadata, err = next.HandleDeserialize(ctx, in) if resp, ok := out.RawResponse.(*smithyhttp.Response); ok { - // Clone the response body - var buf bytes.Buffer - body := resp.Body - tee := io.NopCloser(io.TeeReader(body, &buf)) - - // Replace the body in the response with the cloned body - resp.Body = tee + if resp.StatusCode < http.StatusBadRequest { + return out, metadata, err + } - bodyBytes, err := io.ReadAll(resp.Body) + bodyBytes, err := readBody(resp.Body, resp.ContentLength) if err != nil { log.WithError(err).Error("failed to read response body") - // Return the error to prevent silent failures return out, metadata, fmt.Errorf("reading response body: %w", err) } - // Store the cloned body in metadata metadata.Set(RawResponseKey{}, string(bodyBytes)) - - // Restore the original body for further processing resp.Body = io.NopCloser(bytes.NewReader(bodyBytes)) } return out, metadata, err diff --git a/s3proxy/internal/s3/s3_test.go b/s3proxy/internal/s3/s3_test.go new file mode 100644 index 0000000..27fc203 --- /dev/null +++ b/s3proxy/internal/s3/s3_test.go @@ -0,0 +1,37 @@ +/* +Copyright (c) Intrinsec 2026 + +SPDX-License-Identifier: AGPL-3.0-only +*/ + +package s3 + +import ( + "io" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestReadBodyUsesKnownContentLength(t *testing.T) { + body, err := readBody(strings.NewReader("hello"), 5) + + require.NoError(t, err) + assert.Equal(t, []byte("hello"), body) + assert.Equal(t, 5, cap(body)) +} + +func TestReadBodyFallsBackWhenContentLengthUnknown(t *testing.T) { + body, err := readBody(strings.NewReader("hello"), -1) + + require.NoError(t, err) + assert.Equal(t, []byte("hello"), body) +} + +func TestReadBodyReturnsErrorOnShortBody(t *testing.T) { + _, err := readBody(strings.NewReader("hi"), 5) + + assert.ErrorIs(t, err, io.ErrUnexpectedEOF) +}