From a1460364b5f8259fd4d8e9d051d6da1e1547e119 Mon Sep 17 00:00:00 2001 From: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Date: Sun, 18 Oct 2026 04:33:36 +0530 Subject: [PATCH 1/2] feat(ratelimit): enforce gateway admission Reject exhausted keys before delegation verification. Hold tenant slots until responses and SSE streams finish. Keep admission denials inside the authenticated receipt path. Signed-off-by: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 38 +++- cmd/agentgw/main.go | 7 +- cmd/agentgw/rate_limits.go | 46 +++++ cmd/agentgw/rate_limits_test.go | 75 +++++++ docs/registry.md | 31 ++- internal/gateway/gateway.go | 53 +++-- internal/gateway/gateway_test.go | 7 +- internal/gateway/ratelimit_test.go | 262 ++++++++++++++++++++++++ internal/ratelimit/admission_test.go | 244 ++++++++++++++++++++++ internal/ratelimit/limiter.go | 281 ++++++++++++++++++++++++-- internal/registry/rate_limits_test.go | 99 +++++++++ internal/registry/registry.go | 31 ++- 12 files changed, 1134 insertions(+), 40 deletions(-) create mode 100644 cmd/agentgw/rate_limits.go create mode 100644 cmd/agentgw/rate_limits_test.go create mode 100644 internal/gateway/ratelimit_test.go create mode 100644 internal/ratelimit/admission_test.go create mode 100644 internal/registry/rate_limits_test.go diff --git a/README.md b/README.md index d27a78c..d8ce3ce 100644 --- a/README.md +++ b/README.md @@ -315,6 +315,11 @@ AgentGate is deterministic. Agents never see tokens, calls go only to configured | `AGENTGATE_ADMIN_SECRET` | Secret for admin API access | Yes | | `AGENTGATE_PUBLIC_URL` | Base URL used to build the OAuth callback (default `http://localhost:8080`) | No | | `AGENTGATE_API_KEY_CACHE_TTL` | Verified-key cache duration. Default `30s`; `0` disables caching. | No | +| `AGENTGATE_RATE_KEY_RPS` | Requests per second per API key ID. Default `20`; `0` disables this bucket. | No | +| `AGENTGATE_RATE_KEY_BURST` | Per-key token bucket burst. Default `40`; `0` disables this bucket. | No | +| `AGENTGATE_RATE_TENANT_RPS` | Requests per second per human principal. Default `100`; `0` disables this bucket. | No | +| `AGENTGATE_RATE_TENANT_BURST` | Per-tenant token bucket burst. Default `200`; `0` disables this bucket. | No | +| `AGENTGATE_TENANT_MAX_INFLIGHT` | In-flight upstream requests per human principal. Default `32`; `0` disables this cap. | No | | `AGENTGATE_BISCUIT_MAX_DURATION` | Go duration per Biscuit Datalog world run. Default `250ms`; range `1ms` through `5s`, inclusive. Invalid values stop startup. | No | | `GITHUB_CLIENT_ID` / `GITHUB_CLIENT_SECRET` | GitHub OAuth credentials | For GitHub OAuth | | `SLACK_CLIENT_ID` / `SLACK_CLIENT_SECRET` | Slack OAuth credentials | For Slack OAuth | @@ -332,6 +337,37 @@ SQLite stores a public ID and a peppered HMAC-SHA256, never the secret. Use `POS Legacy `ag_live_` rows keep their bcrypt hashes. Successful validation adds an indexed HMAC without changing the presented key. Revoked rows never enter legacy bcrypt scans. +#### Request admission + +API keys have no organization ID. The organization store binds admins, not agent keys. +Tenant identity is `on_behalf_of`, the requested human principal. +Different keys acting for that principal share its rate bucket and concurrency cap. +Each key's bucket spans services and human principals. + +Rate settings accept finite, non-negative numbers. Bursts and concurrency caps require integers. +Invalid environment values stop startup. Either zero rate or zero burst disables that bucket. +Per-service limits are optional. Set `rate_limit` under a service in the loaded registry YAML. +They apply per API key and service, alongside the global limits. +See [registry rate limits](docs/registry.md#rate-limits) for YAML examples. + +Token buckets run after API key authentication and body validation, before delegation verification or vault access. +Rate denials return HTTP 429 with `rate_limited`. +`Retry-After` rounds the reservation delay up to whole seconds, with a minimum of 1. +Concurrency admission runs immediately before upstream dispatch. +Full tenants return HTTP 429 with `concurrency_limited` and `Retry-After: 1`. +Slots remain held through response completion, including SSE streams and receipt commits. +Failures, disconnects, and stream termination release them. +Every authenticated, schema-valid attempt still commits a receipt before its response, including admission denials. +Both admission errors record policy `rate_limited`. Receipt failures return `receipt_failed` instead. + +Admission state is process-local. Replicas don't share quotas. +The combined bucket and semaphore cache holds at most 100000 entries. +Requests trigger idle sweeps at most once per minute. Entries idle for 10 minutes become eviction candidates. +Only full buckets without in-flight requests can be evicted. Slow buckets retain their quota until fully refilled. +At capacity, a full, inactive LRU entry can be replaced. +If that entry cannot be evicted, admission fails closed with HTTP 429 and `Retry-After: 1`. +Restart the gateway after changing admission settings or service YAML. + #### Biscuit token security caps Tokens are limited to 4096 bytes and 8 attenuation blocks. @@ -344,7 +380,7 @@ Authority checks follow the same caps. These checks run before authorizer construction. Rejections match both `delegation.ErrDenied` and `delegation.ErrTokenShape`. Errors contain fixed reasons, not token contents. -The structural caps bound fact joins and block processing before the request reaches the rate limiter. +The early rate limiter runs before these checks. Structural caps then bound fact joins and block processing. Biscuit-go v2.2.0 deadlines do not cover check evaluation and cannot reliably stop timed-out goroutines. Runtime limits alone do not bound per-request CPU. diff --git a/cmd/agentgw/main.go b/cmd/agentgw/main.go index 4e4ef6f..7d31a0b 100644 --- a/cmd/agentgw/main.go +++ b/cmd/agentgw/main.go @@ -62,6 +62,11 @@ func main() { logger.Error("load Biscuit limits", "error", err) os.Exit(1) } + rateLimits, err := loadRateLimits() + if err != nil { + logger.Error("load rate limits", "error", err) + os.Exit(1) + } // Load service registry. reg := registry.New() @@ -141,7 +146,7 @@ func main() { ledger := receipt.NewLedger(database, signerStore) defer ledger.Close() - limiter := ratelimit.New(nil) // no per-service limits configured yet; Allow() is a pass-through + limiter := ratelimit.New(reg.RateLimits(), ratelimit.WithSettings(rateLimits)) delegationRoot, err := delegation.NewRootStore(database, masterKey) if err != nil { diff --git a/cmd/agentgw/rate_limits.go b/cmd/agentgw/rate_limits.go new file mode 100644 index 0000000..14da0f2 --- /dev/null +++ b/cmd/agentgw/rate_limits.go @@ -0,0 +1,46 @@ +package main + +import ( + "fmt" + "math" + "os" + "strconv" + + "github.com/Clawdlinux/agentgate/internal/ratelimit" +) + +func loadRateLimits() (ratelimit.Settings, error) { + settings := ratelimit.DefaultSettings() + for _, field := range []struct { + name string + target *float64 + }{ + {"AGENTGATE_RATE_KEY_RPS", &settings.KeyRPS}, + {"AGENTGATE_RATE_TENANT_RPS", &settings.TenantRPS}, + } { + if value, ok := os.LookupEnv(field.name); ok { + parsed, err := strconv.ParseFloat(value, 64) + if err != nil || parsed < 0 || math.IsNaN(parsed) || math.IsInf(parsed, 0) { + return ratelimit.Settings{}, fmt.Errorf("%s must be a finite non-negative number", field.name) + } + *field.target = parsed + } + } + for _, field := range []struct { + name string + target *int + }{ + {"AGENTGATE_RATE_KEY_BURST", &settings.KeyBurst}, + {"AGENTGATE_RATE_TENANT_BURST", &settings.TenantBurst}, + {"AGENTGATE_TENANT_MAX_INFLIGHT", &settings.TenantMaxInflight}, + } { + if value, ok := os.LookupEnv(field.name); ok { + parsed, err := strconv.Atoi(value) + if err != nil || parsed < 0 { + return ratelimit.Settings{}, fmt.Errorf("%s must be a non-negative integer", field.name) + } + *field.target = parsed + } + } + return settings, nil +} diff --git a/cmd/agentgw/rate_limits_test.go b/cmd/agentgw/rate_limits_test.go new file mode 100644 index 0000000..e1d8fa8 --- /dev/null +++ b/cmd/agentgw/rate_limits_test.go @@ -0,0 +1,75 @@ +package main + +import ( + "os" + "strings" + "testing" + + "github.com/Clawdlinux/agentgate/internal/ratelimit" +) + +var rateEnvNames = []string{ + "AGENTGATE_RATE_KEY_RPS", "AGENTGATE_RATE_KEY_BURST", + "AGENTGATE_RATE_TENANT_RPS", "AGENTGATE_RATE_TENANT_BURST", + "AGENTGATE_TENANT_MAX_INFLIGHT", +} + +func clearRateEnv(t *testing.T) { + t.Helper() + for _, name := range rateEnvNames { + t.Setenv(name, "") + if err := os.Unsetenv(name); err != nil { + t.Fatal(err) + } + } +} + +func TestLoadRateLimits(t *testing.T) { + clearRateEnv(t) + settings, err := loadRateLimits() + if err != nil || settings != ratelimit.DefaultSettings() { + t.Fatalf("defaults = %+v, err = %v", settings, err) + } + for _, name := range rateEnvNames { + t.Setenv(name, "0") + } + settings, err = loadRateLimits() + if err != nil || settings != (ratelimit.Settings{}) { + t.Fatalf("disabled = %+v, err = %v", settings, err) + } + t.Setenv("AGENTGATE_RATE_KEY_RPS", "0.5") + t.Setenv("AGENTGATE_RATE_KEY_BURST", "2") + t.Setenv("AGENTGATE_RATE_TENANT_RPS", "12.5") + t.Setenv("AGENTGATE_RATE_TENANT_BURST", "25") + t.Setenv("AGENTGATE_TENANT_MAX_INFLIGHT", "7") + settings, err = loadRateLimits() + want := ratelimit.Settings{KeyRPS: 0.5, KeyBurst: 2, TenantRPS: 12.5, TenantBurst: 25, TenantMaxInflight: 7} + if err != nil || settings != want { + t.Fatalf("overrides = %+v, err = %v", settings, err) + } +} + +func TestLoadRateLimitsInvalid(t *testing.T) { + for _, name := range rateEnvNames { + for _, value := range []string{"", "-1", "NaN", "Inf", "1e999", "invalid"} { + t.Run(name+"/"+value, func(t *testing.T) { + clearRateEnv(t) + t.Setenv(name, value) + _, err := loadRateLimits() + if err == nil || !strings.Contains(err.Error(), name) { + t.Fatalf("expected named validation error, got %v", err) + } + }) + } + if strings.Contains(name, "RPS") { + continue + } + t.Run(name+"/fraction", func(t *testing.T) { + clearRateEnv(t) + t.Setenv(name, "1.5") + if _, err := loadRateLimits(); err == nil { + t.Fatal("fractional integer accepted") + } + }) + } +} diff --git a/docs/registry.md b/docs/registry.md index aa3033e..bca34a0 100644 --- a/docs/registry.md +++ b/docs/registry.md @@ -1,8 +1,37 @@ -# Registry response settings +# Registry settings Response limits belong under each service's `upstream` settings. Streaming permission belongs on each action. +## Rate limits + +Optional `rate_limit` settings apply per API key and service. +They combine with the gateway's per-key and per-human-principal buckets. +The gateway loads them from the same registry YAML as actions. +The separate `configs/agentgate.yaml` example isn't loaded by the gateway. +Restart the gateway after changing these settings. + +```yaml +services: + example: + base_url: https://api.example.com + auth: {type: bearer} + rate_limit: + requests_per_second: 25 + burst: 50 + actions: + list: {method: GET, path: /items} +``` + +Use at most 1 positive rate field: `requests_per_second`, `requests_per_minute`, or `requests_per_hour`. +Rates must be finite and non-negative. Fractional rates are accepted. +Absent or zero rates disable the service bucket. Omitted or zero `burst` selects 10 tokens. +Negative bursts and conflicting rates fail registry loading. +Rate rejection returns HTTP 429 with `rate_limited` and a whole-second `Retry-After` header. +The gateway commits its receipt before sending that response. + +## Response settings + ```yaml services: example: diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 1ffd3df..75b157b 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -56,11 +56,11 @@ type AgentAuthorizer interface { Validate(ctx context.Context, plaintext string) (*auth.AgentKey, error) } -// RequestLimiter reports whether a request for (agentKeyID, service) is -// currently allowed. Backed by *ratelimit.Limiter in production. A nil -// Limiter in Config disables rate limiting. +// RequestLimiter applies token buckets before authorization work and semaphore slots before upstream dispatch. +// Zero delay means admission succeeded. Acquire returns an idempotent release function on success. type RequestLimiter interface { - Allow(agentKeyID, service string) bool + Check(agentKeyID, tenant, service string) time.Duration + Acquire(tenant string) (release func(), retryAfter time.Duration) } // DelegationVerifier checks a presented Biscuit delegation token against @@ -201,6 +201,7 @@ type attempt struct { acceptSSE bool delegationToken []byte // raw, not yet verified; set by prepareAttempt, verified by executeAttempt delegationChain []string // set by executeAttempt once delegationToken verifies; empty for a direct grant (DELG-04) + releaseSlot func() } // immediateResponse is written with no receipt: the request never reached @@ -224,6 +225,7 @@ type outcome struct { policyDecision string // "allow", "deny", or "rate_limited" errorCode string // stable code; empty means no error stream *streamResponse + retryAfter time.Duration } func (s *Server) handleAct(w http.ResponseWriter, r *http.Request) { @@ -238,6 +240,9 @@ func (s *Server) handleAct(w http.ResponseWriter, r *http.Request) { att.acceptSSE = false } out := s.executeAttempt(r.Context(), att) + if att.releaseSlot != nil { + defer att.releaseSlot() + } if out.stream != nil { defer out.stream.close() stopCancel := context.AfterFunc(r.Context(), out.stream.close) @@ -293,6 +298,10 @@ func (s *Server) handleAct(w http.ResponseWriter, r *http.Request) { s.writeStream(w, r, att, out.stream) return } + if out.retryAfter > 0 { + seconds := max(int64(1), int64(math.Ceil(out.retryAfter.Seconds()))) + w.Header().Set("Retry-After", strconv.FormatInt(seconds, 10)) + } if actResp, ok := out.body.(ActResponse); ok { actResp.LatencyMS = latencyMS out.body = actResp @@ -349,11 +358,17 @@ func (s *Server) prepareAttempt(r *http.Request) (*attempt, *immediateResponse) } // executeAttempt runs every check and side effect after authentication and -// shape validation: delegation, scope, rate limit, registry lookup, vault +// shape validation: rate limits, delegation, scope, registry lookup, vault // fetch, and upstream dispatch. It never writes to an http.ResponseWriter — // every exit is captured in outcome so handleAct can commit a receipt // before any response is written. func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { + // AgentKey has no organization binding. Tenant identity is the requested human principal. + if s.cfg.Limiter != nil { + if retry := s.cfg.Limiter.Check(att.agentKey.ID, att.req.OnBehalfOf, att.req.Service); retry > 0 { + return limitedOutcome("rate_limited", retry) + } + } if len(att.delegationToken) > 0 { // Verified before registry, vault, or upstream access (DELG-01); a // gateway with no configured verifier fails closed rather than @@ -387,15 +402,6 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { } } - if s.cfg.Limiter != nil && !s.cfg.Limiter.Allow(att.agentKey.ID, att.req.Service) { - return &outcome{ - status: http.StatusTooManyRequests, - body: ErrorResponse{Error: "rate limit exceeded", Code: "rate_limited"}, - policyDecision: "rate_limited", - errorCode: "rate_limited", - } - } - svc, action, err := s.cfg.Registry.GetAction(att.req.Service, att.req.Action) if err != nil { return &outcome{ @@ -485,6 +491,14 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { dispatch = s.upstreams.DoStream } + if s.cfg.Limiter != nil { + release, retry := s.cfg.Limiter.Acquire(att.req.OnBehalfOf) + if retry > 0 { + return limitedOutcome("concurrency_limited", retry) + } + // handleAct releases after response writing, including streams and receipt failures. + att.releaseSlot = release + } resp, err := dispatch(svc.Name, svc.Upstream, upstream, action.Idempotent) if err != nil { status, code, message := http.StatusBadGateway, "upstream_error", "upstream request failed" @@ -545,6 +559,17 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { } } +func limitedOutcome(code string, retry time.Duration) *outcome { + message := "rate limit exceeded" + if code == "concurrency_limited" { + message = "tenant concurrency limit exceeded" + } + return &outcome{ + status: http.StatusTooManyRequests, body: ErrorResponse{Error: message, Code: code}, + policyDecision: "rate_limited", errorCode: code, retryAfter: retry, + } +} + func extractAPIKey(r *http.Request) string { key := r.Header.Get("Authorization") key = strings.TrimPrefix(key, "Bearer ") diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index 3c1e30b..07de358 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -284,6 +284,9 @@ func TestAct_RateLimited(t *testing.T) { if w.Code != http.StatusTooManyRequests { t.Fatalf("status = %d, want 429", w.Code) } + if w.Header().Get("Retry-After") != "1" { + t.Fatalf("Retry-After = %q, want 1", w.Header().Get("Retry-After")) + } if got := receipts.last().PolicyDecision; got != "rate_limited" { t.Fatalf("policy_decision = %s, want rate_limited", got) } @@ -291,7 +294,9 @@ func TestAct_RateLimited(t *testing.T) { type denyAllLimiter struct{} -func (denyAllLimiter) Allow(string, string) bool { return false } +func (denyAllLimiter) Check(string, string, string) time.Duration { return time.Second } + +func (denyAllLimiter) Acquire(string) (func(), time.Duration) { return func() {}, 0 } func TestAct_NoToken(t *testing.T) { t.Parallel() diff --git a/internal/gateway/ratelimit_test.go b/internal/gateway/ratelimit_test.go new file mode 100644 index 0000000..565241f --- /dev/null +++ b/internal/gateway/ratelimit_test.go @@ -0,0 +1,262 @@ +package gateway + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Clawdlinux/agentgate/internal/auth" + "github.com/Clawdlinux/agentgate/internal/ratelimit" + upstreamclient "github.com/Clawdlinux/agentgate/internal/upstream" +) + +type countingDelegation struct{ calls atomic.Int32 } + +func (v *countingDelegation) Verify([]byte, string, string, string, string, time.Time) ([]string, error) { + v.calls.Add(1) + return nil, errors.New("denied") +} + +func limitedAct(srv *Server, key, tenant, service string) *httptest.ResponseRecorder { + req := httptest.NewRequest("POST", "/v1/act", strings.NewReader(fmt.Sprintf( + `{"service":%q,"action":"list_repos","on_behalf_of":%q}`, service, tenant))) + req.Header.Set("Authorization", "Bearer "+key) + req.Header.Set("X-Agentgate-Delegation", "AA==") + w := httptest.NewRecorder() + srv.ServeHTTP(w, req) + return w +} + +func TestAct_EarlyRateLimitsAndReceipts(t *testing.T) { + for _, scope := range []string{"key", "tenant", "service"} { + t.Run(scope, func(t *testing.T) { + now := time.Unix(1000, 0) + settings := ratelimit.Settings{} + configs := map[string]ratelimit.Config{} + switch scope { + case "key": + settings.KeyRPS, settings.KeyBurst = 0.4, 1 + case "tenant": + settings.TenantRPS, settings.TenantBurst = 0.4, 1 + case "service": + configs["github"] = ratelimit.Config{RequestsPerSecond: 0.4, Burst: 1} + } + l := ratelimit.New(configs, ratelimit.WithSettings(settings), + ratelimit.WithClock(func() time.Time { return now })) + l.Check("test-agent", "alice", "github") + srv, _, receipts := testSetup(t) + srv.cfg.Limiter = l + verifier := &countingDelegation{} + srv.cfg.Delegation = verifier + // Any registry, URL, or vault access would panic. + srv.cfg.Registry, srv.cfg.Vault = nil, nil + key := "test-key" + if scope == "tenant" { + key = "second-key" + srv.cfg.Authorizer.(*fakeAuthorizer).keys[key] = &auth.AgentKey{ID: "second-agent"} + } + w := limitedAct(srv, key, "alice", "github") + if w.Code != 429 || w.Header().Get("Retry-After") != "3" || + !strings.Contains(w.Body.String(), `"code":"rate_limited"`) { + t.Fatalf("status=%d retry=%q body=%s", w.Code, w.Header().Get("Retry-After"), w.Body) + } + if verifier.calls.Load() != 0 { + t.Fatalf("delegation calls = %d", verifier.calls.Load()) + } + draft := receipts.last() + if receipts.count() != 1 || draft.StatusCode != 429 || draft.Error != "rate_limited" || + draft.PolicyDecision != "rate_limited" || draft.HumanPrincipal != "alice" { + t.Fatalf("receipt = %+v", draft) + } + receipts.failNext = true + w = limitedAct(srv, key, "alice", "github") + if w.Code != 500 || !strings.Contains(w.Body.String(), `"code":"receipt_failed"`) || + w.Header().Get("Retry-After") != "" { + t.Fatalf("failed receipt response = %d %s", w.Code, w.Body) + } + now = now.Add(2500 * time.Millisecond) + w = limitedAct(srv, key, "alice", "github") + if w.Code != 403 || verifier.calls.Load() != 1 { + t.Fatalf("refill status=%d delegation calls=%d", w.Code, verifier.calls.Load()) + } + }) + } +} + +func TestAct_TenantConcurrencyAndSSE(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) { + entered := make(chan struct{}, 2) + finish := make(chan struct{}) + var calls atomic.Int32 + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + entered <- struct{}{} + if stream { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: first\n\n") + w.(http.Flusher).Flush() + } + select { + case <-finish: + fmt.Fprint(w, "{}") + case <-r.Context().Done(): + } + })) + defer up.Close() + defer close(finish) + srv, receipts := bodyGateway(t, up.URL, "{}", stream) + now := time.Unix(1000, 0) + srv.cfg.Limiter = ratelimit.New(nil, ratelimit.WithSettings(ratelimit.Settings{TenantMaxInflight: 1}), + ratelimit.WithClock(func() time.Time { return now })) + srv.cfg.Authorizer.(*fakeAuthorizer).keys["second-key"] = &auth.AgentKey{ + ID: "second-agent", AllowedServices: []string{"*"}, AllowedUsers: []string{"*"}, + } + firstDone := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Test-First") != "" { + defer close(firstDone) + } + srv.ServeHTTP(w, r) + })) + defer server.Close() + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + type result struct { + resp *http.Response + err error + } + resultCh := make(chan result, 1) + go func() { + req := streamRequest(ctx, server.URL, "text/event-stream") + req.Header.Set("X-Test-First", "1") + resp, err := server.Client().Do(req) + resultCh <- result{resp, err} + }() + awaitSignal(t, entered) + var first result + if stream { + first = <-resultCh + if first.err != nil { + t.Fatal(first.err) + } + defer first.resp.Body.Close() + } + req := streamRequest(ctx, server.URL, "text/event-stream") + req.Header.Set("Authorization", "Bearer second-key") + resp, err := server.Client().Do(req) + if err != nil { + t.Fatal(err) + } + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil || resp.StatusCode != 429 || resp.Header.Get("Retry-After") != "1" || + !strings.Contains(string(body), `"code":"concurrency_limited"`) || calls.Load() != 1 { + t.Fatalf("status=%d body=%s calls=%d err=%v", resp.StatusCode, body, calls.Load(), err) + } + if receipts.last().Error != "concurrency_limited" || receipts.last().PolicyDecision != "rate_limited" { + t.Fatalf("receipt = %+v", receipts.last()) + } + finish <- struct{}{} + if !stream { + first = <-resultCh + if first.err != nil { + t.Fatal(first.err) + } + defer first.resp.Body.Close() + } + if _, err := io.ReadAll(first.resp.Body); err != nil { + t.Fatal(err) + } + awaitSignal(t, firstDone) + if first.resp.StatusCode != 200 { + t.Fatalf("first status = %d", first.resp.StatusCode) + } + release, delay := srv.cfg.Limiter.Acquire("user-1") + if delay != 0 { + t.Fatal("request completion did not release the tenant slot") + } + release() + if receipts.count() != 2 { + t.Fatalf("receipts = %d", receipts.count()) + } + }) + } +} + +func TestAct_ConcurrencyReleaseOnFailures(t *testing.T) { + for _, kind := range []string{"upstream", "body", "receipt", "stream-receipt"} { + t.Run(kind, func(t *testing.T) { + srv, receipts := bodyGateway(t, "http://unused", "{}", kind == "stream-receipt") + srv.cfg.Limiter = ratelimit.New(nil, ratelimit.WithSettings(ratelimit.Settings{TenantMaxInflight: 1})) + srv.upstreams = upstreamclient.New(upstreamclient.Options{Client: &http.Client{ + Transport: responseTransport(func(*http.Request) (*http.Response, error) { + if kind == "upstream" { + return nil, errors.New("upstream failed") + } + var reader io.Reader = strings.NewReader("{}") + if kind == "body" { + reader = brokenResponseReader{} + } + header := make(http.Header) + header.Set("Content-Type", "text/event-stream") + return &http.Response{StatusCode: 200, Header: header, Body: io.NopCloser(reader)}, nil + }), + }}) + receipts.failNext = strings.Contains(kind, "receipt") + req := streamRequest(t.Context(), "http://gateway", "text/event-stream") + w := httptest.NewRecorder() + srv.ServeHTTP(w, req) + if w.Code != 500 && w.Code != 502 { + t.Fatalf("status = %d", w.Code) + } + release, delay := srv.cfg.Limiter.Acquire("user-1") + if delay != 0 { + t.Fatal("failure leaked a concurrency slot") + } + release() + }) + } +} + +func TestAct_SSEDisconnectReleasesTenantSlot(t *testing.T) { + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: first\n\n") + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + defer up.Close() + srv, _ := bodyGateway(t, up.URL, "{}", true) + srv.cfg.Limiter = ratelimit.New(nil, ratelimit.WithSettings(ratelimit.Settings{TenantMaxInflight: 1})) + done := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(done) + srv.ServeHTTP(w, r) + })) + defer server.Close() + ctx, cancel := context.WithTimeout(t.Context(), 3*time.Second) + defer cancel() + resp, err := server.Client().Do(streamRequest(ctx, server.URL, "text/event-stream")) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if _, delay := srv.cfg.Limiter.Acquire("user-1"); delay == 0 { + t.Fatal("stream should hold the tenant slot") + } + cancel() + awaitSignal(t, done) + release, delay := srv.cfg.Limiter.Acquire("user-1") + if delay != 0 { + t.Fatal("disconnect leaked the tenant slot") + } + release() +} diff --git a/internal/ratelimit/admission_test.go b/internal/ratelimit/admission_test.go new file mode 100644 index 0000000..6bad0e8 --- /dev/null +++ b/internal/ratelimit/admission_test.go @@ -0,0 +1,244 @@ +package ratelimit + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestAdmissionBuckets(t *testing.T) { + for _, scope := range []string{"key", "tenant", "service"} { + t.Run(scope, func(t *testing.T) { + now := time.Unix(1000, 0) + settings := Settings{} + configs := map[string]Config{} + key, tenant, service := "key-1", "alice", "github" + switch scope { + case "key": + settings.KeyRPS, settings.KeyBurst = 0.4, 1 + case "tenant": + settings.TenantRPS, settings.TenantBurst = 0.4, 1 + case "service": + configs[service] = Config{RequestsPerSecond: 0.4, Burst: 1} + } + l := New(configs, WithSettings(settings), WithClock(func() time.Time { return now })) + if delay := l.Check(key, tenant, service); delay != 0 { + t.Fatalf("first delay = %v", delay) + } + if scope == "tenant" { + key = "key-2" + } + if delay := l.Check(key, tenant, service); delay != 2500*time.Millisecond { + t.Fatalf("exhausted delay = %v", delay) + } + now = now.Add(500 * time.Millisecond) + if delay := l.Check(key, tenant, service); delay != 2*time.Second { + t.Fatalf("rejected reservations must be canceled, delay = %v", delay) + } + now = now.Add(2 * time.Second) + if delay := l.Check(key, tenant, service); delay != 0 { + t.Fatalf("refilled delay = %v", delay) + } + }) + } +} + +func TestAdmissionIsolationAndDisabled(t *testing.T) { + now := time.Unix(1000, 0) + l := New(map[string]Config{"github": {RequestsPerSecond: 1, Burst: 1}}, + WithSettings(Settings{}), WithClock(func() time.Time { return now })) + for _, pair := range [][2]string{{"a:b", "github"}, {"a", "github"}, {"other", "github"}} { + if l.Check(pair[0], "alice", pair[1]) != 0 { + t.Fatal("separate keys must have separate service buckets") + } + } + if l.Check("a", "alice", "GITHUB") == 0 { + t.Fatal("service case must not bypass the registry limit") + } + for i := 0; i < 100; i++ { + if l.Check("a", "alice", "unconfigured") != 0 { + t.Fatal("disabled key and tenant buckets must allow requests") + } + } + release, delay := l.Acquire("alice") + if delay != 0 { + t.Fatal("zero concurrency cap must disable admission") + } + release() +} + +func TestAdmissionRollsBackOtherBuckets(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1, KeyBurst: 1, TenantRPS: 1, TenantBurst: 1}), + WithClock(func() time.Time { return now })) + l.Check("a", "alice", "github") + if l.Check("b", "alice", "github") == 0 { + t.Fatal("tenant should be exhausted") + } + if l.Check("b", "bob", "github") != 0 { + t.Fatal("denial must not spend the unused key token") + } +} + +func TestAdmissionCapacityAndIdleEviction(t *testing.T) { + now := time.Unix(1000, 0) + l := New(map[string]Config{"github": {RequestsPerSecond: 1, Burst: 1}}, + WithSettings(Settings{KeyRPS: 1, KeyBurst: 1, TenantRPS: 1, TenantBurst: 1}), + WithClock(func() time.Time { return now })) + l.maxEntries = 3 + if l.Check("a", "alice", "github") != 0 || len(l.entries) != 3 { + t.Fatalf("initial entries = %d", len(l.entries)) + } + for i := 0; i < 100; i++ { + if l.Check(fmt.Sprint(i), "alice", "github") == 0 { + t.Fatal("capacity cannot evict a depleted bucket") + } + if len(l.entries) > l.maxEntries { + t.Fatal("entry cap exceeded") + } + } + now = now.Add(11 * time.Minute) + if l.Check("b", "bob", "github") != 0 { + t.Fatal("idle full buckets should be evicted") + } + if len(l.entries) != 3 { + t.Fatalf("entries after sweep = %d", len(l.entries)) + } + if _, ok := l.entries[bucketKey{scope: keyScope, id: "a"}]; ok { + t.Fatal("idle key survived the sweep") + } +} + +func TestAdmissionDoesNotResetSlowBuckets(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1.0 / 3600, KeyBurst: 1}), + WithClock(func() time.Time { return now })) + l.maxEntries = 1 + l.Check("a", "alice", "github") + now = now.Add(11 * time.Minute) + if l.Check("b", "bob", "github") == 0 || l.Check("a", "alice", "github") == 0 { + t.Fatal("idle eviction must not reset an unrefilled bucket") + } + now = now.Add(time.Hour) + if l.Check("b", "bob", "github") != 0 { + t.Fatal("a full least-recently-used bucket can be evicted") + } +} + +func TestConcurrencyCapacityAndRelease(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{TenantMaxInflight: 1}), + WithClock(func() time.Time { return now })) + l.maxEntries = 1 + release, delay := l.Acquire("alice") + if delay != 0 { + t.Fatal(delay) + } + now = now.Add(time.Hour) + if _, delay := l.Acquire("bob"); delay == 0 { + t.Fatal("active semaphore cannot be evicted") + } + if _, delay := l.Acquire("alice"); delay == 0 { + t.Fatal("tenant cap exceeded") + } + release() + release() + next, delay := l.Acquire("bob") + if delay != 0 { + t.Fatal("released slot should permit eviction") + } + next() +} + +func TestAdmissionConcurrent(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1, KeyBurst: 10, TenantRPS: 1, TenantBurst: 5}), + WithClock(func() time.Time { return now })) + var allowed atomic.Int32 + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Go(func() { + if l.Check("a", "alice", "github") == 0 { + allowed.Add(1) + } + }) + } + wg.Wait() + if allowed.Load() != 5 { + t.Fatalf("allowed = %d", allowed.Load()) + } +} + +func TestAdmissionDefaultsAndIndependentScopes(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithClock(func() time.Time { return now })) + for i := 0; i < 40; i++ { + if l.Check("a", "alice", "github") != 0 { + t.Fatalf("default burst rejected request %d", i) + } + } + if l.Check("a", "bob", "stripe") == 0 { + t.Fatal("key limit must span tenants and services") + } + if l.Check("b", "alice", "github") != 0 { + t.Fatal("other keys must retain their own buckets") + } + l = New(nil, WithSettings(Settings{TenantRPS: 1, TenantBurst: 1}), WithClock(func() time.Time { return now })) + if l.Check("a", "alice", "github") != 0 || l.Check("b", "alice", "stripe") == 0 { + t.Fatal("tenant limit must span keys and services") + } + if l.Check("b", "bob", "github") != 0 { + t.Fatal("other tenants must retain their own buckets") + } +} + +func TestZeroRateOrBurstDisablesBucket(t *testing.T) { + for _, settings := range []Settings{ + {KeyRPS: 1, TenantRPS: 1}, + {KeyBurst: 1, TenantBurst: 1}, + } { + l := New(nil, WithSettings(settings)) + for i := 0; i < 100; i++ { + if l.Check("a", "alice", "github") != 0 { + t.Fatal("zero rate or burst should disable the bucket") + } + } + if len(l.entries) != 0 { + t.Fatal("disabled buckets should not allocate entries") + } + } +} + +func TestConcurrencySeparateTenants(t *testing.T) { + l := New(nil, WithSettings(Settings{TenantMaxInflight: 1})) + release, delay := l.Acquire("alice") + if delay != 0 { + t.Fatal(delay) + } + defer release() + releaseOther, delay := l.Acquire("bob") + if delay != 0 { + t.Fatal("separate tenants must have separate semaphores") + } + defer releaseOther() + if _, delay := l.Acquire("alice"); delay == 0 { + t.Fatal("first tenant cap should still apply") + } +} + +func TestKeyRejectionDoesNotAllocateTenantState(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1, KeyBurst: 1, TenantRPS: 1, TenantBurst: 1}), + WithClock(func() time.Time { return now })) + l.Check("a", "alice", "github") + for i := 0; i < 100; i++ { + if l.Check("a", fmt.Sprint(i), "github") == 0 { + t.Fatal("key limit should reject each attempt") + } + } + if len(l.entries) != 2 { + t.Fatalf("key-limited requests allocated tenant state: %d entries", len(l.entries)) + } +} diff --git a/internal/ratelimit/limiter.go b/internal/ratelimit/limiter.go index aa64553..299c5ea 100644 --- a/internal/ratelimit/limiter.go +++ b/internal/ratelimit/limiter.go @@ -1,11 +1,17 @@ package ratelimit import ( + "container/list" "encoding/json" + "fmt" + "math" "net/http" + "strings" "sync" + "time" "golang.org/x/time/rate" + "gopkg.in/yaml.v3" ) // Config defines rate limit for a service. @@ -16,6 +22,50 @@ type Config struct { Burst int `yaml:"burst"` } +// UnmarshalYAML prevents YAML's numeric coercion from silently truncating fractional bursts. +func (c *Config) UnmarshalYAML(node *yaml.Node) error { + type plain Config + var decoded plain + if err := node.Decode(&decoded); err != nil { + return err + } + var fields struct { + Burst yaml.Node `yaml:"burst"` + } + if err := node.Decode(&fields); err != nil { + return err + } + burst := &fields.Burst + for burst.Kind == yaml.AliasNode { + burst = burst.Alias + } + if burst.Kind != 0 && burst.Tag != "!!int" { + return fmt.Errorf("rate_limit burst must be an integer") + } + *c = Config(decoded) + return nil +} + +// Validate rejects ambiguous or invalid service settings. +func (c Config) Validate() error { + rates := 0 + for _, value := range []float64{c.RequestsPerSecond, c.RequestsPerMinute, c.RequestsPerHour} { + if value < 0 || math.IsNaN(value) || math.IsInf(value, 0) { + return fmt.Errorf("rate_limit rates must be finite and non-negative") + } + if value > 0 { + rates++ + } + } + if rates > 1 || c.Burst < 0 { + return fmt.Errorf("rate_limit requires at most 1 positive rate and a non-negative burst") + } + if rates == 1 && c.Rate() <= 0 { + return fmt.Errorf("rate_limit rate is too small") + } + return nil +} + // Rate returns the rate.Limit based on the config. func (c Config) Rate() rate.Limit { if c.RequestsPerSecond > 0 { @@ -38,31 +88,228 @@ func (c Config) BurstSize() int { return 10 } -// Limiter implements per-(agent, service) rate limiting. +// Settings controls gateway-wide admission. Zero rate or burst disables that bucket. +type Settings struct { + KeyRPS float64 + KeyBurst int + TenantRPS float64 + TenantBurst int + TenantMaxInflight int +} + +// DefaultSettings enables per-key and per-tenant admission. +func DefaultSettings() Settings { + return Settings{KeyRPS: 20, KeyBurst: 40, TenantRPS: 100, TenantBurst: 200, TenantMaxInflight: 32} +} + +type bucketScope uint8 + +const ( + keyScope bucketScope = iota + tenantScope + serviceScope +) + +// Separate fields prevent collisions between identities containing separators. +type bucketKey struct { + scope bucketScope + id string + service string +} + +type entry struct { + key bucketKey + bucket *rate.Limiter + slots chan struct{} + lastUsed time.Time +} + +// Limiter shares a bounded LRU across key, tenant, and key-service buckets. +// Active semaphores and buckets that haven't refilled are never evicted. type Limiter struct { - configs map[string]Config // service -> config - limiters sync.Map // "agentKeyID:service" -> *rate.Limiter + mu sync.Mutex + configs map[string]Config + settings Settings + now func() time.Time + entries map[bucketKey]*list.Element + lru *list.List + maxEntries int + nextSweep time.Time } -// New creates a rate limiter from service configs. -func New(configs map[string]Config) *Limiter { - return &Limiter{configs: configs} +type Option func(*Limiter) + +func WithSettings(settings Settings) Option { + return func(l *Limiter) { l.settings = settings } +} + +// WithClock supplies explicit reservation times, including deterministic test clocks. +func WithClock(now func() time.Time) Option { + return func(l *Limiter) { l.now = now } +} + +// New copies service settings. Global defaults are enabled unless overridden. +func New(configs map[string]Config, options ...Option) *Limiter { + l := &Limiter{ + configs: make(map[string]Config, len(configs)), settings: DefaultSettings(), now: time.Now, + entries: make(map[bucketKey]*list.Element), lru: list.New(), maxEntries: 100_000, + } + for name, cfg := range configs { + l.configs[strings.ToLower(name)] = cfg + } + for _, option := range options { + option(l) + } + return l } -// Allow checks if a request is allowed for the given agent and service. +// Allow preserves the standalone middleware's per-key, per-service contract. +// Gateway admission uses Check and Acquire instead. func (l *Limiter) Allow(agentKeyID, service string) bool { - key := agentKeyID + ":" + service - lim, ok := l.limiters.Load(key) - if !ok { - cfg, exists := l.configs[service] - if !exists { - return true // no rate limit configured + l.mu.Lock() + defer l.mu.Unlock() + now := l.now() + service = strings.ToLower(service) + cfg, ok := l.configs[service] + if !ok || cfg.Rate() == rate.Inf { + return true + } + e := l.get(bucketKey{scope: serviceScope, id: agentKeyID, service: service}, cfg.Rate(), cfg.BurstSize(), now) + if e == nil { + return false + } + reservation := e.bucket.ReserveN(now, 1) + if reservation.DelayFrom(now) > 0 { + reservation.CancelAt(now) + return false + } + return true +} + +// Check reserves 1 token in every enabled bucket. Zero means admission succeeded. +// Rejections cancel all reservations, so retries don't postpone recovery. +func (l *Limiter) Check(agentKeyID, tenant, service string) time.Duration { + l.mu.Lock() + defer l.mu.Unlock() + now := l.now() + service = strings.ToLower(service) + type spec struct { + key bucketKey + rps rate.Limit + burst int + } + specs := []spec{ + {bucketKey{scope: keyScope, id: agentKeyID}, rate.Limit(l.settings.KeyRPS), l.settings.KeyBurst}, + {bucketKey{scope: tenantScope, id: tenant}, rate.Limit(l.settings.TenantRPS), l.settings.TenantBurst}, + } + if cfg, ok := l.configs[service]; ok && cfg.Rate() != rate.Inf { + specs = append(specs, spec{ + bucketKey{scope: serviceScope, id: agentKeyID, service: service}, cfg.Rate(), cfg.BurstSize(), + }) + } + var reservations []*rate.Reservation + var retry time.Duration + for _, spec := range specs { + if spec.rps <= 0 || spec.burst <= 0 { + continue } - newLim := rate.NewLimiter(cfg.Rate(), cfg.BurstSize()) - actual, _ := l.limiters.LoadOrStore(key, newLim) - lim = actual + e := l.get(spec.key, spec.rps, spec.burst, now) + if e == nil { + retry = time.Second + break + } + reservation := e.bucket.ReserveN(now, 1) + reservations = append(reservations, reservation) + retry = reservation.DelayFrom(now) + if retry > 0 { + break // Rejected keys must not allocate state for arbitrary tenant names. + } + } + if retry > 0 { + for _, reservation := range reservations { + reservation.CancelAt(now) + } + } + return retry +} + +// Acquire takes a tenant semaphore slot without waiting. The release function is idempotent. +// A concurrency rejection retries after 1s because completion times are unknown. +func (l *Limiter) Acquire(tenant string) (func(), time.Duration) { + l.mu.Lock() + defer l.mu.Unlock() + if l.settings.TenantMaxInflight <= 0 { + return func() {}, 0 + } + now := l.now() + e := l.get(bucketKey{scope: tenantScope, id: tenant}, rate.Limit(l.settings.TenantRPS), l.settings.TenantBurst, now) + if e == nil { + return nil, time.Second } - return lim.(*rate.Limiter).Allow() + select { + case e.slots <- struct{}{}: + var once sync.Once + return func() { + once.Do(func() { + l.mu.Lock() + defer l.mu.Unlock() + <-e.slots + e.lastUsed = l.now() + l.lru.MoveToFront(l.entries[e.key]) + }) + }, 0 + default: + return nil, time.Second + } +} + +func (l *Limiter) get(key bucketKey, rps rate.Limit, burst int, now time.Time) *entry { + if !now.Before(l.nextSweep) { + // Sweep lazily at most once per minute. No background goroutine needs shutdown. + for element := l.lru.Back(); element != nil; { + previous := element.Prev() + e := element.Value.(*entry) + if now.Sub(e.lastUsed) < 10*time.Minute { + break + } + if e.evictable(now) { + l.remove(element) + } + element = previous + } + l.nextSweep = now.Add(time.Minute) + } + if element := l.entries[key]; element != nil { + e := element.Value.(*entry) + e.lastUsed = now + l.lru.MoveToFront(element) + return e + } + if len(l.entries) >= l.maxEntries { + oldest := l.lru.Back() + if oldest == nil || !oldest.Value.(*entry).evictable(now) { + return nil // Fail closed rather than reset an exhausted bucket or active semaphore. + } + l.remove(oldest) + } + e := &entry{key: key, lastUsed: now} + if rps > 0 && burst > 0 { + e.bucket = rate.NewLimiter(rps, burst) + } + if key.scope == tenantScope && l.settings.TenantMaxInflight > 0 { + e.slots = make(chan struct{}, l.settings.TenantMaxInflight) + } + l.entries[key] = l.lru.PushFront(e) + return e +} + +func (e *entry) evictable(now time.Time) bool { + return len(e.slots) == 0 && (e.bucket == nil || e.bucket.TokensAt(now) >= float64(e.bucket.Burst())) +} + +func (l *Limiter) remove(element *list.Element) { + delete(l.entries, element.Value.(*entry).key) + l.lru.Remove(element) } // Middleware wraps an http.Handler with rate limiting. diff --git a/internal/registry/rate_limits_test.go b/internal/registry/rate_limits_test.go new file mode 100644 index 0000000..39adee6 --- /dev/null +++ b/internal/registry/rate_limits_test.go @@ -0,0 +1,99 @@ +package registry + +import ( + "fmt" + "strings" + "testing" + "time" + + "github.com/Clawdlinux/agentgate/internal/ratelimit" +) + +func TestRegistryRateLimitSnapshot(t *testing.T) { + r := New() + if err := r.LoadBytes(sampleYAML); err != nil { + t.Fatal(err) + } + if len(r.RateLimits()) != 0 { + t.Fatal("existing YAML must keep service limits optional") + } + err := r.LoadBytes([]byte(` +services: + github: + base_url: https://api.github.com + rate_limit: {requests_per_minute: 30, burst: 1} + actions: + list_repos: {method: GET, path: /user/repos} +`)) + if err != nil { + t.Fatal(err) + } + snapshot := r.RateLimits() + if snapshot["github"].RequestsPerMinute != 30 || snapshot["github"].Burst != 1 { + t.Fatalf("snapshot = %+v", snapshot) + } + now := time.Unix(1000, 0) + l := ratelimit.New(snapshot, ratelimit.WithSettings(ratelimit.Settings{}), + ratelimit.WithClock(func() time.Time { return now })) + delete(snapshot, "github") + if len(r.RateLimits()) != 1 { + t.Fatal("snapshot mutation changed the registry") + } + if l.Check("a", "alice", "github") != 0 || l.Check("a", "alice", "github") != 2*time.Second { + t.Fatal("registry configuration was not applied to admission") + } + if l.Check("a", "alice", "stripe") != 0 { + t.Fatal("unconfigured services should remain unlimited") + } +} + +func TestRegistryInvalidRateLimits(t *testing.T) { + for _, settings := range []string{ + "requests_per_second: -1", + "requests_per_minute: .inf", + "requests_per_hour: .nan", + "requests_per_second: 2, requests_per_minute: 2", + "requests_per_second: 1, burst: -1", + "requests_per_second: 1, burst: 1.5", + "requests_per_hour: 5e-324", + } { + t.Run(settings, func(t *testing.T) { + r := New() + err := r.LoadBytes([]byte(fmt.Sprintf(` +services: + github: + base_url: https://api.github.com + rate_limit: {%s} +`, settings))) + if err == nil || (!strings.Contains(err.Error(), "github") && !strings.Contains(err.Error(), "parse")) { + t.Fatalf("expected service validation error, got %v", err) + } + if r.Count() != 0 { + t.Fatal("invalid service was partially loaded") + } + }) + } +} + +func TestRegistryRateLimitYAMLAliases(t *testing.T) { + for _, burst := range []string{"2", "1.5"} { + t.Run(burst, func(t *testing.T) { + r := New() + err := r.LoadBytes([]byte(fmt.Sprintf(` +defaults: &limits + requests_per_second: 1 + burst: &burst %s +services: + github: + base_url: https://api.github.com + rate_limit: {<<: *limits} + slack: + base_url: https://slack.com/api + rate_limit: {requests_per_second: 1, burst: *burst} +`, burst))) + if (err != nil) != (burst == "1.5") { + t.Fatalf("burst=%s err=%v", burst, err) + } + }) + } +} diff --git a/internal/registry/registry.go b/internal/registry/registry.go index e18741b..fc6286a 100644 --- a/internal/registry/registry.go +++ b/internal/registry/registry.go @@ -20,6 +20,7 @@ import ( "strings" "sync" + "github.com/Clawdlinux/agentgate/internal/ratelimit" "github.com/Clawdlinux/agentgate/internal/upstream" "gopkg.in/yaml.v3" ) @@ -32,11 +33,12 @@ var ErrActionNotFound = errors.New("registry: action not found") // Service describes one SaaS API endpoint. type Service struct { - Name string `yaml:"name"` - BaseURL string `yaml:"base_url"` - Auth AuthCfg `yaml:"auth"` - Actions map[string]Action `yaml:"actions"` - Upstream upstream.Config `yaml:"upstream,omitempty"` + Name string `yaml:"name"` + BaseURL string `yaml:"base_url"` + Auth AuthCfg `yaml:"auth"` + Actions map[string]Action `yaml:"actions"` + Upstream upstream.Config `yaml:"upstream,omitempty"` + RateLimit *ratelimit.Config `yaml:"rate_limit,omitempty"` } // AuthCfg holds the authentication configuration for a service. @@ -89,9 +91,15 @@ func (r *Registry) LoadBytes(data []byte) error { return fmt.Errorf("registry: parse: %w", err) } for name, svc := range cfg.Services { + if svc.RateLimit != nil { + if err := svc.RateLimit.Validate(); err != nil { + return fmt.Errorf("registry: service %s: %w", name, err) + } + } if err := svc.Upstream.Validate(); err != nil { return fmt.Errorf("registry: service %s: %w", name, err) } + for actionName, action := range svc.Actions { if _, query, ok := strings.Cut(action.Path, "?"); ok && strings.Contains(query, "{") { return fmt.Errorf("registry: service %s action %s: placeholders are not allowed after '?'", name, actionName) @@ -107,6 +115,19 @@ func (r *Registry) LoadBytes(data []byte) error { return nil } +// RateLimits returns a snapshot for gateway admission at startup. +func (r *Registry) RateLimits() map[string]ratelimit.Config { + r.mu.RLock() + defer r.mu.RUnlock() + configs := make(map[string]ratelimit.Config) + for name, svc := range r.services { + if svc.RateLimit != nil { + configs[name] = *svc.RateLimit + } + } + return configs +} + // Get returns a service by name. func (r *Registry) Get(name string) (Service, error) { r.mu.RLock() From 31d945997048532d3712547517f2d7d0559e3c16 Mon Sep 17 00:00:00 2001 From: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Date: Sun, 18 Oct 2026 04:54:43 +0530 Subject: [PATCH 2/2] fix(ratelimit): isolate tenant admission Authorize principal and service scopes before tenant charges to prevent cross-tenant quota exhaustion. Bound tenant identities and scan past pinned cache entries. Refund concurrency rejections without rewinding bucket clocks. Release acquired slots during panic unwinding. Signed-off-by: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 35 ++-- cmd/agentgw/rate_limits.go | 6 + cmd/agentgw/rate_limits_review_test.go | 33 ++++ docs/registry.md | 6 +- internal/gateway/admission.go | 24 +++ internal/gateway/admission_review_test.go | 226 ++++++++++++++++++++++ internal/gateway/gateway.go | 64 +++--- internal/gateway/gateway_test.go | 9 +- internal/gateway/ratelimit_test.go | 20 +- internal/ratelimit/limiter.go | 94 +++------ internal/ratelimit/permit.go | 131 +++++++++++++ internal/ratelimit/review_test.go | 193 ++++++++++++++++++ 12 files changed, 736 insertions(+), 105 deletions(-) create mode 100644 cmd/agentgw/rate_limits_review_test.go create mode 100644 internal/gateway/admission.go create mode 100644 internal/gateway/admission_review_test.go create mode 100644 internal/ratelimit/permit.go create mode 100644 internal/ratelimit/review_test.go diff --git a/README.md b/README.md index d8ce3ce..c1e2a57 100644 --- a/README.md +++ b/README.md @@ -315,11 +315,11 @@ AgentGate is deterministic. Agents never see tokens, calls go only to configured | `AGENTGATE_ADMIN_SECRET` | Secret for admin API access | Yes | | `AGENTGATE_PUBLIC_URL` | Base URL used to build the OAuth callback (default `http://localhost:8080`) | No | | `AGENTGATE_API_KEY_CACHE_TTL` | Verified-key cache duration. Default `30s`; `0` disables caching. | No | -| `AGENTGATE_RATE_KEY_RPS` | Requests per second per API key ID. Default `20`; `0` disables this bucket. | No | -| `AGENTGATE_RATE_KEY_BURST` | Per-key token bucket burst. Default `40`; `0` disables this bucket. | No | -| `AGENTGATE_RATE_TENANT_RPS` | Requests per second per human principal. Default `100`; `0` disables this bucket. | No | +| `AGENTGATE_RATE_KEY_RPS` | Requests per second per API key ID. Default `20`. Zero requires tenant admission disabled. | No | +| `AGENTGATE_RATE_KEY_BURST` | Per-key token bucket burst. Default `40`. Zero requires tenant admission disabled. | No | +| `AGENTGATE_RATE_TENANT_RPS` | Requests per second per tenant. Default `100`; `0` disables this bucket. | No | | `AGENTGATE_RATE_TENANT_BURST` | Per-tenant token bucket burst. Default `200`; `0` disables this bucket. | No | -| `AGENTGATE_TENANT_MAX_INFLIGHT` | In-flight upstream requests per human principal. Default `32`; `0` disables this cap. | No | +| `AGENTGATE_TENANT_MAX_INFLIGHT` | In-flight upstream requests per tenant. Default `32`; `0` disables this cap. | No | | `AGENTGATE_BISCUIT_MAX_DURATION` | Go duration per Biscuit Datalog world run. Default `250ms`; range `1ms` through `5s`, inclusive. Invalid values stop startup. | No | | `GITHUB_CLIENT_ID` / `GITHUB_CLIENT_SECRET` | GitHub OAuth credentials | For GitHub OAuth | | `SLACK_CLIENT_ID` / `SLACK_CLIENT_SECRET` | Slack OAuth credentials | For Slack OAuth | @@ -340,32 +340,45 @@ Revoked rows never enter legacy bcrypt scans. #### Request admission API keys have no organization ID. The organization store binds admins, not agent keys. -Tenant identity is `on_behalf_of`, the requested human principal. -Different keys acting for that principal share its rate bucket and concurrency cap. +For explicitly scoped keys, tenant identity is the authorized `on_behalf_of` principal. +Such keys share that principal's rate bucket and concurrency cap. Each key's bucket spans services and human principals. +Wildcard keys allowing `"*"` use 1 key-scoped tenant for every principal. +Different wildcard keys don't share human quotas. +Each key selects at most 64 tenant identities from its fixed scope list. +When the list exceeds 64 entries, only the first 63 retain shared human identity. +The remaining principals share 1 key-scoped fallback tenant. +Human and fallback identities use separate namespaces. No per-key tenant-tracking map is needed. Rate settings accept finite, non-negative numbers. Bursts and concurrency caps require integers. Invalid environment values stop startup. Either zero rate or zero burst disables that bucket. +Both key settings must stay positive while tenant rate limiting or tenant concurrency is enabled. Per-service limits are optional. Set `rate_limit` under a service in the loaded registry YAML. They apply per API key and service, alongside the global limits. See [registry rate limits](docs/registry.md#rate-limits) for YAML examples. -Token buckets run after API key authentication and body validation, before delegation verification or vault access. +The key bucket runs after API key authentication and body validation. +In-memory service and principal scope checks run next. +Scope denials consume only key tokens. Their receipts retain `scope_denied`. +Tenant and service buckets run after scope checks, before delegation verification or vault access. Rate denials return HTTP 429 with `rate_limited`. `Retry-After` rounds the reservation delay up to whole seconds, with a minimum of 1. Concurrency admission runs immediately before upstream dispatch. Full tenants return HTTP 429 with `concurrency_limited` and `Retry-After: 1`. +Concurrency rejection refunds the key, tenant, and service charges without losing elapsed refill. Slots remain held through response completion, including SSE streams and receipt commits. -Failures, disconnects, and stream termination release them. +Failures, disconnects, stream termination, and dispatch panics release them. Every authenticated, schema-valid attempt still commits a receipt before its response, including admission denials. Both admission errors record policy `rate_limited`. Receipt failures return `receipt_failed` instead. Admission state is process-local. Replicas don't share quotas. The combined bucket and semaphore cache holds at most 100000 entries. Requests trigger idle sweeps at most once per minute. Entries idle for 10 minutes become eviction candidates. -Only full buckets without in-flight requests can be evicted. Slow buckets retain their quota until fully refilled. -At capacity, a full, inactive LRU entry can be replaced. -If that entry cannot be evicted, admission fails closed with HTTP 429 and `Retry-After: 1`. +Only full buckets without pending permits or in-flight requests can be evicted. +Slow buckets retain their quota until fully refilled. +At capacity, admission scans up to 64 entries from the LRU tail. +Pinned or depleted entries move to the front. The first evictable entry is replaced. +If none qualifies, admission fails closed with HTTP 429 and `Retry-After: 1`. Restart the gateway after changing admission settings or service YAML. #### Biscuit token security caps diff --git a/cmd/agentgw/rate_limits.go b/cmd/agentgw/rate_limits.go index 14da0f2..43223fa 100644 --- a/cmd/agentgw/rate_limits.go +++ b/cmd/agentgw/rate_limits.go @@ -42,5 +42,11 @@ func loadRateLimits() (ratelimit.Settings, error) { *field.target = parsed } } + tenantEnabled := (settings.TenantRPS > 0 && settings.TenantBurst > 0) || settings.TenantMaxInflight > 0 + if tenantEnabled && (settings.KeyRPS == 0 || settings.KeyBurst == 0) { + return ratelimit.Settings{}, fmt.Errorf( + "AGENTGATE_RATE_KEY_RPS and AGENTGATE_RATE_KEY_BURST must be positive while tenant admission is enabled", + ) + } return settings, nil } diff --git a/cmd/agentgw/rate_limits_review_test.go b/cmd/agentgw/rate_limits_review_test.go new file mode 100644 index 0000000..63ba29f --- /dev/null +++ b/cmd/agentgw/rate_limits_review_test.go @@ -0,0 +1,33 @@ +package main + +import ( + "strings" + "testing" +) + +func TestTenantAdmissionRequiresKeyLimit(t *testing.T) { + for _, name := range []string{"AGENTGATE_RATE_KEY_RPS", "AGENTGATE_RATE_KEY_BURST"} { + for _, enabled := range []string{"rate", "inflight", "neither"} { + t.Run(name+"/"+enabled, func(t *testing.T) { + clearRateEnv(t) + t.Setenv(name, "0") + t.Setenv("AGENTGATE_RATE_TENANT_RPS", "0") + t.Setenv("AGENTGATE_TENANT_MAX_INFLIGHT", "0") + switch enabled { + case "rate": + t.Setenv("AGENTGATE_RATE_TENANT_RPS", "1") + case "inflight": + t.Setenv("AGENTGATE_TENANT_MAX_INFLIGHT", "1") + } + _, err := loadRateLimits() + if enabled == "neither" { + if err != nil { + t.Fatal(err) + } + } else if err == nil || !strings.Contains(err.Error(), name) { + t.Fatalf("expected key limit validation error, got %v", err) + } + }) + } + } +} diff --git a/docs/registry.md b/docs/registry.md index bca34a0..cd4ac86 100644 --- a/docs/registry.md +++ b/docs/registry.md @@ -6,7 +6,10 @@ Streaming permission belongs on each action. ## Rate limits Optional `rate_limit` settings apply per API key and service. -They combine with the gateway's per-key and per-human-principal buckets. +They combine with the gateway's per-key and per-tenant buckets. +Explicitly scoped keys share human principal tenants. Wildcard keys use key-scoped tenants. +Large scope lists use a key fallback beyond the first 63 principals. +See [request admission](../README.md#request-admission) for identity bounds and startup validation. The gateway loads them from the same registry YAML as actions. The separate `configs/agentgate.yaml` example isn't loaded by the gateway. Restart the gateway after changing these settings. @@ -29,6 +32,7 @@ Absent or zero rates disable the service bucket. Omitted or zero `burst` selects Negative bursts and conflicting rates fail registry loading. Rate rejection returns HTTP 429 with `rate_limited` and a whole-second `Retry-After` header. The gateway commits its receipt before sending that response. +Scope checks run before service and tenant buckets. Concurrency denials refund all token charges. ## Response settings diff --git a/internal/gateway/admission.go b/internal/gateway/admission.go new file mode 100644 index 0000000..f4d7350 --- /dev/null +++ b/internal/gateway/admission.go @@ -0,0 +1,24 @@ +package gateway + +import "github.com/Clawdlinux/agentgate/internal/auth" + +// tenantIdentity bounds each key to 64 tenant identities without a separate tracking cache. +// Wildcard keys share one key-scoped tenant across all requested principals. +func tenantIdentity(key *auth.AgentKey, principal string) string { + fallback := "key:" + key.ID + for _, allowed := range key.AllowedUsers { + if allowed == "*" { + return fallback + } + } + limit := 64 + if len(key.AllowedUsers) > limit { + limit-- // Reserve one identity for the overflow bucket. + } + for _, allowed := range key.AllowedUsers[:min(limit, len(key.AllowedUsers))] { + if allowed == principal { + return "human:" + principal + } + } + return fallback +} diff --git a/internal/gateway/admission_review_test.go b/internal/gateway/admission_review_test.go new file mode 100644 index 0000000..6f2626c --- /dev/null +++ b/internal/gateway/admission_review_test.go @@ -0,0 +1,226 @@ +package gateway + +import ( + "fmt" + "io" + "log" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Clawdlinux/agentgate/internal/auth" + "github.com/Clawdlinux/agentgate/internal/ratelimit" + upstreamclient "github.com/Clawdlinux/agentgate/internal/upstream" + "github.com/Clawdlinux/agentgate/internal/vault" +) + +func directAct(srv *Server, key, principal, service string) *httptest.ResponseRecorder { + req := httptest.NewRequest("POST", "/v1/act", strings.NewReader(fmt.Sprintf( + `{"service":%q,"action":"list_repos","on_behalf_of":%q}`, service, principal))) + req.Header.Set("Authorization", "Bearer "+key) + w := httptest.NewRecorder() + srv.ServeHTTP(w, req) + return w +} + +func installSuccessTransport(srv *Server) { + srv.upstreams = upstreamclient.New(upstreamclient.Options{Client: &http.Client{ + Transport: responseTransport(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("{}"))}, nil + }), + }}) +} + +func TestAct_OutOfScopeCannotDrainAnotherTenant(t *testing.T) { + for _, dimension := range []string{"principal", "service"} { + t.Run(dimension, func(t *testing.T) { + srv, receipts := bodyGateway(t, "http://unused", "{}", false) + now := time.Unix(1000, 0) + srv.cfg.Limiter = ratelimit.New(nil, ratelimit.WithSettings(ratelimit.Settings{ + KeyRPS: 1000, KeyBurst: 1000, TenantRPS: 1, TenantBurst: 1, + }), ratelimit.WithClock(func() time.Time { return now })) + alice := &auth.AgentKey{ID: "alice-key", AllowedUsers: []string{"alice"}, AllowedServices: []string{"github"}} + if dimension == "service" { + alice.AllowedUsers, alice.AllowedServices = []string{"bob"}, []string{"stripe"} + } + srv.cfg.Authorizer = &fakeAuthorizer{keys: map[string]*auth.AgentKey{ + "alice-key": alice, + "bob-key": {ID: "bob-key", AllowedUsers: []string{"bob"}, AllowedServices: []string{"github"}}, + }} + if err := srv.cfg.Vault.Put("bob", "github", vault.Token{AccessToken: "secret"}); err != nil { + t.Fatal(err) + } + installSuccessTransport(srv) + verifier := &countingDelegation{} + srv.cfg.Delegation = verifier + for i := 0; i < 100; i++ { + w := limitedAct(srv, "alice-key", "bob", "github") + if w.Code != 403 || receipts.last().Error != "scope_denied" || receipts.last().PolicyDecision != "deny" { + t.Fatalf("attempt=%d status=%d receipt=%+v", i, w.Code, receipts.last()) + } + if i%10 == 0 { + w = directAct(srv, "bob-key", "bob", "github") + if w.Code != 200 { + t.Fatalf("legitimate request status=%d body=%s", w.Code, w.Body) + } + now = now.Add(time.Second) + } + } + if verifier.calls.Load() != 0 { + t.Fatal("out-of-scope calls reached delegation verification") + } + if receipts.count() != 110 { + t.Fatalf("receipt count = %d", receipts.count()) + } + }) + } +} + +func TestTenantIdentityBound(t *testing.T) { + key := &auth.AgentKey{ID: "key-1"} + for i := 0; i < 200; i++ { + key.AllowedUsers = append(key.AllowedUsers, fmt.Sprint(i)) + } + identities := make(map[string]bool) + for _, principal := range key.AllowedUsers { + identities[tenantIdentity(key, principal)] = true + } + if len(identities) != 64 { + t.Fatalf("one key selected %d tenant entries", len(identities)) + } + if tenantIdentity(key, "190") != tenantIdentity(key, "199") { + t.Fatal("principals beyond the cap must share the key fallback") + } + key.AllowedUsers = []string{"*"} + identities = make(map[string]bool) + for i := 0; i < 200; i++ { + identities[tenantIdentity(key, fmt.Sprint(i))] = true + } + if len(identities) != 1 { + t.Fatalf("wildcard key selected %d tenant entries", len(identities)) + } + other := &auth.AgentKey{ID: "key-2", AllowedUsers: []string{"*"}} + if tenantIdentity(key, "alice") == tenantIdentity(other, "alice") { + t.Fatal("wildcard fallbacks must be key-scoped") + } + key.AllowedUsers, other.AllowedUsers = []string{"alice"}, []string{"alice"} + if tenantIdentity(key, "alice") != tenantIdentity(other, "alice") { + t.Fatal("explicitly scoped keys must share the principal tenant") + } + other.AllowedUsers = []string{"*"} + key.AllowedUsers = []string{"key:key-2"} + if tenantIdentity(key, "key:key-2") == tenantIdentity(other, "alice") { + t.Fatal("principal names must not collide with fallback identities") + } +} + +type advancingLimiter struct { + *ratelimit.Limiter + advance func() +} + +func (l advancingLimiter) Acquire(tenant string) (func(), time.Duration) { + l.advance() + return l.Limiter.Acquire(tenant) +} + +func TestAct_ConcurrencyDenialsRefundAllBuckets(t *testing.T) { + srv, receipts := bodyGateway(t, "http://unused", "{}", false) + installSuccessTransport(srv) + now := time.Unix(1000, 0) + l := ratelimit.New(map[string]ratelimit.Config{"github": {RequestsPerSecond: 1, Burst: 1}}, + ratelimit.WithSettings(ratelimit.Settings{ + KeyRPS: 1, KeyBurst: 1, TenantRPS: 1, TenantBurst: 1, TenantMaxInflight: 1, + }), ratelimit.WithClock(func() time.Time { return now })) + srv.cfg.Limiter = advancingLimiter{l, func() { now = now.Add(10 * time.Millisecond) }} + key := srv.cfg.Authorizer.(*fakeAuthorizer).keys["test-key"] + release, delay := l.Acquire(tenantIdentity(key, "user-1")) + if delay != 0 { + t.Fatal(delay) + } + defer release() + for i := 0; i < 10; i++ { + w := transportAct(t, srv) + if w.Code != 429 || receipts.last().Error != "concurrency_limited" { + t.Fatalf("retry=%d status=%d receipt=%+v", i, w.Code, receipts.last()) + } + } + release() + if w := transportAct(t, srv); w.Code != 200 { + t.Fatalf("concurrency retries drained tokens: status=%d body=%s", w.Code, w.Body) + } + if w := transportAct(t, srv); w.Code != 429 || receipts.last().Error != "rate_limited" { + t.Fatalf("refunds bypassed quota: status=%d receipt=%+v", w.Code, receipts.last()) + } +} + +func TestAct_PanickingDispatchReleasesSlot(t *testing.T) { + srv, _ := bodyGateway(t, "http://unused", "{}", false) + l := ratelimit.New(nil, ratelimit.WithSettings(ratelimit.Settings{TenantMaxInflight: 1})) + srv.cfg.Limiter = l + var calls atomic.Int32 + srv.upstreams = upstreamclient.New(upstreamclient.Options{Client: &http.Client{ + Transport: responseTransport(func(*http.Request) (*http.Response, error) { + calls.Add(1) + panic("dispatch failed") + }), + }}) + done := make(chan struct{}) + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(done) + srv.ServeHTTP(w, r) + })) + server.Config.ErrorLog = log.New(io.Discard, "", 0) + server.Start() + defer server.Close() + req := streamRequest(t.Context(), server.URL, "application/json") + resp, err := server.Client().Do(req) + if resp != nil { + resp.Body.Close() + } + if err == nil || calls.Load() != 1 { + t.Fatalf("dispatch did not panic: calls=%d err=%v", calls.Load(), err) + } + awaitSignal(t, done) + key := srv.cfg.Authorizer.(*fakeAuthorizer).keys["test-key"] + release, delay := l.Acquire(tenantIdentity(key, "user-1")) + if delay != 0 { + t.Fatal("panic leaked the tenant slot") + } + release() +} + +func TestAct_TenantFallbackRateLimit(t *testing.T) { + for _, wildcard := range []bool{false, true} { + t.Run(fmt.Sprintf("wildcard=%t", wildcard), func(t *testing.T) { + srv, _ := bodyGateway(t, "http://unused", "{}", false) + installSuccessTransport(srv) + key := srv.cfg.Authorizer.(*fakeAuthorizer).keys["test-key"] + key.AllowedUsers = []string{"*"} + if !wildcard { + key.AllowedUsers = nil + for i := 0; i < 100; i++ { + key.AllowedUsers = append(key.AllowedUsers, fmt.Sprint(i)) + } + } + for _, principal := range []string{"98", "99"} { + if err := srv.cfg.Vault.Put(principal, "github", vault.Token{AccessToken: "secret"}); err != nil { + t.Fatal(err) + } + } + now := time.Unix(1000, 0) + srv.cfg.Limiter = ratelimit.New(nil, ratelimit.WithSettings(ratelimit.Settings{ + KeyRPS: 100, KeyBurst: 100, TenantRPS: 1, TenantBurst: 1, + }), ratelimit.WithClock(func() time.Time { return now })) + if w := directAct(srv, "test-key", "98", "github"); w.Code != 200 { + t.Fatalf("first principal status=%d body=%s", w.Code, w.Body) + } + if w := directAct(srv, "test-key", "99", "github"); w.Code != 429 { + t.Fatalf("fallback must aggregate principals, status=%d body=%s", w.Code, w.Body) + } + }) + } +} diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 75b157b..4b50bf8 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -37,6 +37,7 @@ import ( "time" "github.com/Clawdlinux/agentgate/internal/auth" + "github.com/Clawdlinux/agentgate/internal/ratelimit" "github.com/Clawdlinux/agentgate/internal/receipt" "github.com/Clawdlinux/agentgate/internal/registry" upstreamclient "github.com/Clawdlinux/agentgate/internal/upstream" @@ -56,10 +57,11 @@ type AgentAuthorizer interface { Validate(ctx context.Context, plaintext string) (*auth.AgentKey, error) } -// RequestLimiter applies token buckets before authorization work and semaphore slots before upstream dispatch. -// Zero delay means admission succeeded. Acquire returns an idempotent release function on success. +// RequestLimiter charges keys before scope checks, then charges authorized tenants and services. +// Successful permits must be committed or canceled. Acquire returns an idempotent slot release function on success. type RequestLimiter interface { - Check(agentKeyID, tenant, service string) time.Duration + CheckKey(agentKeyID string) (*ratelimit.Permit, time.Duration) + CheckTenant(agentKeyID, tenant, service string) (*ratelimit.Permit, time.Duration) Acquire(tenant string) (release func(), retryAfter time.Duration) } @@ -202,6 +204,9 @@ type attempt struct { delegationToken []byte // raw, not yet verified; set by prepareAttempt, verified by executeAttempt delegationChain []string // set by executeAttempt once delegationToken verifies; empty for a direct grant (DELG-04) releaseSlot func() + keyPermit *ratelimit.Permit + tenantPermit *ratelimit.Permit + tenantID string } // immediateResponse is written with no receipt: the request never reached @@ -239,10 +244,14 @@ func (s *Server) handleAct(w http.ResponseWriter, r *http.Request) { if _, ok := w.(http.Flusher); !ok { att.acceptSSE = false } + defer func() { + att.keyPermit.Commit() + att.tenantPermit.Commit() + if att.releaseSlot != nil { + att.releaseSlot() + } + }() out := s.executeAttempt(r.Context(), att) - if att.releaseSlot != nil { - defer att.releaseSlot() - } if out.stream != nil { defer out.stream.close() stopCancel := context.AfterFunc(r.Context(), out.stream.close) @@ -357,17 +366,31 @@ func (s *Server) prepareAttempt(r *http.Request) (*attempt, *immediateResponse) }, nil } -// executeAttempt runs every check and side effect after authentication and -// shape validation: rate limits, delegation, scope, registry lookup, vault -// fetch, and upstream dispatch. It never writes to an http.ResponseWriter — -// every exit is captured in outcome so handleAct can commit a receipt -// before any response is written. +// executeAttempt applies policy checks and dispatches upstream after authentication and shape validation. +// It returns an outcome without writing HTTP responses. handleAct commits a receipt before writing that outcome. func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { - // AgentKey has no organization binding. Tenant identity is the requested human principal. if s.cfg.Limiter != nil { - if retry := s.cfg.Limiter.Check(att.agentKey.ID, att.req.OnBehalfOf, att.req.Service); retry > 0 { + permit, retry := s.cfg.Limiter.CheckKey(att.agentKey.ID) + if retry > 0 { return limitedOutcome("rate_limited", retry) } + att.keyPermit = permit + } + if !att.agentKey.CanAccessService(att.req.Service) || !att.agentKey.CanAccessUser(att.req.OnBehalfOf) { + return &outcome{ + status: http.StatusForbidden, + body: ErrorResponse{Error: "agent key is not scoped for this service or user", Code: "scope_denied"}, + policyDecision: "deny", + errorCode: "scope_denied", + } + } + att.tenantID = tenantIdentity(att.agentKey, att.req.OnBehalfOf) + if s.cfg.Limiter != nil { + permit, retry := s.cfg.Limiter.CheckTenant(att.agentKey.ID, att.tenantID, att.req.Service) + if retry > 0 { + return limitedOutcome("rate_limited", retry) + } + att.tenantPermit = permit } if len(att.delegationToken) > 0 { // Verified before registry, vault, or upstream access (DELG-01); a @@ -393,15 +416,6 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { att.delegationChain = chain } - if !att.agentKey.CanAccessService(att.req.Service) || !att.agentKey.CanAccessUser(att.req.OnBehalfOf) { - return &outcome{ - status: http.StatusForbidden, - body: ErrorResponse{Error: "agent key is not scoped for this service or user", Code: "scope_denied"}, - policyDecision: "deny", - errorCode: "scope_denied", - } - } - svc, action, err := s.cfg.Registry.GetAction(att.req.Service, att.req.Action) if err != nil { return &outcome{ @@ -492,13 +506,17 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { } if s.cfg.Limiter != nil { - release, retry := s.cfg.Limiter.Acquire(att.req.OnBehalfOf) + release, retry := s.cfg.Limiter.Acquire(att.tenantID) if retry > 0 { + att.keyPermit.Cancel() + att.tenantPermit.Cancel() return limitedOutcome("concurrency_limited", retry) } // handleAct releases after response writing, including streams and receipt failures. att.releaseSlot = release } + att.keyPermit.Commit() + att.tenantPermit.Commit() resp, err := dispatch(svc.Name, svc.Upstream, upstream, action.Idempotent) if err != nil { status, code, message := http.StatusBadGateway, "upstream_error", "upstream request failed" diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index 07de358..b12002b 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -14,6 +14,7 @@ import ( "time" "github.com/Clawdlinux/agentgate/internal/auth" + "github.com/Clawdlinux/agentgate/internal/ratelimit" "github.com/Clawdlinux/agentgate/internal/receipt" "github.com/Clawdlinux/agentgate/internal/registry" "github.com/Clawdlinux/agentgate/internal/vault" @@ -294,7 +295,13 @@ func TestAct_RateLimited(t *testing.T) { type denyAllLimiter struct{} -func (denyAllLimiter) Check(string, string, string) time.Duration { return time.Second } +func (denyAllLimiter) CheckKey(string) (*ratelimit.Permit, time.Duration) { + return nil, time.Second +} + +func (denyAllLimiter) CheckTenant(string, string, string) (*ratelimit.Permit, time.Duration) { + return nil, time.Second +} func (denyAllLimiter) Acquire(string) (func(), time.Duration) { return func() {}, 0 } diff --git a/internal/gateway/ratelimit_test.go b/internal/gateway/ratelimit_test.go index 565241f..2a25957 100644 --- a/internal/gateway/ratelimit_test.go +++ b/internal/gateway/ratelimit_test.go @@ -50,8 +50,9 @@ func TestAct_EarlyRateLimitsAndReceipts(t *testing.T) { } l := ratelimit.New(configs, ratelimit.WithSettings(settings), ratelimit.WithClock(func() time.Time { return now })) - l.Check("test-agent", "alice", "github") + l.Check("test-agent", "human:alice", "github") srv, _, receipts := testSetup(t) + srv.cfg.Authorizer.(*fakeAuthorizer).keys["test-key"].AllowedUsers = []string{"alice"} srv.cfg.Limiter = l verifier := &countingDelegation{} srv.cfg.Delegation = verifier @@ -60,7 +61,9 @@ func TestAct_EarlyRateLimitsAndReceipts(t *testing.T) { key := "test-key" if scope == "tenant" { key = "second-key" - srv.cfg.Authorizer.(*fakeAuthorizer).keys[key] = &auth.AgentKey{ID: "second-agent"} + srv.cfg.Authorizer.(*fakeAuthorizer).keys[key] = &auth.AgentKey{ + ID: "second-agent", AllowedUsers: []string{"alice"}, AllowedServices: []string{"github"}, + } } w := limitedAct(srv, key, "alice", "github") if w.Code != 429 || w.Header().Get("Retry-After") != "3" || @@ -113,11 +116,12 @@ func TestAct_TenantConcurrencyAndSSE(t *testing.T) { defer up.Close() defer close(finish) srv, receipts := bodyGateway(t, up.URL, "{}", stream) + srv.cfg.Authorizer.(*fakeAuthorizer).keys["test-key"].AllowedUsers = []string{"user-1"} now := time.Unix(1000, 0) srv.cfg.Limiter = ratelimit.New(nil, ratelimit.WithSettings(ratelimit.Settings{TenantMaxInflight: 1}), ratelimit.WithClock(func() time.Time { return now })) srv.cfg.Authorizer.(*fakeAuthorizer).keys["second-key"] = &auth.AgentKey{ - ID: "second-agent", AllowedServices: []string{"*"}, AllowedUsers: []string{"*"}, + ID: "second-agent", AllowedServices: []string{"*"}, AllowedUsers: []string{"user-1"}, } firstDone := make(chan struct{}) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -179,7 +183,7 @@ func TestAct_TenantConcurrencyAndSSE(t *testing.T) { if first.resp.StatusCode != 200 { t.Fatalf("first status = %d", first.resp.StatusCode) } - release, delay := srv.cfg.Limiter.Acquire("user-1") + release, delay := srv.cfg.Limiter.Acquire("human:user-1") if delay != 0 { t.Fatal("request completion did not release the tenant slot") } @@ -217,7 +221,8 @@ func TestAct_ConcurrencyReleaseOnFailures(t *testing.T) { if w.Code != 500 && w.Code != 502 { t.Fatalf("status = %d", w.Code) } - release, delay := srv.cfg.Limiter.Acquire("user-1") + key := srv.cfg.Authorizer.(*fakeAuthorizer).keys["test-key"] + release, delay := srv.cfg.Limiter.Acquire(tenantIdentity(key, "user-1")) if delay != 0 { t.Fatal("failure leaked a concurrency slot") } @@ -249,12 +254,13 @@ func TestAct_SSEDisconnectReleasesTenantSlot(t *testing.T) { t.Fatal(err) } defer resp.Body.Close() - if _, delay := srv.cfg.Limiter.Acquire("user-1"); delay == 0 { + key := srv.cfg.Authorizer.(*fakeAuthorizer).keys["test-key"] + if _, delay := srv.cfg.Limiter.Acquire(tenantIdentity(key, "user-1")); delay == 0 { t.Fatal("stream should hold the tenant slot") } cancel() awaitSignal(t, done) - release, delay := srv.cfg.Limiter.Acquire("user-1") + release, delay := srv.cfg.Limiter.Acquire(tenantIdentity(key, "user-1")) if delay != 0 { t.Fatal("disconnect leaked the tenant slot") } diff --git a/internal/ratelimit/limiter.go b/internal/ratelimit/limiter.go index 299c5ea..e91a172 100644 --- a/internal/ratelimit/limiter.go +++ b/internal/ratelimit/limiter.go @@ -122,10 +122,12 @@ type entry struct { bucket *rate.Limiter slots chan struct{} lastUsed time.Time + credits float64 + pending int } // Limiter shares a bounded LRU across key, tenant, and key-service buckets. -// Active semaphores and buckets that haven't refilled are never evicted. +// Pending permits and active semaphores pin their entries. Depleted buckets also remain cached. type Limiter struct { mu sync.Mutex configs map[string]Config @@ -166,71 +168,34 @@ func New(configs map[string]Config, options ...Option) *Limiter { // Allow preserves the standalone middleware's per-key, per-service contract. // Gateway admission uses Check and Acquire instead. func (l *Limiter) Allow(agentKeyID, service string) bool { - l.mu.Lock() - defer l.mu.Unlock() - now := l.now() service = strings.ToLower(service) cfg, ok := l.configs[service] if !ok || cfg.Rate() == rate.Inf { return true } - e := l.get(bucketKey{scope: serviceScope, id: agentKeyID, service: service}, cfg.Rate(), cfg.BurstSize(), now) - if e == nil { - return false - } - reservation := e.bucket.ReserveN(now, 1) - if reservation.DelayFrom(now) > 0 { - reservation.CancelAt(now) - return false - } - return true + permit, retry := l.check([]bucketSpec{{ + bucketKey{scope: serviceScope, id: agentKeyID, service: service}, cfg.Rate(), cfg.BurstSize(), + }}) + permit.Commit() + return retry == 0 } // Check reserves 1 token in every enabled bucket. Zero means admission succeeded. // Rejections cancel all reservations, so retries don't postpone recovery. +// Gateway admission uses CheckKey and CheckTenant around scope validation. func (l *Limiter) Check(agentKeyID, tenant, service string) time.Duration { - l.mu.Lock() - defer l.mu.Unlock() - now := l.now() - service = strings.ToLower(service) - type spec struct { - key bucketKey - rps rate.Limit - burst int - } - specs := []spec{ - {bucketKey{scope: keyScope, id: agentKeyID}, rate.Limit(l.settings.KeyRPS), l.settings.KeyBurst}, - {bucketKey{scope: tenantScope, id: tenant}, rate.Limit(l.settings.TenantRPS), l.settings.TenantBurst}, - } - if cfg, ok := l.configs[service]; ok && cfg.Rate() != rate.Inf { - specs = append(specs, spec{ - bucketKey{scope: serviceScope, id: agentKeyID, service: service}, cfg.Rate(), cfg.BurstSize(), - }) - } - var reservations []*rate.Reservation - var retry time.Duration - for _, spec := range specs { - if spec.rps <= 0 || spec.burst <= 0 { - continue - } - e := l.get(spec.key, spec.rps, spec.burst, now) - if e == nil { - retry = time.Second - break - } - reservation := e.bucket.ReserveN(now, 1) - reservations = append(reservations, reservation) - retry = reservation.DelayFrom(now) - if retry > 0 { - break // Rejected keys must not allocate state for arbitrary tenant names. - } + key, retry := l.CheckKey(agentKeyID) + if retry > 0 { + return retry } + tenantPermit, retry := l.CheckTenant(agentKeyID, tenant, service) if retry > 0 { - for _, reservation := range reservations { - reservation.CancelAt(now) - } + key.Cancel() + return retry } - return retry + key.Commit() + tenantPermit.Commit() + return 0 } // Acquire takes a tenant semaphore slot without waiting. The release function is idempotent. @@ -269,10 +234,7 @@ func (l *Limiter) get(key bucketKey, rps rate.Limit, burst int, now time.Time) * for element := l.lru.Back(); element != nil; { previous := element.Prev() e := element.Value.(*entry) - if now.Sub(e.lastUsed) < 10*time.Minute { - break - } - if e.evictable(now) { + if now.Sub(e.lastUsed) >= 10*time.Minute && e.evictable(now) { l.remove(element) } element = previous @@ -286,11 +248,19 @@ func (l *Limiter) get(key bucketKey, rps rate.Limit, burst int, now time.Time) * return e } if len(l.entries) >= l.maxEntries { - oldest := l.lru.Back() - if oldest == nil || !oldest.Value.(*entry).evictable(now) { - return nil // Fail closed rather than reset an exhausted bucket or active semaphore. + evicted := false + for scanned := 0; scanned < min(64, len(l.entries)); scanned++ { + oldest := l.lru.Back() + if oldest.Value.(*entry).evictable(now) { + l.remove(oldest) + evicted = true + break + } + l.lru.MoveToFront(oldest) + } + if !evicted { + return nil } - l.remove(oldest) } e := &entry{key: key, lastUsed: now} if rps > 0 && burst > 0 { @@ -304,7 +274,7 @@ func (l *Limiter) get(key bucketKey, rps rate.Limit, burst int, now time.Time) * } func (e *entry) evictable(now time.Time) bool { - return len(e.slots) == 0 && (e.bucket == nil || e.bucket.TokensAt(now) >= float64(e.bucket.Burst())) + return e.pending == 0 && len(e.slots) == 0 && (e.bucket == nil || e.balance(now) >= float64(e.bucket.Burst())) } func (l *Limiter) remove(element *list.Element) { diff --git a/internal/ratelimit/permit.go b/internal/ratelimit/permit.go new file mode 100644 index 0000000..4e64833 --- /dev/null +++ b/internal/ratelimit/permit.go @@ -0,0 +1,131 @@ +package ratelimit + +import ( + "math" + "strings" + "sync" + "time" + + "golang.org/x/time/rate" +) + +// Permit holds refundable charges until upstream concurrency admission finishes. +// Commit and Cancel are mutually exclusive and idempotent. Both accept a nil receiver. +type Permit struct { + limiter *Limiter + entries []*entry + once sync.Once +} + +func (p *Permit) Commit() { p.finish(false) } + +func (p *Permit) Cancel() { p.finish(true) } + +func (p *Permit) finish(refund bool) { + if p == nil { + return + } + p.once.Do(func() { + p.limiter.mu.Lock() + defer p.limiter.mu.Unlock() + now := p.limiter.now() + for _, e := range p.entries { + if refund { + e.refund(now) + } + e.pending-- + } + }) +} + +type bucketSpec struct { + key bucketKey + rps rate.Limit + burst int +} + +// CheckKey charges only the API key bucket, before scope checks. +func (l *Limiter) CheckKey(agentKeyID string) (*Permit, time.Duration) { + return l.check([]bucketSpec{{ + bucketKey{scope: keyScope, id: agentKeyID}, rate.Limit(l.settings.KeyRPS), l.settings.KeyBurst, + }}) +} + +// CheckTenant charges tenant and service buckets after the gateway validates both scopes. +func (l *Limiter) CheckTenant(agentKeyID, tenant, service string) (*Permit, time.Duration) { + specs := []bucketSpec{{ + bucketKey{scope: tenantScope, id: tenant}, rate.Limit(l.settings.TenantRPS), l.settings.TenantBurst, + }} + service = strings.ToLower(service) + if cfg, ok := l.configs[service]; ok && cfg.Rate() != rate.Inf { + specs = append(specs, bucketSpec{ + bucketKey{scope: serviceScope, id: agentKeyID, service: service}, cfg.Rate(), cfg.BurstSize(), + }) + } + return l.check(specs) +} + +func (l *Limiter) check(specs []bucketSpec) (*Permit, time.Duration) { + l.mu.Lock() + defer l.mu.Unlock() + now := l.now() + var entries []*entry + var retry time.Duration + for _, spec := range specs { + if spec.rps <= 0 || spec.burst <= 0 { + continue + } + e := l.get(spec.key, spec.rps, spec.burst, now) + if e == nil { + retry = time.Second + break + } + if retry = e.take(now); retry > 0 { + break + } + e.pending++ + entries = append(entries, e) + } + if retry > 0 { + for _, e := range entries { + e.refund(now) + e.pending-- + } + return nil, retry + } + if len(entries) == 0 { + return nil, 0 + } + return &Permit{limiter: l, entries: entries}, 0 +} + +func (e *entry) balance(now time.Time) float64 { + tokens := e.bucket.TokensAt(now) + e.credits = min(e.credits, max(0, float64(e.bucket.Burst())-tokens)) + return tokens + e.credits +} + +func (e *entry) refund(now time.Time) { + // CancelAt cannot refund past reservations. Credits preserve elapsed refill without rewinding the limiter's clock. + tokens := e.bucket.TokensAt(now) + e.credits = min(e.credits+1, max(0, float64(e.bucket.Burst())-tokens)) +} + +func (e *entry) take(now time.Time) time.Duration { + available := e.balance(now) + if e.credits >= 1 { + e.credits-- + return 0 + } + reservation := e.bucket.ReserveN(now, 1) + delay := reservation.DelayFrom(now) + if delay == 0 || available >= 1 { + return 0 + } + reservation.CancelAt(now) + if delay == rate.InfDuration { + return delay + } + creditDelay := math.Min(float64(delay), e.credits/float64(e.bucket.Limit())*float64(time.Second)) + return max(time.Nanosecond, delay-time.Duration(creditDelay)) +} diff --git a/internal/ratelimit/review_test.go b/internal/ratelimit/review_test.go new file mode 100644 index 0000000..f0c2f76 --- /dev/null +++ b/internal/ratelimit/review_test.go @@ -0,0 +1,193 @@ +package ratelimit + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestCapacitySkipsPinnedTail(t *testing.T) { + l := New(nil, WithSettings(Settings{TenantMaxInflight: 1})) + l.maxEntries = 2 + pinned, delay := l.Acquire("pinned") + if delay != 0 { + t.Fatal(delay) + } + defer pinned() + idle, delay := l.Acquire("idle") + if delay != 0 { + t.Fatal(delay) + } + idle() + release, delay := l.Acquire("new") + if delay != 0 { + t.Fatal("a pinned tail must not hide an evictable entry") + } + release() + if len(l.entries) != 2 { + t.Fatalf("entries = %d", len(l.entries)) + } +} + +func TestCapacitySkipsUndrainedTail(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1.0 / 3600, KeyBurst: 1, TenantMaxInflight: 1}), + WithClock(func() time.Time { return now })) + l.maxEntries = 2 + l.Check("slow", "alice", "github") + idle, delay := l.Acquire("idle") + if delay != 0 { + t.Fatal(delay) + } + idle() + release, delay := l.Acquire("new") + if delay != 0 { + t.Fatal("an undrained hourly bucket must not hide an evictable entry") + } + release() + if l.Check("slow", "alice", "github") == 0 { + t.Fatal("the scan must not reset the hourly bucket") + } +} + +func TestCapacityScanIsBounded(t *testing.T) { + l := New(nil, WithSettings(Settings{TenantMaxInflight: 1})) + l.maxEntries = 65 + for i := 0; i < 64; i++ { + release, delay := l.Acquire(fmt.Sprint(i)) + if delay != 0 { + t.Fatal(delay) + } + defer release() + } + idle, delay := l.Acquire("idle") + if delay != 0 { + t.Fatal(delay) + } + idle() + if _, delay := l.Acquire("new"); delay == 0 { + t.Fatal("the first scan must stop after 64 pinned entries") + } + release, delay := l.Acquire("new") + if delay != 0 { + t.Fatal("the scan must rotate pinned entries away from the tail") + } + release() +} + +func TestPermitCancellationAfterTimeAdvances(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1, KeyBurst: 1}), + WithClock(func() time.Time { return now })) + permit, delay := l.CheckKey("a") + if delay != 0 { + t.Fatal(delay) + } + now = now.Add(500 * time.Millisecond) + permit.Cancel() + permit.Cancel() + permit.Commit() + next, delay := l.CheckKey("a") + if delay != 0 { + t.Fatal("a delayed cancellation must restore the token") + } + next.Commit() + if _, delay := l.CheckKey("a"); delay != time.Second { + t.Fatalf("cancellation must not exceed burst capacity, retry = %v", delay) + } + now = now.Add(time.Second) + next, delay = l.CheckKey("a") + if delay != 0 { + t.Fatal("refill was lost after refund") + } + next.Commit() +} + +func TestPermitCancellationPreservesOtherCharges(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1, KeyBurst: 2}), + WithClock(func() time.Time { return now })) + first, _ := l.CheckKey("a") + now = now.Add(500 * time.Millisecond) + second, _ := l.CheckKey("a") + second.Commit() + now = now.Add(250 * time.Millisecond) + first.Cancel() + third, delay := l.CheckKey("a") + if delay != 0 { + t.Fatal(delay) + } + third.Commit() + if _, delay := l.CheckKey("a"); delay != 250*time.Millisecond { + t.Fatalf("refund changed another charge or refill, retry = %v", delay) + } +} + +func TestConcurrentPermitRefundsPreserveQuota(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: 1, KeyBurst: 10}), + WithClock(func() time.Time { return now })) + var committed atomic.Int32 + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Go(func() { + permit, delay := l.CheckKey("a") + if delay != 0 { + return + } + if i%2 == 0 { + permit.Cancel() + } else { + permit.Commit() + committed.Add(1) + } + }) + } + wg.Wait() + for { + permit, delay := l.CheckKey("a") + if delay != 0 { + break + } + permit.Commit() + committed.Add(1) + if committed.Load() > 10 { + t.Fatal("refunds minted extra tokens") + } + } + if committed.Load() != 10 { + t.Fatalf("committed charges = %d", committed.Load()) + } + if e := l.entries[bucketKey{scope: keyScope, id: "a"}].Value.(*entry); e.pending != 0 { + t.Fatalf("pending permits = %d", e.pending) + } +} + +func TestRepeatedLateRefundsPreserveFullBucket(t *testing.T) { + for _, rps := range []float64{0.4, 1, 17.3} { + t.Run(fmt.Sprint(rps), func(t *testing.T) { + now := time.Unix(1000, 0) + l := New(nil, WithSettings(Settings{KeyRPS: rps, KeyBurst: 1}), + WithClock(func() time.Time { return now })) + for i := 0; i < 1000; i++ { + permit, delay := l.CheckKey("a") + if delay != 0 { + t.Fatalf("refunded attempt %d drained the bucket, retry = %v", i, delay) + } + now = now.Add(13 * time.Millisecond) + permit.Cancel() + now = now.Add(17 * time.Millisecond) + } + permit, delay := l.CheckKey("a") + if delay != 0 { + t.Fatal(delay) + } + permit.Commit() + if _, delay := l.CheckKey("a"); delay == 0 { + t.Fatal("refunds exceeded burst capacity") + } + }) + } +}