fix: add max bytes request limit to aibridge (#26164)

Adds limit of 32MiB limit to request body in all aibridge endpoints.

---------

Co-authored-by: Danny Kopping <danny@coder.com>

aibridge/bridge.go

@@ -2,6 +2,7 @@ package aibridge
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -32,6 +33,21 @@ const (
 	// The duration after which an async recording will be aborted.
 	recordingTimeout = time.Second * 5
 
+	// maxRequestBodyBytes caps the request body size for AI Bridge
+	// provider endpoints to prevent denial-of-service via memory exhaustion.
+	// Anthropic enforces 32 MB on the direct API, 30 MB on Vertex AI,
+	// and 20 MB on Amazon Bedrock.
+	// See https://docs.anthropic.com/en/api/overview#request-size-limits
+	// OpenAI and GitHub Copilot do not document an equivalent HTTP body size limit.
+	// Using highest documented provider limit (32 MiB).
+	//
+	// NOTE: aibridge does not currently proxy file-upload endpoints
+	// (e.g. /v1/files). Those endpoints accept much larger bodies
+	// (up to 500 MB for Anthropic, 50 MB for OpenAI). If file-upload
+	// routes are added, they will need a per-route limit instead of
+	// this single global cap.
+	maxRequestBodyBytes = 32 << 20 // 32 MiB
+
 	// ErrorCodeProviderDisabled is the code written in the response
 	// body when a request targets a configured-but-disabled provider.
 	// Paired with HTTP 503.
@@ -214,8 +230,12 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
 		interceptor, err := p.CreateInterceptor(w, r.WithContext(ctx), tracer)
 		if err != nil {
 			span.SetStatus(codes.Error, fmt.Sprintf("failed to create interceptor: %v", err))
-			logger.Warn(ctx, "failed to create interceptor", slog.Error(err), slog.F("path", r.URL.Path))
-			http.Error(w, fmt.Sprintf("failed to create %q interceptor", r.URL.Path), http.StatusInternalServerError)
+			if _, ok := errors.AsType[*http.MaxBytesError](err); ok {
+				writeRequestBodyTooLarge(w)
+			} else {
+				logger.Warn(ctx, "failed to create interceptor", slog.Error(err), slog.F("path", r.URL.Path))
+				http.Error(w, fmt.Sprintf("failed to create %q interceptor", r.URL.Path), http.StatusInternalServerError)
+			}
 			return
 		}
 
@@ -328,6 +348,15 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC
 	}
 }
 
+// writeRequestBodyTooLarge writes a human-readable 413 response indicating that
+// the request body exceeded maxRequestBodyBytes.
+func writeRequestBodyTooLarge(w http.ResponseWriter) {
+	http.Error(w, fmt.Sprintf(
+		"Request body too large. The maximum allowed request body size is %dMiB.",
+		maxRequestBodyBytes>>20,
+	), http.StatusRequestEntityTooLarge)
+}
+
 // ServeHTTP exposes the internal http.Handler, which has all [Provider]s' routes registered.
 // It also tracks inflight requests.
 func (b *RequestBridge) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
@@ -350,6 +379,9 @@ func (b *RequestBridge) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 		b.inflightWG.Done()
 	}()
 
+	// Enforce the request body size limit. MaxBytesReader counts bytes as
+	// they are read from the connection and fails when the limit is exceeded.
+	r.Body = http.MaxBytesReader(rw, r.Body, maxRequestBodyBytes)
 	b.mux.ServeHTTP(rw, r.WithContext(ctx))
 }
 

aibridge/bridge_test.go

@@ -1,8 +1,12 @@
 package aibridge_test
 
 import (
+	"bytes"
+	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -206,6 +210,77 @@ func TestPassthroughRoutesForProviders(t *testing.T) {
 	}
 }
 
