feat: added helper methods for storing and retrieving "AuthRequest" objects in KV

.changeset/ripe-dragons-push.md

@@ -0,0 +1,9 @@
+---
+'@cloudflare/workers-oauth-provider': patch
+---
+
+feat: added helper methods for storing and retrieving "AuthRequest" objects in KV
+chore: fix prettier formatting
+fix: allow storing a superset of Partial<AuthRequest> in KV
+fix: added TTL support to storeAuthRequest
+fix: added helper to delete stored auth request

__tests__/oauth-provider.test.ts

@@ -1,5 +1,5 @@
 import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
-import { OAuthProvider, type OAuthHelpers } from '../src/oauth-provider';
+import { AuthRequest, OAuthProvider, type OAuthHelpers } from '../src/oauth-provider';
 import type { ExecutionContext } from '@cloudflare/workers-types';
 // We're importing WorkerEntrypoint from our mock implementation
 // The actual import is mocked in setup.ts
@@ -3471,4 +3471,125 @@ describe('OAuthProvider', () => {
       expect(newTokens.refresh_token).toBeDefined();
     });
   });
+
+  describe('Store and Retrieve AuthRequest', () => {
+    let clientId: string;
+    let clientSecret: string;
+    let redirectUri: string;
+
+    // Helper to create a test client before authorization tests
+    async function createTestClient() {
+      const clientData = {
+        redirect_uris: ['https://client.example.com/callback'],
+        client_name: 'Test Client',
+        token_endpoint_auth_method: 'client_secret_basic',
+      };
+
+      const request = createMockRequest(
+        'https://example.com/oauth/register',
+        'POST',
+        { 'Content-Type': 'application/json' },
+        JSON.stringify(clientData)
+      );
+
+      const response = await oauthProvider.fetch(request, mockEnv, mockCtx);
+      const client = await response.json<any>();
+
+      clientId = client.client_id;
+      clientSecret = client.client_secret;
+      redirectUri = 'https://client.example.com/callback';
+    }
+
+    beforeEach(async () => {
+      await createTestClient();
+    });
+
+    it('should store and retrieve an authorization request', async () => {
+      const authRequestId = 'test-auth-req-123';
+      const authRequestData: AuthRequest = {
+        responseType: 'code',
+        clientId: 'test-client',
+        redirectUri: 'https://client.example.com/callback',
+        scope: ['read', 'write'],
+        state: 'xyz123',
+      };
+
+      await oauthProvider.fetch(createMockRequest('https://example.com/'), mockEnv, mockCtx);
+      const helpers = mockEnv.OAUTH_PROVIDER!;
+
+      await helpers.storeAuthRequest(authRequestId, authRequestData);
+      const retrievedRequest = await helpers.getAuthRequest(authRequestId);
+
+      expect(retrievedRequest).toEqual(authRequestData);
+      const kvData = await mockEnv.OAUTH_KV.get(`authRequest:${authRequestId}`, { type: 'json' });
+      expect(kvData).toEqual(authRequestData);
+    });
+
+    it('getAuthRequest should return null for a non-existent request', async () => {
+      await oauthProvider.fetch(createMockRequest('https://example.com/'), mockEnv, mockCtx);
+      const helpers = mockEnv.OAUTH_PROVIDER!;
+
+      const retrievedRequest = await helpers.getAuthRequest('non-existent-id');
+      expect(retrievedRequest).toBeNull();
+    });
+
+    it('should expire a stored authorization request after the TTL', async () => {
+      const authRequestId = 'test-auth-req-ttl';
+      const authRequestData: AuthRequest = {
+        responseType: 'code',
+        clientId: 'test-client',
+        redirectUri: 'https://client.example.com/callback',
+        scope: ['read'],
+        state: 'ttl-test',
+      };
+
+      await oauthProvider.fetch(createMockRequest('https://example.com/'), mockEnv, mockCtx);
+      const helpers = mockEnv.OAUTH_PROVIDER!;
+
+      // Store with a 1-second TTL
+      await helpers.storeAuthRequest(authRequestId, authRequestData, { ttl: 1 });
+
+      // Should be retrievable immediately
+      const immediateRetrieve = await helpers.getAuthRequest(authRequestId);
+      expect(immediateRetrieve).toEqual(authRequestData);
+
+      // Wait for 2 seconds for the item to expire
+      await new Promise((resolve) => setTimeout(resolve, 2000));
+
+      // Should be null after expiration
+      const expiredRetrieve = await helpers.getAuthRequest(authRequestId);
+      expect(expiredRetrieve).toBeNull();
+
+      // Also check the underlying KV store to be sure
+      const kvData = await mockEnv.OAUTH_KV.get(`authRequest:${authRequestId}`);
+      expect(kvData).toBeNull();
+    });
+
+    it('should delete a stored authorization request', async () => {
+      const authRequestId = 'test-auth-request-456';
+      const authRequestData: AuthRequest = {
+        responseType: 'code',
+        clientId: 'test-client-2',
+        redirectUri: 'https://client.example.com/callback2',
+        scope: ['profile'],
+        state: 'abc456',
+      };
+      await oauthProvider.fetch(createMockRequest('https://example.com/'), mockEnv, mockCtx);
+      const helpers = mockEnv.OAUTH_PROVIDER!;
+
+      // Store the request first
+      await helpers.storeAuthRequest(authRequestId, authRequestData, { ttl: 600 });
+
+      // Verify it's stored
+      const storedRequest = await helpers.getAuthRequest(authRequestId);
+      expect(storedRequest).toEqual(authRequestData);
+
+      // Delete the request
+      await helpers.deleteAuthRequest(authRequestId);
+
+      // Verify it's deleted
+      const deletedRequest = await helpers.getAuthRequest(authRequestId);
+      expect(deletedRequest).toBeNull();
+    });
+  });
 });

