diff --git a/internal/dialect/anthropic.go b/internal/dialect/anthropic.go index 66e64fdc..d2fb7c02 100644 --- a/internal/dialect/anthropic.go +++ b/internal/dialect/anthropic.go @@ -8,6 +8,7 @@ import ( "strings" "gpt-load/internal/health" + platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/protocol" "gpt-load/internal/state" ) @@ -141,7 +142,7 @@ func (d *Anthropic) listModelsPage( return anthropicModelPage{}, fmt.Errorf("create Anthropic model-list request: %w", err) } ApplyCredential(d, request.Header, apiKey, rules) - request.Header.Set("Accept-Encoding", "identity") + platformheader.NormalizeUpstreamRequestRepresentation(request, 0) response, err := d.client.Do(request) if err != nil { diff --git a/internal/dialect/anthropic_test.go b/internal/dialect/anthropic_test.go index 6968b3d7..38155bf4 100644 --- a/internal/dialect/anthropic_test.go +++ b/internal/dialect/anthropic_test.go @@ -238,21 +238,23 @@ func TestAnthropicListModelsPaginatesAndPreservesInputs(t *testing.T) { } } -func TestAnthropicListModelsForcesIdentity(t *testing.T) { +func TestAnthropicListModelsIdentityRepresentation(t *testing.T) { client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - if got := request.Header.Values("Accept-Encoding"); len(got) != 1 || got[0] != "identity" { - t.Fatalf("Accept-Encoding = %#v, want one identity", got) - } + assertOutboundRequestRepresentation(t, request, 0, map[string]string{ + "X-Api-Key": "override-secret", + "Anthropic-Version": "2025-01-01", + "X-Business": "preserved", + }) return anthropicTestResponse(http.StatusOK, `{"data":[]}`), nil })} models, err := NewAnthropic(client).ListModels( context.Background(), "https://api.example.com", "secret", - state.HeaderRules{ - Set: map[string]string{"Accept-Encoding": "gzip"}, - Remove: []string{"Accept-Encoding"}, - }, + legacyRepresentationHeaderRules(map[string]string{ + "X-Api-Key": "override-${API_KEY}", + "Anthropic-Version": "2025-01-01", + }), ) if err != nil || models == nil { t.Fatalf("ListModels() = %#v, %v", models, err) diff --git a/internal/dialect/gemini.go b/internal/dialect/gemini.go index 2b6fa3e7..79e63490 100644 --- a/internal/dialect/gemini.go +++ b/internal/dialect/gemini.go @@ -8,6 +8,7 @@ import ( "strings" "gpt-load/internal/health" + platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/protocol" "gpt-load/internal/state" ) @@ -171,7 +172,7 @@ func (d *Gemini) listModelsPage( return geminiModelPage{}, fmt.Errorf("create Gemini model-list request: %w", err) } ApplyCredential(d, request.Header, apiKey, rules) - request.Header.Set("Accept-Encoding", "identity") + platformheader.NormalizeUpstreamRequestRepresentation(request, 0) response, err := d.client.Do(request) if err != nil { diff --git a/internal/dialect/gemini_test.go b/internal/dialect/gemini_test.go index 2fa30ab4..ae9781f6 100644 --- a/internal/dialect/gemini_test.go +++ b/internal/dialect/gemini_test.go @@ -251,21 +251,19 @@ func TestGeminiListModelsPaginatesPreservesInputsAndParsesNames(t *testing.T) { } } -func TestGeminiListModelsForcesIdentity(t *testing.T) { +func TestGeminiListModelsIdentityRepresentation(t *testing.T) { client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - if got := request.Header.Values("Accept-Encoding"); len(got) != 1 || got[0] != "identity" { - t.Fatalf("Accept-Encoding = %#v, want one identity", got) - } + assertOutboundRequestRepresentation(t, request, 0, map[string]string{ + "X-Goog-Api-Key": "override-secret", + "X-Business": "preserved", + }) return geminiTestResponse(http.StatusOK, `{"models":[]}`), nil })} models, err := NewGemini(client).ListModels( context.Background(), "https://api.example.com", "secret", - state.HeaderRules{ - Set: map[string]string{"Accept-Encoding": "gzip"}, - Remove: []string{"Accept-Encoding"}, - }, + legacyRepresentationHeaderRules(map[string]string{"X-Goog-Api-Key": "override-${API_KEY}"}), ) if err != nil || models == nil { t.Fatalf("ListModels() = %#v, %v", models, err) diff --git a/internal/dialect/openai.go b/internal/dialect/openai.go index 2e3c6b6a..17625174 100644 --- a/internal/dialect/openai.go +++ b/internal/dialect/openai.go @@ -6,6 +6,7 @@ import ( "net/http" "gpt-load/internal/health" + platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/protocol" "gpt-load/internal/state" ) @@ -78,7 +79,7 @@ func (d *OpenAI) ListModels( return nil, fmt.Errorf("create OpenAI model-list request: %w", err) } ApplyCredential(d, req.Header, apiKey, rules) - req.Header.Set("Accept-Encoding", "identity") + platformheader.NormalizeUpstreamRequestRepresentation(req, 0) resp, err := d.client.Do(req) if err != nil { diff --git a/internal/dialect/openai_test.go b/internal/dialect/openai_test.go index bf1bf2c1..f45f2145 100644 --- a/internal/dialect/openai_test.go +++ b/internal/dialect/openai_test.go @@ -384,11 +384,12 @@ func TestOpenAIListModelsAppliesHeaderRuleOverrides(t *testing.T) { } } -func TestOpenAIListModelsForcesIdentityAndCollectsUniqueModels(t *testing.T) { +func TestOpenAIListModelsIdentityRepresentationAndCollectsUniqueModels(t *testing.T) { client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - if got := request.Header.Values("Accept-Encoding"); !reflect.DeepEqual(got, []string{"identity"}) { - t.Fatalf("Accept-Encoding = %#v, want one identity", got) - } + assertOutboundRequestRepresentation(t, request, 0, map[string]string{ + "Authorization": "Token secret", + "X-Business": "preserved", + }) return &http.Response{ StatusCode: http.StatusOK, Header: make(http.Header), @@ -400,10 +401,7 @@ func TestOpenAIListModelsForcesIdentityAndCollectsUniqueModels(t *testing.T) { context.Background(), "https://api.example.com", "secret", - state.HeaderRules{ - Set: map[string]string{"Accept-Encoding": "gzip"}, - Remove: []string{"Accept-Encoding"}, - }, + legacyRepresentationHeaderRules(map[string]string{"Authorization": "Token ${API_KEY}"}), ) if err != nil { t.Fatalf("ListModels() error = %v", err) diff --git a/internal/dialect/probe.go b/internal/dialect/probe.go index 3cf12fa8..87038428 100644 --- a/internal/dialect/probe.go +++ b/internal/dialect/probe.go @@ -8,6 +8,7 @@ import ( "net/http" "strings" + platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/state" ) @@ -48,7 +49,7 @@ func executeProbe( } request.Header.Set("Content-Type", "application/json") ApplyCredential(dialect, request.Header, apiKey, rules) - request.Header.Set("Accept-Encoding", "identity") + platformheader.NormalizeUpstreamRequestRepresentation(request, int64(len(body))) response, err := client.Do(request) if err != nil { diff --git a/internal/dialect/probe_test.go b/internal/dialect/probe_test.go index 5ebed50c..f5520f08 100644 --- a/internal/dialect/probe_test.go +++ b/internal/dialect/probe_test.go @@ -80,14 +80,14 @@ func TestGeminiProbeSendsMinimalAuthenticatedRequest(t *testing.T) { } } -func TestProbeAppliesHeaderRulesAndForcesIdentity(t *testing.T) { +func TestProbeIdentityRepresentationPreservesRulesAndCredential(t *testing.T) { client := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { - if got := request.Header.Get("X-Probe"); got != "prefix-secret" { - t.Errorf("X-Probe = %q, want expanded rule", got) - } - if got := request.Header.Values("Accept-Encoding"); len(got) != 1 || got[0] != "identity" { - t.Errorf("Accept-Encoding = %#v, want one identity", got) - } + assertProbeRequestRepresentation(t, request, map[string]string{ + "Authorization": "Token secret", + "Content-Type": "application/json", + "X-Business": "preserved", + "X-Probe": "prefix-secret", + }) return &http.Response{ StatusCode: http.StatusNoContent, Body: io.NopCloser(strings.NewReader("ignored")), @@ -95,12 +95,11 @@ func TestProbeAppliesHeaderRulesAndForcesIdentity(t *testing.T) { }, nil })} - err := NewOpenAI(client).Probe(context.Background(), "https://api.example.com", "secret", state.HeaderRules{ - Set: map[string]string{ - "X-Probe": "prefix-${API_KEY}", - "Accept-Encoding": "gzip", - }, - }, "gpt-test") + rules := legacyRepresentationHeaderRules(map[string]string{ + "Authorization": "Token ${API_KEY}", + "X-Probe": "prefix-${API_KEY}", + }) + err := NewOpenAI(client).Probe(context.Background(), "https://api.example.com", "secret", rules, "gpt-test") if err != nil { t.Fatalf("Probe() error = %v", err) } diff --git a/internal/dialect/request_representation_test_helpers_test.go b/internal/dialect/request_representation_test_helpers_test.go new file mode 100644 index 00000000..3b0b4a45 --- /dev/null +++ b/internal/dialect/request_representation_test_helpers_test.go @@ -0,0 +1,94 @@ +package dialect + +import ( + "io" + "net/http" + "reflect" + "strings" + "testing" + + "gpt-load/internal/state" +) + +func legacyRepresentationHeaderRules(extra map[string]string) state.HeaderRules { + set := map[string]string{ + "Accept-Encoding": "gzip", + "Content-Encoding": "zstd", + "Content-Length": "999", + "ETag": `"stale"`, + "Digest": "sha-256=stale", + "Content-MD5": "stale-md5", + "Content-Range": "bytes 0-1/2", + "Content-Digest": "sha-256=:c3RhbGU=:", + "Repr-Digest": "sha-256=:c3RhbGU=:", + "Signature": "stale-signature", + "Signature-Input": "stale-signature-input", + "X-Business": "preserved", + } + for name, value := range extra { + set[name] = value + } + return state.HeaderRules{Set: set} +} + +func assertOutboundRequestRepresentation( + t *testing.T, + request *http.Request, + wantBodyLength int64, + wantHeaders map[string]string, +) { + t.Helper() + if values := requestHeaderFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { + t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) + } + for _, name := range []string{ + "Content-Encoding", + "Content-Length", + "ETag", + "Digest", + "Content-MD5", + "Content-Range", + "Content-Digest", + "Repr-Digest", + "Signature", + "Signature-Input", + } { + if values := requestHeaderFieldValues(request.Header, name); values != nil { + t.Errorf("%s values = %#v, want absent", name, values) + } + } + if request.ContentLength != wantBodyLength { + t.Errorf("ContentLength = %d, want %d", request.ContentLength, wantBodyLength) + } + for name, want := range wantHeaders { + if got := request.Header.Get(name); got != want { + t.Errorf("%s = %q, want %q", name, got, want) + } + } +} + +func assertProbeRequestRepresentation( + t *testing.T, + request *http.Request, + wantHeaders map[string]string, +) { + t.Helper() + body, err := io.ReadAll(request.Body) + if err != nil { + t.Fatalf("read probe request body: %v", err) + } + if len(body) == 0 { + t.Fatal("probe request body is empty") + } + assertOutboundRequestRepresentation(t, request, int64(len(body)), wantHeaders) +} + +func requestHeaderFieldValues(headers http.Header, target string) []string { + var values []string + for name, fieldValues := range headers { + if strings.EqualFold(name, target) { + values = append(values, fieldValues...) + } + } + return values +} diff --git a/internal/gateway/content_coding_test_helpers_test.go b/internal/gateway/content_coding_test_helpers_test.go new file mode 100644 index 00000000..c73e006e --- /dev/null +++ b/internal/gateway/content_coding_test_helpers_test.go @@ -0,0 +1,66 @@ +package gateway + +import ( + "bytes" + "compress/gzip" + "compress/zlib" + "testing" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" + + "gpt-load/internal/platform/contentcoding" +) + +func encodeContentCodingForGatewayTest( + t *testing.T, + encoding contentcoding.Encoding, + body []byte, +) []byte { + t.Helper() + if encoding == contentcoding.Identity { + return bytes.Clone(body) + } + + var buffer bytes.Buffer + switch encoding { + case contentcoding.Gzip: + writer := gzip.NewWriter(&buffer) + if _, err := writer.Write(body); err != nil { + t.Fatalf("gzip fixture Write() error = %v", err) + } + if err := writer.Close(); err != nil { + t.Fatalf("gzip fixture Close() error = %v", err) + } + case contentcoding.Brotli: + writer := brotli.NewWriter(&buffer) + if _, err := writer.Write(body); err != nil { + t.Fatalf("brotli fixture Write() error = %v", err) + } + if err := writer.Close(); err != nil { + t.Fatalf("brotli fixture Close() error = %v", err) + } + case contentcoding.Deflate: + writer := zlib.NewWriter(&buffer) + if _, err := writer.Write(body); err != nil { + t.Fatalf("deflate fixture Write() error = %v", err) + } + if err := writer.Close(); err != nil { + t.Fatalf("deflate fixture Close() error = %v", err) + } + case contentcoding.Zstd: + writer, err := zstd.NewWriter(&buffer, zstd.WithEncoderConcurrency(1)) + if err != nil { + t.Fatalf("zstd fixture NewWriter() error = %v", err) + } + if _, err := writer.Write(body); err != nil { + t.Fatalf("zstd fixture Write() error = %v", err) + } + if err := writer.Close(); err != nil { + t.Fatalf("zstd fixture Close() error = %v", err) + } + default: + t.Fatalf("unsupported fixture encoding %q", encoding) + } + return buffer.Bytes() +} diff --git a/internal/gateway/dialects_integration_test.go b/internal/gateway/dialects_integration_test.go index 10602b20..b43e8db4 100644 --- a/internal/gateway/dialects_integration_test.go +++ b/internal/gateway/dialects_integration_test.go @@ -20,6 +20,7 @@ import ( "gpt-load/internal/dialect" "gpt-load/internal/health" + "gpt-load/internal/platform/contentcoding" "gpt-load/internal/platform/encryption" platformhttp "gpt-load/internal/platform/httpclient" "gpt-load/internal/platform/redact" @@ -137,6 +138,98 @@ func newDialectGatewayEngine( return engine, registry } +func TestGatewayDecodesSupportedClientContentCodings(t *testing.T) { + routes := []struct { + name string + protocol protocol.Protocol + dialects dialect.Set + path string + plaintext []byte + }{ + { + name: "OpenAI completions", protocol: protocol.OpenAICompletions, + dialects: dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), + path: "/v1/chat/completions", plaintext: []byte(`{"model":"public-model","stream":false}`), + }, + { + name: "Anthropic messages", protocol: protocol.Anthropic, + dialects: dialect.NewSet(dialect.NewAnthropic(http.DefaultClient)), + path: "/v1/messages", plaintext: []byte(`{"model":"public-model","max_tokens":1,"messages":[]}`), + }, + { + name: "Gemini generation", protocol: protocol.Gemini, + dialects: dialect.NewSet(dialect.NewGemini(http.DefaultClient)), + path: "/v1beta/models/public-model:generateContent", plaintext: []byte(`{"contents":[{"role":"user","parts":[{"text":"ping"}]}]}`), + }, + { + name: "Responses create", protocol: protocol.OpenAIResponses, + dialects: dialect.NewSet(dialect.NewOpenAIResponses(http.DefaultClient)), + path: "/v1/responses", plaintext: []byte(`{"model":"public-model","input":"ping"}`), + }, + } + encodings := []contentcoding.Encoding{ + contentcoding.Identity, + contentcoding.Gzip, + contentcoding.Brotli, + contentcoding.Deflate, + contentcoding.Zstd, + } + + for _, route := range routes { + t.Run(route.name, func(t *testing.T) { + for _, encoding := range encodings { + t.Run(string(encoding), func(t *testing.T) { + var receivedBody []byte + var receivedHeader http.Header + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + receivedHeader = request.Header.Clone() + var err error + receivedBody, err = io.ReadAll(request.Body) + if err != nil { + t.Errorf("ReadAll(upstream body) error = %v", err) + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"ok":true}`)) + })) + defer upstream.Close() + + engine, _ := newDialectGatewayEngine( + t, + route.protocol, + "public-model", + route.dialects, + dialectGatewayGroup{ + id: 1, name: route.name, upstreamURL: upstream.URL, + apiKeys: []string{"sk-provider"}, + }, + ) + request := httptest.NewRequest( + http.MethodPost, + route.path, + bytes.NewReader(encodeContentCodingForGatewayTest(t, encoding, route.plaintext)), + ) + request.Header.Set("Authorization", "Bearer gl-client") + if encoding != contentcoding.Identity { + request.Header.Set("Content-Encoding", string(encoding)) + } + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusOK { + t.Fatalf("response = %d headers=%v body=%s", recorder.Code, recorder.Header(), recorder.Body.String()) + } + if !bytes.Equal(receivedBody, route.plaintext) { + t.Fatalf("upstream body = %s, want plaintext %s", receivedBody, route.plaintext) + } + if got := receivedHeader.Get("Content-Encoding"); got != "" { + t.Fatalf("upstream Content-Encoding = %q, want absent", got) + } + }) + } + }) + } +} + func TestResponsesNamespaceRoutesOrdinaryMethodsAndRejectsDangerousMethodsLocally(t *testing.T) { var ( mu sync.Mutex @@ -694,8 +787,8 @@ func TestTransparentModelRoutePreservesWire(t *testing.T) { if recorder.Code != http.StatusOK || !bytes.Equal(recorder.Body.Bytes(), rawResponse) { t.Fatalf("transparent response = %d %q, want %q", recorder.Code, recorder.Body.Bytes(), rawResponse) } - if receivedHeader.Get("Accept-Encoding") != "gzip" { - t.Fatalf("upstream Accept-Encoding = %q, want HeaderRule gzip", receivedHeader.Get("Accept-Encoding")) + if receivedHeader.Get("Accept-Encoding") != "identity" { + t.Fatalf("upstream Accept-Encoding = %q, want final identity normalization", receivedHeader.Get("Accept-Encoding")) } assertRepresentationMetadata(t, recorder.Header(), true) } diff --git a/internal/gateway/forward.go b/internal/gateway/forward.go index 94e5fc0b..97f40d9d 100644 --- a/internal/gateway/forward.go +++ b/internal/gateway/forward.go @@ -14,17 +14,16 @@ import ( "net/http/httptrace" "net/textproto" "sort" - "strconv" "strings" "sync" "sync/atomic" "time" "gpt-load/internal/dialect" + "gpt-load/internal/platform/contentcoding" platformhttp "gpt-load/internal/platform/httpclient" platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/platform/redact" - "gpt-load/internal/platform/utils" "gpt-load/internal/state" "gpt-load/internal/usage" ) @@ -103,6 +102,34 @@ func (forwarder *Forwarder) Forward(ctx context.Context, input ForwardInput) Ups } defer response.Body.Close() + representationHeaders := response.Header.Clone() + headers := cloneEndToEndHeaders(response.Header) + policy := classifyResponseBody(input.Request.Method, response.StatusCode) + if !policy.readBody { + encoding, parseErr := contentcoding.ParseContentEncoding( + headerFieldValues(representationHeaders, "Content-Encoding"), + ) + if parseErr != nil || encoding != contentcoding.Identity { + return UpstreamResult{ + Err: fmt.Errorf("%w: bodyless response has non-identity Content-Encoding", ErrUpstreamProtocol), + RequestWritten: true, + } + } + safeHeaders := sanitizeForwardResponseHeaders(representationHeaders, input, knownSecrets...) + safeHeaders, _ = normalizeBufferedResponse( + input.Request.Method, + response.StatusCode, + safeHeaders, + nil, + ) + return UpstreamResult{ + StatusCode: response.StatusCode, + Header: safeHeaders, + RequestWritten: true, + Usage: usage.Result{State: usage.StateNotApplicable}, + } + } + success := response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices limit := maxErrorResponseBodyBytes if success { @@ -123,7 +150,6 @@ func (forwarder *Forwarder) Forward(ctx context.Context, input ForwardInput) Ups } } - headers := cloneEndToEndHeaders(response.Header) result := UpstreamResult{ StatusCode: response.StatusCode, Body: body, @@ -133,42 +159,42 @@ func (forwarder *Forwarder) Forward(ctx context.Context, input ForwardInput) Ups if success { prepared, prepareErr := forwarder.prepareSuccessRepresentation( input, - headers, + response.StatusCode, + representationHeaders, body, knownSecrets, ) if prepareErr != nil { return UpstreamResult{Err: prepareErr, RequestWritten: true} } - result.Body = prepared.wire + result.Body = prepared.downstream result.Header = prepared.headers if input.ObserveUsage && forwarder.usageCapture != nil { result.Usage = forwarder.usageCapture.extractNonStreamingPlain( input.Dialect, - prepared.plain, + prepared.inspectable, ) } return result } if !success { - var safeWire, safePlain []byte - if overflow { - safeWire, safePlain = failClosedErrorBody(headers) - } else { - safeWire, safePlain = forwarder.prepareErrorBody( - headers, body, input, request.Header, - ) - } - if nonIdentityEncodingContainsKey(headers, input.APIKey) { - safeWire, safePlain = failClosedErrorBody(headers) - } - result.Body = safeWire - result.ClassificationBody = safePlain - result.ErrorSummary = summarizeErrorBody( - forwarder.redactor, safePlain, "", summarySecrets..., + prepared := forwarder.prepareErrorRepresentation( + input, + representationHeaders, + body, + knownSecrets, ) + if overflow { + prepared = forwarder.failClosedErrorRepresentation(input, headers, knownSecrets) + } + result.Body = prepared.downstream + result.ClassificationBody = prepared.classification + result.ErrorSummary = summarizeErrorBody( + forwarder.redactor, prepared.classification, "", summarySecrets..., + ) + result.Header = prepared.headers + return result } - result.Header = sanitizeForwardResponseHeaders(headers, input, knownSecrets...) return result } @@ -205,6 +231,7 @@ func (forwarder *Forwarder) ForwardStream( streamBody := response.Body defer func() { _ = streamBody.Close() }() + representationHeaders := response.Header.Clone() headers := cloneEndToEndHeaders(response.Header) result = UpstreamResult{ StatusCode: response.StatusCode, @@ -217,23 +244,19 @@ func (forwarder *Forwarder) ForwardStream( result.RetryableBeforeCommit = retryableBeforeCommit(ctx, result.RequestWritten) return result } + prepared := forwarder.prepareErrorRepresentation(input, representationHeaders, body, knownSecrets) if overflow { - result.Body, result.ClassificationBody = failClosedErrorBody(headers) - } else { - result.Body, result.ClassificationBody = forwarder.prepareErrorBody( - headers, body, input, request.Header, - ) - } - if nonIdentityEncodingContainsKey(headers, input.APIKey) { - result.Body, result.ClassificationBody = failClosedErrorBody(headers) + prepared = forwarder.failClosedErrorRepresentation(input, headers, knownSecrets) } + result.Body = prepared.downstream + result.ClassificationBody = prepared.classification result.ErrorSummary = summarizeErrorBody( forwarder.redactor, result.ClassificationBody, "", summarySecrets..., ) - result.Header = sanitizeForwardResponseHeaders(headers, input, knownSecrets...) + result.Header = prepared.headers return result } streamEvents := newStreamEventObserver( @@ -245,8 +268,8 @@ func (forwarder *Forwarder) ForwardStream( ) defer func() { result.Usage = streamEvents.finalizeUsage() }() - if !inspectableStreamEncoding(response.Header) { - result.Err = fmt.Errorf("%w: Content-Encoding %q", ErrUpstreamProtocol, response.Header.Values("Content-Encoding")) + if err := validateStreamContentEncoding(representationHeaders); err != nil { + result.Err = err result.RetryableBeforeCommit = retryableBeforeCommit(ctx, result.RequestWritten) return result } @@ -379,7 +402,9 @@ func (forwarder *Forwarder) ForwardStream( input.OnStreamReady() } - result.Header = sanitizeForwardResponseHeaders(headers, input, knownSecrets...) + result.Header = normalizeStreamResponseHeaders( + sanitizeForwardResponseHeaders(headers, input, knownSecrets...), + ) streamWriter := newStreamWriteController(downstream, forwarder.streamWriteTimeout) defer func() { _ = streamWriter.clear() }() @@ -481,7 +506,6 @@ func (forwarder *Forwarder) newUpstreamRequest( if stream && input.Group.InjectUsageOptions && forwarder != nil && forwarder.usageCapture != nil { parsed = forwarder.usageCapture.injectStreamUsage(input.Dialect, parsed) } - bodyChanged := !bytes.Equal(input.Request.Body, parsed.Body) upstreamURL, err := input.Dialect.BuildUpstreamURL(input.Group.UpstreamURL, parsed) if err != nil { return nil, nil, nil, fmt.Errorf("build upstream URL: %w", err) @@ -495,15 +519,10 @@ func (forwarder *Forwarder) newUpstreamRequest( request.ContentLength = int64(len(parsed.Body)) request.GetBody = nil request.Header = cloneEndToEndHeaders(parsed.Header) - if bodyChanged { - invalidateRewrittenBodyHeaders(request.Header) - } removeDownstreamCredentials(request.Header) dialect.ApplyCredential(input.Dialect, request.Header, input.APIKey, input.Group.HeaderRules) sanitizeUpstreamRequestHeaders(request.Header) - if stream || rewrite { - request.Header.Set("Accept-Encoding", "identity") - } + platformheader.NormalizeUpstreamRequestRepresentation(request, int64(len(parsed.Body))) if _, exists := request.Header["User-Agent"]; !exists { request.Header["User-Agent"] = nil } @@ -520,72 +539,6 @@ func needsModelRewrite(input ForwardInput) bool { input.ExternalModel != input.UpstreamModelID } -func (forwarder *Forwarder) prepareErrorBody( - headers http.Header, - wire []byte, - input ForwardInput, - resolvedHeaders ...http.Header, -) ([]byte, []byte) { - if int64(len(wire)) > maxErrorResponseBodyBytes { - return failClosedErrorBody(headers) - } - encoding, ok := inspectableErrorBodyEncoding(headers, wire) - if !ok { - return failClosedErrorBody(headers) - } - plain, err := utils.DecompressResponseLimited(encoding, wire, maxDecompressedErrorBodyBytes) - if err != nil { - return failClosedErrorBody(headers) - } - var finalHeaders http.Header - if len(resolvedHeaders) > 0 { - finalHeaders = resolvedHeaders[0] - } - safePlain := plain - for _, secret := range resolvedCredentialSecretValues(input, finalHeaders) { - var ok bool - safePlain, ok = rewriteBoundedLiteral( - safePlain, - secret, - redact.Placeholder, - maxDecompressedErrorBodyBytes, - ) - if !ok { - return failClosedErrorBody(headers) - } - } - safePlain = forwarder.redactor.Bytes(safePlain) - if int64(len(safePlain)) > maxDecompressedErrorBodyBytes { - return failClosedErrorBody(headers) - } - downstreamPlain := safePlain - if needsModelRewrite(input) { - downstreamPlain, ok = rewriteBoundedLiteral( - safePlain, - input.UpstreamModelID, - input.ExternalModel, - maxDecompressedErrorBodyBytes, - ) - if !ok { - wire, _ := failClosedErrorBody(headers) - return wire, safePlain - } - } - if bytes.Equal(downstreamPlain, plain) { - return bytes.Clone(wire), safePlain - } - safeWire, err := utils.CompressResponse(encoding, downstreamPlain) - if err != nil || int64(len(safeWire)) > maxErrorResponseBodyBytes { - if needsModelRewrite(input) { - wire, _ := failClosedErrorBody(headers) - return wire, safePlain - } - return failClosedErrorBody(headers) - } - updateRewrittenBodyHeaders(headers, len(safeWire)) - return safeWire, safePlain -} - func isResolvedCredentialHeaderName(selected dialect.Dialect, name string) bool { if platformheader.IsCredentialName(name) { return true @@ -788,38 +741,6 @@ func replacementResultLength(sourceLength, oldLength, replacementLength, count, return sourceLength + count*delta, true } -func inspectableErrorBodyEncoding(headers http.Header, wire []byte) (string, bool) { - values := headers.Values("Content-Encoding") - if len(values) > 1 { - return "", false - } - encoding := "" - if len(values) == 1 { - encoding = strings.ToLower(strings.TrimSpace(values[0])) - } - switch encoding { - case "", "identity": - return encoding, true - case "gzip", "br", "deflate", "zstd": - return encoding, len(wire) > 0 - default: - return "", false - } -} - -func failClosedErrorBody(headers http.Header) ([]byte, []byte) { - body := []byte(redact.Placeholder) - deleteHeaderField(headers, "Content-Encoding") - headers.Set("Content-Type", "text/plain; charset=utf-8") - updateRewrittenBodyHeaders(headers, len(body)) - return bytes.Clone(body), bytes.Clone(body) -} - -func updateRewrittenBodyHeaders(headers http.Header, bodyLength int) { - invalidateRewrittenBodyHeaders(headers) - headers.Set("Content-Length", strconv.Itoa(bodyLength)) -} - var representationMetadataHeaderNames = [...]string{ "Content-Encoding", "Content-Length", @@ -838,7 +759,7 @@ func invalidateRewrittenBodyHeaders(headers http.Header) { if strings.EqualFold(name, "Content-Encoding") { continue } - headers.Del(name) + deleteHeaderField(headers, name) } } @@ -878,16 +799,14 @@ func streamingClientConfig(timeouts state.TimeoutConfig) *platformhttp.Config { } } -func inspectableStreamEncoding(headers http.Header) bool { - values := headers.Values("Content-Encoding") - if len(values) == 0 { - return true +func validateStreamContentEncoding(headers http.Header) error { + encoding, err := contentcoding.ParseContentEncoding( + headerFieldValues(headers, "Content-Encoding"), + ) + if err != nil || encoding != contentcoding.Identity { + return fmt.Errorf("%w: non-identity stream Content-Encoding", ErrUpstreamProtocol) } - if len(values) != 1 { - return false - } - encoding := strings.TrimSpace(values[0]) - return encoding == "" || strings.EqualFold(encoding, "identity") + return nil } func readStreamingErrorBody(body io.Reader) ([]byte, bool, error) { @@ -1029,7 +948,7 @@ func sanitizeUpstreamRequestHeaders(headers http.Header) { } for name := range headers { if isHopByHopHeader(name) || - platformheader.IsForbiddenRequestRuleSetName(name) || + platformheader.IsForbiddenRequestRuleName(name) || isDebugHeader(name) { delete(headers, name) } @@ -1048,25 +967,6 @@ func headerValuesContainLiteral(values []string, literal string) bool { return false } -func nonIdentityEncodingContainsKey(headers http.Header, apiKey string) bool { - var values []string - for name, headerValues := range headers { - if strings.EqualFold(name, "Content-Encoding") { - values = append(values, headerValues...) - } - } - if !headerValuesContainLiteral(values, apiKey) { - return false - } - for _, value := range values { - normalized := strings.TrimSpace(value) - if normalized != "" && !strings.EqualFold(normalized, "identity") { - return true - } - } - return false -} - func sanitizeUpstreamResponseHeaders(source http.Header, apiKey string) http.Header { headers := cloneEndToEndHeaders(source) namesToDelete := make([]string, 0) @@ -1090,6 +990,7 @@ func sanitizeForwardResponseHeaders( additionalSecrets ...string, ) http.Header { headers := cloneEndToEndHeaders(source) + deleted := endToEndHeaderCleanupDeletedField(source, headers) namesToDelete := make([]string, 0) for actualName, values := range headers { deleteField := isResolvedCredentialHeaderName(input.Dialect, actualName) || @@ -1109,13 +1010,17 @@ func sanitizeForwardResponseHeaders( namesToDelete = append(namesToDelete, actualName) } } + deleted = deleted || len(namesToDelete) > 0 for _, name := range namesToDelete { deleteHeaderField(headers, name) } if !needsModelRewrite(input) { + if deleted { + deleteHeaderField(headers, "Signature") + deleteHeaderField(headers, "Signature-Input") + } return headers } - deleted := false for name, values := range headers { nameContainsModel := headerNameContainsLiteral(name, input.UpstreamModelID) valuesContainModel := headerValuesContainLiteral(values, input.UpstreamModelID) @@ -1147,6 +1052,22 @@ func sanitizeForwardResponseHeaders( return headers } +func endToEndHeaderCleanupDeletedField(source, headers http.Header) bool { + for sourceName := range source { + found := false + for headerName := range headers { + if strings.EqualFold(sourceName, headerName) { + found = true + break + } + } + if !found { + return true + } + } + return false +} + func headerNameContainsLiteral(name, literal string) bool { return literal != "" && strings.Contains(strings.ToLower(name), strings.ToLower(literal)) } diff --git a/internal/gateway/forward_test.go b/internal/gateway/forward_test.go index ddc13028..1367311a 100644 --- a/internal/gateway/forward_test.go +++ b/internal/gateway/forward_test.go @@ -24,10 +24,10 @@ import ( "gpt-load/internal/dialect" "gpt-load/internal/health" + "gpt-load/internal/platform/contentcoding" platformhttp "gpt-load/internal/platform/httpclient" platformheader "gpt-load/internal/platform/httpheader" "gpt-load/internal/platform/redact" - "gpt-load/internal/platform/utils" "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/usage" @@ -183,7 +183,9 @@ func TestForwardRequestRulesCannotRestoreReservedHeaders(t *testing.T) { debugHeaderGroup: "restored-group", debugHeaderKey: "restored-key", debugHeaderAttempts: "restored-attempts", - }}, + "Accept-Encoding": "gzip", + "Content-Encoding": "br", + }, Remove: []string{"Accept-Encoding", "Content-Encoding"}}, }, APIKey: apiKey, Request: &dialect.ParsedRequest{ @@ -194,6 +196,8 @@ func TestForwardRequestRulesCannotRestoreReservedHeaders(t *testing.T) { "X-Client-Hop": {"client-hop"}, "Cookie": {"client=fake"}, "Proxy-Client-Hop": {"proxy-client"}, + "accept-encoding": {"zstd"}, + "content-encoding": {"gzip"}, }, Body: []byte(`{"model":"gpt-test"}`), }, @@ -212,6 +216,9 @@ func TestForwardRequestRulesCannotRestoreReservedHeaders(t *testing.T) { if got := request.Header.Get("Authorization"); got != "Token "+apiKey { t.Fatalf("Authorization = %q, want configured credential", got) } + if values := headerFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { + t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) + } for _, name := range []string{ "Connection", "X-Rule-Hop", @@ -225,13 +232,198 @@ func TestForwardRequestRulesCannotRestoreReservedHeaders(t *testing.T) { debugHeaderGroup, debugHeaderKey, debugHeaderAttempts, + "Content-Encoding", } { - if values := request.Header.Values(name); values != nil { + if values := headerFieldValues(request.Header, name); values != nil { t.Errorf("upstream request Header %s survived: %#v", name, values) } } } +func TestForwardAllRequestsUseIdentityRepresentation(t *testing.T) { + tests := []struct { + name string + stream bool + configure func(*ForwardInput) + assert func(*testing.T, []byte) + }{ + { + name: "normal request", + assert: func(t *testing.T, body []byte) { + t.Helper() + if string(body) != `{"model":"public-model","stream":true}` { + t.Fatalf("upstream body = %s, want unchanged plaintext", body) + } + }, + }, + { + name: "alias rewrite", + configure: func(input *ForwardInput) { + input.ExternalModel = "public-model" + input.UpstreamModelID = "provider-model" + }, + assert: func(t *testing.T, body []byte) { + t.Helper() + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil || payload["model"] != "provider-model" { + t.Fatalf("upstream alias body = %s, error = %v", body, err) + } + }, + }, + { + name: "stream usage injection", + stream: true, + configure: func(input *ForwardInput) { + input.Group.InjectUsageOptions = true + }, + assert: func(t *testing.T, body []byte) { + t.Helper() + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + t.Fatalf("decode upstream stream body: %v", err) + } + options, ok := payload["stream_options"].(map[string]any) + if !ok || options["include_usage"] != true { + t.Fatalf("upstream stream body = %s, want include_usage", body) + } + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + input := streamForwardInput("https://api.example.test") + input.Request.Header = http.Header{ + "Accept-Encoding": {"gzip"}, + "accept-encoding": {"br"}, + "Content-Encoding": {"gzip"}, + "content-encoding": {"zstd"}, + "Content-Type": {"application/json"}, + "Accept": {"application/json"}, + "Idempotency-Key": {"request-1"}, + } + input.Request.Body = []byte(`{"model":"public-model","stream":true}`) + if test.configure != nil { + test.configure(&input) + } + + request, _, replay, err := NewForwarder( + platformhttp.NewHTTPClientManager(), + redact.New(), + ).newUpstreamRequest(t.Context(), input, test.stream) + if err != nil { + t.Fatalf("newUpstreamRequest() error = %v", err) + } + t.Cleanup(func() { + _ = request.Body.Close() + replay.release() + }) + body, err := io.ReadAll(request.Body) + if err != nil { + t.Fatalf("read upstream request body: %v", err) + } + test.assert(t, body) + if values := headerFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { + t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) + } + if values := headerFieldValues(request.Header, "Content-Encoding"); values != nil { + t.Fatalf("Content-Encoding values = %#v, want absent", values) + } + if request.ContentLength != int64(len(body)) { + t.Fatalf("ContentLength = %d, want %d", request.ContentLength, len(body)) + } + for name, want := range map[string]string{ + "Content-Type": "application/json", + "Accept": "application/json", + "Idempotency-Key": "request-1", + "Authorization": "Bearer sk-upstream-secret", + } { + if got := request.Header.Get(name); got != want { + t.Errorf("%s = %q, want %q", name, got, want) + } + } + }) + } +} + +func TestForwardNormalizesFinalRequestMetadata(t *testing.T) { + input := streamForwardInput("https://api.example.test") + input.Request.Header = http.Header{ + "Content-Type": {"application/json"}, + "Accept": {"application/json"}, + "Idempotency-Key": {"request-1"}, + "Content-Length": {"999"}, + "content-length": {"998"}, + "Content-Encoding": {"gzip"}, + "content-encoding": {"br"}, + "ETag": {`"client"`}, + "Digest": {"sha-256=client"}, + "Content-MD5": {"client-md5"}, + "Content-Range": {"bytes 0-1/2"}, + "Content-Digest": {"sha-256=:Y2xpZW50:"}, + "Repr-Digest": {"sha-256=:cmVwcg==:"}, + "Signature": {"client-signature"}, + "Signature-Input": {"client-signature-input"}, + "accept-encoding": {"gzip"}, + "Accept-Encoding": {"br"}, + } + input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ + "Authorization": "Token ${API_KEY}", + "Digest": "sha-256=rule", + "Signature": "rule-signature", + "Content-Encoding": "zstd", + "Accept-Encoding": "deflate", + }} + + request, _, replay, err := NewForwarder( + platformhttp.NewHTTPClientManager(), + redact.New(), + ).newUpstreamRequest(t.Context(), input, false) + if err != nil { + t.Fatalf("newUpstreamRequest() error = %v", err) + } + t.Cleanup(func() { + _ = request.Body.Close() + replay.release() + }) + body, err := io.ReadAll(request.Body) + if err != nil { + t.Fatalf("read upstream body: %v", err) + } + if request.ContentLength != int64(len(body)) { + t.Fatalf("ContentLength = %d, want %d", request.ContentLength, len(body)) + } + for _, name := range []string{ + "Content-Length", + "Content-Encoding", + "ETag", + "Digest", + "Content-MD5", + "Content-Range", + "Content-Digest", + "Repr-Digest", + "Signature", + "Signature-Input", + } { + if values := headerFieldValues(request.Header, name); values != nil { + t.Errorf("%s values = %#v, want removed", name, values) + } + } + if values := headerFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { + t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) + } + for name, want := range map[string]string{ + "Content-Type": "application/json", + "Accept": "application/json", + "Idempotency-Key": "request-1", + "Authorization": "Token sk-upstream-secret", + } { + if got := request.Header.Get(name); got != want { + t.Errorf("%s = %q, want %q", name, got, want) + } + } +} + func TestNewUpstreamRequestDoesNotExposeGetBody(t *testing.T) { const payload = `{"model":"gpt-test","messages":[]}` input := ForwardInput{ @@ -352,8 +544,9 @@ func TestForwarderBoundsDecompressedErrorBodies(t *testing.T) { t.Fatalf("size %d result = %#v", size, result) } if size == 1<<20 { - if !bytes.Equal(result.Body, wire) || !bytes.Equal(result.ClassificationBody, plain) || - result.Header.Get("Content-Encoding") != encoding { + if !bytes.Equal(result.Body, plain) || !bytes.Equal(result.ClassificationBody, plain) || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || + result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { t.Fatalf("exact limit result body lengths = %d/%d, headers=%#v", len(result.Body), len(result.ClassificationBody), result.Header) } } else if string(result.Body) != redact.Placeholder || @@ -367,29 +560,27 @@ func TestForwarderBoundsDecompressedErrorBodies(t *testing.T) { func TestForwarderFailsClosedWhenRedactionExpandsErrorBeyondBounds(t *testing.T) { tests := []struct { - name string - encoding string - plain []byte + name string + encoding string + plain []byte + wantPlaceholder bool }{ { - name: "identity wire exceeds limit after redaction", + name: "identity redaction remains within decoded limit", plain: bytes.Repeat([]byte("a"), int(maxErrorResponseBodyBytes)), }, { - name: "gzip classification body exceeds limit after redaction", - encoding: "gzip", - plain: bytes.Repeat([]byte("a"), int(maxDecompressedErrorBodyBytes)), + name: "gzip classification body exceeds limit after redaction", + encoding: "gzip", + plain: bytes.Repeat([]byte("a"), int(maxDecompressedErrorBodyBytes)), + wantPlaceholder: true, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { wire := test.plain if test.encoding != "" { - var err error - wire, err = utils.CompressResponse(test.encoding, test.plain) - if err != nil { - t.Fatal(err) - } + wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(test.encoding), test.plain) } upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { if test.encoding != "" { @@ -402,8 +593,6 @@ func TestForwarderFailsClosedWhenRedactionExpandsErrorBeyondBounds(t *testing.T) result := testForward(t, upstream.URL, "a", time.Second) if result.Err != nil || result.StatusCode != http.StatusUnauthorized || - string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder || result.Header.Get("Content-Encoding") != "" { t.Fatalf( "result status=%d body=%d classification=%d encoding=%q err=%v", @@ -411,21 +600,27 @@ func TestForwarderFailsClosedWhenRedactionExpandsErrorBeyondBounds(t *testing.T) result.Header.Get("Content-Encoding"), result.Err, ) } + if test.wantPlaceholder { + if string(result.Body) != redact.Placeholder || string(result.ClassificationBody) != redact.Placeholder { + t.Fatalf("overflow result = %#v", result) + } + return + } + if len(result.Body) != len(result.ClassificationBody) || int64(len(result.Body)) > maxDecompressedErrorBodyBytes || bytes.Contains(result.Body, []byte("a")) { + t.Fatalf("safe identity result = %#v", result) + } }) } } -func TestPrepareErrorBodyFailsClosedForOversizedUnchangedWire(t *testing.T) { +func TestPrepareErrorRepresentationFailsClosedForOversizedWire(t *testing.T) { plain := make([]byte, 100<<10) const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_." random := mathrand.New(mathrand.NewSource(1)) for index := range plain { plain[index] = alphabet[random.Intn(len(alphabet))] } - gzipWire, err := utils.CompressResponse("gzip", plain) - if err != nil { - t.Fatal(err) - } + gzipWire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) if len(gzipWire) <= int(maxErrorResponseBodyBytes) { t.Fatalf("gzip fixture length = %d, want oversized wire", len(gzipWire)) } @@ -450,12 +645,14 @@ func TestPrepareErrorBodyFailsClosedForOversizedUnchangedWire(t *testing.T) { for _, test := range tests { t.Run(test.name, func(t *testing.T) { headers := test.headers.Clone() - wire, classification := forwarder.prepareErrorBody(headers, test.wire, ForwardInput{}) - if string(wire) != redact.Placeholder || string(classification) != redact.Placeholder || - headers.Get("Content-Encoding") != "" { + prepared := forwarder.prepareErrorRepresentation(ForwardInput{}, headers, test.wire, nil) + if string(prepared.downstream) != redact.Placeholder || + string(prepared.classification) != redact.Placeholder || + len(headerFieldValues(prepared.headers, "Content-Encoding")) != 0 || + !reflect.DeepEqual(headers, test.headers) { t.Fatalf( - "safe body lengths=%d/%d encoding=%q", - len(wire), len(classification), headers.Get("Content-Encoding"), + "safe representation = %#v, source headers = %#v", + prepared, headers, ) } }) @@ -465,11 +662,7 @@ func TestPrepareErrorBodyFailsClosedForOversizedUnchangedWire(t *testing.T) { func compressResponseWithBoundedZstdWindow(t *testing.T, encoding string, plain []byte) []byte { t.Helper() if encoding != "zstd" { - wire, err := utils.CompressResponse(encoding, plain) - if err != nil { - t.Fatal(err) - } - return wire + return encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) } encoder, err := zstd.NewWriter( nil, @@ -825,48 +1018,18 @@ func TestSanitizeForwardResponseHeadersRemovesCaseSensitiveModelFromRawContentTy } } -func TestNonIdentityEncodingContainsKeyHandlesNonCanonicalFieldNames(t *testing.T) { - for _, test := range []struct { - name string - headers http.Header - }{ - { - name: "lowercase encoding triggers success protocol failure", - headers: http.Header{ - "content-encoding": {"gzip"}, - }, - }, - { - name: "duplicate casing aggregates error response encodings", - headers: http.Header{ - "Content-Encoding": {"identity"}, - "cOnTeNt-EnCoDiNg": {"gzip"}, - }, - }, - } { - t.Run(test.name, func(t *testing.T) { - if !nonIdentityEncodingContainsKey(test.headers, "gzip") { - t.Fatalf("nonIdentityEncodingContainsKey(%#v, gzip) = false", test.headers) - } - }) - } -} - -func TestForwarderFailsClosedWhenShortKeyMatchesContentEncoding(t *testing.T) { +func TestForwarderDoesNotTreatShortKeyAsContentCodingFailure(t *testing.T) { tests := []struct { - name string - status int - wantErr bool + name string + status int + wantBody string }{ - {name: "success", status: http.StatusOK, wantErr: true}, - {name: "error", status: http.StatusUnauthorized}, + {name: "success coding header is terminated", status: http.StatusOK, wantBody: `{"error":"safe"}`}, + {name: "error coding header is terminated", status: http.StatusUnauthorized, wantBody: `{"error":"safe"}`}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - wire, err := utils.CompressResponse("gzip", []byte(`{"error":"safe"}`)) - if err != nil { - t.Fatal(err) - } + wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, []byte(`{"error":"safe"}`)) upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Encoding", "gzip") w.WriteHeader(test.status) @@ -874,40 +1037,30 @@ func TestForwarderFailsClosedWhenShortKeyMatchesContentEncoding(t *testing.T) { })) defer upstream.Close() result := testForward(t, upstream.URL, "gzip", time.Second) - if test.wantErr { - if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.RequestWritten { - t.Fatalf("result = %#v", result) - } - return - } if result.Err != nil || result.StatusCode != test.status || - string(result.Body) != redact.Placeholder || result.Header.Get("Content-Encoding") != "" { + string(result.Body) != test.wantBody || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { t.Fatalf("result = %#v", result) } }) } } -func TestShortKeyErrorCollisionRemovesAllContentEncodingCasings(t *testing.T) { +func TestErrorRepresentationRejectsAllContentEncodingCasings(t *testing.T) { headers := http.Header{ "Content-Encoding": {"identity"}, "cOnTeNt-EnCoDiNg": {"gzip"}, } - if !nonIdentityEncodingContainsKey(headers, "identity") { - t.Fatalf("nonIdentityEncodingContainsKey(%#v, identity) = false", headers) + prepared := (&Forwarder{redactor: redact.New()}).prepareErrorRepresentation( + ForwardInput{}, headers, []byte("opaque"), nil, + ) + if string(prepared.downstream) != redact.Placeholder || + string(prepared.classification) != redact.Placeholder { + t.Fatalf("prepared = %#v", prepared) } - - result := UpstreamResult{StatusCode: http.StatusUnauthorized, RequestWritten: true} - result.Body, result.ClassificationBody = failClosedErrorBody(headers) - result.Header = sanitizeUpstreamResponseHeaders(headers, "identity") - - if result.StatusCode != http.StatusUnauthorized || string(result.Body) != redact.Placeholder || - string(result.ClassificationBody) != redact.Placeholder { - t.Fatalf("result = %#v", result) - } - for actualName := range result.Header { + for actualName := range prepared.headers { if strings.EqualFold(actualName, "Content-Encoding") { - t.Fatalf("Content-Encoding survived as %q: %#v", actualName, result.Header[actualName]) + t.Fatalf("Content-Encoding survived as %q: %#v", actualName, prepared.headers[actualName]) } } } @@ -2208,6 +2361,116 @@ func TestForwardStreamReturnsSafeBoundedNonSuccessResponse(t *testing.T) { } } +func TestForwardStreamBuffersCodedNonSuccessRepresentation(t *testing.T) { + plain := []byte(`{"error":{"code":"rate_limited"}}`) + for _, encoding := range []contentcoding.Encoding{ + contentcoding.Identity, + contentcoding.Gzip, + contentcoding.Brotli, + contentcoding.Deflate, + contentcoding.Zstd, + } { + t.Run(string(encoding), func(t *testing.T) { + wire := encodeContentCodingForGatewayTest(t, encoding, plain) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Encoding", string(encoding)) + writer.WriteHeader(http.StatusUnauthorized) + _, _ = writer.Write(wire) + })) + defer upstream.Close() + + downstream := newRecordingResponseWriter() + result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( + context.Background(), streamForwardInput(upstream.URL), downstream, + ) + if result.Err != nil || result.Committed || result.StatusCode != http.StatusUnauthorized || + !bytes.Equal(result.Body, plain) || !bytes.Equal(result.ClassificationBody, plain) || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || + result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { + t.Fatalf("ForwardStream() result = %#v, want buffered plaintext error", result) + } + if downstream.status != 0 || downstream.body.Len() != 0 { + t.Fatalf("ForwardStream() wrote non-success response before Handler verdict: %d/%q", downstream.status, downstream.body.String()) + } + }) + } + + t.Run("multiple content-encoding values fail closed", func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Add("Content-Encoding", "identity") + writer.Header().Add("Content-Encoding", "gzip") + writer.WriteHeader(http.StatusBadGateway) + _, _ = writer.Write([]byte("opaque")) + })) + defer upstream.Close() + + result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( + context.Background(), streamForwardInput(upstream.URL), newRecordingResponseWriter(), + ) + if result.Err != nil || result.Committed || result.StatusCode != http.StatusBadGateway || + string(result.Body) != redact.Placeholder || string(result.ClassificationBody) != redact.Placeholder || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { + t.Fatalf("ForwardStream() result = %#v, want fail-closed buffered error", result) + } + }) +} + +func TestForwardStreamContentCodingBoundary(t *testing.T) { + for _, test := range []struct { + name string + values []string + wantError bool + }{ + {name: "absent"}, + {name: "empty", values: []string{""}}, + {name: "identity", values: []string{"IDENTITY"}}, + {name: "gzip", values: []string{"gzip"}, wantError: true}, + {name: "brotli", values: []string{"br"}, wantError: true}, + {name: "deflate", values: []string{"deflate"}, wantError: true}, + {name: "zstd", values: []string{"zstd"}, wantError: true}, + {name: "stacked", values: []string{"gzip, br"}, wantError: true}, + {name: "unknown", values: []string{"compress"}, wantError: true}, + {name: "multiple", values: []string{"identity", "gzip"}, wantError: true}, + } { + t.Run(test.name, func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + for _, value := range test.values { + writer.Header().Add("Content-Encoding", value) + } + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = writer.Write([]byte("data: first\n\n")) + })) + defer upstream.Close() + + downstream := newRecordingResponseWriter() + result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).ForwardStream( + context.Background(), streamForwardInput(upstream.URL), downstream, + ) + if test.wantError { + if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed || + downstream.status != 0 || downstream.body.Len() != 0 { + t.Fatalf("ForwardStream() result/downstream = %#v/%d/%q", result, downstream.status, downstream.body.String()) + } + return + } + if result.Err != nil || !result.Committed || downstream.status != http.StatusOK || + downstream.body.String() != "data: first\n\n" { + t.Fatalf("ForwardStream() result/downstream = %#v/%d/%q", result, downstream.status, downstream.body.String()) + } + }) + } + + t.Run("case-colliding header map is rejected", func(t *testing.T) { + err := validateStreamContentEncoding(http.Header{ + "Content-Encoding": {"identity"}, + "content-encoding": {"identity"}, + }) + if !errors.Is(err, ErrUpstreamProtocol) { + t.Fatalf("validateStreamContentEncoding() error = %v, want upstream protocol error", err) + } + }) +} + func TestForwardStreamTimesOutBeforeCompleteFirstEvent(t *testing.T) { upstreamCanceled := make(chan struct{}) upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { @@ -2398,6 +2661,16 @@ func TestStreamingClientConfigHasNoTotalTimeout(t *testing.T) { if config.RequestTimeout != 0 { t.Fatalf("stream RequestTimeout = %s, want 0", config.RequestTimeout) } + if !config.DisableCompression { + t.Fatal("stream DisableCompression = false; transport must not implicitly decode while request normalization explicitly negotiates identity") + } +} + +func TestNonStreamingClientConfigDisablesImplicitCompression(t *testing.T) { + config := nonStreamingClientConfig(state.TimeoutConfig{}) + if !config.DisableCompression { + t.Fatal("non-stream DisableCompression = false; transport must not implicitly decode while request normalization explicitly negotiates identity") + } } func TestGatewayClientConfigsDisableRedirects(t *testing.T) { @@ -2503,11 +2776,7 @@ func TestForwarderIsolatesAliasedModelFromNonStreamingErrors(t *testing.T) { t.Run(name, func(t *testing.T) { wire := plain if encoding != "" { - var err error - wire, err = utils.CompressResponse(encoding, plain) - if err != nil { - t.Fatalf("compress fixture: %v", err) - } + wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) } upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { writer.Header().Set("Content-Type", "application/json") @@ -2539,13 +2808,6 @@ func TestForwarderIsolatesAliasedModelFromNonStreamingErrors(t *testing.T) { t.Fatalf("ClassificationBody = %q", result.ClassificationBody) } downstreamBody := result.Body - if encoding != "" { - var err error - downstreamBody, err = utils.DecompressResponse(encoding, result.Body) - if err != nil { - t.Fatalf("decompress downstream body: %v", err) - } - } if bytes.Contains(downstreamBody, []byte(upstreamModel)) || !bytes.Contains(downstreamBody, []byte(externalModel)) || bytes.Contains(downstreamBody, []byte(secret)) || @@ -2553,7 +2815,7 @@ func TestForwarderIsolatesAliasedModelFromNonStreamingErrors(t *testing.T) { t.Fatalf("downstream body = %q", downstreamBody) } if result.Header.Get("Content-Type") != "application/json" || - result.Header.Get("Content-Encoding") != encoding || + result.Header.Get("Content-Encoding") != "" || result.Header.Get("Content-Length") != strconv.Itoa(len(result.Body)) { t.Fatalf("representation headers = %#v", result.Header) } @@ -2573,11 +2835,7 @@ func TestForwarderRewritesEscapedAliasedModelInJSONErrors(t *testing.T) { wire := plain if encoding != "" { name = encoding - var err error - wire, err = utils.CompressResponse(encoding, plain) - if err != nil { - t.Fatalf("compress fixture: %v", err) - } + wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) } t.Run(name, func(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { @@ -2603,13 +2861,6 @@ func TestForwarderRewritesEscapedAliasedModelInJSONErrors(t *testing.T) { t.Fatalf("ClassificationBody = %q, want original safe JSON %q", result.ClassificationBody, plain) } downstreamBody := result.Body - if encoding != "" { - var err error - downstreamBody, err = utils.DecompressResponse(encoding, result.Body) - if err != nil { - t.Fatalf("decompress downstream body: %v", err) - } - } if !json.Valid(downstreamBody) { t.Fatalf("downstream body is invalid JSON: %q", downstreamBody) } @@ -2640,11 +2891,7 @@ func TestForwarderRedactsEscapedAPIKeyInJSONErrors(t *testing.T) { wire := plain if encoding != "" { name = encoding - var err error - wire, err = utils.CompressResponse(encoding, plain) - if err != nil { - t.Fatalf("compress fixture: %v", err) - } + wire = encodeContentCodingForGatewayTest(t, contentcoding.Encoding(encoding), plain) } t.Run(name, func(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { @@ -2666,13 +2913,6 @@ func TestForwarderRedactsEscapedAPIKeyInJSONErrors(t *testing.T) { t.Fatalf("Forward() result = %#v", result) } downstreamBody := result.Body - if encoding != "" { - var err error - downstreamBody, err = utils.DecompressResponse(encoding, result.Body) - if err != nil { - t.Fatalf("decompress downstream body: %v", err) - } - } for _, body := range [][]byte{downstreamBody, result.ClassificationBody} { var decoded map[string]any if err := json.Unmarshal(body, &decoded); err != nil { @@ -2686,27 +2926,27 @@ func TestForwarderRedactsEscapedAPIKeyInJSONErrors(t *testing.T) { } } -func TestPrepareErrorBodyFailsClosedWhenEscapedAPIKeyRewriteCollides(t *testing.T) { +func TestPrepareErrorRepresentationFailsClosedWhenEscapedAPIKeyRewriteCollides(t *testing.T) { const secret = "secret/key" plain := []byte(`{"secret\/key":"leak","[REDACTED]":"safe"}`) forwarder := &Forwarder{redactor: redact.New()} headers := http.Header{"Content-Type": {"application/json"}} - wire, classification := forwarder.prepareErrorBody(headers, plain, ForwardInput{APIKey: secret}) + prepared := forwarder.prepareErrorRepresentation(ForwardInput{APIKey: secret}, headers, plain, []string{secret}) - if string(wire) != redact.Placeholder || string(classification) != redact.Placeholder { - t.Fatalf("collision result wire=%q classification=%q", wire, classification) + if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { + t.Fatalf("collision result = %#v", prepared) } } -func TestPrepareErrorBodyFailsClosedBeforeRawAPIKeyCollisionRedaction(t *testing.T) { +func TestPrepareErrorRepresentationFailsClosedBeforeRawAPIKeyCollisionRedaction(t *testing.T) { const secret = "secret" plain := []byte(`{"secret":"leak","[REDACTED]":"safe"}`) forwarder := &Forwarder{redactor: redact.New()} headers := http.Header{"Content-Type": {"application/json"}} - wire, classification := forwarder.prepareErrorBody(headers, plain, ForwardInput{APIKey: secret}) + prepared := forwarder.prepareErrorRepresentation(ForwardInput{APIKey: secret}, headers, plain, []string{secret}) - if string(wire) != redact.Placeholder || string(classification) != redact.Placeholder { - t.Fatalf("collision result wire=%q classification=%q", wire, classification) + if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { + t.Fatalf("collision result = %#v", prepared) } } @@ -2714,25 +2954,25 @@ func TestForwarderFailsClosedWhenJSONKeyRewriteCollides(t *testing.T) { plain := []byte(`{"provider-model":"first","public-model":"second"}`) forwarder := &Forwarder{redactor: redact.New()} headers := http.Header{"Content-Type": {"application/json"}} - wire, classification := forwarder.prepareErrorBody(headers, plain, ForwardInput{ + prepared := forwarder.prepareErrorRepresentation(ForwardInput{ ExternalModel: "public-model", UpstreamModelID: "provider-model", - }) + }, headers, plain, nil) - if string(wire) != redact.Placeholder || !bytes.Equal(classification, plain) { - t.Fatalf("collision result wire=%q classification=%q", wire, classification) + if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { + t.Fatalf("collision result = %#v", prepared) } } -func TestPrepareErrorBodyFailsClosedWhenModelExpansionExceedsLimit(t *testing.T) { +func TestPrepareErrorRepresentationFailsClosedWhenModelExpansionExceedsLimit(t *testing.T) { plain := bytes.Repeat([]byte("x"), 64<<10) forwarder := &Forwarder{redactor: redact.New()} headers := make(http.Header) - wire, classification := forwarder.prepareErrorBody(headers, plain, ForwardInput{ + prepared := forwarder.prepareErrorRepresentation(ForwardInput{ ExternalModel: strings.Repeat("a", 32), UpstreamModelID: "x", - }) + }, headers, plain, nil) - if string(wire) != redact.Placeholder || !bytes.Equal(classification, plain) { - t.Fatalf("expansion result wire=%q classification length=%d", wire, len(classification)) + if string(prepared.downstream) != redact.Placeholder || string(prepared.classification) != redact.Placeholder { + t.Fatalf("expansion result = %#v", prepared) } } @@ -2803,13 +3043,10 @@ func TestForwarderPreservesErrorModelWithoutAlias(t *testing.T) { } } -func TestForwarderRedactsCompressedErrorAndPreservesEncoding(t *testing.T) { +func TestForwarderRedactsCompressedErrorAsPlaintext(t *testing.T) { const secret = "custom-upstream-secret" plain := []byte(`{"error":{"api_key":"` + secret + `","code":"invalid_api_key"}}`) - encoded, err := utils.CompressResponse("gzip", plain) - if err != nil { - t.Fatalf("compress fixture: %v", err) - } + encoded := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { writer.Header().Set("Content-Encoding", "gzip") setRepresentationMetadata(writer.Header()) @@ -2822,14 +3059,10 @@ func TestForwarderRedactsCompressedErrorAndPreservesEncoding(t *testing.T) { if result.Err != nil || result.StatusCode != http.StatusUnauthorized { t.Fatalf("Forward() result = %#v", result) } - if result.Header.Get("Content-Encoding") != "gzip" { - t.Fatalf("Content-Encoding = %q, want gzip", result.Header.Get("Content-Encoding")) + if result.Header.Get("Content-Encoding") != "" { + t.Fatalf("Content-Encoding = %q, want absent", result.Header.Get("Content-Encoding")) } - decoded, err := utils.DecompressResponse("gzip", result.Body) - if err != nil { - t.Fatalf("decode forwarded body: %v", err) - } - for _, body := range [][]byte{decoded, result.ClassificationBody} { + for _, body := range [][]byte{result.Body, result.ClassificationBody} { if bytes.Contains(body, []byte(secret)) || !bytes.Contains(body, []byte(redact.Placeholder)) { t.Fatalf("safe body = %q, want placeholder and no secret", body) } @@ -2840,12 +3073,121 @@ func TestForwarderRedactsCompressedErrorAndPreservesEncoding(t *testing.T) { assertRepresentationMetadata(t, result.Header, false) } -func TestForwarderPreservesUnchangedCompressedErrorWireBytes(t *testing.T) { +func TestForwarderErrorRepresentationsArePlain(t *testing.T) { plain := []byte(`{"error":{"code":"rate_limited"}}`) - encoded, err := utils.CompressResponse("gzip", plain) - if err != nil { - t.Fatalf("compress fixture: %v", err) + for _, encoding := range []contentcoding.Encoding{ + contentcoding.Identity, + contentcoding.Gzip, + contentcoding.Brotli, + contentcoding.Deflate, + contentcoding.Zstd, + } { + t.Run(string(encoding), func(t *testing.T) { + wire := encodeContentCodingForGatewayTest(t, encoding, plain) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + if encoding != contentcoding.Identity { + writer.Header().Set("Content-Encoding", string(encoding)) + } + writer.WriteHeader(http.StatusTooManyRequests) + _, _ = writer.Write(wire) + })) + defer upstream.Close() + + result := testForward(t, upstream.URL, "safe-upstream-key", time.Second) + if result.Err != nil || result.StatusCode != http.StatusTooManyRequests || + !bytes.Equal(result.Body, plain) || !bytes.Equal(result.ClassificationBody, plain) || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || + result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { + t.Fatalf("Forward() result = %#v, want plaintext error representation", result) + } + }) } +} + +func TestForwarderErrorRepresentationMetadataPolicy(t *testing.T) { + plain := []byte(`{"error":"safe"}`) + for _, test := range []struct { + name string + encoding contentcoding.Encoding + wantMetadataKept bool + }{ + {name: "unchanged identity", encoding: contentcoding.Identity, wantMetadataKept: true}, + {name: "decoded gzip", encoding: contentcoding.Gzip, wantMetadataKept: false}, + } { + t.Run(test.name, func(t *testing.T) { + wire := encodeContentCodingForGatewayTest(t, test.encoding, plain) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Encoding", string(test.encoding)) + setRepresentationMetadata(writer.Header()) + writer.WriteHeader(http.StatusBadGateway) + _, _ = writer.Write(wire) + })) + defer upstream.Close() + + result := testForward(t, upstream.URL, "safe-upstream-key", time.Second) + if result.Err != nil || !bytes.Equal(result.Body, plain) { + t.Fatalf("Forward() result = %#v", result) + } + if test.wantMetadataKept { + for name, want := range map[string]string{ + "ETag": `"wire-v1"`, + "Digest": "sha-256=wire-digest", + "Content-MD5": "d2lyZQ==", + "Content-Range": "bytes 0-9/10", + "Content-Digest": "sha-256=:d2lyZQ==:", + "Repr-Digest": "sha-256=:cmVwcg==:", + } { + if got := result.Header.Get(name); got != want { + t.Fatalf("%s = %q, want %q", name, got, want) + } + } + if result.Header.Get("Signature") != "" || result.Header.Get("Signature-Input") != "" { + t.Fatalf("normalized identity representation retained stale signatures: %#v", result.Header) + } + } else { + assertRepresentationMetadata(t, result.Header, false) + } + }) + } +} + +func TestForwarderErrorContentCodingFailsClosed(t *testing.T) { + for _, test := range []struct { + name string + coding string + body []byte + apiKey string + }{ + {name: "unknown", coding: "unknown", body: []byte("opaque")}, + {name: "stacked", coding: "gzip, br", body: []byte("opaque")}, + {name: "corrupt gzip", coding: "gzip", body: []byte("not-gzip")}, + {name: "short key no special bypass", coding: "gzip", body: encodeContentCodingForGatewayTest(t, contentcoding.Gzip, []byte(`{"error":"safe"}`)), apiKey: "gzip"}, + } { + t.Run(test.name, func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Encoding", test.coding) + writer.WriteHeader(http.StatusUnauthorized) + _, _ = writer.Write(test.body) + })) + defer upstream.Close() + + result := testForward(t, upstream.URL, test.apiKey, time.Second) + want := []byte(redact.Placeholder) + if test.name == "short key no special bypass" { + want = []byte(`{"error":"safe"}`) + } + if result.Err != nil || result.StatusCode != http.StatusUnauthorized || + !bytes.Equal(result.Body, want) || !bytes.Equal(result.ClassificationBody, want) || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { + t.Fatalf("Forward() result = %#v, want safe error representation", result) + } + }) + } +} + +func TestForwarderReturnsUnchangedCompressedErrorAsPlaintext(t *testing.T) { + plain := []byte(`{"error":{"code":"rate_limited"}}`) + encoded := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { writer.Header().Set("Content-Encoding", "gzip") setRepresentationMetadata(writer.Header()) @@ -2858,17 +3200,17 @@ func TestForwarderPreservesUnchangedCompressedErrorWireBytes(t *testing.T) { if result.Err != nil || result.StatusCode != http.StatusTooManyRequests { t.Fatalf("Forward() result = %#v", result) } - if !bytes.Equal(result.Body, encoded) { - t.Fatalf("wire body changed without redaction: got %x want %x", result.Body, encoded) + if !bytes.Equal(result.Body, plain) { + t.Fatalf("downstream body = %q, want %q", result.Body, plain) } if !bytes.Equal(result.ClassificationBody, plain) { t.Fatalf("ClassificationBody = %q, want %q", result.ClassificationBody, plain) } - if result.Header.Get("Content-Encoding") != "gzip" || - result.Header.Get("Content-Length") != strconv.Itoa(len(encoded)) { - t.Fatalf("compressed response headers = %#v", result.Header) + if len(headerFieldValues(result.Header, "Content-Encoding")) != 0 || + result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { + t.Fatalf("plaintext response headers = %#v", result.Header) } - assertRepresentationMetadata(t, result.Header, true) + assertRepresentationMetadata(t, result.Header, false) } func TestForwarderFailsClosedForUndecodableError(t *testing.T) { @@ -2966,6 +3308,173 @@ func TestForwarderFailsClosedForGzipEmptyBody(t *testing.T) { } } +func TestForwarderBodylessResponseSemantics(t *testing.T) { + tests := []struct { + name string + method string + status int + contentEncoding string + contentLength string + wantErr bool + wantContentLength string + }{ + { + name: "head preserves oversized representation length without reading it", + method: http.MethodHead, + status: http.StatusOK, + contentEncoding: "identity", + contentLength: strconv.FormatInt(maxNonStreamingResponseBodyBytes+1, 10), + wantContentLength: strconv.FormatInt(maxNonStreamingResponseBodyBytes+1, 10), + }, + { + name: "not modified is bodyless", + method: http.MethodGet, + status: http.StatusNotModified, + contentEncoding: "identity", + contentLength: "321", + }, + { + name: "no content removes length", + method: http.MethodGet, + status: http.StatusNoContent, + contentEncoding: "identity", + contentLength: "321", + }, + { + name: "reset content keeps zero length only", + method: http.MethodGet, + status: http.StatusResetContent, + contentEncoding: "identity", + contentLength: "0", + wantContentLength: "0", + }, + { + name: "bodyless response rejects compressed encoding", + method: http.MethodHead, + status: http.StatusOK, + contentEncoding: "gzip", + wantErr: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != test.method { + t.Errorf("request method = %s, want %s", request.Method, test.method) + } + if test.contentEncoding != "" { + writer.Header().Set("Content-Encoding", test.contentEncoding) + } + if test.contentLength != "" { + writer.Header().Set("Content-Length", test.contentLength) + } + writer.Header().Set("ETag", `"safe"`) + writer.WriteHeader(test.status) + _, _ = writer.Write([]byte("must not be read")) + })) + defer upstream.Close() + + input := streamForwardInput(upstream.URL) + input.ObserveUsage = false + input.Request.Method = test.method + result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( + context.Background(), input, + ) + + if test.wantErr { + if !errors.Is(result.Err, ErrUpstreamProtocol) || !result.RequestWritten { + t.Fatalf("Forward() result = %#v, want written upstream protocol error", result) + } + return + } + if result.Err != nil || result.StatusCode != test.status || !result.RequestWritten || len(result.Body) != 0 || len(result.ClassificationBody) != 0 { + t.Fatalf("Forward() result = %#v, want bodyless response", result) + } + if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { + t.Fatalf("Content-Encoding values = %#v, want absent", values) + } + if got := result.Header.Get("Content-Length"); got != test.wantContentLength { + t.Fatalf("Content-Length = %q, want %q", got, test.wantContentLength) + } + if result.Header.Get("ETag") != `"safe"` { + t.Fatalf("ETag = %q, want safe validator preserved", result.Header.Get("ETag")) + } + }) + } +} + +func TestForwarderBodylessResponseInvalidatesSignaturesAfterHeaderNormalization(t *testing.T) { + tests := []struct { + name string + method string + status int + contentEncoding string + contentLength string + credentialHeader bool + }{ + { + name: "head removes identity coding", + method: http.MethodHead, + status: http.StatusOK, + contentEncoding: "identity", + contentLength: "12", + }, + { + name: "not modified removes identity coding", + method: http.MethodGet, + status: http.StatusNotModified, + contentEncoding: "identity", + contentLength: "12", + }, + { + name: "reset content removes nonzero length", + method: http.MethodGet, + status: http.StatusResetContent, + contentLength: "12", + }, + { + name: "head sanitizes credential header", + method: http.MethodHead, + status: http.StatusOK, + contentLength: "12", + credentialHeader: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + if test.contentEncoding != "" { + writer.Header().Set("Content-Encoding", test.contentEncoding) + } + if test.contentLength != "" { + writer.Header().Set("Content-Length", test.contentLength) + } + writer.Header().Set("Signature", "sig") + writer.Header().Set("Signature-Input", `sig1=("content-encoding" "content-length")`) + if test.credentialHeader { + writer.Header().Set("Authorization", "Bearer upstream-secret") + } + writer.WriteHeader(test.status) + })) + defer upstream.Close() + + input := streamForwardInput(upstream.URL) + input.ObserveUsage = false + input.Request.Method = test.method + result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( + context.Background(), input, + ) + if result.Err != nil || result.StatusCode != test.status || !result.RequestWritten || + len(headerFieldValues(result.Header, "Signature")) != 0 || + len(headerFieldValues(result.Header, "Signature-Input")) != 0 { + t.Fatalf("Forward() result = %#v, want bodyless result without stale signatures", result) + } + }) + } +} + func TestForwarderMarksConnectionFailureAsNotWritten(t *testing.T) { upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) url := upstream.URL @@ -2990,6 +3499,109 @@ func TestForwarderMarksTimeoutAfterRequestWrite(t *testing.T) { } } +func TestForwarderDecodesContentEncodingListedAsConnectionToken(t *testing.T) { + plain := []byte(`{"id":"safe"}`) + clients := platformhttp.NewHTTPClientManager() + input := ForwardInput{ + Dialect: dialect.NewOpenAI(http.DefaultClient), + Group: state.GroupView{ + ID: 1, UpstreamURL: "https://api.example.test", + Timeouts: state.TimeoutConfig{Connect: time.Second, FirstByte: time.Second, Request: time.Second}, + }, + APIKey: "sk-upstream-secret", + Request: &dialect.ParsedRequest{ + Method: http.MethodPost, + Path: "/v1/chat/completions", + Header: make(http.Header), + Body: []byte(`{"model":"gpt-4o"}`), + }, + } + clients.GetClient(nonStreamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Connection": {"Content-Encoding"}, + "Content-Encoding": {"gzip"}, + }, + Body: io.NopCloser(bytes.NewReader(encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain))), + Request: request, + }, nil + }) + + result := NewForwarder(clients, redact.New()).Forward(context.Background(), input) + if result.Err != nil || !bytes.Equal(result.Body, plain) || len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { + t.Fatalf("Forward() result = %#v, want plaintext response without Content-Encoding", result) + } +} + +func TestForwarderBodylessResponseInvalidatesSignaturesAfterConnectionTokenCleanup(t *testing.T) { + clients := platformhttp.NewHTTPClientManager() + input := streamForwardInput("https://api.example.test") + input.ObserveUsage = false + input.Request.Method = http.MethodHead + clients.GetClient(nonStreamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Connection": {"X-Signed"}, + "X-Signed": {"upstream-only"}, + "Content-Length": {"12"}, + "Signature": {"sig"}, + "Signature-Input": {`sig1=("x-signed" "content-length")`}, + }, + Body: io.NopCloser(strings.NewReader("")), + Request: request, + }, nil + }) + + result := NewForwarder(clients, redact.New()).Forward(context.Background(), input) + if result.Err != nil || result.StatusCode != http.StatusOK || !result.RequestWritten { + t.Fatalf("Forward() result = %#v, want bodyless forwarded response", result) + } + if len(headerFieldValues(result.Header, "X-Signed")) != 0 || + len(headerFieldValues(result.Header, "Signature")) != 0 || + len(headerFieldValues(result.Header, "Signature-Input")) != 0 { + t.Fatalf("bodyless response retained hop-by-hop cleanup or stale signatures: %#v", result.Header) + } + if result.Header.Get("Content-Length") != "12" { + t.Fatalf("Content-Length = %q, want preserved bodyless representation length", result.Header.Get("Content-Length")) + } +} + +func TestForwardStreamRejectsContentEncodingListedAsConnectionToken(t *testing.T) { + plain := []byte("data: first\n\n") + wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) + clients := platformhttp.NewHTTPClientManager() + input := streamForwardInput("https://api.example.test") + clients.GetClient(streamingClientConfig(input.Group.Timeouts)).Transport = forwarderRoundTripper(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Connection": {"Content-Encoding"}, + "Content-Encoding": {"gzip"}, + }, + Body: io.NopCloser(bytes.NewReader(wire)), + ContentLength: int64(len(wire)), + Request: request, + }, nil + }) + + downstream := newRecordingResponseWriter() + result := NewForwarder(clients, redact.New()).ForwardStream(context.Background(), input, downstream) + if !errors.Is(result.Err, ErrUpstreamProtocol) || result.Committed || result.RetryableBeforeCommit { + t.Fatalf("ForwardStream() result = %#v, want terminal pre-commit protocol error", result) + } + if downstream.status != 0 || downstream.body.Len() != 0 || downstream.flushes != 0 { + t.Fatalf("downstream was touched before protocol rejection: %#v", downstream) + } +} + +type forwarderRoundTripper func(*http.Request) (*http.Response, error) + +func (roundTrip forwarderRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + return roundTrip(request) +} + func testForward(t *testing.T, upstreamURL, apiKey string, timeout time.Duration) UpstreamResult { t.Helper() forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) @@ -3146,18 +3758,16 @@ func TestForwardStreamInjectsUsageAfterModelAliasRewrite(t *testing.T) { } } -func TestForwardStreamInvalidatesRequestRepresentationMetadataOnlyWhenBodyChanges(t *testing.T) { +func TestForwardStreamNormalizesFinalRequestMetadata(t *testing.T) { type capturedRequest struct { body []byte headers http.Header contentLength int64 } tests := []struct { - name string - configure func(*ForwardInput) - wantBodyChanged bool - wantMetadataPreserved bool - wantReadded map[string]string + name string + configure func(*ForwardInput) + wantBodyChanged bool }{ { name: "usage injection invalidates stale metadata", @@ -3167,11 +3777,10 @@ func TestForwardStreamInvalidatesRequestRepresentationMetadataOnlyWhenBodyChange wantBodyChanged: true, }, { - name: "disabled injection preserves metadata", - wantMetadataPreserved: true, + name: "disabled injection still removes stale metadata", }, { - name: "failed injection preserves metadata when body is unchanged", + name: "failed injection still removes stale metadata when body is unchanged", configure: func(input *ForwardInput) { input.Group.InjectUsageOptions = true input.Dialect = streamUsageInjectorDialect{ @@ -3181,7 +3790,6 @@ func TestForwardStreamInvalidatesRequestRepresentationMetadataOnlyWhenBodyChange }, } }, - wantMetadataPreserved: true, }, { name: "model alias rewrite invalidates stale metadata", @@ -3193,7 +3801,7 @@ func TestForwardStreamInvalidatesRequestRepresentationMetadataOnlyWhenBodyChange wantBodyChanged: true, }, { - name: "group header rules can add fresh metadata", + name: "group header rules cannot restore stale metadata", configure: func(input *ForwardInput) { input.Group.InjectUsageOptions = true input.Group.HeaderRules = state.HeaderRules{Set: map[string]string{ @@ -3202,10 +3810,6 @@ func TestForwardStreamInvalidatesRequestRepresentationMetadataOnlyWhenBodyChange }} }, wantBodyChanged: true, - wantReadded: map[string]string{ - "Digest": "sha-256=group-digest", - "Signature": "group-signature", - }, }, } @@ -3247,20 +3851,7 @@ func TestForwardStreamInvalidatesRequestRepresentationMetadataOnlyWhenBodyChange t.Fatalf("ContentLength = %d, want %d", got.contentLength, len(got.body)) } - if len(test.wantReadded) == 0 { - assertRepresentationMetadata(t, got.headers, test.wantMetadataPreserved) - return - } - for name, value := range test.wantReadded { - if got.headers.Get(name) != value { - t.Errorf("%s = %q, want fresh Group HeaderRules value %q", name, got.headers.Get(name), value) - } - } - withoutReadded := got.headers.Clone() - for name := range test.wantReadded { - withoutReadded.Del(name) - } - assertRepresentationMetadata(t, withoutReadded, false) + assertRepresentationMetadata(t, got.headers, false) }) } } @@ -3383,11 +3974,20 @@ func TestInvalidateRewrittenBodyHeadersRemovesRepresentationMetadata(t *testing. headers := make(http.Header) headers.Set("Content-Length", "123") setRepresentationMetadata(headers) + for _, name := range []string{ + "content-length", "eTAG", "dIGEST", "content-md5", "content-range", + "content-digest", "repr-digest", "signature", "signature-input", + } { + headers[name] = []string{"case-colliding-stale-value"} + } invalidateRewrittenBodyHeaders(headers) - if headers.Get("Content-Length") != "" { - t.Fatalf("Content-Length = %q, want removed", headers.Get("Content-Length")) + for actualName := range headers { + if isRepresentationMetadataHeader(actualName) && + !strings.EqualFold(actualName, "Content-Encoding") { + t.Fatalf("stale metadata survived as %q: %#v", actualName, headers[actualName]) + } } assertRepresentationMetadata(t, headers, false) } diff --git a/internal/gateway/handler.go b/internal/gateway/handler.go index 4d493d04..d67505d7 100644 --- a/internal/gateway/handler.go +++ b/internal/gateway/handler.go @@ -5,9 +5,11 @@ import ( "errors" "fmt" "io" + "math" "math/rand" "net/http" "strconv" + "strings" "time" "github.com/gin-gonic/gin" @@ -15,6 +17,7 @@ import ( "gpt-load/internal/dialect" "gpt-load/internal/health" + "gpt-load/internal/platform/contentcoding" "gpt-load/internal/platform/encryption" "gpt-load/internal/platform/utils" "gpt-load/internal/pricing" @@ -270,6 +273,12 @@ func (handler *Handler) Handle(ginContext *gin.Context) { return } if selectedRoute.Kind == endpointModels { + if !contentcoding.IdentityAcceptable( + headerFieldValues(ginContext.Request.Header, "Accept-Encoding"), + ) { + handler.completeReason(ginContext, recorder, reasonNotAcceptable) + return + } handler.writeVisibleModelList(ginContext, snapshot, accessKey, selectedRoute.Protocol) return } @@ -283,19 +292,34 @@ func (handler *Handler) Handle(ginContext *gin.Context) { handler.completeReason(ginContext, recorder, reasonEndpointNotFound) return } - - candidateGroupIDs := scheduler.CandidateGroupIDs( - snapshot, - selectedRoute.Protocol, - accessKey, + encoding, err := contentcoding.ParseContentEncoding( + headerFieldValues(ginContext.Request.Header, "Content-Encoding"), ) - capturedRefs := handler.registry.CaptureActiveKeyRefs(candidateGroupIDs) - allowedKeyRefs := make(map[uint]state.KeyRef, len(capturedRefs)) - for _, ref := range capturedRefs { - allowedKeyRefs[ref.ID] = ref + if err != nil { + ginContext.Writer.Header().Set( + "Accept-Encoding", + contentcoding.SupportedRequestEncodings, + ) + handler.completeReason( + ginContext, + recorder, + reasonUnsupportedContentEncoding, + ) + return + } + if !contentcoding.IdentityAcceptable( + headerFieldValues(ginContext.Request.Header, "Accept-Encoding"), + ) { + handler.completeReason(ginContext, recorder, reasonNotAcceptable) + return } - body, err := readRequestBody(ginContext.Request.Body, maxRequestBodyBytes) + body, err := readDecodedRequestBody( + ginContext.Request, + encoding, + maxRequestBodyBytes, + maxRequestBodyBytes, + ) if err != nil { if ginContext.Request.Context().Err() != nil { recorder.completeCanceled(0) @@ -305,14 +329,21 @@ func (handler *Handler) Handle(ginContext *gin.Context) { handler.completeReason(ginContext, recorder, reasonRequestTooLarge) return } + if errors.Is(err, contentcoding.ErrInvalidContentEncoding) { + handler.completeReason(ginContext, recorder, reasonInvalidContentEncoding) + return + } handler.completeReason(ginContext, recorder, reasonInvalidProtocolRequest) return } + requestHeaders := ginContext.Request.Header.Clone() + deleteHeaderField(requestHeaders, "Content-Encoding") + deleteHeaderField(requestHeaders, "Content-Length") parsed := &dialect.ParsedRequest{ Method: ginContext.Request.Method, Path: ginContext.Request.URL.Path, RawQuery: ginContext.Request.URL.RawQuery, - Header: ginContext.Request.Header.Clone(), + Header: requestHeaders, Body: body, } metadata, err := selectedDialect.InspectRequest(parsed) @@ -332,6 +363,20 @@ func (handler *Handler) Handle(ginContext *gin.Context) { handler.completeReason(ginContext, recorder, reasonInvalidProtocolRequest) return } + if ginContext.Request.Context().Err() != nil { + recorder.completeCanceled(0) + return + } + candidateGroupIDs := scheduler.CandidateGroupIDs( + snapshot, + selectedRoute.Protocol, + accessKey, + ) + capturedRefs := handler.registry.CaptureActiveKeyRefs(candidateGroupIDs) + allowedKeyRefs := make(map[uint]state.KeyRef, len(capturedRefs)) + for _, ref := range capturedRefs { + allowedKeyRefs[ref.ID] = ref + } recorder.setClientModel(model) recorder.setUsageApplicable(metadata.ObserveUsage) recorder.setUsageDiagnostics(metadata.UsageDiagnostics) @@ -403,21 +448,50 @@ func (handler *Handler) completeWriteTerminal( recorder.completeDownstreamWrite(selectedStatus) } -func readRequestBody(reader io.Reader, limit int64) ([]byte, error) { - if reader == nil { +func readDecodedRequestBody( + request *http.Request, + encoding contentcoding.Encoding, + encodedLimit int64, + decodedLimit int64, +) ([]byte, error) { + if request == nil || request.Body == nil { return nil, fmt.Errorf("request body is required") } - if limit < 0 { - return nil, fmt.Errorf("request body limit must not be negative") + if encodedLimit < 0 { + return nil, fmt.Errorf("encoded request body limit must not be negative") } - body, err := io.ReadAll(io.LimitReader(reader, limit+1)) + if decodedLimit < 0 { + return nil, fmt.Errorf("decoded request body limit must not be negative") + } + if request.ContentLength > 0 && request.ContentLength > encodedLimit { + return nil, errRequestTooLarge + } + reader := io.Reader(request.Body) + if encodedLimit < math.MaxInt64 { + reader = io.LimitReader(request.Body, encodedLimit+1) + } + encoded, err := io.ReadAll(reader) if err != nil { return nil, fmt.Errorf("read request body: %w", err) } - if int64(len(body)) > limit { + if int64(len(encoded)) > encodedLimit { return nil, errRequestTooLarge } - return body, nil + body, err := contentcoding.DecodeLimited(encoding, encoded, decodedLimit) + if errors.Is(err, contentcoding.ErrDecodedBodyTooLarge) { + return nil, errRequestTooLarge + } + return body, err +} + +func headerFieldValues(headers http.Header, name string) []string { + var values []string + for actualName, fieldValues := range headers { + if strings.EqualFold(actualName, name) { + values = append(values, fieldValues...) + } + } + return values } func (handler *Handler) executeAttempts( @@ -685,7 +759,12 @@ func (handler *Handler) writeBufferedResponse( } }() - for name, values := range cloneEndToEndHeaders(headers) { + method := "" + if ginContext.Request != nil { + method = ginContext.Request.Method + } + normalizedHeaders, writeBody := normalizeBufferedResponse(method, status, headers, body) + for name, values := range normalizedHeaders { for _, value := range values { ginContext.Writer.Header().Add(name, value) } @@ -694,7 +773,7 @@ func (handler *Handler) writeBufferedResponse( return fmt.Errorf("write downstream response headers: %w", err) } ginContext.Writer.WriteHeaderNow() - if len(body) > 0 { + if writeBody && len(body) > 0 { written, writeErr := controlled.write(body) if writeErr != nil { return fmt.Errorf("write downstream response: %w", writeErr) diff --git a/internal/gateway/handler_test.go b/internal/gateway/handler_test.go index d72892b2..090c3050 100644 --- a/internal/gateway/handler_test.go +++ b/internal/gateway/handler_test.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "math" "math/rand" "net/http" "net/http/httptest" @@ -22,11 +23,13 @@ import ( "gpt-load/internal/dialect" "gpt-load/internal/health" "gpt-load/internal/platform/config" + "gpt-load/internal/platform/contentcoding" "gpt-load/internal/platform/encryption" platformhttp "gpt-load/internal/platform/httpclient" "gpt-load/internal/platform/redact" "gpt-load/internal/platform/utils" "gpt-load/internal/protocol" + "gpt-load/internal/ratelimit" "gpt-load/internal/state" "gpt-load/internal/testutil/fakeupstream" "gpt-load/internal/usage" @@ -123,6 +126,33 @@ type recordingRuntimeRegistry struct { clearCalls int } +type captureCountingRuntimeRegistry struct { + runtimeKeyRegistry + captureCalls int +} + +type cancelingSuccessfulInspectDialect struct { + dialect.Dialect + cancel context.CancelFunc +} + +func (value *cancelingSuccessfulInspectDialect) InspectRequest( + request *dialect.ParsedRequest, +) (dialect.RequestMetadata, error) { + metadata, err := value.Dialect.InspectRequest(request) + if err == nil { + value.cancel() + } + return metadata, err +} + +func (registry *captureCountingRuntimeRegistry) CaptureActiveKeyRefs( + groupIDs []uint, +) []state.KeyRef { + registry.captureCalls++ + return registry.runtimeKeyRegistry.CaptureActiveKeyRefs(groupIDs) +} + type gatewayMutationObservation struct { cooldownCalls int clearCalls int @@ -1580,24 +1610,517 @@ func assertDebugHeaders(t *testing.T, headers http.Header, group, attempts strin } } -func TestReadRequestBodyHonorsLimit(t *testing.T) { +func TestReadDecodedRequestBodyHonorsIndependentLimits(t *testing.T) { t.Run("exact limit is accepted", func(t *testing.T) { - body, err := readRequestBody(strings.NewReader("1234"), 4) + request := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("1234")) + body, err := readDecodedRequestBody(request, contentcoding.Identity, 4, 4) if err != nil || string(body) != "1234" { - t.Fatalf("readRequestBody() = %q, %v", body, err) + t.Fatalf("readDecodedRequestBody() = %q, %v", body, err) } }) t.Run("limit plus one is rejected without draining", func(t *testing.T) { reader := &boundedCountingReader{remaining: 100} - body, err := readRequestBody(reader, 4) + request := httptest.NewRequest(http.MethodPost, "/", reader) + request.ContentLength = -1 + body, err := readDecodedRequestBody(request, contentcoding.Identity, 4, 4) if !errors.Is(err, errRequestTooLarge) || body != nil { - t.Fatalf("readRequestBody() = %q, %v, want request too large", body, err) + t.Fatalf("readDecodedRequestBody() = %q, %v, want request too large", body, err) } if reader.read != 5 { t.Fatalf("reader consumed %d bytes, want limit+1 (5)", reader.read) } }) + + t.Run("positive ContentLength rejects before reading", func(t *testing.T) { + reader := &boundedCountingReader{remaining: 100} + request := httptest.NewRequest(http.MethodPost, "/", reader) + request.ContentLength = 5 + body, err := readDecodedRequestBody(request, contentcoding.Identity, 4, 4) + if !errors.Is(err, errRequestTooLarge) || body != nil || reader.read != 0 { + t.Fatalf("readDecodedRequestBody() = %q, %v; read=%d, want early rejection", body, err, reader.read) + } + }) + + t.Run("decoded overflow maps to request too large", func(t *testing.T) { + encoded := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, []byte("12345")) + request := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(encoded)) + body, err := readDecodedRequestBody(request, contentcoding.Gzip, int64(len(encoded)), 4) + if !errors.Is(err, errRequestTooLarge) || body != nil { + t.Fatalf("readDecodedRequestBody() = %q, %v, want decoded overflow", body, err) + } + }) + + t.Run("MaxInt64 encoded limit does not overflow limit plus one", func(t *testing.T) { + request := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("1234")) + body, err := readDecodedRequestBody( + request, + contentcoding.Identity, + math.MaxInt64, + 4, + ) + if err != nil || string(body) != "1234" { + t.Fatalf("readDecodedRequestBody() = %q, %v, want 1234", body, err) + } + }) + + t.Run("negative decoded limit is explicitly rejected", func(t *testing.T) { + request := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("1234")) + body, err := readDecodedRequestBody( + request, + contentcoding.Identity, + 4, + -1, + ) + if body != nil || err == nil || !strings.Contains(err.Error(), "decoded request body limit") { + t.Fatalf("readDecodedRequestBody() = %q, %v, want explicit decoded-limit error", body, err) + } + }) +} + +func TestHandlerCapturesKeyIdentityOnlyAfterDecodedInspection(t *testing.T) { + tests := []struct { + name string + body []byte + encoding string + contentLength int64 + wantStatus int + wantCapture int + }{ + { + name: "malformed coding", body: []byte("not-gzip"), encoding: "gzip", + wantStatus: http.StatusBadRequest, + }, + { + name: "encoded limit", body: []byte(`{"model":"gpt-4o"}`), + contentLength: maxRequestBodyBytes + 1, wantStatus: http.StatusRequestEntityTooLarge, + }, + { + name: "invalid JSON", body: []byte(`{"model":`), + wantStatus: http.StatusBadRequest, + }, + { + name: "valid decoded request", body: []byte(`{"model":"gpt-4o"}`), + wantStatus: http.StatusOK, wantCapture: 1, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + forwarder := &scriptedForwarder{results: []UpstreamResult{{ + StatusCode: http.StatusOK, Header: make(http.Header), + Body: []byte(`{"ok":true}`), RequestWritten: true, + }}} + handler, _, registry := newHandlerForTest(t, forwarder, "sk-upstream") + counting := &captureCountingRuntimeRegistry{runtimeKeyRegistry: registry} + handler.registry = counting + engine := gin.New() + bindGatewayRoutesForTest(t, engine, handler) + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + bytes.NewReader(test.body), + ) + request.Header.Set("Authorization", "Bearer gl-client") + if test.encoding != "" { + request.Header.Set("Content-Encoding", test.encoding) + } + if test.contentLength > 0 { + request.ContentLength = test.contentLength + } + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + + if recorder.Code != test.wantStatus || counting.captureCalls != test.wantCapture { + t.Fatalf( + "response/capture calls = %d/%d, want %d/%d; body=%s", + recorder.Code, + counting.captureCalls, + test.wantStatus, + test.wantCapture, + recorder.Body.String(), + ) + } + }) + } + + t.Run("canceled body read", func(t *testing.T) { + forwarder := &scriptedForwarder{} + handler, _, registry := newHandlerForTest(t, forwarder, "sk-upstream") + counting := &captureCountingRuntimeRegistry{runtimeKeyRegistry: registry} + handler.registry = counting + engine := gin.New() + bindGatewayRoutesForTest(t, engine, handler) + requestContext, cancel := context.WithCancel(context.Background()) + body := newBlockingRequestBody(`{"model":"gpt-4o"}`, func() error { + cancel() + return context.Canceled + }) + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + body, + ).WithContext(requestContext) + request.Header.Set("Authorization", "Bearer gl-client") + recorder := httptest.NewRecorder() + done := make(chan struct{}) + go func() { + engine.ServeHTTP(recorder, request) + close(done) + }() + receiveTestSignal(t, body.started, "post-inspection capture canceled read") + close(body.release) + receiveTestSignal(t, done, "post-inspection capture canceled completion") + if counting.captureCalls != 0 || len(forwarder.inputs)+len(forwarder.streamInputs) != 0 { + t.Fatalf( + "capture/forward calls = %d/%d, want 0/0", + counting.captureCalls, + len(forwarder.inputs)+len(forwarder.streamInputs), + ) + } + }) + + t.Run("canceled after successful inspection", func(t *testing.T) { + forwarder := &scriptedForwarder{} + handler, _, registry := newHandlerForTest(t, forwarder, "sk-upstream") + counting := &captureCountingRuntimeRegistry{runtimeKeyRegistry: registry} + handler.registry = counting + requestContext, cancel := context.WithCancel(context.Background()) + baseDialect := handler.dialects[protocol.OpenAICompletions] + handler.dialects = dialect.NewSet(&cancelingSuccessfulInspectDialect{ + Dialect: baseDialect, + cancel: cancel, + }) + engine := gin.New() + bindGatewayRoutesForTest(t, engine, handler) + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + strings.NewReader(`{"model":"gpt-4o"}`), + ).WithContext(requestContext) + request.Header.Set("Authorization", "Bearer gl-client") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + + if counting.captureCalls != 0 || len(forwarder.inputs)+len(forwarder.streamInputs) != 0 { + t.Fatalf( + "capture/forward calls = %d/%d, want 0/0 after inspection cancellation", + counting.captureCalls, + len(forwarder.inputs)+len(forwarder.streamInputs), + ) + } + if recorder.Body.Len() != 0 { + t.Fatalf("canceled response body = %q, want empty", recorder.Body.String()) + } + }) +} + +func TestHandlerContentCodingErrorContract(t *testing.T) { + tests := []struct { + name string + contentEncodings []string + collidingEncoding string + body []byte + contentLength int64 + wantStatus int + wantCode string + wantAcceptEncoding string + }{ + { + name: "unknown coding", contentEncodings: []string{"snappy"}, + body: []byte(`{"model":"gpt-4o"}`), + wantStatus: http.StatusUnsupportedMediaType, wantCode: "unsupported_content_encoding", + wantAcceptEncoding: contentcoding.SupportedRequestEncodings, + }, + { + name: "multiple case-colliding fields", contentEncodings: []string{"identity"}, + collidingEncoding: "gzip", body: []byte(`{"model":"gpt-4o"}`), + wantStatus: http.StatusUnsupportedMediaType, wantCode: "unsupported_content_encoding", + wantAcceptEncoding: contentcoding.SupportedRequestEncodings, + }, + { + name: "stacked coding", contentEncodings: []string{"gzip, br"}, + body: []byte(`{"model":"gpt-4o"}`), + wantStatus: http.StatusUnsupportedMediaType, wantCode: "unsupported_content_encoding", + wantAcceptEncoding: contentcoding.SupportedRequestEncodings, + }, + { + name: "malformed known coding", contentEncodings: []string{"gzip"}, + body: []byte("not-gzip"), + wantStatus: http.StatusBadRequest, wantCode: "invalid_content_encoding", + }, + { + name: "encoded length fast rejection", contentEncodings: []string{"identity"}, + body: []byte(`{"model":"gpt-4o"}`), contentLength: maxRequestBodyBytes + 1, + wantStatus: http.StatusRequestEntityTooLarge, wantCode: "request_too_large", + }, + { + name: "decoded length rejection", contentEncodings: []string{"gzip"}, + body: encodeContentCodingForGatewayTest( + t, + contentcoding.Gzip, + bytes.Repeat([]byte("x"), int(maxRequestBodyBytes+1)), + ), + wantStatus: http.StatusRequestEntityTooLarge, wantCode: "request_too_large", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + forwarder := &scriptedForwarder{} + engine, _, _ := newHandlerTestRuntime(t, forwarder, "sk-upstream") + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + bytes.NewReader(test.body), + ) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Del("Content-Encoding") + for _, value := range test.contentEncodings { + request.Header.Add("Content-Encoding", value) + } + if test.collidingEncoding != "" { + request.Header["content-encoding"] = []string{test.collidingEncoding} + } + if test.contentLength > 0 { + request.ContentLength = test.contentLength + } + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + + assertGatewayReasonTest(t, recorder, test.wantStatus, test.wantCode) + if got := recorder.Header().Get("Accept-Encoding"); got != test.wantAcceptEncoding { + t.Fatalf("Accept-Encoding = %q, want %q", got, test.wantAcceptEncoding) + } + if len(forwarder.inputs)+len(forwarder.streamInputs) != 0 { + t.Fatal("invalid content coding reached upstream") + } + }) + } +} + +func TestHandlerContentCodingInspectionUsesDecodedBody(t *testing.T) { + plaintext := []byte(`{"model":"gpt-4o","stream":true}`) + for _, encoding := range []contentcoding.Encoding{ + contentcoding.Identity, + contentcoding.Gzip, + contentcoding.Brotli, + contentcoding.Deflate, + contentcoding.Zstd, + } { + t.Run(string(encoding), func(t *testing.T) { + forwarder := &scriptedForwarder{streamResults: []UpstreamResult{{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Committed: true, + RequestWritten: true, + Stream: StreamObservation{EndReason: StreamEndCleanEOF}, + }}} + engine, _, _ := newHandlerTestRuntime(t, forwarder, "sk-upstream") + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + bytes.NewReader(encodeContentCodingForGatewayTest(t, encoding, plaintext)), + ) + request.Header.Set("Authorization", "Bearer gl-client") + if encoding != contentcoding.Identity { + request.Header.Set("Content-Encoding", string(encoding)) + } + request.Header["content-length"] = []string{"stale-client-length"} + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusOK || len(forwarder.streamInputs) != 1 { + t.Fatalf( + "response/stream calls = %d/%d body=%s, want decoded stream request", + recorder.Code, + len(forwarder.streamInputs), + recorder.Body.String(), + ) + } + parsed := forwarder.streamInputs[0].Request + if !bytes.Equal(parsed.Body, plaintext) || forwarder.streamInputs[0].ExternalModel != "gpt-4o" { + t.Fatalf("parsed body/model = %s / %q", parsed.Body, forwarder.streamInputs[0].ExternalModel) + } + if values := headerFieldValues(parsed.Header, "Content-Encoding"); len(values) != 0 { + t.Fatalf("parsed Content-Encoding = %#v, want absent", values) + } + if values := headerFieldValues(parsed.Header, "Content-Length"); len(values) != 0 { + t.Fatalf("parsed Content-Length = %#v, want absent", values) + } + }) + } +} + +func TestHandlerContentCodingErrorPrecedence(t *testing.T) { + t.Run("invalid access key wins", func(t *testing.T) { + engine, _, _ := newHandlerTestRuntime(t, &scriptedForwarder{}, "sk-upstream") + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("invalid")) + request.Header.Set("Authorization", "Bearer wrong") + request.Header.Set("Content-Encoding", "unknown") + request.Header.Set("Accept-Encoding", "identity;q=0") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + assertGatewayReasonTest(t, recorder, http.StatusUnauthorized, reasonInvalidAccessKey.Code) + }) + + for _, test := range []struct { + name string + method string + target string + wantStatus int + wantCode string + }{ + {name: "route wins", method: http.MethodPost, target: "/v1/unknown", wantStatus: http.StatusNotFound, wantCode: reasonEndpointNotFound.Code}, + {name: "method wins", method: http.MethodGet, target: "/v1/chat/completions", wantStatus: http.StatusMethodNotAllowed, wantCode: reasonMethodNotAllowed.Code}, + } { + t.Run(test.name, func(t *testing.T) { + engine, _, _ := newHandlerTestRuntime(t, &scriptedForwarder{}, "sk-upstream") + request := httptest.NewRequest(test.method, test.target, strings.NewReader("invalid")) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Set("Content-Encoding", "unknown") + request.Header.Set("Accept-Encoding", "identity;q=0") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + assertGatewayReasonTest(t, recorder, test.wantStatus, test.wantCode) + }) + } + + t.Run("RPM wins", func(t *testing.T) { + limiter := &recordingAccessKeyRPMLimiter{decisions: []ratelimit.LimitDecision{{ + Allowed: false, RetryAfter: time.Second, + }}} + engine, _, _, _ := newRequestLogHandlerTestRuntime( + t, + &scriptedForwarder{}, + limiter, + &recordingRequestLogSink{}, + "sk-upstream", + ) + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("invalid")) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Set("Content-Encoding", "unknown") + request.Header.Set("Accept-Encoding", "identity;q=0") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + assertGatewayReasonTest(t, recorder, http.StatusTooManyRequests, reasonAccessKeyRateLimited.Code) + }) + + t.Run("415 wins over 406", func(t *testing.T) { + engine, _, _ := newHandlerTestRuntime(t, &scriptedForwarder{}, "sk-upstream") + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("invalid")) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Set("Content-Encoding", "unknown") + request.Header.Set("Accept-Encoding", "identity;q=0") + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + assertGatewayReasonTest(t, recorder, http.StatusUnsupportedMediaType, "unsupported_content_encoding") + }) + + t.Run("406 wins over body and JSON", func(t *testing.T) { + engine, _, _ := newHandlerTestRuntime(t, &scriptedForwarder{}, "sk-upstream") + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("not-gzip")) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Set("Content-Encoding", "gzip") + request.Header.Set("Accept-Encoding", "identity;q=0") + request.ContentLength = maxRequestBodyBytes + 1 + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + assertGatewayReasonTest(t, recorder, http.StatusNotAcceptable, "not_acceptable") + }) +} + +func TestHandlerIdentityNegotiation(t *testing.T) { + t.Run("case-colliding explicit identity exclusion wins over wildcard", func(t *testing.T) { + engine, _, _ := newHandlerTestRuntime(t, &scriptedForwarder{}, "sk-upstream") + request := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + strings.NewReader(`{"model":"gpt-4o"}`), + ) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header["Accept-Encoding"] = []string{"*;q=1"} + request.Header["accept-encoding"] = []string{"identity;q=0"} + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + assertGatewayReasonTest(t, recorder, http.StatusNotAcceptable, "not_acceptable") + }) + + t.Run("model list negotiates without reading body", func(t *testing.T) { + engine := newModelListHandlerEngine(t, state.FilterSet{}) + request := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Set("Accept-Encoding", "identity;q=0") + request.Header.Set("Content-Encoding", "unknown") + request.Body = panicReadCloser{} + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + assertGatewayReasonTest(t, recorder, http.StatusNotAcceptable, "not_acceptable") + }) + + t.Run("model list accepts implicit identity without reading body", func(t *testing.T) { + engine := newModelListHandlerEngine(t, state.FilterSet{}) + request := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Set("Accept-Encoding", "gzip") + request.Header.Set("Content-Encoding", "unknown") + request.Body = panicReadCloser{} + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("response = %d %s, want 200", recorder.Code, recorder.Body.String()) + } + }) +} + +func TestHandlerCancellationWinsOverContentCoding(t *testing.T) { + forwarder := &scriptedForwarder{} + engine, _, _ := newHandlerTestRuntime(t, forwarder, "sk-upstream") + requestContext, cancel := context.WithCancel(context.Background()) + body := newBlockingRequestBody("not-gzip", func() error { + cancel() + return context.Canceled + }) + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", body).WithContext(requestContext) + request.Header.Set("Authorization", "Bearer gl-client") + request.Header.Set("Content-Encoding", "gzip") + recorder := httptest.NewRecorder() + done := make(chan struct{}) + go func() { + engine.ServeHTTP(recorder, request) + close(done) + }() + receiveTestSignal(t, body.started, "canceled body read") + close(body.release) + receiveTestSignal(t, done, "canceled request completion") + + if recorder.Body.Len() != 0 || recorder.Code == http.StatusBadRequest || recorder.Code == http.StatusRequestEntityTooLarge { + t.Fatalf("canceled response = %d %q, want no coding reason", recorder.Code, recorder.Body.String()) + } + if len(forwarder.inputs)+len(forwarder.streamInputs) != 0 { + t.Fatal("canceled request reached upstream") + } +} + +func assertGatewayReasonTest( + t *testing.T, + recorder *httptest.ResponseRecorder, + wantStatus int, + wantCode string, +) { + t.Helper() + if recorder.Code != wantStatus { + t.Fatalf("status = %d, want %d; body=%s", recorder.Code, wantStatus, recorder.Body.String()) + } + var body struct { + Code string `json:"code"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil { + t.Fatalf("decode reason body: %v; body=%q", err, recorder.Body.String()) + } + if body.Code != wantCode { + t.Fatalf("code = %q, want %q; body=%s", body.Code, wantCode, recorder.Body.String()) + } } func TestHandlerRejectsOversizedRequestBody(t *testing.T) { @@ -2776,7 +3299,7 @@ func TestHandlerSkipsCandidateChangedAfterCollection(t *testing.T) { } } -func TestHandlerFreezesKeyIdentityBeforeReadingBody(t *testing.T) { +func TestHandlerFreezesKeyIdentityAfterInspectingBody(t *testing.T) { tests := []struct { name string seedPlain string @@ -2873,20 +3396,19 @@ func TestHandlerFreezesKeyIdentityBeforeReadingBody(t *testing.T) { t.Fatalf("first Read identity mutation error = %v", body.firstReadErr) } - if len(forwarder.inputs) != 0 { + if oldRecorder.Code != http.StatusOK || len(forwarder.inputs) != 1 { t.Fatalf( - "blocked old request forwarded with API key %q, want no newly visible identity", - forwarder.inputs[0].APIKey, + "decoded request response/attempts = %d/%d, want 200/1; body=%s", + oldRecorder.Code, + len(forwarder.inputs), + oldRecorder.Body.String(), ) } - if len(recordingEncryption.ciphertexts) != 0 { - t.Fatalf( - "blocked old request decrypted %d ciphertexts, want 0", - len(recordingEncryption.ciphertexts), - ) + if got := forwarder.inputs[0].APIKey; got != test.currentPlain { + t.Fatalf("decoded request API key = %q, want %q", got, test.currentPlain) } - if oldRecorder.Code != http.StatusServiceUnavailable { - t.Fatalf("blocked old request status = %d, want 503", oldRecorder.Code) + if len(recordingEncryption.ciphertexts) != 1 { + t.Fatalf("decoded request decrypt calls = %d, want 1", len(recordingEncryption.ciphertexts)) } newRequest := httptest.NewRequest( @@ -2898,20 +3420,20 @@ func TestHandlerFreezesKeyIdentityBeforeReadingBody(t *testing.T) { newRecorder := httptest.NewRecorder() engine.ServeHTTP(newRecorder, newRequest) - if newRecorder.Code != http.StatusOK || len(forwarder.inputs) != 1 { + if newRecorder.Code != http.StatusOK || len(forwarder.inputs) != 2 { t.Fatalf( - "new request response/attempts = %d/%d, want 200/1; body=%s", + "new request response/attempts = %d/%d, want 200/2; body=%s", newRecorder.Code, len(forwarder.inputs), newRecorder.Body.String(), ) } - if got := forwarder.inputs[0].APIKey; got != test.currentPlain { + if got := forwarder.inputs[1].APIKey; got != test.currentPlain { t.Fatalf("new request API key = %q, want %q", got, test.currentPlain) } - if len(recordingEncryption.ciphertexts) != 1 { + if len(recordingEncryption.ciphertexts) != 2 { t.Fatalf( - "new request decrypt calls = %d, want 1", + "new request decrypt calls = %d, want 2", len(recordingEncryption.ciphertexts), ) } diff --git a/internal/gateway/reason.go b/internal/gateway/reason.go index 7aeaa9b9..9fd6e798 100644 --- a/internal/gateway/reason.go +++ b/internal/gateway/reason.go @@ -14,18 +14,33 @@ type reason struct { } var ( - reasonInvalidAccessKey = reason{Status: http.StatusUnauthorized, Code: "invalid_access_key", Message: "Invalid access key."} - reasonEndpointNotFound = reason{Status: http.StatusNotFound, Code: "protocol_endpoint_not_found", Message: "Protocol endpoint not found."} - reasonMethodNotAllowed = reason{Status: http.StatusMethodNotAllowed, Code: "method_not_allowed", Message: "Method not allowed."} - reasonInvalidProtocolRequest = reason{Status: http.StatusBadRequest, Code: "invalid_protocol_request", Message: "Invalid protocol request."} - reasonModelRequiredByFilter = reason{Status: http.StatusBadRequest, Code: "model_required_by_filter", Message: "A model is required by the access key filter."} - reasonNoCandidate = reason{Status: http.StatusServiceUnavailable, Code: "no_available_candidate", Message: "No available upstream candidate."} - reasonUpstreamConnect = reason{Status: http.StatusBadGateway, Code: "upstream_connect_failed", Message: "Could not connect to an upstream service."} - reasonUpstreamTimeout = reason{Status: http.StatusGatewayTimeout, Code: "upstream_timeout", Message: "Upstream request timed out."} - reasonUpstreamProtocol = reason{Status: http.StatusBadGateway, Code: "upstream_protocol_error", Message: "Upstream returned an unsupported response."} - reasonRequestTooLarge = reason{Status: http.StatusRequestEntityTooLarge, Code: "request_too_large", Message: "Request body is too large."} - reasonModelListTooLarge = reason{Status: http.StatusInternalServerError, Code: "model_list_too_large", Message: "Model list is too large."} - reasonAccessKeyRateLimited = reason{ + reasonInvalidAccessKey = reason{Status: http.StatusUnauthorized, Code: "invalid_access_key", Message: "Invalid access key."} + reasonEndpointNotFound = reason{Status: http.StatusNotFound, Code: "protocol_endpoint_not_found", Message: "Protocol endpoint not found."} + reasonMethodNotAllowed = reason{Status: http.StatusMethodNotAllowed, Code: "method_not_allowed", Message: "Method not allowed."} + reasonInvalidProtocolRequest = reason{Status: http.StatusBadRequest, Code: "invalid_protocol_request", Message: "Invalid protocol request."} + reasonModelRequiredByFilter = reason{Status: http.StatusBadRequest, Code: "model_required_by_filter", Message: "A model is required by the access key filter."} + reasonNoCandidate = reason{Status: http.StatusServiceUnavailable, Code: "no_available_candidate", Message: "No available upstream candidate."} + reasonUpstreamConnect = reason{Status: http.StatusBadGateway, Code: "upstream_connect_failed", Message: "Could not connect to an upstream service."} + reasonUpstreamTimeout = reason{Status: http.StatusGatewayTimeout, Code: "upstream_timeout", Message: "Upstream request timed out."} + reasonUpstreamProtocol = reason{Status: http.StatusBadGateway, Code: "upstream_protocol_error", Message: "Upstream returned an unsupported response."} + reasonRequestTooLarge = reason{Status: http.StatusRequestEntityTooLarge, Code: "request_too_large", Message: "Request body is too large."} + reasonUnsupportedContentEncoding = reason{ + Status: http.StatusUnsupportedMediaType, + Code: "unsupported_content_encoding", + Message: "Unsupported Content-Encoding.", + } + reasonInvalidContentEncoding = reason{ + Status: http.StatusBadRequest, + Code: "invalid_content_encoding", + Message: "Invalid encoded request body.", + } + reasonNotAcceptable = reason{ + Status: http.StatusNotAcceptable, + Code: "not_acceptable", + Message: "The gateway can only return an identity-encoded response.", + } + reasonModelListTooLarge = reason{Status: http.StatusInternalServerError, Code: "model_list_too_large", Message: "Model list is too large."} + reasonAccessKeyRateLimited = reason{ Status: http.StatusTooManyRequests, Code: "access_key_rate_limited", Message: "Access key rate limit exceeded.", diff --git a/internal/gateway/reason_test.go b/internal/gateway/reason_test.go index 119580c3..0548d031 100644 --- a/internal/gateway/reason_test.go +++ b/internal/gateway/reason_test.go @@ -2,6 +2,7 @@ package gateway import ( "encoding/json" + "net/http" "net/http/httptest" "testing" @@ -19,6 +20,9 @@ func TestWriteReasonUsesStableDataPlaneEnvelope(t *testing.T) { reasonUpstreamTimeout, reasonUpstreamProtocol, reasonRequestTooLarge, + reasonUnsupportedContentEncoding, + reasonInvalidContentEncoding, + reasonNotAcceptable, reasonModelListTooLarge, reasonAccessKeyRateLimited, } @@ -50,6 +54,43 @@ func TestWriteReasonUsesStableDataPlaneEnvelope(t *testing.T) { } } +func TestContentCodingReasonsUseStableContract(t *testing.T) { + tests := []struct { + got reason + want reason + }{ + { + got: reasonUnsupportedContentEncoding, + want: reason{ + Status: http.StatusUnsupportedMediaType, + Code: "unsupported_content_encoding", + Message: "Unsupported Content-Encoding.", + }, + }, + { + got: reasonInvalidContentEncoding, + want: reason{ + Status: http.StatusBadRequest, + Code: "invalid_content_encoding", + Message: "Invalid encoded request body.", + }, + }, + { + got: reasonNotAcceptable, + want: reason{ + Status: http.StatusNotAcceptable, + Code: "not_acceptable", + Message: "The gateway can only return an identity-encoded response.", + }, + }, + } + for _, test := range tests { + if test.got != test.want { + t.Errorf("reason %q = %#v, want %#v", test.want.Code, test.got, test.want) + } + } +} + func TestUpstreamProtocolReasonCoversNonStreamingResponses(t *testing.T) { const want = "Upstream returned an unsupported response." if reasonUpstreamProtocol.Message != want { diff --git a/internal/gateway/response_framing.go b/internal/gateway/response_framing.go new file mode 100644 index 00000000..d40fd57d --- /dev/null +++ b/internal/gateway/response_framing.go @@ -0,0 +1,101 @@ +package gateway + +import ( + "net/http" + "strconv" + "strings" +) + +type responseBodyPolicy struct { + readBody bool + writeBody bool + preserveContentLength bool +} + +func classifyResponseBody(method string, status int) responseBodyPolicy { + if status >= http.StatusContinue && status < http.StatusOK { + return responseBodyPolicy{} + } + switch status { + case http.StatusNoContent, http.StatusResetContent: + return responseBodyPolicy{} + case http.StatusNotModified: + return responseBodyPolicy{preserveContentLength: true} + } + if strings.EqualFold(method, http.MethodHead) { + return responseBodyPolicy{preserveContentLength: true} + } + return responseBodyPolicy{readBody: true, writeBody: true} +} + +func normalizeBufferedResponse( + method string, + status int, + headers http.Header, + body []byte, +) (http.Header, bool) { + normalized := cloneEndToEndHeaders(headers) + policy := classifyResponseBody(method, status) + representationHeadersChanged := len(headerFieldValues(normalized, "Content-Encoding")) > 0 + deleteHeaderField(normalized, "Content-Encoding") + + switch { + case policy.writeBody: + representationHeadersChanged = representationHeadersChanged || + !hasSingleHeaderFieldValue(normalized, "Content-Length", strconv.Itoa(len(body))) + deleteHeaderField(normalized, "Content-Length") + normalized.Set("Content-Length", strconv.Itoa(len(body))) + case status == http.StatusResetContent: + if contentLength, ok := singleNonNegativeContentLength(normalized); ok && contentLength == "0" { + deleteHeaderField(normalized, "Content-Length") + normalized.Set("Content-Length", "0") + } else { + representationHeadersChanged = representationHeadersChanged || + len(headerFieldValues(normalized, "Content-Length")) > 0 + deleteHeaderField(normalized, "Content-Length") + } + case policy.preserveContentLength: + if strings.EqualFold(method, http.MethodHead) && status != http.StatusNotModified && body != nil { + representationHeadersChanged = representationHeadersChanged || + !hasSingleHeaderFieldValue(normalized, "Content-Length", strconv.Itoa(len(body))) + deleteHeaderField(normalized, "Content-Length") + normalized.Set("Content-Length", strconv.Itoa(len(body))) + break + } + if contentLength, ok := singleNonNegativeContentLength(normalized); ok { + deleteHeaderField(normalized, "Content-Length") + normalized.Set("Content-Length", contentLength) + } else { + representationHeadersChanged = representationHeadersChanged || + len(headerFieldValues(normalized, "Content-Length")) > 0 + deleteHeaderField(normalized, "Content-Length") + } + default: + representationHeadersChanged = representationHeadersChanged || + len(headerFieldValues(normalized, "Content-Length")) > 0 + deleteHeaderField(normalized, "Content-Length") + } + if representationHeadersChanged { + deleteHeaderField(normalized, "Signature") + deleteHeaderField(normalized, "Signature-Input") + } + + return normalized, policy.writeBody +} + +func singleNonNegativeContentLength(headers http.Header) (string, bool) { + values := headerFieldValues(headers, "Content-Length") + if len(values) != 1 || values[0] == "" { + return "", false + } + for _, character := range values[0] { + if character < '0' || character > '9' { + return "", false + } + } + value, err := strconv.ParseInt(values[0], 10, 64) + if err != nil || value < 0 { + return "", false + } + return values[0], true +} diff --git a/internal/gateway/response_framing_test.go b/internal/gateway/response_framing_test.go new file mode 100644 index 00000000..315aa50b --- /dev/null +++ b/internal/gateway/response_framing_test.go @@ -0,0 +1,310 @@ +package gateway + +import ( + "net/http" + "net/http/httptest" + "reflect" + "testing" + "time" + + "github.com/gin-gonic/gin" +) + +func TestClassifyResponseBody(t *testing.T) { + tests := []struct { + name string + method string + status int + want responseBodyPolicy + }{ + {name: "ordinary response", method: http.MethodGet, status: http.StatusOK, want: responseBodyPolicy{readBody: true, writeBody: true}}, + {name: "head", method: http.MethodHead, status: http.StatusOK, want: responseBodyPolicy{preserveContentLength: true}}, + {name: "informational status overrides head", method: http.MethodHead, status: http.StatusContinue, want: responseBodyPolicy{}}, + {name: "no content", method: http.MethodGet, status: http.StatusNoContent, want: responseBodyPolicy{}}, + {name: "reset content", method: http.MethodGet, status: http.StatusResetContent, want: responseBodyPolicy{}}, + {name: "not modified", method: http.MethodGet, status: http.StatusNotModified, want: responseBodyPolicy{preserveContentLength: true}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := classifyResponseBody(test.method, test.status); got != test.want { + t.Fatalf("classifyResponseBody(%q, %d) = %#v, want %#v", test.method, test.status, got, test.want) + } + }) + } +} + +func TestNormalizeBufferedResponse(t *testing.T) { + tests := []struct { + name string + method string + status int + headers http.Header + body []byte + wantWriteBody bool + wantContentLength string + }{ + { + name: "bodyful response gets plaintext length", + method: http.MethodGet, + status: http.StatusOK, + headers: http.Header{ + "content-encoding": {"gzip"}, + "Content-Length": {"100"}, + }, + body: []byte("plain"), + wantWriteBody: true, + wantContentLength: "5", + }, + { + name: "forwarded head preserves legal representation length", + method: http.MethodHead, + status: http.StatusOK, + headers: http.Header{ + "Content-Encoding": {"identity"}, + "Content-Length": {"123"}, + }, + wantContentLength: "123", + }, + { + name: "local head representation uses body length", + method: http.MethodHead, + status: http.StatusBadGateway, + headers: http.Header{ + "Content-Length": {"999"}, + }, + body: []byte("reason"), + wantContentLength: "6", + }, + { + name: "informational status deletes length", + method: http.MethodGet, + status: http.StatusContinue, + headers: http.Header{ + "Content-Length": {"8"}, + }, + }, + { + name: "no content deletes length", + method: http.MethodGet, + status: http.StatusNoContent, + headers: http.Header{ + "Content-Length": {"8"}, + }, + }, + { + name: "reset content permits zero length", + method: http.MethodGet, + status: http.StatusResetContent, + headers: http.Header{ + "Content-Length": {"0"}, + }, + wantContentLength: "0", + }, + { + name: "reset content deletes nonzero length", + method: http.MethodGet, + status: http.StatusResetContent, + headers: http.Header{ + "Content-Length": {"8"}, + }, + }, + { + name: "not modified preserves legal representation length", + method: http.MethodGet, + status: http.StatusNotModified, + headers: http.Header{ + "content-encoding": {"identity"}, + "Content-Length": {"456"}, + }, + wantContentLength: "456", + }, + { + name: "preservation rejects duplicate casing", + method: http.MethodHead, + status: http.StatusOK, + headers: http.Header{ + "Content-Length": {"12"}, + "content-length": {"12"}, + }, + }, + { + name: "preservation rejects duplicate values", + method: http.MethodGet, + status: http.StatusNotModified, + headers: http.Header{ + "Content-Length": {"12", "12"}, + }, + }, + { + name: "preservation rejects comma length", + method: http.MethodHead, + status: http.StatusOK, + headers: http.Header{ + "Content-Length": {"12, 12"}, + }, + }, + { + name: "preservation rejects negative length", + method: http.MethodGet, + status: http.StatusNotModified, + headers: http.Header{ + "Content-Length": {"-1"}, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + source := test.headers.Clone() + got, writeBody := normalizeBufferedResponse(test.method, test.status, source, test.body) + if writeBody != test.wantWriteBody { + t.Fatalf("write body = %t, want %t", writeBody, test.wantWriteBody) + } + if values := headerFieldValues(got, "Content-Encoding"); len(values) != 0 { + t.Fatalf("Content-Encoding values = %#v, want absent", values) + } + if values := headerFieldValues(got, "Content-Length"); test.wantContentLength == "" { + if len(values) != 0 { + t.Fatalf("Content-Length values = %#v, want absent", values) + } + } else if !reflect.DeepEqual(values, []string{test.wantContentLength}) { + t.Fatalf("Content-Length values = %#v, want [%s]", values, test.wantContentLength) + } + if !reflect.DeepEqual(source, test.headers) { + t.Fatalf("normalizeBufferedResponse mutated source headers: got %#v, want %#v", source, test.headers) + } + }) + } +} + +func TestNormalizeBufferedResponseInvalidatesSignaturesAfterRepresentationHeaderChanges(t *testing.T) { + tests := []struct { + name string + method string + status int + headers http.Header + body []byte + wantSignaturesGone bool + }{ + { + name: "head removes noncanonical identity coding", + method: http.MethodHead, + status: http.StatusOK, + headers: http.Header{ + "content-encoding": {"identity"}, + "Content-Length": {"12"}, + "sIgNaTuRe": {"sig"}, + "sIgNaTuRe-InPuT": {`sig1=("content-encoding")`}, + }, + wantSignaturesGone: true, + }, + { + name: "not modified removes identity coding", + method: http.MethodGet, + status: http.StatusNotModified, + headers: http.Header{ + "Content-Encoding": {"identity"}, + "Content-Length": {"12"}, + "Signature": {"sig"}, + "Signature-Input": {`sig1=("content-encoding")`}, + }, + wantSignaturesGone: true, + }, + { + name: "no content deletes length", + method: http.MethodGet, + status: http.StatusNoContent, + headers: http.Header{ + "content-length": {"12"}, + "Signature": {"sig"}, + "Signature-Input": {`sig1=("content-length")`}, + }, + wantSignaturesGone: true, + }, + { + name: "reset content deletes nonzero length", + method: http.MethodGet, + status: http.StatusResetContent, + headers: http.Header{ + "Content-Length": {"12"}, + "Signature": {"sig"}, + "Signature-Input": {`sig1=("content-length")`}, + }, + wantSignaturesGone: true, + }, + { + name: "unchanged head framing preserves signatures", + method: http.MethodHead, + status: http.StatusOK, + headers: http.Header{ + "Content-Length": {"12"}, + "Signature": {"sig"}, + "Signature-Input": {`sig1=("content-length")`}, + }, + }, + { + name: "unchanged not modified framing preserves signatures", + method: http.MethodGet, + status: http.StatusNotModified, + headers: http.Header{ + "Content-Length": {"12"}, + "Signature": {"sig"}, + "Signature-Input": {`sig1=("content-length")`}, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, _ := normalizeBufferedResponse(test.method, test.status, test.headers, test.body) + if test.wantSignaturesGone { + if len(headerFieldValues(got, "Signature")) != 0 || len(headerFieldValues(got, "Signature-Input")) != 0 { + t.Fatalf("normalized headers retained invalid signatures: %#v", got) + } + return + } + if got.Get("Signature") != "sig" || got.Get("Signature-Input") != `sig1=("content-length")` { + t.Fatalf("normalized headers discarded unchanged signatures: %#v", got) + } + }) + } +} + +func TestWriteBufferedResponseMethodAndStatusSemantics(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + method string + status int + headers http.Header + body []byte + wantBody string + wantContentLength string + }{ + {name: "ordinary", method: http.MethodGet, status: http.StatusOK, headers: http.Header{"Content-Encoding": {"gzip"}}, body: []byte("plain"), wantBody: "plain", wantContentLength: "5"}, + {name: "head local reason", method: http.MethodHead, status: http.StatusBadGateway, body: []byte("reason"), wantContentLength: "6"}, + {name: "informational", method: http.MethodGet, status: http.StatusContinue, headers: http.Header{"Content-Length": {"7"}}, body: []byte("ignored")}, + {name: "no content", method: http.MethodGet, status: http.StatusNoContent, headers: http.Header{"Content-Length": {"7"}}, body: []byte("ignored")}, + {name: "reset content", method: http.MethodGet, status: http.StatusResetContent, headers: http.Header{"Content-Length": {"0"}}, body: []byte("ignored"), wantContentLength: "0"}, + {name: "not modified", method: http.MethodGet, status: http.StatusNotModified, headers: http.Header{"Content-Length": {"32"}}, body: []byte("ignored"), wantContentLength: "32"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(test.method, "http://gateway.test/", nil) + handler := &Handler{writeTimeout: time.Second} + if err := handler.writeBufferedResponse(context, test.status, test.headers, test.body); err != nil { + t.Fatalf("writeBufferedResponse() error = %v", err) + } + if got := recorder.Body.String(); got != test.wantBody { + t.Fatalf("body = %q, want %q", got, test.wantBody) + } + if got := recorder.Header().Get("Content-Length"); got != test.wantContentLength { + t.Fatalf("Content-Length = %q, want %q", got, test.wantContentLength) + } + }) + } +} diff --git a/internal/gateway/response_representation.go b/internal/gateway/response_representation.go index ab55a256..5b83c118 100644 --- a/internal/gateway/response_representation.go +++ b/internal/gateway/response_representation.go @@ -6,23 +6,33 @@ import ( "fmt" "io" "net/http" + "reflect" "sort" + "strconv" "strings" "gpt-load/internal/dialect" + "gpt-load/internal/platform/contentcoding" "gpt-load/internal/platform/redact" - "gpt-load/internal/platform/utils" ) type preparedSuccessRepresentation struct { - headers http.Header - wire []byte - plain []byte - changed bool + headers http.Header + downstream []byte + inspectable []byte + changed bool +} + +type preparedErrorRepresentation struct { + headers http.Header + downstream []byte + classification []byte + changed bool } func (forwarder *Forwarder) prepareSuccessRepresentation( input ForwardInput, + status int, headers http.Header, wire []byte, secrets []string, @@ -31,11 +41,13 @@ func (forwarder *Forwarder) prepareSuccessRepresentation( return preparedSuccessRepresentation{}, successRepresentationProtocolError("wire body exceeds limit") } - encoding, ok := inspectableSuccessBodyEncoding(headers, wire) - if !ok { + encoding, err := contentcoding.ParseContentEncoding( + headerFieldValues(headers, "Content-Encoding"), + ) + if err != nil { return preparedSuccessRepresentation{}, successRepresentationProtocolError("unsupported or malformed Content-Encoding") } - originalPlain, err := utils.DecompressResponseLimited( + originalPlain, err := contentcoding.DecodeLimited( encoding, wire, maxNonStreamingResponseBodyBytes, @@ -79,52 +91,224 @@ func (forwarder *Forwarder) prepareSuccessRepresentation( } } - changed := !bytes.Equal(originalPlain, downstreamPlain) + changed := encoding != contentcoding.Identity || + !bytes.Equal(originalPlain, downstreamPlain) + if status == http.StatusPartialContent && changed { + return preparedSuccessRepresentation{}, successRepresentationProtocolError("changed partial response") + } + preparedHeaders := headers.Clone() if preparedHeaders == nil { preparedHeaders = make(http.Header) } - preparedWire := bytes.Clone(wire) + representationHeadersChanged := len(headerFieldValues(preparedHeaders, "Content-Encoding")) > 0 || + !hasSingleHeaderFieldValue( + preparedHeaders, + "Content-Length", + strconv.Itoa(len(downstreamPlain)), + ) + deleteHeaderField(preparedHeaders, "Content-Encoding") if changed { - preparedWire, err = utils.CompressResponse(encoding, downstreamPlain) - if err != nil { - return preparedSuccessRepresentation{}, successRepresentationProtocolError("recompress response body") - } - if int64(len(preparedWire)) > maxNonStreamingResponseBodyBytes { - return preparedSuccessRepresentation{}, successRepresentationProtocolError("recompressed response body exceeds limit") - } - updateRewrittenBodyHeaders(preparedHeaders, len(preparedWire)) + invalidateRewrittenBodyHeaders(preparedHeaders) + } + beforeSanitize := preparedHeaders.Clone() + preparedHeaders = sanitizeForwardResponseHeaders(preparedHeaders, input, secrets...) + representationHeadersChanged = representationHeadersChanged || + !reflect.DeepEqual(beforeSanitize, preparedHeaders) + deleteHeaderField(preparedHeaders, "Content-Encoding") + deleteHeaderField(preparedHeaders, "Content-Length") + preparedHeaders.Set("Content-Length", strconv.Itoa(len(downstreamPlain))) + if representationHeadersChanged { + deleteHeaderField(preparedHeaders, "Signature") + deleteHeaderField(preparedHeaders, "Signature-Input") } if representationMetadataContainsSecrets(preparedHeaders, secrets) { return preparedSuccessRepresentation{}, successRepresentationProtocolError("credential collision in rebuilt representation metadata") } - preparedHeaders = sanitizeForwardResponseHeaders(preparedHeaders, input, secrets...) + if status == http.StatusPartialContent && !hasSafeContentRange(preparedHeaders, len(downstreamPlain)) { + return preparedSuccessRepresentation{}, successRepresentationProtocolError("partial response lacks safe Content-Range") + } return preparedSuccessRepresentation{ - headers: preparedHeaders, - wire: preparedWire, - plain: inspectablePlain, - changed: changed, + headers: preparedHeaders, + downstream: bytes.Clone(downstreamPlain), + inspectable: inspectablePlain, + changed: changed, }, nil } -func inspectableSuccessBodyEncoding(headers http.Header, wire []byte) (string, bool) { - values := headers.Values("Content-Encoding") - if len(values) > 1 { - return "", false +func (forwarder *Forwarder) prepareErrorRepresentation( + input ForwardInput, + headers http.Header, + wire []byte, + secrets []string, +) preparedErrorRepresentation { + if int64(len(wire)) > maxErrorResponseBodyBytes { + return forwarder.failClosedErrorRepresentation(input, headers, secrets) } - encoding := "" - if len(values) == 1 { - encoding = strings.ToLower(strings.TrimSpace(values[0])) + + encoding, err := contentcoding.ParseContentEncoding( + headerFieldValues(headers, "Content-Encoding"), + ) + if err != nil { + return forwarder.failClosedErrorRepresentation(input, headers, secrets) } - switch encoding { - case "", "identity": - return encoding, true - case "gzip", "br", "deflate", "zstd": - return encoding, len(wire) > 0 - default: - return "", false + originalPlain, err := contentcoding.DecodeLimited( + encoding, + wire, + maxDecompressedErrorBodyBytes, + ) + if err != nil { + return forwarder.failClosedErrorRepresentation(input, headers, secrets) } + + patternSafePlain := forwarder.redactor.Bytes(originalPlain) + if int64(len(patternSafePlain)) > maxDecompressedErrorBodyBytes { + return forwarder.failClosedErrorRepresentation(input, headers, secrets) + } + safePlain, residualCredential, ok := redactCredentialLiterals( + patternSafePlain, + secrets, + maxDecompressedErrorBodyBytes, + ) + if !ok || residualCredential { + return forwarder.failClosedErrorRepresentation(input, headers, secrets) + } + + classificationPlain := bytes.Clone(safePlain) + downstreamPlain := safePlain + if needsModelRewrite(input) { + downstreamPlain, ok = rewriteBoundedLiteral( + safePlain, + input.UpstreamModelID, + input.ExternalModel, + maxDecompressedErrorBodyBytes, + ) + if !ok || credentialLiteralsRemain(downstreamPlain, secrets) { + return forwarder.failClosedErrorRepresentation(input, headers, secrets) + } + } + + changed := encoding != contentcoding.Identity || + !bytes.Equal(originalPlain, downstreamPlain) + preparedHeaders := headers.Clone() + if preparedHeaders == nil { + preparedHeaders = make(http.Header) + } + representationHeadersChanged := len(headerFieldValues(preparedHeaders, "Content-Encoding")) > 0 || + !hasSingleHeaderFieldValue( + preparedHeaders, + "Content-Length", + strconv.Itoa(len(downstreamPlain)), + ) + deleteHeaderField(preparedHeaders, "Content-Encoding") + if changed { + invalidateRewrittenBodyHeaders(preparedHeaders) + } + beforeSanitize := preparedHeaders.Clone() + preparedHeaders = sanitizeForwardResponseHeaders(preparedHeaders, input, secrets...) + representationHeadersChanged = representationHeadersChanged || + !reflect.DeepEqual(beforeSanitize, preparedHeaders) + deleteHeaderField(preparedHeaders, "Content-Encoding") + deleteHeaderField(preparedHeaders, "Content-Length") + preparedHeaders.Set("Content-Length", strconv.Itoa(len(downstreamPlain))) + if representationHeadersChanged { + deleteHeaderField(preparedHeaders, "Signature") + deleteHeaderField(preparedHeaders, "Signature-Input") + } + if representationMetadataContainsSecrets(preparedHeaders, secrets) { + return forwarder.failClosedErrorRepresentation(input, headers, secrets) + } + + return preparedErrorRepresentation{ + headers: preparedHeaders, + downstream: bytes.Clone(downstreamPlain), + classification: classificationPlain, + changed: changed, + } +} + +func (forwarder *Forwarder) failClosedErrorRepresentation( + input ForwardInput, + headers http.Header, + secrets []string, +) preparedErrorRepresentation { + preparedHeaders := headers.Clone() + if preparedHeaders == nil { + preparedHeaders = make(http.Header) + } + deleteHeaderField(preparedHeaders, "Content-Encoding") + invalidateRewrittenBodyHeaders(preparedHeaders) + preparedHeaders = sanitizeForwardResponseHeaders(preparedHeaders, input, secrets...) + deleteHeaderField(preparedHeaders, "Content-Encoding") + invalidateRewrittenBodyHeaders(preparedHeaders) + body := []byte(redact.Placeholder) + preparedHeaders.Set("Content-Type", "text/plain; charset=utf-8") + preparedHeaders.Set("Content-Length", strconv.Itoa(len(body))) + return preparedErrorRepresentation{ + headers: preparedHeaders, + downstream: bytes.Clone(body), + classification: bytes.Clone(body), + changed: true, + } +} + +func hasSingleHeaderFieldValue(headers http.Header, name, want string) bool { + values := headerFieldValues(headers, name) + return len(values) == 1 && values[0] == want +} + +func hasSafeContentRange(headers http.Header, bodyLength int) bool { + if bodyLength <= 0 { + return false + } + values := headerFieldValues(headers, "Content-Range") + if len(values) != 1 { + return false + } + parts := strings.Fields(values[0]) + if len(parts) != 2 || !strings.EqualFold(parts[0], "bytes") { + return false + } + rangeParts := strings.Split(parts[1], "/") + if len(rangeParts) != 2 { + return false + } + bounds := strings.Split(rangeParts[0], "-") + if len(bounds) != 2 { + return false + } + if !isDecimalContentRangeNumber(bounds[0]) || !isDecimalContentRangeNumber(bounds[1]) { + return false + } + start, startErr := strconv.ParseInt(bounds[0], 10, 64) + end, endErr := strconv.ParseInt(bounds[1], 10, 64) + if startErr != nil || endErr != nil || start < 0 || end < start { + return false + } + if end-start != int64(bodyLength)-1 { + return false + } + if rangeParts[1] == "*" { + return true + } + if !isDecimalContentRangeNumber(rangeParts[1]) { + return false + } + total, totalErr := strconv.ParseInt(rangeParts[1], 10, 64) + return totalErr == nil && total > end +} + +func isDecimalContentRangeNumber(value string) bool { + if value == "" { + return false + } + for _, character := range value { + if character < '0' || character > '9' { + return false + } + } + return true } func representationMetadataContainsSecrets(headers http.Header, secrets []string) bool { diff --git a/internal/gateway/response_representation_test.go b/internal/gateway/response_representation_test.go index d6525bd2..74e17711 100644 --- a/internal/gateway/response_representation_test.go +++ b/internal/gateway/response_representation_test.go @@ -6,18 +6,18 @@ import ( "context" "encoding/json" "errors" + "fmt" "io" "net/http" "net/http/httptest" "strconv" - "strings" "testing" "time" "gpt-load/internal/dialect" + "gpt-load/internal/platform/contentcoding" platformhttp "gpt-load/internal/platform/httpclient" "gpt-load/internal/platform/redact" - "gpt-load/internal/platform/utils" "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/usage" @@ -29,6 +29,347 @@ type successRepresentationEncoding struct { wire []byte } +func assertRepresentationValidatorsPreservedWithoutSignatures(t *testing.T, headers http.Header) { + t.Helper() + for name, want := range map[string]string{ + "ETag": `"wire-v1"`, + "Digest": "sha-256=wire-digest", + "Content-MD5": "d2lyZQ==", + "Content-Range": "bytes 0-9/10", + "Content-Digest": "sha-256=:d2lyZQ==:", + "Repr-Digest": "sha-256=:cmVwcg==:", + } { + if got := headers.Get(name); got != want { + t.Errorf("%s = %q, want preserved value %q", name, got, want) + } + } + if headers.Get("Signature") != "" || headers.Get("Signature-Input") != "" { + t.Errorf("stale signatures survived representation normalization: %#v", headers) + } +} + +func TestPrepareSuccessRepresentationReturnsPlainIdentity(t *testing.T) { + plain := []byte(`{"id":"safe-response","usage":{"prompt_tokens":100,"completion_tokens":30}}`) + encodings := []contentcoding.Encoding{ + contentcoding.Identity, + contentcoding.Gzip, + contentcoding.Brotli, + contentcoding.Deflate, + contentcoding.Zstd, + } + for _, encoding := range encodings { + t.Run(string(encoding), func(t *testing.T) { + result := forwardSuccessRepresentation( + t, + successRepresentationEncoding{ + name: string(encoding), + encoding: string(encoding), + wire: encodeContentCodingForGatewayTest(t, encoding, plain), + }, + "fake-provider-secret-plaintext-boundary", + nil, + ) + + if result.Err != nil || result.StatusCode != http.StatusOK { + t.Fatalf("Forward() result = %#v, want successful response", result) + } + if !bytes.Equal(result.Body, plain) { + t.Fatalf("downstream body = %x, want plaintext %x", result.Body, plain) + } + if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { + t.Fatalf("Content-Encoding values = %#v, want absent", values) + } + if got := result.Header.Get("Content-Length"); got != strconv.Itoa(len(plain)) { + t.Fatalf("Content-Length = %q, want %d", got, len(plain)) + } + }) + } +} + +func TestPrepareSuccessRepresentationMetadataPolicy(t *testing.T) { + plain := []byte(`{"id":"safe-response"}`) + + identity := forwardSuccessRepresentation( + t, + successRepresentationEncoding{ + name: "identity", + encoding: "identity", + wire: bytes.Clone(plain), + }, + "fake-provider-secret-metadata-identity", + nil, + ) + if identity.Err != nil || !bytes.Equal(identity.Body, plain) { + t.Fatalf("identity Forward() result = %#v", identity) + } + assertRepresentationValidatorsPreservedWithoutSignatures(t, identity.Header) + + gzipResult := forwardSuccessRepresentation( + t, + successRepresentationEncoding{ + name: "gzip", + encoding: "gzip", + wire: encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain), + }, + "fake-provider-secret-metadata-gzip", + nil, + ) + if gzipResult.Err != nil || !bytes.Equal(gzipResult.Body, plain) { + t.Fatalf("gzip Forward() result = %#v", gzipResult) + } + assertRepresentationMetadata(t, gzipResult.Header, false) + if values := headerFieldValues(gzipResult.Header, "Content-Encoding"); len(values) != 0 { + t.Fatalf("Content-Encoding values = %#v, want absent", values) + } +} + +func TestPrepareSuccessRepresentationRebuildsContentLengthAcrossHeaderCasings(t *testing.T) { + plain := []byte(`{"id":"safe-response"}`) + input := ForwardInput{ + Dialect: dialect.NewOpenAI(http.DefaultClient), + Request: &dialect.ParsedRequest{Header: make(http.Header)}, + } + headers := http.Header{ + "Content-Encoding": {"identity"}, + "Content-Length": {"999"}, + "content-length": {"998"}, + } + + result, err := NewForwarder( + platformhttp.NewHTTPClientManager(), + redact.New(), + ).prepareSuccessRepresentation(input, http.StatusOK, headers, plain, nil) + if err != nil { + t.Fatalf("prepareSuccessRepresentation() error = %v", err) + } + if values := headerFieldValues(result.headers, "Content-Length"); len(values) != 1 || values[0] != strconv.Itoa(len(plain)) { + t.Fatalf("Content-Length values = %#v, want exactly [%d]", values, len(plain)) + } +} + +func TestPrepareSuccessRepresentationDropsSignaturesAfterRepresentationHeaderNormalization(t *testing.T) { + plain := []byte(`{"id":"safe-response"}`) + headers := http.Header{ + "Content-Encoding": {"identity"}, + "Content-Length": {strconv.Itoa(len(plain))}, + "Signature": {"sig1=:c2lnOg==:"}, + "Signature-Input": {`sig1=("content-encoding" "content-length");created=1`}, + } + input := ForwardInput{ + Dialect: dialect.NewOpenAI(http.DefaultClient), + Request: &dialect.ParsedRequest{Header: make(http.Header)}, + } + + result, err := NewForwarder( + platformhttp.NewHTTPClientManager(), + redact.New(), + ).prepareSuccessRepresentation(input, http.StatusOK, headers, plain, nil) + if err != nil || result.changed || result.headers.Get("Signature") != "" || result.headers.Get("Signature-Input") != "" { + t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want normalized representation without signatures", result, err) + } +} + +func TestPrepareSuccessRepresentationRejectsPartialContentWithoutSafeContentRange(t *testing.T) { + plain := []byte(`{"id":"partial"}`) + input := ForwardInput{ + Dialect: dialect.NewOpenAI(http.DefaultClient), + ExternalModel: "public-model", + UpstreamModelID: "provider-model", + Request: &dialect.ParsedRequest{Header: make(http.Header)}, + } + headers := http.Header{ + "Content-Encoding": {"identity"}, + "Content-Length": {strconv.Itoa(len(plain))}, + "Content-Range": {"bytes provider-model/100"}, + } + + result, err := NewForwarder( + platformhttp.NewHTTPClientManager(), + redact.New(), + ).prepareSuccessRepresentation(input, http.StatusPartialContent, headers, plain, nil) + if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { + t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want unsafe partial response rejection", result, err) + } +} + +func TestPrepareSuccessRepresentationValidatesPartialContentRangeAgainstPlaintext(t *testing.T) { + plain := []byte("abc") + input := ForwardInput{ + Dialect: dialect.NewOpenAI(http.DefaultClient), + Request: &dialect.ParsedRequest{Header: make(http.Header)}, + } + tests := []struct { + name string + headers http.Header + wantErr bool + }{ + { + name: "matching decimal range", + headers: http.Header{ + "Content-Length": {"3"}, + "Content-Range": {"bytes 10-12/100"}, + }, + }, + { + name: "range length differs from plaintext", + headers: http.Header{ + "Content-Length": {"3"}, + "Content-Range": {"bytes 10-13/100"}, + }, + wantErr: true, + }, + { + name: "plus-prefixed range bound", + headers: http.Header{ + "Content-Length": {"3"}, + "Content-Range": {"bytes +10-12/100"}, + }, + wantErr: true, + }, + { + name: "plus-prefixed complete length", + headers: http.Header{ + "Content-Length": {"3"}, + "Content-Range": {"bytes 10-12/+100"}, + }, + wantErr: true, + }, + { + name: "case-colliding duplicate range", + headers: http.Header{ + "Content-Length": {"3"}, + "Content-Range": {"bytes 10-12/100"}, + "content-range": {"bytes 10-12/100"}, + }, + wantErr: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result, err := NewForwarder( + platformhttp.NewHTTPClientManager(), + redact.New(), + ).prepareSuccessRepresentation(input, http.StatusPartialContent, test.headers, plain, nil) + if test.wantErr { + if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { + t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want protocol error", result, err) + } + return + } + if err != nil || !bytes.Equal(result.downstream, plain) || + result.headers.Get("Content-Range") != "bytes 10-12/100" { + t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want valid partial response", result, err) + } + }) + } +} + +func TestPrepareSuccessRepresentationRejectsChangedPartialContent(t *testing.T) { + const ( + upstreamModel = "provider-model" + externalModel = "public-model" + ) + plain := []byte(`{"model":"provider-model","id":"partial"}`) + forwardPartial := func( + t *testing.T, + encoding contentcoding.Encoding, + configure func(*ForwardInput), + ) UpstreamResult { + t.Helper() + wire := encodeContentCodingForGatewayTest(t, encoding, plain) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Encoding", string(encoding)) + writer.Header().Set("Content-Range", fmt.Sprintf("bytes 0-%d/100", len(plain)-1)) + writer.WriteHeader(http.StatusPartialContent) + _, _ = writer.Write(wire) + })) + defer upstream.Close() + + input := usageForwardInput(upstream.URL, dialect.NewOpenAI(http.DefaultClient)) + input.Request.Body = []byte(`{"model":"public-model"}`) + if configure != nil { + configure(&input) + } + return NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward( + context.Background(), + input, + ) + } + + t.Run("identity alias rewrite", func(t *testing.T) { + result := forwardPartial(t, contentcoding.Identity, func(input *ForwardInput) { + input.ExternalModel = externalModel + input.UpstreamModelID = upstreamModel + }) + if !errors.Is(result.Err, ErrUpstreamProtocol) || result.StatusCode != 0 || len(result.Body) != 0 { + t.Fatalf("Forward() result = %#v, want changed 206 protocol error", result) + } + }) + + t.Run("compressed representation", func(t *testing.T) { + result := forwardPartial(t, contentcoding.Gzip, nil) + if !errors.Is(result.Err, ErrUpstreamProtocol) || result.StatusCode != 0 || len(result.Body) != 0 { + t.Fatalf("Forward() result = %#v, want decoded 206 protocol error", result) + } + }) + + t.Run("identity unchanged", func(t *testing.T) { + result := forwardPartial(t, contentcoding.Identity, nil) + if result.Err != nil || result.StatusCode != http.StatusPartialContent || + !bytes.Equal(result.Body, plain) || result.Header.Get("Content-Range") != fmt.Sprintf("bytes 0-%d/100", len(plain)-1) || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { + t.Fatalf("Forward() result = %#v, want unchanged legal 206", result) + } + }) +} + +func TestPrepareSuccessRepresentationBoundsWireAndDecodedBodies(t *testing.T) { + forwarder := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()) + input := ForwardInput{ + Dialect: dialect.NewOpenAI(http.DefaultClient), + Request: &dialect.ParsedRequest{Header: make(http.Header)}, + } + + caseCollidingHeaders := http.Header{ + "Content-Encoding": {"identity"}, + "content-encoding": {"gzip"}, + } + result, err := forwarder.prepareSuccessRepresentation( + input, + http.StatusOK, + caseCollidingHeaders, + []byte(`{"id":"safe-response"}`), + nil, + ) + if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { + t.Fatalf("case-colliding Content-Encoding result/error = %#v/%v, want protocol error", result, err) + } + + overflow := gzipRepeatedByte(t, 'x', maxNonStreamingResponseBodyBytes+1) + result, err = forwarder.prepareSuccessRepresentation( + input, + http.StatusOK, + http.Header{"Content-Encoding": {"gzip"}}, + overflow, + nil, + ) + if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { + t.Fatalf("decoded overflow result/error = %#v/%v, want protocol error", result, err) + } + + result, err = forwarder.prepareSuccessRepresentation( + input, + http.StatusOK, + http.Header{"Content-Encoding": {"identity"}}, + bytes.Repeat([]byte("w"), int(maxNonStreamingResponseBodyBytes)+1), + nil, + ) + if !errors.Is(err, ErrUpstreamProtocol) || result.headers != nil { + t.Fatalf("wire overflow result/error = %#v/%v, want protocol error", result, err) + } +} + func gzipRepeatedByte(t *testing.T, value byte, count int64) []byte { t.Helper() var wire bytes.Buffer @@ -47,10 +388,11 @@ func successRepresentationEncodings(t *testing.T, plain []byte) []successReprese encodings := []string{"", "identity", "gzip", "br", "deflate", "zstd"} result := make([]successRepresentationEncoding, 0, len(encodings)) for _, encoding := range encodings { - wire, err := utils.CompressResponse(encoding, plain) - if err != nil { - t.Fatalf("CompressResponse(%q): %v", encoding, err) + fixtureEncoding := contentcoding.Identity + if encoding != "" { + fixtureEncoding = contentcoding.Encoding(encoding) } + wire := encodeContentCodingForGatewayTest(t, fixtureEncoding, plain) name := encoding if name == "" { name = "empty" @@ -109,7 +451,7 @@ func forwardSuccessRepresentation( ) } -func TestForwardSuccessRepresentationPreservesUnchangedWireAndMetadata(t *testing.T) { +func TestForwardSuccessRepresentationReturnsPlaintextAndConditionalMetadata(t *testing.T) { plain := []byte(`{"id":"safe-response","usage":{"prompt_tokens":100,"completion_tokens":30}}`) for _, representation := range successRepresentationEncodings(t, plain) { t.Run(representation.name, func(t *testing.T) { @@ -123,21 +465,28 @@ func TestForwardSuccessRepresentationPreservesUnchangedWireAndMetadata(t *testin if result.Err != nil || result.StatusCode != http.StatusOK || !result.RequestWritten { t.Fatalf("Forward() result = %#v, want successful written response", result) } - if !bytes.Equal(result.Body, representation.wire) { - t.Fatalf("wire changed without a plain-body change:\n got %x\nwant %x", result.Body, representation.wire) + if !bytes.Equal(result.Body, plain) { + t.Fatalf("downstream body = %x, want plaintext %x", result.Body, plain) } - if result.Header.Get("Content-Encoding") != representation.encoding { - t.Fatalf("Content-Encoding = %q, want %q", result.Header.Get("Content-Encoding"), representation.encoding) + if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { + t.Fatalf("Content-Encoding values = %#v, want absent", values) } - if result.Header.Get("Content-Length") != strconv.Itoa(len(representation.wire)) { - t.Fatalf("Content-Length = %q, want %d", result.Header.Get("Content-Length"), len(representation.wire)) + if result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) { + t.Fatalf("Content-Length = %q, want %d", result.Header.Get("Content-Length"), len(plain)) + } + switch representation.encoding { + case "": + assertRepresentationMetadata(t, result.Header, true) + case "identity": + assertRepresentationValidatorsPreservedWithoutSignatures(t, result.Header) + default: + assertRepresentationMetadata(t, result.Header, false) } - assertRepresentationMetadata(t, result.Header, true) }) } } -func TestPrepareSuccessRepresentationRejectsUnchangedCredentialMetadataCollision(t *testing.T) { +func TestPrepareSuccessRepresentationSanitizesUnchangedCredentialMetadataCollision(t *testing.T) { const credential = "fake-metadata-credential" wire := []byte(`{"id":"safe-response"}`) contentLength := strconv.Itoa(len(wire)) @@ -173,13 +522,25 @@ func TestPrepareSuccessRepresentationRejectsUnchangedCredentialMetadataCollision result, err := NewForwarder( platformhttp.NewHTTPClientManager(), redact.New(), - ).prepareSuccessRepresentation(input, headers, wire, []string{test.secret}) + ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, []string{test.secret}) - if !errors.Is(err, ErrUpstreamProtocol) || - result.wire != nil || - result.headers != nil { + if test.name == "Content-Length" { + if !errors.Is(err, ErrUpstreamProtocol) || + result.downstream != nil || + result.headers != nil { + t.Fatalf( + "prepareSuccessRepresentation() = %#v, %v; want fail-closed regenerated length collision", + result, + err, + ) + } + return + } + if err != nil || !bytes.Equal(result.downstream, wire) || + len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || + result.headers.Get(test.name) != "" { t.Fatalf( - "prepareSuccessRepresentation() = %#v, %v; want fail-closed protocol error", + "prepareSuccessRepresentation() = %#v, %v; want sanitized plaintext response", result, err, ) @@ -200,11 +561,11 @@ func TestPrepareSuccessRepresentationRejectsUnchangedCredentialMetadataCollision result, err := NewForwarder( platformhttp.NewHTTPClientManager(), redact.New(), - ).prepareSuccessRepresentation(input, headers, wire, []string{credential}) + ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, []string{credential}) if err != nil || result.changed || - !bytes.Equal(result.wire, wire) || - result.headers.Get("Content-Encoding") != "identity" || + !bytes.Equal(result.downstream, wire) || + len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || result.headers.Get("Content-Length") != contentLength { t.Fatalf( "non-collision representation = %#v, %v; want unchanged wire and framing", @@ -212,16 +573,13 @@ func TestPrepareSuccessRepresentationRejectsUnchangedCredentialMetadataCollision err, ) } - assertRepresentationMetadata(t, result.headers, true) + assertRepresentationValidatorsPreservedWithoutSignatures(t, result.headers) } func TestPrepareSuccessRepresentationRebuildsChangedBodyDespiteStaleCredentialMetadata(t *testing.T) { const credential = "fake-stale-metadata-credential" plain := []byte(`{"echo":"fake-stale-metadata-credential"}`) - wire, err := utils.CompressResponse("gzip", plain) - if err != nil { - t.Fatalf("compress response fixture: %v", err) - } + wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) headers := http.Header{ "Content-Encoding": {"gzip"}, "Content-Length": {strconv.Itoa(len(wire))}, @@ -237,17 +595,17 @@ func TestPrepareSuccessRepresentationRebuildsChangedBodyDespiteStaleCredentialMe result, err := NewForwarder( platformhttp.NewHTTPClientManager(), redact.New(), - ).prepareSuccessRepresentation(input, headers, wire, []string{credential}) + ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, []string{credential}) if err != nil || !result.changed { t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want rebuilt response", result, err) } - gotPlain, err := utils.DecompressResponse("gzip", result.wire) - if err != nil || string(gotPlain) != `{"echo":"[REDACTED]"}` { - t.Fatalf("rebuilt plain/error = %q/%v", gotPlain, err) + gotPlain := result.downstream + if string(gotPlain) != `{"echo":"[REDACTED]"}` { + t.Fatalf("rebuilt plain = %q", gotPlain) } - if result.headers.Get("Content-Encoding") != "gzip" || - result.headers.Get("Content-Length") != strconv.Itoa(len(result.wire)) { + if len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || + result.headers.Get("Content-Length") != strconv.Itoa(len(result.downstream)) { t.Fatalf("rebuilt framing headers = %#v", result.headers) } assertRepresentationMetadata(t, result.headers, false) @@ -280,12 +638,12 @@ func TestPrepareSuccessRepresentationRemovesAliasedModelFromUnchangedMetadata(t result, err := NewForwarder( platformhttp.NewHTTPClientManager(), redact.New(), - ).prepareSuccessRepresentation(input, headers, wire, nil) + ).prepareSuccessRepresentation(input, http.StatusOK, headers, wire, nil) - if err != nil || result.changed || !bytes.Equal(result.wire, wire) { + if err != nil || result.changed || !bytes.Equal(result.downstream, wire) { t.Fatalf("prepareSuccessRepresentation() = %#v, %v; want unchanged wire", result, err) } - if result.headers.Get("Content-Encoding") != "identity" || + if len(headerFieldValues(result.headers, "Content-Encoding")) != 0 || result.headers.Get("Content-Length") != strconv.Itoa(len(wire)) || result.headers.Get("X-Safe") != "kept" { t.Fatalf("unchanged safe headers = %#v", result.headers) @@ -341,18 +699,15 @@ func TestForwardSuccessRepresentationRedactsAndRebuildsMetadata(t *testing.T) { if result.Err != nil || result.StatusCode != http.StatusOK || !result.RequestWritten { t.Fatalf("Forward() result = %#v, want successful rewritten response", result) } - plain, err := utils.DecompressResponse(representation.encoding, result.Body) - if err != nil { - t.Fatalf("decompress downstream representation: %v", err) - } + plain := result.Body if !bytes.Equal(plain, test.wantPlain) { t.Fatalf("downstream plain = %q, want %q", plain, test.wantPlain) } if bytes.Contains(plain, []byte(apiKey)) { t.Fatalf("downstream plain leaked credential: %q", plain) } - if result.Header.Get("Content-Encoding") != representation.encoding { - t.Fatalf("Content-Encoding = %q, want %q", result.Header.Get("Content-Encoding"), representation.encoding) + if values := headerFieldValues(result.Header, "Content-Encoding"); len(values) != 0 { + t.Fatalf("Content-Encoding values = %#v, want absent", values) } if result.Header.Get("Content-Length") != strconv.Itoa(len(result.Body)) { t.Fatalf("Content-Length = %q, want %d", result.Header.Get("Content-Length"), len(result.Body)) @@ -365,11 +720,6 @@ func TestForwardSuccessRepresentationRedactsAndRebuildsMetadata(t *testing.T) { func TestForwardSuccessRepresentationRejectsUnknownStackedMalformedAndPlainOverflow(t *testing.T) { compressedOverflow := gzipRepeatedByte(t, 'a', maxNonStreamingResponseBodyBytes+1) - validGzip, err := utils.CompressResponse("gzip", []byte(`{"id":"safe-response"}`)) - if err != nil { - t.Fatalf("compress credential collision fixture: %v", err) - } - tests := []struct { name string apiKey string @@ -385,6 +735,10 @@ func TestForwardSuccessRepresentationRejectsUnknownStackedMalformedAndPlainOverf {name: "malformed br", apiKey: "fake-provider-secret-br", encodings: []string{"br"}, wire: []byte{0xff, 0xff, 0xff}}, {name: "malformed deflate", apiKey: "fake-provider-secret-deflate", encodings: []string{"deflate"}, wire: []byte("not-a-deflate-stream")}, {name: "malformed zstd", apiKey: "fake-provider-secret-zstd", encodings: []string{"zstd"}, wire: []byte("not-a-zstd-stream")}, + {name: "empty gzip", apiKey: "fake-provider-secret-empty-gzip", encodings: []string{"gzip"}}, + {name: "empty br", apiKey: "fake-provider-secret-empty-br", encodings: []string{"br"}}, + {name: "empty deflate", apiKey: "fake-provider-secret-empty-deflate", encodings: []string{"deflate"}}, + {name: "empty zstd", apiKey: "fake-provider-secret-empty-zstd", encodings: []string{"zstd"}}, { name: "wire overflow", apiKey: "fake-provider-secret-wire-overflow", @@ -399,12 +753,6 @@ func TestForwardSuccessRepresentationRejectsUnknownStackedMalformedAndPlainOverf encodings: []string{"gzip"}, wire: compressedOverflow, }, - { - name: "Content-Encoding credential collision", - apiKey: "gzip", - encodings: []string{"gzip"}, - wire: validGzip, - }, { name: "Content-Length credential collision", apiKey: "42", @@ -448,6 +796,24 @@ func TestForwardSuccessRepresentationRejectsUnknownStackedMalformedAndPlainOverf } } +func TestForwardSuccessRepresentationDropsContentEncodingCredentialCollision(t *testing.T) { + plain := []byte(`{"id":"safe-response"}`) + result := forwardSuccessRepresentation( + t, + successRepresentationEncoding{ + name: "gzip", + encoding: "gzip", + wire: encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain), + }, + "gzip", + nil, + ) + if result.Err != nil || !bytes.Equal(result.Body, plain) || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { + t.Fatalf("Forward() result = %#v, want plaintext with coding header removed", result) + } +} + func TestForwardSuccessRepresentationFailsClosedOnResidualCredentialTokens(t *testing.T) { tests := []struct { name string @@ -627,10 +993,7 @@ func TestForwardSuccessRepresentationRejectsCredentialIntroducedByAliasRewrite(t for _, test := range tests { t.Run(test.name, func(t *testing.T) { - wire, err := utils.CompressResponse(test.encoding, []byte(providerBody)) - if err != nil { - t.Fatalf("compress provider response: %v", err) - } + wire := encodeContentCodingForGatewayTest(t, contentcoding.Encoding(test.encoding), []byte(providerBody)) result := forwardSuccessRepresentation( t, successRepresentationEncoding{ @@ -676,10 +1039,7 @@ func TestForwardSuccessRepresentationExtractsUsageFromPlainBody(t *testing.T) { `{"model":"provider-model","echo":"fake-provider-secret-usage",` + `"usage":{"prompt_tokens":100,"completion_tokens":30,"prompt_tokens_details":{"cached_tokens":20}}}`, ) - wire, err := utils.CompressResponse("gzip", plain) - if err != nil { - t.Fatalf("compress usage fixture: %v", err) - } + wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) var captured []byte selected := usageExtractorDialect{ Dialect: dialect.NewOpenAI(http.DefaultClient), @@ -726,10 +1086,7 @@ func TestForwardSuccessRepresentationExtractsUsageFromPlainBody(t *testing.T) { bytes.Contains(captured, []byte(externalModel)) { t.Fatalf("usage extractor body = %q, want safe provider plain before alias rewrite", captured) } - downstreamPlain, err := utils.DecompressResponse("gzip", result.Body) - if err != nil { - t.Fatalf("decompress downstream response: %v", err) - } + downstreamPlain := result.Body if bytes.Contains(downstreamPlain, []byte(apiKey)) || !bytes.Contains(downstreamPlain, []byte(redact.Placeholder)) || !bytes.Contains(downstreamPlain, []byte(externalModel)) || @@ -737,7 +1094,7 @@ func TestForwardSuccessRepresentationExtractsUsageFromPlainBody(t *testing.T) { t.Fatalf("downstream plain = %q", downstreamPlain) } if result.Header.Get("Content-Length") != strconv.Itoa(len(result.Body)) || - !strings.EqualFold(result.Header.Get("Content-Encoding"), "gzip") { + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { t.Fatalf("downstream representation headers = %#v", result.Header) } assertRepresentationMetadata(t, result.Header, false) diff --git a/internal/gateway/stream.go b/internal/gateway/stream.go index fa0f2a69..569d8200 100644 --- a/internal/gateway/stream.go +++ b/internal/gateway/stream.go @@ -225,7 +225,7 @@ func commitStream(writer *streamWriteController, status int, headers http.Header if writer == nil || writer.writer == nil { return fmt.Errorf("downstream response writer is required") } - for name, values := range cloneEndToEndHeaders(headers) { + for name, values := range normalizeStreamResponseHeaders(headers) { for _, value := range values { writer.writer.Header().Add(name, value) } @@ -260,6 +260,14 @@ func commitStream(writer *streamWriteController, status int, headers http.Header return nil } +func normalizeStreamResponseHeaders(headers http.Header) http.Header { + normalized := cloneEndToEndHeaders(headers) + for _, name := range representationMetadataHeaderNames { + deleteHeaderField(normalized, name) + } + return normalized +} + type streamWatchdogController interface { reset() interrupt(error) diff --git a/internal/gateway/stream_test.go b/internal/gateway/stream_test.go index d7cdd569..44631e73 100644 --- a/internal/gateway/stream_test.go +++ b/internal/gateway/stream_test.go @@ -6,6 +6,7 @@ import ( "errors" "io" "net/http" + "reflect" "strings" "sync" "testing" @@ -307,6 +308,41 @@ func TestCommitStreamWritesHeadersPrefixAndFlushes(t *testing.T) { } } +func TestCommitStreamStripsRepresentationMetadata(t *testing.T) { + writer := newRecordingResponseWriter() + headers := http.Header{ + "Content-Type": {"text/event-stream"}, + "X-Upstream": {"kept"}, + "content-encoding": {"identity"}, + "CONTENT-LENGTH": {"123"}, + "eTAG": {`"stale"`}, + "dIgEsT": {"sha-256=stale"}, + "content-md5": {"stale"}, + "Content-Range": {"bytes 0-1/2"}, + "content-digest": {"sha-256=:stale:"}, + "Repr-Digest": {"sha-256=:stale:"}, + "signature": {"sig"}, + "SIGNATURE-INPUT": {"sig1=(\"content-length\")"}, + } + original := headers.Clone() + controller := newStreamWriteController(writer, time.Second) + + if err := commitStream(controller, http.StatusOK, headers, []byte("data: first\n\n")); err != nil { + t.Fatalf("commitStream() error = %v", err) + } + if !reflect.DeepEqual(headers, original) { + t.Fatalf("commitStream() mutated source headers: got %#v, want %#v", headers, original) + } + for _, name := range representationMetadataHeaderNames { + if values := headerFieldValues(writer.header, name); len(values) != 0 { + t.Fatalf("committed %s values = %#v, want absent", name, values) + } + } + if writer.header.Get("Content-Type") != "text/event-stream" || writer.header.Get("X-Upstream") != "kept" { + t.Fatalf("safe headers = %#v", writer.header) + } +} + func TestStreamWriteControllerArmsEveryOperationAndClears(t *testing.T) { writer := newRecordingResponseWriter() controller := newStreamWriteController(writer, time.Minute) diff --git a/internal/gateway/usage_capture.go b/internal/gateway/usage_capture.go index 910aa779..867065c4 100644 --- a/internal/gateway/usage_capture.go +++ b/internal/gateway/usage_capture.go @@ -2,7 +2,6 @@ package gateway import ( "bytes" - "net/http" "reflect" "sync" "sync/atomic" @@ -289,32 +288,6 @@ func (capture *streamUsageCapture) finalize() usage.Result { return capture.result } -func (boundary *usageCaptureBoundary) extractNonStreaming( - selected dialect.Dialect, - headers http.Header, - wire []byte, -) usage.Result { - _, ok := selected.(dialect.UsageExtractor) - if !ok { - return missingUsage(false) - } - encoding, ok := inspectableErrorBodyEncoding(headers, wire) - if !ok { - boundary.recordFailure("decompress", selected.Protocol()) - return missingUsage(false) - } - body, err := utils.DecompressResponseLimited( - encoding, - wire, - maxNonStreamingResponseBodyBytes, - ) - if err != nil { - boundary.recordFailure("decompress", selected.Protocol()) - return missingUsage(false) - } - return boundary.extractNonStreamingPlain(selected, body) -} - func (boundary *usageCaptureBoundary) extractNonStreamingPlain( selected dialect.Dialect, plain []byte, diff --git a/internal/gateway/usage_capture_test.go b/internal/gateway/usage_capture_test.go index 663f7168..e78a1112 100644 --- a/internal/gateway/usage_capture_test.go +++ b/internal/gateway/usage_capture_test.go @@ -22,9 +22,9 @@ import ( "gpt-load/internal/dialect" "gpt-load/internal/health" + "gpt-load/internal/platform/contentcoding" platformhttp "gpt-load/internal/platform/httpclient" "gpt-load/internal/platform/redact" - "gpt-load/internal/platform/utils" "gpt-load/internal/protocol" "gpt-load/internal/state" "gpt-load/internal/usage" @@ -881,10 +881,10 @@ func TestUsageCaptureBoundaryExtractsValidNonStreamingResult(t *testing.T) { Tokens: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}, } boundary := newUsageCaptureBoundary() - got := boundary.extractNonStreaming(usageExtractorDialect{ + got := boundary.extractNonStreamingPlain(usageExtractorDialect{ Dialect: dialect.NewOpenAI(http.DefaultClient), extract: func([]byte) (usage.Result, error) { return want, nil }, - }, make(http.Header), []byte(`{"provider":"body"}`)) + }, []byte(`{"provider":"body"}`)) if got != want || boundary.failureTotal.Load() != 0 { t.Fatalf("result/failures = %#v/%d, want %#v/0", got, boundary.failureTotal.Load(), want) @@ -896,7 +896,7 @@ func TestUsageCaptureBoundaryMissingCapabilityIsQuietMissing(t *testing.T) { var logs bytes.Buffer boundary.logger = logrus.New() boundary.logger.SetOutput(&logs) - got := boundary.extractNonStreaming(dialectWithoutUsage{Dialect: dialect.NewOpenAI(http.DefaultClient)}, make(http.Header), []byte(`{}`)) + got := boundary.extractNonStreamingPlain(dialectWithoutUsage{Dialect: dialect.NewOpenAI(http.DefaultClient)}, []byte(`{}`)) if got != (usage.Result{State: usage.StateMissing}) || boundary.failureTotal.Load() != 0 || logs.Len() != 0 { t.Fatalf("result/failures/logs = %#v/%d/%q", got, boundary.failureTotal.Load(), logs.String()) @@ -917,7 +917,7 @@ func TestUsageCaptureBoundaryRecoversExtractErrorAndPanic(t *testing.T) { var logs bytes.Buffer boundary.logger = logrus.New() boundary.logger.SetOutput(&logs) - got := boundary.extractNonStreaming(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: test.extract}, make(http.Header), []byte(`{}`)) + got := boundary.extractNonStreamingPlain(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: test.extract}, []byte(`{}`)) if got.State != usage.StateMissing || !got.Diagnostics.Has(usage.DiagnosticInvalidPayload) || boundary.failureTotal.Load() != 1 || bytes.Contains(logs.Bytes(), []byte(canary)) { t.Fatalf("result/failures/logs = %#v/%d/%q", got, boundary.failureTotal.Load(), logs.String()) @@ -937,7 +937,7 @@ func TestUsageCaptureBoundaryRejectsInvalidResults(t *testing.T) { {State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: math.MaxInt64, Output: 1}}, } { boundary := newUsageCaptureBoundary() - got := boundary.extractNonStreaming(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: func([]byte) (usage.Result, error) { return result, nil }}, make(http.Header), []byte(`{}`)) + got := boundary.extractNonStreamingPlain(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: func([]byte) (usage.Result, error) { return result, nil }}, []byte(`{}`)) if got.State != usage.StateMissing || !got.Diagnostics.Has(usage.DiagnosticInvalidPayload) || boundary.failureTotal.Load() != 1 { t.Fatalf("input/result/failures = %#v/%#v/%d", result, got, boundary.failureTotal.Load()) } @@ -945,65 +945,11 @@ func TestUsageCaptureBoundaryRejectsInvalidResults(t *testing.T) { boundary := newUsageCaptureBoundary() zero := usage.Result{State: usage.StateComplete} - if got := boundary.extractNonStreaming(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: func([]byte) (usage.Result, error) { return zero, nil }}, make(http.Header), []byte(`{}`)); got != zero || boundary.failureTotal.Load() != 0 { + if got := boundary.extractNonStreamingPlain(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: func([]byte) (usage.Result, error) { return zero, nil }}, []byte(`{}`)); got != zero || boundary.failureTotal.Load() != 0 { t.Fatalf("zero complete/failures = %#v/%d", got, boundary.failureTotal.Load()) } } -func TestUsageCaptureBoundaryDecodesSupportedContentEncoding(t *testing.T) { - plain := []byte(`{"canonical":true}`) - want := usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}} - for _, encoding := range []string{"identity", "gzip", "br", "deflate", "zstd"} { - t.Run(encoding, func(t *testing.T) { - wire, err := utils.CompressResponse(encoding, plain) - if err != nil { - t.Fatalf("CompressResponse(%q): %v", encoding, err) - } - originalWire := bytes.Clone(wire) - boundary := newUsageCaptureBoundary() - got := boundary.extractNonStreaming(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: func(body []byte) (usage.Result, error) { - if !bytes.Equal(body, plain) { - t.Fatalf("decoded body = %q, want %q", body, plain) - } - return want, nil - }}, http.Header{"Content-Encoding": {encoding}}, wire) - if got != want || boundary.failureTotal.Load() != 0 || !bytes.Equal(wire, originalWire) { - t.Fatalf("result/failures = %#v/%d", got, boundary.failureTotal.Load()) - } - }) - } -} - -func TestUsageCaptureBoundaryUnsupportedEncodingIsMissing(t *testing.T) { - tooLargePlain := bytes.Repeat([]byte("x"), int(maxNonStreamingResponseBodyBytes)+1) - tooLargeWire, err := utils.CompressResponse("gzip", tooLargePlain) - if err != nil { - t.Fatalf("compress oversized fixture: %v", err) - } - for _, test := range []struct { - name string - headers http.Header - wire []byte - }{ - {name: "multiple values", headers: http.Header{"Content-Encoding": {"identity", "gzip"}}, wire: []byte(`{}`)}, - {name: "comma stacked", headers: http.Header{"Content-Encoding": {"gzip, br"}}, wire: []byte(`{}`)}, - {name: "unknown", headers: http.Header{"Content-Encoding": {"unknown"}}, wire: []byte(`{}`)}, - {name: "corrupt gzip", headers: http.Header{"Content-Encoding": {"gzip"}}, wire: []byte("not-gzip")}, - {name: "decompressed over limit", headers: http.Header{"Content-Encoding": {"gzip"}}, wire: tooLargeWire}, - } { - t.Run(test.name, func(t *testing.T) { - boundary := newUsageCaptureBoundary() - got := boundary.extractNonStreaming(usageExtractorDialect{Dialect: dialect.NewOpenAI(http.DefaultClient), extract: func([]byte) (usage.Result, error) { - t.Fatal("extractor called after decompression failure") - return usage.Result{}, nil - }}, test.headers, test.wire) - if got != (usage.Result{State: usage.StateMissing}) || boundary.failureTotal.Load() != 1 { - t.Fatalf("result/failures = %#v/%d", got, boundary.failureTotal.Load()) - } - }) - } -} - func TestForwardCapturesCanonicalNonStreamingUsage(t *testing.T) { want := usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}} tests := []struct { @@ -1038,12 +984,9 @@ func TestForwardCapturesCanonicalNonStreamingUsage(t *testing.T) { } } -func TestForwardUsageCapturePreservesCompressedSuccessWire(t *testing.T) { +func TestForwarderUsesPreparedPlainBodyForUsage(t *testing.T) { plain := []byte(`{"usage":{"prompt_tokens":100,"completion_tokens":30,"prompt_tokens_details":{"cached_tokens":20}}}`) - wire, err := utils.CompressResponse("gzip", plain) - if err != nil { - t.Fatalf("compress fixture: %v", err) - } + wire := encodeContentCodingForGatewayTest(t, contentcoding.Gzip, plain) upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { writer.Header().Set("Content-Encoding", "gzip") writer.Header().Set("Content-Length", strconv.Itoa(len(wire))) @@ -1059,7 +1002,10 @@ func TestForwardUsageCapturePreservesCompressedSuccessWire(t *testing.T) { result := NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()).Forward(context.Background(), input) baseline := forwardWithCaptureDisabled(input) want := usage.Result{State: usage.StateComplete, Tokens: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}} - if result.Usage != want || baseline.Usage != (usage.Result{State: usage.StateNotApplicable}) || !bytes.Equal(result.Body, wire) || result.Header.Get("Content-Length") != strconv.Itoa(len(wire)) { + if result.Usage != want || baseline.Usage != (usage.Result{State: usage.StateNotApplicable}) || + !bytes.Equal(result.Body, plain) || + result.Header.Get("Content-Length") != strconv.Itoa(len(plain)) || + len(headerFieldValues(result.Header, "Content-Encoding")) != 0 { t.Fatalf("result = %#v", result) } assertForwardWireContract(t, selected, result, baseline) diff --git a/internal/platform/contentcoding/accept.go b/internal/platform/contentcoding/accept.go new file mode 100644 index 00000000..26bbb612 --- /dev/null +++ b/internal/platform/contentcoding/accept.go @@ -0,0 +1,93 @@ +package contentcoding + +import "strings" + +func IdentityAcceptable(values []string) bool { + explicitQuality, explicitFound := -1, false + wildcardQuality, wildcardFound := -1, false + + for _, value := range values { + for _, rawItem := range strings.Split(value, ",") { + token, quality, ok := parseAcceptEncodingItem(rawItem) + if !ok { + continue + } + switch token { + case string(Identity): + explicitFound = true + explicitQuality = max(explicitQuality, quality) + case "*": + wildcardFound = true + wildcardQuality = max(wildcardQuality, quality) + } + } + } + + if explicitFound { + return explicitQuality > 0 + } + if wildcardFound { + return wildcardQuality > 0 + } + return true +} + +func parseAcceptEncodingItem(raw string) (string, int, bool) { + parts := strings.Split(raw, ";") + token := strings.ToLower(strings.TrimSpace(parts[0])) + if token == "" { + return "", 0, false + } + if len(parts) == 1 { + return token, 1000, true + } + if len(parts) != 2 { + return "", 0, false + } + + parameter := strings.SplitN(parts[1], "=", 2) + if len(parameter) != 2 || !strings.EqualFold(strings.TrimSpace(parameter[0]), "q") { + return "", 0, false + } + quality, ok := parseQvalue(strings.TrimSpace(parameter[1])) + if !ok { + return "", 0, false + } + return token, quality, true +} + +func parseQvalue(value string) (int, bool) { + if value == "0" { + return 0, true + } + if value == "1" { + return 1000, true + } + if len(value) < 2 || value[1] != '.' { + return 0, false + } + + fraction := value[2:] + if len(fraction) > 3 { + return 0, false + } + quality := 0 + for _, digit := range fraction { + if digit < '0' || digit > '9' { + return 0, false + } + quality = quality*10 + int(digit-'0') + } + for digits := len(fraction); digits < 3; digits++ { + quality *= 10 + } + + switch value[0] { + case '0': + return quality, true + case '1': + return 1000, quality == 0 + default: + return 0, false + } +} diff --git a/internal/platform/contentcoding/accept_test.go b/internal/platform/contentcoding/accept_test.go new file mode 100644 index 00000000..3d357b67 --- /dev/null +++ b/internal/platform/contentcoding/accept_test.go @@ -0,0 +1,38 @@ +package contentcoding + +import "testing" + +func TestIdentityAcceptable(t *testing.T) { + tests := []struct { + name string + values []string + want bool + }{ + {name: "missing", want: true}, + {name: "empty", values: []string{""}, want: true}, + {name: "unrelated coding defaults to identity", values: []string{"gzip"}, want: true}, + {name: "explicit identity", values: []string{"identity"}, want: true}, + {name: "explicit identity excluded", values: []string{"identity;q=0"}, want: false}, + {name: "wildcard excludes identity", values: []string{"*;q=0"}, want: false}, + {name: "explicit identity overrides wildcard", values: []string{"identity;q=0, *;q=1"}, want: false}, + {name: "duplicate identity uses highest valid quality", values: []string{"identity;q=0.5, identity;q=0"}, want: true}, + {name: "wildcard allows identity", values: []string{"gzip;q=1, *;q=0.5"}, want: true}, + {name: "invalid identity quality does not exclude default", values: []string{"identity;q=1.001"}, want: true}, + {name: "zero with three decimal places", values: []string{"identity;q=0.000"}, want: false}, + {name: "multiple fields", values: []string{"gzip", "*;q=0"}, want: false}, + {name: "case insensitive token and parameter", values: []string{"IDENTITY;Q=0"}, want: false}, + {name: "one with zero decimals is valid", values: []string{"identity;q=1.000"}, want: true}, + {name: "too many decimals is invalid", values: []string{"identity;q=0.0000, *;q=0"}, want: false}, + {name: "one with nonzero decimal is invalid", values: []string{"identity;q=1.010, *;q=0"}, want: false}, + {name: "unknown parameter invalidates item", values: []string{"identity;level=1, *;q=0"}, want: false}, + {name: "duplicate quality parameter invalidates item", values: []string{"identity;q=1;q=0, *;q=0"}, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IdentityAcceptable(tt.values); got != tt.want { + t.Fatalf("IdentityAcceptable(%q) = %t, want %t", tt.values, got, tt.want) + } + }) + } +} diff --git a/internal/platform/contentcoding/contentcoding.go b/internal/platform/contentcoding/contentcoding.go new file mode 100644 index 00000000..db2d501f --- /dev/null +++ b/internal/platform/contentcoding/contentcoding.go @@ -0,0 +1,178 @@ +package contentcoding + +import ( + "bytes" + "compress/gzip" + "compress/zlib" + "errors" + "fmt" + "io" + "math" + "strings" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" +) + +type Encoding string + +const ( + Identity Encoding = "identity" + Gzip Encoding = "gzip" + Brotli Encoding = "br" + Deflate Encoding = "deflate" + Zstd Encoding = "zstd" +) + +const ( + SupportedRequestEncodings = "identity, gzip, br, deflate, zstd" + ZstdHTTPMaxWindowBytes = 8 << 20 +) + +var ( + ErrUnsupportedContentEncoding = errors.New("unsupported content encoding") + ErrInvalidContentEncoding = errors.New("invalid content encoding") + ErrDecodedBodyTooLarge = errors.New("decoded body exceeds size limit") +) + +func ParseContentEncoding(values []string) (Encoding, error) { + if len(values) == 0 { + return Identity, nil + } + if len(values) != 1 { + return "", ErrUnsupportedContentEncoding + } + + value := strings.ToLower(strings.TrimSpace(values[0])) + if value == "" || value == string(Identity) { + return Identity, nil + } + if strings.Contains(value, ",") { + return "", ErrUnsupportedContentEncoding + } + + switch Encoding(value) { + case Gzip, Brotli, Deflate, Zstd: + return Encoding(value), nil + default: + return "", ErrUnsupportedContentEncoding + } +} + +func DecodeLimited(encoding Encoding, encoded []byte, limit int64) ([]byte, error) { + if limit < 0 { + return nil, fmt.Errorf("decoded body limit must not be negative") + } + + switch encoding { + case Identity: + return readAtMost(bytes.NewReader(encoded), limit) + case Gzip: + if len(encoded) == 0 { + return nil, ErrInvalidContentEncoding + } + return decodeGzip(encoded, limit) + case Brotli: + if len(encoded) == 0 { + return nil, ErrInvalidContentEncoding + } + return readDecoded(brotli.NewReader(bytes.NewReader(encoded)), limit) + case Deflate: + if len(encoded) == 0 { + return nil, ErrInvalidContentEncoding + } + return decodeDeflate(encoded, limit) + case Zstd: + if len(encoded) == 0 { + return nil, ErrInvalidContentEncoding + } + return decodeZstd(encoded, limit) + default: + return nil, ErrUnsupportedContentEncoding + } +} + +func decodeGzip(encoded []byte, limit int64) ([]byte, error) { + source := bytes.NewReader(encoded) + reader, err := gzip.NewReader(source) + if err != nil { + return nil, invalidContentEncoding(err) + } + reader.Multistream(true) + + body, readErr := readDecoded(reader, limit) + closeErr := reader.Close() + if readErr != nil { + return nil, readErr + } + if closeErr != nil { + return nil, invalidContentEncoding(closeErr) + } + return body, nil +} + +func decodeDeflate(encoded []byte, limit int64) ([]byte, error) { + source := bytes.NewReader(encoded) + reader, err := zlib.NewReader(source) + if err != nil { + return nil, invalidContentEncoding(err) + } + body, readErr := readDecoded(reader, limit) + closeErr := reader.Close() + if readErr != nil { + return nil, readErr + } + if closeErr != nil { + return nil, invalidContentEncoding(closeErr) + } + if source.Len() != 0 { + return nil, ErrInvalidContentEncoding + } + return body, nil +} + +func decodeZstd(encoded []byte, limit int64) ([]byte, error) { + memoryLimit := uint64(max(limit, int64(ZstdHTTPMaxWindowBytes))) + reader, err := zstd.NewReader( + bytes.NewReader(encoded), + zstd.WithDecoderConcurrency(1), + zstd.WithDecoderLowmem(true), + zstd.WithDecoderMaxMemory(memoryLimit), + zstd.WithDecoderMaxWindow(ZstdHTTPMaxWindowBytes), + ) + if err != nil { + return nil, invalidContentEncoding(err) + } + defer reader.Close() + return readDecoded(reader, limit) +} + +func readDecoded(reader io.Reader, limit int64) ([]byte, error) { + body, err := readAtMost(reader, limit) + if err == nil || errors.Is(err, ErrDecodedBodyTooLarge) { + return body, err + } + return nil, invalidContentEncoding(err) +} + +func readAtMost(reader io.Reader, limit int64) ([]byte, error) { + if limit < 0 { + return nil, fmt.Errorf("decoded body limit must not be negative") + } + limited := io.Reader(reader) + if limit < math.MaxInt64 { + limited = io.LimitReader(reader, limit+1) + } + body, err := io.ReadAll(limited) + if err != nil { + return nil, err + } + if int64(len(body)) > limit { + return nil, ErrDecodedBodyTooLarge + } + return body, nil +} + +func invalidContentEncoding(err error) error { + return fmt.Errorf("%w: %v", ErrInvalidContentEncoding, err) +} diff --git a/internal/platform/contentcoding/contentcoding_test.go b/internal/platform/contentcoding/contentcoding_test.go new file mode 100644 index 00000000..9cf2f28b --- /dev/null +++ b/internal/platform/contentcoding/contentcoding_test.go @@ -0,0 +1,299 @@ +package contentcoding + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "compress/zlib" + "errors" + "io" + "testing" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" +) + +func TestParseContentEncoding(t *testing.T) { + tests := []struct { + name string + values []string + want Encoding + wantErr error + }{ + {name: "missing is identity", want: Identity}, + {name: "empty is identity", values: []string{" "}, want: Identity}, + {name: "identity is case insensitive", values: []string{" IDENTITY "}, want: Identity}, + {name: "gzip", values: []string{"gzip"}, want: Gzip}, + {name: "brotli", values: []string{"BR"}, want: Brotli}, + {name: "deflate", values: []string{"deflate"}, want: Deflate}, + {name: "zstd", values: []string{"zstd"}, want: Zstd}, + {name: "multiple fields", values: []string{"gzip", "br"}, wantErr: ErrUnsupportedContentEncoding}, + {name: "stacked coding", values: []string{"gzip, br"}, wantErr: ErrUnsupportedContentEncoding}, + {name: "unknown coding", values: []string{"compress"}, wantErr: ErrUnsupportedContentEncoding}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ParseContentEncoding(tt.values) + if !errors.Is(err, tt.wantErr) { + t.Fatalf("ParseContentEncoding(%q) error = %v, want errors.Is(_, %v)", tt.values, err, tt.wantErr) + } + if got != tt.want { + t.Fatalf("ParseContentEncoding(%q) = %q, want %q", tt.values, got, tt.want) + } + }) + } +} + +func TestDecodeLimited(t *testing.T) { + plain := []byte("content coding boundary") + tests := []struct { + name string + encoding Encoding + encoded []byte + }{ + {name: "identity", encoding: Identity, encoded: plain}, + {name: "gzip", encoding: Gzip, encoded: encodeGzip(t, plain)}, + {name: "brotli", encoding: Brotli, encoded: encodeBrotli(t, plain)}, + {name: "zlib deflate", encoding: Deflate, encoded: encodeZlib(t, plain)}, + {name: "zstd", encoding: Zstd, encoded: encodeZstd(t, plain, ZstdHTTPMaxWindowBytes)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := DecodeLimited(tt.encoding, tt.encoded, int64(len(plain))) + if err != nil { + t.Fatalf("DecodeLimited(%q) error = %v", tt.encoding, err) + } + if !bytes.Equal(got, plain) { + t.Fatalf("DecodeLimited(%q) = %q, want %q", tt.encoding, got, plain) + } + }) + } +} + +func TestDecodeLimitedRejectsInvalidStreams(t *testing.T) { + plain := []byte("content coding boundary") + tests := []struct { + name string + encoding Encoding + encoded []byte + }{ + {name: "gzip empty", encoding: Gzip}, + {name: "brotli empty", encoding: Brotli}, + {name: "deflate empty", encoding: Deflate}, + {name: "zstd empty", encoding: Zstd}, + {name: "gzip truncated", encoding: Gzip, encoded: truncate(encodeGzip(t, plain))}, + {name: "brotli truncated", encoding: Brotli, encoded: truncate(encodeBrotli(t, plain))}, + {name: "deflate truncated", encoding: Deflate, encoded: truncate(encodeZlib(t, plain))}, + {name: "zstd truncated", encoding: Zstd, encoded: truncate(encodeZstd(t, plain, ZstdHTTPMaxWindowBytes))}, + {name: "gzip corrupt", encoding: Gzip, encoded: []byte("not a gzip stream")}, + {name: "brotli corrupt", encoding: Brotli, encoded: []byte("not a brotli stream")}, + {name: "deflate corrupt", encoding: Deflate, encoded: []byte("not a zlib stream")}, + {name: "zstd corrupt", encoding: Zstd, encoded: []byte("not a zstd stream")}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := DecodeLimited(tt.encoding, tt.encoded, 1024) + if !errors.Is(err, ErrInvalidContentEncoding) || got != nil { + t.Fatalf("DecodeLimited(%q, %q) = %q, %v; want invalid content encoding", tt.encoding, tt.encoded, got, err) + } + }) + } +} + +func TestDecodeLimitedHonorsDecodedLimit(t *testing.T) { + plain := []byte("limit") + tests := []struct { + name string + encoding Encoding + encoded []byte + }{ + {name: "identity", encoding: Identity, encoded: plain}, + {name: "gzip", encoding: Gzip, encoded: encodeGzip(t, plain)}, + {name: "brotli", encoding: Brotli, encoded: encodeBrotli(t, plain)}, + {name: "deflate", encoding: Deflate, encoded: encodeZlib(t, plain)}, + {name: "zstd", encoding: Zstd, encoded: encodeZstd(t, plain, ZstdHTTPMaxWindowBytes)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := DecodeLimited(tt.encoding, tt.encoded, int64(len(plain)-1)) + if !errors.Is(err, ErrDecodedBodyTooLarge) || got != nil { + t.Fatalf("DecodeLimited(%q) = %q, %v; want decoded body too large", tt.encoding, got, err) + } + }) + } +} + +func TestDecodeLimitedIdentityReturnsCopy(t *testing.T) { + encoded := []byte("identity") + got, err := DecodeLimited(Identity, encoded, int64(len(encoded))) + if err != nil { + t.Fatal(err) + } + got[0] = 'I' + if encoded[0] != 'i' { + t.Fatal("DecodeLimited(Identity) returned an aliased slice") + } +} + +func TestDecodeLimitedRejectsRawDeflate(t *testing.T) { + plain := []byte("raw deflate is not HTTP deflate") + var encoded bytes.Buffer + writer, err := flate.NewWriter(&encoded, flate.DefaultCompression) + if err != nil { + t.Fatal(err) + } + if _, err := writer.Write(plain); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + + got, err := DecodeLimited(Deflate, encoded.Bytes(), 1024) + if !errors.Is(err, ErrInvalidContentEncoding) || got != nil { + t.Fatalf("DecodeLimited(raw deflate) = %q, %v; want invalid content encoding", got, err) + } +} + +func TestDecodeLimitedRejectsTrailingData(t *testing.T) { + plain := []byte("trailing bytes must not be ignored") + tests := []struct { + name string + encoding Encoding + encoded []byte + }{ + {name: "gzip", encoding: Gzip, encoded: append(encodeGzip(t, plain), []byte("trailing")...)}, + {name: "zlib deflate", encoding: Deflate, encoded: append(encodeZlib(t, plain), []byte("trailing")...)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := DecodeLimited(tt.encoding, tt.encoded, 1024) + if !errors.Is(err, ErrInvalidContentEncoding) || got != nil { + t.Fatalf("DecodeLimited(%q) = %q, %v; want invalid content encoding", tt.encoding, got, err) + } + }) + } +} + +func TestDecodeLimitedAcceptsConcatenatedGzipMembers(t *testing.T) { + first := []byte("first ") + second := []byte("second") + encoded := append(encodeGzip(t, first), encodeGzip(t, second)...) + + got, err := DecodeLimited(Gzip, encoded, int64(len(first)+len(second))) + if err != nil { + t.Fatalf("DecodeLimited() error = %v", err) + } + if want := append(first, second...); !bytes.Equal(got, want) { + t.Fatalf("DecodeLimited() = %q, want %q", got, want) + } +} + +func TestZstdHTTPWindow(t *testing.T) { + plain := []byte("small body with an explicitly declared zstd window") + t.Run("maximum window is accepted below decoded limit", func(t *testing.T) { + encoded := encodeZstdWithDeclaredWindow(t, plain, ZstdHTTPMaxWindowBytes) + got, err := DecodeLimited(Zstd, encoded, int64(len(plain))) + if err != nil || !bytes.Equal(got, plain) { + t.Fatalf("DecodeLimited() = %q, %v", got, err) + } + }) + t.Run("window above maximum is rejected", func(t *testing.T) { + encoded := encodeZstdWithDeclaredWindow(t, plain, 16<<20) + got, err := DecodeLimited(Zstd, encoded, 1024) + if !errors.Is(err, ErrInvalidContentEncoding) || got != nil { + t.Fatalf("DecodeLimited() = %q, %v; want invalid content encoding", got, err) + } + }) + t.Run("output limit is independent from window limit", func(t *testing.T) { + encoded := encodeZstdWithDeclaredWindow(t, plain, ZstdHTTPMaxWindowBytes) + got, err := DecodeLimited(Zstd, encoded, int64(len(plain)-1)) + if !errors.Is(err, ErrDecodedBodyTooLarge) || got != nil { + t.Fatalf("DecodeLimited() = %q, %v; want decoded body too large", got, err) + } + }) +} + +func encodeZstdWithDeclaredWindow(t *testing.T, plain []byte, window int) []byte { + t.Helper() + encoded := encodeZstd(t, plain, window) + + var header zstd.Header + if err := header.Decode(encoded); err != nil { + t.Fatalf("decode generated zstd header: %v", err) + } + header.SingleSegment = false + header.WindowSize = uint64(window) + rebuilt, err := header.AppendTo(nil) + if err != nil { + t.Fatalf("rebuild zstd header: %v", err) + } + rebuilt = append(rebuilt, encoded[header.HeaderSize:]...) + + var got zstd.Header + if err := got.Decode(rebuilt); err != nil { + t.Fatalf("decode rebuilt zstd header: %v", err) + } + if got.SingleSegment || got.WindowSize != uint64(window) { + t.Fatalf("zstd fixture header = {SingleSegment:%t WindowSize:%d}, want {false %d}", got.SingleSegment, got.WindowSize, window) + } + return rebuilt +} + +func encodeGzip(t *testing.T, plain []byte) []byte { + t.Helper() + var encoded bytes.Buffer + writer := gzip.NewWriter(&encoded) + writeAndClose(t, writer, plain) + return encoded.Bytes() +} + +func encodeBrotli(t *testing.T, plain []byte) []byte { + t.Helper() + var encoded bytes.Buffer + writer := brotli.NewWriter(&encoded) + writeAndClose(t, writer, plain) + return encoded.Bytes() +} + +func encodeZlib(t *testing.T, plain []byte) []byte { + t.Helper() + var encoded bytes.Buffer + writer := zlib.NewWriter(&encoded) + writeAndClose(t, writer, plain) + return encoded.Bytes() +} + +func encodeZstd(t *testing.T, plain []byte, window int) []byte { + t.Helper() + encoder, err := zstd.NewWriter( + nil, + zstd.WithEncoderConcurrency(1), + zstd.WithSingleSegment(false), + zstd.WithWindowSize(window), + ) + if err != nil { + t.Fatal(err) + } + defer encoder.Close() + return encoder.EncodeAll(plain, nil) +} + +func writeAndClose(t *testing.T, writer io.WriteCloser, plain []byte) { + t.Helper() + if _, err := writer.Write(plain); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } +} + +func truncate(encoded []byte) []byte { + return append([]byte(nil), encoded[:len(encoded)-1]...) +} diff --git a/internal/platform/httpheader/policy.go b/internal/platform/httpheader/policy.go index 33401cc8..c810784f 100644 --- a/internal/platform/httpheader/policy.go +++ b/internal/platform/httpheader/policy.go @@ -13,7 +13,7 @@ var credentialNames = map[string]struct{}{ "x-goog-api-key": {}, } -var forbiddenRequestRuleSetNames = map[string]struct{}{ +var forbiddenRequestRuleNames = map[string]struct{}{ "connection": {}, "proxy-connection": {}, "keep-alive": {}, @@ -23,6 +23,8 @@ var forbiddenRequestRuleSetNames = map[string]struct{}{ "upgrade": {}, "cookie": {}, "cookie2": {}, + "accept-encoding": {}, + "content-encoding": {}, } func IsCredentialName(name string) bool { @@ -34,7 +36,7 @@ func IsCredentialName(name string) bool { return exists } -func IsForbiddenRequestRuleSetName(name string) bool { +func IsForbiddenRequestRuleName(name string) bool { normalized := normalizeName(name) if normalized == "" { return false @@ -42,10 +44,14 @@ func IsForbiddenRequestRuleSetName(name string) bool { if strings.HasPrefix(normalized, "proxy-") { return true } - _, exists := forbiddenRequestRuleSetNames[normalized] + _, exists := forbiddenRequestRuleNames[normalized] return exists } +func IsForbiddenRequestRuleSetName(name string) bool { + return IsForbiddenRequestRuleName(name) +} + func normalizeName(name string) string { if name == "" { return "" diff --git a/internal/platform/httpheader/policy_test.go b/internal/platform/httpheader/policy_test.go index a2f80c0f..c8540910 100644 --- a/internal/platform/httpheader/policy_test.go +++ b/internal/platform/httpheader/policy_test.go @@ -41,7 +41,8 @@ func TestCredentialAndForbiddenHeaderPolicy(t *testing.T) { {name: "Proxy-Authorization", want: true}, {name: "pRoXy-Custom", want: true}, {name: "Authorization"}, - {name: "Accept-Encoding"}, + {name: "Accept-Encoding", want: true}, + {name: "Content-Encoding", want: true}, {name: "Proxy"}, {name: ""}, {name: " "}, @@ -58,3 +59,40 @@ func TestCredentialAndForbiddenHeaderPolicy(t *testing.T) { } }) } + +func TestIsForbiddenRequestRuleName(t *testing.T) { + tests := []struct { + name string + want bool + }{ + {name: "Connection", want: true}, + {name: "proxy-connection", want: true}, + {name: "KEEP-ALIVE", want: true}, + {name: "te", want: true}, + {name: "Trailer", want: true}, + {name: "transfer-encoding", want: true}, + {name: "Upgrade", want: true}, + {name: "cookie", want: true}, + {name: "Cookie2", want: true}, + {name: "Proxy-Authorization", want: true}, + {name: "pRoXy-Custom", want: true}, + {name: "Accept-Encoding", want: true}, + {name: "aCcEpT-eNcOdInG", want: true}, + {name: "Content-Encoding", want: true}, + {name: "cOnTeNt-EnCoDiNg", want: true}, + {name: "Authorization"}, + {name: "Accept"}, + {name: "Content-Type"}, + {name: "Proxy"}, + {name: ""}, + {name: " "}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := IsForbiddenRequestRuleName(test.name); got != test.want { + t.Fatalf("IsForbiddenRequestRuleName(%q) = %t, want %t", test.name, got, test.want) + } + }) + } +} diff --git a/internal/platform/httpheader/request_representation.go b/internal/platform/httpheader/request_representation.go new file mode 100644 index 00000000..ab913010 --- /dev/null +++ b/internal/platform/httpheader/request_representation.go @@ -0,0 +1,44 @@ +package httpheader + +import ( + "net/http" + "strings" +) + +var requestRepresentationMetadataNames = [...]string{ + "Content-Encoding", + "Content-Length", + "ETag", + "Digest", + "Content-MD5", + "Content-Range", + "Content-Digest", + "Repr-Digest", + "Signature", + "Signature-Input", +} + +func NormalizeUpstreamRequestRepresentation(request *http.Request, finalBodyLength int64) { + if request == nil { + return + } + if request.Header == nil { + request.Header = make(http.Header) + } else { + request.Header = request.Header.Clone() + } + for _, name := range requestRepresentationMetadataNames { + deleteField(request.Header, name) + } + deleteField(request.Header, "Accept-Encoding") + request.Header.Set("Accept-Encoding", "identity") + request.ContentLength = finalBodyLength +} + +func deleteField(headers http.Header, target string) { + for name := range headers { + if strings.EqualFold(name, target) { + delete(headers, name) + } + } +} diff --git a/internal/platform/httpheader/request_representation_test.go b/internal/platform/httpheader/request_representation_test.go new file mode 100644 index 00000000..9387a63a --- /dev/null +++ b/internal/platform/httpheader/request_representation_test.go @@ -0,0 +1,92 @@ +package httpheader + +import ( + "net/http" + "reflect" + "strings" + "testing" +) + +func TestNormalizeUpstreamRequestRepresentation(t *testing.T) { + source := http.Header{ + "Accept-Encoding": {"gzip"}, + "accept-encoding": {"br"}, + "Content-Encoding": {"zstd"}, + "content-encoding": {"deflate"}, + "Content-Length": {"999"}, + "content-length": {"998"}, + "ETag": {`"stale"`}, + "Digest": {"sha-256=stale"}, + "Content-MD5": {"stale-md5"}, + "Content-Range": {"bytes 0-1/2"}, + "Content-Digest": {"sha-256=:c3RhbGU=:"}, + "Repr-Digest": {"sha-256=:c3RhbGU=:"}, + "Signature": {"stale-signature"}, + "Signature-Input": {"stale-signature-input"}, + "Authorization": {"Bearer secret"}, + "Content-Type": {"application/json"}, + "X-Business": {"preserved"}, + } + original := source.Clone() + request := &http.Request{Header: source, ContentLength: 999} + + NormalizeUpstreamRequestRepresentation(request, 37) + + if request.ContentLength != 37 { + t.Fatalf("ContentLength = %d, want 37", request.ContentLength) + } + if values := headerFieldValues(request.Header, "Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { + t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) + } + for _, name := range []string{ + "Content-Encoding", + "Content-Length", + "ETag", + "Digest", + "Content-MD5", + "Content-Range", + "Content-Digest", + "Repr-Digest", + "Signature", + "Signature-Input", + } { + if values := headerFieldValues(request.Header, name); values != nil { + t.Errorf("%s values = %#v, want absent", name, values) + } + } + for name, want := range map[string]string{ + "Authorization": "Bearer secret", + "Content-Type": "application/json", + "X-Business": "preserved", + } { + if got := request.Header.Get(name); got != want { + t.Errorf("%s = %q, want %q", name, got, want) + } + } + if !reflect.DeepEqual(source, original) { + t.Fatalf("source header map was mutated: got %#v, want %#v", source, original) + } +} + +func TestNormalizeUpstreamRequestRepresentationInitializesNilHeader(t *testing.T) { + request := &http.Request{} + + NormalizeUpstreamRequestRepresentation(request, 0) + + if values := request.Header.Values("Accept-Encoding"); !reflect.DeepEqual(values, []string{"identity"}) { + t.Fatalf("Accept-Encoding values = %#v, want [identity]", values) + } + if request.ContentLength != 0 { + t.Fatalf("ContentLength = %d, want 0", request.ContentLength) + } +} + +func headerFieldValues(headers http.Header, target string) []string { + var values []string + for name, fieldValues := range headers { + if strings.EqualFold(name, target) { + values = append(values, fieldValues...) + } + } + return values +} diff --git a/internal/platform/utils/compression_utils.go b/internal/platform/utils/compression_utils.go deleted file mode 100644 index d330105a..00000000 --- a/internal/platform/utils/compression_utils.go +++ /dev/null @@ -1,219 +0,0 @@ -package utils - -import ( - "bytes" - "compress/gzip" - "compress/zlib" - "errors" - "fmt" - "io" - "math" - "strings" - - "github.com/andybalholm/brotli" - "github.com/klauspost/compress/zstd" - "github.com/sirupsen/logrus" -) - -// Decompressor defines the interface for different decompression algorithms -type Decompressor interface { - Decompress(data []byte, limit int64) ([]byte, error) -} - -var ErrDecompressedResponseTooLarge = errors.New("decompressed response exceeds size limit") - -// decompressorRegistry holds all registered decompressors -var decompressorRegistry = make(map[string]Decompressor) - -// init registers default decompressors -func init() { - RegisterDecompressor("gzip", &GzipDecompressor{}) - RegisterDecompressor("br", &BrotliDecompressor{}) - RegisterDecompressor("deflate", &DeflateDecompressor{}) - RegisterDecompressor("zstd", &ZstdDecompressor{}) -} - -// RegisterDecompressor allows registering new decompression algorithms -func RegisterDecompressor(encoding string, decompressor Decompressor) { - decompressorRegistry[encoding] = decompressor - logrus.Debugf("Registered decompressor for encoding: %s", encoding) -} - -// DecompressResponse automatically decompresses response data based on Content-Encoding header -func DecompressResponse(contentEncoding string, data []byte) ([]byte, error) { - return DecompressResponseLimited(contentEncoding, data, math.MaxInt64) -} - -// DecompressResponseLimited decompresses response data up to limit bytes. -func DecompressResponseLimited(contentEncoding string, data []byte, limit int64) ([]byte, error) { - encoding, err := normalizeContentEncoding(contentEncoding) - if err != nil { - return nil, err - } - if limit < 0 { - return nil, fmt.Errorf("decompressed response limit must not be negative") - } - if encoding == "" || encoding == "identity" || len(data) == 0 { - if int64(len(data)) > limit { - return nil, ErrDecompressedResponseTooLarge - } - return bytes.Clone(data), nil - } - decompressor, exists := decompressorRegistry[encoding] - if !exists { - return nil, fmt.Errorf("unsupported content encoding %q", contentEncoding) - } - return decompressor.Decompress(data, limit) -} - -func readDecompressedAtMost(reader io.Reader, limit int64) ([]byte, error) { - var limited io.Reader = reader - if limit < math.MaxInt64 { - limited = io.LimitReader(reader, limit+1) - } - decoded, err := io.ReadAll(limited) - if err != nil { - return nil, err - } - if int64(len(decoded)) > limit { - return nil, ErrDecompressedResponseTooLarge - } - return decoded, nil -} - -// CompressResponse compresses response data based on Content-Encoding header. -func CompressResponse(contentEncoding string, data []byte) ([]byte, error) { - encoding, err := normalizeContentEncoding(contentEncoding) - if err != nil { - return nil, err - } - if encoding == "" || encoding == "identity" { - return bytes.Clone(data), nil - } - - var buffer bytes.Buffer - var writer io.WriteCloser - switch encoding { - case "gzip": - writer = gzip.NewWriter(&buffer) - case "br": - writer = brotli.NewWriter(&buffer) - case "deflate": - writer = zlib.NewWriter(&buffer) - case "zstd": - encoder, createErr := zstd.NewWriter(&buffer) - if createErr != nil { - return nil, fmt.Errorf("create zstd response writer: %w", createErr) - } - writer = encoder - default: - return nil, fmt.Errorf("unsupported content encoding %q", contentEncoding) - } - - if _, err := writer.Write(data); err != nil { - _ = writer.Close() - return nil, fmt.Errorf("compress %s response: %w", encoding, err) - } - if err := writer.Close(); err != nil { - return nil, fmt.Errorf("close %s response writer: %w", encoding, err) - } - return buffer.Bytes(), nil -} - -func normalizeContentEncoding(value string) (string, error) { - normalized := strings.ToLower(strings.TrimSpace(value)) - if strings.Contains(normalized, ",") { - return "", fmt.Errorf("stacked content encoding %q is not supported", value) - } - return normalized, nil -} - -// GzipDecompressor handles gzip compression -type GzipDecompressor struct{} - -// Decompress implements Decompressor interface for gzip -func (g *GzipDecompressor) Decompress(data []byte, limit int64) ([]byte, error) { - reader, err := gzip.NewReader(bytes.NewReader(data)) - if err != nil { - return nil, fmt.Errorf("failed to create gzip reader: %w", err) - } - defer reader.Close() - - decompressed, err := readDecompressedAtMost(reader, limit) - if err != nil { - return nil, fmt.Errorf("failed to read gzip data: %w", err) - } - - return decompressed, nil -} - -// BrotliDecompressor handles brotli compression -type BrotliDecompressor struct{} - -// Decompress implements Decompressor interface for brotli -func (b *BrotliDecompressor) Decompress(data []byte, limit int64) ([]byte, error) { - reader := brotli.NewReader(bytes.NewReader(data)) - - decompressed, err := readDecompressedAtMost(reader, limit) - if err != nil { - return nil, fmt.Errorf("failed to read brotli data: %w", err) - } - - return decompressed, nil -} - -// DeflateDecompressor handles HTTP deflate's zlib-wrapped DEFLATE stream. -type DeflateDecompressor struct{} - -// Decompress implements Decompressor interface for deflate -func (d *DeflateDecompressor) Decompress(data []byte, limit int64) ([]byte, error) { - reader, err := zlib.NewReader(bytes.NewReader(data)) - if err != nil { - return nil, fmt.Errorf("failed to create deflate reader: %w", err) - } - defer reader.Close() - - decompressed, err := readDecompressedAtMost(reader, limit) - if err != nil { - return nil, fmt.Errorf("failed to read deflate data: %w", err) - } - - return decompressed, nil -} - -// ZstdDecompressor handles Zstandard compression -type ZstdDecompressor struct{} - -// Decompress implements Decompressor interface for zstd -func (z *ZstdDecompressor) Decompress(data []byte, limit int64) ([]byte, error) { - options := make([]zstd.DOption, 0, 4) - if limit < math.MaxInt64 { - memoryLimit := uint64(limit) - if memoryLimit == 0 { - memoryLimit = 1 - } - windowLimit := memoryLimit - if windowLimit < zstd.MinWindowSize { - windowLimit = zstd.MinWindowSize - } - options = append( - options, - zstd.WithDecoderConcurrency(1), - zstd.WithDecoderLowmem(true), - zstd.WithDecoderMaxMemory(memoryLimit), - zstd.WithDecoderMaxWindow(windowLimit), - ) - } - reader, err := zstd.NewReader(bytes.NewReader(data), options...) - if err != nil { - return nil, fmt.Errorf("failed to create zstd reader: %w", err) - } - defer reader.Close() - - decompressed, err := readDecompressedAtMost(reader, limit) - if err != nil { - return nil, fmt.Errorf("failed to read zstd data: %w", err) - } - - return decompressed, nil -} diff --git a/internal/platform/utils/compression_utils_test.go b/internal/platform/utils/compression_utils_test.go deleted file mode 100644 index e4fe8e6a..00000000 --- a/internal/platform/utils/compression_utils_test.go +++ /dev/null @@ -1,145 +0,0 @@ -package utils - -import ( - "bytes" - "compress/zlib" - "errors" - "testing" - - "github.com/klauspost/compress/zstd" -) - -func TestDecompressResponseLimitedRejectsZstdWindowAboveLimit(t *testing.T) { - const ( - limit = 1 << 20 - largeWindow = 8 << 20 - ) - plain := bytes.Repeat([]byte("x"), 2<<20) - var wire bytes.Buffer - encoder, err := zstd.NewWriter( - &wire, - zstd.WithWindowSize(largeWindow), - zstd.WithEncoderConcurrency(1), - ) - if err != nil { - t.Fatal(err) - } - if _, err := encoder.Write(plain); err != nil { - t.Fatal(err) - } - if err := encoder.Close(); err != nil { - t.Fatal(err) - } - var header zstd.Header - if err := header.Decode(wire.Bytes()); err != nil { - t.Fatal(err) - } - if header.WindowSize <= limit || wire.Len() >= zstd.MinWindowSize { - t.Fatalf("fixture window/wire = %d/%d", header.WindowSize, wire.Len()) - } - compatDecoded, err := DecompressResponse("zstd", wire.Bytes()) - if err != nil || !bytes.Equal(compatDecoded, plain) { - t.Fatalf("compatible decode = %d bytes, %v", len(compatDecoded), err) - } - - decoded, err := DecompressResponseLimited("zstd", wire.Bytes(), limit) - if !errors.Is(err, zstd.ErrWindowSizeExceeded) || decoded != nil { - t.Fatalf("limited decode = %q, %v", decoded, err) - } -} - -func TestDecompressResponseLimitedBoundsAllSupportedEncodings(t *testing.T) { - for _, encoding := range []string{"gzip", "br", "deflate", "zstd"} { - t.Run(encoding, func(t *testing.T) { - for _, size := range []int{1 << 20, 1<<20 + 1} { - plain := bytes.Repeat([]byte("x"), size) - wire := compressResponseWithBoundedZstdWindow(t, encoding, plain) - decoded, err := DecompressResponseLimited(encoding, wire, 1<<20) - if size == 1<<20 { - if err != nil || !bytes.Equal(decoded, plain) { - t.Fatalf("exact limit = %d bytes, %v", len(decoded), err) - } - } else if !errors.Is(err, ErrDecompressedResponseTooLarge) || decoded != nil { - t.Fatalf("overflow = %d bytes, %v", len(decoded), err) - } - } - }) - } -} - -func compressResponseWithBoundedZstdWindow(t *testing.T, encoding string, plain []byte) []byte { - t.Helper() - if encoding != "zstd" { - wire, err := CompressResponse(encoding, plain) - if err != nil { - t.Fatal(err) - } - return wire - } - encoder, err := zstd.NewWriter( - nil, - zstd.WithWindowSize(1<<20), - zstd.WithEncoderConcurrency(1), - ) - if err != nil { - t.Fatal(err) - } - defer encoder.Close() - return encoder.EncodeAll(plain, nil) -} - -func TestDecompressResponseSupportsHTTPDeflate(t *testing.T) { - want := []byte("deflate response body") - var encoded bytes.Buffer - writer := zlib.NewWriter(&encoded) - if _, err := writer.Write(want); err != nil { - t.Fatalf("encode deflate body: %v", err) - } - if err := writer.Close(); err != nil { - t.Fatalf("close deflate writer: %v", err) - } - - got, err := DecompressResponse("deflate", encoded.Bytes()) - if err != nil { - t.Fatalf("DecompressResponse() error = %v", err) - } - if !bytes.Equal(got, want) { - t.Fatalf("DecompressResponse() = %q, want %q", got, want) - } -} - -func TestCompressAndDecompressResponseRoundTrip(t *testing.T) { - original := []byte(`{"error":{"api_key":"sk-secret-value"}}`) - for _, encoding := range []string{"", "identity", "gzip", "br", "deflate", "zstd"} { - t.Run(encoding, func(t *testing.T) { - encoded, err := CompressResponse(encoding, original) - if err != nil { - t.Fatalf("CompressResponse(%q) error = %v", encoding, err) - } - decoded, err := DecompressResponse(encoding, encoded) - if err != nil { - t.Fatalf("DecompressResponse(%q) error = %v", encoding, err) - } - if !bytes.Equal(decoded, original) { - t.Fatalf("round trip = %q, want %q", decoded, original) - } - }) - } -} - -func TestResponseCompressionRejectsUnknownOrStackedEncoding(t *testing.T) { - for _, encoding := range []string{"compress", "gzip, br"} { - if _, err := DecompressResponse(encoding, []byte("body")); err == nil { - t.Fatalf("DecompressResponse(%q) error = nil", encoding) - } - if _, err := CompressResponse(encoding, []byte("body")); err == nil { - t.Fatalf("CompressResponse(%q) error = nil", encoding) - } - } -} - -func TestDecompressResponseRejectsMalformedBody(t *testing.T) { - if _, err := DecompressResponse("gzip", []byte("not gzip")); err == nil { - t.Fatal("DecompressResponse(malformed gzip) error = nil") - } -} diff --git a/internal/state/runtime_settings.go b/internal/state/runtime_settings.go index b44ea355..c5e0a530 100644 --- a/internal/state/runtime_settings.go +++ b/internal/state/runtime_settings.go @@ -318,7 +318,7 @@ func parseHeaderRules(value any) (HeaderRules, error) { return HeaderRules{}, fmt.Errorf("header_rules.set contains invalid header name %q", name) } canonicalName := textproto.CanonicalMIMEHeaderKey(name) - if httpheader.IsForbiddenRequestRuleSetName(canonicalName) { + if httpheader.IsForbiddenRequestRuleName(canonicalName) { return HeaderRules{}, fmt.Errorf( "header_rules.set cannot set forbidden header %q", canonicalName, @@ -368,6 +368,12 @@ func parseHeaderRules(value any) (HeaderRules, error) { ) } canonicalName := textproto.CanonicalMIMEHeaderKey(name) + if httpheader.IsForbiddenRequestRuleName(canonicalName) { + return HeaderRules{}, fmt.Errorf( + "header_rules.remove cannot remove forbidden header %q", + canonicalName, + ) + } identity := strings.ToLower(name) if _, duplicate := seen[identity]; duplicate { return HeaderRules{}, fmt.Errorf( diff --git a/internal/state/runtime_settings_test.go b/internal/state/runtime_settings_test.go index 5a32fe09..33d037a6 100644 --- a/internal/state/runtime_settings_test.go +++ b/internal/state/runtime_settings_test.go @@ -91,7 +91,7 @@ func TestResolveRuntimeSettingsRejectsPresentNullHeaderRules(t *testing.T) { } } -func TestParseHeaderRulesRejectsUnsafeSetNames(t *testing.T) { +func TestParseHeaderRulesRejectsUnsafeNames(t *testing.T) { names := []string{ "Connection", "proxy-connection", @@ -106,13 +106,45 @@ func TestParseHeaderRulesRejectsUnsafeSetNames(t *testing.T) { "pRoXy-Custom", } for _, name := range names { - for _, value := range []string{"ordinary", "Bearer ${API_KEY}"} { - err := ValidateRuntimeSetting(SettingHeaderRules, map[string]any{ - "set": map[string]any{name: value}, - }) - if err == nil { - t.Errorf("parseHeaderRules accepted unsafe set %q=%q", name, value) + for _, section := range []string{"set", "remove"} { + values := []string{"ordinary"} + if section == "set" { + values = append(values, "Bearer ${API_KEY}") } + for _, value := range values { + setting := map[string]any{} + if section == "set" { + setting[section] = map[string]any{name: value} + } else { + setting[section] = []any{name} + } + if err := ValidateRuntimeSetting(SettingHeaderRules, setting); err == nil { + t.Errorf("parseHeaderRules accepted unsafe %s %q", section, name) + } + } + } + } +} + +func TestParseHeaderRulesRejectsReservedContentCodingNames(t *testing.T) { + for _, section := range []string{"set", "remove"} { + for _, name := range []string{ + "Accept-Encoding", + "aCcEpT-eNcOdInG", + "Content-Encoding", + "cOnTeNt-EnCoDiNg", + } { + t.Run(section+"/"+name, func(t *testing.T) { + value := map[string]any{} + if section == "set" { + value[section] = map[string]any{name: "gzip"} + } else { + value[section] = []any{name} + } + if err := ValidateRuntimeSetting(SettingHeaderRules, value); err == nil { + t.Fatalf("parseHeaderRules accepted header_rules.%s reserved name %q", section, name) + } + }) } } } @@ -171,7 +203,6 @@ func TestParseHeaderRulesRequiresAPIKeyTemplateForCredentials(t *testing.T) { err := ValidateRuntimeSetting(SettingHeaderRules, map[string]any{ "remove": []any{ "Authorization", - "Proxy-Authorization", "Api-Key", "X-Api-Key", "X-Goog-Api-Key", @@ -183,8 +214,8 @@ func TestParseHeaderRulesRequiresAPIKeyTemplateForCredentials(t *testing.T) { err = ValidateRuntimeSetting(SettingHeaderRules, map[string]any{ "set": map[string]any{ - "Accept-Encoding": "gzip", - "X-Custom": "ordinary", + "Accept": "application/json", + "X-Custom": "ordinary", }, }) if err != nil {