Rate-limit the public API per API key
Go · Go · advanced · modification
Adds per-key rate limiting to `/v1/`. `ratelimit.Limiter` keeps one token bucket per API key in a mutex-guarded map (bursts of 20, refilled at 10 per second), brought up to date lazily from an injected clock on each request, so there is no timer per key. Over-limit requests get a 429 with `Retry-After`. Every minute a background sweep drops the buckets that have refilled to full, and it stops with the server's context. The middleware sits behind `auth.Require` and keys on the authenticated key's id, not on anything the client sends. Tests drive the limiter and the middleware with a fake clock.
Public REST API in Go 1.22 (stdlib `net/http`, `log/slog`). Existing and unchanged: `auth.Require(keys, next)` authenticates `Authorization: Bearer <secret>`, answers 401 itself, and stores the caller's `auth.Key` (fields `ID` and `OrgID`) in the request context; `auth.KeyFrom(ctx)` reads it back (`ok` is false only when no key was stored) and `auth.WithKey` stores one (used by handler tests). `Key.ID` is a non-empty, server-generated identifier, not the secret; there are at most 20,000 API keys. `config.Load`, `store.Open` and `api.Routes` are unchanged; `httpserver.Serve(ctx, addr, h)` serves until `ctx` is cancelled, then drains in-flight requests for up to 10 seconds and returns, or returns the listen error. Error responses on this API are plain text from `http.Error`. The API runs as a single instance.
Requirements
- Every request under `/v1/` is limited per API key with a token bucket: a key can make a burst of up to 20 requests, and its allowance refills at 10 requests per second. Over any stretch of time, a key that starts with a full allowance gets at most 20 + 10 × (seconds elapsed) requests through. `/healthz` is not limited, and requests that `auth.Require` rejects with 401 are never counted.
- A request over the limit is answered `429 Too Many Requests` with a `Retry-After` header in whole seconds, rounded up and at least 1, saying when the key will next have room, and it uses up none of the allowance. A request that is let through counts even if its handler then fails. Limits live in memory: the API runs as a single instance, and every key starting with a full allowance after a restart is acceptable.
- The limiter is shared by all request goroutines and must never let more requests through than the bucket allows, however many arrive at once.
- A key's state is created on its first request. A background sweep runs every minute and drops the state of keys that have refilled to their full allowance, so memory follows the active keys; the sweep must not change any later decision, and it stops when the server shuts down.
- Read time through an injected clock (`func() time.Time`); production passes `time.Now`, whose monotonic reading keeps refills immune to wall-clock changes. Ship tests with a fake clock (table-driven where they fit) for the burst, the refill, the cap, `Retry-After` and the sweep.
Files touched
- cmd/api/main.go
- internal/ratelimit/limiter.go
- internal/ratelimit/middleware.go
- internal/ratelimit/limiter_test.go
- internal/ratelimit/middleware_test.go
--- cmd/api/main.go
"os/signal"
"syscall"
+ "time"
"example.com/platform/internal/api"
"example.com/platform/internal/config"
"example.com/platform/internal/httpserver"
+ "example.com/platform/internal/ratelimit"
"example.com/platform/internal/store"
+)
+
+// Per API key: bursts of up to apiBurst requests, refilled at apiRate per second.
+const (
+ apiRate = 10
+ apiBurst = 20
)
defer db.Close()
+ limiter := ratelimit.New(apiRate, apiBurst, time.Now)
+ go limiter.Sweep(ctx, time.Minute)
+
mux := http.NewServeMux()
- mux.Handle("/v1/", auth.Require(db, api.Routes(db, logger)))
+ mux.Handle("/v1/", auth.Require(db, ratelimit.PerKey(limiter, api.Routes(db, logger))))
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusNoContent)
--- internal/ratelimit/limiter.go
+// Package ratelimit limits requests per API key with token buckets.