src/oauth-provider.ts

@@ -279,6 +279,33 @@ export interface OAuthHelpers {
    */
   parseAuthRequest(request: Request): Promise<AuthRequest>;
 
+  /**
+   * Stores an OAuth authorization request in the internal KV store
+   * @param id - The ID used to identify the authorization request
+   * @param authRequest - The authorization request to store
+   * @param options - Options for the storage operation
+   * @returns A Promise resolving to void
+   */
+  storeAuthRequest<T extends Partial<AuthRequest>>(
+    id: string,
+    authRequest: T,
+    options?: { ttl?: number }
+  ): Promise<void>;
+
+  /**
+   * Looks up an OAuth authorization request by its ID
+   * @param id - The ID of the authorization request to look up
+   * @returns A Promise resolving to the authorization request, or null if not found
+   */
+  getAuthRequest<T extends Partial<AuthRequest>>(id: string): Promise<T | null>;
+
+  /**
+   * Deletes an OAuth authorization request by its ID
+   * @param id - The ID of the authorization request to delete
+   * @returns A Promise resolving to void
+   */
+  deleteAuthRequest(id: string): Promise<void>;
+
   /**
    * Looks up a client by its client ID
    * @param clientId - The client ID to look up
@@ -2523,6 +2550,42 @@ class OAuthHelpersImpl implements OAuthHelpers {
     };
   }
 
+  /**
+   * Stores the authorization request in the KV store
+   * @param id - The unique identifier for the request
+   * @param authRequest - The authorization request to store
+   * @param options - Options for the storage operation
+   * @returns A Promise resolving when the request is stored
+   */
+  async storeAuthRequest<T extends Partial<AuthRequest>>(
+    id: string,
+    authRequest: T,
+    options: { ttl?: number } = { ttl: 60 * 10 }
+  ): Promise<void> {
+    return await this.env.OAUTH_KV.put(`authRequest:${id}`, JSON.stringify(authRequest), {
+      expirationTtl: options?.ttl,
+    });
+  }
+
+  /**
+   * Retrieves an authorization request from the KV store
+   * @param id - The unique identifier for the request
+   * @returns A Promise resolving to the authorization request, or null if not found
+   */
+  async getAuthRequest<T extends Partial<AuthRequest>>(id: string): Promise<T | null> {
+    const authRequestStr = await this.env.OAUTH_KV.get(`authRequest:${id}`);
+    return authRequestStr ? JSON.parse(authRequestStr) : null;
+  }
+
+  /**
+   * Deletes an authorization request from the KV store
+   * @param id - The unique identifier for the request
+   * @returns A Promise resolving when the request is deleted
+   */
+  async deleteAuthRequest(id: string): Promise<void> {
+    return await this.env.OAUTH_KV.delete(`authRequest:${id}`);
+  }
+
   /**
    * Looks up a client by its client ID
    * @param clientId - The client ID to look up