diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index 468cd336..e3d33cc7 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -607,7 +607,12 @@ func (s *Server) handleStreamingResponse(w http.ResponseWriter, r *http.Request, // leaves inference/a2a stuck in an unfinalized state and emits no // SessionResponse row to abctl. defer func() { - finalAction := s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, nil, true) + // Use a detached context for finalization: the client may have + // cancelled the request context after reading the full stream, + // but aggregating plugins (inference-parser, token-budget) still + // need their last=true dispatch to finalize state. + finalCtx := context.WithoutCancel(r.Context()) + finalAction := s.OutboundPipeline.RunResponseFrame(finalCtx, pctx, nil, true) if finalAction.Type == pipeline.Reject { // Headers already sent; we can't promote to 502, but // surface the policy violation so operators see it. diff --git a/authbridge/authlib/pipeline/action.go b/authbridge/authlib/pipeline/action.go index b5079d31..072fdd22 100644 --- a/authbridge/authlib/pipeline/action.go +++ b/authbridge/authlib/pipeline/action.go @@ -72,6 +72,7 @@ var codeToStatus = map[string]int{ "upstream.token-exchange-failed": http.StatusServiceUnavailable, "upstream.timeout": http.StatusGatewayTimeout, "pipeline.cancelled": 499, // client-closed request (nginx convention) + "budget.exceeded": http.StatusForbidden, } // StatusFromCode returns the HTTP status a Violation maps to when its diff --git a/authbridge/authlib/plugins/tokenbudget/e2e_test.go b/authbridge/authlib/plugins/tokenbudget/e2e_test.go new file mode 100644 index 00000000..38b9cc58 --- /dev/null +++ b/authbridge/authlib/plugins/tokenbudget/e2e_test.go @@ -0,0 +1,294 @@ +package tokenbudget + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "net/url" + "sync" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/listener/forwardproxy" + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/session" +) + +func newE2EPlugin(t *testing.T, maxTokens int64, store *memStore) *TokenBudget { + t.Helper() + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: maxTokens, + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = store + go p.refreshLoop(30 * time.Millisecond) + t.Cleanup(func() { close(p.stopCh); <-p.stopped }) + return p +} + +func respond(p *TokenBudget, sessionID string, tokens int) { + p.OnResponseFrame(context.Background(), makePctx(sessionID, tokens), nil, true) +} + +func request(p *TokenBudget, sessionID string) pipeline.Action { + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: sessionID}, + } + return p.OnRequest(context.Background(), pctx) +} + +// TestE2E_HTTPRoundTrip wires token-budget into a real forward proxy. +// Under-budget requests reach the backend; the proxy is functional. +func TestE2E_HTTPRoundTrip(t *testing.T) { + store := newMemStore() + p := newE2EPlugin(t, 1000, store) + + pipe, err := pipeline.New([]pipeline.Plugin{p}) + if err != nil { + t.Fatal(err) + } + sessions := session.New(5*time.Minute, 100, 0) + defer sessions.Close() + + srv, err := forwardproxy.NewServer(pipeline.NewHolder(pipe), sessions, nil) + if err != nil { + t.Fatal(err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + })) + defer backend.Close() + + proxyURL, _ := url.Parse(proxy.URL) + client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyURL)}} + + req, _ := http.NewRequest(http.MethodGet, backend.URL+"/v1/chat/completions", nil) + resp, err := client.Do(req) + if err != nil { + t.Fatalf("request through proxy: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + t.Fatalf("expected 200, got %d: %s", resp.StatusCode, body) + } +} + +// TestE2E_AccumulateAndDeny verifies the full lifecycle: accumulate +// tokens via OnResponseFrame, then OnRequest denies with a 403. +func TestE2E_AccumulateAndDeny(t *testing.T) { + p := newE2EPlugin(t, 150, newMemStore()) + + for i := 0; i < 3; i++ { + respond(p, "sess", 60) + } + + action := request(p, "sess") + if action.Type != pipeline.Reject { + t.Fatalf("expected Reject, got %v", action.Type) + } + status, _, body := action.Violation.Render() + if status != http.StatusForbidden { + t.Errorf("status = %d, want 403", status) + } + var parsed map[string]any + if err := json.Unmarshal(body, &parsed); err != nil { + t.Fatal(err) + } + if parsed["error"] != "budget.exceeded" { + t.Errorf("error = %v, want budget.exceeded", parsed["error"]) + } +} + +// TestE2E_MultiSession verifies independent session budgets. +func TestE2E_MultiSession(t *testing.T) { + p := newE2EPlugin(t, 100, newMemStore()) + + for i := 0; i < 3; i++ { + respond(p, "A", 40) // 120 > 100 + } + respond(p, "B", 20) // 20 < 100 + + if a := request(p, "A"); a.Type != pipeline.Reject { + t.Fatalf("session A: expected Reject, got %v", a.Type) + } + if a := request(p, "B"); a.Type != pipeline.Continue { + t.Fatalf("session B: expected Continue, got %v", a.Type) + } +} + +// TestE2E_LocalCacheEnforcesDuringOutage confirms that a populated +// cache enforces even when the backing store is unreachable. +func TestE2E_LocalCacheEnforcesDuringOutage(t *testing.T) { + p := newE2EPlugin(t, 100, newMemStore()) + p.store = &failingStore{} + + p.mu.Lock() + p.cache["s"] = &counters{tokens: 110, calls: 5, startedAt: time.Now()} + p.mu.Unlock() + + if a := request(p, "s"); a.Type != pipeline.Reject { + t.Fatalf("expected Reject from cache with store down, got %v", a.Type) + } +} + +// TestE2E_RefreshRecovery confirms that refreshCache picks up +// authoritative store values after an outage resolves. +func TestE2E_RefreshRecovery(t *testing.T) { + inner := newMemStore() + cs := &controllableStore{inner: inner} + p := newE2EPlugin(t, 200, newMemStore()) + p.store = cs + + ctx := context.Background() + inner.HashIncr(ctx, "token-budget:s", "tokens", 180) + inner.HashIncr(ctx, "token-budget:s", "calls", 7) + inner.HashSetNX(ctx, "token-budget:s", "started_at", "1700000000") + + p.mu.Lock() + p.cache["s"] = &counters{tokens: 50} + p.mu.Unlock() + + cs.setFailing(true) + p.refreshCache() + p.mu.RLock() + if p.cache["s"].tokens != 50 { + t.Fatalf("during outage: tokens = %d, want 50", p.cache["s"].tokens) + } + p.mu.RUnlock() + + cs.setFailing(false) + p.refreshCache() + p.mu.RLock() + if p.cache["s"].tokens != 180 { + t.Errorf("after recovery: tokens = %d, want 180", p.cache["s"].tokens) + } + p.mu.RUnlock() +} + +// TestE2E_PodRestart verifies that a fresh plugin with an empty cache +// resumes enforcement after refresh picks up pre-existing store counters. +func TestE2E_PodRestart(t *testing.T) { + store := newMemStore() + p := newE2EPlugin(t, 200, store) + + ctx := context.Background() + store.HashIncr(ctx, "token-budget:s", "tokens", 190) + store.HashIncr(ctx, "token-budget:s", "calls", 8) + store.HashSetNX(ctx, "token-budget:s", "started_at", "1700000000") + + // Cold cache — first request passes (overshoot window). + if a := request(p, "s"); a.Type != pipeline.Continue { + t.Fatalf("cold cache: expected Continue, got %v", a.Type) + } + + // Seed cache entry so refresh picks up this session. + respond(p, "s", 15) + + // Wait for refresh (30ms interval, 80ms is ~2.6x margin). + time.Sleep(80 * time.Millisecond) + + if a := request(p, "s"); a.Type != pipeline.Reject { + t.Fatalf("after refresh: expected Reject, got %v", a.Type) + } +} + +// controllableStore delegates to inner memStore but can be toggled to fail. +type controllableStore struct { + inner *memStore + failing bool + mu sync.Mutex +} + +func (c *controllableStore) setFailing(v bool) { c.mu.Lock(); c.failing = v; c.mu.Unlock() } +func (c *controllableStore) isFailing() bool { c.mu.Lock(); defer c.mu.Unlock(); return c.failing } +func (c *controllableStore) err() error { return context.DeadlineExceeded } + +func (c *controllableStore) Get(ctx context.Context, key string) (string, error) { + if c.isFailing() { return "", c.err() } + return c.inner.Get(ctx, key) +} +func (c *controllableStore) Set(ctx context.Context, key, value string, ttl time.Duration) error { + if c.isFailing() { return c.err() } + return c.inner.Set(ctx, key, value, ttl) +} +func (c *controllableStore) Incr(ctx context.Context, key string, delta int64) (int64, error) { + if c.isFailing() { return 0, c.err() } + return c.inner.Incr(ctx, key, delta) +} +func (c *controllableStore) HashIncr(ctx context.Context, key, field string, delta int64) (int64, error) { + if c.isFailing() { return 0, c.err() } + return c.inner.HashIncr(ctx, key, field, delta) +} +func (c *controllableStore) HashGet(ctx context.Context, key string) (map[string]string, error) { + if c.isFailing() { return nil, c.err() } + return c.inner.HashGet(ctx, key) +} +func (c *controllableStore) HashSetNX(ctx context.Context, key, field, value string) (bool, error) { + if c.isFailing() { return false, c.err() } + return c.inner.HashSetNX(ctx, key, field, value) +} +func (c *controllableStore) Expire(ctx context.Context, key string, ttl time.Duration) error { + if c.isFailing() { return c.err() } + return c.inner.Expire(ctx, key, ttl) +} +func (c *controllableStore) Close() error { return nil } + +// TestE2E_ShadowMode verifies that on_exceed=observe allows requests +// through even when budget is exceeded, while still accumulating. +func TestE2E_ShadowMode(t *testing.T) { + p := New() + cfg, _ := json.Marshal(config{ + RedisURL: "mem://test", + MaxTokens: 150, + OnExceed: "observe", + RefreshInterval: "30ms", + RedisUnavailable: "fail_open", + }) + if err := p.Configure(cfg); err != nil { + t.Fatalf("Configure: %v", err) + } + store := newMemStore() + p.store = store + go p.refreshLoop(30 * time.Millisecond) + t.Cleanup(func() { close(p.stopCh); <-p.stopped }) + + for i := 0; i < 3; i++ { + respond(p, "sess", 60) // 180 total > 150 limit + } + + // In observe mode, request should continue (not reject). + action := request(p, "sess") + if action.Type != pipeline.Continue { + t.Fatalf("shadow mode: expected Continue, got %v", action.Type) + } + + // Counters should still accumulate past the limit. + respond(p, "sess", 20) // 200 total + p.mu.RLock() + c := p.cache["sess"] + p.mu.RUnlock() + if c.tokens != 200 { + t.Errorf("tokens = %d, want 200 (accumulation continues in shadow mode)", c.tokens) + } + + // Subsequent requests also continue. + action = request(p, "sess") + if action.Type != pipeline.Continue { + t.Fatalf("shadow mode (2nd request): expected Continue, got %v", action.Type) + } +} diff --git a/authbridge/authlib/plugins/tokenbudget/lifecycle_test.go b/authbridge/authlib/plugins/tokenbudget/lifecycle_test.go new file mode 100644 index 00000000..ee0a129e --- /dev/null +++ b/authbridge/authlib/plugins/tokenbudget/lifecycle_test.go @@ -0,0 +1,79 @@ +package tokenbudget + +import ( + "context" + "encoding/json" + "net/http" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +func TestFullLifecycle_RefreshFromRedis(t *testing.T) { + store := newMemStore() + p := New() + cfg := `{"redis_url":"mem://test","max_tokens":500,"refresh_interval":"50ms","session_ttl_seconds":3600}` + if err := p.Configure(json.RawMessage(cfg)); err != nil { + t.Fatal(err) + } + p.store = store + + ctx := context.Background() + store.HashIncr(ctx, "token-budget:remote-sess", "tokens", 450) + store.HashIncr(ctx, "token-budget:remote-sess", "calls", 10) + store.HashSetNX(ctx, "token-budget:remote-sess", "started_at", "1700000000") + + // Seed cache so refreshCache picks it up. + p.mu.Lock() + p.cache["remote-sess"] = &counters{} + p.mu.Unlock() + + p.refreshCache() + + p.mu.RLock() + c := p.cache["remote-sess"] + p.mu.RUnlock() + if c == nil { + t.Fatal("expected cache entry after refresh") + } + if c.tokens != 450 { + t.Errorf("tokens = %d, want 450", c.tokens) + } + if c.calls != 10 { + t.Errorf("calls = %d, want 10", c.calls) + } + if c.startedAt.Unix() != 1700000000 { + t.Errorf("startedAt = %v, want 1700000000", c.startedAt.Unix()) + } +} + +func TestFullLifecycle_DurationEnforcement(t *testing.T) { + p := newTestPlugin(0, 0, 1) + p.store = newMemStore() + + ctx := context.Background() + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: "dur-sess"}, + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{TotalTokens: 10}, + }, + } + p.OnResponseFrame(ctx, pctx, nil, true) + + // Backdate startedAt past the 1s limit. + p.mu.Lock() + p.cache["dur-sess"].startedAt = time.Now().Add(-2 * time.Second) + p.mu.Unlock() + + action := p.OnRequest(ctx, &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: "dur-sess"}, + }) + if action.Type != pipeline.Reject { + t.Fatalf("expected Reject after duration exceeded, got %v", action.Type) + } +} diff --git a/authbridge/authlib/plugins/tokenbudget/plugin.go b/authbridge/authlib/plugins/tokenbudget/plugin.go new file mode 100644 index 00000000..102c0299 --- /dev/null +++ b/authbridge/authlib/plugins/tokenbudget/plugin.go @@ -0,0 +1,301 @@ +// Package tokenbudget enforces per-session lifetime budgets on tokens, +// inference calls, and wall-clock duration. Must run before inference-parser +// in the declared plugin order (response path is reverse: inference-parser +// finalizes counts first, then this plugin reads them). +package tokenbudget + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "strconv" + "sync" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/plugins" + "github.com/rossoctl/cortex/authbridge/authlib/storage" +) + +type config struct { + RedisURL string `json:"redis_url" required:"true" description:"Redis/Valkey connection URL."` + MaxTokens int64 `json:"max_tokens" description:"Cumulative token ceiling per session. 0 = no limit."` + MaxCalls int64 `json:"max_calls" description:"Max inference calls per session. 0 = no limit."` + MaxDurationSeconds int64 `json:"max_duration_seconds" description:"Wall-clock session lifetime in seconds. 0 = no limit."` + OnExceed string `json:"on_exceed" description:"Action on breach: deny (block) or observe (shadow — log but continue)." default:"deny" enum:"deny,observe"` + SessionTTLSeconds int `json:"session_ttl_seconds" description:"Redis key TTL; should be >= max_duration_seconds." default:"7200"` + RefreshInterval string `json:"refresh_interval" description:"How often to sync local cache from Redis." default:"5s"` + RedisUnavailable string `json:"redis_unavailable" description:"Behavior when Redis is unreachable." default:"fail_open" enum:"fail_open,fail_closed"` +} + +type counters struct { + tokens int64 + calls int64 + startedAt time.Time +} + +// TokenBudget is the plugin state. Redis provides cross-pod durability; +// the local cache provides zero-I/O enforcement on the request path. +type TokenBudget struct { + cfg config + store storage.Store + log *slog.Logger + + mu sync.RWMutex + cache map[string]*counters + stopCh chan struct{} + stopped chan struct{} +} + +func New() *TokenBudget { + return &TokenBudget{ + cache: make(map[string]*counters), + stopCh: make(chan struct{}), + stopped: make(chan struct{}), + log: slog.Default().With("plugin", "token-budget"), + } +} + +func init() { + plugins.RegisterPlugin("token-budget", func() pipeline.Plugin { return New() }) +} + +func (p *TokenBudget) Name() string { return "token-budget" } + +func (p *TokenBudget) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ + Description: "Enforce per-session token, call, and duration budgets via Redis.", + } +} + +func (p *TokenBudget) Configure(raw json.RawMessage) error { + p.cfg = config{ + OnExceed: "deny", + SessionTTLSeconds: 7200, + RefreshInterval: "5s", + RedisUnavailable: "fail_open", + } + if err := json.Unmarshal(raw, &p.cfg); err != nil { + return fmt.Errorf("token-budget config: %w", err) + } + if p.cfg.RedisURL == "" { + return fmt.Errorf("token-budget: redis_url is required") + } + if p.cfg.MaxTokens <= 0 && p.cfg.MaxCalls <= 0 && p.cfg.MaxDurationSeconds <= 0 { + return fmt.Errorf("token-budget: at least one limit (max_tokens, max_calls, max_duration_seconds) must be > 0") + } + return nil +} + +func (p *TokenBudget) Init(_ context.Context) error { + store, err := storage.Open("redis", p.cfg.RedisURL) + if err != nil { + return fmt.Errorf("token-budget: redis connect: %w", err) + } + p.store = store + + interval, err := time.ParseDuration(p.cfg.RefreshInterval) + if err != nil { + interval = 5 * time.Second + } + go p.refreshLoop(interval) + return nil +} + +func (p *TokenBudget) Shutdown(_ context.Context) error { + close(p.stopCh) + <-p.stopped + if p.store != nil { + return p.store.Close() + } + return nil +} + +// OnRequest evaluates cached counters against limits. No I/O. +func (p *TokenBudget) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { + sessionID := p.sessionID(pctx) + if sessionID == "" { + return pipeline.Action{Type: pipeline.Continue} + } + + p.mu.RLock() + c, ok := p.cache[sessionID] + p.mu.RUnlock() + + if !ok { + return pipeline.Action{Type: pipeline.Continue} + } + + if reason := p.evaluate(c); reason != "" { + if p.cfg.OnExceed == "observe" { + pctx.Observe("shadow_budget_exceeded") + p.log.Warn("budget exceeded (shadow mode)", + "session", sessionID, + "reason", reason, + "tokens", c.tokens, + "calls", c.calls) + return pipeline.Action{Type: pipeline.Continue} + } + return pipeline.DenyWithDetails("budget.exceeded", reason, map[string]any{ + "spent_tokens": c.tokens, + "spent_calls": c.calls, + "limit_tokens": p.cfg.MaxTokens, + "limit_calls": p.cfg.MaxCalls, + }) + } + return pipeline.Action{Type: pipeline.Continue} +} + +// OnResponse is a no-op; see OnResponseFrame. +func (p *TokenBudget) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} + +// OnResponseFrame accumulates token counts on finalization (last=true). +func (p *TokenBudget) OnResponseFrame(_ context.Context, pctx *pipeline.Context, _ []byte, last bool) pipeline.Action { + if !last { + return pipeline.Action{Type: pipeline.Continue} + } + + sessionID := p.sessionID(pctx) + if sessionID == "" { + return pipeline.Action{Type: pipeline.Continue} + } + + inf := pctx.Extensions.Inference + if inf == nil || inf.TotalTokens == 0 { + return pipeline.Action{Type: pipeline.Continue} + } + + tokens := int64(inf.TotalTokens) + + go p.accumulate(sessionID, tokens) + + p.mu.Lock() + c, ok := p.cache[sessionID] + if !ok { + c = &counters{startedAt: time.Now()} + p.cache[sessionID] = c + } + c.tokens += tokens + c.calls++ + p.mu.Unlock() + + return pipeline.Action{Type: pipeline.Continue} +} + +func (p *TokenBudget) evaluate(c *counters) string { + if p.cfg.MaxTokens > 0 && c.tokens >= p.cfg.MaxTokens { + return fmt.Sprintf("token limit reached: %d/%d", c.tokens, p.cfg.MaxTokens) + } + if p.cfg.MaxCalls > 0 && c.calls >= p.cfg.MaxCalls { + return fmt.Sprintf("call limit reached: %d/%d", c.calls, p.cfg.MaxCalls) + } + if p.cfg.MaxDurationSeconds > 0 && !c.startedAt.IsZero() { + elapsed := time.Since(c.startedAt).Seconds() + if int64(elapsed) >= p.cfg.MaxDurationSeconds { + return fmt.Sprintf("duration limit reached: %ds/%ds", int64(elapsed), p.cfg.MaxDurationSeconds) + } + } + return "" +} + +// accumulate writes counters to Redis. On failure, writes are dropped (fail-open). +func (p *TokenBudget) accumulate(sessionID string, tokens int64) { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + key := p.redisKey(sessionID) + ttl := time.Duration(p.cfg.SessionTTLSeconds) * time.Second + + _, err := p.store.HashIncr(ctx, key, "tokens", tokens) + if err != nil { + p.log.Warn("redis HashIncr tokens failed", "session", sessionID, "err", err) + return + } + + _, _ = p.store.HashIncr(ctx, key, "calls", 1) + + set, _ := p.store.HashSetNX(ctx, key, "started_at", strconv.FormatInt(time.Now().Unix(), 10)) + if set { + _ = p.store.Expire(ctx, key, ttl) + } +} + +func (p *TokenBudget) refreshLoop(interval time.Duration) { + defer close(p.stopped) + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-p.stopCh: + return + case <-ticker.C: + p.refreshCache() + } + } +} + +// refreshCache replaces local counters with authoritative Redis values. +func (p *TokenBudget) refreshCache() { + p.mu.RLock() + keys := make([]string, 0, len(p.cache)) + for k := range p.cache { + keys = append(keys, k) + } + p.mu.RUnlock() + + for _, sessionID := range keys { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + fields, err := p.store.HashGet(ctx, p.redisKey(sessionID)) + cancel() + + if err != nil { + // TODO: fail_closed should deny requests when Redis is unreachable + // and the local cache has no data. Currently both modes retain stale cache. + if p.cfg.RedisUnavailable != "fail_open" { + p.log.Warn("redis refresh failed", "session", sessionID, "err", err) + } + continue + } + + if len(fields) == 0 { + p.mu.Lock() + delete(p.cache, sessionID) + p.mu.Unlock() + continue + } + + tokens, _ := strconv.ParseInt(fields["tokens"], 10, 64) + calls, _ := strconv.ParseInt(fields["calls"], 10, 64) + var startedAt time.Time + if ts, err := strconv.ParseInt(fields["started_at"], 10, 64); err == nil { + startedAt = time.Unix(ts, 0) + } + + p.mu.Lock() + p.cache[sessionID] = &counters{tokens: tokens, calls: calls, startedAt: startedAt} + p.mu.Unlock() + } +} + +func (p *TokenBudget) sessionID(pctx *pipeline.Context) string { + if pctx.Session != nil && pctx.Session.ID != "" { + return pctx.Session.ID + } + return "" +} + +func (p *TokenBudget) redisKey(sessionID string) string { + return "token-budget:" + sessionID +} + +var ( + _ pipeline.Plugin = (*TokenBudget)(nil) + _ pipeline.Configurable = (*TokenBudget)(nil) + _ pipeline.Initializer = (*TokenBudget)(nil) + _ pipeline.Shutdowner = (*TokenBudget)(nil) + _ pipeline.StreamingResponder = (*TokenBudget)(nil) +) diff --git a/authbridge/authlib/plugins/tokenbudget/plugin_test.go b/authbridge/authlib/plugins/tokenbudget/plugin_test.go new file mode 100644 index 00000000..426c4244 --- /dev/null +++ b/authbridge/authlib/plugins/tokenbudget/plugin_test.go @@ -0,0 +1,345 @@ +package tokenbudget + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/storage" +) + +// memStore is a minimal in-memory storage.Store for testing. +type memStore struct { + mu sync.Mutex + hashes map[string]map[string]string + kvs map[string]string + ttls map[string]time.Duration +} + +func newMemStore() *memStore { + return &memStore{ + hashes: make(map[string]map[string]string), + kvs: make(map[string]string), + ttls: make(map[string]time.Duration), + } +} + +func (m *memStore) Get(_ context.Context, key string) (string, error) { + m.mu.Lock() + defer m.mu.Unlock() + return m.kvs[key], nil +} + +func (m *memStore) Set(_ context.Context, key, value string, ttl time.Duration) error { + m.mu.Lock() + defer m.mu.Unlock() + m.kvs[key] = value + if ttl > 0 { + m.ttls[key] = ttl + } + return nil +} + +func (m *memStore) Incr(_ context.Context, key string, delta int64) (int64, error) { + m.mu.Lock() + defer m.mu.Unlock() + var cur int64 + if v, ok := m.kvs[key]; ok { + fmt.Sscanf(v, "%d", &cur) + } + cur += delta + m.kvs[key] = fmt.Sprintf("%d", cur) + return cur, nil +} + +func (m *memStore) HashIncr(_ context.Context, key, field string, delta int64) (int64, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.hashes[key] == nil { + m.hashes[key] = make(map[string]string) + } + var cur int64 + if v, ok := m.hashes[key][field]; ok { + fmt.Sscanf(v, "%d", &cur) + } + cur += delta + m.hashes[key][field] = fmt.Sprintf("%d", cur) + return cur, nil +} + +func (m *memStore) HashGet(_ context.Context, key string) (map[string]string, error) { + m.mu.Lock() + defer m.mu.Unlock() + h := m.hashes[key] + if h == nil { + return map[string]string{}, nil + } + out := make(map[string]string, len(h)) + for k, v := range h { + out[k] = v + } + return out, nil +} + +func (m *memStore) HashSetNX(_ context.Context, key, field, value string) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.hashes[key] == nil { + m.hashes[key] = make(map[string]string) + } + if _, exists := m.hashes[key][field]; exists { + return false, nil + } + m.hashes[key][field] = value + return true, nil +} + +func (m *memStore) Expire(_ context.Context, key string, ttl time.Duration) error { + m.mu.Lock() + defer m.mu.Unlock() + m.ttls[key] = ttl + return nil +} + +func (m *memStore) Close() error { return nil } + +var _ storage.Store = (*memStore)(nil) + +// failingStore always returns errors (simulates total store unavailability). +type failingStore struct{} + +func (failingStore) Get(context.Context, string) (string, error) { return "", context.DeadlineExceeded } +func (failingStore) Set(context.Context, string, string, time.Duration) error { return context.DeadlineExceeded } +func (failingStore) Incr(context.Context, string, int64) (int64, error) { return 0, context.DeadlineExceeded } +func (failingStore) HashIncr(context.Context, string, string, int64) (int64, error) { return 0, context.DeadlineExceeded } +func (failingStore) HashGet(context.Context, string) (map[string]string, error) { return nil, context.DeadlineExceeded } +func (failingStore) HashSetNX(context.Context, string, string, string) (bool, error) { return false, context.DeadlineExceeded } +func (failingStore) Expire(context.Context, string, time.Duration) error { return context.DeadlineExceeded } +func (failingStore) Close() error { return nil } + +func init() { + storage.Register("mem", func(_ string) (storage.Store, error) { + return newMemStore(), nil + }) +} + +func newTestPlugin(maxTokens, maxCalls, maxDuration int64) *TokenBudget { + p := New() + cfg := fmt.Sprintf(`{ + "redis_url": "mem://test", + "max_tokens": %d, + "max_calls": %d, + "max_duration_seconds": %d, + "refresh_interval": "100ms" + }`, maxTokens, maxCalls, maxDuration) + if err := p.Configure(json.RawMessage(cfg)); err != nil { + panic(err) + } + store := newMemStore() + p.store = store + return p +} + +func makePctx(sessionID string, totalTokens int) *pipeline.Context { + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: sessionID}, + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{ + TotalTokens: totalTokens, + }, + }, + } + return pctx +} + +func TestOnRequest_UnderLimit(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := makePctx("sess-1", 0) + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } +} + + +func TestOnResponseFrame_Accumulates(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := makePctx("sess-1", 42) + + action := p.OnResponseFrame(context.Background(), pctx, nil, true) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + + // Check in-memory cache was updated. + p.mu.RLock() + c := p.cache["sess-1"] + p.mu.RUnlock() + + if c == nil { + t.Fatal("expected cache entry") + } + if c.tokens != 42 { + t.Errorf("tokens = %d, want 42", c.tokens) + } + if c.calls != 1 { + t.Errorf("calls = %d, want 1", c.calls) + } +} + +func TestOnResponseFrame_SkipsNonLast(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := makePctx("sess-1", 42) + + action := p.OnResponseFrame(context.Background(), pctx, []byte("data"), false) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + + p.mu.RLock() + _, ok := p.cache["sess-1"] + p.mu.RUnlock() + if ok { + t.Error("expected no cache entry on non-last frame") + } +} + +func TestOnResponseFrame_NoInference(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + Session: &pipeline.SessionView{ID: "sess-1"}, + } + + action := p.OnResponseFrame(context.Background(), pctx, nil, true) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue, got %v", action.Type) + } + + p.mu.RLock() + _, ok := p.cache["sess-1"] + p.mu.RUnlock() + if ok { + t.Error("expected no cache entry when no inference data") + } +} + +func TestOnRequest_NoSession(t *testing.T) { + p := newTestPlugin(100, 0, 0) + pctx := &pipeline.Context{ + Direction: pipeline.Outbound, + Headers: http.Header{}, + } + + action := p.OnRequest(context.Background(), pctx) + if action.Type != pipeline.Continue { + t.Fatalf("expected Continue for nil session, got %v", action.Type) + } +} + +func TestAccumulate_WritesToStore(t *testing.T) { + p := newTestPlugin(1000, 0, 0) + store := newMemStore() + p.store = store + + p.accumulate("sess-1", 100) + + fields, _ := store.HashGet(context.Background(), "token-budget:sess-1") + if fields["tokens"] != "100" { + t.Errorf("tokens in store = %q, want 100", fields["tokens"]) + } + if fields["calls"] != "1" { + t.Errorf("calls in store = %q, want 1", fields["calls"]) + } + if fields["started_at"] == "" { + t.Error("started_at not set in store") + } +} + +func TestConfigure_Validation(t *testing.T) { + tests := []struct { + name string + cfg string + wantErr bool + }{ + {"valid", `{"redis_url":"redis://localhost","max_tokens":100}`, false}, + {"missing redis_url", `{"max_tokens":100}`, true}, + {"no limits", `{"redis_url":"redis://localhost"}`, true}, + {"invalid json", `{broken}`, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + p := New() + err := p.Configure(json.RawMessage(tt.cfg)) + if (err != nil) != tt.wantErr { + t.Errorf("Configure() error = %v, wantErr = %v", err, tt.wantErr) + } + }) + } +} + +func TestOnRequest_ShadowMode(t *testing.T) { + p := New() + cfg := `{ + "redis_url": "mem://test", + "max_tokens": 100, + "on_exceed": "observe", + "refresh_interval": "100ms" + }` + if err := p.Configure(json.RawMessage(cfg)); err != nil { + t.Fatalf("Configure: %v", err) + } + p.store = newMemStore() + + p.OnResponseFrame(context.Background(), makePctx("sess-1", 60), nil, true) + p.OnResponseFrame(context.Background(), makePctx("sess-1", 60), nil, true) + + p.mu.RLock() + c := p.cache["sess-1"] + p.mu.RUnlock() + if c.tokens != 120 { + t.Errorf("tokens = %d, want 120", c.tokens) + } + + action := p.OnRequest(context.Background(), makePctx("sess-1", 0)) + if action.Type != pipeline.Continue { + t.Fatalf("shadow mode: expected Continue past limit, got %v", action.Type) + } +} + +func TestEvaluate_MultipleLimits(t *testing.T) { + p := newTestPlugin(100, 10, 60) + + tests := []struct { + name string + c *counters + wantDeny bool + }{ + {"all under", &counters{tokens: 50, calls: 5, startedAt: time.Now()}, false}, + {"tokens over", &counters{tokens: 100, calls: 5, startedAt: time.Now()}, true}, + {"calls over", &counters{tokens: 50, calls: 10, startedAt: time.Now()}, true}, + {"duration over", &counters{tokens: 50, calls: 5, startedAt: time.Now().Add(-90 * time.Second)}, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + reason := p.evaluate(tt.c) + if tt.wantDeny && reason == "" { + t.Error("expected denial reason, got empty") + } + if !tt.wantDeny && reason != "" { + t.Errorf("expected no denial, got %q", reason) + } + }) + } +} diff --git a/authbridge/authlib/storage/provider.go b/authbridge/authlib/storage/provider.go new file mode 100644 index 00000000..ec38f7a8 --- /dev/null +++ b/authbridge/authlib/storage/provider.go @@ -0,0 +1,36 @@ +package storage + +import ( + "fmt" + "sync" +) + +// Factory creates a Store from a connection URL. +type Factory func(url string) (Store, error) + +var ( + mu sync.RWMutex + factories = make(map[string]Factory) +) + +// Register makes a store factory available by scheme (e.g. "redis", "valkey"). +// Typically called from init() in the implementation module. +func Register(scheme string, f Factory) { + mu.Lock() + defer mu.Unlock() + if _, dup := factories[scheme]; dup { + panic("storage: Register called twice for scheme " + scheme) + } + factories[scheme] = f +} + +// Open creates a Store using the registered factory for the given scheme. +func Open(scheme, url string) (Store, error) { + mu.RLock() + f, ok := factories[scheme] + mu.RUnlock() + if !ok { + return nil, fmt.Errorf("storage: unknown scheme %q (forgot to import the driver?)", scheme) + } + return f(url) +} diff --git a/authbridge/authlib/storage/provider_test.go b/authbridge/authlib/storage/provider_test.go new file mode 100644 index 00000000..2bf3fd9e --- /dev/null +++ b/authbridge/authlib/storage/provider_test.go @@ -0,0 +1,51 @@ +package storage + +import ( + "context" + "testing" + "time" +) + +type nopStore struct{} + +func (nopStore) Get(context.Context, string) (string, error) { return "", nil } +func (nopStore) Set(context.Context, string, string, time.Duration) error { return nil } +func (nopStore) Incr(context.Context, string, int64) (int64, error) { return 0, nil } +func (nopStore) HashIncr(context.Context, string, string, int64) (int64, error) { return 0, nil } +func (nopStore) HashGet(context.Context, string) (map[string]string, error) { return nil, nil } +func (nopStore) HashSetNX(context.Context, string, string, string) (bool, error) { + return false, nil +} +func (nopStore) Expire(context.Context, string, time.Duration) error { return nil } +func (nopStore) Close() error { return nil } + +func TestRegisterAndOpen(t *testing.T) { + Register("test-scheme", func(url string) (Store, error) { + return nopStore{}, nil + }) + + s, err := Open("test-scheme", "test://addr") + if err != nil { + t.Fatalf("Open: %v", err) + } + if s == nil { + t.Fatal("expected non-nil store") + } +} + +func TestOpen_UnknownScheme(t *testing.T) { + _, err := Open("nonexistent", "foo://bar") + if err == nil { + t.Error("expected error for unknown scheme") + } +} + +func TestRegister_Duplicate(t *testing.T) { + defer func() { + if r := recover(); r == nil { + t.Error("expected panic on duplicate Register") + } + }() + Register("dup-scheme", func(url string) (Store, error) { return nopStore{}, nil }) + Register("dup-scheme", func(url string) (Store, error) { return nopStore{}, nil }) +} diff --git a/authbridge/authlib/storage/store.go b/authbridge/authlib/storage/store.go new file mode 100644 index 00000000..0984a5ea --- /dev/null +++ b/authbridge/authlib/storage/store.go @@ -0,0 +1,36 @@ +package storage + +import ( + "context" + "time" +) + +// Store provides key-value and hash operations for plugins that need +// persistent, cross-pod state (budget counters, session data, caches, +// rate limiting). Implementations live in separate modules to isolate +// their dependencies. +type Store interface { + // Get returns the value for a simple key. Returns "" and nil if the key does not exist. + Get(ctx context.Context, key string) (string, error) + + // Set stores a value with an expiration. Pass 0 for no expiration. + Set(ctx context.Context, key, value string, ttl time.Duration) error + + // Incr atomically increments a simple key by delta and returns the new value. + Incr(ctx context.Context, key string, delta int64) (int64, error) + + // HashIncr atomically increments a field in a hash by delta. + HashIncr(ctx context.Context, key, field string, delta int64) (int64, error) + + // HashGet returns all fields of a hash as a string map. + HashGet(ctx context.Context, key string) (map[string]string, error) + + // HashSetNX sets a field only if it does not already exist. Returns true if set. + HashSetNX(ctx context.Context, key, field, value string) (bool, error) + + // Expire sets a TTL on a key. No-op if the key does not exist. + Expire(ctx context.Context, key string, ttl time.Duration) error + + // Close releases the underlying connection. + Close() error +} diff --git a/authbridge/cmd/authbridge-envoy/go.mod b/authbridge/cmd/authbridge-envoy/go.mod index 501fa746..066430ba 100644 --- a/authbridge/cmd/authbridge-envoy/go.mod +++ b/authbridge/cmd/authbridge-envoy/go.mod @@ -2,11 +2,15 @@ module github.com/rossoctl/cortex/authbridge/cmd/authbridge-envoy go 1.26.4 -replace github.com/rossoctl/cortex/authbridge/authlib => ../../authlib +replace ( + github.com/rossoctl/cortex/authbridge/authlib => ../../authlib + github.com/rossoctl/cortex/authbridge/storage/redis => ../../storage/redis +) require ( github.com/envoyproxy/go-control-plane/envoy v1.37.0 - github.com/rossoctl/cortex/authbridge/authlib v0.0.0-00010101000000-000000000000 + github.com/rossoctl/cortex/authbridge/authlib v0.0.0 + github.com/rossoctl/cortex/authbridge/storage/redis v0.0.0 google.golang.org/grpc v1.82.0 ) @@ -40,6 +44,7 @@ require ( github.com/cloudwego/base64x v0.1.6 // indirect github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.1 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dlclark/regexp2 v1.11.5 // indirect github.com/envoyproxy/protoc-gen-validate v1.3.3 // indirect github.com/fsnotify/fsnotify v1.10.1 // indirect @@ -81,6 +86,7 @@ require ( github.com/prometheus/common v0.67.5 // indirect github.com/prometheus/procfs v0.20.1 // indirect github.com/rcrowley/go-metrics v0.0.0-20250401214520-65e299d6c5c9 // indirect + github.com/redis/go-redis/v9 v9.9.0 // indirect github.com/rossoctl/context-guru v0.0.0-20260720181432-8fc7c7b36563 // indirect github.com/segmentio/asm v1.2.1 // indirect github.com/sirupsen/logrus v1.9.4 // indirect diff --git a/authbridge/cmd/authbridge-envoy/go.sum b/authbridge/cmd/authbridge-envoy/go.sum index e1a9d034..3f7ac221 100644 --- a/authbridge/cmd/authbridge-envoy/go.sum +++ b/authbridge/cmd/authbridge-envoy/go.sum @@ -32,6 +32,8 @@ github.com/alexaandru/go-sitter-forest/tsx v1.9.2 h1:R6CIvxcs6zGF/nwI7YbHzRM1HuR github.com/alexaandru/go-sitter-forest/tsx v1.9.2/go.mod h1:JIGhuMRPOeeUQbqrNiDh3YZV9iMtWTHUmTgSH8WMTIk= github.com/alexaandru/go-sitter-forest/typescript v1.9.4 h1:k+zE1JbmcDjgqPxO0fVnCnsCFj0yWmRaLpp2sbC4MoA= github.com/alexaandru/go-sitter-forest/typescript v1.9.4/go.mod h1:fzlkFeml5odd1gUkYOgiNXK4bF2M6hBcfTitiJPlso8= +github.com/alicebob/miniredis/v2 v2.38.0 h1:nZAzCR+Lj+Vxk4ZXzm2NuKq2O33RXj1XxJ2e2uP9jiw= +github.com/alicebob/miniredis/v2 v2.38.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM= github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0 h1:jfIu9sQUG6Ig+0+Ap1h4unLjW6YQJpKZVmUzxsD4E/Q= @@ -40,6 +42,10 @@ github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPn github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk= github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/bytecodealliance/wasmtime-go/v44 v44.0.0 h1:WRZXnLPIer/TWs5aYPaMlmVcOlzmR6Ur6wjLRIQOhTQ= @@ -68,6 +74,8 @@ github.com/dgraph-io/badger/v4 v4.9.2 h1:Wb5qw8gElqwV1a8msHTeQKova9b1V10heFKMIiP github.com/dgraph-io/badger/v4 v4.9.2/go.mod h1:nJjaJTUOSsQEBhsq209FmwCvMJzEA3e74RjZw6V2pQI= github.com/dgraph-io/ristretto/v2 v2.2.0 h1:bkY3XzJcXoMuELV8F+vS8kzNgicwQFAaGINAEJdWGOM= github.com/dgraph-io/ristretto/v2 v2.2.0/go.mod h1:RZrm63UmcBAaYWC1DotLYBmTvgkrs0+XhBd7Npn7/zI= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dgryski/trifles v0.0.0-20230903005119-f50d829f2e54 h1:SG7nF6SRlWhcT7cNTs5R6Hk4V2lcmLz2NsG2VnInyNo= github.com/dgryski/trifles v0.0.0-20230903005119-f50d829f2e54/go.mod h1:if7Fbed8SFyPtHLHbg49SI7NAdJiC5WIA09pe59rfAA= github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= @@ -187,6 +195,8 @@ github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEy github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= github.com/rcrowley/go-metrics v0.0.0-20250401214520-65e299d6c5c9 h1:bsUq1dX0N8AOIL7EB/X911+m4EHsnWEHeJ0c+3TTBrg= github.com/rcrowley/go-metrics v0.0.0-20250401214520-65e299d6c5c9/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4= +github.com/redis/go-redis/v9 v9.9.0 h1:URbPQ4xVQSQhZ27WMQVmZSo3uT3pL+4IdHVcYq2nVfM= +github.com/redis/go-redis/v9 v9.9.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw= github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/rossoctl/context-guru v0.0.0-20260720181432-8fc7c7b36563 h1:YsOMYgqAxFzy9mxzmSC+1dWfC36riFQrV7vce2AtjUo= @@ -273,6 +283,8 @@ github.com/yashtewari/glob-intersection v0.2.0 h1:8iuHdN88yYuCzCdjt0gDe+6bAhUwBe github.com/yashtewari/glob-intersection v0.2.0/go.mod h1:LK7pIC3piUjovexikBbJ26Yml7g8xa5bsjfx2v1fwok= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= +github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 h1:8tvICD4vSTOOsNrsI4Ljf6C+6UKvpTEH5XY3JMoyPoo= diff --git a/authbridge/cmd/authbridge-envoy/plugins_tokenbudget.go b/authbridge/cmd/authbridge-envoy/plugins_tokenbudget.go new file mode 100644 index 00000000..93a74754 --- /dev/null +++ b/authbridge/cmd/authbridge-envoy/plugins_tokenbudget.go @@ -0,0 +1,10 @@ +//go:build include_plugin_tokenbudget + +// token-budget is opt-IN: it pulls the storage/redis module (go-redis) +// into the binary. Build with -tags include_plugin_tokenbudget to link it. +package main + +import ( + _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/tokenbudget" + _ "github.com/rossoctl/cortex/authbridge/storage/redis" +) diff --git a/authbridge/cmd/authbridge-proxy/Dockerfile b/authbridge/cmd/authbridge-proxy/Dockerfile index f1da5eb4..a27440c8 100644 --- a/authbridge/cmd/authbridge-proxy/Dockerfile +++ b/authbridge/cmd/authbridge-proxy/Dockerfile @@ -28,6 +28,7 @@ RUN apk add --no-cache git WORKDIR /app COPY authlib/ authlib/ +COPY storage/ storage/ COPY cmd/authbridge-proxy/ cmd/authbridge-proxy/ ARG GO_BUILD_TAGS="" diff --git a/authbridge/cmd/authbridge-proxy/go.mod b/authbridge/cmd/authbridge-proxy/go.mod index 7166cca4..1d0a6f8f 100644 --- a/authbridge/cmd/authbridge-proxy/go.mod +++ b/authbridge/cmd/authbridge-proxy/go.mod @@ -2,7 +2,10 @@ module github.com/rossoctl/cortex/authbridge/cmd/authbridge-proxy go 1.26.4 -require github.com/rossoctl/cortex/authbridge/authlib v0.0.0-00010101000000-000000000000 +require ( + github.com/rossoctl/cortex/authbridge/authlib v0.0.0 + github.com/rossoctl/cortex/authbridge/storage/redis v0.0.0 +) require ( github.com/Microsoft/go-winio v0.6.2 // indirect @@ -33,6 +36,7 @@ require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.1 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dlclark/regexp2 v1.11.5 // indirect github.com/fsnotify/fsnotify v1.10.1 // indirect github.com/go-jose/go-jose/v4 v4.1.4 // indirect @@ -72,6 +76,7 @@ require ( github.com/prometheus/common v0.67.5 // indirect github.com/prometheus/procfs v0.20.1 // indirect github.com/rcrowley/go-metrics v0.0.0-20250401214520-65e299d6c5c9 // indirect + github.com/redis/go-redis/v9 v9.9.0 // indirect github.com/rossoctl/context-guru v0.0.0-20260720181432-8fc7c7b36563 // indirect github.com/segmentio/asm v1.2.1 // indirect github.com/sirupsen/logrus v1.9.4 // indirect @@ -118,4 +123,7 @@ require ( sigs.k8s.io/yaml v1.6.0 // indirect ) -replace github.com/rossoctl/cortex/authbridge/authlib => ../../authlib +replace ( + github.com/rossoctl/cortex/authbridge/authlib => ../../authlib + github.com/rossoctl/cortex/authbridge/storage/redis => ../../storage/redis +) diff --git a/authbridge/cmd/authbridge-proxy/go.sum b/authbridge/cmd/authbridge-proxy/go.sum index ce2fa127..f5a34628 100644 --- a/authbridge/cmd/authbridge-proxy/go.sum +++ b/authbridge/cmd/authbridge-proxy/go.sum @@ -32,6 +32,8 @@ github.com/alexaandru/go-sitter-forest/tsx v1.9.2 h1:R6CIvxcs6zGF/nwI7YbHzRM1HuR github.com/alexaandru/go-sitter-forest/tsx v1.9.2/go.mod h1:JIGhuMRPOeeUQbqrNiDh3YZV9iMtWTHUmTgSH8WMTIk= github.com/alexaandru/go-sitter-forest/typescript v1.9.4 h1:k+zE1JbmcDjgqPxO0fVnCnsCFj0yWmRaLpp2sbC4MoA= github.com/alexaandru/go-sitter-forest/typescript v1.9.4/go.mod h1:fzlkFeml5odd1gUkYOgiNXK4bF2M6hBcfTitiJPlso8= +github.com/alicebob/miniredis/v2 v2.38.0 h1:nZAzCR+Lj+Vxk4ZXzm2NuKq2O33RXj1XxJ2e2uP9jiw= +github.com/alicebob/miniredis/v2 v2.38.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM= github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0 h1:jfIu9sQUG6Ig+0+Ap1h4unLjW6YQJpKZVmUzxsD4E/Q= @@ -40,6 +42,10 @@ github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPn github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk= github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/bytecodealliance/wasmtime-go/v44 v44.0.0 h1:WRZXnLPIer/TWs5aYPaMlmVcOlzmR6Ur6wjLRIQOhTQ= @@ -66,6 +72,8 @@ github.com/dgraph-io/badger/v4 v4.9.2 h1:Wb5qw8gElqwV1a8msHTeQKova9b1V10heFKMIiP github.com/dgraph-io/badger/v4 v4.9.2/go.mod h1:nJjaJTUOSsQEBhsq209FmwCvMJzEA3e74RjZw6V2pQI= github.com/dgraph-io/ristretto/v2 v2.2.0 h1:bkY3XzJcXoMuELV8F+vS8kzNgicwQFAaGINAEJdWGOM= github.com/dgraph-io/ristretto/v2 v2.2.0/go.mod h1:RZrm63UmcBAaYWC1DotLYBmTvgkrs0+XhBd7Npn7/zI= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dgryski/trifles v0.0.0-20230903005119-f50d829f2e54 h1:SG7nF6SRlWhcT7cNTs5R6Hk4V2lcmLz2NsG2VnInyNo= github.com/dgryski/trifles v0.0.0-20230903005119-f50d829f2e54/go.mod h1:if7Fbed8SFyPtHLHbg49SI7NAdJiC5WIA09pe59rfAA= github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= @@ -179,6 +187,8 @@ github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEy github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= github.com/rcrowley/go-metrics v0.0.0-20250401214520-65e299d6c5c9 h1:bsUq1dX0N8AOIL7EB/X911+m4EHsnWEHeJ0c+3TTBrg= github.com/rcrowley/go-metrics v0.0.0-20250401214520-65e299d6c5c9/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4= +github.com/redis/go-redis/v9 v9.9.0 h1:URbPQ4xVQSQhZ27WMQVmZSo3uT3pL+4IdHVcYq2nVfM= +github.com/redis/go-redis/v9 v9.9.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw= github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/rossoctl/context-guru v0.0.0-20260720181432-8fc7c7b36563 h1:YsOMYgqAxFzy9mxzmSC+1dWfC36riFQrV7vce2AtjUo= @@ -265,6 +275,8 @@ github.com/yashtewari/glob-intersection v0.2.0 h1:8iuHdN88yYuCzCdjt0gDe+6bAhUwBe github.com/yashtewari/glob-intersection v0.2.0/go.mod h1:LK7pIC3piUjovexikBbJ26Yml7g8xa5bsjfx2v1fwok= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= +github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 h1:8tvICD4vSTOOsNrsI4Ljf6C+6UKvpTEH5XY3JMoyPoo= diff --git a/authbridge/cmd/authbridge-proxy/plugins_tokenbudget.go b/authbridge/cmd/authbridge-proxy/plugins_tokenbudget.go new file mode 100644 index 00000000..93a74754 --- /dev/null +++ b/authbridge/cmd/authbridge-proxy/plugins_tokenbudget.go @@ -0,0 +1,10 @@ +//go:build include_plugin_tokenbudget + +// token-budget is opt-IN: it pulls the storage/redis module (go-redis) +// into the binary. Build with -tags include_plugin_tokenbudget to link it. +package main + +import ( + _ "github.com/rossoctl/cortex/authbridge/authlib/plugins/tokenbudget" + _ "github.com/rossoctl/cortex/authbridge/storage/redis" +) diff --git a/authbridge/docs/token-budget-plugin.md b/authbridge/docs/token-budget-plugin.md new file mode 100644 index 00000000..3c69d236 --- /dev/null +++ b/authbridge/docs/token-budget-plugin.md @@ -0,0 +1,123 @@ +# token-budget Plugin + +Enforces per-session lifetime budgets on tokens, inference calls, and +wall-clock duration. Runs in the outbound pipeline alongside `inference-parser`. +Uses Redis for cross-pod durable counters; evaluates limits from a local +cache with zero I/O on the hot path. Returns HTTP 403 on breach (lifetime +cap is permanently exhausted; retrying won't help). + +## Build Tag + +This plugin is **opt-IN**. Build with `-tags include_plugin_tokenbudget` +to include it (and its `storage/redis` dependency) in the binary: + +```bash +cd authbridge +docker build -f cmd/authbridge-proxy/Dockerfile \ + --build-arg GO_BUILD_TAGS="include_plugin_tokenbudget" \ + -t authbridge:latest . +``` + +Without the tag, neither token-budget nor go-redis are linked. + +## Configuration + +```yaml +pipeline: + outbound: + plugins: + - name: token-exchange + config: { ... } + - name: token-budget + config: + redis_url: "redis://valkey.infra.svc:6379" + max_tokens: 50000 + max_calls: 100 + max_duration_seconds: 1800 + - name: inference-parser +``` + +| Field | Required | Default | Description | +|-------|----------|---------|-------------| +| `redis_url` | yes | — | Redis/Valkey connection URL | +| `max_tokens` | no | 0 | Cumulative token ceiling per session (0 = no limit) | +| `max_calls` | no | 0 | Max inference calls per session (0 = no limit) | +| `max_duration_seconds` | no | 0 | Wall-clock session lifetime in seconds (0 = no limit) | +| `on_exceed` | no | "deny" | `deny` (block with 403) or `observe` (shadow — log but continue) | +| `session_ttl_seconds` | no | 7200 | Redis key TTL; should be >= `max_duration_seconds` | +| `refresh_interval` | no | "5s" | How often to sync local cache from Redis | +| `redis_unavailable` | no | "fail_open" | `fail_open` (suppress refresh warnings) or `fail_closed` (log warnings; stale cache retained until Redis recovers) | + +At least one of `max_tokens`, `max_calls`, or `max_duration_seconds` must be > 0. + +## Shadow Mode + +Set `on_exceed: "observe"` to run the plugin in shadow mode. The plugin +still accumulates counters and evaluates limits, but instead of blocking +requests it logs a WARN and continues the pipeline. Use this to calibrate +limits under real workloads before enabling enforcement. + +Rollout workflow: +1. Deploy with `on_exceed: "observe"` and conservative limits +2. Monitor logs for `"budget exceeded (shadow mode)"` entries +3. Adjust `max_tokens` / `max_calls` / `max_duration_seconds` based on observed patterns +4. Flip to `on_exceed: "deny"` when confident in the thresholds + +## Pipeline Position + +Must be declared **before** `inference-parser` in the outbound plugin list. +The response path runs in reverse order, so inference-parser finalizes token +counts first, then token-budget reads them. Both plugins implement +`StreamingResponder` so they work with streaming SSE responses (Ollama, +LiteLLM, OpenAI) and buffered JSON responses alike. + +## Response Format (403) + +```json +{ + "error": "budget.exceeded", + "message": "token limit reached: 50200/50000", + "details": { + "spent_tokens": 50200, + "spent_calls": 42, + "limit_tokens": 50000, + "limit_calls": 100 + } +} +``` + +## Redis Key Schema + +``` +token-budget: (Hash, TTL = session_ttl_seconds) + tokens int cumulative TotalTokens + calls int inference call count + started_at unix first-call timestamp (set-if-not-exists) +``` + +## Failure Modes + +| Scenario | Behavior | +|----------|----------| +| Redis down at startup | Pod fails to start (`Init` error) | +| Redis fails mid-session | Local cache continues enforcing; writes dropped silently | +| Pod restarts | First request passes (cold cache); refresh picks up Redis counters within one interval | +| `fail_closed` + refresh failure | Stale cache retained; enforcement lags until Redis recovers | +| Provider returns no usage data | `max_tokens` not enforced; `max_calls` and `max_duration_seconds` still work | + +**Note on token counting:** Token accumulation requires the LLM provider to +return `usage` (prompt/completion token counts) in responses. Providers that +omit usage from streaming chunks (e.g. Anthropic via LiteLLM) will show +`promptTokens=0` in inference-parser logs — `max_tokens` enforcement won't +trigger for these providers, but `max_calls` and `max_duration_seconds` still +apply. Ollama, OpenAI, and Azure OpenAI include usage in streaming responses +and work fully. + +## Testing + +```bash +cd authbridge/authlib +go test ./plugins/tokenbudget/... -v -count=1 +``` + +No external dependencies — tests use in-memory stores. diff --git a/authbridge/go.work b/authbridge/go.work index ccedcd8b..03ea6bcb 100644 --- a/authbridge/go.work +++ b/authbridge/go.work @@ -6,4 +6,5 @@ use ( ./cmd/authbridge-cpex ./cmd/authbridge-envoy ./cmd/authbridge-proxy + ./storage/redis ) diff --git a/authbridge/storage/redis/go.mod b/authbridge/storage/redis/go.mod new file mode 100644 index 00000000..4d8532c1 --- /dev/null +++ b/authbridge/storage/redis/go.mod @@ -0,0 +1,17 @@ +module github.com/rossoctl/cortex/authbridge/storage/redis + +go 1.26.4 + +require ( + github.com/alicebob/miniredis/v2 v2.38.0 + github.com/redis/go-redis/v9 v9.9.0 + github.com/rossoctl/cortex/authbridge/authlib v0.0.0 +) + +require ( + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/yuin/gopher-lua v1.1.1 // indirect +) + +replace github.com/rossoctl/cortex/authbridge/authlib => ../../authlib diff --git a/authbridge/storage/redis/go.sum b/authbridge/storage/redis/go.sum new file mode 100644 index 00000000..728c4e5f --- /dev/null +++ b/authbridge/storage/redis/go.sum @@ -0,0 +1,14 @@ +github.com/alicebob/miniredis/v2 v2.38.0 h1:nZAzCR+Lj+Vxk4ZXzm2NuKq2O33RXj1XxJ2e2uP9jiw= +github.com/alicebob/miniredis/v2 v2.38.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/redis/go-redis/v9 v9.9.0 h1:URbPQ4xVQSQhZ27WMQVmZSo3uT3pL+4IdHVcYq2nVfM= +github.com/redis/go-redis/v9 v9.9.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw= +github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= +github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= diff --git a/authbridge/storage/redis/redis.go b/authbridge/storage/redis/redis.go new file mode 100644 index 00000000..dd6e620e --- /dev/null +++ b/authbridge/storage/redis/redis.go @@ -0,0 +1,71 @@ +package redis + +import ( + "context" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/storage" + goredis "github.com/redis/go-redis/v9" +) + +func init() { + storage.Register("redis", func(url string) (storage.Store, error) { + return New(url) + }) +} + +var _ storage.Store = (*Client)(nil) + +// Client implements storage.Store backed by Redis/Valkey. +type Client struct { + rdb *goredis.Client +} + +// New creates a Redis-backed Store from a connection URL. +func New(redisURL string) (*Client, error) { + opts, err := goredis.ParseURL(redisURL) + if err != nil { + return nil, err + } + return &Client{rdb: goredis.NewClient(opts)}, nil +} + +func (c *Client) Get(ctx context.Context, key string) (string, error) { + result, err := c.rdb.Get(ctx, key).Result() + if err == goredis.Nil { + return "", nil + } + return result, err +} + +func (c *Client) Set(ctx context.Context, key, value string, ttl time.Duration) error { + return c.rdb.Set(ctx, key, value, ttl).Err() +} + +func (c *Client) Incr(ctx context.Context, key string, delta int64) (int64, error) { + return c.rdb.IncrBy(ctx, key, delta).Result() +} + +func (c *Client) HashIncr(ctx context.Context, key, field string, delta int64) (int64, error) { + return c.rdb.HIncrBy(ctx, key, field, delta).Result() +} + +func (c *Client) HashGet(ctx context.Context, key string) (map[string]string, error) { + result, err := c.rdb.HGetAll(ctx, key).Result() + if err != nil { + return nil, err + } + return result, nil +} + +func (c *Client) HashSetNX(ctx context.Context, key, field, value string) (bool, error) { + return c.rdb.HSetNX(ctx, key, field, value).Result() +} + +func (c *Client) Expire(ctx context.Context, key string, ttl time.Duration) error { + return c.rdb.Expire(ctx, key, ttl).Err() +} + +func (c *Client) Close() error { + return c.rdb.Close() +} diff --git a/authbridge/storage/redis/redis_test.go b/authbridge/storage/redis/redis_test.go new file mode 100644 index 00000000..5b1659e7 --- /dev/null +++ b/authbridge/storage/redis/redis_test.go @@ -0,0 +1,186 @@ +package redis + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" +) + +func setup(t *testing.T) (*Client, *miniredis.Miniredis) { + t.Helper() + mr := miniredis.RunT(t) + c, err := New("redis://" + mr.Addr()) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { c.Close() }) + return c, mr +} + +func TestGetSet(t *testing.T) { + c, _ := setup(t) + ctx := context.Background() + + val, err := c.Get(ctx, "missing") + if err != nil { + t.Fatalf("Get missing: %v", err) + } + if val != "" { + t.Errorf("Get missing = %q, want empty", val) + } + + if err := c.Set(ctx, "key1", "hello", 0); err != nil { + t.Fatalf("Set: %v", err) + } + + val, err = c.Get(ctx, "key1") + if err != nil { + t.Fatalf("Get: %v", err) + } + if val != "hello" { + t.Errorf("Get = %q, want hello", val) + } +} + +func TestSetWithTTL(t *testing.T) { + c, mr := setup(t) + ctx := context.Background() + + if err := c.Set(ctx, "ephemeral", "value", 10*time.Second); err != nil { + t.Fatalf("Set: %v", err) + } + + mr.FastForward(11 * time.Second) + + val, err := c.Get(ctx, "ephemeral") + if err != nil { + t.Fatalf("Get after TTL: %v", err) + } + if val != "" { + t.Errorf("expected empty after TTL expiry, got %q", val) + } +} + +func TestIncr(t *testing.T) { + c, _ := setup(t) + ctx := context.Background() + + v, err := c.Incr(ctx, "counter", 5) + if err != nil { + t.Fatalf("Incr: %v", err) + } + if v != 5 { + t.Errorf("Incr = %d, want 5", v) + } + + v, err = c.Incr(ctx, "counter", 3) + if err != nil { + t.Fatalf("Incr: %v", err) + } + if v != 8 { + t.Errorf("Incr = %d, want 8", v) + } +} + +func TestHashIncr(t *testing.T) { + c, _ := setup(t) + ctx := context.Background() + + v, err := c.HashIncr(ctx, "h1", "tokens", 100) + if err != nil { + t.Fatalf("HashIncr: %v", err) + } + if v != 100 { + t.Errorf("HashIncr = %d, want 100", v) + } + + v, err = c.HashIncr(ctx, "h1", "tokens", 50) + if err != nil { + t.Fatalf("HashIncr: %v", err) + } + if v != 150 { + t.Errorf("HashIncr = %d, want 150", v) + } +} + +func TestHashGet(t *testing.T) { + c, _ := setup(t) + ctx := context.Background() + + fields, err := c.HashGet(ctx, "empty") + if err != nil { + t.Fatalf("HashGet: %v", err) + } + if len(fields) != 0 { + t.Errorf("HashGet empty = %v, want empty map", fields) + } + + c.HashIncr(ctx, "h1", "tokens", 42) + c.HashIncr(ctx, "h1", "calls", 3) + + fields, err = c.HashGet(ctx, "h1") + if err != nil { + t.Fatalf("HashGet: %v", err) + } + if fields["tokens"] != "42" { + t.Errorf("tokens = %q, want 42", fields["tokens"]) + } + if fields["calls"] != "3" { + t.Errorf("calls = %q, want 3", fields["calls"]) + } +} + +func TestHashSetNX(t *testing.T) { + c, _ := setup(t) + ctx := context.Background() + + set, err := c.HashSetNX(ctx, "h1", "started_at", "1000") + if err != nil { + t.Fatalf("HashSetNX: %v", err) + } + if !set { + t.Error("HashSetNX should return true on first set") + } + + set, err = c.HashSetNX(ctx, "h1", "started_at", "2000") + if err != nil { + t.Fatalf("HashSetNX: %v", err) + } + if set { + t.Error("HashSetNX should return false when field exists") + } + + fields, _ := c.HashGet(ctx, "h1") + if fields["started_at"] != "1000" { + t.Errorf("started_at = %q, want 1000 (first value preserved)", fields["started_at"]) + } +} + +func TestExpire(t *testing.T) { + c, mr := setup(t) + ctx := context.Background() + + c.HashIncr(ctx, "h1", "tokens", 100) + if err := c.Expire(ctx, "h1", 5*time.Second); err != nil { + t.Fatalf("Expire: %v", err) + } + + mr.FastForward(6 * time.Second) + + fields, err := c.HashGet(ctx, "h1") + if err != nil { + t.Fatalf("HashGet after TTL: %v", err) + } + if len(fields) != 0 { + t.Errorf("expected empty after expiry, got %v", fields) + } +} + +func TestNew_InvalidURL(t *testing.T) { + _, err := New("not-a-valid-url") + if err == nil { + t.Error("expected error for invalid URL") + } +}