From 3459131e309bcdf1a64429dd4de7dc48494db7d0 Mon Sep 17 00:00:00 2001 From: Sebastien Tardif Date: Sat, 29 Aug 2026 11:55:54 -0700 Subject: [PATCH] fix(uploadstore): bound R2 Serve body read R2 ServeHTTP copies the object body with Client.Timeout 0. Response headers are bounded, but a stalled body after headers left the download handler blocked. Wrap the Serve response body with an idle read deadline so a hung origin unblocks the request without changing the streaming-safe client timeout. Signed-off-by: Sebastien Tardif --- apps/api/internal/uploadstore/r2.go | 77 ++++++++++++++++++------ apps/api/internal/uploadstore/r2_test.go | 62 +++++++++++++++++++ 2 files changed, 122 insertions(+), 17 deletions(-) diff --git a/apps/api/internal/uploadstore/r2.go b/apps/api/internal/uploadstore/r2.go index fe6629a01..90a3f2a26 100644 --- a/apps/api/internal/uploadstore/r2.go +++ b/apps/api/internal/uploadstore/r2.go @@ -15,10 +15,17 @@ import ( "os" "sort" "strings" + "sync/atomic" "time" ) -const defaultR2ResponseHeaderTimeout = 30 * time.Second +const ( + defaultR2ResponseHeaderTimeout = 30 * time.Second + // Idle read bound for ServeHTTP. Client.Timeout stays 0 for streaming. + defaultR2ServeBodyIdleTimeout = 30 * time.Second +) + +var errServeBodyStalled = errors.New("r2 serve: body read stalled") type R2Config struct { AccountID string @@ -32,14 +39,15 @@ type R2Config struct { } type R2 struct { - accountID string - accessKeyID string - secretAccessKey string - bucket string - prefix string - endpoint string - region string - httpClient *http.Client + accountID string + accessKeyID string + secretAccessKey string + bucket string + prefix string + endpoint string + region string + httpClient *http.Client + serveBodyIdleTimeout time.Duration } func NewR2(cfg R2Config) (*R2, error) { @@ -72,14 +80,15 @@ func NewR2(cfg R2Config) (*R2, error) { prefix += "/" } return &R2{ - accountID: cfg.AccountID, - accessKeyID: cfg.AccessKeyID, - secretAccessKey: cfg.SecretAccessKey, - bucket: cfg.Bucket, - prefix: prefix, - endpoint: endpoint, - region: region, - httpClient: client, + accountID: cfg.AccountID, + accessKeyID: cfg.AccessKeyID, + secretAccessKey: cfg.SecretAccessKey, + bucket: cfg.Bucket, + prefix: prefix, + endpoint: endpoint, + region: region, + httpClient: client, + serveBodyIdleTimeout: defaultR2ServeBodyIdleTimeout, }, nil } @@ -180,6 +189,7 @@ func (s *R2) ServeHTTP(w http.ResponseWriter, r *http.Request, object Object) er if err != nil { return err } + resp.Body = wrapServeBody(resp.Body, s.serveBodyIdleTimeout) defer resp.Body.Close() if resp.StatusCode == http.StatusNotFound { return ErrNotFound @@ -368,6 +378,39 @@ func copyHeader(dst, src http.Header, name string) { } } +func wrapServeBody(body io.ReadCloser, timeout time.Duration) io.ReadCloser { + if timeout <= 0 { + timeout = defaultR2ServeBodyIdleTimeout + } + return &idleTimeoutReadCloser{rc: body, timeout: timeout} +} + +type idleTimeoutReadCloser struct { + rc io.ReadCloser + timeout time.Duration + timedOut atomic.Bool +} + +func (r *idleTimeoutReadCloser) Read(p []byte) (int, error) { + timer := time.AfterFunc(r.timeout, func() { + r.timedOut.Store(true) + _ = r.rc.Close() + }) + defer timer.Stop() + n, err := r.rc.Read(p) + if n > 0 { + return n, err + } + if r.timedOut.Load() { + return 0, errServeBodyStalled + } + return 0, err +} + +func (r *idleTimeoutReadCloser) Close() error { + return r.rc.Close() +} + func responseError(prefix string, resp *http.Response) error { body, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) if len(body) > 0 { diff --git a/apps/api/internal/uploadstore/r2_test.go b/apps/api/internal/uploadstore/r2_test.go index 1ce852d47..c2020da31 100644 --- a/apps/api/internal/uploadstore/r2_test.go +++ b/apps/api/internal/uploadstore/r2_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "strings" "testing" + "time" ) func TestR2SaveServeAndDelete(t *testing.T) { @@ -103,6 +104,67 @@ func TestR2SaveServeAndDelete(t *testing.T) { } } +func TestR2ServeHTTPStalledBodyDoesNotHang(t *testing.T) { + t.Parallel() + headersSent := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("Content-Length", "100") + w.WriteHeader(http.StatusOK) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + close(headersSent) + <-r.Context().Done() + })) + t.Cleanup(func() { + server.CloseClientConnections() + server.Close() + }) + + store, err := NewR2(R2Config{ + AccountID: "account", + AccessKeyID: "access", + SecretAccessKey: "secret", + Bucket: "bucket", + Prefix: "prefix", + Endpoint: server.URL, + }) + if err != nil { + t.Fatal(err) + } + store.serveBodyIdleTimeout = 50 * time.Millisecond + + done := make(chan error, 1) + go func() { + req := httptest.NewRequest(http.MethodGet, "/api/uploads/upl_stalled", nil) + done <- store.ServeHTTP(httptest.NewRecorder(), req, Object{ + Path: "r2://bucket/prefix/upload-stalled", + ContentType: "application/octet-stream", + ByteSize: 100, + }) + }() + + select { + case <-headersSent: + case <-time.After(2 * time.Second): + t.Fatal("r2 handler never sent response headers") + } + + select { + case err := <-done: + if !errors.Is(err, errServeBodyStalled) { + t.Fatalf("expected stalled-body error, got %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("ServeHTTP hung on stalled R2 body") + } +} + func TestR2ConfigValidation(t *testing.T) { t.Parallel() if _, err := NewR2(R2Config{AccessKeyID: "access", SecretAccessKey: "secret", Bucket: "bucket"}); err == nil {