diff --git a/src/auth/session-auth-prehandler.ts b/src/auth/session-auth-prehandler.ts index 335f15ad..7be0aeb7 100644 --- a/src/auth/session-auth-prehandler.ts +++ b/src/auth/session-auth-prehandler.ts @@ -27,12 +27,19 @@ export function createSessionAuthPreHandler ( } // Skip authorization for well-known endpoints - if (request.url.startsWith('/.well-known/')) { + if (request.url.startsWith('/.well-known/') || request.url.startsWith('/mcp/.well-known')) { return } - // Skip authorization for the start of the OAuth authorization flow. - if (request.url.startsWith('/oauth/authorize')) { + // Skip authorization for OAuth flow endpoints (authorize initiates, callback receives code, register is pre-auth) + if (request.url.startsWith('/oauth/authorize') || request.url.startsWith('/oauth/callback') || request.url.startsWith('/oauth/register')) { + return + } + + // Skip authorization for custom excluded paths + if (config.excludedPaths?.some(path => + typeof path === 'string' ? request.url.startsWith(path) : path.test(request.url) + )) { return } diff --git a/src/auth/token-refresh-service.ts b/src/auth/token-refresh-service.ts index 99ebb416..d898dfde 100644 --- a/src/auth/token-refresh-service.ts +++ b/src/auth/token-refresh-service.ts @@ -216,16 +216,38 @@ export class TokenRefreshService { * This is separated from coordination logic for clarity */ private async performTokenRefresh (): Promise { - // This is a simplified implementation placeholder - // In a complete implementation, you would: - // 1. Query session store for sessions with expiring tokens - // 2. Iterate through sessions and refresh tokens as needed - // 3. Update session store with new token information - // 4. Send notifications to affected sessions + let checkedSessions = 0 + let refreshedCandidates = 0 + + for await (const session of this.sessionStore.iterate()) { + checkedSessions += 1 + if ( + !session.authorization || + !session.tokenRefresh || + !shouldAttemptRefresh(session.authorization, session.tokenRefresh) + ) { + continue + } + + refreshedCandidates += 1 + try { + await this.refreshSessionToken(session.id) + } catch (error) { + if (this.fastify) { + this.fastify.log.warn({ + error, + sessionId: session.id, + instanceId: this.instanceId + }, 'Background token refresh failed for session') + } + } + } if (this.fastify) { this.fastify.log.debug({ - instanceId: this.instanceId + instanceId: this.instanceId, + checkedSessions, + refreshedCandidates }, 'Token refresh service check completed') } } diff --git a/src/brokers/redis-message-broker.ts b/src/brokers/redis-message-broker.ts index e47efa98..d874def6 100644 --- a/src/brokers/redis-message-broker.ts +++ b/src/brokers/redis-message-broker.ts @@ -53,7 +53,8 @@ export class RedisMessageBroker implements MessageBroker { host: redis.options.host, password: redis.options.password, db: redis.options.db || 0, - family: redis.options.family || 4 + family: redis.options.family || 4, + enableReadyCheck: false }) this.closeTimeoutMs = options.closeTimeoutMs ?? DEFAULT_CLOSE_TIMEOUT_MS this.onCloseTimeout = options.onCloseTimeout diff --git a/src/index.ts b/src/index.ts index 9af3cef9..959b327a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -18,6 +18,8 @@ import routes from './routes/mcp.ts' import wellKnownRoutes from './routes/well-known.ts' import { TokenValidator } from './auth/token-validator.ts' import { createAuthPreHandler } from './auth/prehandler.ts' +import { createSessionAuthPreHandler } from './auth/session-auth-prehandler.ts' +import { registerTokenRefreshService } from './auth/token-refresh-service.ts' import oauthClientPlugin from './auth/oauth-client.ts' import authRoutesPlugin from './routes/auth-routes.ts' import { quitWithTimeout } from './redis-quit-with-timeout.ts' @@ -139,12 +141,22 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption if (opts.authorization?.enabled) { tokenValidator = new TokenValidator(opts.authorization, app) - // Register authorization preHandler for all routes - app.addHook('preHandler', createAuthPreHandler(opts.authorization, tokenValidator)) - // Register OAuth client plugin if configured if (opts.authorization.oauth2Client) { await app.register(oauthClientPlugin, opts.authorization.oauth2Client) + + app.addHook( + 'preHandler', + createSessionAuthPreHandler({ + config: opts.authorization, + tokenValidator, + sessionStore, + oauthClient: app.oauthClient + }) + ) + } else { + // Register authorization preHandler for all routes + app.addHook('preHandler', createAuthPreHandler(opts.authorization, tokenValidator)) } } @@ -157,8 +169,16 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption if (opts.authorization?.enabled && opts.authorization?.oauth2Client) { await app.register(authRoutesPlugin, { sessionStore, + tokenValidator: tokenValidator ?? undefined, dcrHooks: opts.authorization.dcrHooks }) + + await registerTokenRefreshService(app, { + sessionStore, + messageBroker, + oauthClient: app.oauthClient, + ...(redis ? { redis } : {}) + }) } // AJV instance and compiled-schema cache scoped to this plugin registration @@ -216,7 +236,9 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption } streams.clear() // Collect unsubscribe promises for parallel execution - unsubscribePromises.push(messageBroker.unsubscribe(`mcp/session/${sessionId}/message`)) + unsubscribePromises.push( + messageBroker.unsubscribe(`mcp/session/${sessionId}/message`) + ) } localStreams.clear() @@ -300,8 +322,8 @@ export type { // Export authorization types export type { AuthorizationConfig, - TokenValidationResult, ProtectedResourceMetadata, + TokenValidationResult, TokenIntrospectionResponse, IntrospectionAuthConfig, DCRRequest, diff --git a/src/routes/auth-routes.ts b/src/routes/auth-routes.ts index c08b4b8a..9678e12f 100644 --- a/src/routes/auth-routes.ts +++ b/src/routes/auth-routes.ts @@ -2,6 +2,8 @@ import type { FastifyInstance, FastifyPluginAsync } from 'fastify' import fp from 'fastify-plugin' import { Type } from '@sinclair/typebox' import type { PKCEChallenge } from '../auth/oauth-client.ts' +import type { TokenValidator } from '../auth/token-validator.ts' +import { createAuthorizationContext, createTokenRefreshInfo } from '../auth/token-utils.ts' import type { SessionStore } from '../stores/session-store.ts' import type { DCRHooks, DCRRequest, DCRResponse } from '../types/auth-types.ts' @@ -26,6 +28,7 @@ export interface TokenRefreshBody { export interface AuthRoutesOptions { sessionStore: SessionStore + tokenValidator?: TokenValidator /** DCR hooks for custom request/response processing */ dcrHooks?: DCRHooks } @@ -101,7 +104,7 @@ const LogoutResponse = Type.Object({ }) const authRoutesPlugin: FastifyPluginAsync = async (fastify: FastifyInstance, opts) => { - const { sessionStore } = opts + const { sessionStore, tokenValidator } = opts // Initiate OAuth authorization flow fastify.get('/oauth/authorize', { schema: { @@ -199,9 +202,6 @@ const authRoutesPlugin: FastifyPluginAsync = async (fastify: const sessionData = sessionMetadata.authSession as AuthSession - // Clean up session data - await sessionStore.delete(state) - // Exchange authorization code for tokens (include redirect_uri for OIDC compliance) const tokens = await fastify.oauthClient.exchangeCodeForToken( code, @@ -211,6 +211,36 @@ const authRoutesPlugin: FastifyPluginAsync = async (fastify: sessionData.callbackUrl ) + if (tokenValidator) { + const validationResult = await tokenValidator.validateToken(tokens.access_token) + if (!validationResult.valid || !validationResult.payload) { + return reply.status(500).send({ + error: 'token_validation_failed', + error_description: validationResult.error || 'Access token validation failed' + }) + } + + const authContext = createAuthorizationContext( + validationResult.payload, + tokens.access_token, + { + refreshToken: tokens.refresh_token, + authorizationServer: validationResult.payload.iss + } + ) + + const refreshInfo = tokens.refresh_token + ? createTokenRefreshInfo( + tokens.refresh_token, + authContext.clientId || '', + authContext.authorizationServer || '', + authContext.scopes || [] + ) + : undefined + + await sessionStore.updateAuthorization(state, authContext, refreshInfo) + } + // Return tokens to client or redirect with tokens if (sessionData.originalUrl) { const redirectUrl = new URL(sessionData.originalUrl) diff --git a/src/stores/memory-session-store.ts b/src/stores/memory-session-store.ts index 55f9cc5f..ebe39bfb 100644 --- a/src/stores/memory-session-store.ts +++ b/src/stores/memory-session-store.ts @@ -41,6 +41,12 @@ export class MemorySessionStore implements SessionStore { return session ? { ...session } : null } + async * iterate (): AsyncIterable { + for (const session of this.sessions.values()) { + yield { ...session } + } + } + async delete (sessionId: string): Promise { // Clean up token mappings for this session const session = this.sessions.get(sessionId) diff --git a/src/stores/redis-session-store.ts b/src/stores/redis-session-store.ts index 9dd80284..9a1fe676 100644 --- a/src/stores/redis-session-store.ts +++ b/src/stores/redis-session-store.ts @@ -3,6 +3,21 @@ import type { JSONRPCMessage } from '../schema.ts' import type { SessionStore, SessionMetadata } from './session-store.ts' import type { AuthorizationContext, TokenRefreshInfo } from '../types/auth-types.ts' +function normalizeAuthorizationContext (authorization: AuthorizationContext): AuthorizationContext { + return { + ...authorization, + expiresAt: authorization.expiresAt ? new Date(authorization.expiresAt) : undefined, + issuedAt: authorization.issuedAt ? new Date(authorization.issuedAt) : undefined, + } +} + +function normalizeTokenRefreshInfo (tokenRefresh: TokenRefreshInfo): TokenRefreshInfo { + return { + ...tokenRefresh, + lastRefreshAt: tokenRefresh.lastRefreshAt ? new Date(tokenRefresh.lastRefreshAt) : undefined, + } +} + export class RedisSessionStore implements SessionStore { private redis: Redis private maxMessages: number @@ -79,6 +94,10 @@ export class RedisSessionStore implements SessionStore { const sessionKey = `session:${sessionId}` const result = await this.redis.hgetall(sessionKey) + return this.deserializeSession(result) + } + + private deserializeSession (result: Record): SessionMetadata | null { if (!result.id) { return null } @@ -95,7 +114,7 @@ export class RedisSessionStore implements SessionStore { // Parse authorization context if present if (result.authorization) { try { - metadata.authorization = JSON.parse(result.authorization) + metadata.authorization = normalizeAuthorizationContext(JSON.parse(result.authorization)) } catch (error) { // Ignore parsing errors for authorization context } @@ -103,7 +122,7 @@ export class RedisSessionStore implements SessionStore { if (result.tokenRefresh) { try { - metadata.tokenRefresh = JSON.parse(result.tokenRefresh) + metadata.tokenRefresh = normalizeTokenRefreshInfo(JSON.parse(result.tokenRefresh)) } catch (error) { // Ignore parsing errors for token refresh } @@ -120,6 +139,43 @@ export class RedisSessionStore implements SessionStore { return metadata } + async * iterate (): AsyncIterable { + let cursor = '0' + + do { + const [nextCursor, keys] = await this.redis.scan(cursor, 'MATCH', 'session:*', 'COUNT', 100) + cursor = nextCursor + const sessionKeys = keys.filter(key => !key.endsWith(':history')) + + if (sessionKeys.length === 0) { + continue + } + + // Fetch one SCAN page at a time to bound memory while avoiding one + // Redis round trip per session. + const pipeline = this.redis.pipeline() + for (const key of sessionKeys) { + pipeline.hgetall(key) + } + + const results = await pipeline.exec() + if (!results) { + continue + } + + for (const [error, result] of results) { + if (error || !result) { + continue + } + + const session = this.deserializeSession(result as Record) + if (session) { + yield session + } + } + } while (cursor !== '0') + } + async delete (sessionId: string): Promise { const sessionKey = `session:${sessionId}` const historyKey = `session:${sessionId}:history` diff --git a/src/stores/session-store.ts b/src/stores/session-store.ts index 1a2083c7..ec9ec716 100644 --- a/src/stores/session-store.ts +++ b/src/stores/session-store.ts @@ -24,6 +24,7 @@ export interface SessionStore { */ update(metadata: SessionMetadata): Promise get(sessionId: string): Promise + iterate(): AsyncIterable delete(sessionId: string): Promise cleanup(): Promise diff --git a/test/oauth-routes.test.ts b/test/oauth-routes.test.ts index 05964a18..3b260d9c 100644 --- a/test/oauth-routes.test.ts +++ b/test/oauth-routes.test.ts @@ -5,6 +5,7 @@ import Fastify from 'fastify' import oauthClientPlugin from '../src/auth/oauth-client.ts' import authRoutesPlugin from '../src/routes/auth-routes.ts' import { MemorySessionStore } from '../src/stores/memory-session-store.ts' +import type { TokenValidator } from '../src/auth/token-validator.ts' const originalDispatcher = getGlobalDispatcher() const mockAgent = new MockAgent() @@ -98,6 +99,81 @@ describe('OAuth Routes', () => { assert.strictEqual(body.token_type, 'Bearer') }) + test('should persist authorization and refresh metadata on OAuth callback', async (t) => { + const mockPool = mockAgent.get('https://auth.example.com') + mockPool.intercept({ + path: '/oauth/token', + method: 'POST' + }).reply(200, { + access_token: 'callback-access-token', + token_type: 'Bearer', + expires_in: 3600, + refresh_token: 'callback-refresh-token', + scope: 'read write' + }) + + const fastify = Fastify() + t.after(async () => { + await fastify.close() + }) + + const config = { + clientId: 'test-client', + clientSecret: 'test-secret', + authorizationServer: 'https://auth.example.com' + } + + const tokenValidator = { + validateToken: async (token: string) => { + assert.strictEqual(token, 'callback-access-token') + return { + valid: true, + payload: { + sub: 'user-123', + client_id: 'test-client', + scope: 'read write', + aud: 'https://mcp.example.com', + exp: Math.floor(Date.now() / 1000) + 3600, + iat: Math.floor(Date.now() / 1000), + iss: 'https://auth.example.com' + } + } + } + } as TokenValidator + + await fastify.register(oauthClientPlugin, config) + const sessionStore = new MemorySessionStore(100) + await fastify.register(authRoutesPlugin, { sessionStore, tokenValidator }) + + const authResponse = await fastify.inject({ + method: 'GET', + url: '/oauth/authorize' + }) + + const locationHeader = authResponse.headers.location + assert.ok(locationHeader, 'Location header should be present') + const location = new URL(locationHeader!) + const state = location.searchParams.get('state') + assert.ok(state, 'State should be present') + + const callbackResponse = await fastify.inject({ + method: 'GET', + url: `/oauth/callback?code=test-code&state=${state}` + }) + + assert.strictEqual(callbackResponse.statusCode, 200) + + const persistedSession = await sessionStore.get(state!) + assert.ok(persistedSession) + assert.strictEqual(persistedSession.authorization?.userId, 'user-123') + assert.strictEqual(persistedSession.authorization?.clientId, 'test-client') + assert.deepStrictEqual(persistedSession.authorization?.scopes, ['read', 'write']) + assert.ok(persistedSession.authorization?.tokenHash) + assert.strictEqual(persistedSession.tokenRefresh?.refreshToken, 'callback-refresh-token') + assert.strictEqual(persistedSession.tokenRefresh?.clientId, 'test-client') + assert.deepStrictEqual(persistedSession.tokenRefresh?.scopes, ['read', 'write']) + }) + test('should handle OAuth callback with error', async (t) => { const fastify = Fastify() t.after(async () => { diff --git a/test/redis-message-broker.test.ts b/test/redis-message-broker.test.ts index f233ce71..464bdd5e 100644 --- a/test/redis-message-broker.test.ts +++ b/test/redis-message-broker.test.ts @@ -258,4 +258,19 @@ describe('RedisMessageBroker', () => { // Close should not throw - will be handled by t.after() }) + + testWithRedis('should disable ready checks on pub/sub connections', async (redis, t) => { + const broker = new RedisMessageBroker(redis) + t.after(() => broker.close()) + + const internalBroker = broker as unknown as { + emitter: { + subConn: { options: { enableReadyCheck: boolean } } + pubConn: { options: { enableReadyCheck: boolean } } + } + } + + assert.strictEqual(internalBroker.emitter.subConn.options.enableReadyCheck, false) + assert.strictEqual(internalBroker.emitter.pubConn.options.enableReadyCheck, false) + }) }) diff --git a/test/redis-session-store.test.ts b/test/redis-session-store.test.ts index 03f1f510..dfc15b92 100644 --- a/test/redis-session-store.test.ts +++ b/test/redis-session-store.test.ts @@ -35,6 +35,35 @@ describe('RedisSessionStore', () => { assert.strictEqual(result, null) }) + testWithRedis('should iterate sessions without returning message histories', async (redis) => { + const store = new RedisSessionStore({ redis, maxMessages: 100 }) + const now = new Date() + + await store.create({ + id: 'iterated-session-1', + eventId: 0, + createdAt: now, + lastActivity: now + }) + await store.create({ + id: 'iterated-session-2', + eventId: 0, + createdAt: now, + lastActivity: now + }) + await store.addMessage('iterated-session-1', '1', { + jsonrpc: '2.0', + method: 'test' + }) + + const ids: string[] = [] + for await (const session of store.iterate()) { + ids.push(session.id) + } + + assert.deepStrictEqual(ids.sort(), ['iterated-session-1', 'iterated-session-2']) + }) + testWithRedis('should delete session and its history', async (redis) => { const store = new RedisSessionStore({ redis, maxMessages: 100 }) diff --git a/test/session-auth.test.ts b/test/session-auth.test.ts index 74727f11..6c0831af 100644 --- a/test/session-auth.test.ts +++ b/test/session-auth.test.ts @@ -298,6 +298,44 @@ describe('Session-Based Authorization', () => { assert.strictEqual(body.message, 'public') }) + test('should skip authorization for MCP well-known endpoints', async (t) => { + const fastify = Fastify() + t.after(async () => { + await fastify.close() + }) + + const config = { + enabled: true, + authorizationServers: ['https://auth.example.com'], + resourceUri: 'https://api.example.com', + tokenValidation: { + jwksUri: 'https://auth.example.com/.well-known/jwks.json', + validateAudience: true + } + } + + const sessionStore = new MemorySessionStore(100) + const tokenValidator = new TokenValidator(config, fastify) + + const preHandler = createSessionAuthPreHandler({ + config, + tokenValidator, + sessionStore + }) + + fastify.addHook('preHandler', preHandler) + fastify.get('/mcp/.well-known/oauth-protected-resource', async () => ({ message: 'public' })) + + const response = await fastify.inject({ + method: 'GET', + url: '/mcp/.well-known/oauth-protected-resource' + }) + + assert.strictEqual(response.statusCode, 200) + const body = JSON.parse(response.body) + assert.strictEqual(body.message, 'public') + }) + test('should skip authorization for the start of the OAuth authorization flow', async (t) => { const fastify = Fastify() t.after(async () => { @@ -337,6 +375,99 @@ describe('Session-Based Authorization', () => { assert.strictEqual(body.message, 'public') }) + test('should skip authorization for OAuth callback and registration endpoints', async (t) => { + const fastify = Fastify() + t.after(async () => { + await fastify.close() + }) + + const config = { + enabled: true, + authorizationServers: ['https://auth.example.com'], + resourceUri: 'https://api.example.com', + tokenValidation: { + jwksUri: 'https://auth.example.com/.well-known/jwks.json', + validateAudience: true + } + } + + const sessionStore = new MemorySessionStore(100) + const tokenValidator = new TokenValidator(config, fastify) + + const preHandler = createSessionAuthPreHandler({ + config, + tokenValidator, + sessionStore + }) + + fastify.addHook('preHandler', preHandler) + fastify.get('/oauth/callback', async () => ({ message: 'callback-public' })) + fastify.post('/oauth/register', async () => ({ message: 'register-public' })) + + const callbackResponse = await fastify.inject({ + method: 'GET', + url: '/oauth/callback' + }) + + assert.strictEqual(callbackResponse.statusCode, 200) + assert.strictEqual(JSON.parse(callbackResponse.body).message, 'callback-public') + + const registerResponse = await fastify.inject({ + method: 'POST', + url: '/oauth/register' + }) + + assert.strictEqual(registerResponse.statusCode, 200) + assert.strictEqual(JSON.parse(registerResponse.body).message, 'register-public') + }) + + test('should skip authorization for excluded paths', async (t) => { + const fastify = Fastify() + t.after(async () => { + await fastify.close() + }) + + const config = { + enabled: true, + authorizationServers: ['https://auth.example.com'], + resourceUri: 'https://api.example.com', + excludedPaths: ['/health', /^\/public\//], + tokenValidation: { + jwksUri: 'https://auth.example.com/.well-known/jwks.json', + validateAudience: true + } + } + + const sessionStore = new MemorySessionStore(100) + const tokenValidator = new TokenValidator(config, fastify) + + const preHandler = createSessionAuthPreHandler({ + config, + tokenValidator, + sessionStore + }) + + fastify.addHook('preHandler', preHandler) + fastify.get('/health/ready', async () => ({ message: 'ok' })) + fastify.get('/public/info', async () => ({ message: 'public' })) + + const healthResponse = await fastify.inject({ + method: 'GET', + url: '/health/ready' + }) + + assert.strictEqual(healthResponse.statusCode, 200) + assert.strictEqual(JSON.parse(healthResponse.body).message, 'ok') + + const publicResponse = await fastify.inject({ + method: 'GET', + url: '/public/info' + }) + + assert.strictEqual(publicResponse.statusCode, 200) + assert.strictEqual(JSON.parse(publicResponse.body).message, 'public') + }) + test('should handle invalid bearer token format', async (t) => { const fastify = Fastify() t.after(async () => { diff --git a/test/token-refresh.test.ts b/test/token-refresh.test.ts index b02d13b7..87d66cf6 100644 --- a/test/token-refresh.test.ts +++ b/test/token-refresh.test.ts @@ -46,6 +46,74 @@ describe('Token Refresh Service', () => { service.stop() }) + test('should refresh expiring sessions during background checks', async (t) => { + const fastify = Fastify() + t.after(async () => { + await fastify.close() + }) + + const sessionStore = new MemorySessionStore(100) + const messageBroker = new MemoryMessageBroker() + let refreshCalls = 0 + + const mockOAuthClient = { + refreshToken: async (refreshToken: string) => { + refreshCalls += 1 + assert.strictEqual(refreshToken, 'refresh-123') + return { + access_token: `background-access-token-${refreshCalls}`, + token_type: 'Bearer', + expires_in: 3600, + refresh_token: `background-refresh-token-${refreshCalls}`, + scope: 'read write' + } + } + } + + const service = new TokenRefreshService({ + sessionStore, + messageBroker, + oauthClient: mockOAuthClient, + checkIntervalMs: 25, + refreshBufferMinutes: 1 + }) + + const session = { + id: 'session-background', + eventId: 0, + createdAt: new Date(), + lastActivity: new Date() + } + + await sessionStore.create(session) + await sessionStore.updateAuthorization(session.id, { + userId: 'user123', + tokenHash: hashToken('old-token'), + expiresAt: new Date(Date.now() + 30000) + }, { + refreshToken: 'refresh-123', + clientId: 'client456', + authorizationServer: 'https://auth.example.com', + scopes: ['read', 'write'], + refreshAttempts: 0 + }) + + service.start(fastify) + t.after(async () => { + await service.stop() + }) + + await new Promise(resolve => setTimeout(resolve, 80)) + + assert.strictEqual(refreshCalls, 1) + + const updatedSession = await sessionStore.get(session.id) + assert.ok(updatedSession?.authorization?.tokenHash) + assert.strictEqual(updatedSession?.tokenRefresh?.refreshToken, 'background-refresh-token-1') + assert.ok(updatedSession?.authorization?.expiresAt instanceof Date) + assert.ok(updatedSession!.authorization!.expiresAt! > new Date(Date.now() + 3000000)) + }) + test('should handle manual token refresh', async (_t) => { const sessionStore = new MemorySessionStore(100) const messageBroker = new MemoryMessageBroker()