@cloudflare/workers-oauth-provider 0.0.1 → 0.0.3

package/README.md

@@ -2,9 +2,9 @@
 
 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.
 
-## EXPERIMENTAL
+## Beta
 
-As of March, 2025, this library is very new. Please use with caution as additional security reviews are still in progress.
+As of March, 2025, this library is very new, prerelease software. The API is still subject to change.
 
 ## Benefits of this library
 
@@ -196,11 +196,100 @@ The `env.OAUTH_PROVIDER` object available to the fetch handlers provides some me
 
 See the `OAuthHelpers` interface definition for full API details.
 
-## Written by Claude
+## Token Exchange Callback
 
-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 Cloudflare 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.
+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 additional processing when tokens are issued or refreshed.
 
-(@kentonv, the lead engineer, was actually an AI skeptic, and started this project with the intent to prove that LLMs cannot code. He ended up deciding he had proven himself wrong.)
+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 callback can be used to update the props for both the grant record and specific access tokens.
+
+To use this feature, provide a `tokenExchangeCallback` in your OAuthProvider options:
+
+```ts
+new OAuthProvider({
+  // ... other options ...
+  tokenExchangeCallback: async (options) => {
+    // options.grantType is either 'authorization_code' or 'refresh_token'
+    // options.props contains the current props
+    // options.clientId, options.userId, and options.scope are also available
+
+    if (options.grantType === 'authorization_code') {
+      // For authorization code exchange, might want to obtain upstream tokens
+      const upstreamTokens = await exchangeUpstreamToken(options.props.someCode);
+
+      return {
+        // Update the props stored in the access token
+        accessTokenProps: {
+          ...options.props,
+          upstreamAccessToken: upstreamTokens.access_token
+        },
+        // Update the props stored in the grant (for future token refreshes)
+        newProps: {
+          ...options.props,
+          upstreamRefreshToken: upstreamTokens.refresh_token
+        }
+      };
+    }
+
+    if (options.grantType === 'refresh_token') {
+      // For refresh token exchanges, might want to refresh upstream tokens too
+      const upstreamTokens = await refreshUpstreamToken(options.props.upstreamRefreshToken);
+
+      return {
+        accessTokenProps: {
+          ...options.props,
+          upstreamAccessToken: upstreamTokens.access_token
+        },
+        newProps: {
+          ...options.props,
+          upstreamRefreshToken: upstreamTokens.refresh_token || options.props.upstreamRefreshToken
+        },
+        // Optionally override the default access token TTL to match the upstream token
+        accessTokenTTL: upstreamTokens.expires_in
+      };
+    }
+  }
+});
+```
+
+The callback can:
+- Return both `accessTokenProps` and `newProps` to update both
+- Return only `accessTokenProps` to update just the current access token
+- Return only `newProps` to update both the grant and access token (the access token inherits these props)
+- Return `accessTokenTTL` to override the default TTL for this specific access token
+- Return nothing to keep the original props unchanged
+
+The `accessTokenTTL` override is particularly useful when the application is also an OAuth client to another service and wants to match its access token TTL to the upstream access token TTL. This helps prevent situations where the downstream token is still valid but the upstream token has expired.
+
+The `props` values are end-to-end encrypted, so they can safely contain sensitive information.
+
+## Custom Error Responses
+
+By using the `onError` option, you can emit notifications or take other actions when an error response was to be emitted:
+
+```ts
+new OAuthProvider({
+  // ... other options ...
+  onError({ code, description, status, headers }) {
+    Sentry.captureMessage(/* ... */)
+  }
+})
+```
+
+By returning a `Response` you can also override what the OAuthProvider returns to your users:
+
+```ts
+new OAuthProvider({
+  // ... other options ...
+  onError({ code, description, status, headers }) {
+    if (code === 'unsupported_grant_type') {
+      return new Response('...', { status, headers })
+    }
+    // returning undefined (i.e. void) uses the default Response generation
+  }
+})
+```
+
+By default, the `onError` callback is set to ``({ status, code, description }) => console.warn(`OAuth error response: ${status} ${code} - ${description}`)``.
 
 ## Implementation Notes
 
@@ -220,3 +309,17 @@ OAuth 2.1 requires that refresh tokens are either "cryptographically bound" to t
 This requirement is seemingly fundamentally flawed as it assumes that every refresh request will complete with no errors. In the real world, a transient network error, machine failure, or software fault could mean that the client fails to store the new refresh token after a refresh request. In this case, the client would be permanently unable to make any further requests, as the only token it has is no longer valid.
 
 This library implements a compromise: At any particular time, a grant may have two valid refresh tokens. When the client uses one of them, the other one is invalidated, and a new one is generated and returned. Thus, if the client correctly uses the new refresh token each time, then older refresh tokens are continuously invalidated. But if a transient failure prevents the client from updating its token, it can always retry the request with the token it used previously.