+func TestRequestBodySizeLimit(t *testing.T) {
+	t.Parallel()
+
+	newOpenAI := func(baseURL string) provider.Provider {
+		return aibridge.NewOpenAIProvider(config.OpenAI{Name: "openai", BaseURL: baseURL})
+	}
+	newAnthropic := func(baseURL string) provider.Provider {
+		return aibridge.NewAnthropicProvider(config.Anthropic{Name: "anthropic", BaseURL: baseURL}, nil)
+	}
+	newCopilot := func(baseURL string) provider.Provider {
+		return aibridge.NewCopilotProvider(config.Copilot{Name: "copilot", BaseURL: baseURL})
+	}
+
+	// Each body is a well-formed, schema-valid request for its provider, with
+	// an oversized message content that pushes it past the 32 MiB limit.
+	filler := strings.Repeat("A", 32<<20)
+	chatCompletionsBody := fmt.Appendf(nil, `{"model":"gpt-4","messages":[{"role":"user","content":"%s"}]}`, filler)
+	responsesBody := fmt.Appendf(nil, `{"model":"gpt-4","input":"%s"}`, filler)
+	messagesBody := fmt.Appendf(nil, `{"model":"claude-3-5-sonnet-latest","max_tokens":1024,"messages":[{"role":"user","content":"%s"}]}`, filler)
+
+	tests := []struct {
+		name     string
+		provider func(baseURL string) provider.Provider
+		path     string
+		body     []byte
+	}{
+		{name: "openai_passthrough", provider: newOpenAI, path: "/openai/v1/models", body: chatCompletionsBody},
+		{name: "openai_chat_completions", provider: newOpenAI, path: "/openai/v1/chat/completions", body: chatCompletionsBody},
+		{name: "openai_responses", provider: newOpenAI, path: "/openai/v1/responses", body: responsesBody},
+		{name: "anthropic_passthrough", provider: newAnthropic, path: "/anthropic/v1/models", body: messagesBody},
+		{name: "anthropic_messages", provider: newAnthropic, path: "/anthropic/v1/messages", body: messagesBody},
+		{name: "copilot_passthrough", provider: newCopilot, path: "/copilot/models", body: chatCompletionsBody},
+		{name: "copilot_chat_completions", provider: newCopilot, path: "/copilot/chat/completions", body: chatCompletionsBody},
+		{name: "copilot_responses", provider: newCopilot, path: "/copilot/responses", body: responsesBody},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			logger := slogtest.Make(t, nil)
+
+			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				_, _ = io.ReadAll(r.Body)
+				w.WriteHeader(http.StatusOK)
+			}))
+			t.Cleanup(upstream.Close)
+
+			prov := tc.provider(upstream.URL)
+			bridge, err := aibridge.NewRequestBridge(
+				t.Context(),
+				[]provider.Provider{prov},
+				nil, nil, logger, nil, bridgeTestTracer,
+			)
+			require.NoError(t, err)
+
+			req := httptest.NewRequest(http.MethodPost, tc.path, bytes.NewReader(tc.body))
+			// Unknown Content-Length
+			req.ContentLength = -1
+			// Copilot's bridged route checks Authorization before reading the
+			// body, so provide a token to reach the read path.
+			req.Header.Set("Authorization", "Bearer test-key")
+			resp := httptest.NewRecorder()
+			bridge.ServeHTTP(resp, req)
+
+			assert.Equal(t, http.StatusRequestEntityTooLarge, resp.Code)
+			assert.Contains(t, resp.Body.String(), "Request body too large")
+		})
+	}
+}
+
 // TestDisabledProviderHandler asserts that requests to a disabled
 // provider return a 503 with an ErrorCodeProviderDisabled body and
 // that a sibling enabled provider keeps routing normally.

aibridge/passthrough.go

@@ -2,6 +2,7 @@ package aibridge
 
 import (
 	"context"
+	"errors"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -54,8 +55,12 @@ func newPassthroughRouter(prov provider.Provider, logger slog.Logger, m *metrics
 			prov.KeyFailoverConfig(logger),
 		),
 		ErrorHandler: func(rw http.ResponseWriter, req *http.Request, e error) {
-			logger.Warn(req.Context(), "reverse proxy error", slog.Error(e), slog.F("path", req.URL.Path))
-			http.Error(rw, "upstream proxy error", http.StatusBadGateway)
+			if _, ok := errors.AsType[*http.MaxBytesError](e); ok {
+				writeRequestBodyTooLarge(rw)
+			} else {
+				logger.Warn(req.Context(), "reverse proxy error", slog.Error(e), slog.F("path", req.URL.Path))
+				http.Error(rw, "upstream proxy error", http.StatusBadGateway)
+			}
 		},
 	}