diff --git a/auth.go b/auth.go index 3a5dc5c..5710ada 100644 --- a/auth.go +++ b/auth.go @@ -118,6 +118,9 @@ func newAuthStore(cfg *Config) (*authStore, error) { nowFn: time.Now, } for name, gc := range cfg.Groups { + if err := gc.Limits.validate(); err != nil { + return nil, fmt.Errorf("llmgateway: group %q: %w", name, err) + } a.groups[name] = &group{ limits: gc.Limits, name: name, @@ -155,6 +158,9 @@ func (a *authStore) buildEntry(uc *UserConfig) (*authEntry, error) { if !ok { return nil, fmt.Errorf("llmgateway: user %q references unknown group %q", uc.Name, uc.Group) } + if err := uc.Limits.validate(); err != nil { + return nil, fmt.Errorf("llmgateway: user %q: %w", uc.Name, err) + } key, err := resolveSecret(uc.APIKey) if err != nil { return nil, fmt.Errorf("llmgateway: user %q: %w", uc.Name, err) diff --git a/limits.go b/limits.go index d7ae3ed..c2a7144 100644 --- a/limits.go +++ b/limits.go @@ -7,6 +7,12 @@ import ( "time" ) +// storeErrorLogEvery bounds how often the limiter logs a configured +// store's operation error: once per this interval, regardless of how many +// operations fail in between. A backend outage under load must not flood +// the log with one line per request. +const storeErrorLogEvery = 30 * time.Second + // Window identifiers used throughout windowKey, counter metrics, and // retry-after calculation. const ( @@ -47,6 +53,52 @@ func usdToMicros(usd float64) int64 { return int64(math.Round(usd * usdToMicroFactor)) } +// validate reports an error if lc holds a negative int64 count limit, or a +// cost limit that is negative, NaN, or ±Inf. A nil lc (no limits +// configured) is valid. Without this check, a negative int64 limit would +// silently behave as unlimited — every enforcement path below treats +// limit<=0 as "no limit" — and a NaN/Inf cost limit would propagate +// through usdToMicros into meaningless comparisons instead of failing +// loudly at construction time. Called by newAuthStore for every group's +// and every user's LimitsConfig, so a malformed limit is a constructor +// error rather than a silent misconfiguration. +func (lc *LimitsConfig) validate() error { + if lc == nil { + return nil + } + counts := []struct { + name string + v int64 + }{ + {"requestsPerMinute", lc.RequestsPerMinute}, + {"requestsPerDay", lc.RequestsPerDay}, + {"tokensPerDay", lc.TokensPerDay}, + {"tokensPerMonth", lc.TokensPerMonth}, + } + for _, c := range counts { + if c.v < 0 { + return fmt.Errorf("llmgateway: limits: %s must not be negative, got %d", c.name, c.v) + } + } + + costs := []struct { + name string + v float64 + }{ + {"costPerDayUSD", lc.CostPerDayUSD}, + {"costPerMonthUSD", lc.CostPerMonthUSD}, + } + for _, c := range costs { + if math.IsNaN(c.v) || math.IsInf(c.v, 0) { + return fmt.Errorf("llmgateway: limits: %s must be a finite number, got %v", c.name, c.v) + } + if c.v < 0 { + return fmt.Errorf("llmgateway: limits: %s must not be negative, got %v", c.name, c.v) + } + } + return nil +} + // counterStore is the storage backend the limiter uses for atomic windowed // counters. memoryStore (below) is the in-process fallback; a distributed // (e.g. Redis-backed) implementation is wired in a later task. @@ -171,33 +223,41 @@ type limitScope struct { type limitViolation struct { message string // retryAfter is the number of seconds until the violated window ends - // — until the client can plausibly succeed again. + // — until the client can plausibly succeed again. Left at its zero + // value for a storeDown violation, which carries no meaningful window. retryAfter int // storeDown reports whether the violation was manufactured because the - // configured counterStore was unreachable, rather than an actual limit - // breach. Always false in this task; a later task wires failOpen / - // fail-closed handling for a real remote store. + // configured counterStore was unreachable and failOpen is false, + // rather than an actual limit breach. Callers map this to an HTTP 503 + // instead of 429 (Task 12). storeDown bool } // limiter enforces per-minute/day/month request, token, and cost limits // using fixed windows keyed by windowKey. type limiter struct { - store counterStore // nil means always use fallback (see newLimiter) - fallback *memoryStore // in-process counter store, always available - nowFn func() time.Time // injected for tests; defaults to time.Now - failOpen bool // wired for a later task's store-down handling + store counterStore // configured backend; nil means always use fallback (see newLimiter) + fallback *memoryStore // in-process counter store, always available + nowFn func() time.Time // injected for tests; defaults to time.Now + logf func(format string, args ...any) // injectable store-error log; defaults to a no-op + lastLogAt time.Time // guarded by logMu; last time a store error was logged + logMu sync.Mutex + failOpen bool // store-error policy: true falls back to fallback, false refuses the request } // newLimiter returns a limiter. A nil store means every operation uses the // limiter's own in-process fallback memoryStore — the plugin still // enforces limits with no distributed backend configured, just without -// sharing counters across Traefik instances. +// sharing counters across Traefik instances. failOpen governs what happens +// when a non-nil store errors: true silently falls back to the in-process +// memoryStore for that operation, false refuses the request (see +// storeIncrBy / storeGet). func newLimiter(store counterStore, failOpen bool) *limiter { return &limiter{ store: store, fallback: newMemoryStore(), nowFn: time.Now, + logf: func(string, ...any) {}, failOpen: failOpen, } } @@ -207,13 +267,71 @@ func (l *limiter) now() time.Time { return l.nowFn() } -// activeStore returns the store operations should use: the configured -// store when set, otherwise the in-process fallback. -func (l *limiter) activeStore() counterStore { - if l.store != nil { - return l.store +// logStoreError logs err via l.logf, rate-limited to once per +// storeErrorLogEvery regardless of how many store operations fail in +// between. +func (l *limiter) logStoreError(err error) { + now := l.now() + l.logMu.Lock() + shouldLog := now.Sub(l.lastLogAt) >= storeErrorLogEvery + if shouldLog { + l.lastLogAt = now } - return l.fallback + l.logMu.Unlock() + if shouldLog { + l.logf("limit store error: %v", err) + } +} + +// storeIncrBy increments key by n with the given ttl, applying the +// limiter's fail-open/fail-closed policy when a configured store errors. +// ok is false only when the store errored and failOpen is false — the +// fail-closed case a caller must refuse the request for, rather than +// treating a zero value as a real counter reading. +func (l *limiter) storeIncrBy(key string, n int64, ttl time.Duration) (v int64, ok bool) { + if l.store == nil { + v, _ = l.fallback.incrBy(key, n, ttl) // fallback never errors + return v, true + } + v, err := l.store.incrBy(key, n, ttl) + if err == nil { + return v, true + } + l.logStoreError(err) + if !l.failOpen { + return 0, false + } + v, _ = l.fallback.incrBy(key, n, ttl) + return v, true +} + +// storeGet mirrors storeIncrBy for a read: it applies the same fail-open +// (transparently reads the fallback) / fail-closed (ok=false) policy on a +// configured store's error. +func (l *limiter) storeGet(key string) (v int64, ok bool) { + if l.store == nil { + v, _ = l.fallback.get(key) + return v, true + } + v, err := l.store.get(key) + if err == nil { + return v, true + } + l.logStoreError(err) + if !l.failOpen { + return 0, false + } + v, _ = l.fallback.get(key) + return v, true +} + +// storeDownViolation is the violation checkAndCount/budgetViolation return +// when the configured store is unreachable and failOpen is false: the +// request is refused instead of silently enforcing limits against a +// non-shared fallback, so a backend outage cannot let every configured +// limit go unenforced across a fleet of gateway instances. +func storeDownViolation() *limitViolation { + return &limitViolation{message: "limit store unavailable", storeDown: true} } // windowKey builds the counterStore key for one (kind, id, metric, window) @@ -265,20 +383,19 @@ func retryAfterSeconds(t time.Time, window string) int { } // incrCounter increments the (kind, id, metric, window) counter at t by n -// and returns its new value. A store error is treated as "no increment -// observed" — the fallback memoryStore never errors, and a real store's -// fail-open/fail-closed policy is wired in a later task. -func (l *limiter) incrCounter(kind, id, metric, window string, t time.Time, n int64, ttl time.Duration) int64 { - v, err := l.activeStore().incrBy(windowKey(kind, id, metric, window, t), n, ttl) - if err != nil { - return 0 - } - return v +// and returns its new value, together with whether the operation +// succeeded under the limiter's fail-open/fail-closed policy (see +// storeIncrBy). ok is false only in the fail-closed case: a configured +// store errored and failOpen is false. +func (l *limiter) incrCounter(kind, id, metric, window string, t time.Time, n int64, ttl time.Duration) (int64, bool) { + return l.storeIncrBy(windowKey(kind, id, metric, window, t), n, ttl) } -// getCounter returns the (kind, id, metric, window) counter's value at t. -func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int64, error) { - return l.activeStore().get(windowKey(kind, id, metric, window, t)) +// getCounter returns the (kind, id, metric, window) counter's value at t, +// together with whether the read succeeded under the limiter's fail-open/ +// fail-closed policy (see storeGet). +func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int64, bool) { + return l.storeGet(windowKey(kind, id, metric, window, t)) } // checkAndCount increments every scope's req:min and req:day counters — @@ -286,7 +403,10 @@ func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int6 // gets refused by one scope's limit still counts toward every other // scope's rate tracking — then evaluates each scope's set limits in order // and returns the first violation found, or nil if the request may -// proceed. +// proceed. If any store operation fails closed (a configured store errored +// and failOpen is false), it returns a storeDownViolation immediately — +// the request is refused rather than evaluated against partial or +// fallback-only counters. // // Requests limits compare against the value just incremented in this call. // Token and cost limits compare against the value already accumulated by @@ -299,8 +419,17 @@ func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation { minCounts := make([]int64, len(scopes)) dayCounts := make([]int64, len(scopes)) for i, sc := range scopes { - minCounts[i] = l.incrCounter(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL) - dayCounts[i] = l.incrCounter(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL) + v, ok := l.incrCounter(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL) + if !ok { + return storeDownViolation() + } + minCounts[i] = v + + v, ok = l.incrCounter(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL) + if !ok { + return storeDownViolation() + } + dayCounts[i] = v } for i, sc := range scopes { @@ -357,15 +486,19 @@ func requestLimitViolation(sc limitScope, name string, limit, count int64, windo // budgetViolation reports a violation when the scope's accumulated // counter for metric/window is already at or above limit. limit<=0 means -// unlimited. A store read error fails open — no violation reported — -// consistent with this task's storeDown always being false; a later task -// wires strict fail-closed handling for a real remote store. +// unlimited. A failed-closed read (configured store errored, failOpen +// false) returns storeDownViolation rather than silently passing the +// check — a store outage must never look the same as staying under +// budget. func (l *limiter) budgetViolation(sc limitScope, metric, name string, limit int64, window string, now time.Time) *limitViolation { if limit <= 0 { return nil } - used, err := l.getCounter(sc.kind, sc.id, metric, window, now) - if err != nil || used < limit { + used, ok := l.getCounter(sc.kind, sc.id, metric, window, now) + if !ok { + return storeDownViolation() + } + if used < limit { return nil } return &limitViolation{ @@ -376,7 +509,11 @@ func (l *limiter) budgetViolation(sc limitScope, metric, name string, limit int6 // account records u's total tokens and costMicros against every scope's // day and month counters. A zero-valued metric is skipped entirely — no -// store write for a metric this call has nothing to report. +// store write for a metric this call has nothing to report. A sample that +// hits a fail-closed store error (see storeIncrBy) is dropped silently: +// account has no error return to signal it, and dropping — rather than +// counting it in the fallback — avoids double-counting once the store +// recovers. func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) { now := l.now() total := u.total() diff --git a/limits_test.go b/limits_test.go index 69ccd5d..b8d4226 100644 --- a/limits_test.go +++ b/limits_test.go @@ -1,6 +1,9 @@ package traefikllmgateway import ( + "errors" + "fmt" + "math" "testing" "time" ) @@ -308,12 +311,233 @@ func TestAccountSkipsZeroMetrics(t *testing.T) { l.account(scopes, usage{}, 0) - tok, err := l.getCounter("user", "z", metricTok, windowDay, now) - if err != nil || tok != 0 { - t.Errorf("tok counter = %d, %v, want 0, nil", tok, err) + tok, ok := l.getCounter("user", "z", metricTok, windowDay, now) + if !ok || tok != 0 { + t.Errorf("tok counter = %d, ok=%v, want 0, true", tok, ok) } - cost, err := l.getCounter("user", "z", metricCost, windowDay, now) - if err != nil || cost != 0 { - t.Errorf("cost counter = %d, %v, want 0, nil", cost, err) + cost, ok := l.getCounter("user", "z", metricCost, windowDay, now) + if !ok || cost != 0 { + t.Errorf("cost counter = %d, ok=%v, want 0, true", cost, ok) + } +} + +// TestLimiterRequestPerDayLimit is carried-item (b): a per-day request +// limit blocks once the day's count exceeds it, independent of the +// (unset, unlimited) per-minute limit. +func TestLimiterRequestPerDayLimit(t *testing.T) { + now := time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC) + l := newLimiter(nil, true) + l.nowFn = func() time.Time { return now } + scopes := []limitScope{{kind: "user", id: "a", limits: &LimitsConfig{RequestsPerDay: 2}}} + + if v := l.checkAndCount(scopes); v != nil { + t.Fatal("1st should pass") + } + if v := l.checkAndCount(scopes); v != nil { + t.Fatal("2nd should pass") + } + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("3rd should violate the day limit") + } + if v.storeDown { + t.Error("storeDown must be false for a real limit breach") + } +} + +// TestLimiterTokensPerMonthBudget mirrors TestLimiterTokenBudget for the +// month window: accounting tokens past a month budget makes the next +// check refuse. +func TestLimiterTokensPerMonthBudget(t *testing.T) { + l := newLimiter(nil, true) + scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerMonth: 100}}} + l.account(scopes, usage{prompt: 60, completion: 50}, 0) // 110 > 100 + if v := l.checkAndCount(scopes); v == nil { + t.Fatal("over month token budget must refuse") + } +} + +// TestLimiterCostPerMonthBudget mirrors TestLimiterCostBudget for the +// month window, and checks the violation carries the month window's +// Retry-After. +func TestLimiterCostPerMonthBudget(t *testing.T) { + now := time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC) + l := newLimiter(nil, true) + l.nowFn = func() time.Time { return now } + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{CostPerMonthUSD: 1.00}}} + + l.account(scopes, usage{}, 1_500_000) // $1.50 > $1.00 budget + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("over month cost budget must refuse") + } + wantRetry := int(time.Date(2026, 9, 1, 0, 0, 0, 0, time.UTC).Sub(now).Seconds()) + if v.retryAfter != wantRetry { + t.Errorf("retryAfter = %d, want %d (seconds to next UTC month start)", v.retryAfter, wantRetry) + } + if v.storeDown { + t.Error("storeDown must be false for a real limit breach") + } +} + +// erroringStore is a counterStore stub that always fails, used to test the +// limiter's fail-open/fail-closed handling on a store error without a +// real network dependency (see redis_store_test.go for the equivalent +// end-to-end test against an unreachable redisStore). +type erroringStore struct { + err error +} + +func (s *erroringStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, s.err } +func (s *erroringStore) get(string) (int64, error) { return 0, s.err } + +// TestLimiter_FailOpen_StoreErrorUsesFallback is carried-item (c): a store +// error with failOpen=true must not block the request — the operation +// transparently falls back to the limiter's in-process memoryStore, and +// the fallback counter actually advances (proving the write really +// happened somewhere, not just that no violation was reported). +func TestLimiter_FailOpen_StoreErrorUsesFallback(t *testing.T) { + store := &erroringStore{err: errors.New("boom")} + l := newLimiter(store, true) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}} + + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("want no violation with failOpen=true on a store error, got %+v", v) + } + got, err := l.fallback.get(windowKey("user", "u", metricReq, windowDay, l.now())) + if err != nil { + t.Fatalf("fallback.get: %v", err) + } + if got != 1 { + t.Errorf("fallback req:day counter = %d, want 1", got) + } +} + +// TestLimiter_FailClosed_StoreErrorReturnsStoreDownViolation is +// carried-item (c): failOpen=false refuses the request outright on a +// store error, rather than silently using the fallback. +func TestLimiter_FailClosed_StoreErrorReturnsStoreDownViolation(t *testing.T) { + store := &erroringStore{err: errors.New("boom")} + l := newLimiter(store, false) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}} + + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("want a storeDown violation on a store error with failOpen=false") + } + if !v.storeDown { + t.Error("v.storeDown = false, want true") + } + if v.message != "limit store unavailable" { + t.Errorf("v.message = %q, want %q", v.message, "limit store unavailable") + } +} + +// TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation covers the +// budgetViolation fail-closed path specifically (as opposed to the +// request-counter incrCounter path covered above): a store error on a +// token/cost budget read also refuses the request when failOpen is false. +func TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation(t *testing.T) { + store := &erroringStore{err: errors.New("boom")} + l := newLimiter(store, false) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{TokensPerDay: 100}}} + + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("want a storeDown violation on a store error with failOpen=false") + } + if !v.storeDown { + t.Error("v.storeDown = false, want true") + } +} + +// TestLimiter_FailClosed_AccountDropsSampleOnStoreError asserts account +// drops a sample entirely — writes it nowhere — when failOpen=false and +// the store errors, rather than falling back to counting it locally (which +// would double-count once the store recovers). account has no error +// return to signal the drop, so this is observed via the fallback staying +// at 0. +func TestLimiter_FailClosed_AccountDropsSampleOnStoreError(t *testing.T) { + store := &erroringStore{err: errors.New("boom")} + l := newLimiter(store, false) + scopes := []limitScope{{kind: "user", id: "u", limits: nil}} + + l.account(scopes, usage{prompt: 10, completion: 10}, 500) + + got, err := l.fallback.get(windowKey("user", "u", metricTok, windowDay, l.now())) + if err != nil { + t.Fatalf("fallback.get: %v", err) + } + if got != 0 { + t.Errorf("fallback tok:day counter = %d, want 0 (sample must be dropped, not counted locally)", got) + } +} + +// TestLimiter_LogsStoreErrorOncePerRateLimit asserts a store error is +// logged via the limiter's injectable log func, rate-limited to once per +// storeErrorLogEvery — a store outage under load must not flood the log +// with one line per request, however many store operations fail within +// the window. +func TestLimiter_LogsStoreErrorOncePerRateLimit(t *testing.T) { + store := &erroringStore{err: errors.New("boom")} + l := newLimiter(store, true) + now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC) + l.nowFn = func() time.Time { return now } + + var logged []string + l.logf = func(format string, args ...any) { logged = append(logged, fmt.Sprintf(format, args...)) } + + scopes := []limitScope{{kind: "user", id: "u", limits: nil}} + + l.checkAndCount(scopes) // two store ops (min, day) within one call + if len(logged) != 1 { + t.Fatalf("logged = %d lines after the first checkAndCount, want 1 (rate-limited)", len(logged)) + } + + l.checkAndCount(scopes) // still inside the rate-limit window + if len(logged) != 1 { + t.Fatalf("logged = %d lines after a second call inside the window, want 1", len(logged)) + } + + now = now.Add(storeErrorLogEvery) + l.checkAndCount(scopes) + if len(logged) != 2 { + t.Fatalf("logged = %d lines after the rate-limit window elapsed, want 2", len(logged)) + } +} + +// TestLimitsConfigValidate is carried-item (a): validate rejects +// NaN/±Inf/negative cost limits and negative int64 count limits, accepts +// zero/positive values, and treats a nil receiver (no limits configured +// at all) as valid. +func TestLimitsConfigValidate(t *testing.T) { + cases := []struct { + lc *LimitsConfig + name string + wantErr bool + }{ + {name: "nil is valid", lc: nil, wantErr: false}, + {name: "zero value is valid", lc: &LimitsConfig{}, wantErr: false}, + {name: "positive values are valid", lc: &LimitsConfig{ + RequestsPerMinute: 1, RequestsPerDay: 1, TokensPerDay: 1, TokensPerMonth: 1, + CostPerDayUSD: 1.5, CostPerMonthUSD: 1.5, + }, wantErr: false}, + {name: "negative RequestsPerMinute", lc: &LimitsConfig{RequestsPerMinute: -1}, wantErr: true}, + {name: "negative RequestsPerDay", lc: &LimitsConfig{RequestsPerDay: -1}, wantErr: true}, + {name: "negative TokensPerDay", lc: &LimitsConfig{TokensPerDay: -1}, wantErr: true}, + {name: "negative TokensPerMonth", lc: &LimitsConfig{TokensPerMonth: -1}, wantErr: true}, + {name: "negative CostPerDayUSD", lc: &LimitsConfig{CostPerDayUSD: -0.01}, wantErr: true}, + {name: "negative CostPerMonthUSD", lc: &LimitsConfig{CostPerMonthUSD: -0.01}, wantErr: true}, + {name: "NaN CostPerDayUSD", lc: &LimitsConfig{CostPerDayUSD: math.NaN()}, wantErr: true}, + {name: "+Inf CostPerMonthUSD", lc: &LimitsConfig{CostPerMonthUSD: math.Inf(1)}, wantErr: true}, + {name: "-Inf CostPerDayUSD", lc: &LimitsConfig{CostPerDayUSD: math.Inf(-1)}, wantErr: true}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + err := c.lc.validate() + if (err != nil) != c.wantErr { + t.Errorf("validate() error = %v, wantErr %v", err, c.wantErr) + } + }) } } diff --git a/llmgateway.go b/llmgateway.go index fac0c5b..e1b18e4 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -101,12 +101,13 @@ func CreateConfig() *Config { // Gateway is the Traefik middleware handler. type Gateway struct { - next http.Handler - cfg *Config - auth *authStore - name string - // wired in later tasks: limiter *limiter, registry *modelRegistry, - // adapters map[string]providerAdapter + next http.Handler + cfg *Config + auth *authStore + limiter *limiter + name string + // wired in later tasks: registry *modelRegistry, adapters + // map[string]providerAdapter } // New creates the middleware. NOTE: no tail call — Yaegi zeroes @@ -130,14 +131,58 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str 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 { + if err = attachUsersFile(auth, config.Users.File, g); err != nil { return nil, err } } + lim, err := newConfiguredLimiter(config) + if err != nil { + return nil, err + } + lim.logf = g.errorf + g.limiter = lim + + setPricingWarnFn(func(msg string) { + if msg == warnCapMessage { + g.logf("%s", msg) + return + } + g.logf("pricing: no price configured for model %q; cost will be recorded as 0", msg) + }) + return g, nil } +// newConfiguredLimiter builds the limiter for config, wiring config.Redis +// when set. A configured Redis block wires a respClient-backed redisStore +// behind AUTH/SELECT and resolveSecret'd credentials; otherwise the +// limiter runs on its in-process fallback alone. failOpen defaults to +// true (a Redis outage must not take the whole gateway down unless an +// operator opts into strict enforcement) unless RedisConfig.FailOpen is +// explicitly set. +func newConfiguredLimiter(config *Config) (*limiter, error) { + if config.Redis == nil { + return newLimiter(nil, true), nil + } + if config.Redis.Address == "" { + return nil, errors.New("llmgateway: redis: address must not be empty") + } + + password, err := resolveSecret(config.Redis.Password) + if err != nil { + return nil, fmt.Errorf("llmgateway: redis: %w", err) + } + + failOpen := true + if config.Redis.FailOpen != nil { + failOpen = *config.Redis.FailOpen + } + + client := newRESPClient(config.Redis.Address, password, config.Redis.DB) + return newLimiter(newRedisStore(client), failOpen), 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 diff --git a/llmgateway_test.go b/llmgateway_test.go index 2435fbe..0076295 100644 --- a/llmgateway_test.go +++ b/llmgateway_test.go @@ -416,3 +416,222 @@ func TestServeHTTP_TriggersUsersFileReload(t *testing.T) { t.Fatal("want ServeHTTP's entry-time maybeReload call to have picked up the file change") } } + +// TestNewGateway_RejectsInvalidGroupLimits is carried-item (a): a negative +// limit value in a group's LimitsConfig must fail plugin construction, not +// silently behave as unlimited. +func TestNewGateway_RejectsInvalidGroupLimits(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Groups = map[string]*GroupConfig{"default": {Limits: &LimitsConfig{RequestsPerDay: -1}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + if _, err := New(context.Background(), next, cfg, "llmgw"); err == nil { + t.Fatal("want constructor error for a negative group limit") + } +} + +// TestNewGateway_RejectsInvalidUserLimits mirrors the group case for an +// inline user's own limits override. +func TestNewGateway_RejectsInvalidUserLimits(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{ + {Name: "a", Group: "default", APIKey: "sk-a", Limits: &LimitsConfig{CostPerMonthUSD: -1}}, + }} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + if _, err := New(context.Background(), next, cfg, "llmgw"); err == nil { + t.Fatal("want constructor error for a negative user limit") + } +} + +// TestNewGateway_RejectsInvalidFileUserLimits covers the file-sourced user +// path (attachUsersFile -> replaceFileUsers -> buildEntry), the other +// caller of LimitsConfig.validate besides newAuthStore's inline loop. +func TestNewGateway_RejectsInvalidFileUserLimits(t *testing.T) { + dir := t.TempDir() + fp := filepath.Join(dir, "users.json") + writeUsersDoc(t, fp, []*UserConfig{ + {Name: "f", Group: "default", APIKey: "sk-f", Limits: &LimitsConfig{TokensPerDay: -1}}, + }, 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) {}) + if _, err := New(context.Background(), next, cfg, "llmgw"); err == nil { + t.Fatal("want constructor error for a negative file-sourced user limit") + } +} + +// TestNewGateway_NoRedis_UsesMemoryLimiterFailOpenDefaultTrue is +// carried-item (d)'s "else" branch: no Redis configured wires +// newLimiter(nil, true) — a nil store (memoryStore fallback only) and +// failOpen defaulted true. +func TestNewGateway_NoRedis_UsesMemoryLimiterFailOpenDefaultTrue(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + 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 := h.(*Gateway) + if gw.limiter == nil { + t.Fatal("want a non-nil limiter") + } + if gw.limiter.store != nil { + t.Errorf("gw.limiter.store = %#v, want nil (no Redis configured)", gw.limiter.store) + } + if !gw.limiter.failOpen { + t.Error("gw.limiter.failOpen = false, want true (default)") + } +} + +// TestNewGateway_RedisConfigured_WiresRedisStoreAndFailOpen is +// carried-item (d)'s main branch: a configured Redis block wires a +// redisStore (backed by a respClient) into the limiter, and an explicit +// FailOpen=false is honored rather than overridden by the default. +func TestNewGateway_RedisConfigured_WiresRedisStoreAndFailOpen(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + failOpen := false + cfg.Redis = &RedisConfig{Address: "127.0.0.1:0", FailOpen: &failOpen} + 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 := h.(*Gateway) + if gw.limiter == nil { + t.Fatal("want a non-nil limiter") + } + if _, ok := gw.limiter.store.(*redisStore); !ok { + t.Errorf("gw.limiter.store = %#v (%T), want *redisStore", gw.limiter.store, gw.limiter.store) + } + if gw.limiter.failOpen { + t.Error("gw.limiter.failOpen = true, want false (explicit config)") + } +} + +// TestNewGateway_RedisFailOpenNilDefaultsTrue asserts a configured Redis +// block with FailOpen left nil defaults to true, matching the no-Redis +// default. +func TestNewGateway_RedisFailOpenNilDefaultsTrue(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Redis = &RedisConfig{Address: "127.0.0.1:0"} + 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 := h.(*Gateway) + if !gw.limiter.failOpen { + t.Error("gw.limiter.failOpen = false, want true (FailOpen left nil)") + } +} + +// TestNewGateway_RedisConfigured_EmptyAddress_ReturnsConstructorError +// asserts a Redis block with no address fails construction immediately, +// rather than deferring the error to the first lazy-connect attempt on +// the request path. +func TestNewGateway_RedisConfigured_EmptyAddress_ReturnsConstructorError(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Redis = &RedisConfig{} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + if _, err := New(context.Background(), next, cfg, "llmgw"); err == nil { + t.Fatal("want constructor error for a Redis block with an empty address") + } +} + +// TestNewGateway_RedisPassword_ResolvedViaSecret is an end-to-end proof +// that newGateway resolves Config.Redis.Password through resolveSecret +// (env:/file:/literal) before handing it to respClient: an "env:" password +// must arrive at the fake server as its resolved value, not the literal +// "env:..." string. +func TestNewGateway_RedisPassword_ResolvedViaSecret(t *testing.T) { + t.Setenv("LLMGW_TEST_REDIS_PASSWORD", "resolved-pw") + + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"AUTH", "resolved-pw"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, reply: []byte("$-1\r\n")}, + }) + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Redis = &RedisConfig{Address: ln.Addr().String(), Password: "env:LLMGW_TEST_REDIS_PASSWORD"} // #nosec G101 -- not a credential, a resolveSecret "env:" reference to an env var name + 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 := h.(*Gateway) + store, ok := gw.limiter.store.(*redisStore) + if !ok { + t.Fatalf("gw.limiter.store = %#v (%T), want *redisStore", gw.limiter.store, gw.limiter.store) + } + if _, err := store.get("k"); err != nil { + t.Fatalf("store.get: %v (want the resolved password to authenticate against the fake server)", err) + } +} + +// TestNewGateway_RedisPassword_UnresolvableSecret_ReturnsConstructorError +// asserts an "env:" password referencing an unset variable fails +// construction with resolveSecret's error, rather than deferring to a +// confusing AUTH failure at request time. +func TestNewGateway_RedisPassword_UnresolvableSecret_ReturnsConstructorError(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Redis = &RedisConfig{Address: "127.0.0.1:0", Password: "env:LLMGW_TEST_REDIS_PASSWORD_UNSET"} // #nosec G101 -- not a credential, a resolveSecret "env:" reference to an env var name + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + if _, err := New(context.Background(), next, cfg, "llmgw"); err == nil { + t.Fatal("want constructor error for an unresolvable Redis password secret") + } +} + +// TestNewGateway_SetsPricingWarnFn is carried-item (d)'s pricing wiring: +// New must call setPricingWarnFn so an unpriced model's warning reaches +// the gateway's own log, not pricing.go's default no-op. +func TestNewGateway_SetsPricingWarnFn(t *testing.T) { + warnedModelsMu.Lock() + prevWarned, prevCapNotified := warnedModels, warnCapNotified + warnedModels = map[string]bool{} + warnCapNotified = false + warnedModelsMu.Unlock() + t.Cleanup(func() { + warnedModelsMu.Lock() + warnedModels, warnCapNotified = prevWarned, prevCapNotified + warnedModelsMu.Unlock() + setPricingWarnFn(func(string) {}) + }) + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "mygw") + if err != nil { + t.Fatalf("New: %v", err) + } + _ = h.(*Gateway) + + origStderr := os.Stderr + r, w, _ := os.Pipe() + os.Stderr = w + + costMicros("totally-unpriced-model-for-newgateway-test", usage{prompt: 1}, nil) + + _ = w.Close() // closing the pipe write end to unblock the read; error not actionable in a test + os.Stderr = origStderr + var buf bytes.Buffer + if _, err := io.Copy(&buf, r); err != nil { + t.Fatalf("io.Copy: %v", err) + } + if !strings.Contains(buf.String(), "llmgw[mygw]") { + t.Fatalf("want the pricing warning routed through the gateway's own logf, got %q", buf.String()) + } +} diff --git a/redis_store.go b/redis_store.go new file mode 100644 index 0000000..c1b6866 --- /dev/null +++ b/redis_store.go @@ -0,0 +1,81 @@ +package traefikllmgateway + +import ( + "fmt" + "strconv" + "time" +) + +// redisStore is a counterStore backed by a Redis-compatible server over +// respClient, so multiple Traefik instances share the same limit counters +// instead of each enforcing limits against its own in-process state. +type redisStore struct { + client *respClient +} + +// newRedisStore returns a redisStore using client for storage. +func newRedisStore(client *respClient) *redisStore { + return &redisStore{client: client} +} + +// incrBy implements counterStore: INCRBY key n, then EXPIRE key +// ttlSeconds, sent as one pipeline so the two commands share a single round +// trip. It returns INCRBY's resulting counter value; EXPIRE's reply is not +// otherwise inspected — a failed EXPIRE right after a successful INCRBY on +// the same key would only leave that key without a fresh TTL, not corrupt +// the count. ttl is rounded up to whole seconds (Redis EXPIRE's unit), +// with a floor of 1s so a sub-second ttl never turns into EXPIRE 0 (an +// immediate delete). +func (s *redisStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) { + ttlSeconds := int64(ttl / time.Second) + if ttl%time.Second != 0 { + ttlSeconds++ + } + if ttlSeconds < 1 { + ttlSeconds = 1 + } + + replies, err := s.client.pipeline([][]string{ + {"INCRBY", key, strconv.FormatInt(n, 10)}, + {"EXPIRE", key, strconv.FormatInt(ttlSeconds, 10)}, + }) + if err != nil { + return 0, fmt.Errorf("redisStore: incrBy %q: %w", key, err) + } + if len(replies) == 0 { + return 0, fmt.Errorf("redisStore: incrBy %q: empty pipeline reply", key) + } + if e, ok := replies[0].(error); ok { + return 0, fmt.Errorf("redisStore: incrBy %q: INCRBY failed: %w", key, e) + } + v, ok := replies[0].(int64) + if !ok { + return 0, fmt.Errorf("redisStore: incrBy %q: unexpected INCRBY reply type %T", key, replies[0]) + } + return v, nil +} + +// get implements counterStore: GET key, treating a missing key (RESP null +// bulk reply) as 0, matching memoryStore's behavior for an absent or +// expired counter. +func (s *redisStore) get(key string) (int64, error) { + reply, err := s.client.do("GET", key) + if err != nil { + return 0, fmt.Errorf("redisStore: get %q: %w", key, err) + } + if reply == nil { + return 0, nil + } + if e, ok := reply.(error); ok { + return 0, fmt.Errorf("redisStore: get %q: %w", key, e) + } + b, ok := reply.([]byte) + if !ok { + return 0, fmt.Errorf("redisStore: get %q: unexpected reply type %T", key, reply) + } + v, err := strconv.ParseInt(string(b), 10, 64) + if err != nil { + return 0, fmt.Errorf("redisStore: get %q: non-integer value %q: %w", key, b, err) + } + return v, nil +} diff --git a/redis_store_test.go b/redis_store_test.go new file mode 100644 index 0000000..447330b --- /dev/null +++ b/redis_store_test.go @@ -0,0 +1,138 @@ +package traefikllmgateway + +import ( + "testing" + "time" +) + +// TestRedisStore_IncrBy_PipelinesIncrbyAndExpire is the brief's Step-1 +// pipeline case: incrBy sends INCRBY and EXPIRE as one pipeline and +// returns INCRBY's reply value. +func TestRedisStore_IncrBy_PipelinesIncrbyAndExpire(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"INCRBY", "llmgw:user:a:req:day:20260820", "1"}, reply: []byte(":5\r\n")}, + {wantArgs: []string{"EXPIRE", "llmgw:user:a:req:day:20260820", "60"}, reply: []byte(":1\r\n")}, + }) + + store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0)) + v, err := store.incrBy("llmgw:user:a:req:day:20260820", 1, 60*time.Second) + if err != nil { + t.Fatalf("incrBy: %v", err) + } + if v != 5 { + t.Errorf("incrBy = %d, want 5", v) + } +} + +// TestRedisStore_IncrBy_RoundsSubSecondTTLUp asserts a ttl under one +// second is never sent to EXPIRE as 0 (which would delete the key +// immediately) — it is rounded up to a 1s floor. +func TestRedisStore_IncrBy_RoundsSubSecondTTLUp(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"INCRBY", "k", "1"}, reply: []byte(":1\r\n")}, + {wantArgs: []string{"EXPIRE", "k", "1"}, reply: []byte(":1\r\n")}, + }) + + store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0)) + if _, err := store.incrBy("k", 1, 10*time.Millisecond); err != nil { + t.Fatalf("incrBy: %v", err) + } +} + +// TestRedisStore_Get_MissingKeyReturnsZero is the brief's Step-1 GET case: +// a null bulk reply ("$-1") decodes to 0 with no error. +func TestRedisStore_Get_MissingKeyReturnsZero(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "llmgw:user:a:req:day:20260820"}, reply: []byte("$-1\r\n")}, + }) + + store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0)) + v, err := store.get("llmgw:user:a:req:day:20260820") + if err != nil { + t.Fatalf("get: %v", err) + } + if v != 0 { + t.Errorf("get = %d, want 0", v) + } +} + +// TestRedisStore_Get_ParsesIntegerValue asserts a present key's bulk-string +// value is parsed as a base-10 integer. +func TestRedisStore_Get_ParsesIntegerValue(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, reply: []byte("$3\r\n123\r\n")}, + }) + + store := newRedisStore(newRESPClient(ln.Addr().String(), "", 0)) + v, err := store.get("k") + if err != nil { + t.Fatalf("get: %v", err) + } + if v != 123 { + t.Errorf("get = %d, want 123", v) + } +} + +// TestLimiter_FailOpenFalse_DeadRedisAddress_ReturnsStoreDownViolation is +// the brief's Step-1 case: a limiter backed by a redisStore pointed at an +// address nothing listens on, with failOpen=false, must refuse the +// request with a storeDown violation rather than silently falling back to +// an in-process counter or panicking. +func TestLimiter_FailOpenFalse_DeadRedisAddress_ReturnsStoreDownViolation(t *testing.T) { + ln := newFakeListener(t) + deadAddr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatalf("Close: %v", err) + } // nothing listens at deadAddr from here on + + store := newRedisStore(newRESPClient(deadAddr, "", 0)) + l := newLimiter(store, false) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}} + + v := l.checkAndCount(scopes) + if v == nil { + t.Fatal("want a violation when the store is unreachable and failOpen is false") + } + if !v.storeDown { + t.Errorf("v.storeDown = false, want true") + } + if v.message != "limit store unavailable" { + t.Errorf("v.message = %q, want %q", v.message, "limit store unavailable") + } +} + +// TestLimiter_FailOpenTrue_DeadRedisAddress_FallsBackAndAllows mirrors the +// fail-closed case above with failOpen=true: the same dead address must +// not block the request — the limiter transparently counts against its +// in-process fallback instead. +func TestLimiter_FailOpenTrue_DeadRedisAddress_FallsBackAndAllows(t *testing.T) { + ln := newFakeListener(t) + deadAddr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + store := newRedisStore(newRESPClient(deadAddr, "", 0)) + l := newLimiter(store, true) + scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}} + + if v := l.checkAndCount(scopes); v != nil { + t.Fatalf("want no violation with failOpen=true and a dead store, got %+v", v) + } + + got, err := l.fallback.get(windowKey("user", "u", metricReq, windowDay, l.now())) + if err != nil { + t.Fatalf("fallback.get: %v", err) + } + if got != 1 { + t.Errorf("fallback req:day counter = %d, want 1 (the request must have been counted in the fallback)", got) + } +} diff --git a/resp.go b/resp.go new file mode 100644 index 0000000..5ec5b15 --- /dev/null +++ b/resp.go @@ -0,0 +1,292 @@ +package traefikllmgateway + +import ( + "bufio" + "fmt" + "io" + "net" + "strconv" + "strings" + "sync" + "time" +) + +// respDialTimeout, respReadTimeout, and respWriteTimeout bound every +// network operation respClient performs: connecting, writing a pipeline, +// and reading each reply. A hung or unreachable Redis-compatible server +// must never block a request indefinitely — the limiter's fail-open/ +// fail-closed handling (limits.go) depends on a store call returning +// within a bounded time. +const ( + respDialTimeout = 2 * time.Second + respReadTimeout = 2 * time.Second + respWriteTimeout = 2 * time.Second +) + +// respErr is a RESP2 error reply ("-message\r\n"). It implements error so a +// caller inspecting a decoded reply can type-assert for it the same way as +// any other error. +type respErr string + +// Error implements the error interface. +func (e respErr) Error() string { return string(e) } + +// respClient is a minimal stdlib-only RESP2 client for a single Redis- +// compatible server: one mutex-guarded TCP connection, lazily dialled on +// first use and reconnected at most once when a command round trip fails. +// It exists because the plugin runs interpreted under Yaegi with only the +// Go standard library available, so a full Redis client library is not an +// option. +// +// mu serializes every operation — do and pipeline hold it for their whole +// round trip, including any reconnect — since RESP is not multiplexed: +// interleaving two callers' writes on one connection would corrupt both +// replies. +type respClient struct { + conn net.Conn + r *bufio.Reader + addr string + password string + mu sync.Mutex + db int +} + +// newRESPClient returns a respClient for addr. It does not connect until +// the first do or pipeline call. password, if non-empty, is sent via AUTH +// on every new connection; db is always sent via SELECT on every new +// connection (including db 0), since a Redis-compatible proxy is not +// guaranteed to default a fresh connection to database 0. +func newRESPClient(addr, password string, db int) *respClient { + return &respClient{addr: addr, password: password, db: db} +} + +// do sends one command and returns its decoded reply. It is equivalent to +// pipeline with a single command. +func (c *respClient) do(args ...string) (any, error) { + replies, err := c.pipeline([][]string{args}) + if err != nil { + return nil, err + } + return replies[0], nil +} + +// pipeline sends every command in cmds on one round trip and returns their +// decoded replies in the same order. On any I/O error — including a +// failure to (re)connect — it closes the current connection, reconnects +// once, and retries the whole pipeline from scratch; a second failure is +// returned to the caller with the connection left closed, so the next call +// lazily reconnects fresh rather than retrying against a connection +// already known bad. +func (c *respClient) pipeline(cmds [][]string) ([]any, error) { + c.mu.Lock() + defer c.mu.Unlock() + + replies, err := c.attemptPipelineLocked(cmds) + if err == nil { + return replies, nil + } + + c.closeLocked() + replies, err = c.attemptPipelineLocked(cmds) + if err != nil { + c.closeLocked() + return nil, err + } + return replies, nil +} + +// attemptPipelineLocked runs one full attempt of cmds over c's connection, +// connecting first if needed. Callers must hold c.mu. +func (c *respClient) attemptPipelineLocked(cmds [][]string) ([]any, error) { + if err := c.ensureConnLocked(); err != nil { + return nil, err + } + + var buf strings.Builder + for _, args := range cmds { + buf.WriteString(encodeCommand(args)) + } + if err := c.conn.SetWriteDeadline(time.Now().Add(respWriteTimeout)); err != nil { + return nil, fmt.Errorf("resp: set write deadline: %w", err) + } + if _, err := io.WriteString(c.conn, buf.String()); err != nil { + return nil, fmt.Errorf("resp: write: %w", err) + } + + replies := make([]any, len(cmds)) + for i := range cmds { + if err := c.conn.SetReadDeadline(time.Now().Add(respReadTimeout)); err != nil { + return nil, fmt.Errorf("resp: set read deadline: %w", err) + } + v, err := decodeReply(c.r) + if err != nil { + return nil, fmt.Errorf("resp: read reply %d/%d: %w", i+1, len(cmds), err) + } + replies[i] = v + } + return replies, nil +} + +// ensureConnLocked dials, authenticates, and selects the database if c has +// no live connection. Callers must hold c.mu. +func (c *respClient) ensureConnLocked() error { + if c.conn != nil { + return nil + } + + conn, err := net.DialTimeout("tcp", c.addr, respDialTimeout) + if err != nil { + return fmt.Errorf("resp: dial %q: %w", c.addr, err) + } + c.conn = conn + c.r = bufio.NewReader(conn) + + if c.password != "" { + if err := c.handshakeLocked("AUTH", c.password); err != nil { + c.closeLocked() + return err + } + } + if err := c.handshakeLocked("SELECT", strconv.Itoa(c.db)); err != nil { + c.closeLocked() + return err + } + return nil +} + +// handshakeLocked sends one connection-setup command (AUTH or SELECT) and +// requires a non-error reply. Callers must hold c.mu and have a live +// c.conn. +func (c *respClient) handshakeLocked(cmd, arg string) error { + if err := c.conn.SetWriteDeadline(time.Now().Add(respWriteTimeout)); err != nil { + return fmt.Errorf("resp: set write deadline: %w", err) + } + if _, err := io.WriteString(c.conn, encodeCommand([]string{cmd, arg})); err != nil { + return fmt.Errorf("resp: %s: write: %w", cmd, err) + } + if err := c.conn.SetReadDeadline(time.Now().Add(respReadTimeout)); err != nil { + return fmt.Errorf("resp: set read deadline: %w", err) + } + v, err := decodeReply(c.r) + if err != nil { + return fmt.Errorf("resp: %s: read reply: %w", cmd, err) + } + if e, ok := v.(respErr); ok { + return fmt.Errorf("resp: %s failed: %w", cmd, e) + } + return nil +} + +// closeLocked closes and clears c's connection, if any. Callers must hold +// c.mu. +func (c *respClient) closeLocked() { + if c.conn == nil { + return + } + _ = c.conn.Close() + c.conn = nil + c.r = nil +} + +// encodeCommand renders args as a RESP2 command array: +// "*N\r\n$len\r\narg\r\n..." for each arg. +func encodeCommand(args []string) string { + var b strings.Builder + fmt.Fprintf(&b, "*%d\r\n", len(args)) + for _, a := range args { + fmt.Fprintf(&b, "$%d\r\n%s\r\n", len(a), a) + } + return b.String() +} + +// decodeReply reads and decodes one RESP2 reply from r: +// +// - "+simple\r\n" -> string +// - "-error\r\n" -> respErr (implements error; not returned as err) +// - ":123\r\n" -> int64 +// - "$len\r\n...\r\n" -> []byte, or nil for a null bulk ("$-1\r\n") +// - "*n\r\n..." -> []any of n decoded elements, or nil for a null +// array ("*-1\r\n") +// +// The returned error is non-nil only for a transport or protocol failure +// (short read, malformed line, unknown type byte) — a well-formed RESP +// error reply decodes successfully to a respErr value, not to a non-nil +// error, so a caller pipelining several commands can still read every +// reply after one of them fails server-side. +func decodeReply(r *bufio.Reader) (any, error) { + line, err := readLine(r) + if err != nil { + return nil, err + } + if line == "" { + return nil, fmt.Errorf("resp: empty reply line") + } + + prefix, rest := line[0], line[1:] + switch prefix { + case '+': + return rest, nil + case '-': + return respErr(rest), nil + case ':': + n, err := strconv.ParseInt(rest, 10, 64) + if err != nil { + return nil, fmt.Errorf("resp: malformed integer reply %q: %w", rest, err) + } + return n, nil + case '$': + return decodeBulk(r, rest) + case '*': + return decodeArray(r, rest) + default: + return nil, fmt.Errorf("resp: unknown reply type %q", prefix) + } +} + +// decodeBulk reads a bulk string body given its already-parsed "$" length +// field lenField, returning nil for a null bulk ("$-1"). +func decodeBulk(r *bufio.Reader, lenField string) (any, error) { + n, err := strconv.Atoi(lenField) + if err != nil { + return nil, fmt.Errorf("resp: malformed bulk length %q: %w", lenField, err) + } + if n < 0 { + return nil, nil + } + buf := make([]byte, n+2) // payload plus trailing CRLF + if _, err := io.ReadFull(r, buf); err != nil { + return nil, fmt.Errorf("resp: read bulk body: %w", err) + } + return buf[:n], nil +} + +// decodeArray reads n elements given the already-parsed "*" length field +// lenField, returning nil for a null array ("*-1"). +func decodeArray(r *bufio.Reader, lenField string) (any, error) { + n, err := strconv.Atoi(lenField) + if err != nil { + return nil, fmt.Errorf("resp: malformed array length %q: %w", lenField, err) + } + if n < 0 { + return nil, nil + } + arr := make([]any, n) + for i := 0; i < n; i++ { + v, err := decodeReply(r) + if err != nil { + return nil, fmt.Errorf("resp: array element %d/%d: %w", i+1, n, err) + } + arr[i] = v + } + return arr, nil +} + +// readLine reads one CRLF-terminated line from r, with the trailing CRLF +// (or bare LF) stripped. +func readLine(r *bufio.Reader) (string, error) { + line, err := r.ReadString('\n') + if err != nil { + return "", fmt.Errorf("resp: read line: %w", err) + } + return strings.TrimRight(line, "\r\n"), nil +} diff --git a/resp_test.go b/resp_test.go new file mode 100644 index 0000000..31e3aca --- /dev/null +++ b/resp_test.go @@ -0,0 +1,347 @@ +package traefikllmgateway + +import ( + "bufio" + "bytes" + "io" + "net" + "strconv" + "strings" + "testing" +) + +// TestEncodeCommand pins the exact RESP2 wire format every command must +// use: "*N\r\n$len\r\narg\r\n..." — a fake or real server parses this +// framing byte-for-byte, so a drift here breaks every other test silently. +func TestEncodeCommand(t *testing.T) { + cases := []struct { + name string + want string + args []string + }{ + {name: "single arg", args: []string{"GET"}, want: "*1\r\n$3\r\nGET\r\n"}, + {name: "multiple args", args: []string{"SET", "k", "v"}, want: "*3\r\n$3\r\nSET\r\n$1\r\nk\r\n$1\r\nv\r\n"}, + {name: "empty arg", args: []string{"SELECT", ""}, want: "*2\r\n$6\r\nSELECT\r\n$0\r\n\r\n"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := encodeCommand(c.args); got != c.want { + t.Errorf("encodeCommand(%v) = %q, want %q", c.args, got, c.want) + } + }) + } +} + +// TestDecodeReply covers every RESP2 type respClient must decode: simple +// string, error, integer, bulk string (including the null bulk "$-1"), and +// array (including the null array "*-1" and nesting). +func TestDecodeReply(t *testing.T) { + t.Run("simple string", func(t *testing.T) { + v, err := decodeReply(newReader("+OK\r\n")) + if err != nil { + t.Fatalf("decodeReply: %v", err) + } + if v != "OK" { + t.Errorf("v = %#v, want %q", v, "OK") + } + }) + + t.Run("error reply decodes to respErr, not a decode error", func(t *testing.T) { + v, err := decodeReply(newReader("-ERR wrong number of arguments\r\n")) + if err != nil { + t.Fatalf("decodeReply returned an error for a well-formed RESP error reply: %v", err) + } + e, ok := v.(respErr) + if !ok { + t.Fatalf("v = %#v (%T), want respErr", v, v) + } + if e.Error() != "ERR wrong number of arguments" { + t.Errorf("e.Error() = %q, want %q", e.Error(), "ERR wrong number of arguments") + } + }) + + t.Run("integer", func(t *testing.T) { + v, err := decodeReply(newReader(":42\r\n")) + if err != nil { + t.Fatalf("decodeReply: %v", err) + } + if v != int64(42) { + t.Errorf("v = %#v, want int64(42)", v) + } + }) + + t.Run("negative integer", func(t *testing.T) { + v, err := decodeReply(newReader(":-7\r\n")) + if err != nil { + t.Fatalf("decodeReply: %v", err) + } + if v != int64(-7) { + t.Errorf("v = %#v, want int64(-7)", v) + } + }) + + t.Run("bulk string", func(t *testing.T) { + v, err := decodeReply(newReader("$5\r\nhello\r\n")) + if err != nil { + t.Fatalf("decodeReply: %v", err) + } + b, ok := v.([]byte) + if !ok || !bytes.Equal(b, []byte("hello")) { + t.Errorf("v = %#v, want []byte(\"hello\")", v) + } + }) + + t.Run("null bulk (missing key)", func(t *testing.T) { + v, err := decodeReply(newReader("$-1\r\n")) + if err != nil { + t.Fatalf("decodeReply: %v", err) + } + if v != nil { + t.Errorf("v = %#v, want nil", v) + } + }) + + t.Run("array", func(t *testing.T) { + v, err := decodeReply(newReader("*2\r\n:5\r\n:1\r\n")) + if err != nil { + t.Fatalf("decodeReply: %v", err) + } + arr, ok := v.([]any) + if !ok || len(arr) != 2 || arr[0] != int64(5) || arr[1] != int64(1) { + t.Errorf("v = %#v, want []any{int64(5), int64(1)}", v) + } + }) + + t.Run("null array", func(t *testing.T) { + v, err := decodeReply(newReader("*-1\r\n")) + if err != nil { + t.Fatalf("decodeReply: %v", err) + } + if v != nil { + t.Errorf("v = %#v, want nil", v) + } + }) + + t.Run("malformed type byte is a decode error", func(t *testing.T) { + if _, err := decodeReply(newReader("?nope\r\n")); err == nil { + t.Fatal("want error for an unknown reply type byte") + } + }) +} + +// newReader wraps s in a *bufio.Reader, matching what respClient reads +// from its connection. +func newReader(s string) *bufio.Reader { + return bufio.NewReader(strings.NewReader(s)) +} + +// respStep is one entry in a fakeRESPServer script. +type respStep struct { + // wantArgs, if non-nil, asserts the next command read equals this + // exactly before replying. + wantArgs []string + // reply is written back verbatim; ignored when closeConn is true. + reply []byte + // closeConn closes the connection instead of replying, simulating a + // server that drops mid-command. The next step (if any) accepts a new + // connection. + closeConn bool +} + +// runFakeRESPServer executes script against ln in a background goroutine. +// It accepts a new connection whenever the previous one ends (either +// because a step set closeConn, or because the client itself closed it), +// so a script can simulate a server that drops a connection mid-command +// and expect the client to reconnect for the remaining steps. Mismatches +// are reported via t.Errorf, which is safe to call from a goroutine. +func runFakeRESPServer(t *testing.T, ln net.Listener, script []respStep) { + t.Helper() + go func() { + var conn net.Conn + var r *bufio.Reader + for _, step := range script { + if conn == nil { + c, err := ln.Accept() + if err != nil { + return // listener closed by t.Cleanup; test is finishing + } + conn = c + r = bufio.NewReader(conn) + } + + args, err := readRESPCommand(r) + if err != nil { + t.Errorf("fake RESP server: read command: %v", err) + return + } + if step.wantArgs != nil && !equalStrSlices(args, step.wantArgs) { + t.Errorf("fake RESP server: got command %v, want %v", args, step.wantArgs) + } + + if step.closeConn { + _ = conn.Close() + conn, r = nil, nil + continue + } + if _, err := conn.Write(step.reply); err != nil { + t.Errorf("fake RESP server: write reply: %v", err) + return + } + } + }() +} + +// readRESPCommand reads one RESP2 command array ("*N\r\n$len\r\narg\r\n...") +// from r — the same framing encodeCommand produces — and returns its args. +func readRESPCommand(r *bufio.Reader) ([]string, error) { + line, err := readLine(r) + if err != nil { + return nil, err + } + if !strings.HasPrefix(line, "*") { + return nil, &strconvLikeError{"want array header, got " + strconv.Quote(line)} + } + n, err := strconv.Atoi(line[1:]) + if err != nil { + return nil, err + } + args := make([]string, n) + for i := 0; i < n; i++ { + lenLine, err := readLine(r) + if err != nil { + return nil, err + } + if !strings.HasPrefix(lenLine, "$") { + return nil, &strconvLikeError{"want bulk header, got " + strconv.Quote(lenLine)} + } + l, err := strconv.Atoi(lenLine[1:]) + if err != nil { + return nil, err + } + buf := make([]byte, l+2) + if _, err := io.ReadFull(r, buf); err != nil { + return nil, err + } + args[i] = string(buf[:l]) + } + return args, nil +} + +// strconvLikeError is a trivial error string type for readRESPCommand's +// framing-mismatch messages. +type strconvLikeError struct{ msg string } + +func (e *strconvLikeError) Error() string { return e.msg } + +// equalStrSlices reports whether a and b hold the same strings in the same +// order. +func equalStrSlices(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +// newFakeListener returns a TCP listener on 127.0.0.1 with an OS-assigned +// port, closed automatically at test cleanup. +func newFakeListener(t *testing.T) net.Listener { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("net.Listen: %v", err) + } + t.Cleanup(func() { _ = ln.Close() }) + return ln +} + +// TestRESPClient_ConnectSendsAuthThenSelect is the brief's Step-1 +// connection-setup case: a configured password sends AUTH before SELECT, +// both requiring a non-error reply, and SELECT is sent with the +// configured db even when db is 0 — a fresh connection to a Redis- +// compatible proxy is not guaranteed to already sit on database 0. +func TestRESPClient_ConnectSendsAuthThenSelect(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"AUTH", "pw"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "missing-key"}, reply: []byte("$-1\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "pw", 0) + v, err := c.do("GET", "missing-key") + if err != nil { + t.Fatalf("do: %v", err) + } + if v != nil { + t.Errorf("v = %#v, want nil (null bulk for a missing key)", v) + } +} + +// TestRESPClient_NoPasswordSkipsAuth asserts an empty password never sends +// AUTH — the first command on the wire is SELECT. +func TestRESPClient_NoPasswordSkipsAuth(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + {wantArgs: []string{"SELECT", "3"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, reply: []byte("$1\r\nv\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 3) + v, err := c.do("GET", "k") + if err != nil { + t.Fatalf("do: %v", err) + } + if b, ok := v.([]byte); !ok || string(b) != "v" { + t.Errorf("v = %#v, want []byte(\"v\")", v) + } +} + +// TestRESPClient_ReconnectsOnceOnMidCommandClose is the brief's Step-1 +// reconnect case: the server drops the connection after reading a command +// but before replying (simulating a network failure mid-round-trip). The +// client must reconnect once — replaying the connection handshake — and +// retry the command, succeeding on the second attempt. +func TestRESPClient_ReconnectsOnceOnMidCommandClose(t *testing.T) { + ln := newFakeListener(t) + runFakeRESPServer(t, ln, []respStep{ + // First connection: handshake succeeds, then GET is read but the + // server closes without ever writing a reply. + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, closeConn: true}, + // Second connection: the client must redo the handshake, then + // resend GET and this time get a reply. + {wantArgs: []string{"SELECT", "0"}, reply: []byte("+OK\r\n")}, + {wantArgs: []string{"GET", "k"}, reply: []byte("$2\r\nv2\r\n")}, + }) + + c := newRESPClient(ln.Addr().String(), "", 0) + v, err := c.do("GET", "k") + if err != nil { + t.Fatalf("do: %v (want the client to reconnect once and succeed)", err) + } + if b, ok := v.([]byte); !ok || string(b) != "v2" { + t.Errorf("v = %#v, want []byte(\"v2\")", v) + } +} + +// TestRESPClient_ReconnectFailsTwice_ReturnsError asserts a client gives +// up (rather than retrying forever) once the retried attempt also fails: +// dialing a dead address never succeeds, so both the first attempt and the +// one reconnect attempt fail, and do returns an error. +func TestRESPClient_ReconnectFailsTwice_ReturnsError(t *testing.T) { + ln := newFakeListener(t) + deadAddr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatalf("Close: %v", err) + } // nothing listens at deadAddr from here on + + c := newRESPClient(deadAddr, "", 0) + if _, err := c.do("GET", "k"); err == nil { + t.Fatal("want an error dialing a dead address") + } +}