295 lines
10 KiB
Go
295 lines
10 KiB
Go
package httpserver_test
|
|
|
|
// Multi-client rate-limit isolation, end to end.
|
|
//
|
|
// WHAT THIS IS FOR
|
|
//
|
|
// clientip_test.go proves the address RESOLVER picks the right string. It does
|
|
// not prove the string reaches Postgres as a distinct bucket, that the mounted
|
|
// route uses the resolver at all, or that the OAuth registration limit is
|
|
// actually per-client once it does. Those are different failures — a correct
|
|
// resolver wired to nothing looks identical from a unit test — and they are
|
|
// what broke in production, so they are tested here against the real handler
|
|
// and the real limiter.
|
|
//
|
|
// Every test drives httpserver.Handler() through the full middleware stack with
|
|
// a real database behind it. Nothing is stubbed.
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/netip"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/krow/krow-backend/go-api/internal/config"
|
|
"github.com/krow/krow-backend/go-api/internal/db"
|
|
"github.com/krow/krow-backend/go-api/internal/httpserver"
|
|
"github.com/krow/krow-backend/go-api/internal/ratelimit"
|
|
"github.com/krow/krow-backend/go-api/internal/testutil"
|
|
)
|
|
|
|
// The proxy this fixture's deployment sits behind, and an address inside it.
|
|
const (
|
|
proxyNetwork = "10.42.0.0/16"
|
|
proxyAddr = "10.42.0.1:9999"
|
|
)
|
|
|
|
// proxiedAPI is newOAuthAPI with a trusted proxy configured.
|
|
//
|
|
// Deliberately not a flag on newOAuthAPI: every existing test in this package
|
|
// must keep running with an EMPTY trusted set, because that is the default
|
|
// posture and a regression in it is the thing worth catching.
|
|
type proxiedAPI struct {
|
|
handler http.Handler
|
|
h *testutil.Harness
|
|
}
|
|
|
|
func newProxiedAPI(t *testing.T) *proxiedAPI {
|
|
t.Helper()
|
|
h := testutil.New(t)
|
|
|
|
network, err := netip.ParsePrefix(proxyNetwork)
|
|
if err != nil {
|
|
t.Fatalf("bad test network: %v", err)
|
|
}
|
|
|
|
cfg := &config.Config{
|
|
AppEnv: "development",
|
|
HTTP: config.HTTPConfig{
|
|
Host: "127.0.0.1", Port: 0, ShutdownTimeout: time.Second,
|
|
TrustedProxies: []netip.Prefix{network.Masked()},
|
|
},
|
|
DB: config.DBConfig{Schema: "public"},
|
|
OAuth: config.OAuthConfig{
|
|
Issuer: testOAuthIssuer,
|
|
Resource: testMCPResource,
|
|
LoginPath: "/login",
|
|
},
|
|
}
|
|
srv, err := httpserver.New(cfg, &db.DB{Pool: h.Pool, Schema: "public"},
|
|
slog.New(slog.NewTextHandler(io.Discard, nil)))
|
|
if err != nil {
|
|
t.Fatalf("build the server: %v", err)
|
|
}
|
|
return &proxiedAPI{handler: srv.Handler(), h: h}
|
|
}
|
|
|
|
// register performs one DCR as a caller arriving via the proxy.
|
|
//
|
|
// peer is what net/http would report as RemoteAddr; forwarded is the
|
|
// X-Forwarded-For the proxy appended. An empty forwarded value sends no header.
|
|
func (a *proxiedAPI) register(t *testing.T, peer, forwarded, name string) int {
|
|
t.Helper()
|
|
body := fmt.Sprintf(
|
|
`{"client_name":%q,"redirect_uris":["https://client.example.test/cb"]}`, name)
|
|
req, err := http.NewRequest("POST", "/oauth/register", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatalf("build request: %v", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.RemoteAddr = peer
|
|
if forwarded != "" {
|
|
req.Header.Set("X-Forwarded-For", forwarded)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
a.handler.ServeHTTP(rec, req)
|
|
return rec.Code
|
|
}
|
|
|
|
// exhaust registers until the limit refuses, and returns how many succeeded.
|
|
// It stops well past the limit so a failure reports a number rather than hanging.
|
|
func (a *proxiedAPI) exhaust(t *testing.T, peer, forwarded, name string) int {
|
|
t.Helper()
|
|
allowed := 0
|
|
for i := 0; i < ratelimit.OAuthRegister.Limit*3; i++ {
|
|
code := a.register(t, peer, forwarded, fmt.Sprintf("%s-%d", name, i))
|
|
if code == http.StatusTooManyRequests {
|
|
return allowed
|
|
}
|
|
if code != http.StatusCreated {
|
|
t.Fatalf("%s attempt %d: unexpected status %d", name, i, code)
|
|
}
|
|
allowed++
|
|
}
|
|
t.Fatalf("%s was never refused after %d registrations; the limit is not applied",
|
|
name, ratelimit.OAuthRegister.Limit*3)
|
|
return allowed
|
|
}
|
|
|
|
/* ── Ten independent clients ────────────────────────────────────────────── */
|
|
|
|
// The headline requirement: many users behind one proxy each get their own
|
|
// budget. Before this change every one of these shared a bucket and the
|
|
// eleventh registration on the list would have been refused.
|
|
func TestTenClientsBehindOneProxyDoNotShareABucket(t *testing.T) {
|
|
a := newProxiedAPI(t)
|
|
|
|
const clients = 10
|
|
for i := 0; i < clients; i++ {
|
|
client := fmt.Sprintf("203.0.113.%d", i+1)
|
|
// Each client registers TWICE — Claude mints a new client per connect,
|
|
// so a reconnect must not count against anybody else.
|
|
for attempt := 0; attempt < 2; attempt++ {
|
|
code := a.register(t, proxyAddr, client+", "+"10.42.0.1",
|
|
fmt.Sprintf("client-%d-%d", i, attempt))
|
|
if code != http.StatusCreated {
|
|
t.Fatalf("client %s attempt %d: status %d, want 201 — clients are sharing a bucket",
|
|
client, attempt, code)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Twenty registrations went through on a limit of ten per subject, which is
|
|
// only possible if the subject really is the forwarded client.
|
|
var buckets int
|
|
if err := a.h.Pool.QueryRow(t.Context(),
|
|
`SELECT count(*) FROM rate_limits WHERE bucket LIKE 'oauth.register:%'`).Scan(&buckets); err != nil {
|
|
t.Fatalf("count buckets: %v", err)
|
|
}
|
|
if buckets != clients {
|
|
t.Errorf("%d distinct oauth.register buckets, want %d — one per client", buckets, clients)
|
|
}
|
|
}
|
|
|
|
/* ── One client's exhaustion is its own ─────────────────────────────────── */
|
|
|
|
// Client A burns its whole budget; client B is unaffected. This is the property
|
|
// that failed in production, where A's retries refused B outright.
|
|
func TestOneClientExhaustingDoesNotBlockAnother(t *testing.T) {
|
|
a := newProxiedAPI(t)
|
|
|
|
const clientA, clientB = "203.0.113.50", "203.0.113.51"
|
|
|
|
allowed := a.exhaust(t, proxyAddr, clientA, "A")
|
|
if allowed != ratelimit.OAuthRegister.Limit {
|
|
t.Errorf("client A got %d registrations, want %d", allowed, ratelimit.OAuthRegister.Limit)
|
|
}
|
|
|
|
// A is now refused.
|
|
if code := a.register(t, proxyAddr, clientA, "A-again"); code != http.StatusTooManyRequests {
|
|
t.Errorf("client A after exhausting: status %d, want 429", code)
|
|
}
|
|
// B is not.
|
|
if code := a.register(t, proxyAddr, clientB, "B"); code != http.StatusCreated {
|
|
t.Errorf("client B: status %d, want 201 — A's exhaustion blocked B", code)
|
|
}
|
|
}
|
|
|
|
// A single client reconnecting repeatedly spends only its own budget, which is
|
|
// what Claude actually does: a new DCR client on every connect.
|
|
func TestRepeatedReconnectConsumesOnlyThatClientsBudget(t *testing.T) {
|
|
a := newProxiedAPI(t)
|
|
|
|
const reconnecting = "203.0.113.60"
|
|
a.exhaust(t, proxyAddr, reconnecting, "reconnector")
|
|
|
|
// Nine other clients are untouched by it.
|
|
for i := 0; i < 9; i++ {
|
|
other := fmt.Sprintf("198.51.100.%d", i+1)
|
|
if code := a.register(t, proxyAddr, other, fmt.Sprintf("other-%d", i)); code != http.StatusCreated {
|
|
t.Fatalf("client %s: status %d, want 201", other, code)
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ── Spoofing still buys nothing ────────────────────────────────────────── */
|
|
|
|
// A caller reaching the API directly — not through the proxy — cannot escape
|
|
// its bucket by varying X-Forwarded-For. It gets ONE budget however many
|
|
// different values it sends.
|
|
func TestUntrustedCallerCannotEscapeItsBucketByForging(t *testing.T) {
|
|
a := newProxiedAPI(t)
|
|
|
|
const direct = "198.51.100.200:40000" // outside proxyNetwork
|
|
|
|
allowed := 0
|
|
for i := 0; i < ratelimit.OAuthRegister.Limit*2; i++ {
|
|
// A different forged client on every single request.
|
|
code := a.register(t, direct, fmt.Sprintf("203.0.113.%d", i+100), fmt.Sprintf("forger-%d", i))
|
|
if code == http.StatusTooManyRequests {
|
|
break
|
|
}
|
|
if code != http.StatusCreated {
|
|
t.Fatalf("attempt %d: unexpected status %d", i, code)
|
|
}
|
|
allowed++
|
|
}
|
|
|
|
if allowed != ratelimit.OAuthRegister.Limit {
|
|
t.Errorf("a forging caller got %d registrations, want %d — the header bought extra budget",
|
|
allowed, ratelimit.OAuthRegister.Limit)
|
|
}
|
|
|
|
var buckets int
|
|
if err := a.h.Pool.QueryRow(t.Context(),
|
|
`SELECT count(*) FROM rate_limits WHERE bucket LIKE 'oauth.register:%'`).Scan(&buckets); err != nil {
|
|
t.Fatalf("count buckets: %v", err)
|
|
}
|
|
if buckets != 1 {
|
|
t.Errorf("a forging caller produced %d buckets, want exactly 1", buckets)
|
|
}
|
|
}
|
|
|
|
// Claiming to be the trusted proxy does not make a caller trusted.
|
|
func TestClaimingToBeTheProxyDoesNotWork(t *testing.T) {
|
|
a := newProxiedAPI(t)
|
|
|
|
const direct = "198.51.100.201:40000"
|
|
// The forged chain ends in the proxy's own address, which is what an
|
|
// attacker who has read this file would try.
|
|
if code := a.register(t, direct, "203.0.113.9, 10.42.0.1", "impostor"); code != http.StatusCreated {
|
|
t.Fatalf("setup: status %d", code)
|
|
}
|
|
|
|
var bucket string
|
|
if err := a.h.Pool.QueryRow(t.Context(),
|
|
`SELECT bucket FROM rate_limits WHERE bucket LIKE 'oauth.register:%'`).Scan(&bucket); err != nil {
|
|
t.Fatalf("read bucket: %v", err)
|
|
}
|
|
// The bucket must be the DIRECT caller's address, not the forged one.
|
|
want := "oauth.register:" + ratelimit.Subject("198.51.100.201")
|
|
if bucket != want {
|
|
t.Errorf("bucket = %q, want %q — a forged chain was believed", bucket, want)
|
|
}
|
|
}
|
|
|
|
/* ── The default posture is unchanged ───────────────────────────────────── */
|
|
|
|
// With no trusted proxies configured — the default, and how every other test in
|
|
// this package runs — the forwarded header is ignored and callers share the
|
|
// peer's bucket exactly as before.
|
|
func TestWithoutTrustedProxiesCallersShareThePeerBucket(t *testing.T) {
|
|
a := newOAuthAPI(t) // no TrustedProxies in its config
|
|
|
|
for i := 0; i < 3; i++ {
|
|
body := fmt.Sprintf(
|
|
`{"client_name":"unproxied-%d","redirect_uris":["https://client.example.test/cb"]}`, i)
|
|
req, err := http.NewRequest("POST", "/oauth/register", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatalf("build request: %v", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.RemoteAddr = "192.0.2.10:5000"
|
|
req.Header.Set("X-Forwarded-For", fmt.Sprintf("203.0.113.%d", i+1))
|
|
rec := httptest.NewRecorder()
|
|
a.handler.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusCreated {
|
|
t.Fatalf("attempt %d: status %d", i, rec.Code)
|
|
}
|
|
}
|
|
|
|
var buckets int
|
|
if err := a.h.Pool.QueryRow(t.Context(),
|
|
`SELECT count(*) FROM rate_limits WHERE bucket LIKE 'oauth.register:%'`).Scan(&buckets); err != nil {
|
|
t.Fatalf("count buckets: %v", err)
|
|
}
|
|
if buckets != 1 {
|
|
t.Errorf("%d buckets with no trusted proxy configured, want 1 — the header was read", buckets)
|
|
}
|
|
}
|