Warp/server/internal/ratelimit/ratelimit.go

// Package ratelimit is the per-endpoint rate limiting the Client-Server API
// needs: a token bucket per caller, answered with M_LIMIT_EXCEEDED and a
// retry_after_ms the client is expected to honour.
//
// Each rate-limited route gets its own Limiter, so a client hammering
// /login cannot use up the budget of /profile. Buckets are keyed per caller
// within a limiter; the key is supplied by the caller (the client address
// for unauthenticated endpoints, the user ID once a request is
// authenticated).
package ratelimit

import (
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/httputil"
)

// Policy is one bucket's shape: Burst requests may arrive back to back, and
// one slot returns every Interval.
type Policy struct {
	Burst    int
	Interval time.Duration
}

// Limiter is a set of token buckets, one per key.
//
// The implementation stores a single timestamp per key — the moment the
// bucket would next be completely full — rather than a count and a refill
// time. That is the whole state a leaky bucket needs: how far that moment is
// in the future says how many slots are in use, and the arithmetic cannot
// drift the way a periodically refilled counter can.
type Limiter struct {
	policy Policy

	// span is Burst × Interval: the longest the "full again" moment can be
	// pushed into the future, which is what makes the bucket finite.
	span time.Duration

	now func() time.Time

	mu       sync.Mutex
	fullAt   map[string]time.Time
	sweptAt  time.Time
	disabled bool
}

// sweepInterval is how often idle buckets are collected. Without it the map
// would keep one entry per address that ever connected.
const sweepInterval = time.Minute

// New builds a limiter. A policy with a non-positive burst or interval
// disables it, so that a limiter is always safe to call even when the
// configuration turned rate limiting off.
func New(policy Policy) *Limiter {
	l := &Limiter{policy: policy, now: time.Now, fullAt: map[string]time.Time{}}
	if policy.Burst <= 0 || policy.Interval <= 0 {
		l.disabled = true
		return l
	}
	l.span = time.Duration(policy.Burst) * policy.Interval
	return l
}

// Disabled returns a limiter that allows everything, for a deployment with
// rate limiting switched off.
func Disabled() *Limiter { return &Limiter{disabled: true} }

// Allow records one request against key. It returns how long the caller
// should wait; a zero duration means the request may proceed.
func (l *Limiter) Allow(key string) time.Duration {
	if l.disabled {
		return 0
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	now := l.now()
	l.sweep(now)

	// An unseen or long-idle bucket starts full, which is the same as
	// "would be completely full as of now".
	fullAt := l.fullAt[key]
	if fullAt.Before(now) {
		fullAt = now
	}

	// Taking a slot pushes the "full again" moment one interval further
	// out. It is allowed as long as that keeps it within the bucket's span;
	// beyond that, every slot is in use and the overshoot is exactly how
	// long the caller has to wait for the next one.
	next := fullAt.Add(l.policy.Interval)
	if overshoot := next.Sub(now) - l.span; overshoot > 0 {
		return overshoot
	}
	l.fullAt[key] = next
	return 0
}

// sweep drops buckets that have been full — and therefore idle — for a
// while. It runs at most once per sweepInterval and under the caller's lock,
// so it costs nothing on a busy server beyond one comparison per request.
func (l *Limiter) sweep(now time.Time) {
	if now.Sub(l.sweptAt) < sweepInterval {
		return
	}
	l.sweptAt = now
	for key, fullAt := range l.fullAt {
		if fullAt.Before(now) {
			delete(l.fullAt, key)
		}
	}
}

// KeyFunc derives the bucket key for a request.
type KeyFunc func(*http.Request) string

// Middleware rejects requests over the limit with M_LIMIT_EXCEEDED.
func (l *Limiter) Middleware(key KeyFunc) httputil.Middleware {
	return l.MiddlewareUnless(key, nil)
}

// MiddlewareUnless is Middleware with an exemption: a request exempt says
// true for is let through without being counted. Nil exempts nothing.
//
// The exemption runs before anything else, so it must be cheap — an
// in-memory lookup, never a database query.
func (l *Limiter) MiddlewareUnless(key KeyFunc, exempt func(*http.Request) bool) httputil.Middleware {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if exempt != nil && exempt(r) {
				next.ServeHTTP(w, r)
				return
			}
			if retryAfter := l.Allow(key(r)); retryAfter > 0 {
				// Rounded up: a client that waits exactly retry_after_ms
				// and is told to wait again learns nothing useful.
				ms := retryAfter.Milliseconds()
				if retryAfter%time.Millisecond != 0 {
					ms++
				}
				httputil.WriteError(w, r, httputil.LimitExceeded(ms))
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// ClientIP returns the address a request came from, as a rate limiting key.
//
// trustForwardedFor takes the left-most X-Forwarded-For entry instead of the
// socket address. It must only be turned on behind a proxy that overwrites
// the header: anyone can send it, and a spoofable key is no key at all.
func ClientIP(trustForwardedFor bool) KeyFunc {
	return func(r *http.Request) string {
		if trustForwardedFor {
			if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" {
				first, _, _ := strings.Cut(forwarded, ",")
				if first = strings.TrimSpace(first); first != "" {
					return first
				}
			}
		}
		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			// RemoteAddr is not host:port for a unix socket or a synthetic
			// request; use it whole rather than losing the distinction.
			return r.RemoteAddr
		}
		return host
	}
}