feat(api): emit rate limit headers on success and reject

Adds X-RateLimit-Limit, X-RateLimit-Remaining, and X-RateLimit-Reset to successful responses and 429 rejections, and documents the headers in OpenAPI.

Co-Authored-By: Virgil <virgil@lethean.io>
This commit is contained in:
Virgil 2026-04-01 16:01:09 +00:00
parent 28f9540fa8
commit e713fb9f56
6 changed files with 165 additions and 12 deletions

View file

@ -153,7 +153,7 @@ They execute after `gin.Recovery()` but before any route handler. The `Option` t
| `WithRequestID()` | `X-Request-ID` propagation | Preserves client-supplied IDs; generates 16-byte hex otherwise |
| `WithResponseMeta()` | Request metadata in JSON envelopes | Merges `request_id` and `duration` into standard responses |
| `WithCORS(origins...)` | CORS policy | `"*"` enables `AllowAllOrigins`; 12-hour `MaxAge` |
| `WithRateLimit(limit)` | Per-IP token-bucket rate limiting | `429 Too Many Requests`; `Retry-After` on rejection; zero or negative disables |
| `WithRateLimit(limit)` | Per-IP token-bucket rate limiting | `429 Too Many Requests`; `X-RateLimit-*` on success; `Retry-After` on rejection; zero or negative disables |
| `WithMiddleware(mw...)` | Arbitrary Gin middleware | Escape hatch for custom middleware |
| `WithStatic(prefix, root)` | Static file serving | Directory listing disabled |
| `WithWSHandler(h)` | WebSocket at `/ws` | Wraps any `http.Handler` |

View file

