Files
traefik-llmgateway/users_file_test.go
T
lukaszraczylo 4c01c75276 fix(gateway): close billing, limit, routing and translation gaps
Billing:
- Passthrough SSE streams now bill the usage the upstream reports.
  The gateway falls back to a request-size estimate only on a 2xx
  stream or a truncated JSON body.
- A unified stream that fails after it starts is estimated, not
  billed at zero.
- Billing uses the built-in LiteLLM price table and honours
  modelMeta free: true.
- Pricing entries are validated. Cost arithmetic rounds the total and
  saturates on overflow.
- Gemini thinking tokens count as completion tokens.

Limits and storage:
- A rejected request rolls back its counts on the other scopes, so one
  user cannot lock out a group.
- Redis EXPIRE replies are checked. Keys recreated without a TTL are
  repaired at once.
- A mutating pipeline is not retried after reply bytes arrive.
- History, model totals and target usage report a store outage.

Routing and translation:
- Discovered ids with a provider-like prefix are listed only in
  owner/id form, so listing and routing agree.
- Anthropic and Gemini discovery follows pagination.
- The developer role maps to system.
- stream_options.include_usage is honoured on translated streams.
- tool_choice none keeps tools when history holds tool blocks.
- refusal maps to content_filter. Data URI parameters are stripped.

Streaming, retry and failover:
- SSE accepts CR, LF and CRLF line ends. This closes a CR injection
  path.
- Retry-After is never jittered. A cancelled wait returns the
  context error.
- The watchdog timer stops at end of body.
- Oversized bodies return 413.
- A failed stream ends with an error event.
- Failover restores pre-existing headers. Client cancels do not count
  as provider successes.

Also: "admin" is a reserved provider name. metrics.path must start
with "/". Pricing warnings go to the right gateway instance. The admin
targets view sends access as null or [] so the two states differ.
Upstream organisation headers are stripped. The users-file reload
error is logged once per error. Dead code is removed, and test-only
helpers move to test files.

Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
2026-09-23 09:44:31 +01:00

517 lines
17 KiB
Go

