fix(uploadstore): bound R2 Serve body read

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 {

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 {