mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-13 04:58:50 +00:00
feat: users file with throttled hot reload
Add usersFile (newUsersFile/load) for the JSON users-file shape
{"users":[{name,group,apiKey,limits?}]}, and authStore.maybeReload
for throttled hot reload: checks mtime at most every reloadEvery
(5s), reloads on change, keeps the last good set and logs via
errorf on any error (stat, parse, or a rejected replace), and logs
via logf when the reload changes the file-sourced user count.
nowFn is injectable for tests.
Wire-up in llmgateway.go: newGateway does a synchronous initial
load when Users.File is set (a bad initial file fails the
constructor); ServeHTTP calls g.auth.maybeReload() at entry
(no-op when no file is configured).
Also fixes replaceFileUsers's merge semantics: a file user now
overrides an inline user of the same name, even when their keys
differ, dropping the inline entry from the rebuilt set. The prior
digest-only keying let a same-named inline/file pair with
different keys coexist instead of the file version winning.
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"path"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// user is an authenticated API-key holder resolved by authStore.identify.
|
||||
@@ -84,18 +85,26 @@ type authEntry struct {
|
||||
// their group. Raw API keys are never retained after construction — only
|
||||
// their SHA-256 digests.
|
||||
//
|
||||
// mu guards byDigest for readers (identify) and the final swap in
|
||||
// replaceFileUsers. buildMu serializes the whole body of replaceFileUsers —
|
||||
// including buildEntry's potentially slow file-backed secret resolution —
|
||||
// so concurrent reload attempts never interleave or overwrite each other.
|
||||
// buildMu is never held at the same time as mu, so a rebuild in progress
|
||||
// never blocks identify's readers.
|
||||
// mu guards byDigest and fileUserCount for readers (identify) and the final
|
||||
// swap in replaceFileUsers. buildMu serializes the whole body of
|
||||
// replaceFileUsers — including buildEntry's potentially slow file-backed
|
||||
// secret resolution — so concurrent reload attempts never interleave or
|
||||
// overwrite each other. buildMu is never held at the same time as mu, so a
|
||||
// rebuild in progress never blocks identify's readers. reloadMu (Task 4)
|
||||
// serializes maybeReload's throttle check and the reload it may trigger.
|
||||
type authStore struct {
|
||||
groups map[string]*group
|
||||
inline map[[32]byte]*authEntry // built once by newAuthStore; never mutated afterward
|
||||
byDigest map[[32]byte]*authEntry // inline entries plus the current file-sourced set; guarded by mu
|
||||
mu sync.RWMutex
|
||||
buildMu sync.Mutex
|
||||
groups map[string]*group
|
||||
inline map[[32]byte]*authEntry // built once by newAuthStore; never mutated afterward
|
||||
byDigest map[[32]byte]*authEntry // inline entries plus the current file-sourced set; guarded by mu
|
||||
usersFile *usersFile // nil when Users.File is not configured; maybeReload no-ops
|
||||
log gatewayLogger // set alongside usersFile; unused when usersFile is nil
|
||||
nowFn func() time.Time // injected for tests; defaults to time.Now
|
||||
lastCheck time.Time // guarded by reloadMu
|
||||
lastModTime time.Time // guarded by reloadMu
|
||||
fileUserCount int // guarded by mu; size of the current file-sourced user set
|
||||
mu sync.RWMutex
|
||||
buildMu sync.Mutex
|
||||
reloadMu sync.Mutex
|
||||
}
|
||||
|
||||
// newAuthStore builds an authStore from cfg's groups and inline users. It
|
||||
@@ -106,6 +115,7 @@ func newAuthStore(cfg *Config) (*authStore, error) {
|
||||
groups: make(map[string]*group, len(cfg.Groups)),
|
||||
inline: make(map[[32]byte]*authEntry),
|
||||
byDigest: make(map[[32]byte]*authEntry),
|
||||
nowFn: time.Now,
|
||||
}
|
||||
for name, gc := range cfg.Groups {
|
||||
a.groups[name] = &group{
|
||||
@@ -160,11 +170,16 @@ func (a *authStore) buildEntry(uc *UserConfig) (*authEntry, error) {
|
||||
}, nil
|
||||
}
|
||||
|
||||
// replaceFileUsers rebuilds the file-sourced portion of the store from us,
|
||||
// leaving the inline users untouched. It is used by the config-reload path
|
||||
// (Task 4) to pick up changes to Users.File without restarting the plugin.
|
||||
// On error the store is left exactly as it was before the call — the new
|
||||
// set is validated and built in full before it replaces the old one.
|
||||
// replaceFileUsers rebuilds the file-sourced portion of the store from us. A
|
||||
// file user overrides an inline user of the same name — the inline entry is
|
||||
// dropped from the rebuilt set, even if its API key differs from the file
|
||||
// user's — so an operator can promote or rotate a statically-configured user
|
||||
// through the hot-reloadable file. Inline users with no same-named file
|
||||
// counterpart are carried forward unchanged. This is used by the
|
||||
// config-reload path (Task 4) to pick up changes to Users.File without
|
||||
// restarting the plugin. On error the store is left exactly as it was before
|
||||
// the call — the new set is validated and built in full before it replaces
|
||||
// the old one.
|
||||
//
|
||||
// buildMu is held for the whole rebuild, not just the final swap — that
|
||||
// serializes concurrent callers (a reload timer must never race itself), so
|
||||
@@ -177,8 +192,18 @@ func (a *authStore) replaceFileUsers(us []*UserConfig) error {
|
||||
a.buildMu.Lock()
|
||||
defer a.buildMu.Unlock()
|
||||
|
||||
fileNames := make(map[string]bool, len(us))
|
||||
for _, uc := range us {
|
||||
if uc != nil {
|
||||
fileNames[uc.Name] = true
|
||||
}
|
||||
}
|
||||
|
||||
next := make(map[[32]byte]*authEntry, len(a.inline)+len(us))
|
||||
for digest, entry := range a.inline {
|
||||
if fileNames[entry.user.name] {
|
||||
continue // a file user of the same name overrides this inline user
|
||||
}
|
||||
next[digest] = entry
|
||||
}
|
||||
for _, uc := range us {
|
||||
@@ -194,6 +219,7 @@ func (a *authStore) replaceFileUsers(us []*UserConfig) error {
|
||||
|
||||
a.mu.Lock()
|
||||
a.byDigest = next
|
||||
a.fileUserCount = len(us)
|
||||
a.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
+35
-1
@@ -6,7 +6,9 @@ package traefikllmgateway
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
)
|
||||
|
||||
// Config is the plugin's dynamic configuration, populated by Traefik from
|
||||
@@ -125,13 +127,45 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Gateway{next: next, name: name, cfg: config, auth: auth}, nil
|
||||
g := &Gateway{next: next, name: name, cfg: config, auth: auth}
|
||||
|
||||
if config.Users != nil && config.Users.File != "" {
|
||||
if err := attachUsersFile(auth, config.Users.File, g); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return g, nil
|
||||
}
|
||||
|
||||
// attachUsersFile performs the synchronous initial load of a Users.File path
|
||||
// and wires auth up for throttled hot reload. A failing initial load is a
|
||||
// constructor error — fail fast on bad config rather than start with an
|
||||
// empty file-sourced user set.
|
||||
func attachUsersFile(auth *authStore, path string, log gatewayLogger) error {
|
||||
uf := newUsersFile(path)
|
||||
initial, err := uf.load()
|
||||
if err != nil {
|
||||
return fmt.Errorf("llmgateway: cannot load initial users file %q: %w", path, err)
|
||||
}
|
||||
if err := auth.replaceFileUsers(initial); err != nil {
|
||||
return fmt.Errorf("llmgateway: initial users file %q: %w", path, err)
|
||||
}
|
||||
|
||||
auth.usersFile = uf
|
||||
auth.log = log
|
||||
if info, statErr := os.Stat(path); statErr == nil {
|
||||
auth.lastModTime = info.ModTime()
|
||||
}
|
||||
auth.lastCheck = auth.nowFn()
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServeHTTP is the internal router. It grows in later tasks.
|
||||
func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
sw := &statusTrackingWriter{ResponseWriter: w}
|
||||
defer recoverPanic(sw, g)
|
||||
g.auth.maybeReload()
|
||||
if g.cfg.PassthroughUnknown {
|
||||
g.next.ServeHTTP(sw, r)
|
||||
return
|
||||
|
||||
@@ -8,9 +8,11 @@ import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCreateConfig_ReturnsEmptyConfig(t *testing.T) {
|
||||
@@ -287,3 +289,130 @@ func TestServeHTTP_PanicAfterHeadersCommitted_DoesNotOverwriteResponse(t *testin
|
||||
t.Fatalf("want panic still logged even though no envelope was written, got %q", logBuf.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewGateway_UsersFile_InitialLoadFailure_ReturnsConstructorError is the
|
||||
// fail-fast case: a malformed users file must fail plugin construction, not
|
||||
// leave the plugin running with an empty file-sourced user set.
|
||||
func TestNewGateway_UsersFile_InitialLoadFailure_ReturnsConstructorError(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)
|
||||
}
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||
cfg.Users = &UsersConfig{File: fp}
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err == nil {
|
||||
t.Fatal("want constructor error for a malformed initial users file")
|
||||
}
|
||||
if h != nil {
|
||||
t.Fatal("want nil handler on error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewGateway_UsersFile_InitialLoadMissingFile_ReturnsConstructorError
|
||||
// covers the other fail-fast path: a configured file that does not exist.
|
||||
func TestNewGateway_UsersFile_InitialLoadMissingFile_ReturnsConstructorError(t *testing.T) {
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||
cfg.Users = &UsersConfig{File: filepath.Join(t.TempDir(), "nope.json")}
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err == nil {
|
||||
t.Fatal("want constructor error for a missing initial users file")
|
||||
}
|
||||
if h != nil {
|
||||
t.Fatal("want nil handler on error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewGateway_UsersFile_ValidFile_LoadsUsersSynchronously verifies the
|
||||
// synchronous initial load: a file-sourced user must be identifiable
|
||||
// immediately after New returns, with no reload wait.
|
||||
func TestNewGateway_UsersFile_ValidFile_LoadsUsersSynchronously(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "users.json")
|
||||
writeUsersDoc(t, fp, []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}, time.Time{})
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{File: fp}
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
gw, ok := h.(*Gateway)
|
||||
if !ok {
|
||||
t.Fatal("handler is not *Gateway")
|
||||
}
|
||||
if _, _, ok := identifyWithKey(gw.auth, "sk-alice"); !ok {
|
||||
t.Fatal("want file-sourced user identifiable immediately after construction")
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeHTTP_NoUsersFile_MaybeReloadIsNoOp verifies that entry-time
|
||||
// maybeReload is a cheap no-op — and, crucially, does not panic — when no
|
||||
// users file is configured.
|
||||
func TestServeHTTP_NoUsersFile_MaybeReloadIsNoOp(t *testing.T) {
|
||||
called := false
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
called = true
|
||||
})
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||
cfg.PassthroughUnknown = true
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
if !called {
|
||||
t.Fatal("want next handler called")
|
||||
}
|
||||
}
|
||||
|
||||
// TestServeHTTP_TriggersUsersFileReload is the wire-up regression: ServeHTTP
|
||||
// must call g.auth.maybeReload() at entry, not just newGateway's initial
|
||||
// load. It drives a real request through ServeHTTP (never calling
|
||||
// maybeReload directly) and expects the file change to have been picked up.
|
||||
func TestServeHTTP_TriggersUsersFileReload(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
fp := filepath.Join(dir, "users.json")
|
||||
writeUsersDoc(t, fp, []*UserConfig{{Name: "f1", Group: "default", APIKey: "sk-f1"}}, time.Time{})
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{File: fp}
|
||||
cfg.PassthroughUnknown = true
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
gw, ok := h.(*Gateway)
|
||||
if !ok {
|
||||
t.Fatal("handler is not *Gateway")
|
||||
}
|
||||
|
||||
// Force the throttle window open and give the file an unambiguous
|
||||
// future mtime, then drive one request through ServeHTTP.
|
||||
clock := &fakeClock{now: gw.auth.nowFn().Add(reloadEvery + time.Second)}
|
||||
gw.auth.nowFn = clock.Now
|
||||
writeUsersDoc(t, fp, []*UserConfig{{Name: "f2", Group: "default", APIKey: "sk-f2"}}, time.Now().Add(time.Hour))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
if _, _, ok := identifyWithKey(gw.auth, "sk-f2"); !ok {
|
||||
t.Fatal("want ServeHTTP's entry-time maybeReload call to have picked up the file change")
|
||||
}
|
||||
}
|
||||
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
)
|
||||
|
||||
// reloadEvery is the minimum interval between users-file mtime checks on the
|
||||
// request path. maybeReload never stats the file more often than this, even
|
||||
// under heavy concurrent traffic.
|
||||
const reloadEvery = 5 * time.Second
|
||||
|
||||
// usersFile is a JSON-encoded file of API-key users at path, in the shape
|
||||
// {"users":[{"name":...,"group":...,"apiKey":...,"limits":{...}}]}.
|
||||
type usersFile struct {
|
||||
path string
|
||||
}
|
||||
|
||||
// newUsersFile returns a usersFile reading from path. It does not touch the
|
||||
// filesystem — call load to read and parse the file.
|
||||
func newUsersFile(path string) *usersFile {
|
||||
return &usersFile{path: path}
|
||||
}
|
||||
|
||||
// usersFileDoc is the on-disk shape of a users file.
|
||||
type usersFileDoc struct {
|
||||
Users []*UserConfig `json:"users"`
|
||||
}
|
||||
|
||||
// load reads and parses uf's file. path is operator-supplied middleware
|
||||
// configuration, not untrusted request input.
|
||||
func (uf *usersFile) load() ([]*UserConfig, error) {
|
||||
b, err := os.ReadFile(uf.path) // #nosec G304 -- operator-supplied users file path from middleware config
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("llmgateway: cannot read users file %q: %w", uf.path, err)
|
||||
}
|
||||
var doc usersFileDoc
|
||||
if err := json.Unmarshal(b, &doc); err != nil {
|
||||
return nil, fmt.Errorf("llmgateway: cannot parse users file %q: %w", uf.path, err)
|
||||
}
|
||||
return doc.Users, nil
|
||||
}
|
||||
|
||||
// gatewayLogger is the logging surface maybeReload needs. *Gateway
|
||||
// implements it (see logger.go); authStore depends on this small interface
|
||||
// rather than *Gateway directly, so the reload path can be tested without an
|
||||
// HTTP handler chain.
|
||||
type gatewayLogger interface {
|
||||
logf(format string, args ...any)
|
||||
errorf(format string, args ...any)
|
||||
}
|
||||
|
||||
// maybeReload reloads a's file-sourced users when a.usersFile is configured,
|
||||
// reloadEvery has elapsed since the last check, and the file's mtime has
|
||||
// changed since the last successful load. It runs on the request path
|
||||
// (ServeHTTP calls it at entry), so it must stay cheap when nothing changed
|
||||
// and must never panic: on any error — stat, load, or replaceFileUsers
|
||||
// rejecting the new set — it logs via a.log.errorf and keeps the last good
|
||||
// set. A successful reload that changes the file-sourced user count logs via
|
||||
// a.log.logf.
|
||||
func (a *authStore) maybeReload() {
|
||||
if a.usersFile == nil {
|
||||
return
|
||||
}
|
||||
|
||||
a.reloadMu.Lock()
|
||||
defer a.reloadMu.Unlock()
|
||||
|
||||
now := a.nowFn()
|
||||
if now.Sub(a.lastCheck) < reloadEvery {
|
||||
return
|
||||
}
|
||||
a.lastCheck = now
|
||||
|
||||
info, err := os.Stat(a.usersFile.path)
|
||||
if err != nil {
|
||||
a.log.errorf("llmgateway: cannot stat users file %q: %v", a.usersFile.path, err)
|
||||
return
|
||||
}
|
||||
if !info.ModTime().After(a.lastModTime) {
|
||||
return
|
||||
}
|
||||
|
||||
users, err := a.usersFile.load()
|
||||
if err != nil {
|
||||
a.log.errorf("llmgateway: users file %q reload failed: %v", a.usersFile.path, err)
|
||||
return
|
||||
}
|
||||
|
||||
a.mu.RLock()
|
||||
prevCount := a.fileUserCount
|
||||
a.mu.RUnlock()
|
||||
|
||||
if err := a.replaceFileUsers(users); err != nil {
|
||||
a.log.errorf("llmgateway: users file %q reload rejected: %v", a.usersFile.path, err)
|
||||
return
|
||||
}
|
||||
a.lastModTime = info.ModTime()
|
||||
|
||||
if len(users) != prevCount {
|
||||
a.log.logf("llmgateway: users file %q reload: user count changed from %d to %d", a.usersFile.path, prevCount, len(users))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,310 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"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)
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user