refactor(oauth): simplify state management and enhance state generation for OAuth

Author: OXeu

server/src/services/user.ts

@@ -6,26 +6,32 @@ import base from "../base";
 export const UserService = base()
     .group('/user', (group) =>
         group
-            .get("/github", ({ oauth2, set, headers: { referer }, cookie: { redirect_to } }) => {
+            .get("/github", ({ oauth2, set, headers: { referer }, cookie: { redirect_to, state } }) => {
                 if (!referer) {
                     return 'Referer not found'
                 }
                 redirect_to.value = `${referer}`
-                set.headers['Location'] = oauth2.createRedirectUrl("GitHub")
+                const genState = oauth2.generateState();
+                // Store state in cookie to verify later
+                state.value = genState;
+                set.headers['Location'] = oauth2.createRedirectUrl(genState, "GitHub")
                 set.status = 302
             })
-            .get("/github/callback", async ({ jwt, oauth2, set, store, query, cookie: { token, redirect_to } }) => {
+            .get("/github/callback", async ({ jwt, oauth2, set, store, query, cookie: { token, redirect_to, state } }) => {
                 const { db } = store;
 
-                console.log('p_state', query.state)
+                console.log('param_state', query.state)
+                console.log('cookie_state', state.value)
 
                 // Verify state to prevent CSRF attacks
-                if (!oauth2.verifyState(query.state)) {
+                if (query.state != state.value) {
                     set.status = 400;
                     return 'Invalid state parameter';
                 }
+                // Clear state cookie
+                state.value = '';
 
-                oauth2.removeState(query.state);
+                // Exchange code for access token
                 const gh_token = await oauth2.authorize("GitHub", query.code);
                 if (!gh_token) {
                     set.status = 400;

server/src/utils/oauth.ts

@@ -49,49 +49,32 @@ export interface OAuth2Methods {
 }
 
 export function createOAuthPlugin(providers: Record<string, OAuthProvider>) {
-    const stateStore = new Map<string, { provider: string; timestamp: number }>();
-
-    // Clean up expired states (older than 10 minutes)
-    const cleanupExpiredStates = () => {
-        const now = Date.now();
-        const expiryTime = 10 * 60 * 1000; // 10 minutes
-        for (const [state, data] of stateStore.entries()) {
-            if (now - data.timestamp > expiryTime) {
-                stateStore.delete(state);
-            }
-        }
-    };
 
     // Generate random state string
-    const generateState = (): string => {
-        const array = new Uint8Array(32);
-        crypto.getRandomValues(array);
-        return Array.from(array, (byte) => byte.toString(16).padStart(2, "0")).join("");
-    };
+
 
     return new Elysia({ name: "oauth2" })
         .decorate("oauth2", {
-            createRedirectUrl: (providerName: string): string => {
+            generateState: () => {
+                const array = new Uint8Array(32);
+                crypto.getRandomValues(array);
+                return Array.from(array, (byte) => byte.toString(16).padStart(2, "0")).join("");
+            },
+            createRedirectUrl: (state: string, providerName: string): string => {
                 const provider = providers[providerName];
                 if (!provider) {
                     throw new Error(`OAuth provider "${providerName}" not found`);
                 }
 
-                cleanupExpiredStates();
 
-                const state = generateState();
-                stateStore.set(state, {
-                    provider: providerName,
-                    timestamp: Date.now(),
-                });
 
                 const params = new URLSearchParams({
                     client_id: provider.clientId,
                     state: state,
                 });
 
                 // if (provider.redirectUri) {
-                    // params.set("redirect_uri", provider.redirectUri);
+                // params.set("redirect_uri", provider.redirectUri);
                 // }
 
                 return `${provider.authorizeUrl}?${params.toString()}`;
@@ -146,18 +129,5 @@ export function createOAuthPlugin(providers: Record<string, OAuthProvider>) {
                     refreshToken: data.refresh_token,
                 };
             },
-
-            verifyState: (state: string): boolean => {
-                cleanupExpiredStates();
-                return stateStore.has(state);
-            },
-
-            getStateData: (state: string) => {
-                return stateStore.get(state);
-            },
-
-            removeState: (state: string) => {
-                stateStore.delete(state);
-            },
         });
 }