Warp/server/internal/ratelimit/ratelimit_test.go

package ratelimit

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// clock is a hand-wound clock, so the tests do not sleep.
type clock struct{ at time.Time }

func (c *clock) now() time.Time      { return c.at }
func (c *clock) add(d time.Duration) { c.at = c.at.Add(d) }

func newTestLimiter(policy Policy) (*Limiter, *clock) {
	c := &clock{at: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)}
	l := New(policy)
	l.now = c.now
	return l, c
}

func TestLimiterAllowsBurstThenRefills(t *testing.T) {
	l, c := newTestLimiter(Policy{Burst: 3, Interval: time.Second})

	for i := range 3 {
		if wait := l.Allow("alice"); wait != 0 {
			t.Fatalf("request %d of the burst was limited (wait %v)", i+1, wait)
		}
	}

	wait := l.Allow("alice")
	if wait <= 0 {
		t.Fatal("the fourth request should have been limited")
	}
	if wait != time.Second {
		t.Errorf("retry after = %v, want 1s (one interval)", wait)
	}

	// Half an interval is not enough for a slot to come back.
	c.add(500 * time.Millisecond)
	if wait := l.Allow("alice"); wait != 500*time.Millisecond {
		t.Errorf("retry after half an interval = %v, want 500ms", wait)
	}

	c.add(500 * time.Millisecond)
	if wait := l.Allow("alice"); wait != 0 {
		t.Errorf("a slot should have come back after one interval, got wait %v", wait)
	}
}

func TestLimiterKeysAreIndependent(t *testing.T) {
	l, _ := newTestLimiter(Policy{Burst: 1, Interval: time.Second})

	if wait := l.Allow("alice"); wait != 0 {
		t.Fatalf("alice was limited straight away: %v", wait)
	}
	if wait := l.Allow("alice"); wait == 0 {
		t.Fatal("alice should be out of budget")
	}
	if wait := l.Allow("bob"); wait != 0 {
		t.Errorf("bob was limited by alice's traffic: %v", wait)
	}
}

func TestLimiterRecoversFullyWhenIdle(t *testing.T) {
	l, c := newTestLimiter(Policy{Burst: 2, Interval: time.Second})

	l.Allow("alice")
	l.Allow("alice")
	c.add(time.Hour)

	for i := range 2 {
		if wait := l.Allow("alice"); wait != 0 {
			t.Errorf("after an idle hour, request %d was limited: %v", i+1, wait)
		}
	}
}

func TestLimiterSweepsIdleBuckets(t *testing.T) {
	l, c := newTestLimiter(Policy{Burst: 1, Interval: time.Second})

	l.Allow("alice")
	c.add(sweepInterval + time.Second)
	// This call triggers the sweep, which must drop alice's long-full bucket
	// rather than keeping one entry per address that ever connected.
	l.Allow("bob")

	l.mu.Lock()
	defer l.mu.Unlock()
	if _, present := l.fullAt["alice"]; present {
		t.Error("an idle bucket survived the sweep")
	}
	if _, present := l.fullAt["bob"]; !present {
		t.Error("the active bucket was swept away")
	}
}

func TestDisabledLimiterAllowsEverything(t *testing.T) {
	for name, l := range map[string]*Limiter{
		"explicitly disabled": Disabled(),
		"zero burst":          New(Policy{Burst: 0, Interval: time.Second}),
		"zero interval":       New(Policy{Burst: 10, Interval: 0}),
	} {
		t.Run(name, func(t *testing.T) {
			for range 100 {
				if wait := l.Allow("alice"); wait != 0 {
					t.Fatalf("a disabled limiter limited a request: %v", wait)
				}
			}
		})
	}
}

func TestMiddlewareReturnsLimitExceeded(t *testing.T) {
	l, _ := newTestLimiter(Policy{Burst: 1, Interval: 1500 * time.Millisecond})
	handler := l.Middleware(ClientIP(false))(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }))

	call := func() *httptest.ResponseRecorder {
		req := httptest.NewRequest(http.MethodGet, "/_matrix/client/v3/login", nil)
		req.RemoteAddr = "192.0.2.10:1234"
		rec := httptest.NewRecorder()
		handler.ServeHTTP(rec, req)
		return rec
	}

	if rec := call(); rec.Code != http.StatusOK {
		t.Fatalf("first request status = %d, want 200", rec.Code)
	}

	rec := call()
	if rec.Code != http.StatusTooManyRequests {
		t.Fatalf("second request status = %d, want 429", rec.Code)
	}
	if ct := rec.Header().Get("Content-Type"); ct != "application/json" {
		t.Errorf("Content-Type = %q, want application/json", ct)
	}
	body := rec.Body.String()
	for _, want := range []string{`"errcode":"M_LIMIT_EXCEEDED"`, `"retry_after_ms":1500`} {
		if !strings.Contains(body, want) {
			t.Errorf("body %s is missing %s", body, want)
		}
	}
}

func TestClientIPUsesForwardedForOnlyWhenTrusted(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "192.0.2.10:1234"
	req.Header.Set("X-Forwarded-For", "203.0.113.7, 192.0.2.1")

	if got := ClientIP(false)(req); got != "192.0.2.10" {
		t.Errorf("untrusted X-Forwarded-For was used: key = %q", got)
	}
	if got := ClientIP(true)(req); got != "203.0.113.7" {
		t.Errorf("trusted X-Forwarded-For key = %q, want the left-most entry", got)
	}

	// A request without the header falls back to the socket either way.
	req.Header.Del("X-Forwarded-For")
	if got := ClientIP(true)(req); got != "192.0.2.10" {
		t.Errorf("key without the header = %q, want the socket address", got)
	}
}