mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-13 04:58:50 +00:00
fix(passthrough): fail closed on a model peek that never closes the object
Close the residual gap in the duplicate-"model" fix (security review
finding 2, round 3, coordinator ruling): round 2 made a duplicate
top-level "model" key WITHIN the 64KiB peek window fail closed, but a
client can still pad the body so a second "model" sits just PAST
maxModelPeekBytes -- `{"model":"cheap", <64KiB padding>,
"model":"expensive"}`. The peek reports "cheap" (the only key it saw),
the body is forwarded byte-for-byte, and the upstream (last-wins on
every mainstream JSON parser) executes "expensive".
scanTopLevelModel now also reports whether it actually confirmed the
top-level object's closing '}' within the scanned window (closed).
peekPassthroughModel combines that with its own knowledge of whether the
read hit the byte cap (more of the body remains unread) to report
truncated: true only when both hold, meaning a later duplicate cannot be
ruled out. handlePassthrough denies (403, a message distinct from "model
could not be determined") whenever truncated is true, checked before
even looking at the found model -- a found model under truncation can
never be trusted.
Scoped exactly as ruled: only affects a group with a non-empty Models
list (grp.hasModelRestriction) whose body is large enough to hit
maxModelPeekBytes -- an unrestricted group never calls
peekPassthroughModel at all, unchanged from round 2 and pinned by a new
test. No live group is model-restricted today. The unified /v1/* route
remains the supported path for a legitimately large restricted-group
body.
Tests: the coordinator's own exploit body end-to-end (403, upstream
never called); the general non-duplicate case (any restricted-group body
that never closes within the window denies, not just a provably
duplicate one); the unrestricted-group-unaffected case explicitly, using
the identical exploit body; and closed's own truncation-tracking table
against scanTopLevelModel directly. One pre-existing test
(TestHandlePassthrough_BodyIntegrity_ByteIdentical_LargeBody) used a
restricted group incidentally while testing an unrelated property (byte-
for-byte forwarding) -- switched to an unrestricted group so it tests
only what its name says.
This commit is contained in:
+107
-22
@@ -390,34 +390,59 @@ func extractPassthroughUsage(typeName, providerName string, body []byte) (usage,
|
||||
// Unmarshal requires a complete, valid document — it would reject a
|
||||
// legitimately-truncated-but-found "model" as if it were malformed JSON.
|
||||
//
|
||||
// truncated is the round-3 fix (coordinator ruling, 2026-08-22): it is
|
||||
// true exactly when BOTH (a) the read hit the maxModelPeekBytes cap
|
||||
// itself — io.ReadFull returned n == len(head) with no error, meaning
|
||||
// unread bytes genuinely remain on r.Body beyond what was scanned — AND
|
||||
// (b) scanTopLevelModel could not confirm the top-level JSON object
|
||||
// actually closed within that window. Only that combination means a
|
||||
// LATER "model" key could exist past the peek and this function cannot
|
||||
// rule it out; a body that is merely shorter than the window and
|
||||
// happens to be malformed (never closes, but nothing more of it exists
|
||||
// either) is NOT truncated — there is no hidden byte range to distrust,
|
||||
// and returning found=false there already denies via the caller's
|
||||
// existing "model could not be determined" branch. See handlePassthrough's
|
||||
// own doc comment for what a caller does with truncated=true: DENY,
|
||||
// fail-closed, regardless of what model/hasModel came back, because a
|
||||
// found model under truncation cannot be trusted — RFC 8259 permits
|
||||
// duplicate keys and every mainstream JSON parser an upstream might use
|
||||
// resolves them last-wins, so a second "model" sitting just past
|
||||
// maxModelPeekBytes would authorize against the first and execute the
|
||||
// second: `{"model":"cheap", <64KiB of padding>, "model":"expensive"}`.
|
||||
//
|
||||
// err is non-nil only for a genuine body-read failure (client disconnect,
|
||||
// deadline) distinct from merely reaching the end of a body shorter than
|
||||
// the peek window — the caller maps err to the same 400 runUnified's own
|
||||
// body-read failure uses (routes_unified.go).
|
||||
func peekPassthroughModel(r *http.Request) (model string, hasModel bool, err error) {
|
||||
func peekPassthroughModel(r *http.Request) (model string, hasModel, truncated bool, err error) {
|
||||
if r.Body == nil || isPassthroughBinaryContentType(r.Header.Get("Content-Type")) {
|
||||
return "", false, nil
|
||||
return "", false, false, nil
|
||||
}
|
||||
|
||||
head := make([]byte, maxModelPeekBytes)
|
||||
n, readErr := io.ReadFull(r.Body, head)
|
||||
capped := false
|
||||
switch readErr { //nolint:errorlint // io.ReadFull returns these sentinels bare, never wrapped — see its own doc comment
|
||||
case nil:
|
||||
// Read exactly maxModelPeekBytes: the body has at least that much
|
||||
// left unread on r.Body, restored below via io.MultiReader.
|
||||
// left unread on r.Body, restored below via io.MultiReader — a
|
||||
// LATER "model" key beyond this window cannot be ruled out unless
|
||||
// scanTopLevelModel confirms the object closed within it.
|
||||
capped = true
|
||||
case io.EOF, io.ErrUnexpectedEOF:
|
||||
// Body was shorter than the peek window (io.EOF: empty; io.
|
||||
// ErrUnexpectedEOF: 0 < n < len(head)) — head[:n] IS the whole
|
||||
// body; r.Body is now exhausted, so appending it below is a
|
||||
// harmless immediate EOF.
|
||||
// harmless immediate EOF. No bytes exist beyond what was scanned,
|
||||
// so "truncated" (in the hidden-duplicate sense) cannot apply.
|
||||
default:
|
||||
return "", false, readErr
|
||||
return "", false, false, readErr
|
||||
}
|
||||
head = head[:n]
|
||||
r.Body = io.NopCloser(io.MultiReader(bytes.NewReader(head), r.Body))
|
||||
|
||||
model, hasModel = scanTopLevelModel(head)
|
||||
return model, hasModel, nil
|
||||
model, hasModel, closed := scanTopLevelModel(head)
|
||||
return model, hasModel, capped && !closed, nil
|
||||
}
|
||||
|
||||
// scanTopLevelModel walks head — a bounded, possibly-truncated PREFIX of
|
||||
@@ -432,28 +457,45 @@ func peekPassthroughModel(r *http.Request) (model string, hasModel bool, err err
|
||||
// not an error — for anything else: head is not a JSON object at all,
|
||||
// "model" never appears among its top-level keys before head runs out,
|
||||
// or "model" is present but its value is not a string.
|
||||
func scanTopLevelModel(head []byte) (model string, found bool) {
|
||||
//
|
||||
// closed reports whether this walk actually confirmed the top-level
|
||||
// object's closing '}' within head — round-3 fix (coordinator ruling,
|
||||
// 2026-08-22): every early-return branch below (a token error, a
|
||||
// non-string key, a duplicate "model", a nested value that never
|
||||
// finishes) sets closed=false, and the ONLY way to reach closed=true is
|
||||
// to fall all the way through the main loop and then successfully
|
||||
// consume a literal '}' token afterward. dec.More() alone cannot make
|
||||
// this distinction — its own doc comment says it "reports whether there
|
||||
// is another element", but internally it treats "the next byte is the
|
||||
// closing delimiter" and "the underlying reader ran out of bytes mid-
|
||||
// object" identically (both make its own peek fail and it just returns
|
||||
// false either way) — so this function does its own explicit
|
||||
// closing-token check rather than trusting a loop-exited-via-More()
|
||||
// alone. peekPassthroughModel combines closed with its own knowledge of
|
||||
// whether the read actually hit the byte cap (there is more of r.Body
|
||||
// left unread) to decide whether a "found" model can be trusted.
|
||||
func scanTopLevelModel(head []byte) (model string, found, closed bool) {
|
||||
seenModel := false
|
||||
dec := json.NewDecoder(bytes.NewReader(head))
|
||||
tok, err := dec.Token()
|
||||
if err != nil {
|
||||
return "", false
|
||||
return "", false, false
|
||||
}
|
||||
if d, ok := tok.(json.Delim); !ok || d != '{' {
|
||||
return "", false
|
||||
return "", false, false
|
||||
}
|
||||
for dec.More() {
|
||||
keyTok, err := dec.Token()
|
||||
if err != nil {
|
||||
return model, found
|
||||
keyTok, keyErr := dec.Token()
|
||||
if keyErr != nil {
|
||||
return model, found, false
|
||||
}
|
||||
key, ok := keyTok.(string)
|
||||
if !ok {
|
||||
return model, found // malformed: expected an object key
|
||||
return model, found, false // malformed: expected an object key
|
||||
}
|
||||
valTok, err := dec.Token()
|
||||
if err != nil {
|
||||
return model, found
|
||||
valTok, valErr := dec.Token()
|
||||
if valErr != nil {
|
||||
return model, found, false
|
||||
}
|
||||
if key == "model" {
|
||||
// A DUPLICATE top-level "model" fails closed, matching
|
||||
@@ -469,7 +511,7 @@ func scanTopLevelModel(head []byte) (model string, found bool) {
|
||||
// resolvable here without rewriting the client's bytes, so the
|
||||
// only safe answer is to refuse to answer.
|
||||
if seenModel {
|
||||
return "", false
|
||||
return "", false, false
|
||||
}
|
||||
seenModel = true
|
||||
if s, isStr := valTok.(string); isStr && s != "" {
|
||||
@@ -479,14 +521,27 @@ func scanTopLevelModel(head []byte) (model string, found bool) {
|
||||
continue
|
||||
}
|
||||
if d, ok := valTok.(json.Delim); ok && (d == '{' || d == '[') {
|
||||
if err := skipJSONValue(dec); err != nil {
|
||||
return model, found
|
||||
if skipErr := skipJSONValue(dec); skipErr != nil {
|
||||
return model, found, false
|
||||
}
|
||||
}
|
||||
// Otherwise valTok was a scalar (string/float64/bool/nil) other
|
||||
// than "model" — nothing further to do; loop to the next key.
|
||||
}
|
||||
return model, found
|
||||
// dec.More() returned false above: either the object legitimately
|
||||
// closed, or the decoder simply ran out of bytes searching for the
|
||||
// next token — see this function's own doc comment for why More()
|
||||
// cannot tell those apart. Explicitly consume the closing '}' to find
|
||||
// out which; only a real, present '}' confirms the object closed
|
||||
// within head.
|
||||
closeTok, err := dec.Token()
|
||||
if err != nil {
|
||||
return model, found, false
|
||||
}
|
||||
if d, ok := closeTok.(json.Delim); !ok || d != '}' {
|
||||
return model, found, false
|
||||
}
|
||||
return model, found, true
|
||||
}
|
||||
|
||||
// skipJSONValue consumes the remainder of a nested JSON array/object
|
||||
@@ -587,6 +642,25 @@ func (g *Gateway) allowsPassthroughModel(grp *group, providerName, model string)
|
||||
// the Gemini-passthrough interaction this creates (its model id lives in
|
||||
// the URL, never the body) and the operator-facing contract.
|
||||
//
|
||||
// PADDED-DUPLICATE MODEL (round 3, 2026-08-22, coordinator ruling — closes
|
||||
// the round-2 fix's own residual gap): round 2 made a duplicate top-level
|
||||
// "model" WITHIN the peek window fail closed (scanTopLevelModel), but a
|
||||
// client can still pad the body so a second "model" key sits just PAST
|
||||
// maxModelPeekBytes: `{"model":"cheap", <64KiB padding>,
|
||||
// "model":"expensive"}` — the peek reports "cheap" (the only one it saw),
|
||||
// the body is forwarded byte-for-byte, and the upstream (last-wins on
|
||||
// every mainstream JSON parser) executes "expensive". Fixed the same way:
|
||||
// peekPassthroughModel now also reports truncated=true whenever the read
|
||||
// hit the byte cap AND the scanned window never confirmed the top-level
|
||||
// object actually closed — meaning a later duplicate cannot be ruled
|
||||
// out — and this function denies (403, a message distinct from "model
|
||||
// could not be determined") BEFORE even looking at hasModel/model in that
|
||||
// case. This only ever applies to a restricted group whose body was
|
||||
// large enough to hit maxModelPeekBytes in the first place; an
|
||||
// unrestricted group never calls peekPassthroughModel at all (unchanged
|
||||
// from round 2), and the unified /v1/* route remains the supported path
|
||||
// for a legitimately large restricted-group request.
|
||||
//
|
||||
// PATH ALLOWLIST (same review): GroupConfig.PassthroughPaths, when
|
||||
// non-empty, additionally restricts which rest path this group's
|
||||
// passthrough requests may address (grp.allowsPassthroughPath). Empty
|
||||
@@ -608,11 +682,22 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
||||
}
|
||||
|
||||
if grp.hasModelRestriction() {
|
||||
model, hasModel, err := peekPassthroughModel(r)
|
||||
model, hasModel, truncated, err := peekPassthroughModel(r)
|
||||
if err != nil {
|
||||
writeOAIError(w, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
|
||||
return
|
||||
}
|
||||
// Checked BEFORE hasModel (round-3 coordinator ruling, 2026-08-22):
|
||||
// a truncated peek means a "model" found within the window cannot
|
||||
// be trusted — a duplicate top-level "model" beyond
|
||||
// maxModelPeekBytes may still exist and would win upstream
|
||||
// (last-wins parsing) — so this must fail closed regardless of
|
||||
// whatever hasModel/model peekPassthroughModel also returned. See
|
||||
// peekPassthroughModel's own doc comment for the full mechanism.
|
||||
if truncated {
|
||||
writeOAIError(w, http.StatusForbidden, "invalid_request_error", "request body exceeds the model-enforcement window; a top-level model field beyond it cannot be safely authorized")
|
||||
return
|
||||
}
|
||||
if !hasModel {
|
||||
writeOAIError(w, http.StatusForbidden, "invalid_request_error", "model could not be determined")
|
||||
return
|
||||
|
||||
+190
-18
@@ -180,7 +180,7 @@ func TestScanTopLevelModel(t *testing.T) {
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
gotModel, gotFound := scanTopLevelModel([]byte(tt.head))
|
||||
gotModel, gotFound, _ := scanTopLevelModel([]byte(tt.head))
|
||||
if gotFound != tt.wantFound {
|
||||
t.Errorf("found = %v, want %v", gotFound, tt.wantFound)
|
||||
}
|
||||
@@ -1420,7 +1420,17 @@ func TestHandlePassthrough_BodyIntegrity_ByteIdentical_LargeBody(t *testing.T) {
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {Models: []string{"gpt-allowed"}}}
|
||||
// Deliberately UNRESTRICTED (no Models list): this test's own purpose is
|
||||
// byte-for-byte forwarding integrity for a large body, independent of
|
||||
// model enforcement — a RESTRICTED group's own behavior for a body
|
||||
// this large is covered separately (round-3 fix, security review
|
||||
// 2026-08-22: TestHandlePassthrough_PaddedDuplicateModelBeyondPeekWindow_Returns403
|
||||
// and TestHandlePassthrough_RestrictedGroup_LargeBodyNeverClosesWindow_Returns403
|
||||
// below), and now deliberately DENIES a body this large that never
|
||||
// closes within the peek window — the opposite of what this test
|
||||
// checks. Using an unrestricted group here keeps this test's own
|
||||
// single concern (byte integrity) unaffected by that unrelated fix.
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
@@ -1719,27 +1729,189 @@ func TestHandlePassthrough_ZeroConfig_NewFieldsUnset_BehavesIdenticallyToBefore(
|
||||
// instead, matching extractMultipartModel's errDuplicateModelField rule.
|
||||
func TestScanTopLevelModel_DuplicateModelKeyFailsClosed(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
body string
|
||||
wantModel string
|
||||
wantFound bool
|
||||
name string
|
||||
body string
|
||||
wantModel string
|
||||
wantFound bool
|
||||
wantClosed bool
|
||||
}{
|
||||
{"single model", `{"model":"allowed","messages":[]}`, "allowed", true},
|
||||
{"model after other keys", `{"stream":true,"model":"allowed"}`, "allowed", true},
|
||||
{"model found then body truncated", `{"model":"allowed","messages":[{"role":"user"`, "allowed", true},
|
||||
{"duplicate model", `{"model":"allowed","model":"expensive","messages":[]}`, "", false},
|
||||
{"duplicate across a nested object", `{"model":"allowed","opts":{"a":1},"model":"expensive"}`, "", false},
|
||||
{"duplicate with identical values", `{"model":"same","model":"same"}`, "", false},
|
||||
{"no model", `{"messages":[]}`, "", false},
|
||||
{"nested model only", `{"opts":{"model":"x"}}`, "", false},
|
||||
{"single model", `{"model":"allowed","messages":[]}`, "allowed", true, true},
|
||||
{"model after other keys", `{"stream":true,"model":"allowed"}`, "allowed", true, true},
|
||||
{"model found then body truncated", `{"model":"allowed","messages":[{"role":"user"`, "allowed", true, false},
|
||||
{"duplicate model", `{"model":"allowed","model":"expensive","messages":[]}`, "", false, false},
|
||||
{"duplicate across a nested object", `{"model":"allowed","opts":{"a":1},"model":"expensive"}`, "", false, false},
|
||||
{"duplicate with identical values", `{"model":"same","model":"same"}`, "", false, false},
|
||||
{"no model", `{"messages":[]}`, "", false, true},
|
||||
{"nested model only", `{"opts":{"model":"x"}}`, "", false, true},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
gotModel, gotFound := scanTopLevelModel([]byte(tc.body))
|
||||
if gotModel != tc.wantModel || gotFound != tc.wantFound {
|
||||
t.Fatalf("scanTopLevelModel(%s) = (%q, %v), want (%q, %v)",
|
||||
tc.body, gotModel, gotFound, tc.wantModel, tc.wantFound)
|
||||
gotModel, gotFound, gotClosed := scanTopLevelModel([]byte(tc.body))
|
||||
if gotModel != tc.wantModel || gotFound != tc.wantFound || gotClosed != tc.wantClosed {
|
||||
t.Fatalf("scanTopLevelModel(%s) = (%q, %v, %v), want (%q, %v, %v)",
|
||||
tc.body, gotModel, gotFound, gotClosed, tc.wantModel, tc.wantFound, tc.wantClosed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestScanTopLevelModel_ClosedTracksWindowTruncation is the round-3
|
||||
// regression test (coordinator ruling, 2026-08-22): closed must be false
|
||||
// whenever the scanned window ends before the top-level object's closing
|
||||
// '}' is actually consumed — even when a valid "model" was already found —
|
||||
// and true whenever it genuinely was.
|
||||
func TestScanTopLevelModel_ClosedTracksWindowTruncation(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
body string
|
||||
wantClosed bool
|
||||
}{
|
||||
{"cut mid-string value of a later key", `{"model":"cheap","padding":"AAAA`, false},
|
||||
{"cut exactly after a complete key:value pair, no closing brace", `{"model":"cheap","padding":"x"`, false},
|
||||
{"cut inside a nested array", `{"model":"cheap","messages":[{"role":"user"`, false},
|
||||
{"complete, closes properly", `{"model":"cheap","messages":[]}`, true},
|
||||
{"complete with trailing whitespace inside cap", `{"model":"cheap"} `, true},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, _, gotClosed := scanTopLevelModel([]byte(tc.body))
|
||||
if gotClosed != tc.wantClosed {
|
||||
t.Errorf("scanTopLevelModel(%s) closed = %v, want %v", tc.body, gotClosed, tc.wantClosed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// --- security review finding 2, round 3, 2026-08-22: padded-duplicate
|
||||
// model beyond the peek window ---
|
||||
|
||||
// TestHandlePassthrough_PaddedDuplicateModelBeyondPeekWindow_Returns403
|
||||
// is the coordinator's own demonstrated exploit: a restricted group's
|
||||
// glob authorizes "cheap" but not "expensive"; the body's SECOND "model"
|
||||
// key sits past maxModelPeekBytes, hidden behind padding the peek never
|
||||
// reaches. Round 2 alone (fail-closed only on a duplicate WITHIN the
|
||||
// window) let this straight through, reporting "cheap" and forwarding
|
||||
// the whole body — where every mainstream upstream JSON parser resolves
|
||||
// duplicate keys last-wins and would have executed "expensive". This
|
||||
// must now be denied with 403 before the upstream is ever called.
|
||||
func TestHandlePassthrough_PaddedDuplicateModelBeyondPeekWindow_Returns403(t *testing.T) {
|
||||
var upstreamCalled bool
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCalled = true
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {Models: []string{"cheap"}}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
padding := strings.Repeat("A", maxModelPeekBytes) // pushes the second "model" key past the peek window
|
||||
body := `{"model":"cheap","padding":"` + padding + `","model":"expensive"}`
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want 403 (a duplicate model beyond the peek window must fail closed), body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if upstreamCalled {
|
||||
t.Error("upstream must never be called: the padded-duplicate exploit must be denied before forwarding")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandlePassthrough_RestrictedGroup_LargeBodyNeverClosesWindow_Returns403
|
||||
// covers the general (non-duplicate) case the coordinator ruling actually
|
||||
// implements: ANY restricted-group body that never closes its top-level
|
||||
// object within maxModelPeekBytes is denied, not just one carrying a
|
||||
// provably duplicate key — because a later duplicate cannot be ruled out
|
||||
// either way once the window is exhausted mid-object.
|
||||
func TestHandlePassthrough_RestrictedGroup_LargeBodyNeverClosesWindow_Returns403(t *testing.T) {
|
||||
var upstreamCalled bool
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCalled = true
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {Models: []string{"gpt-allowed"}}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
padding := strings.Repeat("x", maxModelPeekBytes*2) // single, genuine "model" — but far past the peek window
|
||||
body := `{"model":"gpt-allowed","padding":"` + padding + `"}`
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want 403 (fail closed: the object never closes within the peek window, so a later duplicate cannot be ruled out), body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if upstreamCalled {
|
||||
t.Error("upstream must never be called for a restricted group whose body exceeds the peek window without closing")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandlePassthrough_UnrestrictedGroup_PaddedDuplicateModel_NeverPeeked
|
||||
// is the GATE's own explicit requirement: an UNRESTRICTED group
|
||||
// (grp.hasModelRestriction() false) must be COMPLETELY unaffected by
|
||||
// this fix — the exact padded-duplicate body from the exploit test above
|
||||
// must still be forwarded untouched, because peekPassthroughModel is
|
||||
// never even called for such a group (unchanged since round 2).
|
||||
func TestHandlePassthrough_UnrestrictedGroup_PaddedDuplicateModel_NeverPeeked(t *testing.T) {
|
||||
var upstreamCalled bool
|
||||
var gotBody []byte
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamCalled = true
|
||||
gotBody, _ = io.ReadAll(r.Body)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}} // no Models restriction
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
padding := strings.Repeat("A", maxModelPeekBytes)
|
||||
body := `{"model":"cheap","padding":"` + padding + `","model":"expensive"}`
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 — an unrestricted group must never be affected by model-peek fixes, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if !upstreamCalled {
|
||||
t.Error("upstream must be called: an unrestricted group never peeks the body at all")
|
||||
}
|
||||
if string(gotBody) != body {
|
||||
t.Error("upstream body must be forwarded byte-for-byte unchanged for an unrestricted group")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user