@ -170,7 +170,7 @@ func operationResponses(dataSchema map[string]any) map[string]any {
"schema": envelopeSchema(dataSchema),
},
},
"headers": standardResponseHeaders(),
"headers": mergeHeaders(standardResponseHeaders(), rateLimitSuccessHeaders()),
},
"400": map[string]any{
"description": "Bad request",
@ -231,7 +231,7 @@ func healthResponses() map[string]any {
"schema": envelopeSchema(map[string]any{"type": "string"}),
},
},
"headers": standardResponseHeaders(),
"headers": mergeHeaders(standardResponseHeaders(), rateLimitSuccessHeaders()),
},
"429": map[string]any{
"description": "Too many requests",
@ -316,6 +316,27 @@ func envelopeSchema(dataSchema map[string]any) map[string]any {
// rejects a request.
func rateLimitHeaders() map[string]any {
return map[string]any{
"X-RateLimit-Limit": map[string]any{
"description": "Maximum number of requests allowed in the current window",
"schema": map[string]any{
"type": "integer",
"minimum": 1,
},
},
"X-RateLimit-Remaining": map[string]any{
"description": "Number of requests remaining in the current window",
"schema": map[string]any{
"type": "integer",
"minimum": 0,
},
},
"X-RateLimit-Reset": map[string]any{
"description": "Unix timestamp when the rate limit window resets",
"schema": map[string]any{
"type": "integer",
"minimum": 1,
},
},
"Retry-After": map[string]any{
"description": "Seconds until the rate limit resets",
"schema": map[string]any{
@ -326,6 +347,34 @@ func rateLimitHeaders() map[string]any {
}
}
// rateLimitSuccessHeaders documents the response headers emitted on
// successful requests when rate limiting is enabled.
func rateLimitSuccessHeaders() map[string]any {
return map[string]any{
"X-RateLimit-Limit": map[string]any{
"description": "Maximum number of requests allowed in the current window",
"schema": map[string]any{
"type": "integer",
"minimum": 1,
},
},
"X-RateLimit-Remaining": map[string]any{
"description": "Number of requests remaining in the current window",
"schema": map[string]any{
"type": "integer",
"minimum": 0,
},
},
"X-RateLimit-Reset": map[string]any{
"description": "Unix timestamp when the rate limit window resets",
"schema": map[string]any{
"type": "integer",
"minimum": 1,
},
},
}
}
// standardResponseHeaders documents headers emitted by the response envelope
// middleware on all responses when request IDs are enabled.
func standardResponseHeaders() map[string]any {

View file

@ -76,6 +76,15 @@ func TestSpecBuilder_Good_EmptyGroups(t *testing.T) {
if _, ok := headers["X-Request-ID"]; !ok {
t.Fatal("expected X-Request-ID header on /health 429 response")
}
if _, ok := headers["X-RateLimit-Limit"]; !ok {
t.Fatal("expected X-RateLimit-Limit header on /health 429 response")
}
if _, ok := headers["X-RateLimit-Remaining"]; !ok {
t.Fatal("expected X-RateLimit-Remaining header on /health 429 response")
}
if _, ok := headers["X-RateLimit-Reset"]; !ok {
t.Fatal("expected X-RateLimit-Reset header on /health 429 response")
}
// Verify system tag exists.
tags := spec["tags"].([]any)
@ -256,6 +265,15 @@ func TestSpecBuilder_Good_SecuredResponses(t *testing.T) {
if _, ok := headers["X-Request-ID"]; !ok {
t.Fatal("expected X-Request-ID header in secured operation 429 response")
}
if _, ok := headers["X-RateLimit-Limit"]; !ok {
t.Fatal("expected X-RateLimit-Limit header in secured operation 429 response")
}
if _, ok := headers["X-RateLimit-Remaining"]; !ok {
t.Fatal("expected X-RateLimit-Remaining header in secured operation 429 response")
}
if _, ok := headers["X-RateLimit-Reset"]; !ok {
t.Fatal("expected X-RateLimit-Reset header in secured operation 429 response")
}
}
func TestSpecBuilder_Good_EnvelopeWrapping(t *testing.T) {
@ -302,6 +320,15 @@ func TestSpecBuilder_Good_EnvelopeWrapping(t *testing.T) {
if _, ok := headers["X-Request-ID"]; !ok {
t.Fatal("expected X-Request-ID header on 200 response")
}
if _, ok := headers["X-RateLimit-Limit"]; !ok {
t.Fatal("expected X-RateLimit-Limit header on 200 response")
}
if _, ok := headers["X-RateLimit-Remaining"]; !ok {
t.Fatal("expected X-RateLimit-Remaining header on 200 response")
}
if _, ok := headers["X-RateLimit-Reset"]; !ok {
t.Fatal("expected X-RateLimit-Reset header on 200 response")
}
content := resp200["content"].(map[string]any)
appJSON := content["application/json"].(map[string]any)
schema := appJSON["schema"].(map[string]any)

View file

@ -260,8 +260,10 @@ func WithCache(ttl time.Duration, maxEntries ...int) Option {
}
// WithRateLimit adds per-IP token-bucket rate limiting middleware.
// Requests exceeding the configured limit per second are rejected with
// 429 Too Many Requests and the standard Fail() error envelope.
// Requests that pass are annotated with X-RateLimit-Limit,
// X-RateLimit-Remaining, and X-RateLimit-Reset headers. Requests
// exceeding the configured limit are rejected with 429 Too Many
// Requests, Retry-After, and the standard Fail() error envelope.
// A zero or negative limit disables rate limiting.
func WithRateLimit(limit int) Option {
return func(e *Engine) {

View file

@ -3,6 +3,7 @@
package api
import (
"math"
"net/http"
"strconv"
"sync"
@ -30,6 +31,14 @@ type rateLimitBucket struct {
lastSeen time.Time
}
type rateLimitDecision struct {
allowed bool
retryAfter time.Duration
limit int
remaining int
resetAt time.Time
}
func newRateLimitStore(limit int) *rateLimitStore {
now := time.Now()
return &rateLimitStore{
@ -39,7 +48,7 @@ func newRateLimitStore(limit int) *rateLimitStore {
}
}
func (s *rateLimitStore) allow(key string) (bool, time.Duration) {
func (s *rateLimitStore) allow(key string) rateLimitDecision {
now := time.Now()
s.mu.Lock()
@ -81,7 +90,12 @@ func (s *rateLimitStore) allow(key string) (bool, time.Duration) {
if bucket.tokens >= 1 {
bucket.tokens--
return true, 0
return rateLimitDecision{
allowed: true,
limit: s.limit,
remaining: int(math.Floor(bucket.tokens)),
resetAt: now.Add(timeUntilFull(bucket.tokens, s.limit)),
}
}
deficit := 1 - bucket.tokens
@ -93,7 +107,13 @@ func (s *rateLimitStore) allow(key string) (bool, time.Duration) {
}
}
return false, wait
return rateLimitDecision{
allowed: false,
retryAfter: wait,
limit: s.limit,
remaining: 0,
resetAt: now.Add(wait),
}
}
func rateLimitMiddleware(limit int) gin.HandlerFunc {
@ -107,15 +127,16 @@ func rateLimitMiddleware(limit int) gin.HandlerFunc {
return func(c *gin.Context) {
key := clientRateLimitKey(c)
allowed, retryAfter := store.allow(key)
if !allowed {
secs := int(retryAfter / time.Second)
if retryAfter%time.Second != 0 {
decision := store.allow(key)
if !decision.allowed {
secs := int(decision.retryAfter / time.Second)
if decision.retryAfter%time.Second != 0 {
secs++
}
if secs < 1 {
secs = 1
}
setRateLimitHeaders(c, decision.limit, decision.remaining, decision.resetAt)
c.Header("Retry-After", strconv.Itoa(secs))
c.AbortWithStatusJSON(http.StatusTooManyRequests, Fail(
"rate_limit_exceeded",
@ -125,9 +146,42 @@ func rateLimitMiddleware(limit int) gin.HandlerFunc {
}
c.Next()
setRateLimitHeaders(c, decision.limit, decision.remaining, decision.resetAt)
}
}
func setRateLimitHeaders(c *gin.Context, limit, remaining int, resetAt time.Time) {
if limit > 0 {
c.Header("X-RateLimit-Limit", strconv.Itoa(limit))
}
if remaining < 0 {
remaining = 0
}
c.Header("X-RateLimit-Remaining", strconv.Itoa(remaining))
if !resetAt.IsZero() {
reset := resetAt.Unix()
if reset <= time.Now().Unix() {
reset = time.Now().Add(time.Second).Unix()
}
c.Header("X-RateLimit-Reset", strconv.FormatInt(reset, 10))
}
}
func timeUntilFull(tokens float64, limit int) time.Duration {
if limit <= 0 {
return 0
}
missing := float64(limit) - tokens
if missing <= 0 {
return 0
}
seconds := missing / float64(limit)
if seconds <= 0 {
return 0
}
return time.Duration(math.Ceil(seconds * float64(time.Second)))
}
func clientRateLimitKey(c *gin.Context) string {
if ip := c.ClientIP(); ip != "" {
return ip

View file

@ -38,6 +38,15 @@ func TestWithRateLimit_Good_AllowsBurstThenRejects(t *testing.T) {
if w1.Code != http.StatusOK {
t.Fatalf("expected first request to succeed, got %d", w1.Code)
}
if got := w1.Header().Get("X-RateLimit-Limit"); got != "2" {
t.Fatalf("expected X-RateLimit-Limit=2, got %q", got)
}
if got := w1.Header().Get("X-RateLimit-Remaining"); got != "1" {
t.Fatalf("expected X-RateLimit-Remaining=1, got %q", got)
}
if got := w1.Header().Get("X-RateLimit-Reset"); got == "" {
t.Fatal("expected X-RateLimit-Reset on successful response")
}
w2 := httptest.NewRecorder()
req2, _ := http.NewRequest(http.MethodGet, "/rate/ping", nil)
@ -58,6 +67,15 @@ func TestWithRateLimit_Good_AllowsBurstThenRejects(t *testing.T) {
if got := w3.Header().Get("Retry-After"); got == "" {
t.Fatal("expected Retry-After header on 429 response")
}
if got := w3.Header().Get("X-RateLimit-Limit"); got != "2" {
t.Fatalf("expected X-RateLimit-Limit=2 on 429, got %q", got)
}
if got := w3.Header().Get("X-RateLimit-Remaining"); got != "0" {
t.Fatalf("expected X-RateLimit-Remaining=0 on 429, got %q", got)
}
if got := w3.Header().Get("X-RateLimit-Reset"); got == "" {
t.Fatal("expected X-RateLimit-Reset on 429 response")
}
var resp api.Response[any]
if err := json.Unmarshal(w3.Body.Bytes(), &resp); err != nil {
@ -85,6 +103,9 @@ func TestWithRateLimit_Good_IsolatesPerIP(t *testing.T) {
if w1.Code != http.StatusOK {
t.Fatalf("expected first IP to succeed, got %d", w1.Code)
}
if got := w1.Header().Get("X-RateLimit-Limit"); got != "1" {
t.Fatalf("expected X-RateLimit-Limit=1, got %q", got)
}
w2 := httptest.NewRecorder()
req2, _ := http.NewRequest(http.MethodGet, "/rate/ping", nil)