Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 34 additions & 2 deletions aibridge/bridge.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package aibridge

import (
"context"
"errors"
"fmt"
"net/http"
"net/url"
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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)
Comment thread
pawbana marked this conversation as resolved.
}

// 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) {
Expand All @@ -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))
}

Expand Down
75 changes: 75 additions & 0 deletions aibridge/bridge_test.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
package aibridge_test

import (
"bytes"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"

"github.com/stretchr/testify/assert"
Expand Down Expand Up @@ -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.
Expand Down
9 changes: 7 additions & 2 deletions aibridge/passthrough.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package aibridge

import (
"context"
"errors"
"net/http"
"net/http/httputil"
"net/url"
Expand Down Expand Up @@ -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)
}
},
}

Expand Down
Loading