╭────────────────────────────────────────────╮
│ ✻ Welcome to Claude Code research preview! │
│ │
│ /help for help │
│ │
│ cwd: /home/kenton/src/oauth-provider │
╰────────────────────────────────────────────╯
> I want to add a new feature to this OAuth library. Each time a token exchange occurs, the application may wish to update the `props` value. For example, if the application itself is also a
client to some other OAuth API, it may want to perform the equivalent upstream token exchange, and store the result back into the `props`. The callback should actually be able to return two
new `props` values, one to store with the access token and one to replace the props stored in the grant itself.
● I'll help implement this feature. Let's first understand the current codebase structure and then make the necessary changes.
● Search(pattern: "src/*.ts")…
⎿ Found 1 file
● Read(file_path: "src/oauth-provider.ts")…
⎿ File content (21834 tokens) exceeds maximum allowed tokens (20000). Please use offset and limit parameters to read specific portions of the file, or use the GrepTool to search for specific
content.
● Let me search for the token exchange functionality first:
● Search(pattern: "token exchange|tokenExchange", path: "src")…
⎿ Found 0 files
● Search(pattern: "exchangeToken|exchange", path: "src")…
⎿ Found 1 file
● Task(Search in the file /home/kenton/src/oauth-provider/src/oauth-provider.ts for:…)…
⎿ Done (27 tool uses · 25.4k tokens · 1m 13.2s)
● I need to look at the tests to understand how the library is being used:
● Read(file_path: "README.md")…
⎿ # OAuth 2.1 Provider Framework for Cloudflare Workers
This is a TypeScript library that implements the provider side of the OAuth 2.1 protocol with PKCE support. The library is intended to be used on Cloudflare Workers.
... (+220 lines)
● Read(file_path: "__tests__/oauth-provider.test.ts")…
⎿ import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
import { OAuthProvider, ClientInfo, AuthRequest, CompleteAuthorizationOptions } from '../src/oauth-provider';
import { ExecutionContext } from '@cloudflare/workers-types';
... (+1107 lines)
● Let's look at the types and structure of the CompleteAuthorizationOptions to understand the current structure:
● Search(pattern: "interface CompleteAuthorizationOptions", path: "src")…
⎿ Found 1 file
● Search(pattern: "export interface CompleteAuthorizationOptions", path: "src")…
⎿ Found 1 file
● Search(pattern: "handleAuthorizationCodeGrant", path: "src")…
⎿ Found 1 file
● Now let's try to understand the current token exchange functionality:
● Task(Search for the token exchange functionality in /home/kenton/src/oauth-provider/src/oauth-provider.ts - specifically:…)…
⎿ Done (15 tool uses · 25.9k tokens · 46.6s)
● Let's look for any existing callbacks related to token handling:
● Search(pattern: "callback|onToken|handler", path: "src")…
⎿ Found 1 file
● Let's now look at how the props are stored and used:
● Task(Search in the file /home/kenton/src/oauth-provider/src/oauth-provider.ts for how props are stored and used, specifically:…)…
⎿ Done (45 tool uses · 31.1k tokens · 2m 1.2s)
● Now I have a good understanding of how the code works. Let's implement the feature that allows updating the props during token exchange:
1. First, let's modify the type definitions to include a callback function for token exchange
● Search(pattern: "interface OAuthProviderOptions", path: "src")…
⎿ Found 1 file
● Read(file_path: "src/oauth-provider.ts", offset: 100, limit: 400)…
⎿ // Using ExportedHandler from Cloudflare Workers Types for both API and default handlers
// This is Cloudflare's built-in type for Workers handlers with a fetch method
... (+397 lines)
● Search(pattern: "export interface OAuthProviderOptions", path: "src")…
⎿ Found 1 file
● Call(Examine provider options and token exchange implementation)…
⎿ Read(file_path: "src/oauth-provider.ts", offset: 600, limit: 350)…
⎿ Search(pattern: "handleAuthorizationCodeGrant", path: "src")…
⎿ Search(pattern: "handleRefreshTokenGrant", path: "src")…
⎿ Done (3 tool uses · 0.0s)
● Now let's examine the token exchange handlers:
● Read(file_path: "src/oauth-provider.ts", offset: 1000, limit: 400)…
⎿ 'invalid_client',
'Client authentication failed: client has no registered secret',
401
... (+397 lines)
● Let's look at the API handler code to understand how props are used:
● Search(pattern: "handleApiRequest", path: "src")…
⎿ Found 1 file
● Read(file_path: "src/oauth-provider.ts", offset: 1500, limit: 300)…
⎿ jwksUri: validateStringField(clientMetadata.jwks_uri),
contacts: validateStringArray(clientMetadata.contacts),
grantTypes: validateStringArray(clientMetadata.grant_types) || ['authorization_code', 'refresh_token'],
... (+297 lines)
● Based on my analysis of the codebase, I'm ready to implement the new feature. I'll need to:
1. Add a tokenExchangeCallback option to the OAuthProviderOptions interface
2. Update the handleAuthorizationCodeGrant and handleRefreshTokenGrant methods to call this callback
3. Update the types to support the callback returning new props values
Let's begin implementing these changes:
● Search(pattern: "export interface OAuthProviderOptions", path: "src")…
⎿ Found 1 file
● Read(file_path: "src/oauth-provider.ts", offset: 500, limit: 150)…
⎿ cursor?: string;
}
... (+147 lines)
● Now let's find the OAuthProviderOptions:
● Task(Find and retrieve the OAuthProviderOptions interface definition in /home/kenton/src/oauth-provider/src/oauth-provider.ts and return it in full with its line numbers.)…
⎿ Done (4 tool uses · 6.2k tokens · 26.7s)
● Now I'll define the token exchange callback structure and update the interfaces:
● Call(Update interfaces with token exchange callback)…
⎿ Updated src/oauth-provider.ts with 61 additions
31 /**
32 * Configuration options for the OAuth Provider
33 */
34 /**
35 * Result of a token exchange callback function.
36 * Allows updating the props stored in both the access token and the grant.
37 */
38 export interface TokenExchangeCallbackResult {
39 /**
40 * New props to be stored with the access token.
41 * If not provided, the original props will be used.
42 */
43 tokenProps?: any;
44
45 /**
46 * New props to replace the props stored in the grant itself.
47 * If not provided, the original props will be used.
48 */
49 grantProps?: any;
50 }
51
52 /**
53 * Options for token exchange callback functions
54 */
55 export interface TokenExchangeCallbackOptions {
56 /**
57 * The type of grant being processed.
58 * 'authorization_code' for initial code exchange,
59 * 'refresh_token' for refresh token exchange.
60 */
61 grantType: 'authorization_code' | 'refresh_token';
62
63 /**
64 * Client that received this grant
65 */
66 clientId: string;
67
68 /**
69 * User who authorized this grant
70 */
71 userId: string;
72
73 /**
74 * List of scopes that were granted
75 */
76 scope: string[];
77
78 /**
79 * Application-specific properties currently associated with this grant
80 */
81 props: any;
82 }
83
84 export interface OAuthProviderOptions {
85 /**
86 * URL(s) for API routes. Requests with URLs starting with any of these prefixes
...
96 * Defaults to false.
97 */
98 disallowPublicClientRegistration?: boolean;
99
100 /**
101 * Optional callback function that is called during token exchange.
102 * This allows updating the props stored in both the access token and the grant.
103 * For example, if the application itself is also a client to some other OAuth API,
104 * it may want to perform the equivalent upstream token exchange, and store the result in the props.
105 *
106 * The callback can return new props values that will be stored with the token or grant.
107 * If the callback returns nothing or undefined for a props field, the original props will be used.
108 */
109 tokenExchangeCallback?: (options: TokenExchangeCallbackOptions) => Promise<TokenExchangeCallbackResult | void> | TokenExchangeCallbackResult | void;
110 }
111
112 // Using ExportedHandler from Cloudflare Workers Types for both API and default handlers
⎿ Done (1 tool use · 3m 11.1s)
● Now, let's update the handleAuthorizationCodeGrant method to use the callback:
● Update(file_path: src/oauth-provider.ts)…
⎿ User rejected update to src/oauth-provider.ts
1228
1229 // Get the encryption key for props by unwrapping it using the auth code
1230 const encryptionKey = await unwrapKeyWithToken(code, grantData.authCodeWrappedKey!);
1231
1231
1232 // Decrypt the existing props to provide them to the callback
1233 let decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
1234
1235 // Execute token exchange callback if provided
1236 let tokenProps = decryptedProps;
1237 let grantProps = decryptedProps;
1238
1239 if (this.options.tokenExchangeCallback) {
1240 try {
1241 const callbackOptions: TokenExchangeCallbackOptions = {
1242 grantType: 'authorization_code',
1243 clientId: clientInfo.clientId,
1244 userId: userId,
1245 scope: grantData.scope,
1246 props: decryptedProps
1247 };
1248
1249 const callbackResult = await Promise.resolve(this.options.tokenExchangeCallback(callbackOptions));
1250
1251 if (callbackResult) {
1252 // Use the returned props if provided, otherwise keep the original props
1253 if (callbackResult.tokenProps) {
1254 tokenProps = callbackResult.tokenProps;
1255 }
1256
1257 if (callbackResult.grantProps) {
1258 grantProps = callbackResult.grantProps;
1259 }
1260 }
1261 } catch (error) {
1262 // Log the error but continue with the original props
1263 console.error('Error executing token exchange callback:', error);
1264 }
1265 }
1266
1267 // Encrypt the potentially updated grant props
1268 const { encryptedData, key: newEncryptionKey } = await encryptProps(grantProps);
1269
1270 // Update the encrypted props in the grant data
1271 grantData.encryptedProps = encryptedData;
1272
1273 // Wrap the key for both the new access token and refresh token
1274 const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, encryptionKey);
1275 const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, encryptionKey);
1274 const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, newEncryptionKey);
1275 const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, newEncryptionKey);
1276
1277 // Update the grant:
1278 // - Remove the auth code hash (it's single-use)
> Let's only decrypt the original props at all if we have a token exchange callback, otherwise there's no need. Also please don't add a try/catch, if the callback throws an exception we should
just let it propagate normally.
● You're right, let's make those changes. We'll only decrypt the props if we have a callback, and we'll let any exceptions from the callback propagate normally.
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 41 additions and 3 removals
1228
1229 // Get the encryption key for props by unwrapping it using the auth code
1230 const encryptionKey = await unwrapKeyWithToken(code, grantData.authCodeWrappedKey!);
1231
1231
1232 let newEncryptionKey = encryptionKey;
1233
1234 // Process token exchange callback if provided
1235 if (this.options.tokenExchangeCallback) {
1236 // Decrypt the existing props to provide them to the callback
1237 const decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
1238
1239 // Default to using the original props
1240 let tokenProps = decryptedProps;
1241 let grantProps = decryptedProps;
1242
1243 const callbackOptions: TokenExchangeCallbackOptions = {
1244 grantType: 'authorization_code',
1245 clientId: clientInfo.clientId,
1246 userId: userId,
1247 scope: grantData.scope,
1248 props: decryptedProps
1249 };
1250
1251 const callbackResult = await Promise.resolve(this.options.tokenExchangeCallback(callbackOptions));
1252
1253 if (callbackResult) {
1254 // Use the returned props if provided, otherwise keep the original props
1255 if (callbackResult.tokenProps) {
1256 tokenProps = callbackResult.tokenProps;
1257 }
1258
1259 if (callbackResult.grantProps) {
1260 grantProps = callbackResult.grantProps;
1261 }
1262 }
1263
1264 // Re-encrypt the potentially updated grant props
1265 const result = await encryptProps(grantProps);
1266 grantData.encryptedProps = result.encryptedData;
1267 newEncryptionKey = result.key;
1268 }
1269
1270 // Wrap the key for both the new access token and refresh token
1271 const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, encryptionKey);
1272 const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, encryptionKey);
1271 const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, newEncryptionKey);
1272 const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, newEncryptionKey);
1273
1274 // Update the grant:
1275 // - Remove the auth code hash (it's single-use)
● Now we need to update how we handle the token props in the handleAuthorizationCodeGrant method, after the encryption key management:
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 7 additions
1303 encryptedProps: grantData.encryptedProps
1304 }
1305 };
1306
1307 // If a token exchange callback was executed and returned token-specific props, store them separately
1308 if (this.options.tokenExchangeCallback && tokenProps !== undefined) {
1309 // Re-encrypt the token-specific props if they were provided
1310 const { encryptedData } = await encryptProps(tokenProps);
1311 accessTokenData.grant.encryptedProps = encryptedData;
1312 }
1313
1314 // Save access token with TTL
1315 await env.OAUTH_KV.put(
● There's an issue with the code - we need to define tokenProps outside of the if block so it's in scope. Let's fix that:
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 2 additions and 1 removal
1230 const encryptionKey = await unwrapKeyWithToken(code, grantData.authCodeWrappedKey!);
1231
1232 let newEncryptionKey = encryptionKey;
1233 let tokenProps: any = undefined;
1234
1235 // Process token exchange callback if provided
1236 if (this.options.tokenExchangeCallback) {
...
1237 const decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
1238
1239 // Default to using the original props
1240 let tokenProps = decryptedProps;
1240 tokenProps = decryptedProps;
1241 let grantProps = decryptedProps;
1242
1243 const callbackOptions: TokenExchangeCallbackOptions = {
● Now let's implement the same changes in the handleRefreshTokenGrant method:
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 41 additions and 2 removals
1420
1421 // Unwrap the encryption key using the refresh token
1422 const encryptionKey = await unwrapKeyWithToken(refreshToken, wrappedKeyToUse);
1423
1424 let newEncryptionKey = encryptionKey;
1425 let tokenProps: any = undefined;
1426
1427 // Process token exchange callback if provided
1428 if (this.options.tokenExchangeCallback) {
1429 // Decrypt the existing props to provide them to the callback
1430 const decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
1431
1432 // Default to using the original props
1433 tokenProps = decryptedProps;
1434 let grantProps = decryptedProps;
1435
1436 const callbackOptions: TokenExchangeCallbackOptions = {
1437 grantType: 'refresh_token',
1438 clientId: clientInfo.clientId,
1439 userId: userId,
1440 scope: grantData.scope,
1441 props: decryptedProps
1442 };
1443
1444 const callbackResult = await Promise.resolve(this.options.tokenExchangeCallback(callbackOptions));
1445
1446 if (callbackResult) {
1447 // Use the returned props if provided, otherwise keep the original props
1448 if (callbackResult.tokenProps) {
1449 tokenProps = callbackResult.tokenProps;
1450 }
1451
1452 if (callbackResult.grantProps) {
1453 grantProps = callbackResult.grantProps;
1454 }
1455 }
1456
1457 // Re-encrypt the potentially updated grant props
1458 const result = await encryptProps(grantProps);
1459 grantData.encryptedProps = result.encryptedData;
1460 newEncryptionKey = result.key;
1461 }
1462
1463 // Wrap the key for both the new access token and refresh token
1464 const accessTokenWrappedKey = await wrapKeyWithToken(newAccessToken, encryptionKey);
1465 const newRefreshTokenWrappedKey = await wrapKeyWithToken(newRefreshToken, encryptionKey);
1464 const accessTokenWrappedKey = await wrapKeyWithToken(newAccessToken, newEncryptionKey);
1465 const newRefreshTokenWrappedKey = await wrapKeyWithToken(newRefreshToken, newEncryptionKey);
1466
1467 // Update the grant with the token rotation information
1468
● Finally, let's update the access token storage part in the handleRefreshTokenGrant method to use the token-specific props if provided:
● Update(file_path: src/oauth-provider.ts)…
⎿ User rejected update to src/oauth-provider.ts
1500 encryptedProps: grantData.encryptedProps
1501 }
1502 };
1503
1504 // If a token exchange callback was executed and returned token-specific props, store them separately
1505 if (this.options.tokenExchangeCallback && tokenProps !== undefined) {
1506 // Re-encrypt the token-specific props if they were provided
1507 const { encryptedData } = await encryptProps(tokenProps);
1508 accessTokenData.grant.encryptedProps = encryptedData;
1509 }
1510
1511 // Save access token with TTL
1512 await env.OAUTH_KV.put(
> There's an issue here -- encryptProps() returns a new encryption key, which you will need to wrap and store into the access token record. This means that the access token record and grant
record have different keys, but that's OK, they just have to match the corresponding props.
● You're right, I need to handle the encryption key properly for both token props and grant props. Let's update the implementation:
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 7 additions and 1 removal
1308 // If a token exchange callback was executed and returned token-specific props, store them separately
1309 if (this.options.tokenExchangeCallback && tokenProps !== undefined) {
1310 // Re-encrypt the token-specific props if they were provided
1311 const { encryptedData } = await encryptProps(tokenProps);
1311 const { encryptedData, key: tokenEncryptionKey } = await encryptProps(tokenProps);
1312
1313 // Re-wrap this separate encryption key with the access token
1314 const tokenWrappedKey = await wrapKeyWithToken(accessToken, tokenEncryptionKey);
1315
1316 // Update the access token data with token-specific props and their key
1317 accessTokenData.grant.encryptedProps = encryptedData;
1318 accessTokenData.wrappedEncryptionKey = tokenWrappedKey;
1319 }
1320
1321 // Save access token with TTL
● Now let's do the same update for the refresh token flow:
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 13 additions
1506 encryptedProps: grantData.encryptedProps
1507 }
1508 };
1509
1510 // If a token exchange callback was executed and returned token-specific props, store them separately
1511 if (this.options.tokenExchangeCallback && tokenProps !== undefined) {
1512 // Re-encrypt the token-specific props if they were provided
1513 const { encryptedData, key: tokenEncryptionKey } = await encryptProps(tokenProps);
1514
1515 // Re-wrap this separate encryption key with the access token
1516 const tokenWrappedKey = await wrapKeyWithToken(newAccessToken, tokenEncryptionKey);
1517
1518 // Update the access token data with token-specific props and their key
1519 accessTokenData.grant.encryptedProps = encryptedData;
1520 accessTokenData.wrappedEncryptionKey = tokenWrappedKey;
1521 }
1522
1523 // Save access token with TTL
1524 await env.OAUTH_KV.put(
● The implementation is now complete. Let's summarize the changes we've made:
1. Added interfaces for token exchange callback:
- TokenExchangeCallbackOptions - The data passed to the callback
- TokenExchangeCallbackResult - The result returned by the callback, with optional tokenProps and grantProps
2. Added a tokenExchangeCallback option to the OAuthProviderOptions interface
3. Updated handleAuthorizationCodeGrant and handleRefreshTokenGrant methods to:
- Call the callback if provided
- Update the grant's props if new grant props were returned
- Store separate token-specific props if returned by the callback
- Properly handle encryption keys for both grant and token props
These changes allow applications to update the props stored with the grant and/or token during token exchange, which is useful for scenarios like upstream token exchange or refreshing
additional data during token issuance.
> Related to that last issue I spotted, it looks like we're now potentially wrapping the access token key twice, once redundantly, when the access token has unique props. Perhaps the initial
block handling tokenExchangeCallback should actually handle encrypting both sets of props, and instead of assigning `tokenProps` (declared outside the block), it assigns
`encryptedAccessTokenProps` and `newAccessTokenKey`?
● You're absolutely right. We're handling this inefficiently by re-wrapping the keys. Let's refactor to avoid redundant operations and make the code clearer:
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 24 additions and 11 removals
1229 // Get the encryption key for props by unwrapping it using the auth code
1230 const encryptionKey = await unwrapKeyWithToken(code, grantData.authCodeWrappedKey!);
1231
1232 let newEncryptionKey = encryptionKey;
1233 let tokenProps: any = undefined;
1232 // Default to using the same encryption key and props for both grant and access token
1233 let grantEncryptionKey = encryptionKey;
1234 let accessTokenEncryptionKey = encryptionKey;
1235 let encryptedAccessTokenProps = grantData.encryptedProps;
1236
1237 // Process token exchange callback if provided
1238 if (this.options.tokenExchangeCallback) {
1239 // Decrypt the existing props to provide them to the callback
1240 const decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
1241
1242 // Default to using the original props
1243 tokenProps = decryptedProps;
1242 // Default to using the original props for both grant and token
1243 let grantProps = decryptedProps;
1244 let accessTokenProps = decryptedProps;
1245
1246 const callbackOptions: TokenExchangeCallbackOptions = {
1247 grantType: 'authorization_code',
...
1254 if (callbackResult) {
1255 // Use the returned props if provided, otherwise keep the original props
1256 if (callbackResult.tokenProps) {
1257 tokenProps = callbackResult.tokenProps;
1257 accessTokenProps = callbackResult.tokenProps;
1258 }
1259
1260 if (callbackResult.grantProps) {
...
1263 }
1264
1265 // Re-encrypt the potentially updated grant props
1266 const result = await encryptProps(grantProps);
1267 grantData.encryptedProps = result.encryptedData;
1268 newEncryptionKey = result.key;
1266 const grantResult = await encryptProps(grantProps);
1267 grantData.encryptedProps = grantResult.encryptedData;
1268 grantEncryptionKey = grantResult.key;
1269
1270 // Re-encrypt the access token props if they're different from grant props
1271 if (accessTokenProps !== grantProps) {
1272 const tokenResult = await encryptProps(accessTokenProps);
1273 encryptedAccessTokenProps = tokenResult.encryptedData;
1274 accessTokenEncryptionKey = tokenResult.key;
1275 } else {
1276 // If they're the same, use the grant's encrypted data and key
1277 encryptedAccessTokenProps = grantData.encryptedProps;
1278 accessTokenEncryptionKey = grantEncryptionKey;
1279 }
1280 }
1281
1282 // Wrap the key for both the new access token and refresh token
1283 const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, newEncryptionKey);
1284 const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, newEncryptionKey);
1282 // Wrap the keys for the new tokens
1283 const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, accessTokenEncryptionKey);
1284 const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, grantEncryptionKey);
1285
1286 // Update the grant:
1287 // - Remove the auth code hash (it's single-use)
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 1 addition and 14 removals
1314 grant: {
1315 clientId: grantData.clientId,
1316 scope: grantData.scope,
1317 encryptedProps: grantData.encryptedProps
1317 encryptedProps: encryptedAccessTokenProps
1318 }
1319 };
1320
1321 // If a token exchange callback was executed and returned token-specific props, store them separately
1322 if (this.options.tokenExchangeCallback && tokenProps !== undefined) {
1323 // Re-encrypt the token-specific props if they were provided
1324 const { encryptedData, key: tokenEncryptionKey } = await encryptProps(tokenProps);
1325
1326 // Re-wrap this separate encryption key with the access token
1327 const tokenWrappedKey = await wrapKeyWithToken(accessToken, tokenEncryptionKey);
1328
1329 // Update the access token data with token-specific props and their key
1330 accessTokenData.grant.encryptedProps = encryptedData;
1331 accessTokenData.wrappedEncryptionKey = tokenWrappedKey;
1332 }
1320
1321 // Save access token with TTL
1322 await env.OAUTH_KV.put(
● Now let's update the refresh token flow similarly:
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 23 additions and 10 removals
1427 // Unwrap the encryption key using the refresh token
1428 const encryptionKey = await unwrapKeyWithToken(refreshToken, wrappedKeyToUse);
1429
1430 let newEncryptionKey = encryptionKey;
1431 let tokenProps: any = undefined;
1430 // Default to using the same encryption key and props for both grant and access token
1431 let grantEncryptionKey = encryptionKey;
1432 let accessTokenEncryptionKey = encryptionKey;
1433 let encryptedAccessTokenProps = grantData.encryptedProps;
1434
1435 // Process token exchange callback if provided
1436 if (this.options.tokenExchangeCallback) {
1437 // Decrypt the existing props to provide them to the callback
1438 const decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
1439
1440 // Default to using the original props
1441 tokenProps = decryptedProps;
1440 // Default to using the original props for both grant and token
1441 let grantProps = decryptedProps;
1442 let accessTokenProps = decryptedProps;
1443
1444 const callbackOptions: TokenExchangeCallbackOptions = {
1445 grantType: 'refresh_token',
...
1452 if (callbackResult) {
1453 // Use the returned props if provided, otherwise keep the original props
1454 if (callbackResult.tokenProps) {
1455 tokenProps = callbackResult.tokenProps;
1455 accessTokenProps = callbackResult.tokenProps;
1456 }
1457
1458 if (callbackResult.grantProps) {
...
1461 }
1462
1463 // Re-encrypt the potentially updated grant props
1464 const result = await encryptProps(grantProps);
1465 grantData.encryptedProps = result.encryptedData;
1466 newEncryptionKey = result.key;
1464 const grantResult = await encryptProps(grantProps);
1465 grantData.encryptedProps = grantResult.encryptedData;
1466 grantEncryptionKey = grantResult.key;
1467
1468 // Re-encrypt the access token props if they're different from grant props
1469 if (accessTokenProps !== grantProps) {
1470 const tokenResult = await encryptProps(accessTokenProps);
1471 encryptedAccessTokenProps = tokenResult.encryptedData;
1472 accessTokenEncryptionKey = tokenResult.key;
1473 } else {
1474 // If they're the same, use the grant's encrypted data and key
1475 encryptedAccessTokenProps = grantData.encryptedProps;
1476 accessTokenEncryptionKey = grantEncryptionKey;
1477 }
1478 }
1479
1480 // Wrap the key for both the new access token and refresh token
1481 const accessTokenWrappedKey = await wrapKeyWithToken(newAccessToken, newEncryptionKey);
1482 const newRefreshTokenWrappedKey = await wrapKeyWithToken(newRefreshToken, newEncryptionKey);
1481 const accessTokenWrappedKey = await wrapKeyWithToken(newAccessToken, accessTokenEncryptionKey);
1482 const newRefreshTokenWrappedKey = await wrapKeyWithToken(newRefreshToken, grantEncryptionKey);
1483
1484 // Update the grant with the token rotation information
1485
● Update(file_path: src/oauth-provider.ts)…
⎿ Updated src/oauth-provider.ts with 1 addition and 14 removals
1516 grant: {
1517 clientId: grantData.clientId,
1518 scope: grantData.scope,
1519 encryptedProps: grantData.encryptedProps
1519 encryptedProps: encryptedAccessTokenProps
1520 }
1521 };
1522
1523 // If a token exchange callback was executed and returned token-specific props, store them separately
1524 if (this.options.tokenExchangeCallback && tokenProps !== undefined) {
1525 // Re-encrypt the token-specific props if they were provided
1526 const { encryptedData, key: tokenEncryptionKey } = await encryptProps(tokenProps);
1527
1528 // Re-wrap this separate encryption key with the access token
1529 const tokenWrappedKey = await wrapKeyWithToken(newAccessToken, tokenEncryptionKey);
1530
1531 // Update the access token data with token-specific props and their key
1532 accessTokenData.grant.encryptedProps = encryptedData;
1533 accessTokenData.wrappedEncryptionKey = tokenWrappedKey;
1534 }
1522
1523 // Save access token with TTL
1524 await env.OAUTH_KV.put(
● The implementation is now complete and more efficient. We've made the following improvements:
1. Fixed the issue where we were potentially wrapping keys redundantly
2. Separated the encryption of grant props and token props more clearly
3. Made the code more efficient by only encrypting token props separately when they differ from grant props
4. Ensured proper handling of all encryption keys and encrypted data
This implementation correctly handles the token exchange callback feature, allowing applications to update either the grant props, the token props, or both during token exchanges. The code
is now more efficient and clearer in its intent.
> Let's add some tests for this feature.
● Let's add tests to verify the token exchange callback functionality. We'll create tests for both authorization code flow and refresh token flow to ensure the callback works properly in both
cases.
● Read(file_path: "__tests__/oauth-provider.test.ts", offset: 600, limit: 150)…
⎿ const authResponse = await oauthProvider.fetch(authRequest, mockEnv, mockCtx);
const location = authResponse.headers.get('Location')!;
const url = new URL(location);
... (+147 lines)
● Let's add a new test suite for testing the token exchange callback functionality:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 286 additions
1019 });
1020 });
1021
1022 describe('Token Exchange Callback', () => {
1023 // Test with provider that has token exchange callback
1024 let oauthProviderWithCallback: OAuthProvider;
1025 let callbackInvocations: any[] = [];
1026 let mockEnv: ReturnType<typeof createMockEnv>;
1027 let mockCtx: MockExecutionContext;
1028
1029 // Helper function to create a test OAuth provider with a token exchange callback
1030 function createProviderWithCallback() {
1031 callbackInvocations = [];
1032
1033 const tokenExchangeCallback = async (options: any) => {
1034 // Record that the callback was called and with what arguments
1035 callbackInvocations.push({...options});
1036
1037 // Return different props based on the grant type
1038 if (options.grantType === 'authorization_code') {
1039 return {
1040 tokenProps: {
1041 ...options.props,
1042 tokenSpecific: true,
1043 tokenUpdatedAt: 'auth_code_flow'
1044 },
1045 grantProps: {
1046 ...options.props,
1047 grantUpdated: true
1048 }
1049 };
1050 } else if (options.grantType === 'refresh_token') {
1051 return {
1052 tokenProps: {
1053 ...options.props,
1054 tokenSpecific: true,
1055 tokenUpdatedAt: 'refresh_token_flow'
1056 },
1057 grantProps: {
1058 ...options.props,
1059 grantUpdated: true,
1060 refreshCount: (options.props.refreshCount || 0) + 1
1061 }
1062 };
1063 }
1064 };
1065
1066 return new OAuthProvider({
1067 apiRoute: ['/api/', 'https://api.example.com/'],
1068 apiHandler: TestApiHandler,
1069 defaultHandler: testDefaultHandler,
1070 authorizeEndpoint: '/authorize',
1071 tokenEndpoint: '/oauth/token',
1072 clientRegistrationEndpoint: '/oauth/register',
1073 scopesSupported: ['read', 'write', 'profile'],
1074 accessTokenTTL: 3600,
1075 allowImplicitFlow: true,
1076 tokenExchangeCallback
1077 });
1078 }
1079
1080 let clientId: string;
1081 let clientSecret: string;
1082 let redirectUri: string;
1083
1084 // Helper to create a test client
1085 async function createTestClient() {
1086 const clientData = {
1087 redirect_uris: ['https://client.example.com/callback'],
1088 client_name: 'Test Client',
1089 token_endpoint_auth_method: 'client_secret_basic'
1090 };
1091
1092 const request = createMockRequest(
1093 'https://example.com/oauth/register',
1094 'POST',
1095 { 'Content-Type': 'application/json' },
1096 JSON.stringify(clientData)
1097 );
1098
1099 const response = await oauthProviderWithCallback.fetch(request, mockEnv, mockCtx);
1100 const client = await response.json();
1101
1102 clientId = client.client_id;
1103 clientSecret = client.client_secret;
1104 redirectUri = 'https://client.example.com/callback';
1105 }
1106
1107 beforeEach(async () => {
1108 // Reset mocks before each test
1109 vi.resetAllMocks();
1110
1111 // Create fresh instances for each test
1112 mockEnv = createMockEnv();
1113 mockCtx = new MockExecutionContext();
1114
1115 // Create OAuth provider with test configuration and callback
1116 oauthProviderWithCallback = createProviderWithCallback();
1117
1118 // Create a test client
1119 await createTestClient();
1120 });
1121
1122 afterEach(() => {
1123 // Clean up KV storage after each test
1124 mockEnv.OAUTH_KV.clear();
1125 });
1126
1127 it('should call the callback during authorization code flow', async () => {
1128 // First get an auth code
1129 const authRequest = createMockRequest(
1130 `https://example.com/authorize?response_type=code&client_id=${clientId}` +
1131 `&redirect_uri=${encodeURIComponent(redirectUri)}` +
1132 `&scope=read%20write&state=xyz123`
1133 );
1134
1135 const authResponse = await oauthProviderWithCallback.fetch(authRequest, mockEnv, mockCtx);
1136 const location = authResponse.headers.get('Location')!;
1137 const code = new URL(location).searchParams.get('code')!;
1138
1139 // Reset callback invocations tracking before token exchange
1140 callbackInvocations = [];
1141
1142 // Exchange code for tokens
1143 const params = new URLSearchParams();
1144 params.append('grant_type', 'authorization_code');
1145 params.append('code', code);
1146 params.append('redirect_uri', redirectUri);
1147 params.append('client_id', clientId);
1148 params.append('client_secret', clientSecret);
1149
1150 const tokenRequest = createMockRequest(
1151 'https://example.com/oauth/token',
1152 'POST',
1153 { 'Content-Type': 'application/x-www-form-urlencoded' },
1154 params.toString()
1155 );
1156
1157 const tokenResponse = await oauthProviderWithCallback.fetch(tokenRequest, mockEnv, mockCtx);
1158
1159 // Check that the token exchange was successful
1160 expect(tokenResponse.status).toBe(200);
1161 const tokens = await tokenResponse.json();
1162 expect(tokens.access_token).toBeDefined();
1163
1164 // Check that the callback was called once
1165 expect(callbackInvocations.length).toBe(1);
1166
1167 // Check that callback was called with correct arguments
1168 const callbackArgs = callbackInvocations[0];
1169 expect(callbackArgs.grantType).toBe('authorization_code');
1170 expect(callbackArgs.clientId).toBe(clientId);
1171 expect(callbackArgs.props).toEqual({ userId: "test-user-123", username: "TestUser" });
1172
1173 // Use the token to access API
1174 const apiRequest = createMockRequest(
1175 'https://example.com/api/test',
1176 'GET',
1177 { 'Authorization': `Bearer ${tokens.access_token}` }
1178 );
1179
1180 const apiResponse = await oauthProviderWithCallback.fetch(apiRequest, mockEnv, mockCtx);
1181 expect(apiResponse.status).toBe(200);
1182
1183 // Check that the API received the token-specific props from the callback
1184 const apiData = await apiResponse.json();
1185 expect(apiData.user).toEqual({
1186 userId: "test-user-123",
1187 username: "TestUser",
1188 tokenSpecific: true,
1189 tokenUpdatedAt: 'auth_code_flow'
1190 });
1191 });
1192
1193 it('should call the callback during refresh token flow', async () => {
1194 // First get an auth code and exchange it for tokens
1195 const authRequest = createMockRequest(
1196 `https://example.com/authorize?response_type=code&client_id=${clientId}` +
1197 `&redirect_uri=${encodeURIComponent(redirectUri)}` +
1198 `&scope=read%20write&state=xyz123`
1199 );
1200
1201 const authResponse = await oauthProviderWithCallback.fetch(authRequest, mockEnv, mockCtx);
1202 const location = authResponse.headers.get('Location')!;
1203 const code = new URL(location).searchParams.get('code')!;
1204
1205 // Exchange code for tokens
1206 const codeParams = new URLSearchParams();
1207 codeParams.append('grant_type', 'authorization_code');
1208 codeParams.append('code', code);
1209 codeParams.append('redirect_uri', redirectUri);
1210 codeParams.append('client_id', clientId);
1211 codeParams.append('client_secret', clientSecret);
1212
1213 const tokenRequest = createMockRequest(
1214 'https://example.com/oauth/token',
1215 'POST',
1216 { 'Content-Type': 'application/x-www-form-urlencoded' },
1217 codeParams.toString()
1218 );
1219
1220 const tokenResponse = await oauthProviderWithCallback.fetch(tokenRequest, mockEnv, mockCtx);
1221 const tokens = await tokenResponse.json();
1222
1223 // Reset the callback invocations tracking before refresh
1224 callbackInvocations = [];
1225
1226 // Now use the refresh token
1227 const refreshParams = new URLSearchParams();
1228 refreshParams.append('grant_type', 'refresh_token');
1229 refreshParams.append('refresh_token', tokens.refresh_token);
1230 refreshParams.append('client_id', clientId);
1231 refreshParams.append('client_secret', clientSecret);
1232
1233 const refreshRequest = createMockRequest(
1234 'https://example.com/oauth/token',
1235 'POST',
1236 { 'Content-Type': 'application/x-www-form-urlencoded' },
1237 refreshParams.toString()
1238 );
1239
1240 const refreshResponse = await oauthProviderWithCallback.fetch(refreshRequest, mockEnv, mockCtx);
1241
1242 // Check that the refresh was successful
1243 expect(refreshResponse.status).toBe(200);
1244 const newTokens = await refreshResponse.json();
1245 expect(newTokens.access_token).toBeDefined();
1246
1247 // Check that the callback was called once
1248 expect(callbackInvocations.length).toBe(1);
1249
1250 // Check that callback was called with correct arguments
1251 const callbackArgs = callbackInvocations[0];
1252 expect(callbackArgs.grantType).toBe('refresh_token');
1253 expect(callbackArgs.clientId).toBe(clientId);
1254
1255 // The props are from the updated grant during auth code flow
1256 expect(callbackArgs.props).toEqual({
1257 userId: "test-user-123",
1258 username: "TestUser",
1259 grantUpdated: true
1260 });
1261
1262 // Use the new token to access API
1263 const apiRequest = createMockRequest(
1264 'https://example.com/api/test',
1265 'GET',
1266 { 'Authorization': `Bearer ${newTokens.access_token}` }
1267 );
1268
1269 const apiResponse = await oauthProviderWithCallback.fetch(apiRequest, mockEnv, mockCtx);
1270 expect(apiResponse.status).toBe(200);
1271
1272 // Check that the API received the token-specific props from the refresh callback
1273 const apiData = await apiResponse.json();
1274 expect(apiData.user).toEqual({
1275 userId: "test-user-123",
1276 username: "TestUser",
1277 grantUpdated: true,
1278 tokenSpecific: true,
1279 tokenUpdatedAt: 'refresh_token_flow'
1280 });
1281
1282 // Do a second refresh to verify that grant props are properly updated
1283 const refresh2Params = new URLSearchParams();
1284 refresh2Params.append('grant_type', 'refresh_token');
1285 refresh2Params.append('refresh_token', newTokens.refresh_token);
1286 refresh2Params.append('client_id', clientId);
1287 refresh2Params.append('client_secret', clientSecret);
1288
1289 // Reset the callback invocations before second refresh
1290 callbackInvocations = [];
1291
1292 const refresh2Request = createMockRequest(
1293 'https://example.com/oauth/token',
1294 'POST',
1295 { 'Content-Type': 'application/x-www-form-urlencoded' },
1296 refresh2Params.toString()
1297 );
1298
1299 const refresh2Response = await oauthProviderWithCallback.fetch(refresh2Request, mockEnv, mockCtx);
1300 const newerTokens = await refresh2Response.json();
1301
1302 // Check that the refresh count was incremented in the grant props
1303 expect(callbackInvocations.length).toBe(1);
1304 expect(callbackInvocations[0].props.refreshCount).toBe(1);
1305 });
1306 });
1307
1308 describe('OAuthHelpers', () => {
1309 it('should allow listing and revoking grants', async () => {
1310 // Create a client
● Let's also add a test to verify that the callback works when it only returns one of the two props types:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 116 additions
1303 expect(callbackInvocations.length).toBe(1);
1304 expect(callbackInvocations[0].props.refreshCount).toBe(1);
1305 });
1306
1307 it('should handle callback that returns only tokenProps or only grantProps', async () => {
1308 // Create a provider with a callback that returns only tokenProps for auth code
1309 // and only grantProps for refresh token
1310 const tokenPropsOnlyCallback = async (options: any) => {
1311 if (options.grantType === 'authorization_code') {
1312 return {
1313 tokenProps: { ...options.props, tokenOnly: true }
1314 };
1315 } else if (options.grantType === 'refresh_token') {
1316 return {
1317 grantProps: { ...options.props, grantOnly: true }
1318 };
1319 }
1320 };
1321
1322 const specialProvider = new OAuthProvider({
1323 apiRoute: ['/api/'],
1324 apiHandler: TestApiHandler,
1325 defaultHandler: testDefaultHandler,
1326 authorizeEndpoint: '/authorize',
1327 tokenEndpoint: '/oauth/token',
1328 scopesSupported: ['read', 'write'],
1329 tokenExchangeCallback: tokenPropsOnlyCallback
1330 });
1331
1332 // Create a client
1333 const clientData = {
1334 redirect_uris: ['https://client.example.com/callback'],
1335 client_name: 'Token Props Only Test',
1336 token_endpoint_auth_method: 'client_secret_basic'
1337 };
1338
1339 const registerRequest = createMockRequest(
1340 'https://example.com/oauth/register',
1341 'POST',
1342 { 'Content-Type': 'application/json' },
1343 JSON.stringify(clientData)
1344 );
1345
1346 const registerResponse = await specialProvider.fetch(registerRequest, mockEnv, mockCtx);
1347 const client = await registerResponse.json();
1348 const testClientId = client.client_id;
1349 const testClientSecret = client.client_secret;
1350 const testRedirectUri = 'https://client.example.com/callback';
1351
1352 // Get an auth code
1353 const authRequest = createMockRequest(
1354 `https://example.com/authorize?response_type=code&client_id=${testClientId}` +
1355 `&redirect_uri=${encodeURIComponent(testRedirectUri)}` +
1356 `&scope=read%20write&state=xyz123`
1357 );
1358
1359 const authResponse = await specialProvider.fetch(authRequest, mockEnv, mockCtx);
1360 const code = new URL(authResponse.headers.get('Location')!).searchParams.get('code')!;
1361
1362 // Exchange code for tokens
1363 const params = new URLSearchParams();
1364 params.append('grant_type', 'authorization_code');
1365 params.append('code', code);
1366 params.append('redirect_uri', testRedirectUri);
1367 params.append('client_id', testClientId);
1368 params.append('client_secret', testClientSecret);
1369
1370 const tokenRequest = createMockRequest(
1371 'https://example.com/oauth/token',
1372 'POST',
1373 { 'Content-Type': 'application/x-www-form-urlencoded' },
1374 params.toString()
1375 );
1376
1377 const tokenResponse = await specialProvider.fetch(tokenRequest, mockEnv, mockCtx);
1378 const tokens = await tokenResponse.json();
1379
1380 // Verify the token has the tokenOnly property when used for API access
1381 const apiRequest = createMockRequest(
1382 'https://example.com/api/test',
1383 'GET',
1384 { 'Authorization': `Bearer ${tokens.access_token}` }
1385 );
1386
1387 const apiResponse = await specialProvider.fetch(apiRequest, mockEnv, mockCtx);
1388 const apiData = await apiResponse.json();
1389 expect(apiData.user.tokenOnly).toBe(true);
1390
1391 // Now do a refresh token exchange
1392 const refreshParams = new URLSearchParams();
1393 refreshParams.append('grant_type', 'refresh_token');
1394 refreshParams.append('refresh_token', tokens.refresh_token);
1395 refreshParams.append('client_id', testClientId);
1396 refreshParams.append('client_secret', testClientSecret);
1397
1398 const refreshRequest = createMockRequest(
1399 'https://example.com/oauth/token',
1400 'POST',
1401 { 'Content-Type': 'application/x-www-form-urlencoded' },
1402 refreshParams.toString()
1403 );
1404
1405 const refreshResponse = await specialProvider.fetch(refreshRequest, mockEnv, mockCtx);
1406 const newTokens = await refreshResponse.json();
1407
1408 // Use the new token to access API
1409 const api2Request = createMockRequest(
1410 'https://example.com/api/test',
1411 'GET',
1412 { 'Authorization': `Bearer ${newTokens.access_token}` }
1413 );
1414
1415 const api2Response = await specialProvider.fetch(api2Request, mockEnv, mockCtx);
1416 const api2Data = await api2Response.json();
1417
1418 // The token should have the grantOnly property from the grant, not tokenOnly
1419 expect(api2Data.user.grantOnly).toBe(true);
1420 expect(api2Data.user.tokenOnly).toBeUndefined();
1421 });
1422 });
1423
1424 describe('OAuthHelpers', () => {
● Let's add one more test to cover the case where the callback returns undefined:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 79 additions
1419 expect(api2Data.user.grantOnly).toBe(true);
1420 expect(api2Data.user.tokenOnly).toBeUndefined();
1421 });
1422
1423 it('should handle callback that returns undefined (keeping original props)', async () => {
1424 // Create a provider with a callback that returns undefined
1425 const noopCallback = async (options: any) => {
1426 // Don't return anything, which should keep the original props
1427 return undefined;
1428 };
1429
1430 const noopProvider = new OAuthProvider({
1431 apiRoute: ['/api/'],
1432 apiHandler: TestApiHandler,
1433 defaultHandler: testDefaultHandler,
1434 authorizeEndpoint: '/authorize',
1435 tokenEndpoint: '/oauth/token',
1436 scopesSupported: ['read', 'write'],
1437 tokenExchangeCallback: noopCallback
1438 });
1439
1440 // Create a client
1441 const clientData = {
1442 redirect_uris: ['https://client.example.com/callback'],
1443 client_name: 'Noop Callback Test',
1444 token_endpoint_auth_method: 'client_secret_basic'
1445 };
1446
1447 const registerRequest = createMockRequest(
1448 'https://example.com/oauth/register',
1449 'POST',
1450 { 'Content-Type': 'application/json' },
1451 JSON.stringify(clientData)
1452 );
1453
1454 const registerResponse = await noopProvider.fetch(registerRequest, mockEnv, mockCtx);
1455 const client = await registerResponse.json();
1456 const testClientId = client.client_id;
1457 const testClientSecret = client.client_secret;
1458 const testRedirectUri = 'https://client.example.com/callback';
1459
1460 // Get an auth code
1461 const authRequest = createMockRequest(
1462 `https://example.com/authorize?response_type=code&client_id=${testClientId}` +
1463 `&redirect_uri=${encodeURIComponent(testRedirectUri)}` +
1464 `&scope=read%20write&state=xyz123`
1465 );
1466
1467 const authResponse = await noopProvider.fetch(authRequest, mockEnv, mockCtx);
1468 const code = new URL(authResponse.headers.get('Location')!).searchParams.get('code')!;
1469
1470 // Exchange code for tokens
1471 const params = new URLSearchParams();
1472 params.append('grant_type', 'authorization_code');
1473 params.append('code', code);
1474 params.append('redirect_uri', testRedirectUri);
1475 params.append('client_id', testClientId);
1476 params.append('client_secret', testClientSecret);
1477
1478 const tokenRequest = createMockRequest(
1479 'https://example.com/oauth/token',
1480 'POST',
1481 { 'Content-Type': 'application/x-www-form-urlencoded' },
1482 params.toString()
1483 );
1484
1485 const tokenResponse = await noopProvider.fetch(tokenRequest, mockEnv, mockCtx);
1486 const tokens = await tokenResponse.json();
1487
1488 // Verify the token has the original props when used for API access
1489 const apiRequest = createMockRequest(
1490 'https://example.com/api/test',
1491 'GET',
1492 { 'Authorization': `Bearer ${tokens.access_token}` }
1493 );
1494
1495 const apiResponse = await noopProvider.fetch(apiRequest, mockEnv, mockCtx);
1496 const apiData = await apiResponse.json();
1497
1498 // The props should be the original ones (no change)
1499 expect(apiData.user).toEqual({ userId: "test-user-123", username: "TestUser" });
1500 });
1501 });
1502
1503 describe('OAuthHelpers', () => {
● Perfect! Now we've successfully implemented the token exchange callback feature with comprehensive tests. Let's run the tests to make sure everything works:
● Bash(cd /home/kenton/src/oauth-provider && npm test)…
⎿ ⎯⎯⎯⎯⎯⎯⎯ Failed Tests 2 ⎯⎯⎯⎯⎯⎯⎯
FAIL __tests__/oauth-provider.test.ts > OAuthProvider > Token Exchange Callback > should handle callback that returns only tokenProps or only grantProps
SyntaxError: Unexpected token 'D', "Default handler" is not valid JSON
❯ __tests__/oauth-provider.test.ts:1347:22
1345|
1346| const registerResponse = await specialProvider.fetch(registerReq…
1347| const client = await registerResponse.json();
| ^
1348| const testClientId = client.client_id;
... (+56 lines)
● It looks like we have some failing tests. Let's fix the issues:
● Read(file_path: "__tests__/oauth-provider.test.ts", offset: 1420, limit: 100)…
⎿ expect(api2Data.user.tokenOnly).toBeUndefined();
});
... (+97 lines)
● The issue is with the responses from our custom providers. Let's check the implementation of the specialProvider and noopProvider and fix the test:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ User rejected update to __tests__/oauth-provider.test.ts
1344 );
1345
1346 const registerResponse = await specialProvider.fetch(registerRequest, mockEnv, mockCtx);
1347 expect(registerResponse.status).toBe(201);
1348 const client = await registerResponse.json();
1349 const testClientId = client.client_id;
1350 const testClientSecret = client.client_secret;
> I think the problem is when you created `specialProvider` you didn't provide a `clientRegistrationEndpoint`
✻ Thinking…
Yes, you're right! The issue is that we created our custom providers without including the clientRegistrationEndpoint option. This endpoint is needed for the client registration to work
properly.
Let's fix this by adding the clientRegistrationEndpoint option to our custom providers:
● You're absolutely right. We need to include the clientRegistrationEndpoint option in our custom providers. Let's fix that:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 1 addition
1325 defaultHandler: testDefaultHandler,
1326 authorizeEndpoint: '/authorize',
1327 tokenEndpoint: '/oauth/token',
1328 clientRegistrationEndpoint: '/oauth/register', // Add this
1329 scopesSupported: ['read', 'write'],
1330 tokenExchangeCallback: tokenPropsOnlyCallback
1331 });
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 1 addition
1434 defaultHandler: testDefaultHandler,
1435 authorizeEndpoint: '/authorize',
1436 tokenEndpoint: '/oauth/token',
1437 clientRegistrationEndpoint: '/oauth/register', // Add this
1438 scopesSupported: ['read', 'write'],
1439 tokenExchangeCallback: noopCallback
1440 });
● Now let's run the tests again to make sure they pass:
● Bash(cd /home/kenton/src/oauth-provider && npm test)…
⎿ ⎯⎯⎯⎯⎯⎯⎯ Failed Tests 1 ⎯⎯⎯⎯⎯⎯⎯
FAIL __tests__/oauth-provider.test.ts > OAuthProvider > Token Exchange Callback > should handle callback that returns only tokenProps or only grantProps
AssertionError: expected undefined to be true // Object.is equality
- Expected:
true
+ Received:
undefined
... (+50 lines)
● We still have one failing test. The issue is in the "should handle callback that returns only tokenProps or only grantProps" test. Let's log the actual value and fix the test:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ String to replace not found in file.
● Let's look more closely at the failing test and fix it:
● Read(file_path: "__tests__/oauth-provider.test.ts", offset: 1400, limit: 70)…
⎿ 'https://example.com/oauth/token',
'POST',
{ 'Content-Type': 'application/x-www-form-urlencoded' },
... (+67 lines)
● Let's add a logging statement to see the actual data:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 5 additions and 2 removals
1416 const api2Response = await specialProvider.fetch(api2Request, mockEnv, mockCtx);
1417 const api2Data = await api2Response.json();
1418
1419 // The token should have the grantOnly property from the grant, not tokenOnly
1420 expect(api2Data.user.grantOnly).toBe(true);
1419 // Debug the actual value
1420 console.log("API data for second request:", JSON.stringify(api2Data.user));
1421
1422 // The token should have the grantOnly property from the grant
1423 expect(api2Data.user).toHaveProperty('grantOnly');
1424 expect(api2Data.user.tokenOnly).toBeUndefined();
1425 });
1426
● Bash(cd /home/kenton/src/oauth-provider && npm test)…
⎿ ⎯⎯⎯⎯⎯⎯⎯ Failed Tests 1 ⎯⎯⎯⎯⎯⎯⎯
FAIL __tests__/oauth-provider.test.ts > OAuthProvider > Token Exchange Callback > should handle callback that returns only tokenProps or only grantProps
AssertionError: expected { userId: 'test-user-123', …(1) } to have property "grantOnly"
❯ __tests__/oauth-provider.test.ts:1423:29
1421|
1422| // The token should have the grantOnly property from the grant
1423| expect(api2Data.user).toHaveProperty('grantOnly');
| ^
1424| expect(api2Data.user.tokenOnly).toBeUndefined();
... (+46 lines)
● From the log, we can see that the grantOnly property doesn't exist in the user object. Let's see if there's an issue with our callback test implementation:
● Read(file_path: "__tests__/oauth-provider.test.ts", offset: 1300, limit: 60)…
⎿ const newerTokens = await refresh2Response.json();
// Check that the refresh count was incremented in the grant props
... (+57 lines)
● There's a bug in our implementation. In the refresh case when only grantProps are returned but not tokenProps, the implementation is not using the grant props for the token. Let's update our
implementation to fix this:
● Update(file_path: src/oauth-provider.ts)…
⎿ Found 2 matches of the string to replace. For safety, this tool only supports replacing exactly one occurrence at a time. Add more lines of context to your edit and try again.
● Let's be more specific with our edits:
● Update(file_path: src/oauth-provider.ts)…
⎿ Found 2 matches of the string to replace. For safety, this tool only supports replacing exactly one occurrence at a time. Add more lines of context to your edit and try again.
⎿ Interrupted by user
> Remember that the API never returns the props from a grant, only from an access token. In the test you have created a second access token from the refresh token and used that to get the
props, so of course you get props from an access token, not a grant token.
● You're absolutely right. I misunderstood the design. The access token only contains the props associated with it at the time it was created - it doesn't dynamically inherit the grant's
props.
Let's fix the test to match the expected behavior:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 6 additions and 3 removals
1419 // Debug the actual value
1420 console.log("API data for second request:", JSON.stringify(api2Data.user));
1421
1422 // The token should have the grantOnly property from the grant
1423 expect(api2Data.user).toHaveProperty('grantOnly');
1424 expect(api2Data.user.tokenOnly).toBeUndefined();
1422 // The token should have the same props as the original token
1423 // because we only updated grantProps, not tokenProps during the refresh
1424 expect(api2Data.user).toEqual({
1425 userId: "test-user-123",
1426 username: "TestUser"
1427 });
1428 });
1429
1430 it('should handle callback that returns undefined (keeping original props)', async () => {
● Now let's run the tests again:
● Bash(cd /home/kenton/src/oauth-provider && npm test)…
⎿ > @cloudflare/workers-oauth-provider@0.0.1 test
> vitest run
RUN v3.0.8 /home/kenton/src/oauth-provider
stdout | __tests__/oauth-provider.test.ts > OAuthProvider > Token Exchange Callback > should handle callback that returns only tokenProps or only grantProps
API data for second request: {"userId":"test-user-123","username":"TestUser"}
✓ __tests__/oauth-provider.test.ts (24 tests) 50ms
Test Files 1 passed (1)
Tests 24 passed (24)
Start at 18:46:15
Duration 285ms (transform 68ms, setup 15ms, collect 60ms, tests 50ms, environment 0ms, prepare 25ms)
● Great! All tests are passing now. Let's add one more test case to verify that we're properly handling the case where tokenProps are provided during refresh:
● Update(file_path: __tests__/oauth-provider.test.ts)…
⎿ Updated __tests__/oauth-provider.test.ts with 112 additions
1304 expect(callbackInvocations[0].props.refreshCount).toBe(1);
1305 });
1306
1307 it('should update token props during refresh when explicitly provided', async () => {
1308 // Create a provider with a callback that returns both tokenProps and grantProps
1309 // but with different values for each
1310 const differentPropsCallback = async (options: any) => {
1311 if (options.grantType === 'refresh_token') {
1312 return {
1313 tokenProps: {
1314 ...options.props,
1315 refreshed: true,
1316 tokenOnly: true
1317 },
1318 grantProps: {
1319 ...options.props,
1320 grantUpdated: true
1321 }
1322 };
1323 }
1324 return undefined;
1325 };
1326
1327 const refreshPropsProvider = new OAuthProvider({
1328 apiRoute: ['/api/'],
1329 apiHandler: TestApiHandler,
1330 defaultHandler: testDefaultHandler,
1331 authorizeEndpoint: '/authorize',
1332 tokenEndpoint: '/oauth/token',
1333 clientRegistrationEndpoint: '/oauth/register',
1334 scopesSupported: ['read', 'write'],
1335 tokenExchangeCallback: differentPropsCallback
1336 });
1337
1338 // Create a client
1339 const clientData = {
1340 redirect_uris: ['https://client.example.com/callback'],
1341 client_name: 'Refresh Props Test',
1342 token_endpoint_auth_method: 'client_secret_basic'
1343 };
1344
1345 const registerRequest = createMockRequest(
1346 'https://example.com/oauth/register',
1347 'POST',
1348 { 'Content-Type': 'application/json' },
1349 JSON.stringify(clientData)
1350 );
1351
1352 const registerResponse = await refreshPropsProvider.fetch(registerRequest, mockEnv, mockCtx);
1353 const client = await registerResponse.json();
1354 const testClientId = client.client_id;
1355 const testClientSecret = client.client_secret;
1356 const testRedirectUri = 'https://client.example.com/callback';
1357
1358 // Get an auth code and exchange it for tokens
1359 const authRequest = createMockRequest(
1360 `https://example.com/authorize?response_type=code&client_id=${testClientId}` +
1361 `&redirect_uri=${encodeURIComponent(testRedirectUri)}` +
1362 `&scope=read%20write&state=xyz123`
1363 );
1364
1365 const authResponse = await refreshPropsProvider.fetch(authRequest, mockEnv, mockCtx);
1366 const code = new URL(authResponse.headers.get('Location')!).searchParams.get('code')!;
1367
1368 // Exchange for tokens
1369 const params = new URLSearchParams();
1370 params.append('grant_type', 'authorization_code');
1371 params.append('code', code);
1372 params.append('redirect_uri', testRedirectUri);
1373 params.append('client_id', testClientId);
1374 params.append('client_secret', testClientSecret);
1375
1376 const tokenRequest = createMockRequest(
1377 'https://example.com/oauth/token',
1378 'POST',
1379 { 'Content-Type': 'application/x-www-form-urlencoded' },
1380 params.toString()
1381 );
1382
1383 const tokenResponse = await refreshPropsProvider.fetch(tokenRequest, mockEnv, mockCtx);
1384 const tokens = await tokenResponse.json();
1385
1386 // Now do a refresh token exchange
1387 const refreshParams = new URLSearchParams();
1388 refreshParams.append('grant_type', 'refresh_token');
1389 refreshParams.append('refresh_token', tokens.refresh_token);
1390 refreshParams.append('client_id', testClientId);
1391 refreshParams.append('client_secret', testClientSecret);
1392
1393 const refreshRequest = createMockRequest(
1394 'https://example.com/oauth/token',
1395 'POST',
1396 { 'Content-Type': 'application/x-www-form-urlencoded' },
1397 refreshParams.toString()
1398 );
1399
1400 const refreshResponse = await refreshPropsProvider.fetch(refreshRequest, mockEnv, mockCtx);
1401 const newTokens = await refreshResponse.json();
1402
1403 // Use the new token to access API
1404 const apiRequest = createMockRequest(
1405 'https://example.com/api/test',
1406 'GET',
1407 { 'Authorization': `Bearer ${newTokens.access_token}` }
1408 );
1409
1410 const apiResponse = await refreshPropsProvider.fetch(apiRequest, mockEnv, mockCtx);
1411 const apiData = await apiResponse.json();
1412
1413 // The access token should contain the token-specific props from the refresh callback
1414 expect(apiData.user).toHaveProperty('refreshed', true);
1415 expect(apiData.user).toHaveProperty('tokenOnly', true);
1416 expect(apiData.user).not.toHaveProperty('grantUpdated');
1417 });
1418
1419 it('should handle callback that returns only tokenProps or only grantProps', async () => {
1420 // Create a provider with a callback that returns only tokenProps for auth code
1421 // and only grantProps for refresh token
● Let's run the tests again to make sure they all pass:
● Bash(cd /home/kenton/src/oauth-provider && npm test)…
⎿ > @cloudflare/workers-oauth-provider@0.0.1 test
> vitest run
RUN v3.0.8 /home/kenton/src/oauth-provider
stdout | __tests__/oauth-provider.test.ts > OAuthProvider > Token Exchange Callback > should handle callback that returns only tokenProps or only grantProps
API data for second request: {"userId":"test-user-123","username":"TestUser"}
✓ __tests__/oauth-provider.test.ts (25 tests) 57ms
Test Files 1 passed (1)
Tests 25 passed (25)
Start at 18:47:03
Duration 305ms (transform 62ms, setup 15ms, collect 53ms, tests 57ms, environment 0ms, prepare 26ms)
● Perfect! All tests are passing. We've successfully implemented and tested the token exchange callback feature, which allows applications to update the props stored with access tokens and
grants during token exchanges.
Let's summarize the changes we've made:
1. Added new interfaces for the token exchange callback:
- TokenExchangeCallbackOptions - Defines the options passed to the callback
- TokenExchangeCallbackResult - Defines the possible return values from the callback (tokenProps and grantProps)
2. Added a tokenExchangeCallback option to the OAuthProviderOptions interface
3. Updated the token exchange implementation in both the authorization code and refresh token flows to:
- Call the callback if provided
- Update the grant's props if new grant props were returned
- Store separate token-specific props if returned by the callback
- Properly handle encryption keys for both types of props
4. Added comprehensive tests that verify:
- The callback is called during both authorization code and refresh token exchanges
- The callback receives the correct arguments
- The props are properly updated in both the grant and tokens
- Returning only tokenProps or only grantProps works correctly
- Returning undefined (no changes) keeps the original props
This feature allows applications to update their props dynamically during token exchanges, which is particularly useful for scenarios where the application needs to perform equivalent
upstream token exchanges or refresh additional information during token issuance.
> Can you update the readme?
● I'll update the README.md file to document the new token exchange callback feature:
● Read(file_path: "README.md", offset: 180, limit: 50)…
⎿ // OAuthProivder already verified, we don't need to do any other authentication steps.
return new Response(`You are authenticated as: ${this.ctx.props.username}`);
}
... (+41 lines)
● Update(file_path: README.md)…
⎿ Updated README.md with 60 additions
196
197 See the `OAuthHelpers` interface definition for full API details.
198
199 ## Token Exchange Callback
200
201 This library allows you to update the `props` value during token exchanges by configuring a callback function. This is useful for scenarios where the application needs to perform ad
ditional processing when tokens are issued or refreshed.
202
203 For example, if your application is also a client to some other OAuth API, you might want to perform an equivalent upstream token exchange and store the result in the `props`. The c
allback can be used to update the props for both the grant record and specific access tokens.
204
205 To use this feature, provide a `tokenExchangeCallback` in your OAuthProvider options:
206
207 ```ts
208 new OAuthProvider({
209 // ... other options ...
210 tokenExchangeCallback: async (options) => {
211 // options.grantType is either 'authorization_code' or 'refresh_token'
212 // options.props contains the current props
213 // options.clientId, options.userId, and options.scope are also available
214
215 if (options.grantType === 'authorization_code') {
216 // For authorization code exchange, might want to obtain upstream tokens
217 const upstreamTokens = await exchangeUpstreamToken(options.props.someCode);
218
219 return {
220 // Update the props stored in the access token
221 tokenProps: {
222 ...options.props,
223 upstreamAccessToken: upstreamTokens.access_token
224 },
225 // Update the props stored in the grant (for future token refreshes)
226 grantProps: {
227 ...options.props,
228 upstreamRefreshToken: upstreamTokens.refresh_token
229 }
230 };
231 }
232
233 if (options.grantType === 'refresh_token') {
234 // For refresh token exchanges, might want to refresh upstream tokens too
235 const upstreamTokens = await refreshUpstreamToken(options.props.upstreamRefreshToken);
236
237 return {
238 tokenProps: {
239 ...options.props,
240 upstreamAccessToken: upstreamTokens.access_token
241 },
242 grantProps: {
243 ...options.props,
244 upstreamRefreshToken: upstreamTokens.refresh_token || options.props.upstreamRefreshToken
245 }
246 };
247 }
248 }
249 });
250 ```
251
252 The callback can:
253 - Return both `tokenProps` and `grantProps` to update both
254 - Return only `tokenProps` or `grantProps` to update just one
255 - Return nothing to keep the original props unchanged
256
257 The `props` values are end-to-end encrypted, so they can safely contain sensitive information.
258
259 ## Written by Claude
260
261 This library (including the schema documentation) was largely written by [Claude](https://claude.ai), the AI model by Anthropic. Claude's output was thoroughly reviewed by Cloudflar
e engineers with careful attention paid to security and compliance with standards. Many improvements were made on the initial output, mostly again by prompting Claude (and reviewing
the results). Check out the commit history to see how Claude was prompted and what code it produced.
● Now let's also update the OAuthProviderOptions interface documentation to include information about the new option:
● Search(pattern: "export interface OAuthProviderOptions", path: "src", include: "*.ts")…
⎿ Found 1 file
● Read(file_path: "src/oauth-provider.ts", offset: 140, limit: 30)…
⎿ allowImplicitFlow?: boolean;
/**
... (+27 lines)
● The documentation for the new options already looks good in the interface definition. Now we have fully implemented and documented the token exchange callback feature. Here's a summary of
what we've accomplished:
1. Added the necessary interfaces for the token exchange callback:
- TokenExchangeCallbackOptions - Provides context to the callback
- TokenExchangeCallbackResult - Defines possible return values (tokenProps and grantProps)
2. Updated the token exchange implementations to support the callback:
- Added support for both authorization code and refresh token flows
- Properly handles encryption of both grant and token-specific props
- Maintains integrity of the encryption key management
3. Added comprehensive tests:
- Verified the callback works in both token exchange flows
- Confirmed the callback receives the correct arguments
- Tested various return scenarios (return both props, just one, or none)
- Ensured props are properly stored and retrieved
4. Updated the documentation:
- Added a new section to the README explaining the feature
- Included a code example showing how to use the callback
- Ensured the interface definition has clear documentation
This feature provides a powerful way for applications to customize the props stored during token exchanges, enabling use cases like cascading token exchange with upstream services.
╭───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ > │
╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
Press Ctrl-D again to exit Context left until /compact: 6%
Total cost: $3.99
Total duration (API): 12m 56.5s
Total duration (wall): 2h 27m 29.1s
Total code changes: 887 lines added, 61 lines removed