package gateway import ( "encoding/json" "fmt" "net/http" "net/http/httptest" "os" "strings" "testing" "llmsproxy/internal/config" "llmsproxy/internal/core" ) // adminGateway builds a gateway whose admin key can call /api/keys. func adminGateway(t *testing.T, keys ...config.GWKey) *Gateway { t.Helper() td := t.TempDir() cfgPath := td + "/config.yaml" if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil { t.Fatal(err) } cfg := &config.Config{ Path: cfgPath, AdapterDir: td + "/adapters", RuntimeFile: td + "/runtime.json", Keys: keys, } if err := cfg.ApplyDefaults(); err != nil { t.Fatal(err) } c, err := core.NewFromConfig(cfg) if err != nil { t.Fatalf("core: %v", err) } t.Cleanup(c.Close) g, err := New(c, []string{"sk-admin"}) if err != nil { t.Fatalf("gateway: %v", err) } return g } func adminReq(t *testing.T, g *Gateway, method, path, body string) *httptest.ResponseRecorder { t.Helper() req, _ := http.NewRequest(method, path, strings.NewReader(body)) req.Header.Set("Authorization", "Bearer sk-admin") req.Header.Set("Content-Type", "application/json") rr := httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) return rr } // keyRecord pulls one key's stored record out of the admin list. func keyRecord(t *testing.T, g *Gateway, secret string) config.GWKey { t.Helper() rr := adminReq(t, g, "GET", "/api/keys", "") if rr.Code != 200 { t.Fatalf("GET /api/keys: %d %s", rr.Code, rr.Body.String()) } var out struct { Keys []config.GWKey `json:"keys"` } if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil { t.Fatalf("decode: %v (%s)", err, rr.Body.String()) } for _, k := range out.Keys { if k.Key == secret { return k } } t.Fatalf("key %q not found in %s", secret, rr.Body.String()) return config.GWKey{} } func TestKeyAPIStoresQuota(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}) rr := adminReq(t, g, "POST", "/api/keys", `{"name":"agent-x","role":"user","token_quota":50000,"req_quota":200,"period":"nhour","hours":6,"models":[{"model":"m1"}]}`) if rr.Code != 200 { t.Fatalf("create: %d %s", rr.Code, rr.Body.String()) } var created struct { Key config.GWKey `json:"key"` } if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil { t.Fatalf("decode: %v", err) } if created.Key.TokenQuota != 50000 || created.Key.ReqQuota != 200 || created.Key.Period != "nhour" || created.Key.Hours != 6 { t.Fatalf("created key did not carry the caps: %+v", created.Key) } // and it must survive a read-back (persisted, not just echoed) back := keyRecord(t, g, created.Key.Key) if back.TokenQuota != 50000 || back.Period != "nhour" || back.Hours != 6 { t.Errorf("read-back lost the caps: %+v", back) } } // Editing only the model scope must not silently clear a key's budget: the // caps are pointers precisely so "absent" is not "zero". func TestKeyAPIUpdateKeepsQuotaWhenOmitted(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}) rr := adminReq(t, g, "POST", "/api/keys", `{"name":"agent-x","role":"user","token_quota":50000,"period":"day-typo-free","models":[{"model":"m1"}]}`) rr = adminReq(t, g, "POST", "/api/keys", `{"name":"y","role":"user","token_quota":50000,"period":"hour","models":[{"model":"m1"}]}`) if rr.Code != 200 { t.Fatalf("setup create: %d %s", rr.Code, rr.Body.String()) } var created struct { Key config.GWKey `json:"key"` } _ = json.Unmarshal(rr.Body.Bytes(), &created) // a scope-only edit rr = adminReq(t, g, "PUT", "/api/keys/"+created.Key.Key, `{"name":"agent-y","models":[{"model":"m1"},{"model":"m2"}]}`) if rr.Code != 200 { t.Fatalf("update: %d %s", rr.Code, rr.Body.String()) } back := keyRecord(t, g, created.Key.Key) if back.TokenQuota != 50000 { t.Errorf("token_quota was cleared by a scope-only edit: %d", back.TokenQuota) } if back.Period != "hour" { t.Errorf("period was cleared by a scope-only edit: %q", back.Period) } if len(back.Models) != 2 { t.Errorf("scope edit did not apply: %+v", back.Models) } } // Sending 0 explicitly must lift the cap, not be treated as "absent". func TestKeyAPIUpdateZeroLiftsCap(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}) rr := adminReq(t, g, "POST", "/api/keys", `{"name":"z","role":"user","token_quota":1000,"period":"hour"}`) if rr.Code != 200 { t.Fatalf("create: %d %s", rr.Code, rr.Body.String()) } var created struct { Key config.GWKey `json:"key"` } _ = json.Unmarshal(rr.Body.Bytes(), &created) rr = adminReq(t, g, "PUT", "/api/keys/"+created.Key.Key, `{"token_quota":0}`) if rr.Code != 200 { t.Fatalf("lift: %d %s", rr.Code, rr.Body.String()) } if back := keyRecord(t, g, created.Key.Key); back.TokenQuota != 0 { t.Errorf("token_quota = %d, want 0 (cap lifted)", back.TokenQuota) } } // A misspelled period must be refused, not quietly turned into an all-time // quota — which is the exact opposite of what the operator typed. func TestKeyAPIRejectsBadPeriod(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}) rr := adminReq(t, g, "POST", "/api/keys", `{"name":"bad","role":"user","token_quota":1000,"period":"houre"}`) if rr.Code != http.StatusBadRequest { t.Fatalf("want 400 for a bad period, got %d %s", rr.Code, rr.Body.String()) } if !strings.Contains(rr.Body.String(), "period") { t.Errorf("error should name the period field: %s", rr.Body.String()) } } func TestKeyAPIRejectsNegativeQuota(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}) rr := adminReq(t, g, "POST", "/api/keys", `{"name":"bad","role":"user","token_quota":-5}`) if rr.Code != http.StatusBadRequest { t.Fatalf("want 400 for a negative quota, got %d %s", rr.Code, rr.Body.String()) } } // A non-admin key must not be able to set or read another key's budget. func TestKeyAPIQuotaIsAdminOnly(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}, config.GWKey{Key: "sk-u", Role: "user", TokenQuota: 10, Period: "hour"}, ) req, _ := http.NewRequest("POST", "/api/keys", strings.NewReader(`{"name":"x","role":"admin","token_quota":0}`)) req.Header.Set("Authorization", "Bearer sk-u") req.Header.Set("Content-Type", "application/json") rr := httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != http.StatusForbidden { t.Fatalf("non-admin create: want 403, got %d %s", rr.Code, rr.Body.String()) } // /api/v1/keys echoes the caps but never the secret rr = adminReq(t, g, "GET", "/api/v1/keys", "") if rr.Code != 200 { t.Fatalf("GET /api/v1/keys: %d", rr.Code) } if strings.Contains(rr.Body.String(), "sk-u") { t.Error("/api/v1/keys leaked a key secret") } if !strings.Contains(rr.Body.String(), `"token_quota":10`) { t.Errorf("/api/v1/keys should expose the cap: %s", rr.Body.String()) } } // The AUTO scope entry must honour its reset window: usage that aged out of // the window must not count against a per-key cap. func TestAutoScopeQuotaHonoursWindow(t *testing.T) { g, _ := quotaGateway(t, config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{ {Model: "AUTO", TokenQuota: 1000, Period: "hour"}, }}, config.GWKey{Key: "sk-b", Role: "user"}, ) ctx := quotaCtx(t, g, "sk-a") sc := config.ModelScope{Model: "AUTO", TokenQuota: 1000, Period: "hour"} // aged-out usage: 2 days old, 5M tokens — must be invisible to a 1h window g.stats.Record(Req{Time: nowMSOffset(-48 * 3600 * 1000), Key: keyID("sk-a"), Model: "m1", Source: "up", Prompt: 2500000, Compl: 2500000, OK: true, Status: 200}) if used := g.scopeTokens(ctx, sc); used != 0 { t.Fatalf("AUTO scope saw %d tokens outside its 1h window; the period is being ignored", used) } // in-window usage counts g.stats.Record(Req{Time: nowMSOffset(0), Key: keyID("sk-a"), Model: "m1", Source: "up", Prompt: 400, Compl: 400, OK: true, Status: 200}) if used := g.scopeTokens(ctx, sc); used != 800 { t.Fatalf("AUTO scope used = %d, want 800", used) } } var _ = fmt.Sprintf // A user must be able to see their own budget: /api/keys/me is the only key // view a non-admin gets, so a cap missing from it is invisible to the very // client it constrains. func TestKeyMeExposesOwnQuota(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}, config.GWKey{Key: "sk-u", Role: "user", Name: "agent", TokenQuota: 123456, ReqQuota: 42, Period: "week", Hours: 0}, ) req, _ := http.NewRequest("GET", "/api/keys/me", nil) req.Header.Set("Authorization", "Bearer sk-u") rr := httptest.NewRecorder() g.Handler().ServeHTTP(rr, req) if rr.Code != 200 { t.Fatalf("GET /api/keys/me: %d %s", rr.Code, rr.Body.String()) } // the endpoint wraps the record: {"key": {...}} var wrap struct { Key config.GWKey `json:"key"` } if err := json.Unmarshal(rr.Body.Bytes(), &wrap); err != nil { t.Fatalf("decode: %v (%s)", err, rr.Body.String()) } me := wrap.Key if me.TokenQuota != 123456 || me.ReqQuota != 42 || me.Period != "week" { t.Errorf("own quota not visible to the key's owner: %+v", me) } }