diff --git a/README.md b/README.md index d27a78c..c1e2a57 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`. 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 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 | @@ -332,6 +337,50 @@ 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. +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. + +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, 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 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 Tokens are limited to 4096 bytes and 8 attenuation blocks. @@ -344,7 +393,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..43223fa --- /dev/null +++ b/cmd/agentgw/rate_limits.go @@ -0,0 +1,52 @@ +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 + } + } + 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/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..cd4ac86 100644 --- a/docs/registry.md +++ b/docs/registry.md @@ -1,8 +1,41 @@ -# 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-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. + +```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. +Scope checks run before service and tenant buckets. Concurrency denials refund all token charges. + +## Response settings + ```yaml services: example: 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 1ffd3df..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,11 +57,12 @@ 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 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 { - Allow(agentKeyID, service string) bool + CheckKey(agentKeyID string) (*ratelimit.Permit, time.Duration) + CheckTenant(agentKeyID, tenant, service string) (*ratelimit.Permit, time.Duration) + Acquire(tenant string) (release func(), retryAfter time.Duration) } // DelegationVerifier checks a presented Biscuit delegation token against @@ -201,6 +203,10 @@ 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() + keyPermit *ratelimit.Permit + tenantPermit *ratelimit.Permit + tenantID string } // immediateResponse is written with no receipt: the request never reached @@ -224,6 +230,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) { @@ -237,6 +244,13 @@ 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 out.stream != nil { defer out.stream.close() @@ -293,6 +307,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 @@ -348,12 +366,32 @@ func (s *Server) prepareAttempt(r *http.Request) (*attempt, *immediateResponse) }, nil } -// executeAttempt runs every check and side effect after authentication and -// shape validation: delegation, scope, rate limit, 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 { + if s.cfg.Limiter != nil { + 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 // gateway with no configured verifier fails closed rather than @@ -378,24 +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", - } - } - - 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 +505,18 @@ 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.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" @@ -545,6 +577,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..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" @@ -284,6 +285,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 +295,15 @@ func TestAct_RateLimited(t *testing.T) { type denyAllLimiter struct{} -func (denyAllLimiter) Allow(string, string) bool { return false } +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 } 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..2a25957 --- /dev/null +++ b/internal/gateway/ratelimit_test.go @@ -0,0 +1,268 @@ +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", "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 + // 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", AllowedUsers: []string{"alice"}, AllowedServices: []string{"github"}, + } + } + 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) + 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{"user-1"}, + } + 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("human: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) + } + 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") + } + 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() + 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(tenantIdentity(key, "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..e91a172 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,198 @@ 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 + credits float64 + pending int +} + +// Limiter shares a bounded LRU across key, tenant, and key-service buckets. +// Pending permits and active semaphores pin their entries. Depleted buckets also remain cached. 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 } } -// Allow checks if a request is allowed for the given agent and service. +// 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 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 + service = strings.ToLower(service) + cfg, ok := l.configs[service] + if !ok || cfg.Rate() == rate.Inf { + 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 { + key, retry := l.CheckKey(agentKeyID) + if retry > 0 { + return retry + } + tenantPermit, retry := l.CheckTenant(agentKeyID, tenant, service) + if retry > 0 { + key.Cancel() + return retry + } + key.Commit() + tenantPermit.Commit() + return 0 +} + +// 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 + } + 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 && 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 { + 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) } - newLim := rate.NewLimiter(cfg.Rate(), cfg.BurstSize()) - actual, _ := l.limiters.LoadOrStore(key, newLim) - lim = actual + if !evicted { + return nil + } + } + e := &entry{key: key, lastUsed: now} + if rps > 0 && burst > 0 { + e.bucket = rate.NewLimiter(rps, burst) } - return lim.(*rate.Limiter).Allow() + 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 e.pending == 0 && len(e.slots) == 0 && (e.bucket == nil || e.balance(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/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") + } + }) + } +} 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()