feat(gateway): add configurable response write timeout

examples/gateway/common/handler.go

@@ -2,6 +2,7 @@ package common
 
 import (
 	"net/http"
+	"time"
 
 	"github.com/ipfs/boxo/gateway"
 	"github.com/ipfs/boxo/gateway/assets"
@@ -52,6 +53,8 @@ func NewHandler(gwAPI gateway.IPFSBackend) http.Handler {
 				Title: "Boxo",
 			},
 		},
+
+		ResponseWriteTimeout: 30 * time.Second, // Set default timeout to 30 seconds
 	}
 
 	// Creates a mux to serve the gateway paths. This is not strictly necessary

gateway/gateway.go

@@ -48,6 +48,11 @@ type Config struct {
 	// directory listings, DAG previews and errors. These will be displayed to the
 	// right of "About IPFS" and "Install IPFS".
 	Menu []assets.MenuItem
+
+	// ResponseWriteTimeout is the maximum duration the gateway will wait for a response
+	// to be written before timing out and returning a 504 Gateway Timeout error.
+	// If not set, a default timeout will be used.
+	ResponseWriteTimeout time.Duration
 }
 
 // PublicGateway is the specification of an IPFS Public Gateway.

gateway/handler.go

@@ -13,6 +13,7 @@
 	"regexp"
 	"runtime/debug"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/ipfs/boxo/gateway/assets"
@@ -68,7 +69,83 @@
 //
 // [IPFS HTTP Gateway]: https://specs.ipfs.tech/http-gateways/
 func NewHandler(c Config, backend IPFSBackend) http.Handler {
-	return newHandlerWithMetrics(&c, backend)
+	handler := newHandlerWithMetrics(&c, backend)
+
+	// Use the configured timeout or fall back to a default value
+	timeout := c.ResponseWriteTimeout
+	if timeout == 0 {
+		timeout = 30 * time.Second // Default timeout of 30 seconds
+	}
+
+	// Apply the timeout middleware
+	return WithResponseWriteTimeout(handler, timeout)
+}
+
+// timeoutResponseWriter implements http.ResponseWriter with timeout control
+type timeoutResponseWriter struct {
+	http.ResponseWriter                    // Embedded standard response writer
+	timer               *time.Timer        // Timeout tracking mechanism
+	timeout             time.Duration      // Configured timeout duration
+	mu                  sync.Mutex         // Mutex for concurrent access protection
+	done                bool               // Completion state flag
+	requestCtx          context.Context    // Original request context
+	cancel              context.CancelFunc // Context cancellation handler
+}
+
+// Write implements io.Writer interface with timeout reset functionality
+func (w *timeoutResponseWriter) Write(data []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if w.done {
+		return 0, http.ErrHandlerTimeout
+	}
+
+	// Reset timer on successful write attempt
+	w.timer.Reset(w.timeout)
+	return w.ResponseWriter.Write(data)
+}
+
+// cancelRequest handles timeout termination sequence
+func (w *timeoutResponseWriter) cancelRequest() {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if !w.done {
+		w.done = true
+		w.cancel() // Propagate context cancellation
+		w.ResponseWriter.WriteHeader(http.StatusGatewayTimeout)
+	}
+}
+
+// WithResponseWriteTimeout creates middleware for response write timeout handling
+func WithResponseWriteTimeout(next http.Handler, timeout time.Duration) http.Handler {
+	return http.HandlerFunc(func(origWriter http.ResponseWriter, r *http.Request) {
+		// Create derived context with cancellation capability
+		ctx, cancel := context.WithCancel(r.Context())
+
+		// Initialize enhanced response writer
+		tw := &timeoutResponseWriter{
+			ResponseWriter: origWriter,
+			timeout:        timeout,
+			timer:          time.NewTimer(timeout),
+			requestCtx:     ctx,
+			cancel:         cancel,
+		}
+		defer tw.timer.Stop() // Ensure timer resource cleanup
+
+		// Concurrent timeout monitor
+		go func() {
+			select {
+			case <-tw.timer.C: // Timeout expiration
+				tw.cancelRequest()
+			case <-ctx.Done(): // Normal completion
+			}
+		}()
+
+		// Execute handler chain with wrapped context
+		next.ServeHTTP(tw, r.WithContext(ctx))
+	})
 }
 
 // serveContent replies to the request using the content in the provided Reader

gateway/handler_test.go

@@ -1,7 +1,11 @@
 package gateway
 
 import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -28,3 +32,90 @@
 		assert.Equalf(t, test.expected, result, "etagMatch(%q, %q, %q)", test.header, test.cidEtag, test.dirEtag)
 	}
 }
+
+
+
+func TestWithResponseWriteTimeout(t *testing.T) {
+	tests := []struct {
+		name           string
+		handler        http.HandlerFunc // Test scenario handler implementation
+		timeout        time.Duration    // Timeout configuration for middleware
+		expectStatus   int              // Anticipated HTTP response code
+		expectedChunks int              // Expected number of completed writes
+	}{
+		{
+			name: "Normal completion - write within timeout",
+			handler: func(w http.ResponseWriter, r *http.Request) {
+				select {
+				case <-time.After(100 * time.Millisecond):
+					w.Write([]byte("chunk\n"))
+				case <-r.Context().Done():
+				}
+			},
+			timeout:        300 * time.Millisecond,
+			expectStatus:   http.StatusOK,
+			expectedChunks: 1,
+		},
+		{
+			name: "Timeout triggered - write after deadline",
+			handler: func(w http.ResponseWriter, r *http.Request) {
+				select {
+				case <-time.After(400 * time.Millisecond): // Exceeds 300ms timeout
+					w.Write([]byte("chunk\n")) // Should be post-timeout
+				case <-r.Context().Done():
+				}
+			},
+			timeout:        300 * time.Millisecond,
+			expectStatus:   http.StatusGatewayTimeout,
+			expectedChunks: 0,
+		},
+		{
+			name: "Timer reset with staggered writes",
+			handler: func(w http.ResponseWriter, r *http.Request) {
+				for i := 0; i < 3; i++ {
+					select {
+					case <-time.After(200 * time.Millisecond): // Each write within timeout window
+						w.Write([]byte("chunk\n")) // Resets timer on each write
+					case <-r.Context().Done():
+						return
+					}
+				}
+			},
+			timeout:        300 * time.Millisecond,
+			expectStatus:   http.StatusOK,
+			expectedChunks: 3,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Initialize test components
+			req := httptest.NewRequest("GET", "http://example.com", nil)
+			rec := httptest.NewRecorder()
+
+			// Configure middleware chain
+			wrappedHandler := WithResponseWriteTimeout(tt.handler, tt.timeout)
+			wrappedHandler.ServeHTTP(rec, req)
+
+			// Validate HTTP status code
+			if rec.Code != tt.expectStatus {
+				t.Errorf("Status code validation failed: expected %d, received %d",
+					tt.expectStatus, rec.Code)
+			}
+
+			// Verify response body integrity
+			actualChunks := strings.Count(rec.Body.String(), "chunk")
+			if actualChunks != tt.expectedChunks {
+				t.Errorf("Data chunk mismatch: anticipated %d, obtained %d",
+					tt.expectedChunks, actualChunks)
+			}
+
+			// Ensure timeout responses don't set headers
+			if tt.expectStatus == http.StatusGatewayTimeout {
+				if contentType := rec.Header().Get("Content-Type"); contentType != "" {
+					t.Errorf("Timeout response contains unexpected Content-Type: %s", contentType)
+				}
+			}
+		})
+	}
+}