From 966792a2fea432c47f07fd40cb672ade683f3994 Mon Sep 17 00:00:00 2001 From: notset Date: Thu, 3 Sep 2026 15:57:44 +0800 Subject: [PATCH] Implement rate limiting with configurable options --- .env.example | 33 +++++++ README.md | 76 +++++++++++++++ package.json | 3 +- pnpm-lock.yaml | 95 ++++++++++++++++++ src/config/quota.js | 43 +++++++++ src/config/rateLimit.js | 104 ++++++++++++++++++++ src/rateLimit.js | 208 ++++++++++++++++++++++++++++++++++++++++ src/server.js | 2 + src/utils/jwt.js | 56 +++++++++++ 9 files changed, 619 insertions(+), 1 deletion(-) create mode 100644 src/config/quota.js create mode 100644 src/config/rateLimit.js create mode 100644 src/rateLimit.js create mode 100644 src/utils/jwt.js diff --git a/.env.example b/.env.example index 1371901..4971ea8 100644 --- a/.env.example +++ b/.env.example @@ -9,3 +9,36 @@ DOT_TREASURY_GRAPHQL_URL=https://eco-api.dotreasury.com/graphql # Identifiers for the SubSquare API endpoints SUBSQUARE_IDENTITY_SERVER_URL=https://id2.statescan.io + +# ---- Rate limiting --------------------------------------------------------- +# Fixed-window rate limit on the /mcp endpoint. Rules are optional: every +# variable has a default, so you can omit them entirely. Defaults match the +# SubSquare GraphQL gateway rate limit plugin. + +# Toggle rate limiting on/off (default: true) +RATE_LIMIT_ENABLED=true + +# Requests carrying a valid API token are limited per user (token `sub`) using +# the built-in quota plan named by the token's `quota` claim (standard / +# premium — see README). The knobs below only apply to tokenless (anonymous) +# requests, which are limited per IP. + +# Max requests allowed per anonymous IP within the window (default: 100) +RATE_LIMIT_MAX=100 + +# Window length in seconds (default: 1) +RATE_LIMIT_WINDOW=1 + +# Semicolon-separated IPs exempted from rate limiting (default: none) +RATE_LIMIT_WHITELIST_IPS= + +# Max number of distinct clients tracked in the in-memory store (default: 10000) +RATE_LIMIT_STORE_SIZE=10000 + +# HS256 secret used to verify API tokens sent as `Authorization: Bearer `. +# DEDICATED to MCP tokens — independent from subsquare-backend's own +# user-session JWT_SECRET_KEY. It must match the key the backend uses to sign +# MCP tokens. When set, the token's `sub` (user ID) and `quota` (plan) claims +# drive rate limiting; invalid or expired tokens are rejected with 401. Leave +# empty to disable token verification (per-IP limiting only). +MCP_API_JWT_SECRET_KEY= diff --git a/README.md b/README.md index 5e35ff6..b9e59a8 100644 --- a/README.md +++ b/README.md @@ -37,3 +37,79 @@ Available `--scope` values: ``` http://127.0.0.1:3210/mcp ``` + +# Rate limiting + +The HTTP endpoint is protected by a fixed-window rate limit, applied before +the MCP transport handles a request. Rules mirror the SubSquare GraphQL gateway +and are configurable through environment variables (see `.env.example`); the +server must be restarted after changing them. + +## Anonymous requests (per IP) + +Requests without a token are treated as anonymous and limited per client IP +using the operator-tunable knobs below (defaults match the `standard` quota +plan). + +| Env var | Default | Description | +| -------------------------- | ------- | --------------------------------------------- | +| `RATE_LIMIT_ENABLED` | `true` | Set to `0`/`false` to disable rate limiting | +| `RATE_LIMIT_MAX` | `100` | Max requests per anonymous IP per window | +| `RATE_LIMIT_WINDOW` | `1` | Window length in seconds | +| `RATE_LIMIT_WHITELIST_IPS` | (none) | Semicolon-separated IPs exempt from the limit | +| `RATE_LIMIT_STORE_SIZE` | `10000` | Max clients tracked in the in-memory store | +| `MCP_API_JWT_SECRET_KEY` | (unset) | HS256 secret verifying API tokens (see below) | + +The client IP is derived from reverse-proxy headers when present +(`cf-connecting-ip`, then `x-real-ip`, then `x-forwarded-for`), falling back to +the TCP peer address. Put the server behind a trusted proxy (e.g. Cloudflare) +if those headers can be spoofed in your network. + +## API token requests (per user) + +When `MCP_API_JWT_SECRET_KEY` is set, a client can present an HS256 JWT as +`Authorization: Bearer `. Valid tokens are rate limited **per user** +(`sub` claim) instead of per IP, using the quota plan named by the token's +`quota` claim. An invalid or expired token is rejected with `401 Unauthorized`. + +`MCP_API_JWT_SECRET_KEY` is a **dedicated key** for MCP tokens: it is separate +from subsquare-backend's own user-session `JWT_SECRET_KEY`, and must match the +key the backend signs MCP tokens with. Tokens are issued by subsquare-backend +using the same [`jsonwebtoken`](https://www.npmjs.com/package/jsonwebtoken) +package and conventions — `jwt.sign(content, secret, { expiresIn })` (HS256), +verified here with `jwt.verify(token, secret)`. + +Token payload: + +```json +{ + "sub": "user_9527", + "jti": "mcp_tk_abc123", + "exp": 1716283700, + "quota": "premium" +} +``` + +| Claim | Required | Meaning | +| ------- | -------- | ------------------------------------------------ | +| `sub` | yes | Unique user ID; the rate limit bucket key | +| `exp` | yes | Expiry timestamp (seconds); hard safety floor | +| `jti` | no | Unique token ID (reserved for future revocation) | +| `quota` | no | Rate limit plan; defaults to `standard` | + +Built-in quota plans (hardcoded for now): + +| Plan | Max requests per window | Window | +| ---------- | ----------------------- | ------ | +| `standard` | 100 | 1s | +| `premium` | 600 | 1s | + +An unknown or missing `quota` falls back to `standard`. + +## Responses + +A client that exceeds its limit receives `429 Too Many Requests` with +`Retry-After` and `X-RateLimit-*` headers, plus a JSON-RPC error body +(`code -32029`) that MCP clients can surface as a retryable error. An invalid +or expired token receives `401 Unauthorized` (`code -32002`, plus +`WWW-Authenticate: Bearer`). diff --git a/package.json b/package.json index 200be45..8d14086 100644 --- a/package.json +++ b/package.json @@ -17,9 +17,10 @@ "@koa/router": "^15.7.0", "@modelcontextprotocol/sdk": "^1.30.0", "bignumber.js": "^11.1.5", + "dotenv": "^17.2.3", + "jsonwebtoken": "^8.5.1", "koa": "^3.2.1", "lodash": "^4.18.1", - "dotenv": "^17.2.3", "lru-cache": "^11.5.2", "p-limit": "^7.3.2", "undici": "^6.28.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 767c549..e9f8d13 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -20,6 +20,9 @@ importers: dotenv: specifier: ^17.2.3 version: 17.4.2 + jsonwebtoken: + specifier: ^8.5.1 + version: 8.5.1 koa: specifier: ^3.2.1 version: 3.2.1 @@ -89,6 +92,9 @@ packages: resolution: {integrity: sha512-2cGmJupaNgg+QUwVLAucDuWuoMZ6EX9iHDRswZ5lsNYEmwPaRknMPCLZz07yTzVq/83p4o/wzbDZbBrTvGGTIw==} engines: {node: '>=18'} + buffer-equal-constant-time@1.0.1: + resolution: {integrity: sha512-zRpUiDwd/xk6ADqPMATG8vc9VPrkck7T07OIx0gnjmJAnHnTVXNQG3vfvWNuiZIkwu9KrKdA1iJKfsfTVxE6NA==} + bytes@3.1.2: resolution: {integrity: sha512-/Nf7TyzTx6S3yRJObOAV7956r8cr2+Oj8AC5dt8wSP3BQAoeX58NoHyCU8P8zGkNXStjTSi6fzO6F0pBdcYbEg==} engines: {node: '>= 0.8'} @@ -172,6 +178,9 @@ packages: resolution: {integrity: sha512-KIN/nDJBQRcXw0MLVhZE9iQHmG68qAVIBg9CqmUYjmQIhgij9U5MFvrqkUL5FbtyyzZuOeOt0zdeRe4UY7ct+A==} engines: {node: '>= 0.4'} + ecdsa-sig-formatter@1.0.11: + resolution: {integrity: sha512-nagl3RYrbNv6kQkeJIpt6NJZy8twLB/2vtz6yN9Z4vRKHN4/QZJIEbqohALSgwKdnksuY3k5Addp5lg8sVoVcQ==} + ee-first@1.1.1: resolution: {integrity: sha512-WMwm9LhRUo+WUaRN+vRuETqG89IgZphVSNkdFgeb6sS/E4OrDIN7t48CAewSHXc6C8lefD8KKfr5vY61brQlow==} @@ -307,6 +316,16 @@ packages: json-schema-typed@8.0.2: resolution: {integrity: sha512-fQhoXdcvc3V28x7C7BMs4P5+kNlgUURe2jmUT1T//oBRMDrqy1QPelJimwZGo7Hg9VPV3EQV5Bnq4hbFy2vetA==} + jsonwebtoken@8.5.1: + resolution: {integrity: sha512-XjwVfRS6jTMsqYs0EsuJ4LGxXV14zQybNd4L2r0UvbVnSF9Af8x7p5MzbJ90Ioz/9TI41/hTCvznF/loiSzn8w==} + engines: {node: '>=4', npm: '>=1.4.28'} + + jwa@1.4.2: + resolution: {integrity: sha512-eeH5JO+21J78qMvTIDdBXidBd6nG2kZjg5Ohz/1fpa28Z4CcsWUzJ1ZZyFq/3z3N17aZy+ZuBoHljASbL1WfOw==} + + jws@3.2.3: + resolution: {integrity: sha512-byiJ0FLRdLdSVSReO/U4E7RoEyOCKnEnEPMjq3HxWtvzLsV08/i5RQKsFVNkCldrCaPr2vDNAOMsfs8T/Hze7g==} + keygrip@1.1.0: resolution: {integrity: sha512-iYSchDJ+liQ8iwbSI2QqsQOvqv58eJCEanyJPJi+Khyu8smkcKSFUCbPwzFcL7YVtZ6eONjqRX/38caJ7QjRAQ==} engines: {node: '>= 0.6'} @@ -318,6 +337,27 @@ packages: resolution: {integrity: sha512-e7IpWJrnanNUroVK2taAgMxoEZvHLXdQiNjeExSu/DEIWm83jaKGBgb7tLmu2rMYpA027qFB3iLR/k3AVpFRnA==} engines: {node: '>= 18'} + lodash.includes@4.3.0: + resolution: {integrity: sha512-W3Bx6mdkRTGtlJISOvVD/lbqjTlPPUDTMnlXZFnVwi9NKJ6tiAk6LVdlhZMm17VZisqhKcgzpO5Wz91PCt5b0w==} + + lodash.isboolean@3.0.3: + resolution: {integrity: sha512-Bz5mupy2SVbPHURB98VAcw+aHh4vRV5IPNhILUCsOzRmsTmSQ17jIuqopAentWoehktxGd9e/hbIXq980/1QJg==} + + lodash.isinteger@4.0.4: + resolution: {integrity: sha512-DBwtEWN2caHQ9/imiNeEA5ys1JoRtRfY3d7V9wkqtbycnAmTvRRmbHKDV4a0EYc678/dia0jrte4tjYwVBaZUA==} + + lodash.isnumber@3.0.3: + resolution: {integrity: sha512-QYqzpfwO3/CWf3XP+Z+tkQsfaLL/EnUlXWVkIk5FUPc4sBdTehEqZONuyRt2P67PXAk+NXmTBcc97zw9t1FQrw==} + + lodash.isplainobject@4.0.6: + resolution: {integrity: sha512-oSXzaWypCMHkPC3NvBEaPHf0KsA5mvPrOPgQWDsbg8n7orZ290M0BmC/jgRZ4vcJ6DTAhjrsSYgdsW/F+MFOBA==} + + lodash.isstring@4.0.1: + resolution: {integrity: sha512-0wJxfxH1wgO3GrbuP+dTTk7op+6L41QCXbGINEmD+ny/G/eCqGzxyCsh7159S+mgDDcoarnBw6PC1PS5+wUGgw==} + + lodash.once@4.1.1: + resolution: {integrity: sha512-Sb487aTOCr9drQVL8pIxOzVhafOjZN9UU54hiN8PU3uAiSV7lx1yYNpbNmex2PK6dSJoNTSJUUswT651yww3Mg==} + lodash@4.18.1: resolution: {integrity: sha512-dMInicTPVE8d1e5otfwmmjlxkZoUpiVLwyeTdUsi/Caj/gfzzblBcCE5sRHV/AsjuCmxWrte2TNGSYuCeCq+0Q==} @@ -422,9 +462,16 @@ packages: resolution: {integrity: sha512-nLTrUKm2UyiL7rlhapu/Zl45FwNgkZGaCpZbIHajDYgwlJCOzLSk+cIPAnsEqV955GjILJnKbdQC1nVPz+gAYQ==} engines: {node: '>= 18'} + safe-buffer@5.2.1: + resolution: {integrity: sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ==} + safer-buffer@2.1.2: resolution: {integrity: sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==} + semver@5.7.2: + resolution: {integrity: sha512-cBznnQ9KjJqU67B52RMC65CMarK2600WFnbkcaiwWq3xy/5haFJlshgnpjovMVJ+Hff49d8GEn0b87C5pDQ10g==} + hasBin: true + send@1.2.1: resolution: {integrity: sha512-1gnZf7DFcoIcajTjTwjwuDjzuz4PPcY2StKPlsGAQ1+YH20IRVrBaXSWmdjowTJ6u8Rc01PoYOGHXfP1mYcZNQ==} engines: {node: '>= 18'} @@ -587,6 +634,8 @@ snapshots: transitivePeerDependencies: - supports-color + buffer-equal-constant-time@1.0.1: {} + bytes@3.1.2: {} call-bind-apply-helpers@1.0.2: @@ -649,6 +698,10 @@ snapshots: es-errors: 1.3.0 gopd: 1.2.0 + ecdsa-sig-formatter@1.0.11: + dependencies: + safe-buffer: 5.2.1 + ee-first@1.1.1: {} encodeurl@2.0.0: {} @@ -804,6 +857,30 @@ snapshots: json-schema-typed@8.0.2: {} + jsonwebtoken@8.5.1: + dependencies: + jws: 3.2.3 + lodash.includes: 4.3.0 + lodash.isboolean: 3.0.3 + lodash.isinteger: 4.0.4 + lodash.isnumber: 3.0.3 + lodash.isplainobject: 4.0.6 + lodash.isstring: 4.0.1 + lodash.once: 4.1.1 + ms: 2.1.3 + semver: 5.7.2 + + jwa@1.4.2: + dependencies: + buffer-equal-constant-time: 1.0.1 + ecdsa-sig-formatter: 1.0.11 + safe-buffer: 5.2.1 + + jws@3.2.3: + dependencies: + jwa: 1.4.2 + safe-buffer: 5.2.1 + keygrip@1.1.0: dependencies: tsscmp: 1.0.6 @@ -831,6 +908,20 @@ snapshots: type-is: 2.1.0 vary: 1.1.2 + lodash.includes@4.3.0: {} + + lodash.isboolean@3.0.3: {} + + lodash.isinteger@4.0.4: {} + + lodash.isnumber@3.0.3: {} + + lodash.isplainobject@4.0.6: {} + + lodash.isstring@4.0.1: {} + + lodash.once@4.1.1: {} + lodash@4.18.1: {} lru-cache@11.5.2: {} @@ -916,8 +1007,12 @@ snapshots: transitivePeerDependencies: - supports-color + safe-buffer@5.2.1: {} + safer-buffer@2.1.2: {} + semver@5.7.2: {} + send@1.2.1: dependencies: debug: 4.4.3 diff --git a/src/config/quota.js b/src/config/quota.js new file mode 100644 index 0000000..8184474 --- /dev/null +++ b/src/config/quota.js @@ -0,0 +1,43 @@ +// Built-in rate limit quota plans. +// +// An API token's `quota` claim (see README) selects one of these plans, which +// defines how many requests that user may send per window. Two tiers are +// hardcoded for now; make this configurable when more tiers are needed. +// +// | Plan | Max requests per window | Window | +// |------------|-------------------------|--------| +// | standard | 100 | 1s | +// | premium | 600 | 1s | +// +// Tokenless (anonymous) requests do not carry a quota claim; they use +// `anonymousPlan` below, derived from the operator-tunable rate limit knobs. + +import { rateLimitConfig } from "./rateLimit.js"; + +export const QUOTA_PLANS = Object.freeze({ + standard: Object.freeze({ maxRequests: 100, windowSeconds: 1 }), + premium: Object.freeze({ maxRequests: 600, windowSeconds: 1 }), +}); + +// Plan applied when a token carries no quota claim or an unknown one. +export const DEFAULT_QUOTA = "standard"; + +export function getQuotaPlan(quota) { + const name = Object.hasOwn(QUOTA_PLANS, quota) ? quota : DEFAULT_QUOTA; + const plan = QUOTA_PLANS[name]; + + return Object.freeze({ + name, + maxRequests: plan.maxRequests, + windowMs: plan.windowSeconds * 1000, + }); +} + +// Requests without a token are treated as anonymous and limited per IP with +// the operator-tunable environment knobs (their defaults match the `standard` +// quota plan). +export const anonymousPlan = Object.freeze({ + name: "anonymous", + maxRequests: rateLimitConfig.maxRequests, + windowMs: rateLimitConfig.windowMs, +}); diff --git a/src/config/rateLimit.js b/src/config/rateLimit.js new file mode 100644 index 0000000..6963c6b --- /dev/null +++ b/src/config/rateLimit.js @@ -0,0 +1,104 @@ +// Rate limiting rules for the MCP HTTP endpoint. +// +// Mirrors the rate limit plugin of the SubSquare GraphQL gateway +// (packages/graphql-server/src/rateLimitPlugin.js) so both public endpoints +// follow the same model: fixed-window counting over an in-memory LRU store, +// with exempted IPs. +// +// Clients that present a valid API token (Authorization: Bearer ) are +// limited per user (`sub` claim) using a built-in quota plan selected by the +// token's `quota` claim — see src/config/quota.js. Tokenless requests are +// treated as anonymous and limited per IP using the knobs below. +// +// Every rule is exposed as an environment variable with a default value, so +// operators can tune it without touching code. Restart the server after +// changing any of them. +// +// | Env var | Meaning | Default | +// |--------------------------|------------------------------------|---------| +// | RATE_LIMIT_ENABLED | Toggle rate limiting on/off | true | +// | RATE_LIMIT_MAX | Max requests per anonymous IP/window | 100 | +// | RATE_LIMIT_WINDOW | Window length in seconds (anonymous) | 1 | +// | RATE_LIMIT_WHITELIST_IPS | Semicolon-separated exempt IPs | (none) | +// | RATE_LIMIT_STORE_SIZE | Max clients kept in the LRU store | 10000 | +// | MCP_API_JWT_SECRET_KEY | HS256 secret verifying MCP API tokens | (unset) | +// +// MCP_API_JWT_SECRET_KEY is a DEDICATED key for MCP tokens: it is independent +// from subsquare-backend's own user-session JWT_SECRET_KEY. The backend signs +// MCP tokens with this same dedicated value, and the MCP server only verifies +// them (it never signs). + +const DEFAULTS = Object.freeze({ + enabled: true, + maxRequests: 100, + windowSeconds: 1, + whitelistIps: [], + storeSize: 10000, +}); + +// Accepts 1/true/yes/on (and 0/false/no/off) regardless of case. +function parseBool(value, fallback) { + const normalized = String(value ?? "") + .trim() + .toLowerCase(); + if (normalized === "") { + return fallback; + } + return !["0", "false", "no", "off"].includes(normalized); +} + +function parsePositiveInt(value, fallback) { + const parsed = parseInt(value, 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback; +} + +function parseWhitelistIps(value) { + const ips = String(value ?? "") + .split(";") + .map((ip) => ip.trim()) + .filter(Boolean); + return new Set(ips); +} + +function loadRateLimitConfig() { + const enabled = parseBool(process.env.RATE_LIMIT_ENABLED, DEFAULTS.enabled); + const maxRequests = parsePositiveInt( + process.env.RATE_LIMIT_MAX, + DEFAULTS.maxRequests, + ); + const windowSeconds = parsePositiveInt( + process.env.RATE_LIMIT_WINDOW, + DEFAULTS.windowSeconds, + ); + const whitelistIps = parseWhitelistIps(process.env.RATE_LIMIT_WHITELIST_IPS); + const storeSize = parsePositiveInt( + process.env.RATE_LIMIT_STORE_SIZE, + DEFAULTS.storeSize, + ); + // Secret for verifying HS256 API tokens (Authorization: Bearer ). + // Undefined when unset, in which case tokens are ignored and all requests + // are limited per IP. + const jwtSecretKey = + (process.env.MCP_API_JWT_SECRET_KEY || "").trim() || undefined; + + return Object.freeze({ + enabled, + maxRequests, + windowMs: windowSeconds * 1000, + whitelistIps, + storeSize, + jwtSecretKey, + }); +} + +export const rateLimitConfig = loadRateLimitConfig(); + +if (rateLimitConfig.enabled) { + console.log("Rate limit:", { + maxRequests: rateLimitConfig.maxRequests, + windowMs: rateLimitConfig.windowMs, + whitelistIps: [...rateLimitConfig.whitelistIps], + storeSize: rateLimitConfig.storeSize, + apiTokenVerification: Boolean(rateLimitConfig.jwtSecretKey), + }); +} diff --git a/src/rateLimit.js b/src/rateLimit.js new file mode 100644 index 0000000..143800e --- /dev/null +++ b/src/rateLimit.js @@ -0,0 +1,208 @@ +import { LRUCache } from "lru-cache"; + +import { rateLimitConfig } from "./config/rateLimit.js"; +import { anonymousPlan, getQuotaPlan } from "./config/quota.js"; +import { verifyJwtToken } from "./utils/jwt.js"; + +// JSON-RPC server error codes for HTTP-layer rejections. The MCP SDK emits +// errors in the -32000..-32099 "server error" range; these codes let MCP +// clients distinguish a rate limit or auth failure from other transport +// errors while still parsing the standard { jsonrpc, error, id } envelope. +const RATE_LIMIT_ERROR_CODE = -32029; +const UNAUTHORIZED_ERROR_CODE = -32002; + +// Per-client request store. Each entry counts the requests made by one client +// (a user identified by its API token, or an anonymous IP) within the current +// window. Entries carry a per-entry TTL equal to the applicable plan's window +// length, so they expire automatically and the store size stays bounded. +const rateLimitStore = new LRUCache({ + max: rateLimitConfig.storeSize, +}); + +// Localhost loopback addresses can be reported in several forms; collapse them +// so a single client is not split across keys. +function normalizeIP(ip) { + if (ip === "::1" || ip === "::ffff:127.0.0.1") { + return "127.0.0.1"; + } + return ip; +} + +// Derives the client identity from the request, preferring reverse-proxy +// headers (Cloudflare, then generic proxies) before falling back to the TCP +// peer address. Same order as the GraphQL gateway rate limit plugin. +function extractClientIP(ctx) { + const headers = ctx.headers; + const ip = + headers["cf-connecting-ip"] || + headers["x-real-ip"] || + headers["x-forwarded-for"]?.split(",")[0]?.trim() || + ctx.req?.socket?.remoteAddress || + "unknown"; + + return normalizeIP(ip); +} + +// Fixed-window counter scoped to a bucket key and its plan: the first request +// of a window creates a record, each further request increments it, and the +// window resets once the per-entry LRU TTL expires. +function checkRateLimit(bucketKey, plan) { + const now = Date.now(); + const record = rateLimitStore.get(bucketKey); + + if (!record) { + rateLimitStore.set( + bucketKey, + { count: 1, resetTime: now + plan.windowMs }, + { ttl: plan.windowMs }, + ); + return { allowed: true }; + } + + record.count++; + + if (record.count > plan.maxRequests) { + return { + allowed: false, + limit: plan.maxRequests, + retryAfter: Math.max(1, Math.ceil((record.resetTime - now) / 1000)), + }; + } + + return { allowed: true }; +} + +// Extracts a Bearer token from the Authorization header, if any. +function extractBearerToken(ctx) { + const header = ctx.headers["authorization"]; + if (!header) { + return null; + } + + const match = /^Bearer\s+(.+)$/i.exec(header); + return match ? match[1].trim() : null; +} + +// Resolves who is making this request and which rate limit applies: +// +// - A valid API token (Authorization: Bearer ) → the token's `sub` +// (unique user ID) becomes the limiting key and its `quota` claim selects +// a built-in plan. A malformed or expired token is rejected with 401. +// - No token (or no secret configured) → the client IP is the limiting key +// with the anonymous (env-tunable) plan, preserving per-IP behaviour for +// local and unauthenticated usage. +// +// Returns { bucketKey, plan, identity } or { error: "invalid" | "expired" }. +function resolveRateLimitContext(ctx, ip) { + const token = extractBearerToken(ctx); + + if (!rateLimitConfig.jwtSecretKey || !token) { + return { + bucketKey: `ip:${ip}`, + plan: anonymousPlan, + identity: { kind: "ip", id: ip }, + }; + } + + try { + const payload = verifyJwtToken(token, rateLimitConfig.jwtSecretKey); + const plan = getQuotaPlan(payload.quota); + + return { + bucketKey: `user:${payload.sub}`, + plan, + identity: { kind: "user", id: payload.sub, quota: plan.name }, + }; + } catch (error) { + return { error: error.code === "TOKEN_EXPIRED" ? "expired" : "invalid" }; + } +} + +function describeClient(identity) { + return identity.kind === "user" + ? `user ${identity.id} (quota: ${identity.quota})` + : `IP ${identity.id}`; +} + +function sendRateLimited(ctx, limitResult) { + const retryAfter = limitResult.retryAfter; + + ctx.status = 429; + ctx.body = { + jsonrpc: "2.0", + error: { + code: RATE_LIMIT_ERROR_CODE, + message: "Too Many Requests", + data: { retryAfter }, + }, + id: null, + }; + ctx.set("Retry-After", String(retryAfter)); + ctx.set("X-RateLimit-Limit", String(limitResult.limit)); + ctx.set("X-RateLimit-Remaining", "0"); + ctx.set( + "X-RateLimit-Reset", + String(Math.ceil(Date.now() / 1000) + retryAfter), + ); +} + +function sendUnauthorized(ctx, reason) { + const message = reason === "expired" ? "Token expired" : "Invalid token"; + + ctx.status = 401; + ctx.body = { + jsonrpc: "2.0", + error: { + code: UNAUTHORIZED_ERROR_CODE, + message, + data: { reason }, + }, + id: null, + }; + ctx.set("WWW-Authenticate", "Bearer"); +} + +// Koa middleware. Runs before the router, so it guards every request to the +// MCP endpoint (POST JSON-RPC messages and GET SSE streams alike). Clients +// are identified by their API token's user ID when one is presented, or by IP +// otherwise. Exceeding the applicable plan yields a 429 with standard rate +// limit headers and a JSON-RPC error body, so MCP clients surface it as a +// retryable error. An invalid or expired token is rejected with 401. +export function rateLimitMiddleware() { + if (!rateLimitConfig.enabled) { + return async function rateLimitDisabled(ctx, next) { + await next(); + }; + } + + return async function enforceRateLimit(ctx, next) { + const ip = extractClientIP(ctx); + + // Network-level exemption: whitelisted IPs are never limited. + if (rateLimitConfig.whitelistIps.has(ip)) { + await next(); + return; + } + + const context = resolveRateLimitContext(ctx, ip); + if (context.error) { + console.warn( + "Rejected request with", + context.error, + "token from IP:", + ip, + ); + sendUnauthorized(ctx, context.error); + return; + } + + const limitResult = checkRateLimit(context.bucketKey, context.plan); + if (!limitResult.allowed) { + console.warn("Too many requests from", describeClient(context.identity)); + sendRateLimited(ctx, limitResult); + return; + } + + await next(); + }; +} diff --git a/src/server.js b/src/server.js index 41b0707..e1ff805 100644 --- a/src/server.js +++ b/src/server.js @@ -1,9 +1,11 @@ import { createServer } from "node:http"; import Koa from "koa"; +import { rateLimitMiddleware } from "./rateLimit.js"; import router from "./routes.js"; const app = new Koa(); +app.use(rateLimitMiddleware()); app.use(router.routes()); app.use(router.allowedMethods()); diff --git a/src/utils/jwt.js b/src/utils/jwt.js new file mode 100644 index 0000000..dd07568 --- /dev/null +++ b/src/utils/jwt.js @@ -0,0 +1,56 @@ +import jwt from "jsonwebtoken"; + +// API token verification for the MCP endpoint. +// +// Backed by the same jsonwebtoken package and usage pattern as +// subsquare-backend (jwt.sign / jwt.verify over an HS256 secret, with `exp` +// enforced by the library), so both sides stay interchangeable. The MCP server +// only decodes — it never signs — tokens that subsquare-backend issues with a +// DEDICATED key for MCP (see MCP_API_JWT_SECRET_KEY in config/rateLimit.js), +// separate from the backend's own user-session JWT_SECRET_KEY. The secret here +// must match the key the backend signs these MCP tokens with. +// +// Verified tokens must carry: +// sub string user ID, used as the rate limit bucket key (required) +// exp numeric expiry (seconds), set by jsonwebtoken via expiresIn and +// enforced on verify; required here as the hard safety floor +// jti optional unique token ID (reserved for future revocation) +// quota optional rate limit plan name; falls back to the default plan + +function invalidTokenError(message = "Invalid token") { + const error = new Error(message); + error.code = "INVALID_TOKEN"; + return error; +} + +// Verifies an HS256 JWT against `secret` and returns its decoded payload. +// Throws an Error whose `code` is "INVALID_TOKEN" for a malformed, wrongly +// signed, or structurally-invalid token, and "TOKEN_EXPIRED" when jsonwebtoken +// reports the token's exp claim is in the past. +export function verifyJwtToken(token, secret) { + let payload; + try { + payload = jwt.verify(token, secret); + } catch (error) { + if (error && error.name === "TokenExpiredError") { + const expiredError = new Error("Token expired"); + expiredError.code = "TOKEN_EXPIRED"; + throw expiredError; + } + throw invalidTokenError(); + } + + if (!payload || typeof payload !== "object" || Array.isArray(payload)) { + throw invalidTokenError(); + } + if (typeof payload.sub !== "string" || payload.sub === "") { + throw invalidTokenError("Token is missing a sub claim"); + } + // exp is the hard safety floor: jsonwebtoken only enforces it when present, + // so require it explicitly to guarantee every token expires. + if (typeof payload.exp !== "number") { + throw invalidTokenError("Token is missing an exp claim"); + } + + return payload; +}