Skip to content

Commit 38412eb

Browse files
authored
fix(storage): proxy S3 attachments through authenticated routes (#6210)
1 parent c9324f9 commit 38412eb

76 files changed

Lines changed: 527 additions & 624 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

AGENTS.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,12 @@ cd proto && buf format -w # Format proto files
5656
| Path | Purpose |
5757
| --- | --- |
5858
| `cmd/memos/main.go` | Cobra/Viper CLI setup and server startup |
59-
| `server/server.go` | Echo HTTP server and background runner wiring |
59+
| `server/server.go` | Echo HTTP server, router wiring, and graceful shutdown |
6060
| `server/auth/` | JWT access tokens, refresh tokens, PAT handling |
6161
| `server/router/api/v1/` | Connect/gRPC-Gateway services, ACL config, SSE hub |
6262
| `server/router/frontend/` | Static SPA serving |
6363
| `server/router/fileserver/` | Native HTTP file serving, thumbnails, range requests |
64-
| `server/runner/` | Background memo processing and S3 presign refresh |
64+
| `server/runner/` | Memo payload rebuilding |
6565
| `store/` | Store facade, cache, migrations, driver interface |
6666
| `store/db/{sqlite,mysql,postgres}/` | Database-specific drivers and SQL |
6767
| `proto/api/v1/` | Public API service definitions |

cmd/memos/main.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ var (
8686
// which is taken as the graceful shutdown signal for many systems, eg., Kubernetes, Gunicorn.
8787
signal.Notify(c, os.Interrupt, syscall.SIGTERM)
8888

89-
if err := s.Start(ctx); err != nil {
89+
if err := s.Start(); err != nil {
9090
if err != http.ErrServerClosed {
9191
slog.Error("failed to start server", "error", err)
9292
cancel()

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ require (
145145
github.com/aws/aws-sdk-go-v2/service/sso v1.33.6 // indirect
146146
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.6 // indirect
147147
github.com/aws/aws-sdk-go-v2/service/sts v1.45.6 // indirect
148-
github.com/aws/smithy-go v1.27.8 // indirect
148+
github.com/aws/smithy-go v1.27.8
149149
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
150150
github.com/disintegration/imaging v1.6.2
151151
github.com/golang-jwt/jwt/v5 v5.3.1

internal/storage/driver.go

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,25 @@ import (
1010
storepb "github.com/usememos/memos/proto/gen/store"
1111
)
1212

13+
// ErrRangeNotSatisfiable reports a ranged read whose byte range falls outside
14+
// the object, so HTTP handlers can answer 416 instead of 500.
15+
var ErrRangeNotSatisfiable = s3.ErrRangeNotSatisfiable
16+
17+
// RangeNotSatisfiableError carries response metadata for an unsatisfied range.
18+
type RangeNotSatisfiableError = s3.RangeNotSatisfiableError
19+
20+
// ObjectStream is object content with the metadata needed to answer HTTP
21+
// range requests. Aliased so driver consumers never import a concrete backend.
22+
type ObjectStream = s3.ObjectStream
23+
1324
// Driver provides object operations for a configured attachment storage.
1425
type Driver interface {
1526
UploadObject(ctx context.Context, key string, fileType string, content io.Reader) (string, error)
16-
PresignGetObject(ctx context.Context, key string) (string, error)
1727
GetObject(ctx context.Context, key string) ([]byte, error)
18-
GetObjectStream(ctx context.Context, key string) (io.ReadCloser, error)
28+
// GetObjectStream streams an object; a non-empty byteRange is forwarded as
29+
// an HTTP Range header and yields a partial object. byteRange must request at
30+
// most one range because S3 does not support multipart range responses.
31+
GetObjectStream(ctx context.Context, key string, byteRange string) (*ObjectStream, error)
1932
DeleteObject(ctx context.Context, key string) error
2033
}
2134

internal/storage/s3/s3.go

Lines changed: 129 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,52 @@ package s3
33
import (
44
"context"
55
"crypto/tls"
6+
"fmt"
67
"io"
78
"net/http"
8-
"time"
99

1010
"github.com/aws/aws-sdk-go-v2/aws"
11+
v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" //nolint:revive // goimports insists on aliasing versioned import paths
1112
awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
1213
"github.com/aws/aws-sdk-go-v2/config"
1314
"github.com/aws/aws-sdk-go-v2/credentials"
1415
"github.com/aws/aws-sdk-go-v2/service/s3"
16+
"github.com/aws/smithy-go"
17+
"github.com/aws/smithy-go/middleware"
18+
smithyhttp "github.com/aws/smithy-go/transport/http"
1519
"github.com/pkg/errors"
1620

1721
storepb "github.com/usememos/memos/proto/gen/store"
1822
)
1923

24+
// ErrRangeNotSatisfiable reports a ranged read whose byte range falls outside
25+
// the object, so HTTP handlers can answer 416 instead of 500.
26+
var ErrRangeNotSatisfiable = errors.New("requested range not satisfiable")
27+
28+
// RangeNotSatisfiableError carries response metadata for an unsatisfied range.
29+
type RangeNotSatisfiableError struct {
30+
ContentRange string
31+
}
32+
33+
func (*RangeNotSatisfiableError) Error() string {
34+
return ErrRangeNotSatisfiable.Error()
35+
}
36+
37+
func (*RangeNotSatisfiableError) Unwrap() error {
38+
return ErrRangeNotSatisfiable
39+
}
40+
41+
// ObjectStream is object content with the metadata needed to answer HTTP
42+
// range requests.
43+
type ObjectStream struct {
44+
Body io.ReadCloser
45+
// ContentLength is the number of bytes in Body, or -1 when unknown.
46+
ContentLength int64
47+
// ContentRange echoes the backend's Content-Range header for partial reads;
48+
// empty when the whole object is returned.
49+
ContentRange string
50+
}
51+
2052
// Driver stores attachment objects in an S3-compatible object store.
2153
type Driver struct {
2254
Client *s3.Client
@@ -48,13 +80,65 @@ func NewDriver(ctx context.Context, s3Config *storepb.StorageS3Config) (*Driver,
4880
o.UsePathStyle = s3Config.UsePathStyle
4981
o.RequestChecksumCalculation = aws.RequestChecksumCalculationWhenRequired
5082
o.ResponseChecksumValidation = aws.ResponseChecksumValidationWhenRequired
83+
o.APIOptions = append(o.APIOptions, excludeAcceptEncodingFromSigning, forceSignedPayload)
5184
})
5285
return &Driver{
5386
Client: client,
5487
Bucket: aws.String(s3Config.Bucket),
5588
}, nil
5689
}
5790

91+
type acceptEncodingKey struct{}
92+
93+
type dropAcceptEncoding struct{}
94+
95+
func (dropAcceptEncoding) ID() string { return "MemosDropAcceptEncoding" }
96+
97+
func (dropAcceptEncoding) HandleFinalize(ctx context.Context, in middleware.FinalizeInput, next middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
98+
if req, ok := in.Request.(*smithyhttp.Request); ok {
99+
if values := req.Header.Values("Accept-Encoding"); len(values) > 0 {
100+
ctx = context.WithValue(ctx, acceptEncodingKey{}, values)
101+
req.Header.Del("Accept-Encoding")
102+
}
103+
}
104+
return next.HandleFinalize(ctx, in)
105+
}
106+
107+
type restoreAcceptEncoding struct{}
108+
109+
func (restoreAcceptEncoding) ID() string { return "MemosRestoreAcceptEncoding" }
110+
111+
func (restoreAcceptEncoding) HandleFinalize(ctx context.Context, in middleware.FinalizeInput, next middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
112+
if req, ok := in.Request.(*smithyhttp.Request); ok {
113+
if values, ok := ctx.Value(acceptEncodingKey{}).([]string); ok {
114+
for _, value := range values {
115+
req.Header.Add("Accept-Encoding", value)
116+
}
117+
}
118+
}
119+
return next.HandleFinalize(ctx, in)
120+
}
121+
122+
// excludeAcceptEncodingFromSigning keeps the Accept-Encoding header out of the
123+
// SigV4 signature while still sending it on the wire. Some S3-compatible
124+
// providers (notably Google Cloud Storage) rewrite the header in transit, so a
125+
// signature covering it never verifies and requests fail with
126+
// SignatureDoesNotMatch.
127+
func excludeAcceptEncodingFromSigning(stack *middleware.Stack) error {
128+
if err := stack.Finalize.Insert(dropAcceptEncoding{}, "Signing", middleware.Before); err != nil {
129+
return err
130+
}
131+
return stack.Finalize.Insert(restoreAcceptEncoding{}, "Signing", middleware.After)
132+
}
133+
134+
// forceSignedPayload signs request payloads with their real SHA-256 instead of
135+
// the UNSIGNED-PAYLOAD marker the SDK uses over TLS, which some S3-compatible
136+
// providers (notably Google Cloud Storage) reject.
137+
func forceSignedPayload(stack *middleware.Stack) error {
138+
_, err := stack.Finalize.Swap((*v4.ComputePayloadSHA256)(nil).ID(), &v4.ComputePayloadSHA256{})
139+
return err
140+
}
141+
58142
// UploadObject uploads an object to S3.
59143
func (c *Driver) UploadObject(ctx context.Context, key string, fileType string, content io.Reader) (string, error) {
60144
putInput := s3.PutObjectInput{
@@ -69,50 +153,65 @@ func (c *Driver) UploadObject(ctx context.Context, key string, fileType string,
69153
return key, nil
70154
}
71155

72-
// PresignGetObject presigns an object in S3.
73-
func (c *Driver) PresignGetObject(ctx context.Context, key string) (string, error) {
74-
presignClient := s3.NewPresignClient(c.Client)
75-
presignResult, err := presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
76-
Bucket: aws.String(*c.Bucket),
77-
Key: aws.String(key),
78-
}, func(opts *s3.PresignOptions) {
79-
// Set the expiration time of the presigned URL to 5 days.
80-
// Reference: https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-query-string-auth.html
81-
opts.Expires = time.Duration(5 * 24 * time.Hour)
82-
})
83-
if err != nil {
84-
return "", errors.Wrap(err, "failed to presign get object")
85-
}
86-
return presignResult.URL, nil
87-
}
88-
89156
// GetObject retrieves an object from S3.
90157
func (c *Driver) GetObject(ctx context.Context, key string) ([]byte, error) {
91-
output, err := c.Client.GetObject(ctx, &s3.GetObjectInput{
92-
Bucket: c.Bucket,
93-
Key: aws.String(key),
94-
})
158+
stream, err := c.GetObjectStream(ctx, key, "")
95159
if err != nil {
96-
return nil, errors.Wrap(err, "failed to download object")
160+
return nil, err
97161
}
98-
defer output.Body.Close()
99-
data, err := io.ReadAll(output.Body)
162+
defer stream.Body.Close()
163+
data, err := io.ReadAll(stream.Body)
100164
if err != nil {
101165
return nil, errors.Wrap(err, "failed to read object body")
102166
}
103167
return data, nil
104168
}
105169

106-
// GetObjectStream retrieves an object from S3 as a stream.
107-
func (c *Driver) GetObjectStream(ctx context.Context, key string) (io.ReadCloser, error) {
108-
output, err := c.Client.GetObject(ctx, &s3.GetObjectInput{
170+
// GetObjectStream retrieves an object from S3 as a stream. A non-empty
171+
// byteRange is forwarded as an HTTP Range header (e.g. "bytes=0-1023") and
172+
// yields a partial object with ContentRange set. Callers must supply at most
173+
// one range because S3 does not support multipart range responses.
174+
func (c *Driver) GetObjectStream(ctx context.Context, key string, byteRange string) (*ObjectStream, error) {
175+
input := &s3.GetObjectInput{
109176
Bucket: c.Bucket,
110177
Key: aws.String(key),
111-
})
178+
}
179+
if byteRange != "" {
180+
input.Range = aws.String(byteRange)
181+
}
182+
output, err := c.Client.GetObject(ctx, input)
112183
if err != nil {
184+
var apiErr smithy.APIError
185+
if errors.As(err, &apiErr) && apiErr.ErrorCode() == "InvalidRange" {
186+
rangeErr := &RangeNotSatisfiableError{}
187+
var responseErr *smithyhttp.ResponseError
188+
if errors.As(err, &responseErr) && responseErr.Response != nil && responseErr.Response.Response != nil {
189+
rangeErr.ContentRange = responseErr.Response.Header.Get("Content-Range")
190+
}
191+
if rangeErr.ContentRange == "" {
192+
head, headErr := c.Client.HeadObject(ctx, &s3.HeadObjectInput{
193+
Bucket: c.Bucket,
194+
Key: aws.String(key),
195+
})
196+
if headErr == nil && head.ContentLength != nil {
197+
rangeErr.ContentRange = fmt.Sprintf("bytes */%d", *head.ContentLength)
198+
}
199+
}
200+
return nil, rangeErr
201+
}
113202
return nil, errors.Wrap(err, "failed to get object")
114203
}
115-
return output.Body, nil
204+
stream := &ObjectStream{
205+
Body: output.Body,
206+
ContentLength: -1,
207+
}
208+
if output.ContentLength != nil {
209+
stream.ContentLength = *output.ContentLength
210+
}
211+
if output.ContentRange != nil {
212+
stream.ContentRange = *output.ContentRange
213+
}
214+
return stream, nil
116215
}
117216

118217
// DeleteObject deletes an object in S3.

internal/storage/s3/s3_test.go

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package s3
33
import (
44
"bytes"
55
"context"
6+
"fmt"
67
"io"
78
"net/http"
89
"net/http/httptest"
@@ -16,7 +17,7 @@ import (
1617
)
1718

1819
// assertObjectLifecycle exercises the full driver contract — upload, download,
19-
// stream, presign + fetch, delete — against whichever backend the driver targets.
20+
// stream (full and ranged), delete — against whichever backend the driver targets.
2021
func assertObjectLifecycle(ctx context.Context, t *testing.T, driver *Driver, key string, content []byte) {
2122
t.Helper()
2223

@@ -28,25 +29,29 @@ func assertObjectLifecycle(ctx context.Context, t *testing.T, driver *Driver, ke
2829
require.NoError(t, err)
2930
require.Equal(t, content, downloaded)
3031

31-
stream, err := driver.GetObjectStream(ctx, key)
32+
stream, err := driver.GetObjectStream(ctx, key, "")
3233
require.NoError(t, err)
33-
streamed, err := io.ReadAll(stream)
34+
streamed, err := io.ReadAll(stream.Body)
3435
require.NoError(t, err)
35-
require.NoError(t, stream.Close())
36+
require.NoError(t, stream.Body.Close())
3637
require.Equal(t, content, streamed)
38+
require.Equal(t, int64(len(content)), stream.ContentLength)
39+
require.Empty(t, stream.ContentRange)
3740

38-
presignedURL, err := driver.PresignGetObject(ctx, key)
41+
partial, err := driver.GetObjectStream(ctx, key, "bytes=4-9")
3942
require.NoError(t, err)
40-
require.Contains(t, presignedURL, "X-Amz-Signature=")
41-
request, err := http.NewRequestWithContext(ctx, http.MethodGet, presignedURL, nil)
43+
partialContent, err := io.ReadAll(partial.Body)
4244
require.NoError(t, err)
43-
response, err := http.DefaultClient.Do(request)
44-
require.NoError(t, err)
45-
defer response.Body.Close()
46-
require.Equal(t, http.StatusOK, response.StatusCode)
47-
presignedContent, err := io.ReadAll(response.Body)
48-
require.NoError(t, err)
49-
require.Equal(t, content, presignedContent)
45+
require.NoError(t, partial.Body.Close())
46+
require.Equal(t, content[4:10], partialContent)
47+
require.Equal(t, int64(len(partialContent)), partial.ContentLength)
48+
require.Equal(t, fmt.Sprintf("bytes 4-9/%d", len(content)), partial.ContentRange)
49+
50+
_, err = driver.GetObjectStream(ctx, key, fmt.Sprintf("bytes=%d-", len(content)*2))
51+
require.ErrorIs(t, err, ErrRangeNotSatisfiable)
52+
var rangeErr *RangeNotSatisfiableError
53+
require.ErrorAs(t, err, &rangeErr)
54+
require.Equal(t, fmt.Sprintf("bytes */%d", len(content)), rangeErr.ContentRange)
5055

5156
require.NoError(t, driver.DeleteObject(ctx, key))
5257
_, err = driver.GetObject(ctx, key)

0 commit comments

Comments
 (0)