refactor(gateway): 统一请求响应内容编码边界
This commit is contained in:
37 files changed
+3850
-1042
No files matched your search
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
@@ -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))
|
||||
}
|
||||
|
||||
+799
-199
File diff suppressed because it is too large.
Load diff
+99
-20
@@ -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)
|
||||
|
||||
@@ -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
@@ -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.",
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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]...)
|
||||
}
|
||||
@@ -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 ""
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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(
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in new issue
Block a user