Enhance authentication handling in withMcpAuth

Author: allenzhou101

src/next/auth-wrapper.ts

@@ -1,40 +1,96 @@
+import { InvalidTokenError, InsufficientScopeError, ServerError } from '@modelcontextprotocol/sdk/server/auth/errors';
+import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types';
+
+export interface McpAuthOptions {
+  /**
+   * Optional, scopes that the token must have.
+   */
+  requiredScopes?: string[];
+
+  /**
+   * Optional, resource metadata path to include in WWW-Authenticate header.
+   */
+  resourceMetadataPath?: string;
+}
+
 export function withMcpAuth(
   handler: (req: Request) => Promise<Response>,
-  verifyToken: (req: Request, token: string) => Promise<boolean>,
-  oauthResourcePath = "/.well-known/oauth-protected-resource"
+  verifyToken: (req: Request, token: string) => Promise<AuthInfo>,
+  options: McpAuthOptions = {
+    resourceMetadataPath: "/.well-known/oauth-protected-resource"
+  }
 ) {
   return async (req: Request) => {
-    const origin = new URL(req.url).origin;
-
-    if (!req.headers.get("Authorization")) {
-      return new Response(null, {
-        status: 401,
-        headers: {
-          "WWW-Authenticate": `Bearer resource_metadata=${origin}${oauthResourcePath}`,
-        },
-      });
-    }
+    try {
+      if (!req.headers.get("Authorization")) {
+        throw new InvalidTokenError("Missing Authorization header");
+      }
 
-    const authHeader = req.headers.get("Authorization");
-    const token = authHeader?.split(" ")[1];
+      const authHeader = req.headers.get("Authorization");
+      const [type, token] = authHeader?.split(" ") || [];
 
-    if (!token) {
-      throw new Error(
-        `Invalid authorization header value, expected Bearer <token>, received ${authHeader}`
-      );
-    }
+      if (type?.toLowerCase() !== "bearer" || !token) {
+        throw new InvalidTokenError("Invalid Authorization header format, expected 'Bearer TOKEN'");
+      }
 
-    const isAuthenticated = await verifyToken(req, token);
+      const authInfo = await verifyToken(req, token);
 
-    if (!isAuthenticated) {
-      return new Response(JSON.stringify({ error: "Unauthorized" }), {
-        status: 401,
-        headers: {
-          "WWW-Authenticate": `Bearer resource_metadata=${origin}${oauthResourcePath}`,
-        },
-      });
-    }
+      // Check if token has the required scopes (if any)
+      if (options.requiredScopes?.length) {
+        const hasAllScopes = options.requiredScopes.every(scope =>
+          authInfo.scopes.includes(scope)
+        );
+
+        if (!hasAllScopes) {
+          throw new InsufficientScopeError("Insufficient scope");
+        }
+      }
+
+      // Check if the token is expired
+      if (authInfo.expiresAt && authInfo.expiresAt < Date.now() / 1000) {
+        throw new InvalidTokenError("Token has expired");
+      }
 
-    return handler(req);
+      // Set auth info on the request object after successful verification
+      (req as any).auth = authInfo;
+
+      return handler(req);
+    } catch (error) {
+      const origin = new URL(req.url).origin;
+      const resourceMetadataUrl = options.resourceMetadataPath || `${origin}/.well-known/oauth-protected-resource`;
+      if (error instanceof InvalidTokenError) {
+        return new Response(JSON.stringify(error.toResponseObject()), {
+          status: 401,
+          headers: {
+            "WWW-Authenticate": `Bearer error="${error.errorCode}", error_description="${error.message}", resource_metadata="${resourceMetadataUrl}"`,
+            "Content-Type": "application/json"
+          }
+        });
+      } else if (error instanceof InsufficientScopeError) {
+        return new Response(JSON.stringify(error.toResponseObject()), {
+          status: 403,
+          headers: {
+            "WWW-Authenticate": `Bearer error="${error.errorCode}", error_description="${error.message}", resource_metadata="${resourceMetadataUrl}"`,
+            "Content-Type": "application/json"
+          }
+        });
+      } else if (error instanceof ServerError) {
+        return new Response(JSON.stringify(error.toResponseObject()), {
+          status: 500,
+          headers: {
+            "Content-Type": "application/json"
+          }
+        });
+      } else {
+        console.error("Unexpected error authenticating bearer token:", error);
+        const serverError = new ServerError("Internal Server Error");
+        return new Response(JSON.stringify(serverError.toResponseObject()), {
+          status: 500,
+          headers: {
+            "Content-Type": "application/json"
+          }
+        });
+      }
+    }
   };
 }

src/next/mcp-api-handler.ts

@@ -20,6 +20,7 @@ import type {
 } from "../lib/log-helper";
 import { createEvent } from "../lib/log-helper";
 import { EventEmittingResponse } from "../lib/event-emitter.js";
+import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types";
 
 interface SerializedRequest {
   requestId: string;
@@ -310,7 +311,9 @@ export function initializeMcpApiHandler(
           url: req.url,
           headers: Object.fromEntries(req.headers),
           body: bodyContent,
+          auth: (req as any).auth, // Use the auth info that should already be set by withMcpAuth
         });
+        
 
         // Create a response that will emit events
         const wrappedRes = new EventEmittingResponse(
@@ -618,19 +621,21 @@ interface FakeIncomingMessageOptions {
   url?: string;
   headers?: IncomingHttpHeaders;
   body?: BodyType;
+  auth?: AuthInfo;
   socket?: Socket;
 }
 
 // Create a fake IncomingMessage
 function createFakeIncomingMessage(
   options: FakeIncomingMessageOptions = {}
-): IncomingMessage {
+): IncomingMessage & { auth?: AuthInfo } {
   const {
     method = "GET",
     url = "/",
     headers = {},
     body = null,
     socket = new Socket(),
+    auth,
   } = options;
 
   // Create a readable stream that will be used as the base for IncomingMessage
@@ -654,12 +659,15 @@ function createFakeIncomingMessage(
   }
 
   // Create the IncomingMessage instance
-  const req = new IncomingMessage(socket);
+  const req = new IncomingMessage(socket) as IncomingMessage & { auth?: AuthInfo };
 
   // Set the properties
   req.method = method;
   req.url = url;
   req.headers = headers;
+  if (auth) {
+    req.auth = auth;
+  }
 
   // Copy over the stream methods
   req.push = readable.push.bind(readable);