diff --git a/internal/dialect/pipeline.go b/internal/dialect/pipeline.go index 85ff9a09a..4a18ad845 100644 --- a/internal/dialect/pipeline.go +++ b/internal/dialect/pipeline.go @@ -4,6 +4,7 @@ import ( "net/http" "strings" + "gpt-load/internal/platform/httpheader" "gpt-load/internal/state" ) @@ -26,6 +27,9 @@ func ApplyCredential( headers.Set(name, strings.ReplaceAll(value, "${API_KEY}", apiKey)) } for _, name := range rules.Remove { + if httpheader.IsForbiddenRequestRuleSetName(name) { + continue + } headers.Del(name) } } diff --git a/internal/dialect/pipeline_test.go b/internal/dialect/pipeline_test.go index 0560d474d..36023f58a 100644 --- a/internal/dialect/pipeline_test.go +++ b/internal/dialect/pipeline_test.go @@ -9,14 +9,7 @@ import ( func TestApplyCredentialUsesDialectDefaultWithEmptyRules(t *testing.T) { headers := make(http.Header) - - ApplyCredential( - NewOpenAI(http.DefaultClient), - headers, - "sk-default", - state.HeaderRules{}, - ) - + ApplyCredential(NewOpenAI(http.DefaultClient), headers, "sk-default", state.HeaderRules{}) if got := headers.Get("Authorization"); got != "Bearer sk-default" { t.Fatalf("Authorization = %q, want default Bearer credential", got) } @@ -31,14 +24,7 @@ func TestApplyCredentialExpandsSetRulesAfterDefault(t *testing.T) { }, Remove: []string{"X-Remove-Me"}, } - - ApplyCredential( - NewOpenAI(http.DefaultClient), - headers, - "sk-custom", - rules, - ) - + ApplyCredential(NewOpenAI(http.DefaultClient), headers, "sk-custom", rules) if got := headers.Get("Authorization"); got != "Token sk-custom" { t.Fatalf("Authorization = %q, want custom override", got) } @@ -50,20 +36,22 @@ func TestApplyCredentialExpandsSetRulesAfterDefault(t *testing.T) { } } +func TestApplyCredentialCannotRemoveSystemOwnedContentCodingHeaders(t *testing.T) { + headers := http.Header{"Accept-Encoding": {"identity"}} + rules := state.HeaderRules{Remove: []string{"Accept-Encoding", "Content-Encoding", "Content-Length"}} + ApplyCredential(NewOpenAI(http.DefaultClient), headers, "sk-system", rules) + if got := headers.Get("Accept-Encoding"); got != "identity" { + t.Fatalf("Accept-Encoding = %q, want identity", got) + } +} + func TestApplyCredentialRemoveWinsOverSet(t *testing.T) { headers := make(http.Header) rules := state.HeaderRules{ Set: map[string]string{"Authorization": "Token ${API_KEY}"}, Remove: []string{"Authorization"}, } - - ApplyCredential( - NewOpenAI(http.DefaultClient), - headers, - "sk-removed", - rules, - ) - + ApplyCredential(NewOpenAI(http.DefaultClient), headers, "sk-removed", rules) if got := headers.Get("Authorization"); got != "" { t.Fatalf("Authorization = %q, want final remove to win", got) } diff --git a/internal/gateway/content_coding.go b/internal/gateway/content_coding.go new file mode 100644 index 000000000..382f10908 --- /dev/null +++ b/internal/gateway/content_coding.go @@ -0,0 +1,53 @@ +package gateway + +import ( + "fmt" + "net/http" + "strconv" + + "gpt-load/internal/platform/contentcoding" +) + +func readDecodedRequestBody( + request *http.Request, + encodedLimit int64, + decodedLimit int64, +) ([]byte, http.Header, error) { + if request == nil || request.Body == nil { + return nil, nil, fmt.Errorf("%w: request body is required", contentcoding.ErrInvalidEncoding) + } + if request.ContentLength > encodedLimit && encodedLimit >= 0 { + return nil, nil, contentcoding.ErrEncodedTooLarge + } + body, err := contentcoding.ReadDecodedBody( + request.Body, + request.Header.Values("Content-Encoding"), + encodedLimit, + decodedLimit, + ) + if err != nil { + return nil, nil, err + } + headers := cloneEndToEndHeaders(request.Header) + stripRepresentationMetadata(headers) + headers.Set("Accept-Encoding", "identity") + return body, headers, nil +} + +func stripRepresentationMetadata(headers http.Header) { + if headers == nil { + return + } + for _, name := range representationMetadataHeaderNames { + deleteHeaderField(headers, name) + } +} + +func rebuildPlainBufferedResponseHeaders(headers http.Header, bodyLength int) { + stripRepresentationMetadata(headers) + headers.Set("Content-Length", strconv.Itoa(bodyLength)) +} + +func normalizePlainStreamingResponseHeaders(headers http.Header) { + stripRepresentationMetadata(headers) +} diff --git a/internal/gateway/content_coding_reason_test.go b/internal/gateway/content_coding_reason_test.go new file mode 100644 index 000000000..0ffa50f48 --- /dev/null +++ b/internal/gateway/content_coding_reason_test.go @@ -0,0 +1,27 @@ +package gateway + +import ( + "net/http" + "testing" +) + +// These public reason codes are part of the data-plane compatibility contract. +func TestContentCodingReasonsRemainStable(t *testing.T) { + tests := []struct { + name string + reason reason + status int + code string + }{ + {name: "invalid encoding", reason: reasonInvalidContentEncoding, status: http.StatusBadRequest, code: "invalid_content_encoding"}, + {name: "unsupported encoding", reason: reasonUnsupportedContentEncoding, status: http.StatusUnsupportedMediaType, code: "unsupported_content_encoding"}, + {name: "identity rejected", reason: reasonNotAcceptable, status: http.StatusNotAcceptable, code: "not_acceptable"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if test.reason.Status != test.status || test.reason.Code != test.code { + t.Fatalf("reason = %#v, want status/code %d/%q", test.reason, test.status, test.code) + } + }) + } +} diff --git a/internal/gateway/content_coding_test.go b/internal/gateway/content_coding_test.go new file mode 100644 index 000000000..bc69e820c --- /dev/null +++ b/internal/gateway/content_coding_test.go @@ -0,0 +1,174 @@ +package gateway + +import ( + "bytes" + "compress/gzip" + "compress/zlib" + "errors" + "io" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/andybalholm/brotli" + "github.com/gin-gonic/gin" + "github.com/klauspost/compress/zstd" + + "gpt-load/internal/platform/contentcoding" +) + +func TestReadDecodedRequestBodyNormalizesSupportedEncodings(t *testing.T) { + plain := []byte(`{"model":"gpt-test","messages":[]}`) + for _, encoding := range []string{"", "identity", "gzip", "br", "deflate", "zstd"} { + t.Run(encoding, func(t *testing.T) { + wire := encodeGatewayContentCodingFixture(t, encoding, plain) + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(wire)) + if encoding != "" { + request.Header.Set("Content-Encoding", encoding) + } + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Accept-Encoding", "gzip") + request.Header.Set("Digest", "sha-256=stale") + + body, headers, err := readDecodedRequestBody(request, 1<<20, 1<<20) + if err != nil || !bytes.Equal(body, plain) { + t.Fatalf("readDecodedRequestBody(%q) = %q, %v", encoding, body, err) + } + for _, name := range []string{"Content-Encoding", "Content-Length", "Digest"} { + if values := headers.Values(name); values != nil { + t.Errorf("normalized header %s survived: %#v", name, values) + } + } + if headers.Get("Accept-Encoding") != "identity" || headers.Get("Content-Type") != "application/json" { + t.Fatalf("normalized headers = %#v", headers) + } + }) + } +} + +func TestReadDecodedRequestBodyReturnsStableContentCodingErrors(t *testing.T) { + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader([]byte("body"))) + request.Header.Set("Content-Encoding", "gzip, br") + if _, _, err := readDecodedRequestBody(request, 1<<20, 1<<20); !errors.Is(err, contentcoding.ErrUnsupportedEncoding) { + t.Fatalf("stacked encoding error = %v", err) + } + + request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader([]byte("not-gzip"))) + request.Header.Set("Content-Encoding", "gzip") + if _, _, err := readDecodedRequestBody(request, 1<<20, 1<<20); !errors.Is(err, contentcoding.ErrInvalidEncoding) { + t.Fatalf("malformed encoding error = %v", err) + } + + plain := bytes.Repeat([]byte("x"), 1025) + wire := encodeGatewayContentCodingFixture(t, "gzip", plain) + request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(wire)) + request.Header.Set("Content-Encoding", "gzip") + if _, _, err := readDecodedRequestBody(request, 1<<20, 1024); !errors.Is(err, contentcoding.ErrDecodedTooLarge) { + t.Fatalf("decoded overflow error = %v", err) + } +} + +func TestPlaintextResponseWriterDecodesCompressedBufferedResponse(t *testing.T) { + gin.SetMode(gin.TestMode) + plain := []byte(`{"id":"response","model":"gpt-test"}`) + for _, encoding := range []string{"gzip", "br", "deflate", "zstd"} { + t.Run(encoding, func(t *testing.T) { + wire := encodeGatewayContentCodingFixture(t, encoding, plain) + recorder := httptest.NewRecorder() + engine := gin.New() + engine.GET("/", (&Handler{}).normalizeDownstreamContentCoding(), func(context *gin.Context) { + context.Header("Content-Type", "application/json") + context.Header("Content-Encoding", encoding) + context.Header("Content-Length", strconv.Itoa(len(wire))) + context.Header("ETag", "stale-compressed-etag") + context.Status(http.StatusOK) + _, _ = context.Writer.Write(wire) + context.Writer.Flush() + }) + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil)) + if recorder.Code != http.StatusOK || !bytes.Equal(recorder.Body.Bytes(), plain) { + t.Fatalf("response = %d %q, want %q", recorder.Code, recorder.Body.Bytes(), plain) + } + if recorder.Header().Get("Content-Encoding") != "" || recorder.Header().Get("ETag") != "" || + recorder.Header().Get("Content-Length") != strconv.Itoa(len(plain)) { + t.Fatalf("plaintext headers = %#v", recorder.Header()) + } + }) + } +} + +func TestPlaintextResponseWriterPreservesValidIdentityMetadata(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"ok":true}`) + recorder := httptest.NewRecorder() + engine := gin.New() + engine.GET("/", (&Handler{}).normalizeDownstreamContentCoding(), func(context *gin.Context) { + context.Header("Content-Encoding", "identity") + context.Header("ETag", "valid-identity-etag") + context.Status(http.StatusOK) + _, _ = context.Writer.Write(body) + context.Writer.Flush() + }) + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil)) + if !bytes.Equal(recorder.Body.Bytes(), body) || recorder.Header().Get("Content-Encoding") != "" || + recorder.Header().Get("ETag") != "valid-identity-etag" { + t.Fatalf("identity response = %q headers=%#v", recorder.Body.Bytes(), recorder.Header()) + } +} + +func TestPlaintextResponseWriterStreamsIdentitySSE(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + engine := gin.New() + engine.GET("/", (&Handler{}).normalizeDownstreamContentCoding(), func(context *gin.Context) { + context.Header("Content-Type", "text/event-stream") + context.Header("Content-Encoding", "identity") + context.Header("ETag", "stale-stream-etag") + context.Status(http.StatusOK) + _, _ = context.Writer.WriteString("data: one\n\n") + context.Writer.Flush() + _, _ = context.Writer.WriteString("data: two\n\n") + context.Writer.Flush() + }) + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil)) + if recorder.Body.String() != "data: one\n\ndata: two\n\n" { + t.Fatalf("stream body = %q", recorder.Body.String()) + } + if recorder.Header().Get("Content-Encoding") != "" || recorder.Header().Get("Content-Length") != "" || + recorder.Header().Get("ETag") != "" { + t.Fatalf("stream headers = %#v", recorder.Header()) + } +} + +func encodeGatewayContentCodingFixture(t *testing.T, encoding string, plain []byte) []byte { + t.Helper() + if encoding == "" || encoding == "identity" { + return bytes.Clone(plain) + } + 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, err := zstd.NewWriter(&buffer, zstd.WithEncoderConcurrency(1)) + if err != nil { + t.Fatal(err) + } + writer = encoder + default: + t.Fatalf("unsupported fixture encoding %q", encoding) + } + if _, err := writer.Write(plain); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + return buffer.Bytes() +} diff --git a/internal/gateway/downstream_content_coding.go b/internal/gateway/downstream_content_coding.go new file mode 100644 index 000000000..9b59869a2 --- /dev/null +++ b/internal/gateway/downstream_content_coding.go @@ -0,0 +1,179 @@ +package gateway + +import ( + "bytes" + "fmt" + "io" + "net/http" + "strconv" + "strings" + + "github.com/gin-gonic/gin" + + "gpt-load/internal/platform/contentcoding" +) + +type plaintextResponseWriter struct { + gin.ResponseWriter + header http.Header + status int + size int + wroteHeader bool + committed bool + streaming bool + body bytes.Buffer +} + +func newPlaintextResponseWriter(base gin.ResponseWriter) *plaintextResponseWriter { + header := make(http.Header) + if base != nil { + header = base.Header().Clone() + for name := range base.Header() { + base.Header().Del(name) + } + } + return &plaintextResponseWriter{ResponseWriter: base, header: header, status: http.StatusOK} +} + +func (writer *plaintextResponseWriter) Header() http.Header { + return writer.header +} + +func (writer *plaintextResponseWriter) WriteHeader(statusCode int) { + if writer.wroteHeader || writer.committed { + return + } + writer.status = statusCode + writer.wroteHeader = true +} + +func (writer *plaintextResponseWriter) WriteHeaderNow() { + if !writer.wroteHeader { + writer.WriteHeader(http.StatusOK) + } +} + +func (writer *plaintextResponseWriter) Write(value []byte) (int, error) { + if writer.committed { + written, err := writer.ResponseWriter.Write(value) + writer.size += written + return written, err + } + writer.WriteHeaderNow() + written, err := writer.body.Write(value) + writer.size += written + return written, err +} + +func (writer *plaintextResponseWriter) WriteString(value string) (int, error) { + return writer.Write([]byte(value)) +} + +func (writer *plaintextResponseWriter) Status() int { + return writer.status +} + +func (writer *plaintextResponseWriter) Size() int { + return writer.size +} + +func (writer *plaintextResponseWriter) Written() bool { + return writer.wroteHeader || writer.committed || writer.body.Len() > 0 +} + +func (writer *plaintextResponseWriter) Flush() { + _ = writer.FlushError() +} + +func (writer *plaintextResponseWriter) FlushError() error { + if writer.committed { + return http.NewResponseController(writer.ResponseWriter).Flush() + } + writer.WriteHeaderNow() + if isEventStreamContentType(writer.header.Get("Content-Type")) { + encoding, err := contentcoding.ParseContentEncoding(writer.header.Values("Content-Encoding")) + if err != nil || encoding != contentcoding.EncodingIdentity { + if err == nil { + err = fmt.Errorf("%w: compressed event stream", contentcoding.ErrUnsupportedEncoding) + } + return err + } + normalizePlainStreamingResponseHeaders(writer.header) + writer.streaming = true + return writer.commit(writer.body.Bytes()) + } + + encoding, err := contentcoding.ParseContentEncoding(writer.header.Values("Content-Encoding")) + if err != nil { + return err + } + plain, err := contentcoding.DecodeBytesLimited( + encoding, + writer.body.Bytes(), + maxNonStreamingResponseBodyBytes, + ) + if err != nil { + return err + } + if encoding == contentcoding.EncodingIdentity { + writer.header.Del("Content-Encoding") + writer.header.Del("Content-Length") + writer.header.Set("Content-Length", strconv.Itoa(len(plain))) + } else { + rebuildPlainBufferedResponseHeaders(writer.header, len(plain)) + } + return writer.commit(plain) +} + +func (writer *plaintextResponseWriter) Unwrap() http.ResponseWriter { + return writer.ResponseWriter +} + +func (writer *plaintextResponseWriter) commit(body []byte) error { + if writer.committed { + return nil + } + if writer.ResponseWriter == nil { + return fmt.Errorf("downstream response writer is required") + } + for name := range writer.ResponseWriter.Header() { + writer.ResponseWriter.Header().Del(name) + } + for name, values := range writer.header { + for _, value := range values { + writer.ResponseWriter.Header().Add(name, value) + } + } + writer.ResponseWriter.WriteHeader(writer.status) + writer.ResponseWriter.WriteHeaderNow() + if len(body) > 0 { + written, err := writer.ResponseWriter.Write(body) + if err != nil { + return err + } + if written != len(body) { + return io.ErrShortWrite + } + } + writer.committed = true + writer.body.Reset() + return http.NewResponseController(writer.ResponseWriter).Flush() +} + +func isEventStreamContentType(value string) bool { + mediaType := strings.TrimSpace(strings.SplitN(value, ";", 2)[0]) + return strings.EqualFold(mediaType, "text/event-stream") +} + +func (handler *Handler) normalizeDownstreamContentCoding() gin.HandlerFunc { + return func(context *gin.Context) { + wrapped := newPlaintextResponseWriter(context.Writer) + context.Writer = wrapped + context.Next() + if wrapped.Written() && !wrapped.committed { + if err := wrapped.FlushError(); err != nil { + _ = context.Error(err) + } + } + } +} diff --git a/internal/gateway/handler.go b/internal/gateway/handler.go index 4d493d04a..a34e8df12 100644 --- a/internal/gateway/handler.go +++ b/internal/gateway/handler.go @@ -15,6 +15,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" @@ -269,6 +270,10 @@ func (handler *Handler) Handle(ginContext *gin.Context) { handler.completeReason(ginContext, recorder, reasonAccessKeyRateLimited) return } + if !contentcoding.AcceptsIdentity(ginContext.Request.Header.Values("Accept-Encoding")) { + handler.completeReason(ginContext, recorder, reasonNotAcceptable) + return + } if selectedRoute.Kind == endpointModels { handler.writeVisibleModelList(ginContext, snapshot, accessKey, selectedRoute.Protocol) return @@ -295,24 +300,35 @@ func (handler *Handler) Handle(ginContext *gin.Context) { allowedKeyRefs[ref.ID] = ref } - body, err := readRequestBody(ginContext.Request.Body, maxRequestBodyBytes) + body, requestHeaders, err := readDecodedRequestBody( + ginContext.Request, + maxRequestBodyBytes, + maxRequestBodyBytes, + ) if err != nil { if ginContext.Request.Context().Err() != nil { recorder.completeCanceled(0) return } - if errors.Is(err, errRequestTooLarge) { + switch { + case errors.Is(err, contentcoding.ErrEncodedTooLarge), + errors.Is(err, contentcoding.ErrDecodedTooLarge), + errors.Is(err, errRequestTooLarge): handler.completeReason(ginContext, recorder, reasonRequestTooLarge) - return + case errors.Is(err, contentcoding.ErrUnsupportedEncoding): + handler.completeReason(ginContext, recorder, reasonUnsupportedContentEncoding) + case errors.Is(err, contentcoding.ErrInvalidEncoding): + handler.completeReason(ginContext, recorder, reasonInvalidContentEncoding) + default: + handler.completeReason(ginContext, recorder, reasonInvalidProtocolRequest) } - handler.completeReason(ginContext, recorder, reasonInvalidProtocolRequest) return } 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) diff --git a/internal/gateway/http_routes.go b/internal/gateway/http_routes.go index c03961643..5aedf9344 100644 --- a/internal/gateway/http_routes.go +++ b/internal/gateway/http_routes.go @@ -34,7 +34,10 @@ func (handler *Handler) HTTPModule() httproute.Module { Prepare: gin.HandlersChain{ handler.prepareDataPlaneRequest(selectedEndpoint), }, - Handlers: gin.HandlersChain{handler.Handle}, + Handlers: gin.HandlersChain{ + handler.normalizeDownstreamContentCoding(), + handler.Handle, + }, }) } diff --git a/internal/gateway/reason.go b/internal/gateway/reason.go index 7aeaa9b9b..51c37524e 100644 --- a/internal/gateway/reason.go +++ b/internal/gateway/reason.go @@ -14,18 +14,21 @@ 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."} + reasonInvalidContentEncoding = reason{Status: http.StatusBadRequest, Code: "invalid_content_encoding", Message: "Invalid request content encoding."} + reasonUnsupportedContentEncoding = reason{Status: http.StatusUnsupportedMediaType, Code: "unsupported_content_encoding", Message: "Unsupported request content encoding."} + reasonNotAcceptable = reason{Status: http.StatusNotAcceptable, Code: "not_acceptable", Message: "An identity response representation is required."} + 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{ Status: http.StatusTooManyRequests, Code: "access_key_rate_limited", Message: "Access key rate limit exceeded.", diff --git a/internal/platform/contentcoding/contentcoding.go b/internal/platform/contentcoding/contentcoding.go new file mode 100644 index 000000000..861c29e54 --- /dev/null +++ b/internal/platform/contentcoding/contentcoding.go @@ -0,0 +1,240 @@ +package contentcoding + +import ( + "bytes" + "compress/gzip" + "compress/zlib" + "errors" + "fmt" + "io" + "math" + "strconv" + "strings" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" +) + +// Encoding identifies one supported HTTP content coding. +type Encoding string + +const ( + EncodingIdentity Encoding = "identity" + EncodingGzip Encoding = "gzip" + EncodingBrotli Encoding = "br" + EncodingDeflate Encoding = "deflate" + EncodingZstd Encoding = "zstd" +) + +var ( + ErrUnsupportedEncoding = errors.New("unsupported content encoding") + ErrInvalidEncoding = errors.New("invalid content encoding") + ErrEncodedTooLarge = errors.New("encoded body exceeds size limit") + ErrDecodedTooLarge = errors.New("decoded body exceeds size limit") +) + +// ParseContentEncoding accepts an absent/identity coding or one supported +// coding. Stacked encodings and repeated fields are intentionally unsupported. +func ParseContentEncoding(values []string) (Encoding, error) { + if len(values) == 0 { + return EncodingIdentity, nil + } + if len(values) != 1 { + return "", fmt.Errorf("%w: multiple Content-Encoding fields", ErrUnsupportedEncoding) + } + value := strings.ToLower(strings.TrimSpace(values[0])) + if strings.Contains(value, ",") { + return "", fmt.Errorf("%w: stacked Content-Encoding %q", ErrUnsupportedEncoding, values[0]) + } + switch value { + case "", string(EncodingIdentity): + return EncodingIdentity, nil + case string(EncodingGzip): + return EncodingGzip, nil + case string(EncodingBrotli): + return EncodingBrotli, nil + case string(EncodingDeflate): + return EncodingDeflate, nil + case string(EncodingZstd): + return EncodingZstd, nil + default: + return "", fmt.Errorf("%w: %q", ErrUnsupportedEncoding, values[0]) + } +} + +// ReadDecodedBody bounds both the received representation and its decoded +// form. It never returns a partial body when either limit is exceeded. +func ReadDecodedBody( + reader io.Reader, + contentEncodingValues []string, + encodedLimit int64, + decodedLimit int64, +) ([]byte, error) { + if reader == nil { + return nil, fmt.Errorf("%w: request body is nil", ErrInvalidEncoding) + } + encoding, err := ParseContentEncoding(contentEncodingValues) + if err != nil { + return nil, err + } + wire, err := readAtMost(reader, encodedLimit, ErrEncodedTooLarge) + if err != nil { + return nil, err + } + return DecodeBytesLimited(encoding, wire, decodedLimit) +} + +// DecodeBytesLimited decodes one representation while bounding decoded output +// and decoder memory/window use where the codec supports it. +func DecodeBytesLimited(encoding Encoding, wire []byte, limit int64) ([]byte, error) { + if limit < 0 { + return nil, fmt.Errorf("decoded body limit must not be negative") + } + if encoding == EncodingIdentity { + if int64(len(wire)) > limit { + return nil, ErrDecodedTooLarge + } + return bytes.Clone(wire), nil + } + if len(wire) == 0 { + return nil, fmt.Errorf("%w: empty %s body", ErrInvalidEncoding, encoding) + } + + var reader io.Reader + var closeReader func() error + switch encoding { + case EncodingGzip: + decoded, err := gzip.NewReader(bytes.NewReader(wire)) + if err != nil { + return nil, errors.Join(ErrInvalidEncoding, err) + } + reader = decoded + closeReader = decoded.Close + case EncodingBrotli: + reader = brotli.NewReader(bytes.NewReader(wire)) + case EncodingDeflate: + decoded, err := zlib.NewReader(bytes.NewReader(wire)) + if err != nil { + return nil, errors.Join(ErrInvalidEncoding, err) + } + reader = decoded + closeReader = decoded.Close + case EncodingZstd: + options := []zstd.DOption{ + zstd.WithDecoderConcurrency(1), + zstd.WithDecoderLowmem(true), + } + if limit < math.MaxInt64 { + memoryLimit := uint64(limit) + const minimumDecoderBudget = uint64(1 << 20) + if memoryLimit < minimumDecoderBudget { + memoryLimit = minimumDecoderBudget + } + options = append( + options, + zstd.WithDecoderMaxMemory(memoryLimit), + zstd.WithDecoderMaxWindow(memoryLimit), + ) + } + decoded, err := zstd.NewReader(bytes.NewReader(wire), options...) + if err != nil { + return nil, errors.Join(ErrInvalidEncoding, err) + } + reader = decoded + closeReader = func() error { + decoded.Close() + return nil + } + default: + return nil, fmt.Errorf("%w: %q", ErrUnsupportedEncoding, encoding) + } + if closeReader != nil { + defer func() { _ = closeReader() }() + } + decoded, err := readAtMost(reader, limit, ErrDecodedTooLarge) + if err != nil { + if errors.Is(err, ErrDecodedTooLarge) { + return nil, err + } + return nil, errors.Join(ErrInvalidEncoding, err) + } + return decoded, nil +} + +func readAtMost(reader io.Reader, limit int64, overflow error) ([]byte, error) { + if limit < 0 { + return nil, fmt.Errorf("body limit must not be negative") + } + var 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, overflow + } + return body, nil +} + +// AcceptsIdentity reports whether an identity response is acceptable. Invalid +// preference syntax is compatibility-safe: only a valid explicit rejection +// can reject a plaintext response. +func AcceptsIdentity(values []string) bool { + if len(values) == 0 { + return true + } + explicitSeen := false + explicitQ := -1.0 + wildcardSeen := false + wildcardQ := -1.0 + for _, field := range values { + for _, rawItem := range strings.Split(field, ",") { + item := strings.TrimSpace(rawItem) + if item == "" { + continue + } + parts := strings.Split(item, ";") + name := strings.ToLower(strings.TrimSpace(parts[0])) + q := 1.0 + valid := true + for _, rawParam := range parts[1:] { + param := strings.TrimSpace(rawParam) + key, value, found := strings.Cut(param, "=") + if !found || !strings.EqualFold(strings.TrimSpace(key), "q") { + continue + } + parsed, err := strconv.ParseFloat(strings.TrimSpace(value), 64) + if err != nil || parsed < 0 || parsed > 1 { + valid = false + break + } + q = parsed + } + if !valid { + if name == string(EncodingIdentity) || name == "*" { + return true + } + continue + } + switch name { + case string(EncodingIdentity): + explicitSeen = true + if q > explicitQ { + explicitQ = q + } + case "*": + wildcardSeen = true + if q > wildcardQ { + wildcardQ = q + } + } + } + } + if explicitSeen { + return explicitQ > 0 + } + return !wildcardSeen || wildcardQ > 0 +} diff --git a/internal/platform/contentcoding/contentcoding_test.go b/internal/platform/contentcoding/contentcoding_test.go new file mode 100644 index 000000000..60938d1b6 --- /dev/null +++ b/internal/platform/contentcoding/contentcoding_test.go @@ -0,0 +1,140 @@ +package contentcoding + +import ( + "bytes" + "compress/gzip" + "compress/zlib" + "errors" + "io" + "strings" + "testing" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" +) + +func TestParseContentEncodingAcceptsSupportedSingleValues(t *testing.T) { + tests := []struct { + name string + values []string + want Encoding + }{ + {name: "missing", want: EncodingIdentity}, + {name: "empty", values: []string{""}, want: EncodingIdentity}, + {name: "identity", values: []string{" Identity "}, want: EncodingIdentity}, + {name: "gzip", values: []string{"GZip"}, want: EncodingGzip}, + {name: "brotli", values: []string{"br"}, want: EncodingBrotli}, + {name: "deflate", values: []string{"deflate"}, want: EncodingDeflate}, + {name: "zstd", values: []string{"zstd"}, want: EncodingZstd}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := ParseContentEncoding(test.values) + if err != nil || got != test.want { + t.Fatalf("ParseContentEncoding(%#v) = %q, %v; want %q", test.values, got, err, test.want) + } + }) + } +} + +func TestParseContentEncodingRejectsUnsupportedOrStackedValues(t *testing.T) { + for _, values := range [][]string{{"compress"}, {"gzip, br"}, {"gzip", "br"}} { + if _, err := ParseContentEncoding(values); !errors.Is(err, ErrUnsupportedEncoding) { + t.Errorf("ParseContentEncoding(%#v) error = %v, want ErrUnsupportedEncoding", values, err) + } + } +} + +func TestReadDecodedBodySupportsAllSingleEncodings(t *testing.T) { + plain := []byte(`{"model":"test-model","stream":false}`) + for _, encoding := range []Encoding{EncodingIdentity, EncodingGzip, EncodingBrotli, EncodingDeflate, EncodingZstd} { + t.Run(string(encoding), func(t *testing.T) { + wire := encodeContentCodingFixture(t, encoding, plain) + values := []string(nil) + if encoding != EncodingIdentity { + values = []string{string(encoding)} + } + got, err := ReadDecodedBody(bytes.NewReader(wire), values, int64(len(wire)), int64(len(plain))) + if err != nil || !bytes.Equal(got, plain) { + t.Fatalf("ReadDecodedBody(%q) = %q, %v; want %q", encoding, got, err, plain) + } + }) + } +} + +func TestReadDecodedBodyRejectsMalformedEncodedData(t *testing.T) { + for _, encoding := range []Encoding{EncodingGzip, EncodingBrotli, EncodingDeflate, EncodingZstd} { + t.Run(string(encoding), func(t *testing.T) { + if _, err := ReadDecodedBody(strings.NewReader("not-valid-encoded-data"), []string{string(encoding)}, 1<<20, 1<<20); !errors.Is(err, ErrInvalidEncoding) { + t.Fatalf("ReadDecodedBody(%q) error = %v, want ErrInvalidEncoding", encoding, err) + } + }) + } +} + +func TestReadDecodedBodyEnforcesEncodedAndDecodedLimits(t *testing.T) { + if body, err := ReadDecodedBody(strings.NewReader("12345"), nil, 4, 16); !errors.Is(err, ErrEncodedTooLarge) || body != nil { + t.Fatalf("encoded overflow = %q, %v", body, err) + } + plain := bytes.Repeat([]byte("x"), 1<<20+1) + wire := encodeContentCodingFixture(t, EncodingGzip, plain) + if body, err := ReadDecodedBody(bytes.NewReader(wire), []string{"gzip"}, int64(len(wire)), 1<<20); !errors.Is(err, ErrDecodedTooLarge) || body != nil { + t.Fatalf("decoded overflow = %d bytes, %v", len(body), err) + } +} + +func TestAcceptsIdentityHonorsExplicitRejection(t *testing.T) { + tests := []struct { + name string + values []string + want bool + }{ + {name: "missing", want: true}, + {name: "empty", values: []string{""}, want: true}, + {name: "ordinary compression preference", values: []string{"gzip, br"}, want: true}, + {name: "identity positive", values: []string{"identity;q=0.5, *;q=0"}, want: true}, + {name: "identity rejected", values: []string{"gzip, identity;q=0"}}, + {name: "wildcard rejected", values: []string{"gzip, *;q=0"}}, + {name: "malformed q is compatibility safe", values: []string{"identity;q=broken"}, want: true}, + {name: "malformed identity defeats wildcard rejection", values: []string{"identity;q=broken, *;q=0"}, want: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := AcceptsIdentity(test.values); got != test.want { + t.Fatalf("AcceptsIdentity(%#v) = %t, want %t", test.values, got, test.want) + } + }) + } +} + +func encodeContentCodingFixture(t *testing.T, encoding Encoding, plain []byte) []byte { + t.Helper() + if encoding == EncodingIdentity { + return bytes.Clone(plain) + } + var buffer bytes.Buffer + var writer io.WriteCloser + switch encoding { + case EncodingGzip: + writer = gzip.NewWriter(&buffer) + case EncodingBrotli: + writer = brotli.NewWriter(&buffer) + case EncodingDeflate: + writer = zlib.NewWriter(&buffer) + case EncodingZstd: + encoder, err := zstd.NewWriter(&buffer, zstd.WithEncoderConcurrency(1)) + if err != nil { + t.Fatal(err) + } + writer = encoder + default: + t.Fatalf("unsupported fixture encoding %q", encoding) + } + if _, err := writer.Write(plain); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + return buffer.Bytes() +} diff --git a/internal/platform/httpheader/policy.go b/internal/platform/httpheader/policy.go index 33401cc8c..6b08558b0 100644 --- a/internal/platform/httpheader/policy.go +++ b/internal/platform/httpheader/policy.go @@ -23,6 +23,9 @@ var forbiddenRequestRuleSetNames = map[string]struct{}{ "upgrade": {}, "cookie": {}, "cookie2": {}, + "accept-encoding": {}, + "content-encoding": {}, + "content-length": {}, } func IsCredentialName(name string) bool { diff --git a/internal/platform/httpheader/policy_test.go b/internal/platform/httpheader/policy_test.go index a2f80c0fc..44c0a9604 100644 --- a/internal/platform/httpheader/policy_test.go +++ b/internal/platform/httpheader/policy_test.go @@ -40,8 +40,10 @@ func TestCredentialAndForbiddenHeaderPolicy(t *testing.T) { {name: "Cookie2", want: true}, {name: "Proxy-Authorization", want: true}, {name: "pRoXy-Custom", want: true}, + {name: "Accept-Encoding", want: true}, + {name: "Content-Encoding", want: true}, + {name: "Content-Length", want: true}, {name: "Authorization"}, - {name: "Accept-Encoding"}, {name: "Proxy"}, {name: ""}, {name: " "},