diff --git a/aibridge/bridge.go b/aibridge/bridge.go index 65d822069bdc8..835c125b6c9e3 100644 --- a/aibridge/bridge.go +++ b/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)) } diff --git a/aibridge/bridge_test.go b/aibridge/bridge_test.go index 93beb82de9abf..9ac7ea9ec3ddb 100644 --- a/aibridge/bridge_test.go +++ b/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. diff --git a/aibridge/passthrough.go b/aibridge/passthrough.go index 0dc6beb480ef0..c84802bc52a7e 100644 --- a/aibridge/passthrough.go +++ b/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) + } }, }