diff --git a/docs/subsystems/gateway-routing.md b/docs/subsystems/gateway-routing.md index 80c42ffd..57f6ce0e 100644 --- a/docs/subsystems/gateway-routing.md +++ b/docs/subsystems/gateway-routing.md @@ -28,7 +28,7 @@ served by a plugin are in [Provider and plugin ownership](provider-plugins.md). ## Runtime path -1. `handle(proto)` reads the body and the model (`requestBody`, `requestModel`), then calls `serveAgent`. `serveAgent` runs the [middleware](gateway-middleware.md) chain, then `serve`. +1. `handle(proto)` reads the body and the model (`requestBody`, `requestModel`), then calls `serveAgent`. Every route that reads an agent's body (each API, count_tokens, Gemini's, embeddings, images, videos, System One, the Codex backend path) reads it through `readBoundedRequestBody` in [`request_bounds.go`](../../internal/gateway/request_bounds.go). It decodes a `Content-Encoding` of `gzip` or `zstd` (`decodedBody`; Codex sends zstd to a custom provider's `/v1/responses` too, #1223) and drops the header, so the provider is sent plain JSON. The size limit holds for the decoded bytes (413), and any other encoding is refused 415. `keyLimited` reads a limited key's compressed body the same way, so its reservation is the decoded size. `serveAgent` runs the [middleware](gateway-middleware.md) chain, then `serve`. 2. `serve` redacts secrets in the request (`redacted`), resolves the model to a provider or a group, and refuses what can't be served. A request whose provider and fallbacks, or whose group's models, all skip redaction (`unredactedRoute`, `provider.SkipsRedaction`: set *Send requests unmasked*, `Unredacted`, with every address on this machine or the local network) goes to them as the agent wrote it; the log and recent calls keep the masked body. count_tokens and embeddings follow the same rule. - `DecidesModel` gives a 400. - A gateway key held to some models (`keyHolds`, #882) gets a 403 for any other model. diff --git a/internal/gateway/claude_subscription.go b/internal/gateway/claude_subscription.go index 2db151ee..dbfa3153 100644 --- a/internal/gateway/claude_subscription.go +++ b/internal/gateway/claude_subscription.go @@ -2699,7 +2699,7 @@ func (b *subscriptionBridge) mcpCall(w http.ResponseWriter, r *http.Request) { http.Error(w, "unknown or expired Claude run", http.StatusNotFound) return } - body, status, err := readBoundedRequestBody(w, r, requestLimits{body: 16 << 20}, nil) + body, status, err := readBoundedRequestBody(w, r, requestLimits{body: 16 << 20}) if err != nil { http.Error(w, err.Error(), status) return diff --git a/internal/gateway/codex_backend.go b/internal/gateway/codex_backend.go index 5d90efa0..92007d88 100644 --- a/internal/gateway/codex_backend.go +++ b/internal/gateway/codex_backend.go @@ -3,7 +3,6 @@ package gateway import ( "bytes" "cmp" - "compress/gzip" "context" "encoding/base64" "encoding/json" @@ -17,7 +16,6 @@ import ( "strings" "time" - "github.com/klauspost/compress/zstd" "github.com/yetone/magpie/internal/access" "github.com/yetone/magpie/internal/catalog" "github.com/yetone/magpie/internal/codexcat" @@ -68,7 +66,7 @@ func (s *Server) codexBackend(w http.ResponseWriter, r *http.Request) { // newer: the backend serves a model only to a client that knows it provider.SawCodexClient(r.Header) rest := strings.TrimPrefix(r.URL.Path, CodexPath) - body, ok := s.readRequestBody(w, r, provider.Responses, codexReader, 0) + body, ok := s.readRequestBody(w, r, provider.Responses, 0) if !ok { return } @@ -273,29 +271,6 @@ func hasSealedAgentMessage(body []byte) bool { return false } -func codexReader(r *http.Request) (io.ReadCloser, error) { - var rd io.ReadCloser = io.NopCloser(r.Body) - switch enc := strings.ToLower(strings.TrimSpace(r.Header.Get("Content-Encoding"))); enc { - case "", "identity": - case "zstd": - d, err := zstd.NewReader(r.Body, zstd.WithDecoderMaxMemory(defaultBodyLimit)) - if err != nil { - return nil, err - } - rd = d.IOReadCloser() - case "gzip": - g, err := gzip.NewReader(r.Body) - if err != nil { - return nil, err - } - rd = g - default: - return nil, fmt.Errorf("magpie can't read a %s body", enc) - } - r.Header.Del("Content-Encoding") - return rd, nil -} - // compactSigninHeld is the calling key when its accounts (#905) hold a // native compaction off the account Codex is signed in to. Compaction // relays as it came (#876, the encrypted history in it readable only by diff --git a/internal/gateway/decide.go b/internal/gateway/decide.go index a83298e3..5dba5b88 100644 --- a/internal/gateway/decide.go +++ b/internal/gateway/decide.go @@ -400,7 +400,7 @@ const maxSystemOneBody = 16 << 20 // conversation is: the Routing view and the day's jsonl would otherwise // never see Jev, which answers no /v1/chat/completions. func (s *Server) serveSystemOne(w http.ResponseWriter, r *http.Request) { - body, ok := s.readRequestBody(w, r, provider.Chat, nil, maxSystemOneBody) + body, ok := s.readRequestBody(w, r, provider.Chat, maxSystemOneBody) if !ok { return } diff --git a/internal/gateway/key_limit.go b/internal/gateway/key_limit.go index fac71a27..1920d978 100644 --- a/internal/gateway/key_limit.go +++ b/internal/gateway/key_limit.go @@ -32,14 +32,14 @@ func keyLimited(next http.Handler) http.Handler { return } size, model := r.ContentLength, "" - if who.Limit.Cost > 0 || size < 0 { - b, err := io.ReadAll(r.Body) - r.Body.Close() + // a compressed body (#1223) is reserved at its decoded size and model + if enc := r.Header.Get("Content-Encoding"); who.Limit.Cost > 0 || size < 0 || (enc != "" && !strings.EqualFold(enc, "identity")) { + b, status, err := readBoundedRequestBody(w, r, requestLimits{}) if err != nil { - writeError(w, protoOfPath(r.URL.Path), http.StatusBadRequest, err.Error()) + writeError(w, protoOfPath(r.URL.Path), status, err.Error()) return } - r.Body = io.NopCloser(bytes.NewReader(b)) + r.Body, r.ContentLength = io.NopCloser(bytes.NewReader(b)), int64(len(b)) size, model = int64(len(b)), budget.ModelOf(b) } now := limitClock() diff --git a/internal/gateway/request_body_test.go b/internal/gateway/request_body_test.go index 80622316..c7ced26f 100644 --- a/internal/gateway/request_body_test.go +++ b/internal/gateway/request_body_test.go @@ -50,7 +50,7 @@ func TestCodexBodyReadsPastOldLimit(t *testing.T) { // Hide the reader's length, as with a chunked request. r := httptest.NewRequest("POST", CodexPath+"/responses", struct{ io.Reader }{bytes.NewReader(encodedBody(t, body, encoding))}) r.Header.Set("Content-Encoding", encoding) - got, ok := New().readRequestBody(httptest.NewRecorder(), r, provider.Responses, codexReader, 0) + got, ok := New().readRequestBody(httptest.NewRecorder(), r, provider.Responses, 0) if !ok { t.Fatal("body rejected") } diff --git a/internal/gateway/request_bounds.go b/internal/gateway/request_bounds.go index 108af3d4..125e7358 100644 --- a/internal/gateway/request_bounds.go +++ b/internal/gateway/request_bounds.go @@ -1,12 +1,16 @@ package gateway import ( + "compress/gzip" "errors" + "fmt" "io" "net" "net/http" + "strings" "time" + "github.com/klauspost/compress/zstd" "github.com/yetone/magpie/internal/provider" ) @@ -23,15 +27,15 @@ type requestLimits struct { } func (s *Server) requestBody(w http.ResponseWriter, r *http.Request, from provider.Protocol) ([]byte, bool) { - return s.readRequestBody(w, r, from, nil, 0) + return s.readRequestBody(w, r, from, 0) } -func (s *Server) readRequestBody(w http.ResponseWriter, r *http.Request, from provider.Protocol, decode func(*http.Request) (io.ReadCloser, error), limit int64) ([]byte, bool) { +func (s *Server) readRequestBody(w http.ResponseWriter, r *http.Request, from provider.Protocol, limit int64) ([]byte, bool) { limits := s.requestLimits if limit > 0 { limits.body = limit } - body, status, err := readBoundedRequestBody(w, r, limits, decode) + body, status, err := readBoundedRequestBody(w, r, limits) if err != nil { writeError(w, from, status, err.Error()) return nil, false @@ -57,8 +61,39 @@ func (b *idleBody) Read(p []byte) (int, error) { return n, err } -// Read one bounded body; callers choose their protocol's error envelope. -func readBoundedRequestBody(w http.ResponseWriter, r *http.Request, limits requestLimits, decode func(*http.Request) (io.ReadCloser, error)) ([]byte, int, error) { +// errBodyEncoding is a Content-Encoding magpie can't decode. +var errBodyEncoding = errors.New("unsupported Content-Encoding") + +// decodedBody reads r.Body through its Content-Encoding: Codex sends its +// bodies zstd, to the ChatGPT backend path and to a custom provider's +// /v1/responses alike (#1223), and other clients gzip. The header is +// dropped once decoded, so nothing after this sees or forwards it. +func decodedBody(r *http.Request) (io.ReadCloser, error) { + var rd io.ReadCloser = io.NopCloser(r.Body) + switch enc := strings.ToLower(strings.TrimSpace(r.Header.Get("Content-Encoding"))); enc { + case "", "identity": + case "zstd": + d, err := zstd.NewReader(r.Body, zstd.WithDecoderMaxMemory(defaultBodyLimit)) + if err != nil { + return nil, err + } + rd = d.IOReadCloser() + case "gzip", "x-gzip": + g, err := gzip.NewReader(r.Body) + if err != nil { + return nil, fmt.Errorf("request body is not gzip: %w", err) + } + rd = g + default: + return nil, fmt.Errorf("%w: magpie can't read a %q body; send it as identity, gzip or zstd", errBodyEncoding, enc) + } + r.Header.Del("Content-Encoding") + return rd, nil +} + +// Read one bounded body, decoded; callers choose their protocol's error +// envelope. The limit holds for the decoded bytes too. +func readBoundedRequestBody(w http.ResponseWriter, r *http.Request, limits requestLimits) ([]byte, int, error) { controller := http.NewResponseController(w) reject := func(status int, err error) ([]byte, int, error) { // Do not let net/http drain a rejected body from a stalled client. @@ -93,15 +128,14 @@ func readBoundedRequestBody(w http.ResponseWriter, r *http.Request, limits reque original := r.Body r.Body = &idleBody{original, controller, limits.readTimeout} defer func() { r.Body = original }() - rd := r.Body - if decode != nil { - var err error - rd, err = decode(r) - if err != nil { - return readFailure(err) - } - defer rd.Close() + rd, err := decodedBody(r) + if errors.Is(err, errBodyEncoding) { + return reject(http.StatusUnsupportedMediaType, err) } + if err != nil { + return readFailure(err) + } + defer rd.Close() // Limit decoded bytes as well, without allocating from Content-Length. body, err := io.ReadAll(io.LimitReader(rd, limits.body+1)) if int64(len(body)) > limits.body { diff --git a/internal/gateway/request_encoding_test.go b/internal/gateway/request_encoding_test.go new file mode 100644 index 00000000..a984cde1 --- /dev/null +++ b/internal/gateway/request_encoding_test.go @@ -0,0 +1,222 @@ +package gateway + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/yetone/magpie/internal/access" + "github.com/yetone/magpie/internal/budget" + "github.com/yetone/magpie/internal/provider" +) + +// postEncoded sends body to the running gateway at path with the given +// Content-Encoding, as Codex sends a custom provider's /v1/responses zstd +// (#1223); Go's client leaves a body it is handed as it is. +func postEncoded(t *testing.T, gw *httptest.Server, path string, body []byte, enc string, header http.Header) (int, string) { + t.Helper() + req, err := http.NewRequest("POST", gw.URL+path, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + for k, v := range header { + req.Header[k] = v + } + req.Header.Set("Content-Type", "application/json") + if enc != "" { + req.Header.Set("Content-Encoding", enc) + } + res, err := gw.Client().Do(req) + if err != nil { + t.Fatal(err) + } + defer res.Body.Close() + b, _ := io.ReadAll(res.Body) + return res.StatusCode, string(b) +} + +// Every route that reads an agent's body reads it gzip or zstd as well as +// plain, and sends the provider the plain JSON with no Content-Encoding +// (#1223: Codex's custom provider on /v1/responses got "expected a JSON +// object"). +func TestEncodedBodiesOnEveryRoute(t *testing.T) { + const payload = `{"model":"fake/m1","input":"END!","messages":[{"role":"user","content":"END!"}],"contents":[{"role":"user","parts":[{"text":"END!"}]}]}` + for _, tc := range []struct { + path string + proto provider.Protocol + count bool + }{ + {"/v1/responses", provider.Responses, false}, + {"/responses", provider.Responses, false}, + {"/v1/chat/completions", provider.Chat, false}, + {"/v1/messages", provider.Anthropic, false}, + {"/v1/messages/count_tokens", provider.Anthropic, true}, + {"/v1beta/models/fake/m1:generateContent", provider.Chat, false}, + {"/v1beta/models/fake/m1:countTokens", provider.Chat, true}, + {CodexPath + "/responses", provider.Responses, false}, + } { + for _, enc := range []string{"", "identity", "gzip", "zstd"} { + t.Run(tc.path+"/"+enc, func(t *testing.T) { + f := &fake{t: t, ctype: "application/json", reply: `{"id":"ok","input_tokens":7,"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"output":[],"content":[]}`} + if strings.HasSuffix(tc.path, ":generateContent") { + f.ctype = "text/event-stream" + f.reply = sse(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":null}]}`, `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`, `data: [DONE]`) + } + setup(t, tc.proto, f) + gw := httptest.NewServer(New().Handler()) + defer gw.Close() + control, want := postEncoded(t, gw, tc.path, []byte(payload), "", nil) + if control != 200 { + t.Fatalf("plain: %d %s", control, want) + } + f.calls, f.got, f.head = 0, nil, nil + code, got := postEncoded(t, gw, tc.path, encodedBody(t, []byte(payload), enc), enc, nil) + if code != 200 { + t.Fatalf("%s body: %d %s", enc, code, got) + } + if tc.count { + if got != want { + t.Fatalf("%s count = %s, plain = %s", enc, got, want) + } + return + } + if f.calls != 1 || !bytes.Contains(f.got, []byte("END!")) || !bytes.HasPrefix(bytes.TrimSpace(f.got), []byte("{")) { + t.Fatalf("provider got %d calls, body %q", f.calls, f.got) + } + if ce := f.head.Get("Content-Encoding"); ce != "" { + t.Fatalf("provider was told Content-Encoding %q of a decoded body", ce) + } + }) + } + } +} + +// Embeddings and rerank read the agent's body the same way. +func TestEncodedEmbeddingsBody(t *testing.T) { + for _, enc := range []string{"gzip", "zstd"} { + t.Run(enc, func(t *testing.T) { + s, l := shelved(t) + gw := httptest.NewServer(s.Handler()) + defer gw.Close() + body := []byte(`{"model":"lib/embed-1","input":["END!","two"]}`) + if code, got := postEncoded(t, gw, "/v1/embeddings", encodedBody(t, body, enc), enc, nil); code != 200 { + t.Fatalf("%s embeddings: %d %s", enc, code, got) + } + if sent := l.got("/v1/embeddings"); len(sent) != 1 || !strings.Contains(sent[0], `"END!"`) { + t.Fatalf("provider got %q", sent) + } + }) + } +} + +// The gateway's limit holds for the decoded body, so a small compressed +// body that inflates past it is refused, on every route. +func TestEncodedBodyInflatedPastLimit(t *testing.T) { + big := []byte(`{"model":"fake/m1","input":"` + strings.Repeat("x", 4096) + `"}`) + for _, path := range []string{"/v1/responses", "/v1/chat/completions", "/v1/messages", "/v1/messages/count_tokens", "/v1beta/models/fake/m1:generateContent", "/v1/embeddings", "/v1/images/generations", "/v1/videos", CodexPath + "/responses"} { + for _, enc := range []string{"gzip", "zstd"} { + t.Run(path+"/"+enc, func(t *testing.T) { + f := &fake{t: t} + setup(t, provider.Chat, f) + s := New() + s.requestLimits.body = 1024 + gw := httptest.NewServer(s.Handler()) + defer gw.Close() + sent := encodedBody(t, big, enc) + if len(sent) >= 1024 { + t.Fatalf("compressed body is %d bytes, not under the limit", len(sent)) + } + if code, got := postEncoded(t, gw, path, sent, enc, nil); code != http.StatusRequestEntityTooLarge { + t.Fatalf("inflated %s body: %d %s", enc, code, got) + } + if f.calls != 0 { + t.Fatal("an oversize body reached the provider") + } + }) + } + } +} + +// A body in an encoding magpie can't read is refused 415, saying which, +// rather than read as broken JSON. +func TestUnknownBodyEncodingRefused(t *testing.T) { + for _, path := range []string{"/v1/responses", "/v1/chat/completions", "/v1/messages", "/v1beta/models/fake/m1:generateContent", CodexPath + "/responses"} { + t.Run(path, func(t *testing.T) { + f := &fake{t: t} + setup(t, provider.Chat, f) + gw := httptest.NewServer(New().Handler()) + defer gw.Close() + code, got := postEncoded(t, gw, path, []byte(`{"model":"fake/m1"}`), "br", nil) + if code != http.StatusUnsupportedMediaType || !strings.Contains(got, `\"br\"`) || !strings.Contains(got, "zstd") { + t.Fatalf("br body: %d %s", code, got) + } + if f.calls != 0 { + t.Fatal("an unreadable body reached the provider") + } + }) + } +} + +// A gateway key's limit holds a compressed request at what it decodes to, +// not at its compressed size. +func TestKeyLimitReservesDecodedBody(t *testing.T) { + fresh(t) + budget.Forget() + now := pinLimitClock(t) + release := make(chan struct{}) + entered := make(chan struct{}, 1) + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + entered <- struct{}{} + <-release + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":300,"completion_tokens":50}}`) + })) + defer up.Close() + if err := provider.Save(provider.Provider{ID: "plan", Name: "Plan", Key: "upstream-secret", Chat: up.URL + "/v1", Models: []string{"m1"}}); err != nil { + t.Fatal(err) + } + keys, secrets := newCaller(t, "Zipped") + limit := &access.Limit{Period: "month", Tokens: 1 << 30} + setLimit(t, keys[0].ID, limit) + gw := httptest.NewServer(New().Handler()) + defer gw.Close() + body := []byte(`{"model":"plan/m1","messages":[{"role":"user","content":"` + strings.Repeat("x", 40000) + `"}]}`) + sent := encodedBody(t, body, "zstd") + var once sync.Once + free := func() { once.Do(func() { close(release) }) } + defer free() + done := make(chan string, 1) + go func() { + req, _ := http.NewRequest("POST", gw.URL+"/v1/chat/completions", bytes.NewReader(sent)) + req.Header.Set("Authorization", "Bearer "+secrets[0]) + req.Header.Set("Content-Encoding", "zstd") + res, err := gw.Client().Do(req) + if err != nil { + done <- err.Error() + return + } + b, _ := io.ReadAll(res.Body) + res.Body.Close() + done <- res.Status + " " + string(b) + }() + select { + case <-entered: + case got := <-done: + t.Fatalf("the request never reached the provider: %s", got) + case <-time.After(10 * time.Second): + t.Fatal("the request never reached the provider") + } + st := budget.Of(access.Key{ID: keys[0].ID, Limit: limit}, now) + free() + if got := <-done; !strings.HasPrefix(got, "200 ") { + t.Fatalf("zstd request: %s", got) + } + if st.InFlight != 1 || st.Reserved < int64(len(body))/4 { + t.Fatalf("held at %d tokens for a %d-byte body sent as %d bytes", st.Reserved, len(body), len(sent)) + } +} diff --git a/internal/gateway/request_idle_test.go b/internal/gateway/request_idle_test.go index d7f6a5f7..4b194603 100644 --- a/internal/gateway/request_idle_test.go +++ b/internal/gateway/request_idle_test.go @@ -23,7 +23,7 @@ func TestRequestBodyProgressRenewsIdleDeadline(t *testing.T) { s := New() s.requestLimits.readTimeout = 200 * time.Millisecond gw := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - body, ok := s.readRequestBody(w, r, provider.Responses, codexReader, 0) + body, ok := s.readRequestBody(w, r, provider.Responses, 0) if !ok { return } diff --git a/internal/gateway/sealed_agent_live_test.go b/internal/gateway/sealed_agent_live_test.go index 3aa0eb72..2643efee 100644 --- a/internal/gateway/sealed_agent_live_test.go +++ b/internal/gateway/sealed_agent_live_test.go @@ -42,7 +42,7 @@ func TestLiveCodexSealedAgentGuidance(t *testing.T) { gateway := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { isChild := false if r.Method == http.MethodPost && r.URL.Path == CodexPath+"/responses" { - body, ok := s.readRequestBody(w, r, provider.Responses, codexReader, 0) + body, ok := s.readRequestBody(w, r, provider.Responses, 0) if !ok { return }