package server import ( "ai-router/internal/protocol" "context" "encoding/json" "fmt" "net/http/httptest" "strings" "testing" "time" "unicode/utf8" ) func TestSecurityStatsSQLSemantics(t *testing.T) { s := testServer(t) now := time.Now().UTC().Format(time.RFC3339Nano) for _, l := range []requestLog{ {CreatedAt: now, KeyID: "a", KeyName: "old", Status: 200, Usage: protocol.Usage{Known: true, Input: 8}}, {CreatedAt: now, KeyID: "a", KeyName: "new", Status: 200, Error: "failed", Usage: protocol.Usage{Raw: []byte("{}"), Input: 100}}, {CreatedAt: now, KeyID: "b", Status: 400}, {CreatedAt: now, KeyID: "c", Status: 204, Usage: protocol.Usage{Known: true}}, } { s.record(l) } if err := s.flushLogs(); err != nil { t.Fatal(err) } for _, tc := range []struct { query string requests, success, unknown int64 input *int64 }{ {"", 4, 2, 2, new(int64(8))}, {"?key_id=b", 1, 0, 1, nil}, {"?key_id=c", 1, 1, 0, new(int64(0))}, {"?key_id=missing", 0, 0, 0, nil}, } { w := httptest.NewRecorder() s.stats(w, httptest.NewRequest("GET", "/admin/api/stats"+tc.query, nil)) if w.Code != 200 { t.Fatal(w.Body.String()) } var got struct { Summary summary ByKey []struct{ Name string } `json:"by_key"` } if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil { t.Fatal(err) } if got.Summary.Requests != tc.requests || got.Summary.Success != tc.success || got.Summary.Unknown != tc.unknown { t.Fatal(w.Body.String()) } if (got.Summary.Input == nil) != (tc.input == nil) || tc.input != nil && *got.Summary.Input != *tc.input { t.Fatal(w.Body.String()) } if tc.query == "" && got.ByKey[0].Name != "new" { t.Fatal(w.Body.String()) } } } func TestSecurityLogFieldsBounded(t *testing.T) { s := testServer(t) huge := strings.Repeat("界", 10000) s.record(requestLog{RequestModel: huge, ActualModel: huge, Channel: huge, KeyID: huge, KeyName: huge, Upstream: huge, Downstream: huge, Error: huge, Usage: protocol.Usage{Known: true, Input: 12, Raw: []byte(huge)}}) if err := s.flushLogs(); err != nil { t.Fatal(err) } var model, actual, channel, key, name, up, down, message, raw string var input int64 if err := s.db.QueryRow("SELECT request_model,actual_model,channel,key_id,key_name,upstream,downstream,error,raw_usage,input_tokens FROM logs").Scan(&model, &actual, &channel, &key, &name, &up, &down, &message, &raw, &input); err != nil { t.Fatal(err) } for _, field := range []string{model, actual, channel, key, name, up, down} { if len(field) > maxLogFieldBytes || !utf8.ValidString(field) { t.Fatal("unbounded or invalid field") } } if len(message) > maxLogErrorBytes || !utf8.ValidString(message) || len(raw) > maxLogUsageBytes || !json.Valid([]byte(raw)) || input != 12 { t.Fatal("invalid bounded log") } } func TestSecurityLogRetentionOptInAndBatch(t *testing.T) { s := testServer(t) old := time.Now().Add(-91 * 24 * time.Hour).Format(time.RFC3339Nano) for i := 0; i < logCleanupBatch+3; i++ { s.record(requestLog{CreatedAt: old}) } s.record(requestLog{CreatedAt: time.Now().Format(time.RFC3339Nano)}) if err := s.flushLogs(); err != nil { t.Fatal(err) } if s.logRetentionDays.Load() != 0 { t.Fatal("retention enabled by default") } if err := s.purgeExpiredLogs(context.Background(), 0); err != nil { t.Fatal(err) } if n := logCount(t, s.db); n != logCleanupBatch+4 { t.Fatal(n) } s.SetLogRetention(90) if s.logRetentionDays.Load() != 90 { t.Fatal("retention not set") } if err := s.purgeExpiredLogs(context.Background(), 90); err != nil { t.Fatal(err) } if n := logCount(t, s.db); n != 4 { t.Fatal("unbounded purge", n) } if err := s.purgeExpiredLogs(context.Background(), 90); err != nil { t.Fatal(err) } if n := logCount(t, s.db); n != 1 { t.Fatal(n) } s.SetLogRetention(0) if s.logRetentionDays.Load() != 0 { t.Fatal("retention not disabled") } s.SetLogRetention(90) if err := s.Close(); err != nil { t.Fatal(err) } } func TestStatsGroupLimitAndCancellation(t *testing.T) { s := testServer(t) tx, err := s.db.Begin() if err != nil { t.Fatal(err) } for i := 0; i < 10001; i++ { if err := insertLog(context.Background(), tx, requestLog{CreatedAt: "2026-01-01T00:00:00Z", KeyID: fmt.Sprint(i)}); err != nil { tx.Rollback() t.Fatal(err) } } if err := tx.Commit(); err != nil { t.Fatal(err) } w := httptest.NewRecorder() s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil)) if w.Code != 422 { t.Fatal("unbounded group result", w.Code) } ctx, cancel := context.WithCancel(context.Background()) cancel() w = httptest.NewRecorder() s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil).WithContext(ctx)) if w.Code == 200 { t.Fatal("cancelled query succeeded") } var day string if err := s.db.QueryRow("SELECT router_log_day(NULL, 'UTC')").Scan(&day); err == nil { t.Fatal("invalid timestamp accepted") } }