From bff65e67466fd395e58ce7c10b6387b458f21496 Mon Sep 17 00:00:00 2001
From: Justin Spahr-Summers <justin@anthropic.com>
Date: Tue, 18 Feb 2025 09:44:25 +0000
Subject: [PATCH] Add client information to all OAuth provider methods

---
 src/server/auth/handlers/authorize.ts |  3 +--
 src/server/auth/handlers/revoke.ts    |  9 ++++++++-
 src/server/auth/handlers/token.ts     | 13 ++++++++++---
 src/server/auth/provider.ts           | 12 +++++-------
 4 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/server/auth/handlers/authorize.ts b/src/server/auth/handlers/authorize.ts
index be3ac49a..c8c005d5 100644
--- a/src/server/auth/handlers/authorize.ts
+++ b/src/server/auth/handlers/authorize.ts
@@ -83,8 +83,7 @@ export function authorizationHandler({ provider }: AuthorizationHandlerOptions):
       }
     }
 
-    await provider.authorize({
-      client,
+    await provider.authorize(client, {
       state: params.state,
       scopes: requestedScopes,
       redirectUri: redirect_uri,
diff --git a/src/server/auth/handlers/revoke.ts b/src/server/auth/handlers/revoke.ts
index 08b1b01e..77c33109 100644
--- a/src/server/auth/handlers/revoke.ts
+++ b/src/server/auth/handlers/revoke.ts
@@ -35,7 +35,14 @@ export function revocationHandler({ provider }: RevocationHandlerOptions): Reque
       return;
     }
 
-    await provider.revokeToken!(revocationRequest);
+    const client = req.client;
+    if (!client) {
+      console.error("Missing client information after authentication");
+      res.status(500).end("Internal Server Error");
+      return;
+    }
+
+    await provider.revokeToken!(client, revocationRequest);
   });
 
   return router;
diff --git a/src/server/auth/handlers/token.ts b/src/server/auth/handlers/token.ts
index 56a3a475..5bc8300e 100644
--- a/src/server/auth/handlers/token.ts
+++ b/src/server/auth/handlers/token.ts
@@ -46,6 +46,13 @@ export function tokenHandler({ provider }: TokenHandlerOptions): RequestHandler
       return;
     }
 
+    const client = req.client;
+    if (!client) {
+      console.error("Missing client information after authentication");
+      res.status(500).end("Internal Server Error");
+      return;
+    }
+
     switch (grant_type) {
       case "authorization_code": {
         let grant;
@@ -59,7 +66,7 @@ export function tokenHandler({ provider }: TokenHandlerOptions): RequestHandler
           return;
         }
 
-        const codeChallenge = await provider.challengeForAuthorizationCode(grant.code);
+        const codeChallenge = await provider.challengeForAuthorizationCode(client, grant.code);
         if (!(await verifyChallenge(grant.code_verifier, codeChallenge))) {
           res.status(400).json({
             error: "invalid_request",
@@ -69,7 +76,7 @@ export function tokenHandler({ provider }: TokenHandlerOptions): RequestHandler
           return;
         }
 
-        const tokens = await provider.exchangeAuthorizationCode(grant.code);
+        const tokens = await provider.exchangeAuthorizationCode(client, grant.code);
         res.status(200).json(tokens);
         break;
       }
@@ -87,7 +94,7 @@ export function tokenHandler({ provider }: TokenHandlerOptions): RequestHandler
         }
 
         const scopes = grant.scope ? grant.scope.split(" ") : undefined;
-        const tokens = await provider.exchangeRefreshToken(grant.refresh_token, scopes);
+        const tokens = await provider.exchangeRefreshToken(client, grant.refresh_token, scopes);
         res.status(200).json(tokens);
         break;
       }
diff --git a/src/server/auth/provider.ts b/src/server/auth/provider.ts
index 5285b398..78da8bc4 100644
--- a/src/server/auth/provider.ts
+++ b/src/server/auth/provider.ts
@@ -3,7 +3,6 @@ import { OAuthRegisteredClientsStore } from "./clients.js";
 import { OAuthClientInformationFull, OAuthTokenRevocationRequest, OAuthTokens } from "../../shared/auth.js";
 
 export type AuthorizationParams = {
-  client: OAuthClientInformationFull;
   state?: string;
   scopes?: string[];
   codeChallenge: string;
@@ -28,28 +27,27 @@ export interface OAuthServerProvider {
    * - In the successful case, the redirect MUST include the `code` and `state` (if present) query parameters.
    * - In the error case, the redirect MUST include the `error` query parameter, and MAY include an optional `error_description` query parameter.
    */
-  authorize(params: AuthorizationParams, res: Response): Promise<void>;
+  authorize(client: OAuthClientInformationFull, params: AuthorizationParams, res: Response): Promise<void>;
 
   /**
    * Returns the `codeChallenge` that was used when the indicated authorization began.
    */
-  // TODO: Add `client` to this and below methods
-  challengeForAuthorizationCode(authorizationCode: string): Promise<string>;
+  challengeForAuthorizationCode(client: OAuthClientInformationFull, authorizationCode: string): Promise<string>;
 
   /**
    * Exchanges an authorization code for an access token.
    */
-  exchangeAuthorizationCode(authorizationCode: string): Promise<OAuthTokens>;
+  exchangeAuthorizationCode(client: OAuthClientInformationFull, authorizationCode: string): Promise<OAuthTokens>;
 
   /**
    * Exchanges a refresh token for an access token.
    */
-  exchangeRefreshToken(refreshToken: string, scopes?: string[]): Promise<OAuthTokens>;
+  exchangeRefreshToken(client: OAuthClientInformationFull, refreshToken: string, scopes?: string[]): Promise<OAuthTokens>;
 
   /**
    * Revokes an access or refresh token. If unimplemented, token revocation is not supported (not recommended).
    * 
    * If the given token is invalid or already revoked, this method should do nothing.
    */
-  revokeToken?(request: OAuthTokenRevocationRequest): Promise<void>;
+  revokeToken?(client: OAuthClientInformationFull, request: OAuthTokenRevocationRequest): Promise<void>;
 }
\ No newline at end of file