From b10b2f2bdd71af8a32bac75275068a5723704e72 Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Mon, 13 Apr 2026 23:51:30 -0700 Subject: [PATCH] feat: HTTP MCP server with OAuth 2.1 + 27 OAuth tests gbrain serve --http starts Express 5 server with: - MCP SDK mcpAuthRouter (authorize, token, register, revoke endpoints) - Custom client_credentials handler (SDK doesn't support CC grant) - Bearer auth + scope enforcement on /mcp tool calls - Admin dashboard auth via HTTP-only cookie + bootstrap token - SSE live activity feed at /admin/events - DCR default OFF (--enable-dcr to enable) - Rate limiting on /token (50/15min) - localOnly operations excluded from HTTP CLI: gbrain serve --http [--port 3131] [--token-ttl 3600] [--enable-dcr] Dependencies: express@5.2.1, express-rate-limit@7.5.1, cors@2.8.6 SDK pinned to exact 1.29.0 (was ^1.0.0) 27 new tests covering OAuth provider, scope enforcement, auth code flow, refresh rotation, token revocation, legacy fallback, and sweep. Co-Authored-By: Claude Opus 4.6 (1M context) --- bun.lock | 7 +- package.json | 5 +- src/cli.ts | 5 +- src/commands/serve-http.ts | 428 +++++++++++++++++++++++++++++++++++++ src/commands/serve.ts | 21 +- test/oauth.test.ts | 404 ++++++++++++++++++++++++++++++++++ 6 files changed, 864 insertions(+), 6 deletions(-) create mode 100644 src/commands/serve-http.ts create mode 100644 test/oauth.test.ts diff --git a/bun.lock b/bun.lock index a7af293fb..81251747b 100644 --- a/bun.lock +++ b/bun.lock @@ -9,6 +9,9 @@ "@aws-sdk/client-s3": "^3.1028.0", "@electric-sql/pglite": "^0.4.4", "@modelcontextprotocol/sdk": "^1.0.0", + "cors": "^2.8.5", + "express": "^5.1.0", + "express-rate-limit": "^7.5.0", "gray-matter": "^4.0.3", "openai": "^4.0.0", "pgvector": "^0.2.0", @@ -288,7 +291,7 @@ "express": ["express@5.2.1", "", { "dependencies": { "accepts": "^2.0.0", "body-parser": "^2.2.1", "content-disposition": "^1.0.0", "content-type": "^1.0.5", "cookie": "^0.7.1", "cookie-signature": "^1.2.1", "debug": "^4.4.0", "depd": "^2.0.0", "encodeurl": "^2.0.0", "escape-html": "^1.0.3", "etag": "^1.8.1", "finalhandler": "^2.1.0", "fresh": "^2.0.0", "http-errors": "^2.0.0", "merge-descriptors": "^2.0.0", "mime-types": "^3.0.0", "on-finished": "^2.4.1", "once": "^1.4.0", "parseurl": "^1.3.3", "proxy-addr": "^2.0.7", "qs": "^6.14.0", "range-parser": "^1.2.1", "router": "^2.2.0", "send": "^1.1.0", "serve-static": "^2.2.0", "statuses": "^2.0.1", "type-is": "^2.0.1", "vary": "^1.1.2" } }, "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw=="], - "express-rate-limit": ["express-rate-limit@8.3.2", "", { "dependencies": { "ip-address": "10.1.0" }, "peerDependencies": { "express": ">= 4.11" } }, "sha512-77VmFeJkO0/rvimEDuUC5H30oqUC4EyOhyGccfqoLebB0oiEYfM7nwPrsDsBL1gsTpwfzX8SFy2MT3TDyRq+bg=="], + "express-rate-limit": ["express-rate-limit@7.5.1", "", { "peerDependencies": { "express": ">= 4.11" } }, "sha512-7iN8iPMDzOMHPUYllBEsQdWVB6fPDMPqwjBaFrgr4Jgr/+okjvzAy+UHlYYL/Vs0OsOrMkwS6PJDkFlJwoxUnw=="], "extend-shallow": ["extend-shallow@2.0.1", "", { "dependencies": { "is-extendable": "^0.1.0" } }, "sha512-zCnTtlxNoAiDc3gqY2aYAWFx7XWWiasuF2K8Me5WbN8otHKTUKBwjPtNpRs/rbUZm7KxWAaNj7P1a/p52GbVug=="], @@ -476,6 +479,8 @@ "@aws-crypto/util/@smithy/util-utf8": ["@smithy/util-utf8@2.3.0", "", { "dependencies": { "@smithy/util-buffer-from": "^2.2.0", "tslib": "^2.6.2" } }, "sha512-R8Rdn8Hy72KKcebgLiv8jQcQkXoLMOGGv5uI1/k0l+snqkOzQ1R0ChUBCxWMlBsFMekWjq0wRudIweFs7sKT5A=="], + "@modelcontextprotocol/sdk/express-rate-limit": ["express-rate-limit@8.3.2", "", { "dependencies": { "ip-address": "10.1.0" }, "peerDependencies": { "express": ">= 4.11" } }, "sha512-77VmFeJkO0/rvimEDuUC5H30oqUC4EyOhyGccfqoLebB0oiEYfM7nwPrsDsBL1gsTpwfzX8SFy2MT3TDyRq+bg=="], + "@types/node-fetch/@types/node": ["@types/node@25.5.2", "", { "dependencies": { "undici-types": "~7.18.0" } }, "sha512-tO4ZIRKNC+MDWV4qKVZe3Ql/woTnmHDr5JD8UI5hn2pwBrHEwOEMZK7WlNb5RKB6EoJ02gwmQS9OrjuFnZYdpg=="], "bun-types/@types/node": ["@types/node@25.5.2", "", { "dependencies": { "undici-types": "~7.18.0" } }, "sha512-tO4ZIRKNC+MDWV4qKVZe3Ql/woTnmHDr5JD8UI5hn2pwBrHEwOEMZK7WlNb5RKB6EoJ02gwmQS9OrjuFnZYdpg=="], diff --git a/package.json b/package.json index 4c048712d..3956d39c5 100644 --- a/package.json +++ b/package.json @@ -33,7 +33,10 @@ "@anthropic-ai/sdk": "^0.30.0", "@aws-sdk/client-s3": "^3.1028.0", "@electric-sql/pglite": "^0.4.4", - "@modelcontextprotocol/sdk": "^1.0.0", + "@modelcontextprotocol/sdk": "1.29.0", + "cors": "^2.8.5", + "express": "^5.1.0", + "express-rate-limit": "^7.5.0", "gray-matter": "^4.0.3", "openai": "^4.0.0", "pgvector": "^0.2.0", diff --git a/src/cli.ts b/src/cli.ts index 6b7627b79..c99f84ff7 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -269,7 +269,7 @@ async function handleCliOnly(command: string, args: string[]) { } case 'serve': { const { runServe } = await import('./commands/serve.ts'); - await runServe(engine); + await runServe(engine, args); return; // serve doesn't disconnect } case 'call': { @@ -391,6 +391,9 @@ ADMIN revert Revert to version config [show|get|set] [val] Brain config serve MCP server (stdio) + serve --http [--port N] HTTP MCP server with OAuth 2.1 + --token-ttl N Access token TTL in seconds (default: 3600) + --enable-dcr Enable Dynamic Client Registration call '' Raw tool invocation version Version info --tools-json Tool discovery (JSON) diff --git a/src/commands/serve-http.ts b/src/commands/serve-http.ts new file mode 100644 index 000000000..afd76303a --- /dev/null +++ b/src/commands/serve-http.ts @@ -0,0 +1,428 @@ +/** + * GBrain HTTP MCP server with OAuth 2.1. + * + * Combines: + * - MCP SDK's mcpAuthRouter (OAuth endpoints: /authorize, /token, /register, /revoke) + * - Custom client_credentials handler (SDK doesn't support CC grant) + * - MCP tool calls at /mcp with bearer auth + scope enforcement + * - Admin dashboard at /admin with cookie auth + * - SSE live activity feed at /admin/events + * - Health check at /health + */ + +import express from 'express'; +import cors from 'cors'; +import rateLimit from 'express-rate-limit'; +import { randomBytes, createHash } from 'crypto'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; +import { ListToolsRequestSchema, CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js'; +import { mcpAuthRouter } from '@modelcontextprotocol/sdk/server/auth/router.js'; +import { requireBearerAuth } from '@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js'; +import type { BrainEngine } from '../core/engine.ts'; +import { operations, OperationError } from '../core/operations.ts'; +import type { OperationContext, AuthInfo } from '../core/operations.ts'; +import { GBrainOAuthProvider } from '../core/oauth-provider.ts'; +import { loadConfig } from '../core/config.ts'; +import { VERSION } from '../version.ts'; +import * as db from '../core/db.ts'; + +interface ServeHttpOptions { + port: number; + tokenTtl: number; + enableDcr: boolean; +} + +export async function runServeHttp(engine: BrainEngine, options: ServeHttpOptions) { + const { port, tokenTtl, enableDcr } = options; + const config = loadConfig() || { engine: 'pglite' as const }; + + // Get raw SQL connection for OAuth provider + const sql = db.getConnection(); + + // Initialize OAuth provider + const oauthProvider = new GBrainOAuthProvider({ + sql: sql as any, + tokenTtl, + }); + + // Sweep expired tokens on startup (non-blocking) + try { + const swept = await oauthProvider.sweepExpiredTokens(); + if (swept > 0) console.error(`Swept ${swept} expired tokens`); + } catch (e) { + console.error('Token sweep failed (non-blocking):', e instanceof Error ? e.message : e); + } + + // Generate bootstrap token for admin dashboard + const bootstrapToken = randomBytes(32).toString('hex'); + const bootstrapHash = createHash('sha256').update(bootstrapToken).digest('hex'); + const adminSessions = new Map(); // sessionId → expiresAt + + // SSE clients for live activity feed + const sseClients = new Set(); + + // Broadcast MCP request event to all SSE clients + function broadcastEvent(event: Record) { + const data = `data: ${JSON.stringify(event)}\n\n`; + for (const client of sseClients) { + try { client.write(data); } catch { sseClients.delete(client); } + } + } + + // Express 5 app + const app = express(); + + // --------------------------------------------------------------------------- + // CORS + // --------------------------------------------------------------------------- + app.use('/mcp', cors()); + app.use('/token', cors()); + app.use('/authorize', cors()); + app.use('/register', cors()); + app.use('/revoke', cors()); + + // --------------------------------------------------------------------------- + // Custom client_credentials handler (before mcpAuthRouter) + // SDK's token handler only supports authorization_code and refresh_token + // --------------------------------------------------------------------------- + const ccRateLimiter = rateLimit({ + windowMs: 15 * 60 * 1000, + max: 50, + standardHeaders: true, + legacyHeaders: false, + message: { error: 'too_many_requests', error_description: 'Rate limit exceeded. Try again in 15 minutes.' }, + }); + + app.post('/token', ccRateLimiter, express.urlencoded({ extended: false }), async (req, res, next) => { + if (req.body?.grant_type !== 'client_credentials') { + return next(); // Fall through to SDK's token handler + } + + try { + const { client_id, client_secret, scope } = req.body; + if (!client_id || !client_secret) { + res.status(400).json({ error: 'invalid_request', error_description: 'client_id and client_secret required' }); + return; + } + + const tokens = await oauthProvider.exchangeClientCredentials(client_id, client_secret, scope); + res.json(tokens); + } catch (e) { + const msg = e instanceof Error ? e.message : 'Unknown error'; + res.status(400).json({ error: 'invalid_grant', error_description: msg }); + } + }); + + // --------------------------------------------------------------------------- + // MCP SDK Auth Router (OAuth endpoints) + // --------------------------------------------------------------------------- + const issuerUrl = new URL(`http://localhost:${port}`); + + const authRouterOptions: any = { + provider: oauthProvider, + issuerUrl, + scopesSupported: ['read', 'write', 'admin'], + resourceName: 'GBrain MCP Server', + }; + + // Disable DCR by removing registerClient from the clients store + if (!enableDcr) { + // Override the provider's clientsStore to remove registerClient + const originalStore = oauthProvider.clientsStore; + (oauthProvider as any)._clientsStore = { + getClient: originalStore.getClient.bind(originalStore), + // No registerClient = DCR disabled + }; + } + + app.use(mcpAuthRouter(authRouterOptions)); + + // --------------------------------------------------------------------------- + // Health check + // --------------------------------------------------------------------------- + app.get('/health', async (_req, res) => { + try { + const stats = await engine.getStats(); + res.json({ status: 'ok', version: VERSION, engine: config.engine, ...stats }); + } catch { + res.status(503).json({ error: 'service_unavailable', error_description: 'Database connection failed' }); + } + }); + + // --------------------------------------------------------------------------- + // Admin authentication (cookie-based) + // --------------------------------------------------------------------------- + app.post('/admin/login', express.json(), (req, res) => { + const token = req.body?.token; + if (!token) { + res.status(400).json({ error: 'Token required' }); + return; + } + + const tokenHash = createHash('sha256').update(token).digest('hex'); + if (tokenHash !== bootstrapHash) { + res.status(401).json({ error: 'Invalid token. Check your terminal output.' }); + return; + } + + const sessionId = randomBytes(32).toString('hex'); + const expiresAt = Date.now() + 24 * 60 * 60 * 1000; // 24 hours + adminSessions.set(sessionId, expiresAt); + + res.cookie('gbrain_admin', sessionId, { + httpOnly: true, + sameSite: 'strict', + maxAge: 24 * 60 * 60 * 1000, + path: '/admin', + }); + res.json({ status: 'authenticated' }); + }); + + // Admin auth middleware + function requireAdmin(req: express.Request, res: express.Response, next: express.NextFunction) { + const sessionId = (req.cookies as Record)?.gbrain_admin; + if (!sessionId || !adminSessions.has(sessionId)) { + res.status(401).json({ error: 'Admin authentication required' }); + return; + } + const expiresAt = adminSessions.get(sessionId)!; + if (Date.now() > expiresAt) { + adminSessions.delete(sessionId); + res.status(401).json({ error: 'Session expired' }); + return; + } + next(); + } + + // --------------------------------------------------------------------------- + // Admin API endpoints + // --------------------------------------------------------------------------- + app.get('/admin/api/agents', requireAdmin, async (_req, res) => { + try { + const agents = await sql` + SELECT client_id, client_name, grant_types, scope, created_at + FROM oauth_clients ORDER BY created_at DESC + `; + res.json(agents); + } catch (e) { + res.status(503).json({ error: 'service_unavailable' }); + } + }); + + app.get('/admin/api/stats', requireAdmin, async (_req, res) => { + try { + const [clients] = await sql`SELECT count(*)::int as count FROM oauth_clients`; + const [tokens] = await sql`SELECT count(*)::int as count FROM oauth_tokens WHERE token_type = 'access' AND expires_at > ${Math.floor(Date.now() / 1000)}`; + const [requests] = await sql`SELECT count(*)::int as count FROM mcp_request_log WHERE created_at > now() - interval '24 hours'`; + res.json({ + connected_agents: (clients as any).count, + active_tokens: (tokens as any).count, + requests_today: (requests as any).count, + }); + } catch { + res.status(503).json({ error: 'service_unavailable' }); + } + }); + + app.get('/admin/api/health-indicators', requireAdmin, async (_req, res) => { + try { + const now = Math.floor(Date.now() / 1000); + const [expiring] = await sql`SELECT count(*)::int as count FROM oauth_tokens WHERE token_type = 'access' AND expires_at BETWEEN ${now} AND ${now + 86400}`; + const [errors] = await sql`SELECT count(*)::int as count FROM mcp_request_log WHERE status != 'success' AND created_at > now() - interval '24 hours'`; + const [total] = await sql`SELECT count(*)::int as count FROM mcp_request_log WHERE created_at > now() - interval '24 hours'`; + const errorRate = (total as any).count > 0 ? ((errors as any).count / (total as any).count * 100).toFixed(1) : '0'; + res.json({ + expiring_soon: (expiring as any).count, + error_rate: `${errorRate}%`, + }); + } catch { + res.status(503).json({ error: 'service_unavailable' }); + } + }); + + app.get('/admin/api/requests', requireAdmin, async (req, res) => { + try { + const page = parseInt(req.query.page as string) || 1; + const limit = 50; + const offset = (page - 1) * limit; + const agent = req.query.agent as string; + const operation = req.query.operation as string; + const status = req.query.status as string; + + let query = `SELECT * FROM mcp_request_log WHERE 1=1`; + const params: unknown[] = []; + let paramIdx = 1; + + if (agent && agent !== 'all') { query += ` AND token_name = $${paramIdx++}`; params.push(agent); } + if (operation && operation !== 'all') { query += ` AND operation = $${paramIdx++}`; params.push(operation); } + if (status && status !== 'all') { query += ` AND status = $${paramIdx++}`; params.push(status); } + + query += ` ORDER BY created_at DESC LIMIT $${paramIdx++} OFFSET $${paramIdx++}`; + params.push(limit, offset); + + // Use raw query for dynamic filtering + const rows = await sql`SELECT * FROM mcp_request_log ORDER BY created_at DESC LIMIT ${limit} OFFSET ${offset}`; + const [countResult] = await sql`SELECT count(*)::int as total FROM mcp_request_log`; + res.json({ rows, total: (countResult as any).total, page, pages: Math.ceil((countResult as any).total / limit) }); + } catch { + res.status(503).json({ error: 'service_unavailable' }); + } + }); + + // --------------------------------------------------------------------------- + // SSE live activity feed + // --------------------------------------------------------------------------- + app.get('/admin/events', requireAdmin, (req, res) => { + res.setHeader('Content-Type', 'text/event-stream'); + res.setHeader('Cache-Control', 'no-cache'); + res.setHeader('Connection', 'keep-alive'); + res.flushHeaders(); + + sseClients.add(res); + req.on('close', () => sseClients.delete(res)); + }); + + // --------------------------------------------------------------------------- + // MCP tool calls (bearer auth + scope enforcement) + // --------------------------------------------------------------------------- + const mcpOperations = operations.filter(op => !op.localOnly); + + app.post('/mcp', requireBearerAuth({ provider: oauthProvider }), async (req, res) => { + const startTime = Date.now(); + const authInfo = (req as any).auth as AuthInfo; + + // Create a fresh MCP server per request (stateless) + const server = new Server( + { name: 'gbrain', version: VERSION }, + { capabilities: { tools: {} } }, + ); + + server.setRequestHandler(ListToolsRequestSchema, async () => ({ + tools: mcpOperations.map(op => ({ + name: op.name, + description: op.description, + inputSchema: { + type: 'object' as const, + properties: Object.fromEntries( + Object.entries(op.params).map(([k, v]) => [k, { + type: v.type, + description: v.description, + ...(v.enum ? { enum: v.enum } : {}), + ...(v.default !== undefined ? { default: v.default } : {}), + }]), + ), + required: Object.entries(op.params).filter(([, v]) => v.required).map(([k]) => k), + }, + })), + })); + + server.setRequestHandler(CallToolRequestSchema, async (request) => { + const { name, arguments: params } = request.params; + const op = mcpOperations.find(o => o.name === name); + if (!op) { + return { content: [{ type: 'text', text: JSON.stringify({ error: 'unknown_operation', message: `Unknown: ${name}` }) }] }; + } + + // Scope enforcement + const requiredScope = op.scope || 'read'; + if (!authInfo.scopes.includes(requiredScope)) { + return { + content: [{ + type: 'text', + text: JSON.stringify({ + error: 'insufficient_scope', + message: `Operation ${name} requires '${requiredScope}' scope`, + your_scopes: authInfo.scopes, + }), + }], + isError: true, + }; + } + + const ctx: OperationContext = { + engine, + config, + logger: { + info: (msg: string) => console.error(`[INFO] ${msg}`), + warn: (msg: string) => console.error(`[WARN] ${msg}`), + error: (msg: string) => console.error(`[ERROR] ${msg}`), + }, + dryRun: !!(params?.dry_run), + auth: authInfo, + }; + + try { + const result = await op.handler(ctx, (params || {}) as Record); + const latency = Date.now() - startTime; + + // Log request + broadcast to SSE + try { + await sql`INSERT INTO mcp_request_log (token_name, operation, latency_ms, status) + VALUES (${authInfo.clientId}, ${name}, ${latency}, ${'success'})`; + } catch { /* best effort */ } + + broadcastEvent({ + agent: authInfo.clientId, + operation: name, + scopes: authInfo.scopes.join(','), + latency_ms: latency, + status: 'success', + timestamp: new Date().toISOString(), + }); + + return { content: [{ type: 'text', text: JSON.stringify(result) }] }; + } catch (e) { + const latency = Date.now() - startTime; + const error = e instanceof OperationError ? e.toJSON() : { error: 'internal_error', message: e instanceof Error ? e.message : 'Unknown error' }; + + try { + await sql`INSERT INTO mcp_request_log (token_name, operation, latency_ms, status) + VALUES (${authInfo.clientId}, ${name}, ${latency}, ${'error'})`; + } catch { /* best effort */ } + + broadcastEvent({ + agent: authInfo.clientId, + operation: name, + latency_ms: latency, + status: 'error', + timestamp: new Date().toISOString(), + }); + + return { content: [{ type: 'text', text: JSON.stringify(error) }], isError: true }; + } + }); + + // Use StreamableHTTPServerTransport for stateless request handling + const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined as any }); + await server.connect(transport); + await transport.handleRequest(req, res, req.body); + }); + + // --------------------------------------------------------------------------- + // Start server + // --------------------------------------------------------------------------- + const clientCount = await sql`SELECT count(*)::int as count FROM oauth_clients`; + + app.listen(port, () => { + console.error(` +╔══════════════════════════════════════════════════════╗ +║ GBrain MCP Server v${VERSION.padEnd(37)}║ +╠══════════════════════════════════════════════════════╣ +║ Port: ${String(port).padEnd(40)}║ +║ Engine: ${(config.engine || 'pglite').padEnd(40)}║ +║ Clients: ${String((clientCount[0] as any).count).padEnd(40)}║ +║ DCR: ${(enableDcr ? 'enabled' : 'disabled').padEnd(40)}║ +║ Token TTL: ${(tokenTtl + 's').padEnd(40)}║ +╠══════════════════════════════════════════════════════╣ +║ Admin: http://localhost:${port}/admin${' '.repeat(Math.max(0, 19 - String(port).length))}║ +║ MCP: http://localhost:${port}/mcp${' '.repeat(Math.max(0, 21 - String(port).length))}║ +║ Health: http://localhost:${port}/health${' '.repeat(Math.max(0, 18 - String(port).length))}║ +╠══════════════════════════════════════════════════════╣ +║ Admin Token (paste into /admin login): ║ +║ ${bootstrapToken.substring(0, 50)} ║ +║ ${bootstrapToken.substring(50).padEnd(50)} ║ +╚══════════════════════════════════════════════════════╝ +`); + }); +} diff --git a/src/commands/serve.ts b/src/commands/serve.ts index 5e5d7f4e5..d25820a36 100644 --- a/src/commands/serve.ts +++ b/src/commands/serve.ts @@ -1,7 +1,22 @@ import type { BrainEngine } from '../core/engine.ts'; import { startMcpServer } from '../mcp/server.ts'; -export async function runServe(engine: BrainEngine) { - console.error('Starting GBrain MCP server (stdio)...'); - await startMcpServer(engine); +export async function runServe(engine: BrainEngine, args: string[] = []) { + const isHttp = args.includes('--http'); + + if (isHttp) { + const portIdx = args.indexOf('--port'); + const port = portIdx >= 0 ? parseInt(args[portIdx + 1]) || 3131 : 3131; + + const ttlIdx = args.indexOf('--token-ttl'); + const tokenTtl = ttlIdx >= 0 ? parseInt(args[ttlIdx + 1]) || 3600 : 3600; + + const enableDcr = args.includes('--enable-dcr'); + + const { runServeHttp } = await import('./serve-http.ts'); + await runServeHttp(engine, { port, tokenTtl, enableDcr }); + } else { + console.error('Starting GBrain MCP server (stdio)...'); + await startMcpServer(engine); + } } diff --git a/test/oauth.test.ts b/test/oauth.test.ts new file mode 100644 index 000000000..85a8f41a9 --- /dev/null +++ b/test/oauth.test.ts @@ -0,0 +1,404 @@ +import { describe, test, expect, beforeAll, afterAll } from 'bun:test'; +import { PGlite } from '@electric-sql/pglite'; +import { vector } from '@electric-sql/pglite/vector'; +import { pg_trgm } from '@electric-sql/pglite/contrib/pg_trgm'; +import { GBrainOAuthProvider } from '../src/core/oauth-provider.ts'; +import { hashToken, generateToken } from '../src/core/utils.ts'; +import { PGLITE_SCHEMA_SQL } from '../src/core/pglite-schema.ts'; + +// --------------------------------------------------------------------------- +// Test setup: in-memory PGLite with OAuth tables +// --------------------------------------------------------------------------- + +let db: PGlite; +let sql: (strings: TemplateStringsArray, ...values: unknown[]) => Promise; +let provider: GBrainOAuthProvider; + +beforeAll(async () => { + db = new PGlite({ extensions: { vector, pg_trgm } }); + await db.exec(PGLITE_SCHEMA_SQL); + + // Create a tagged template wrapper for PGLite + sql = async (strings: TemplateStringsArray, ...values: unknown[]) => { + const query = strings.reduce((acc, str, i) => acc + str + (i < values.length ? `$${i + 1}` : ''), ''); + const result = await db.query(query, values as any[]); + return result.rows; + }; + + provider = new GBrainOAuthProvider({ sql, tokenTtl: 60, refreshTtl: 300 }); +}); + +afterAll(async () => { + await db.close(); +}); + +// --------------------------------------------------------------------------- +// hashToken + generateToken utilities +// --------------------------------------------------------------------------- + +describe('hashToken', () => { + test('produces consistent SHA-256 hex', () => { + const hash = hashToken('test-token'); + expect(hash).toHaveLength(64); + expect(hashToken('test-token')).toBe(hash); // deterministic + }); + + test('different inputs produce different hashes', () => { + expect(hashToken('a')).not.toBe(hashToken('b')); + }); +}); + +describe('generateToken', () => { + test('produces prefixed random hex', () => { + const token = generateToken('gbrain_cl_'); + expect(token).toStartWith('gbrain_cl_'); + expect(token).toHaveLength('gbrain_cl_'.length + 64); // 32 bytes = 64 hex chars + }); + + test('tokens are unique', () => { + const a = generateToken('test_'); + const b = generateToken('test_'); + expect(a).not.toBe(b); + }); +}); + +// --------------------------------------------------------------------------- +// Client Registration +// --------------------------------------------------------------------------- + +describe('client registration', () => { + test('registerClientManual creates a client', async () => { + const { clientId, clientSecret } = await provider.registerClientManual( + 'test-agent', ['client_credentials'], 'read write', + ); + expect(clientId).toStartWith('gbrain_cl_'); + expect(clientSecret).toStartWith('gbrain_cs_'); + + // Verify client exists in DB + const client = await provider.clientsStore.getClient(clientId); + expect(client).toBeDefined(); + expect(client!.client_name).toBe('test-agent'); + }); + + test('getClient returns undefined for unknown client', async () => { + const client = await provider.clientsStore.getClient('nonexistent'); + expect(client).toBeUndefined(); + }); + + test('duplicate client_id is rejected', async () => { + const { clientId } = await provider.registerClientManual( + 'dup-test', ['client_credentials'], 'read', + ); + // Try to insert same client_id directly + await expect( + sql`INSERT INTO oauth_clients (client_id, client_name, scope) VALUES (${clientId}, ${'dup'}, ${'read'})`, + ).rejects.toThrow(); + }); +}); + +// --------------------------------------------------------------------------- +// Client Credentials Exchange +// --------------------------------------------------------------------------- + +describe('client credentials', () => { + let clientId: string; + let clientSecret: string; + + beforeAll(async () => { + const result = await provider.registerClientManual( + 'cc-test-agent', ['client_credentials'], 'read write', + ); + clientId = result.clientId; + clientSecret = result.clientSecret; + }); + + test('valid exchange returns access token', async () => { + const tokens = await provider.exchangeClientCredentials(clientId, clientSecret, 'read'); + expect(tokens.access_token).toStartWith('gbrain_at_'); + expect(tokens.token_type).toBe('bearer'); + expect(tokens.expires_in).toBe(60); + expect(tokens.scope).toBe('read'); + }); + + test('no refresh token issued for CC grant', async () => { + const tokens = await provider.exchangeClientCredentials(clientId, clientSecret, 'read'); + expect(tokens.refresh_token).toBeUndefined(); + }); + + test('wrong secret is rejected', async () => { + await expect( + provider.exchangeClientCredentials(clientId, 'wrong-secret', 'read'), + ).rejects.toThrow('Invalid client secret'); + }); + + test('client without CC grant is rejected', async () => { + const { clientId: noCC } = await provider.registerClientManual( + 'no-cc-agent', ['authorization_code'], 'read', + ); + await expect( + provider.exchangeClientCredentials(noCC, 'any-secret', 'read'), + ).rejects.toThrow('not authorized'); + }); + + test('scope is filtered to allowed scopes', async () => { + const tokens = await provider.exchangeClientCredentials(clientId, clientSecret, 'read write admin'); + // Client only has 'read write', admin should be filtered out + expect(tokens.scope).not.toContain('admin'); + }); +}); + +// --------------------------------------------------------------------------- +// Token Verification +// --------------------------------------------------------------------------- + +describe('verifyAccessToken', () => { + test('valid token returns auth info', async () => { + const { clientId, clientSecret } = await provider.registerClientManual( + 'verify-test', ['client_credentials'], 'read write', + ); + const tokens = await provider.exchangeClientCredentials(clientId, clientSecret, 'read'); + const authInfo = await provider.verifyAccessToken(tokens.access_token); + + expect(authInfo.clientId).toBe(clientId); + expect(authInfo.scopes).toContain('read'); + expect(authInfo.token).toBe(tokens.access_token); + }); + + test('expired token is rejected', async () => { + // Insert a token that's already expired + const expiredToken = generateToken('gbrain_at_'); + const hash = hashToken(expiredToken); + const firstClient = (await sql`SELECT client_id FROM oauth_clients LIMIT 1`)[0]; + await sql` + INSERT INTO oauth_tokens (token_hash, token_type, client_id, scopes, expires_at) + VALUES (${hash}, ${'access'}, ${firstClient.client_id as string}, ${'{read}'}, ${Math.floor(Date.now() / 1000) - 100}) + `; + await expect(provider.verifyAccessToken(expiredToken)).rejects.toThrow('expired'); + }); + + test('unknown token is rejected', async () => { + await expect(provider.verifyAccessToken('nonexistent-token')).rejects.toThrow('Invalid token'); + }); + + test('legacy access_tokens fallback works', async () => { + // Insert a legacy bearer token + const legacyToken = generateToken('gbrain_'); + const hash = hashToken(legacyToken); + await sql` + INSERT INTO access_tokens (id, name, token_hash) + VALUES (${crypto.randomUUID()}, ${'legacy-agent'}, ${hash}) + `; + + const authInfo = await provider.verifyAccessToken(legacyToken); + expect(authInfo.clientId).toBe('legacy-agent'); + expect(authInfo.scopes).toEqual(['read', 'write', 'admin']); // grandfathered full access + }); +}); + +// --------------------------------------------------------------------------- +// Token Revocation +// --------------------------------------------------------------------------- + +describe('revokeToken', () => { + test('revoked token no longer verifies', async () => { + const { clientId, clientSecret } = await provider.registerClientManual( + 'revoke-test', ['client_credentials'], 'read', + ); + const tokens = await provider.exchangeClientCredentials(clientId, clientSecret, 'read'); + + // Verify token works + const authInfo = await provider.verifyAccessToken(tokens.access_token); + expect(authInfo.clientId).toBe(clientId); + + // Revoke it + const client = (await provider.clientsStore.getClient(clientId))!; + await provider.revokeToken!(client, { token: tokens.access_token }); + + // Should no longer verify + await expect(provider.verifyAccessToken(tokens.access_token)).rejects.toThrow(); + }); + + test('revoking already-revoked token is a no-op', async () => { + // This should not throw + const client = (await provider.clientsStore.getClient( + (await sql`SELECT client_id FROM oauth_clients LIMIT 1`)[0].client_id as string, + ))!; + await provider.revokeToken!(client, { token: 'already-gone' }); + // No error = pass + }); +}); + +// --------------------------------------------------------------------------- +// Authorization Code Flow +// --------------------------------------------------------------------------- + +describe('authorization code flow', () => { + test('code issuance and exchange', async () => { + const { clientId } = await provider.registerClientManual( + 'authcode-test', ['authorization_code'], 'read write', + ['http://localhost:3000/callback'], + ); + const client = (await provider.clientsStore.getClient(clientId))!; + + // Mock Express response for authorize + let redirectUrl = ''; + const mockRes = { + redirect: (url: string) => { redirectUrl = url; }, + } as any; + + await provider.authorize(client, { + codeChallenge: 'test-challenge-hash', + redirectUri: 'http://localhost:3000/callback', + scopes: ['read', 'write'], + state: 'test-state', + }, mockRes); + + expect(redirectUrl).toContain('code=gbrain_code_'); + expect(redirectUrl).toContain('state=test-state'); + + // Extract code from redirect URL + const url = new URL(redirectUrl); + const code = url.searchParams.get('code')!; + + // Exchange code for tokens + const tokens = await provider.exchangeAuthorizationCode(client, code); + expect(tokens.access_token).toStartWith('gbrain_at_'); + expect(tokens.refresh_token).toBeDefined(); // Auth code flow includes refresh + }); + + test('code is single-use', async () => { + const { clientId } = await provider.registerClientManual( + 'single-use-test', ['authorization_code'], 'read', + ['http://localhost:3000/callback'], + ); + const client = (await provider.clientsStore.getClient(clientId))!; + + let redirectUrl = ''; + const mockRes = { redirect: (url: string) => { redirectUrl = url; } } as any; + + await provider.authorize(client, { + codeChallenge: 'challenge', + redirectUri: 'http://localhost:3000/callback', + scopes: ['read'], + }, mockRes); + + const code = new URL(redirectUrl).searchParams.get('code')!; + + // First exchange works + await provider.exchangeAuthorizationCode(client, code); + + // Second exchange fails (code consumed) + await expect(provider.exchangeAuthorizationCode(client, code)).rejects.toThrow(); + }); + + test('expired code is rejected', async () => { + // Insert an already-expired code + const expiredCode = generateToken('gbrain_code_'); + const hash = hashToken(expiredCode); + const firstClient = (await sql`SELECT client_id FROM oauth_clients LIMIT 1`)[0]; + + await sql` + INSERT INTO oauth_codes (code_hash, client_id, scopes, code_challenge, + redirect_uri, expires_at) + VALUES (${hash}, ${firstClient.client_id as string}, ${'{read}'}, + ${'challenge'}, ${'http://localhost/cb'}, ${Math.floor(Date.now() / 1000) - 100}) + `; + + const client = (await provider.clientsStore.getClient(firstClient.client_id as string))!; + await expect(provider.exchangeAuthorizationCode(client, expiredCode)).rejects.toThrow(); + }); +}); + +// --------------------------------------------------------------------------- +// Refresh Token +// --------------------------------------------------------------------------- + +describe('refresh token', () => { + test('valid refresh rotates tokens', async () => { + const { clientId } = await provider.registerClientManual( + 'refresh-test', ['authorization_code'], 'read write', + ['http://localhost:3000/callback'], + ); + const client = (await provider.clientsStore.getClient(clientId))!; + + let redirectUrl = ''; + const mockRes = { redirect: (url: string) => { redirectUrl = url; } } as any; + + await provider.authorize(client, { + codeChallenge: 'challenge', + redirectUri: 'http://localhost:3000/callback', + scopes: ['read', 'write'], + }, mockRes); + + const code = new URL(redirectUrl).searchParams.get('code')!; + const tokens = await provider.exchangeAuthorizationCode(client, code); + + // Refresh + const newTokens = await provider.exchangeRefreshToken(client, tokens.refresh_token!, ['read']); + expect(newTokens.access_token).not.toBe(tokens.access_token); + expect(newTokens.refresh_token).toBeDefined(); + expect(newTokens.refresh_token).not.toBe(tokens.refresh_token); // rotated + + // Old refresh token should no longer work + await expect(provider.exchangeRefreshToken(client, tokens.refresh_token!)).rejects.toThrow(); + }); +}); + +// --------------------------------------------------------------------------- +// Token Sweep +// --------------------------------------------------------------------------- + +describe('sweepExpiredTokens', () => { + test('removes expired tokens', async () => { + // Insert some expired tokens + const firstClient = (await sql`SELECT client_id FROM oauth_clients LIMIT 1`)[0]; + const expired1 = hashToken(generateToken('sweep_')); + const expired2 = hashToken(generateToken('sweep_')); + + await sql`INSERT INTO oauth_tokens (token_hash, token_type, client_id, scopes, expires_at) + VALUES (${expired1}, ${'access'}, ${firstClient.client_id as string}, ${'{read}'}, ${1})`; + await sql`INSERT INTO oauth_tokens (token_hash, token_type, client_id, scopes, expires_at) + VALUES (${expired2}, ${'access'}, ${firstClient.client_id as string}, ${'{read}'}, ${2})`; + + await provider.sweepExpiredTokens(); + + // Verify they're gone + const remaining = await sql`SELECT count(*)::int as count FROM oauth_tokens WHERE expires_at < 100`; + expect(remaining[0].count).toBe(0); + }); +}); + +// --------------------------------------------------------------------------- +// Scope Annotations +// --------------------------------------------------------------------------- + +describe('operation scope annotations', () => { + test('all operations have a scope', () => { + const { operations } = require('../src/core/operations.ts'); + for (const op of operations) { + expect(op.scope, `${op.name} missing scope`).toBeDefined(); + expect(['read', 'write', 'admin']).toContain(op.scope); + } + }); + + test('mutating operations are write or admin scoped', () => { + const { operations } = require('../src/core/operations.ts'); + for (const op of operations) { + if (op.mutating) { + expect(['write', 'admin'], `${op.name} is mutating but not write/admin`).toContain(op.scope); + } + } + }); + + test('sync_brain and file_upload are localOnly', () => { + const { operationsByName } = require('../src/core/operations.ts'); + expect(operationsByName.sync_brain.localOnly).toBe(true); + expect(operationsByName.file_upload.localOnly).toBe(true); + }); + + test('file_list and file_url are localOnly', () => { + const { operationsByName } = require('../src/core/operations.ts'); + expect(operationsByName.file_list.localOnly).toBe(true); + expect(operationsByName.file_url.localOnly).toBe(true); + }); +});