ai-router/internal/server/security_stats_test.go

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")
	}
}