package traefikllmgateway
import (
"encoding/json"
"fmt"
"net/http/httptest"
"os"
"path/filepath"
"slices"
"sync"
"testing"
"time"
)
// --- test helpers shared with llmgateway_test.go ---
// writeUsersDoc marshals users to the {"users":[...]} shape and writes it to
// fp. When mtime is non-zero, it also forces the file's mtime to that value
// via os.Chtimes — tests use this to make an mtime change unambiguous
// regardless of filesystem timestamp resolution, instead of relying on real
// wall-clock write timing.
func writeUsersDoc(t *testing.T, fp string, users []*UserConfig, mtime time.Time) {
t.Helper()
doc := struct {
Users []*UserConfig `json:"users"`
}{Users: users}
b, err := json.Marshal(doc)
if err != nil {
t.Fatalf("json.Marshal: %v", err)
}
if err := os.WriteFile(fp, b, 0o600); err != nil {
t.Fatalf("WriteFile: %v", err)
}
if !mtime.IsZero() {
if err := os.Chtimes(fp, mtime, mtime); err != nil {
t.Fatalf("Chtimes: %v", err)
}
}
}
// writeUsersFile creates a new users file under dir/name and returns its path.
func writeUsersFile(t *testing.T, dir, name string, users []*UserConfig) string {
t.Helper()
fp := filepath.Join(dir, name)
writeUsersDoc(t, fp, users, time.Time{})
return fp
}
// identifyWithKey builds a Bearer-authenticated request for key and runs it
// through a.identify.
func identifyWithKey(a *authStore, key string) (*user, *group, bool) {
r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
r.Header.Set("Authorization", "Bearer "+key)
return a.identify(r)
}
// stubLogger is a gatewayLogger that records calls instead of writing to
// stderr, for tests that exercise maybeReload directly.
type stubLogger struct {
infos []string
errors []string
mu sync.Mutex
}
func (s *stubLogger) logf(format string, args ...any) {
s.mu.Lock()
defer s.mu.Unlock()
s.infos = append(s.infos, fmt.Sprintf(format, args...))
}
func (s *stubLogger) errorf(format string, args ...any) {
s.mu.Lock()
defer s.mu.Unlock()
s.errors = append(s.errors, fmt.Sprintf(format, args...))
}
func (s *stubLogger) infoCount() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.infos)
}
func (s *stubLogger) errorCount() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.errors)
}
// fakeClock is an injectable, manually-advanced clock for nowFn.
type fakeClock struct {
now time.Time
mu sync.Mutex
}
func (c *fakeClock) Now() time.Time {
c.mu.Lock()
defer c.mu.Unlock()
return c.now
}
func (c *fakeClock) Advance(d time.Duration) {
c.mu.Lock()
defer c.mu.Unlock()
c.now = c.now.Add(d)
}
// newReloadableAuthStore builds an authStore with inline users, performs the
// same synchronous initial load newGateway would perform against a users
// file seeded with initialFileUsers, and wires a stub logger plus a fake
// clock for throttle control.
func newReloadableAuthStore(t *testing.T, inline, initialFileUsers []*UserConfig) (a *authStore, fp string, log *stubLogger, clock *fakeClock) {
t.Helper()
dir := t.TempDir()
fp = writeUsersFile(t, dir, "users.json", initialFileUsers)
cfg := &Config{
Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}},
Groups: map[string]*GroupConfig{"eng": {}},
Users: &UsersConfig{Inline: inline},
}
a, err := newAuthStore(cfg)
if err != nil {
t.Fatalf("newAuthStore: %v", err)
}
uf := newUsersFile(fp)
initial, err := uf.load()
if err != nil {
t.Fatalf("initial load: %v", err)
}
if err := a.replaceFileUsers(initial); err != nil {
t.Fatalf("initial replaceFileUsers: %v", err)
}
clock = &fakeClock{now: time.Unix(1_700_000_000, 0)}
log = &stubLogger{}
a.usersFile = uf
a.log = log
a.nowFn = clock.Now
if info, statErr := os.Stat(fp); statErr == nil {
a.lastModTime = info.ModTime()
}
a.lastCheck = clock.Now()
return a, fp, log, clock
}
// --- usersFile.load ---
func TestUsersFile_Load_ValidFile(t *testing.T) {
dir := t.TempDir()
fp := writeUsersFile(t, dir, "users.json", []*UserConfig{
{Name: "alice", Group: "eng", APIKey: "sk-alice", Limits: &LimitsConfig{RequestsPerDay: 100}},
})
uf := newUsersFile(fp)
users, err := uf.load()
if err != nil {
t.Fatalf("load: %v", err)
}
if len(users) != 1 {
t.Fatalf("want 1 user, got %d", len(users))
}
got := users[0]
if got.Name != "alice" || got.Group != "eng" || got.APIKey != "sk-alice" {
t.Fatalf("unexpected user: %+v", got)
}
if got.Limits == nil || got.Limits.RequestsPerDay != 100 {
t.Fatalf("want limits.requestsPerDay=100, got %+v", got.Limits)
}
}
// TestUsersFile_Load_GroupsProvidersModels proves the multi-group/
// personal-grant fields (UserConfig.Groups/Providers/Models,
// llmgateway.go) round-trip through the users file's JSON shape exactly
// like every other UserConfig field.
func TestUsersFile_Load_GroupsProvidersModels(t *testing.T) {
dir := t.TempDir()
fp := writeUsersFile(t, dir, "users.json", []*UserConfig{
{
Name: "carol",
Group: "friends",
Groups: []string{"home"},
APIKey: "sk-carol",
Providers: []string{"gx10"},
Models: []string{"gx10/GLM-5.3-Flash-EXL3"},
},
})
uf := newUsersFile(fp)
users, err := uf.load()
if err != nil {
t.Fatalf("load: %v", err)
}
if len(users) != 1 {
t.Fatalf("want 1 user, got %d", len(users))
}
got := users[0]
if got.Group != "friends" {
t.Errorf("Group = %q, want %q", got.Group, "friends")
}
if !slices.Equal(got.Groups, []string{"home"}) {
t.Errorf("Groups = %v, want [home]", got.Groups)
}
if !slices.Equal(got.Providers, []string{"gx10"}) {
t.Errorf("Providers = %v, want [gx10]", got.Providers)
}
if !slices.Equal(got.Models, []string{"gx10/GLM-5.3-Flash-EXL3"}) {
t.Errorf("Models = %v, want [gx10/GLM-5.3-Flash-EXL3]", got.Models)
}
}
func TestUsersFile_Load_MalformedJSON_ReturnsError(t *testing.T) {
dir := t.TempDir()
fp := filepath.Join(dir, "bad.json")
if err := os.WriteFile(fp, []byte("{not json"), 0o600); err != nil {
t.Fatalf("WriteFile: %v", err)
}
uf := newUsersFile(fp)
if _, err := uf.load(); err == nil {
t.Fatal("want error for malformed JSON")
}
}
func TestUsersFile_Load_MissingFile_ReturnsError(t *testing.T) {
uf := newUsersFile(filepath.Join(t.TempDir(), "nope.json"))
if _, err := uf.load(); err == nil {
t.Fatal("want error for a missing file")
}
}
// --- authStore.maybeReload ---
func TestAuthStore_MaybeReload_NoFileConfigured_NoOp(t *testing.T) {
a, err := newAuthStore(testAuthCfg())
if err != nil {
t.Fatal(err)
}
a.maybeReload() // must not panic when a.usersFile is nil
if _, _, ok := identifyWithKey(a, "sk-secret"); !ok {
t.Fatal("want inline user still identifiable after a no-op maybeReload")
}
}
func TestAuthStore_MaybeReload_WithinThrottleWindow_DoesNotReload(t *testing.T) {
a, fp, _, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
writeUsersDoc(t, fp, []*UserConfig{{Name: "f2", Group: "eng", APIKey: "sk-f2"}}, time.Now().Add(time.Hour))
clock.Advance(reloadEvery - time.Second)
a.maybeReload()
if _, _, ok := identifyWithKey(a, "sk-f1"); !ok {
t.Fatal("want f1 still identifiable before the throttle window elapses")
}
if _, _, ok := identifyWithKey(a, "sk-f2"); ok {
t.Fatal("want f2 not yet identifiable before the throttle window elapses")
}
}
func TestAuthStore_MaybeReload_AfterThrottleWindow_PicksUpMtimeChange(t *testing.T) {
a, fp, _, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
writeUsersDoc(t, fp, []*UserConfig{{Name: "f2", Group: "eng", APIKey: "sk-f2"}}, time.Now().Add(time.Hour))
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if _, _, ok := identifyWithKey(a, "sk-f2"); !ok {
t.Fatal("want f2 identifiable once the throttle window elapses and mtime changed")
}
if _, _, ok := identifyWithKey(a, "sk-f1"); ok {
t.Fatal("want f1 no longer identifiable once the reload replaced the file-sourced set")
}
}
func TestAuthStore_MaybeReload_MalformedJSON_KeepsLastGoodSetAndLogsError(t *testing.T) {
a, fp, log, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
if err := os.WriteFile(fp, []byte("{not json"), 0o600); err != nil {
t.Fatalf("WriteFile: %v", err)
}
future := time.Now().Add(time.Hour)
if err := os.Chtimes(fp, future, future); err != nil {
t.Fatalf("Chtimes: %v", err)
}
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if _, _, ok := identifyWithKey(a, "sk-f1"); !ok {
t.Fatal("want the last good file-sourced set retained after a malformed reload")
}
if log.errorCount() == 0 {
t.Fatal("want an errorf call recorded for the malformed reload")
}
}
// TestAuthStore_MaybeReload_BrokenFile_ErrorLogDedupedUntilMtimeOrTextChanges
// is review-auth finding F11, 2026-09 audit: a users file left broken
// must not re-log the identical failure on every reloadEvery throttle
// window for as long as it stays broken — before this fix, lastModTime
// was never advanced on a failed reload, so every throttle window saw
// the mtime as still different from the last observed one and retried
// AND re-logged, unbounded (~17k lines/day at the default 5s throttle).
func TestAuthStore_MaybeReload_BrokenFile_ErrorLogDedupedUntilMtimeOrTextChanges(t *testing.T) {
a, fp, log, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
if err := os.WriteFile(fp, []byte("{not json"), 0o600); err != nil {
t.Fatalf("WriteFile: %v", err)
}
broken := time.Now().Add(time.Hour)
if err := os.Chtimes(fp, broken, broken); err != nil {
t.Fatalf("Chtimes: %v", err)
}
// 3 throttle windows, same broken file, same mtime: only the FIRST
// attempt may log.
for i := 0; i < 3; i++ {
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
}
if got := log.errorCount(); got != 1 {
t.Fatalf("errorCount = %d after 3 identical-failure attempts, want 1 (deduped)", got)
}
// Rewrite the SAME broken content at a NEW mtime (an operator edited
// the file, even to an identically broken state) — this must log
// again, once.
broken2 := broken.Add(time.Hour)
if err := os.Chtimes(fp, broken2, broken2); err != nil {
t.Fatalf("Chtimes: %v", err)
}
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if got := log.errorCount(); got != 2 {
t.Fatalf("errorCount = %d after the mtime changed, want 2 (a new attempt against a new mtime logs again)", got)
}
// Fix the file: a successful reload clears the dedup key.
writeUsersDoc(t, fp, []*UserConfig{{Name: "f2", Group: "eng", APIKey: "sk-f2"}}, broken2.Add(time.Hour))
clock.Advance(reloadEvery + time.Second)
if !a.maybeReload() {
t.Fatal("want the reload to succeed once the file is valid again")
}
if got := log.errorCount(); got != 2 {
t.Fatalf("errorCount = %d after a successful reload, want unchanged at 2", got)
}
// Break it again, differently: after a success, the very next failure
// must log immediately, not stay suppressed by the earlier dedup key.
if err := os.WriteFile(fp, []byte("{also not json"), 0o600); err != nil {
t.Fatalf("WriteFile: %v", err)
}
broken3 := broken2.Add(2 * time.Hour)
if err := os.Chtimes(fp, broken3, broken3); err != nil {
t.Fatalf("Chtimes: %v", err)
}
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if got := log.errorCount(); got != 3 {
t.Fatalf("errorCount = %d after a fresh failure post-success, want 3 (dedup key was cleared on success)", got)
}
}
func TestAuthStore_MaybeReload_UnchangedMtime_DoesNotReload(t *testing.T) {
a, fp, log, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
_ = fp
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if _, _, ok := identifyWithKey(a, "sk-f1"); !ok {
t.Fatal("want f1 still identifiable when the file did not change")
}
if log.infoCount() != 0 {
t.Fatal("want no count-change log when the file did not change")
}
}
// TestAuthStore_MaybeReload_BackwardMtimeMove_TriggersReload is the
// regression for the Equal-vs-After mtime check: a file restored from a
// backup, or copied with `cp -p` from an older file, can move the mtime
// backward. That must still trigger a reload — only an mtime equal to the
// last observed one may skip it.
func TestAuthStore_MaybeReload_BackwardMtimeMove_TriggersReload(t *testing.T) {
a, fp, _, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
past := time.Now().Add(-time.Hour)
writeUsersDoc(t, fp, []*UserConfig{{Name: "f2", Group: "eng", APIKey: "sk-f2"}}, past)
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if _, _, ok := identifyWithKey(a, "sk-f2"); !ok {
t.Fatal("want a backward mtime move to still trigger a reload")
}
if _, _, ok := identifyWithKey(a, "sk-f1"); ok {
t.Fatal("want f1 no longer identifiable once the backward-mtime reload replaced the file-sourced set")
}
}
func TestAuthStore_MaybeReload_UserCountChange_LogsViaLogf(t *testing.T) {
a, fp, log, clock := newReloadableAuthStore(t, nil, []*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-f1"}})
writeUsersDoc(t, fp, []*UserConfig{
{Name: "f1", Group: "eng", APIKey: "sk-f1"},
{Name: "f2", Group: "eng", APIKey: "sk-f2"},
}, time.Now().Add(time.Hour))
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if log.infoCount() == 0 {
t.Fatal("want a logf call recorded for the user-count change (1 -> 2)")
}
}
// TestAuthStore_MaybeReload_FileUserSameNameAsInline_FileWins is the merge
// semantic required by Task 4: a file user overrides an inline user of the
// same name, even when their API keys differ. The inline entry's key must
// stop working and the file entry's key must resolve to the file's data.
func TestAuthStore_MaybeReload_FileUserSameNameAsInline_FileWins(t *testing.T) {
inline := []*UserConfig{{Name: "alice", Group: "eng", APIKey: "sk-inline-alice"}}
a, fp, _, clock := newReloadableAuthStore(t, inline, nil)
writeUsersDoc(t, fp, []*UserConfig{{Name: "alice", Group: "eng", APIKey: "sk-file-alice"}}, time.Now().Add(time.Hour))
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
if _, _, ok := identifyWithKey(a, "sk-inline-alice"); ok {
t.Fatal("want the inline alice overridden once the file defines a same-named user")
}
u, _, ok := identifyWithKey(a, "sk-file-alice")
if !ok || u.name != "alice" {
t.Fatalf("want file alice identifiable, got %v ok=%v", u, ok)
}
}
// TestAuthStore_MaybeReload_MultiGroupUser_RebuiltPrincipalStillAuthorizesCorrectly
// is LOW-7's regression coverage (review round 2): a users-file hot
// reload rebuilds EVERY principal from scratch (replaceFileUsers calls
// buildEntry fresh for each user) — a multi-group user's rebuilt
// principal must still resolve exactly the same union of member-group
// authorization after a reload as before it. This is functional
// correctness, distinct from MEDIUM-3's own modelsCache-growth-bound
// test: that test proves the CACHE stays bounded across reloads; this
// one proves the REBUILT principal itself still authorizes correctly.
func TestAuthStore_MaybeReload_MultiGroupUser_RebuiltPrincipalStillAuthorizesCorrectly(t *testing.T) {
dir := t.TempDir()
fp := writeUsersFile(t, dir, "users.json", []*UserConfig{
{Name: "carol", Group: "eng", Groups: []string{"ops"}, APIKey: "sk-carol"},
})
cfg := &Config{
Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}},
Groups: map[string]*GroupConfig{
"eng": {Providers: []string{"alpha"}},
"ops": {Providers: []string{"beta"}},
},
}
a, err := newAuthStore(cfg)
if err != nil {
t.Fatalf("newAuthStore: %v", err)
}
uf := newUsersFile(fp)
initial, err := uf.load()
if err != nil {
t.Fatalf("initial load: %v", err)
}
if err := a.replaceFileUsers(initial); err != nil {
t.Fatalf("initial replaceFileUsers: %v", err)
}
clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
a.usersFile = uf
a.log = &stubLogger{}
a.nowFn = clock.Now
if info, statErr := os.Stat(fp); statErr == nil {
a.lastModTime = info.ModTime()
}
a.lastCheck = clock.Now()
_, grpBefore, ok := identifyWithKey(a, "sk-carol")
if !ok {
t.Fatal("identify failed before reload")
}
if !grpBefore.allowsProvider("alpha") || !grpBefore.allowsProvider("beta") {
t.Fatalf("sanity: want both alpha and beta reachable before reload, providers=%v", grpBefore.providers)
}
// Force a real reload of the IDENTICAL content — buildEntry rebuilds
// carol's principal as a brand-new *group either way (effectiveGroup
// never caches or reuses a synthetic principal across calls).
writeUsersDoc(t, fp, []*UserConfig{
{Name: "carol", Group: "eng", Groups: []string{"ops"}, APIKey: "sk-carol"},
}, time.Now().Add(time.Hour))
clock.Advance(reloadEvery + time.Second)
a.maybeReload()
_, grpAfter, ok := identifyWithKey(a, "sk-carol")
if !ok {
t.Fatal("identify failed after reload")
}
if grpAfter == grpBefore {
t.Fatal("sanity: want a NEW *group pointer after reload (buildEntry rebuilds every principal from scratch) — otherwise this test never exercises the rebuild path at all")
}
if !grpAfter.allowsProvider("alpha") || !grpAfter.allowsProvider("beta") {
t.Errorf("want both alpha and beta still reachable after reload, providers=%v", grpAfter.providers)
}
if grpAfter.allowsProvider("gamma") {
t.Error("want gamma still denied after reload")
}
}