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) 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 scopeOf(t *testing.T, k config.GWKey, model string) config.ModelScope { t.Helper() for _, m := range k.Models { if m.Model == model { return m } } t.Fatalf("scope %q not found in %+v", model, k.Models) return config.ModelScope{} } // Quotas live on the scope entries, not on the key: creating a key with a // budget means creating scopes that carry it, and they must survive a // read-back (persisted, not just echoed). func TestKeyAPICreatesPerModelQuota(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","models":[{"model":"m1","token_quota":50000,"req_quota":200,"period":"nhour","hours":6}]}`) 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) } sc := scopeOf(t, created.Key, "m1") if sc.TokenQuota != 50000 || sc.ReqQuota != 200 || sc.Period != "nhour" || sc.Hours != 6 { t.Fatalf("created scope did not carry the caps: %+v", sc) } back := scopeOf(t, keyRecord(t, g, created.Key.Key), "m1") if back.TokenQuota != 50000 || back.ReqQuota != 200 || back.Period != "nhour" || back.Hours != 6 { t.Errorf("read-back lost the caps: %+v", back) } } // Two models on one key carry independent budgets. func TestKeyAPIKeepsPerModelQuotaIndependent(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}) rr := adminReq(t, g, "POST", "/api/keys", `{"name":"agent-y","role":"user","models":[{"model":"m1","token_quota":1000,"period":"hour"},{"model":"m2","token_quota":9999,"req_quota":7,"period":"week"}]}`) 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) back := keyRecord(t, g, created.Key.Key) if a := scopeOf(t, back, "m1"); a.TokenQuota != 1000 || a.ReqQuota != 0 || a.Period != "hour" { t.Errorf("m1 caps wrong: %+v", a) } if b := scopeOf(t, back, "m2"); b.TokenQuota != 9999 || b.ReqQuota != 7 || b.Period != "week" { t.Errorf("m2 caps wrong: %+v", b) } } // Sending 0 explicitly lifts that model's cap. 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","models":[{"model":"m1","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) secret := created.Key.Key rr = adminReq(t, g, "PUT", "/api/keys/"+secret, `{"models":[{"model":"m1","token_quota":0,"req_quota":0,"period":""}]}`) if rr.Code != 200 { t.Fatalf("lift: %d %s", rr.Code, rr.Body.String()) } if back := scopeOf(t, keyRecord(t, g, secret), "m1"); back.TokenQuota != 0 || back.ReqQuota != 0 || back.Period != "" { t.Errorf("caps not lifted: %+v", back) } } // 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","models":[{"model":"m1","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","models":[{"model":"m1","token_quota":-5}]}`) if rr.Code != http.StatusBadRequest { t.Fatalf("want 400 for a negative quota, got %d %s", rr.Code, rr.Body.String()) } } // The rejection must say WHICH model is over budget, so an operator looking at // a key with a dozen scopes can tell which one to raise. func TestKeyAPIRejectionNamesTheModel(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}) rr := adminReq(t, g, "POST", "/api/keys", `{"name":"bad","role":"user","models":[{"model":"m1","req_quota":-1}]}`) if rr.Code != http.StatusBadRequest { t.Fatalf("want 400, got %d", rr.Code) } if !strings.Contains(rr.Body.String(), "m1") { t.Errorf("error should name the offending model: %s", rr.Body.String()) } } // A non-admin key must not be able to mint keys. func TestKeyAPIQuotaIsAdminOnly(t *testing.T) { g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"}, config.GWKey{Key: "sk-u", Role: "user", Models: []config.ModelScope{{Model: "m1", TokenQuota: 10, Period: "hour"}}}, ) req, _ := http.NewRequest("POST", "/api/keys", strings.NewReader(`{"name":"x","role":"admin","models":[{"model":"m1"}]}`)) 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 exposes the per-model caps but never a 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 per-model caps: %s", rr.Body.String()) } } // 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", Models: []config.ModelScope{{Model: "m1", TokenQuota: 123456, ReqQuota: 42, Period: "week"}}}, ) 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()) } sc := scopeOf(t, wrap.Key, "m1") if sc.TokenQuota != 123456 || sc.ReqQuota != 42 || sc.Period != "week" { t.Errorf("own quota not visible to the key's owner: %+v", sc) } } var _ = fmt.Sprintf