Merge pull request #133 from QuantGeekDev/feat/auth-context

Feat/auth context

Author: QuantGeekDev

package.json

@@ -1,6 +1,6 @@
 {
   "name": "mcp-framework",
-  "version": "0.2.16",
+  "version": "0.2.17-beta.2",
   "description": "Framework for building Model Context Protocol (MCP) servers in Typescript",
   "type": "module",
   "author": "Alex Andru <[email protected]>",

src/auth/providers/oauth.ts

@@ -81,7 +81,7 @@ export class OAuthAuthProvider implements AuthProvider {
       logger.debug(`Token claims - sub: ${claims.sub}, scope: ${claims.scope || 'N/A'}`);
 
       return {
-        data: claims,
+        data: { ...claims, token },
       };
     } catch (error) {
       if (error instanceof Error) {

src/auth/validators/jwt-validator.ts

@@ -112,9 +112,11 @@ export class JWTValidator {
       };
 
       // Only validate audience if not set to wildcard
-      if (this.config.audience !== '*') {
-        options.audience = this.config.audience;
-      }
+      // For Cognito Access Tokens, 'aud' is missing but 'client_id' is present.
+      // We disable the library's strict check and handle it manually in the callback.
+      // if (this.config.audience !== '*') {
+      //   options.audience = this.config.audience;
+      // }
 
       jwt.verify(token, publicKey, options, (err, decoded) => {
         if (err) {
@@ -151,10 +153,32 @@ export class JWTValidator {
           return;
         }
 
-        // Only require aud claim if not set to wildcard
-        if (this.config.audience !== '*' && !claims.aud) {
-          reject(new Error('Token missing required claim: aud'));
-          return;
+        // Only require aud/client_id claim if not set to wildcard
+        if (this.config.audience !== '*') {
+          const aud = claims.aud;
+          const clientId = claims.client_id as string | undefined;
+          const expectedAudience = this.config.audience;
+
+          let isValidAudience = false;
+
+          // Check 'aud' claim (ID Tokens)
+          if (aud) {
+            if (Array.isArray(aud)) {
+              if (aud.includes(expectedAudience)) isValidAudience = true;
+            } else {
+              if (aud === expectedAudience) isValidAudience = true;
+            }
+          }
+          
+          // Check 'client_id' claim (Access Tokens)
+          if (!isValidAudience && clientId) {
+            if (clientId === expectedAudience) isValidAudience = true;
+          }
+
+          if (!isValidAudience) {
+             reject(new Error(`Token audience mismatch. Expected ${expectedAudience}, got aud: ${aud}, client_id: ${clientId}`));
+             return;
+          }
         }
 
         if (!claims.exp) {

src/index.ts

@@ -10,3 +10,6 @@ export * from './auth/index.js';
 export type { SSETransportConfig } from './transports/sse/types.js';
 export type { HttpStreamTransportConfig } from './transports/http/types.js';
 export { HttpStreamTransport } from './transports/http/server.js';
+
+export { requestContext, getRequestContext, runInRequestContext } from './utils/requestContext.js';
+export type { RequestContextData } from './utils/requestContext.js';

src/transports/http/server.ts

@@ -8,6 +8,8 @@ import { logger } from '../../core/Logger.js';
 import { ProtectedResourceMetadata } from '../../auth/metadata/protected-resource.js';
 import { handleAuthentication } from '../utils/auth-handler.js';
 import { initializeOAuthMetadata } from '../utils/oauth-metadata.js';
+import { requestContext, RequestContextData } from '../../utils/requestContext.js';
+import { AuthResult } from '../../auth/types.js';
 
 export class HttpStreamTransport extends AbstractTransport {
   readonly type = 'http-stream';
@@ -120,14 +122,17 @@ export class HttpStreamTransport extends AbstractTransport {
 
     // Perform authentication check once at the beginning
     const authEndpoint = isInitialize ? 'sse' : 'messages';
+    let authData: RequestContextData = {};
+
     if (this._config.auth?.endpoints?.[authEndpoint] !== false) {
-      const isAuthenticated = await handleAuthentication(
+      const authResult = await handleAuthentication(
         req,
         res,
         this._config.auth,
         isInitialize ? 'initialize' : 'message'
       );
-      if (!isAuthenticated) return;
+      if (!authResult) return;
+      authData = (authResult as AuthResult).data as RequestContextData || {};
     }
 
     // Handle different request scenarios
@@ -168,7 +173,9 @@ export class HttpStreamTransport extends AbstractTransport {
         }
       };
 
-      await transport.handleRequest(req, res, body);
+      await requestContext.run(authData, async () => {
+        await transport.handleRequest(req, res, body);
+      });
       return;
     } else if (!sessionId) {
       // No session ID and not an initialize request
@@ -181,7 +188,9 @@ export class HttpStreamTransport extends AbstractTransport {
     }
 
     // Existing session - handle request
-    await transport.handleRequest(req, res, body);
+    await requestContext.run(authData, async () => {
+      await transport.handleRequest(req, res, body);
+    });
   }
 
   private async readRequestBody(req: IncomingMessage): Promise<any> {

src/transports/sse/server.ts

@@ -11,6 +11,8 @@ import { PING_SSE_MESSAGE } from "../utils/ping-message.js";
 import { ProtectedResourceMetadata } from "../../auth/metadata/protected-resource.js";
 import { handleAuthentication } from "../utils/auth-handler.js";
 import { initializeOAuthMetadata } from "../utils/oauth-metadata.js";
+import { requestContext, RequestContextData } from "../../utils/requestContext.js";
+import { AuthResult } from "../../auth/types.js";
 
 interface ExtendedIncomingMessage extends IncomingMessage {
   body?: ClientRequest;
@@ -167,9 +169,12 @@ export class SSEServerTransport extends AbstractTransport {
     }
 
     if (req.method === "POST" && url.pathname === this._config.messageEndpoint) {
+      let authData: RequestContextData = {};
+
       if (this._config.auth?.endpoints?.messages !== false) {
-        const isAuthenticated = await handleAuthentication(req, res, this._config.auth, "message")
-        if (!isAuthenticated) return
+        const authResult = await handleAuthentication(req, res, this._config.auth, "message")
+        if (!authResult) return
+        authData = (authResult as AuthResult).data as RequestContextData || {};
       }
 
       // **Connection Validation (User Requested):**
@@ -183,7 +188,7 @@ export class SSEServerTransport extends AbstractTransport {
           return;
       }
 
-      await this.handlePostMessage(req, res)
+      await this.handlePostMessage(req, res, authData)
       return
     }
 
@@ -250,7 +255,7 @@ export class SSEServerTransport extends AbstractTransport {
     logger.info(`SSE connection established successfully: ${connectionId}`);
   }
 
-  private async handlePostMessage(req: IncomingMessage, res: ServerResponse): Promise<void> {
+  private async handlePostMessage(req: IncomingMessage, res: ServerResponse, authData: RequestContextData = {}): Promise<void> {
     // Check if *any* connection is active, not just the old single _sseResponse
     if (this._connections.size === 0) {
         logger.warn(`Rejecting message: no active SSE connections for server session ${this._sessionId}`);
@@ -301,7 +306,9 @@ export class SSEServerTransport extends AbstractTransport {
         throw new Error("No message handler registered")
       }
 
-      await this._onmessage(rpcMessage)
+      await requestContext.run(authData, async () => {
+        await this._onmessage!(rpcMessage)
+      })
 
       res.writeHead(202).end("Accepted")
 

src/transports/utils/auth-handler.ts

@@ -1,5 +1,5 @@
 import { IncomingMessage, ServerResponse } from 'node:http';
-import { AuthConfig } from '../../auth/types.js';
+import { AuthConfig, AuthResult } from '../../auth/types.js';
 import { APIKeyAuthProvider } from '../../auth/providers/apikey.js';
 import { OAuthAuthProvider } from '../../auth/providers/oauth.js';
 import { DEFAULT_AUTH_ERROR } from '../../auth/types.js';
@@ -14,16 +14,16 @@ import { logger } from '../../core/Logger.js';
  * @param res - HTTP response object
  * @param authConfig - Authentication configuration from transport
  * @param context - Description of the context (e.g., "initialize", "message", "SSE connection")
- * @returns True if authenticated, false if authentication failed (response already sent)
+ * @returns AuthResult if authenticated, null if authentication failed (response already sent)
  */
 export async function handleAuthentication(
   req: IncomingMessage,
   res: ServerResponse,
   authConfig: AuthConfig | undefined,
   context: string
-): Promise<boolean> {
+): Promise<AuthResult | null> {
   if (!authConfig?.provider) {
-    return true;
+    return { data: {} };
   }
 
   const isApiKey = authConfig.provider instanceof APIKeyAuthProvider;
@@ -43,7 +43,7 @@ export async function handleAuthentication(
           type: 'authentication_error',
         })
       );
-      return false;
+      return null;
     }
   }
 
@@ -72,12 +72,16 @@ export async function handleAuthentication(
         type: 'authentication_error',
       })
     );
-    return false;
+    return null;
   }
 
   // Authentication successful
   logger.info(`Authentication successful for ${context}:`);
   logger.info(`- Client IP: ${req.socket.remoteAddress}`);
   logger.info(`- Auth Type: ${authConfig.provider.constructor.name}`);
-  return true;
+  
+  if (typeof authResult === 'boolean') {
+    return { data: {} };
+  }
+  return authResult;
 }

src/utils/requestContext.ts

@@ -0,0 +1,26 @@
+import { AsyncLocalStorage } from 'async_hooks';
+
+export interface RequestContextData {
+  token?: string; // The raw token
+  user?: Record<string, unknown>; // Decoded user data/claims
+  [key: string]: unknown;
+}
+
+export const requestContext = new AsyncLocalStorage<RequestContextData>();
+
+/**
+ * Get the current request context.
+ * Returns undefined if called outside of a request context.
+ */
+export function getRequestContext(): RequestContextData | undefined {
+  return requestContext.getStore();
+}
+
+/**
+ * Run a function within a request context.
+ */
+export function runInRequestContext<T>(context: RequestContextData, fn: () => T): T {
+  return requestContext.run(context, fn);
+}
+
+