mcp connection

go-api/internal/httpserver/mcplimit.go

@@ -0,0 +1,230 @@
+package httpserver
+
+import (
+	"context"
+	"net/http"
+	"strconv"
+	"time"
+
+	"github.com/krow/krow-backend/go-api/internal/ratelimit"
+)
+
+// Rate limiting for the mounted OAuth and MCP routes.
+//
+// A middleware rather than a change inside either package, for one reason: the
+// SUBJECT of a limit is an HTTP concept. Which IP, which bearer token, which
+// form field names the client — none of that is knowable from inside
+// internal/oauth, and handing those packages a request so they could work it
+// out would put transport details in a layer that has none.
+//
+// NOTHING SENSITIVE BECOMES A BUCKET KEY. Every subject goes through
+// ratelimit.Subject, which hashes it. A bucket naming a bearer token would
+// write that token into a table and into any log line mentioning the bucket.
+
+// limited wraps a handler with one rule, keyed by a subject derived per request.
+//
+// The subject function returns "" to mean "not limitable" — no token on the
+// request, say — and the request passes through. That is correct rather than
+// lax: a request with no identifiable subject is refused by the handler itself
+// a moment later, and inventing a shared bucket for all of them would let one
+// caller exhaust a budget that everybody else then queues behind.
+func (s *Server) limited(rule ratelimit.Rule, subject func(*http.Request) string, next http.Handler) http.Handler {
+	if s.limiter == nil {
+		return next
+	}
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		raw := subject(r)
+		if raw == "" {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		decision, err := s.limiter.Allow(r.Context(), rule, ratelimit.Subject(raw))
+		if err != nil {
+			// The limiter could not count. It has already decided whether that
+			// permits the request — fail closed by default — and this logs the
+			// fault without naming the subject, which is a hash of a
+			// credential.
+			s.log.Error("rate limiter unavailable",
+				"rule", rule.Name, "allowed", decision.Allowed, "error", err)
+		}
+
+		// Headers on every response, not only refusals, so a well-behaved
+		// client can slow down before it is refused rather than after.
+		w.Header().Set("RateLimit-Limit", strconv.Itoa(decision.Limit))
+		w.Header().Set("RateLimit-Remaining", strconv.Itoa(decision.Remaining))
+
+		if !decision.Allowed {
+			// Retry-After in seconds, rounded up and never zero — "Retry-After:
+			// 0" invites an immediate retry, which is the one thing a limited
+			// client must not do. The same reasoning as retryAfterSeconds in
+			// ratelimit.go, and the same rounding.
+			w.Header().Set("Retry-After", retryAfterSeconds(decision.RetryAfter))
+			s.log.Warn("rate limit exceeded", "rule", rule.Name, "path", r.URL.Path)
+
+			w.Header().Set("Content-Type", "application/json; charset=utf-8")
+			w.WriteHeader(http.StatusTooManyRequests)
+			_, _ = w.Write([]byte(`{"error":"rate_limited",` +
+				`"error_description":"too many requests; retry after the interval in the Retry-After header"}`))
+			return
+		}
+
+		next.ServeHTTP(w, r)
+	})
+}
+
+/* ── Subjects ───────────────────────────────────────────────────────────── */
+
+// byClientAddr keys by the caller's address, for endpoints with no credential.
+//
+// A method, not a free function, because the address is no longer a property of
+// the request alone: resolving it needs the trusted-proxy set, which is wired
+// onto the server. See clientip.go.
+func (s *Server) byClientAddr(r *http.Request) string { return s.trust.clientAddr(r) }
+
+// byAddrAndUser keys the authorization endpoint by address AND signed-in user.
+//
+// Both, because either alone is wrong: by user only, an attacker could exhaust
+// somebody else's budget by naming them; by address only, an office behind one
+// NAT shares one person's allowance.
+func (s *Server) byAddrAndUser(r *http.Request) string {
+	subject := s.trust.clientAddr(r)
+	if identity, ok := (sessionResolver{s}).CurrentUser(r); ok {
+		subject += "|" + identity.UserID
+	}
+	return subject
+}
+
+// byFormClientID keys the token endpoint by the client_id it names.
+//
+// Reading a form value means parsing the body, which the handler then parses
+// again — ParseForm caches on the request, so the second call is free.
+func (s *Server) byFormClientID(r *http.Request) string {
+	if err := r.ParseForm(); err != nil {
+		return ""
+	}
+	if id := r.PostFormValue("client_id"); id != "" {
+		return id
+	}
+	// No client_id: the handler will refuse it. Fall back to the address so a
+	// caller cannot dodge the limit by omitting the field.
+	return s.trust.clientAddr(r)
+}
+
+// byRefreshFamily keys refresh by the token being presented.
+//
+// Keyed by the TOKEN's hash, not the family id, because the family is not
+// knowable without a database read this middleware has no business doing. The
+// effect is very nearly the same: a rotation produces a new token and therefore
+// a new bucket, so the practical limit is per-token-per-window rather than
+// per-family — which bounds a loop just as well, since a loop presenting the
+// SAME token is exactly what the limit is for.
+func byRefreshToken(r *http.Request) string {
+	if err := r.ParseForm(); err != nil {
+		return ""
+	}
+	return r.PostFormValue("refresh_token")
+}
+
+// byBearerToken keys MCP by the presented access token.
+//
+// The narrowest identity available on an MCP request, and the right one: it is
+// one connection from one client for one user. Keying by user would let a
+// person's second client eat the first's budget; keying by IP would make
+// Claude's shared egress one bucket for every customer.
+func byBearerToken(r *http.Request) string {
+	const prefix = "Bearer "
+	header := r.Header.Get("Authorization")
+	if len(header) <= len(prefix) {
+		return ""
+	}
+	// Case-insensitive prefix, matching mcpserver's own parsing.
+	if !equalFoldASCII(header[:len(prefix)], prefix) {
+		return ""
+	}
+	return header[len(prefix):]
+}
+
+func equalFoldASCII(a, b string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := 0; i < len(a); i++ {
+		ca, cb := a[i], b[i]
+		if 'A' <= ca && ca <= 'Z' {
+			ca += 'a' - 'A'
+		}
+		if 'A' <= cb && cb <= 'Z' {
+			cb += 'a' - 'A'
+		}
+		if ca != cb {
+			return false
+		}
+	}
+	return true
+}
+
+// mcpLimited applies BOTH tool-call limits to the MCP endpoint.
+//
+// Two rules stacked rather than one, because they stop different things: the
+// per-minute rule bounds a spike, and the per-hour rule bounds a slow drain
+// that would sit under the per-minute rule indefinitely. Checked minute-first
+// so the cheaper refusal happens earlier.
+func (s *Server) mcpLimited(next http.Handler) http.Handler {
+	return s.limited(ratelimit.MCPToolCallPerMinute, byBearerToken,
+		s.limited(ratelimit.MCPToolCallPerHour, byBearerToken, next))
+}
+
+// tokenLimited applies the right rule for the grant type being requested.
+//
+// One endpoint, two grant types, two different abuse shapes — so one limit
+// keyed one way would be wrong for the other. A code exchange is bounded per
+// client; a refresh is bounded per presented token, so one connection looping
+// cannot spend a second connection's budget.
+func (s *Server) tokenLimited(next http.Handler) http.Handler {
+	exchange := s.limited(ratelimit.OAuthToken, s.byFormClientID, next)
+	refresh := s.limited(ratelimit.OAuthRefresh, byRefreshToken, next)
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// ParseForm caches on the request, so the handler's own call is free.
+		if err := r.ParseForm(); err != nil {
+			next.ServeHTTP(w, r) // let the handler produce the proper error
+			return
+		}
+		if r.PostFormValue("grant_type") == "refresh_token" {
+			refresh.ServeHTTP(w, r)
+			return
+		}
+		exchange.ServeHTTP(w, r)
+	})
+}
+
+/* ── The per-organisation ceiling ───────────────────────────────────────── */
+
+// orgLimiter adapts the shared limiter to mcpserver.OrgLimiter.
+//
+// It is the ONE limit that cannot live in the middleware above, because an
+// organisation is not knowable until the bearer token has been resolved to a
+// user and that user's row read. A middleware running before authentication
+// could only key by something the client supplied — which is exactly the
+// identity this surface refuses to trust.
+//
+// So mcpserver calls this from inside dispatch, after it has an
+// authctx.Identity, and passes the org id from that identity. This type has no
+// access to the request and therefore no way to be handed a different one.
+type orgLimiter struct{ s *Server }
+
+// AllowOrg counts one call against the organisation's hourly ceiling.
+//
+// The org id is hashed like every other subject. It is not a secret, but the
+// bucket format is uniform and a uuid in a table of counters is one more place
+// a tenant identifier exists for no reason.
+func (o orgLimiter) AllowOrg(ctx context.Context, orgID string) (bool, time.Duration, error) {
+	if o.s.limiter == nil || orgID == "" {
+		// No limiter, or no organisation — the latter cannot happen, because
+		// mcpserver refuses an identity without one before it reaches here.
+		return true, 0, nil
+	}
+	d, err := o.s.limiter.Allow(ctx, ratelimit.MCPPerOrgPerHour, ratelimit.Subject(orgID))
+	return d.Allowed, d.RetryAfter, err
+}