mcp connection

2026-09-22 10:58:02 +05:30
parent 4e1f746b22
commit f2aa3b3ad8
53 changed files with 12515 additions and 37 deletions
--- a/go-api/internal/httpserver/proxybuckets_test.go
+++ b/go-api/internal/httpserver/proxybuckets_test.go
@@ -0,0 +1,294 @@
+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)
+	}
+}