+
+## Written using Claude
+
+This library (including the schema documentation) was largely written with the help of [Claude](https://claude.ai), the AI model by Anthropic. Claude's output was thoroughly reviewed by Cloudflare 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.
+
+**"NOOOOOOOO!!!! You can't just use an LLM to write an auth library!"**
+
+"haha gpus go brrr"
+
+In all seriousness, two months ago (January 2025), I ([@kentonv](https://github.com/kentonv)) would have agreed. I was an AI skeptic. I thoughts LLMs were glorified Markov chain generators that didn't actually understand code and couldn't produce anything novel. I started this project on a lark, fully expecting the AI to produce terrible code for me to laugh at. And then, uh... the code actually looked pretty good. Not perfect, but I just told the AI to fix things, and it did. I was shocked.
+
+To emphasize, **this is not "vibe coded"**. Every line was thoroughly reviewed and cross-referenced with relevant RFCs, by security experts with previous experience with those RFCs. I was *trying* to validate my skepticism. I ended up proving myself wrong.
+
+Again, please check out the commit history -- especially early commits -- to understand how this went.

package/dist/oauth-provider.d.ts

@@ -9,6 +9,59 @@ type WorkerEntrypointWithFetch = WorkerEntrypoint & Pick<Required<WorkerEntrypoi
 /**
  * Configuration options for the OAuth Provider
  */
+/**
+ * Result of a token exchange callback function.
+ * Allows updating the props stored in both the access token and the grant.
+ */
+interface TokenExchangeCallbackResult {
+    /**
+     * New props to be stored specifically with the access token.
+     * If not provided but newProps is, the access token will use newProps.
+     * If neither is provided, the original props will be used.
+     */
+    accessTokenProps?: any;
+    /**
+     * New props to replace the props stored in the grant itself.
+     * These props will be used for all future token refreshes.
+     * If accessTokenProps is not provided, these props will also be used for the current access token.
+     * If not provided, the original props will be used.
+     */
+    newProps?: any;
+    /**
+     * Override the default access token TTL (time-to-live) for this specific token.
+     * This is especially useful when the application is also an OAuth client to another service
+     * and wants to match its access token TTL to the upstream access token TTL.
+     * Value should be in seconds.
+     */
+    accessTokenTTL?: number;
+}
+/**
+ * Options for token exchange callback functions
+ */
+interface TokenExchangeCallbackOptions {
+    /**
+     * The type of grant being processed.
+     * 'authorization_code' for initial code exchange,
+     * 'refresh_token' for refresh token exchange.
+     */
+    grantType: 'authorization_code' | 'refresh_token';
+    /**
+     * Client that received this grant
+     */
+    clientId: string;
+    /**
+     * User who authorized this grant
+     */
+    userId: string;
+    /**
+     * List of scopes that were granted
+     */
+    scope: string[];
+    /**
+     * Application-specific properties currently associated with this grant
+     */
+    props: any;
+}
 interface OAuthProviderOptions {
     /**
      * URL(s) for API routes. Requests with URLs starting with any of these prefixes
@@ -65,6 +118,28 @@ interface OAuthProviderOptions {
      * Defaults to false.
      */
     disallowPublicClientRegistration?: boolean;
+    /**
+     * Optional callback function that is called during token exchange.
+     * This allows updating the props stored in both the access token and the grant.
+     * 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 in the props.
+     *
+     * The callback can return new props values that will be stored with the token or grant.
+     * If the callback returns nothing or undefined for a props field, the original props will be used.
+     */
+    tokenExchangeCallback?: (options: TokenExchangeCallbackOptions) => Promise<TokenExchangeCallbackResult | void> | TokenExchangeCallbackResult | void;
+    /**
+     * Optional callback function that is called whenever the OAuthProvider returns an error response
+     * This allows the client to emit notifications or perform other actions when an error occurs.
+     *
+     * If the function returns a Response, that will be used in place of the OAuthProvider's default one.
+     */
+    onError?: (error: {
+        code: string;
+        description: string;
+        status: number;
+        headers: Record<string, string>;
+    }) => Response | void;
 }
 /**
  * Helper methods for OAuth operations provided to handler functions
@@ -456,4 +531,4 @@ declare class OAuthProvider {
     fetch(request: Request, env: any, ctx: ExecutionContext): Promise<Response>;
 }
 
-export { type AuthRequest, type ClientInfo, type CompleteAuthorizationOptions, type Grant, type GrantSummary, type ListOptions, type ListResult, type OAuthHelpers, OAuthProvider, type OAuthProviderOptions, type Token, OAuthProvider as default };
+export { type AuthRequest, type ClientInfo, type CompleteAuthorizationOptions, type Grant, type GrantSummary, type ListOptions, type ListResult, type OAuthHelpers, OAuthProvider, type OAuthProviderOptions, type Token, type TokenExchangeCallbackOptions, type TokenExchangeCallbackResult, OAuthProvider as default };

package/dist/oauth-provider.js

@@ -52,8 +52,9 @@ var OAuthProviderImpl = class {
       this.validateEndpoint(options.clientRegistrationEndpoint, "clientRegistrationEndpoint");
     }
     this.options = {
-      ...options,
-      accessTokenTTL: options.accessTokenTTL || DEFAULT_ACCESS_TOKEN_TTL
+      accessTokenTTL: DEFAULT_ACCESS_TOKEN_TTL,
+      onError: ({ status, code, description }) => console.warn(`OAuth error response: ${status} ${code} - ${description}`),
+      ...options
     };
   }
   /**
@@ -89,7 +90,9 @@ var OAuthProviderImpl = class {
     if (typeof handler === "function" && handler.prototype instanceof WorkerEntrypoint) {
       return { type: 1 /* WORKER_ENTRYPOINT */, handler };
     }
-    throw new TypeError(`${name} must be either an ExportedHandler object with a fetch method or a class extending WorkerEntrypoint`);
+    throw new TypeError(
+      `${name} must be either an ExportedHandler object with a fetch method or a class extending WorkerEntrypoint`
+    );
   }
   /**
    * Main fetch handler for the Worker
@@ -132,7 +135,11 @@ var OAuthProviderImpl = class {
       env.OAUTH_PROVIDER = this.createOAuthHelpers(env);
     }
     if (this.typedDefaultHandler.type === 0 /* EXPORTED_HANDLER */) {
-      return this.typedDefaultHandler.handler.fetch(request, env, ctx);
+      return this.typedDefaultHandler.handler.fetch(
+        request,
+        env,
+        ctx
+      );
     } else {
       const handler = new this.typedDefaultHandler.handler(ctx, env);
       return handler.fetch(request);
@@ -292,20 +299,12 @@ var OAuthProviderImpl = class {
    */
   async handleTokenRequest(request, env) {
     if (request.method !== "POST") {
-      return createErrorResponse(
-        "invalid_request",
-        "Method not allowed",
-        405
-      );
+      return this.createErrorResponse("invalid_request", "Method not allowed", 405);
     }
     let contentType = request.headers.get("Content-Type") || "";
     let body = {};
     if (!contentType.includes("application/x-www-form-urlencoded")) {
-      return createErrorResponse(
-        "invalid_request",
-        "Content-Type must be application/x-www-form-urlencoded",
-        400
-      );
+      return this.createErrorResponse("invalid_request", "Content-Type must be application/x-www-form-urlencoded", 400);
     }
     const formData = await request.formData();
     for (const [key, value] of formData.entries()) {
@@ -324,31 +323,19 @@ var OAuthProviderImpl = class {
       clientSecret = body.client_secret || "";
     }
     if (!clientId) {
-      return createErrorResponse(
-        "invalid_client",
-        "Client ID is required",
-        401
-      );
+      return this.createErrorResponse("invalid_client", "Client ID is required", 401);
     }
     const clientInfo = await this.getClient(env, clientId);
     if (!clientInfo) {
-      return createErrorResponse(
-        "invalid_client",
-        "Client not found",
-        401
-      );
+      return this.createErrorResponse("invalid_client", "Client not found", 401);
     }
     const isPublicClient = clientInfo.tokenEndpointAuthMethod === "none";
     if (!isPublicClient) {
       if (!clientSecret) {
-        return createErrorResponse(
-          "invalid_client",
-          "Client authentication failed: missing client_secret",
-          401
-        );
+        return this.createErrorResponse("invalid_client", "Client authentication failed: missing client_secret", 401);
       }
       if (!clientInfo.clientSecret) {
-        return createErrorResponse(
+        return this.createErrorResponse(
           "invalid_client",
           "Client authentication failed: client has no registered secret",
           401
@@ -356,11 +343,7 @@ var OAuthProviderImpl = class {
       }
       const providedSecretHash = await hashSecret(clientSecret);
       if (providedSecretHash !== clientInfo.clientSecret) {
-        return createErrorResponse(
-          "invalid_client",
-          "Client authentication failed: invalid client_secret",
-          401
-        );
+        return this.createErrorResponse("invalid_client", "Client authentication failed: invalid client_secret", 401);
       }
     }
     const grantType = body.grant_type;
@@ -369,10 +352,7 @@ var OAuthProviderImpl = class {
     } else if (grantType === "refresh_token") {
       return this.handleRefreshTokenGrant(body, clientInfo, env);
     } else {
-      return createErrorResponse(
-        "unsupported_grant_type",
-        "Grant type not supported"
-      );
+      return this.createErrorResponse("unsupported_grant_type", "Grant type not supported");
     }
   }
   /**
@@ -388,65 +368,38 @@ var OAuthProviderImpl = class {
     const redirectUri = body.redirect_uri;
     const codeVerifier = body.code_verifier;
     if (!code) {
-      return createErrorResponse(
-        "invalid_request",
-        "Authorization code is required"
-      );
+      return this.createErrorResponse("invalid_request", "Authorization code is required");
     }
     const codeParts = code.split(":");
     if (codeParts.length !== 3) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Invalid authorization code format"
-      );
+      return this.createErrorResponse("invalid_grant", "Invalid authorization code format");
     }
     const [userId, grantId, _] = codeParts;
     const grantKey = `grant:${userId}:${grantId}`;
     const grantData = await env.OAUTH_KV.get(grantKey, { type: "json" });
     if (!grantData) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Grant not found or authorization code expired"
-      );
+      return this.createErrorResponse("invalid_grant", "Grant not found or authorization code expired");
     }
     if (!grantData.authCodeId) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Authorization code already used"
-      );
+      return this.createErrorResponse("invalid_grant", "Authorization code already used");
     }
     const codeHash = await hashSecret(code);
     if (codeHash !== grantData.authCodeId) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Invalid authorization code"
-      );
+      return this.createErrorResponse("invalid_grant", "Invalid authorization code");
     }
     if (grantData.clientId !== clientInfo.clientId) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Client ID mismatch"
-      );
+      return this.createErrorResponse("invalid_grant", "Client ID mismatch");
     }
     const isPkceEnabled = !!grantData.codeChallenge;
     if (!redirectUri && !isPkceEnabled) {
-      return createErrorResponse(
-        "invalid_request",
-        "redirect_uri is required when not using PKCE"
-      );
+      return this.createErrorResponse("invalid_request", "redirect_uri is required when not using PKCE");
     }
     if (redirectUri && !clientInfo.redirectUris.includes(redirectUri)) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Invalid redirect URI"
-      );
+      return this.createErrorResponse("invalid_grant", "Invalid redirect URI");
     }
     if (isPkceEnabled) {
       if (!codeVerifier) {
-        return createErrorResponse(
-          "invalid_request",
-          "code_verifier is required for PKCE"
-        );
+        return this.createErrorResponse("invalid_request", "code_verifier is required for PKCE");
       }
       let calculatedChallenge;
       if (grantData.codeChallengeMethod === "S256") {
@@ -459,10 +412,7 @@ var OAuthProviderImpl = class {
         calculatedChallenge = codeVerifier;
       }
       if (calculatedChallenge !== grantData.codeChallenge) {
-        return createErrorResponse(
-          "invalid_grant",
-          "Invalid PKCE code_verifier"
-        );
+        return this.createErrorResponse("invalid_grant", "Invalid PKCE code_verifier");
       }
     }
     const accessTokenSecret = generateRandomString(TOKEN_LENGTH);
@@ -471,11 +421,53 @@ var OAuthProviderImpl = class {
     const refreshToken = `${userId}:${grantId}:${refreshTokenSecret}`;
     const accessTokenId = await generateTokenId(accessToken);
     const refreshTokenId = await generateTokenId(refreshToken);
-    const now = Math.floor(Date.now() / 1e3);
-    const accessTokenExpiresAt = now + this.options.accessTokenTTL;
+    let accessTokenTTL = this.options.accessTokenTTL;
     const encryptionKey = await unwrapKeyWithToken(code, grantData.authCodeWrappedKey);
-    const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, encryptionKey);
-    const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, encryptionKey);
+    let grantEncryptionKey = encryptionKey;
+    let accessTokenEncryptionKey = encryptionKey;
+    let encryptedAccessTokenProps = grantData.encryptedProps;
+    if (this.options.tokenExchangeCallback) {
+      const decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
+      let grantProps = decryptedProps;
+      let accessTokenProps = decryptedProps;
+      const callbackOptions = {
+        grantType: "authorization_code",
+        clientId: clientInfo.clientId,
+        userId,
+        scope: grantData.scope,
+        props: decryptedProps
+      };
+      const callbackResult = await Promise.resolve(this.options.tokenExchangeCallback(callbackOptions));
+      if (callbackResult) {
+        if (callbackResult.newProps) {
+          grantProps = callbackResult.newProps;
+          if (!callbackResult.accessTokenProps) {
+            accessTokenProps = callbackResult.newProps;
+          }
+        }
+        if (callbackResult.accessTokenProps) {
+          accessTokenProps = callbackResult.accessTokenProps;
+        }
+        if (callbackResult.accessTokenTTL !== void 0) {
+          accessTokenTTL = callbackResult.accessTokenTTL;
+        }
+      }
+      const grantResult = await encryptProps(grantProps);
+      grantData.encryptedProps = grantResult.encryptedData;
+      grantEncryptionKey = grantResult.key;
+      if (accessTokenProps !== grantProps) {
+        const tokenResult = await encryptProps(accessTokenProps);
+        encryptedAccessTokenProps = tokenResult.encryptedData;
+        accessTokenEncryptionKey = tokenResult.key;
+      } else {
+        encryptedAccessTokenProps = grantData.encryptedProps;
+        accessTokenEncryptionKey = grantEncryptionKey;
+      }
+    }
+    const now = Math.floor(Date.now() / 1e3);
+    const accessTokenExpiresAt = now + accessTokenTTL;
+    const accessTokenWrappedKey = await wrapKeyWithToken(accessToken, accessTokenEncryptionKey);
+    const refreshTokenWrappedKey = await wrapKeyWithToken(refreshToken, grantEncryptionKey);
     delete grantData.authCodeId;
     delete grantData.codeChallenge;
     delete grantData.codeChallengeMethod;
@@ -495,23 +487,24 @@ var OAuthProviderImpl = class {
       grant: {
         clientId: grantData.clientId,
         scope: grantData.scope,
-        encryptedProps: grantData.encryptedProps
+        encryptedProps: encryptedAccessTokenProps
       }
     };
-    await env.OAUTH_KV.put(
-      `token:${userId}:${grantId}:${accessTokenId}`,
-      JSON.stringify(accessTokenData),
-      { expirationTtl: this.options.accessTokenTTL }
-    );
-    return new Response(JSON.stringify({
-      access_token: accessToken,
-      token_type: "bearer",
-      expires_in: this.options.accessTokenTTL,
-      refresh_token: refreshToken,
-      scope: grantData.scope.join(" ")
-    }), {
-      headers: { "Content-Type": "application/json" }
+    await env.OAUTH_KV.put(`token:${userId}:${grantId}:${accessTokenId}`, JSON.stringify(accessTokenData), {
+      expirationTtl: accessTokenTTL
     });
+    return new Response(
+      JSON.stringify({
+        access_token: accessToken,
+        token_type: "bearer",
+        expires_in: accessTokenTTL,
+        refresh_token: refreshToken,
+        scope: grantData.scope.join(" ")
+      }),
+      {
+        headers: { "Content-Type": "application/json" }
+      }
+    );
   }
   /**
    * Handles the refresh token grant type
@@ -524,41 +517,26 @@ var OAuthProviderImpl = class {
   async handleRefreshTokenGrant(body, clientInfo, env) {
     const refreshToken = body.refresh_token;
     if (!refreshToken) {
-      return createErrorResponse(
-        "invalid_request",
-        "Refresh token is required"
-      );
+      return this.createErrorResponse("invalid_request", "Refresh token is required");
     }
     const tokenParts = refreshToken.split(":");
     if (tokenParts.length !== 3) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Invalid token format"
-      );
+      return this.createErrorResponse("invalid_grant", "Invalid token format");
     }
     const [userId, grantId, _] = tokenParts;
     const providedTokenHash = await generateTokenId(refreshToken);
     const grantKey = `grant:${userId}:${grantId}`;
     const grantData = await env.OAUTH_KV.get(grantKey, { type: "json" });
     if (!grantData) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Grant not found"
-      );
+      return this.createErrorResponse("invalid_grant", "Grant not found");
     }
     const isCurrentToken = grantData.refreshTokenId === providedTokenHash;
     const isPreviousToken = grantData.previousRefreshTokenId === providedTokenHash;
     if (!isCurrentToken && !isPreviousToken) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Invalid refresh token"
-      );
+      return this.createErrorResponse("invalid_grant", "Invalid refresh token");
     }
     if (grantData.clientId !== clientInfo.clientId) {
-      return createErrorResponse(
-        "invalid_grant",
-        "Client ID mismatch"
-      );
+      return this.createErrorResponse("invalid_grant", "Client ID mismatch");
     }
     const accessTokenSecret = generateRandomString(TOKEN_LENGTH);
     const newAccessToken = `${userId}:${grantId}:${accessTokenSecret}`;
@@ -566,8 +544,7 @@ var OAuthProviderImpl = class {
     const refreshTokenSecret = generateRandomString(TOKEN_LENGTH);
     const newRefreshToken = `${userId}:${grantId}:${refreshTokenSecret}`;
     const newRefreshTokenId = await generateTokenId(newRefreshToken);
-    const now = Math.floor(Date.now() / 1e3);
-    const accessTokenExpiresAt = now + this.options.accessTokenTTL;
+    let accessTokenTTL = this.options.accessTokenTTL;
     let wrappedKeyToUse;
     if (isCurrentToken) {
       wrappedKeyToUse = grantData.refreshTokenWrappedKey;
@@ -575,8 +552,60 @@ var OAuthProviderImpl = class {
       wrappedKeyToUse = grantData.previousRefreshTokenWrappedKey;
     }
     const encryptionKey = await unwrapKeyWithToken(refreshToken, wrappedKeyToUse);
-    const accessTokenWrappedKey = await wrapKeyWithToken(newAccessToken, encryptionKey);
-    const newRefreshTokenWrappedKey = await wrapKeyWithToken(newRefreshToken, encryptionKey);
+    let grantEncryptionKey = encryptionKey;
+    let accessTokenEncryptionKey = encryptionKey;
+    let encryptedAccessTokenProps = grantData.encryptedProps;
+    if (this.options.tokenExchangeCallback) {
+      const decryptedProps = await decryptProps(encryptionKey, grantData.encryptedProps);
+      let grantProps = decryptedProps;
+      let accessTokenProps = decryptedProps;
+      const callbackOptions = {
+        grantType: "refresh_token",
+        clientId: clientInfo.clientId,
+        userId,
+        scope: grantData.scope,
+        props: decryptedProps
+      };
+      const callbackResult = await Promise.resolve(this.options.tokenExchangeCallback(callbackOptions));
+      let grantPropsChanged = false;
+      if (callbackResult) {
+        if (callbackResult.newProps) {
+          grantProps = callbackResult.newProps;
+          grantPropsChanged = true;
+          if (!callbackResult.accessTokenProps) {
+            accessTokenProps = callbackResult.newProps;
+          }
+        }
+        if (callbackResult.accessTokenProps) {
+          accessTokenProps = callbackResult.accessTokenProps;
+        }
+        if (callbackResult.accessTokenTTL !== void 0) {
+          accessTokenTTL = callbackResult.accessTokenTTL;
+        }
+      }
+      if (grantPropsChanged) {
+        const grantResult = await encryptProps(grantProps);
+        grantData.encryptedProps = grantResult.encryptedData;
+        if (grantResult.key !== encryptionKey) {
+          grantEncryptionKey = grantResult.key;
+          wrappedKeyToUse = await wrapKeyWithToken(refreshToken, grantEncryptionKey);
+        } else {
+          grantEncryptionKey = grantResult.key;
+        }
+      }
+      if (accessTokenProps !== grantProps) {
+        const tokenResult = await encryptProps(accessTokenProps);
+        encryptedAccessTokenProps = tokenResult.encryptedData;
+        accessTokenEncryptionKey = tokenResult.key;
+      } else {
+        encryptedAccessTokenProps = grantData.encryptedProps;
+        accessTokenEncryptionKey = grantEncryptionKey;
+      }
+    }
+    const now = Math.floor(Date.now() / 1e3);
+    const accessTokenExpiresAt = now + accessTokenTTL;
+    const accessTokenWrappedKey = await wrapKeyWithToken(newAccessToken, accessTokenEncryptionKey);
+    const newRefreshTokenWrappedKey = await wrapKeyWithToken(newRefreshToken, grantEncryptionKey);
     grantData.previousRefreshTokenId = providedTokenHash;
     grantData.previousRefreshTokenWrappedKey = wrappedKeyToUse;
     grantData.refreshTokenId = newRefreshTokenId;
@@ -592,23 +621,24 @@ var OAuthProviderImpl = class {
       grant: {
         clientId: grantData.clientId,
         scope: grantData.scope,
-        encryptedProps: grantData.encryptedProps
+        encryptedProps: encryptedAccessTokenProps
       }
     };
-    await env.OAUTH_KV.put(
-      `token:${userId}:${grantId}:${accessTokenId}`,
-      JSON.stringify(accessTokenData),
-      { expirationTtl: this.options.accessTokenTTL }
-    );
-    return new Response(JSON.stringify({
-      access_token: newAccessToken,
-      token_type: "bearer",
-      expires_in: this.options.accessTokenTTL,
-      refresh_token: newRefreshToken,
-      scope: grantData.scope.join(" ")
-    }), {
-      headers: { "Content-Type": "application/json" }
+    await env.OAUTH_KV.put(`token:${userId}:${grantId}:${accessTokenId}`, JSON.stringify(accessTokenData), {
+      expirationTtl: accessTokenTTL
     });
+    return new Response(
+      JSON.stringify({
+        access_token: newAccessToken,
+        token_type: "bearer",
+        expires_in: accessTokenTTL,
+        refresh_token: newRefreshToken,
+        scope: grantData.scope.join(" ")
+      }),
+      {
+        headers: { "Content-Type": "application/json" }
+      }
+    );
   }
   /**
    * Handles the dynamic client registration endpoint (RFC 7591)
@@ -618,44 +648,24 @@ var OAuthProviderImpl = class {
    */
   async handleClientRegistration(request, env) {
     if (!this.options.clientRegistrationEndpoint) {
-      return createErrorResponse(
-        "not_implemented",
-        "Client registration is not enabled",
-        501
-      );
+      return this.createErrorResponse("not_implemented", "Client registration is not enabled", 501);
     }
     if (request.method !== "POST") {
-      return createErrorResponse(
-        "invalid_request",
-        "Method not allowed",
-        405
-      );
+      return this.createErrorResponse("invalid_request", "Method not allowed", 405);
     }
     const contentLength = parseInt(request.headers.get("Content-Length") || "0", 10);
     if (contentLength > 1048576) {
-      return createErrorResponse(
-        "invalid_request",
-        "Request payload too large, must be under 1 MiB",
-        413
-      );
+      return this.createErrorResponse("invalid_request", "Request payload too large, must be under 1 MiB", 413);
     }
     let clientMetadata;
     try {
       const text = await request.text();
       if (text.length > 1048576) {
-        return createErrorResponse(
-          "invalid_request",
-          "Request payload too large, must be under 1 MiB",
-          413
-        );
+        return this.createErrorResponse("invalid_request", "Request payload too large, must be under 1 MiB", 413);
       }
       clientMetadata = JSON.parse(text);
     } catch (error) {
-      return createErrorResponse(
-        "invalid_request",
-        "Invalid JSON payload",
-        400
-      );
+      return this.createErrorResponse("invalid_request", "Invalid JSON payload", 400);
     }
     const validateStringField = (field) => {
       if (field === void 0) {
@@ -683,10 +693,7 @@ var OAuthProviderImpl = class {
     const authMethod = validateStringField(clientMetadata.token_endpoint_auth_method) || "client_secret_basic";
     const isPublicClient = authMethod === "none";
     if (isPublicClient && this.options.disallowPublicClientRegistration) {
-      return createErrorResponse(
-        "invalid_client_metadata",
-        "Public client registration is not allowed"
-      );
+      return this.createErrorResponse("invalid_client_metadata", "Public client registration is not allowed");
     }
     const clientId = generateRandomString(16);
     let clientSecret;
@@ -720,7 +727,7 @@ var OAuthProviderImpl = class {
         clientInfo.clientSecret = hashedSecret;
       }
     } catch (error) {
-      return createErrorResponse(
+      return this.createErrorResponse(
         "invalid_client_metadata",
         error instanceof Error ? error.message : "Invalid client metadata"
       );
@@ -760,49 +767,34 @@ var OAuthProviderImpl = class {
   async handleApiRequest(request, env, ctx) {
     const authHeader = request.headers.get("Authorization");
     if (!authHeader || !authHeader.startsWith("Bearer ")) {
-      return createErrorResponse(
-        "invalid_token",
-        "Missing or invalid access token",
-        401,
-        { "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token", error_description="Missing or invalid access token"' }
-      );
+      return this.createErrorResponse("invalid_token", "Missing or invalid access token", 401, {
+        "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token", error_description="Missing or invalid access token"'
+      });
     }
     const accessToken = authHeader.substring(7);
     const tokenParts = accessToken.split(":");
     if (tokenParts.length !== 3) {
-      return createErrorResponse(
-        "invalid_token",
-        "Invalid token format",
-        401,
-        { "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token"' }
-      );
+      return this.createErrorResponse("invalid_token", "Invalid token format", 401, {
+        "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token"'
+      });
     }
     const [userId, grantId, _] = tokenParts;
     const accessTokenId = await generateTokenId(accessToken);
     const tokenKey = `token:${userId}:${grantId}:${accessTokenId}`;
     const tokenData = await env.OAUTH_KV.get(tokenKey, { type: "json" });
     if (!tokenData) {
-      return createErrorResponse(
-        "invalid_token",
-        "Invalid access token",
-        401,
-        { "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token"' }
-      );
+      return this.createErrorResponse("invalid_token", "Invalid access token", 401, {
+        "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token"'
+      });
     }
     const now = Math.floor(Date.now() / 1e3);
     if (tokenData.expiresAt < now) {
-      return createErrorResponse(
-        "invalid_token",
-        "Access token expired",
-        401,
-        { "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token"' }
-      );
+      return this.createErrorResponse("invalid_token", "Access token expired", 401, {
+        "WWW-Authenticate": 'Bearer realm="OAuth", error="invalid_token"'
+      });
     }
     const encryptionKey = await unwrapKeyWithToken(accessToken, tokenData.wrappedEncryptionKey);
-    const decryptedProps = await decryptProps(
-      encryptionKey,
-      tokenData.grant.encryptedProps
-    );
+    const decryptedProps = await decryptProps(encryptionKey, tokenData.grant.encryptedProps);
     ctx.props = decryptedProps;
     if (!env.OAUTH_PROVIDER) {
       env.OAUTH_PROVIDER = this.createOAuthHelpers(env);
@@ -836,22 +828,32 @@ var OAuthProviderImpl = class {
     const clientKey = `client:${clientId}`;
     return env.OAUTH_KV.get(clientKey, { type: "json" });
   }
+  /**
+   * Helper function to create OAuth error responses
+   * @param code - OAuth error code (e.g., 'invalid_request', 'invalid_token')
+   * @param description - Human-readable error description
+   * @param status - HTTP status code (default: 400)
+   * @param headers - Additional headers to include
+   * @returns A Response object with the error
+   */
+  createErrorResponse(code, description, status = 400, headers = {}) {
+    const customErrorResponse = this.options.onError?.({ code, description, status, headers });
+    if (customErrorResponse) return customErrorResponse;
+    const body = JSON.stringify({
+      error: code,
+      error_description: description
+    });
+    return new Response(body, {
+      status,
+      headers: {
+        "Content-Type": "application/json",
+        ...headers
+      }
+    });
+  }
 };
 var DEFAULT_ACCESS_TOKEN_TTL = 60 * 60;
 var TOKEN_LENGTH = 32;
-function createErrorResponse(code, description, status = 400, headers = {}) {
-  const body = JSON.stringify({
-    error: code,
-    error_description: description
-  });
-  return new Response(body, {
-    status,
-    headers: {
-      "Content-Type": "application/json",
-      ...headers
-    }
-  });
-}
 async function hashSecret(secret) {
   return generateTokenId(secret);
 }
@@ -972,11 +974,7 @@ async function deriveKeyFromToken(tokenStr) {
     false,
     ["sign"]
   );
-  const hmacResult = await crypto.subtle.sign(
-    "HMAC",
-    hmacKey,
-    encoder.encode(tokenStr)
-  );
+  const hmacResult = await crypto.subtle.sign("HMAC", hmacKey, encoder.encode(tokenStr));
   return await crypto.subtle.importKey(
     "raw",
     hmacResult,
@@ -988,12 +986,7 @@ async function deriveKeyFromToken(tokenStr) {
 }
 async function wrapKeyWithToken(tokenStr, keyToWrap) {
   const wrappingKey = await deriveKeyFromToken(tokenStr);
-  const wrappedKeyBuffer = await crypto.subtle.wrapKey(
-    "raw",
-    keyToWrap,
-    wrappingKey,
-    { name: "AES-KW" }
-  );
+  const wrappedKeyBuffer = await crypto.subtle.wrapKey("raw", keyToWrap, wrappingKey, { name: "AES-KW" });
   return arrayBufferToBase64(wrappedKeyBuffer);
 }
 async function unwrapKeyWithToken(tokenStr, wrappedKeyBase64) {
@@ -1319,9 +1312,11 @@ var OAuthHelpersImpl = class {
       }
       const result = await this.env.OAUTH_KV.list(listOptions);
       if (result.keys.length > 0) {
-        await Promise.all(result.keys.map((key) => {
-          return this.env.OAUTH_KV.delete(key.name);
-        }));
+        await Promise.all(
+          result.keys.map((key) => {
+            return this.env.OAUTH_KV.delete(key.name);
+          })
+        );
       }
       if (result.list_complete) {
         allTokensDeleted = true;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@cloudflare/workers-oauth-provider",
-  "version": "0.0.1",
+  "version": "0.0.3",
   "description": "OAuth provider for Cloudflare Workers",
   "main": "dist/oauth-provider.js",
   "types": "dist/oauth-provider.d.ts",
@@ -11,19 +11,23 @@
   ],
   "type": "module",
   "publishConfig": {
-    "access": "restricted"
+    "access": "public"
   },
   "scripts": {
     "build": "tsup",
-    "test": "vitest",
-    "prepublishOnly": "npm run build"
+    "build:watch": "tsup --watch",
+    "test": "vitest run",
+    "test:watch": "vitest",
+    "prepublishOnly": "npm run build",
+    "prettier": "prettier -w ."
   },
   "dependencies": {
     "@cloudflare/workers-types": "^4.20250311.0"
   },
   "devDependencies": {
+    "prettier": "^3.5.3",
     "tsup": "^8.4.0",
     "typescript": "^5.8.2",
-    "vitest": "^1.2.2"
+    "vitest": "^3.0.8"
   }
 }