feat: stdlib RESP client and redis-backed limit store with fail-open

Add resp.go: a minimal stdlib-only RESP2 client (respClient) for a
single Redis-compatible server. It lazily dials one mutex-guarded
connection, applies 2s dial/read/write deadlines, sends AUTH (if a
password is set) then SELECT on every new connection, and reconnects
once on any I/O error before giving up. do and pipeline encode
commands as "*N\r\n$len\r\n..." and decode +/-/:/$/* replies.

Add redis_store.go: redisStore implements counterStore over
respClient. incrBy pipelines INCRBY and EXPIRE in one round trip and
returns INCRBY's value; get treats a missing key (RESP null bulk) as
0.

Wire the limiter's fail-open/fail-closed policy in limits.go, which
Task 5 stored but never read: a configured store's error now routes
through storeIncrBy/storeGet. failOpen=true falls back to the
in-process memoryStore transparently; failOpen=false makes
checkAndCount return a storeDown violation and makes account drop the
sample. Both paths log through an injectable, 30s-rate-limited logf.

Wire newGateway (llmgateway.go) to build a redisStore-backed limiter
when Config.Redis is set (password through resolveSecret, failOpen
defaulting to true unless RedisConfig.FailOpen says otherwise) and a
memory-only limiter otherwise, and to route pricing.go's unknown-model
warnings through the gateway's own logf.

Validate LimitsConfig at construction: a new LimitsConfig.validate
rejects a negative int64 count limit or a negative/NaN/Inf cost limit,
called from newAuthStore for every group and every user (inline and
file-sourced) so a malformed limit is a constructor error instead of
silently behaving as unlimited.

Add RequestsPerDay, TokensPerMonth, and CostPerMonthUSD enforcement
tests alongside the existing per-minute/day/cost-per-day cases, plus
an erroringStore-backed test pair proving the fail-open (fallback
counted) and fail-closed (storeDown violation) branches directly.

Folded the validation work into this single commit rather than a
second one: it interleaves with the fail-open changes across
limits.go, limits_test.go, and llmgateway_test.go closely enough that
splitting by hunk would be artificial.
This commit is contained in:
2026-08-20 02:00:26 +01:00
parent ed1b1d7b2f
commit 5eab51bdef
9 changed files with 1538 additions and 49 deletions
+6
View File
@@ -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)
+173 -36
View File
@@ -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()
+230 -6
View File
@@ -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)
}
})
}
}
+52 -7
View File
@@ -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
+219
View File
@@ -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())
}
}
+81
View File
@@ -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
}
+138
View File
@@ -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)
}
}
+292
View File
@@ -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
}
+347
View File
@@ -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")
}
}