refactor(gateway): 统一请求响应内容编码边界

This commit is contained in:
tbphp committed 2026-08-04 07:39:49 +08:00
1 parent 230b29f0c4
commit 0c33cd4a57
37 files changed
+3850 -1042

No files matched your search

+2 -1
View File
@@ -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 {
+10 -8
View File
@@ -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)
+2 -1
View File
@@ -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 {
+6 -8
View File
@@ -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)
+2 -1
View File
@@ -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 {
+6 -8
View File
@@ -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)
+2 -1
View File
@@ -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 {
+12 -13
View File
@@ -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)
}
@@ -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
}
@@ -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()
}
+95 -2
View File
@@ -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)
}
+91 -170
View File
@@ -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))
}
File diff suppressed because it is too large. Load diff
+99 -20
View File
@@ -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)
+543 -21
View File
@@ -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),
)
}
+27 -12
View File
@@ -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.",
+41
View File
@@ -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 {
+101
View File
@@ -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
}
+310
View File
@@ -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)
}
})
}
}
+221 -37
View File
@@ -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 {
+423 -66
View File
@@ -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)
+9 -1
View File
@@ -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)
+36
View File
@@ -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)
-27
View File
@@ -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,
+13 -67
View File
@@ -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)
+93
View File
@@ -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
}
}
@@ -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)
}
})
}
}
@@ -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)
}
@@ -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]...)
}
+9 -3
View File
@@ -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 ""
+39 -1
View File
@@ -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)
}
})
}
}
@@ -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)
}
}
}
@@ -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
}
@@ -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
}
@@ -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")
}
}
+7 -1
View File
@@ -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(
+41 -10
View File
@@ -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 {