feat: add storeStateStrategy (#5470)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

packages/better-auth/src/__snapshots__/init.test.ts.snap

@@ -111,6 +111,10 @@ exports[`init > should match config 1`] = `
     "warn": [Function],
   },
   "newSession": null,
+  "oauthConfig": {
+    "skipStateCookieCheck": false,
+    "storeStateStrategy": "database",
+  },
   "options": {
     "basePath": "/api/auth",
     "baseURL": "http://localhost:3000",

packages/better-auth/src/init.ts

@@ -101,6 +101,12 @@ export const init = async (options: BetterAuthOptions) => {
 		appName: options.appName || "Better Auth",
 		socialProviders: providers,
 		options,
+		oauthConfig: {
+			storeStateStrategy:
+				options.advanced?.oauthConfig?.storeStateStrategy || "database",
+			skipStateCookieCheck:
+				!!options.advanced?.oauthConfig?.skipStateCookieCheck,
+		},
 		tables,
 		trustedOrigins: getTrustedOrigins(options),
 		baseURL: baseURL || "",

packages/better-auth/src/oauth2/state.ts

@@ -1,6 +1,7 @@
 import * as z from "zod";
 import { APIError } from "better-call";
 import { generateRandomString } from "../crypto";
+import { symmetricEncrypt, symmetricDecrypt } from "../crypto";
 import type { GenericEndpointContext } from "@better-auth/core";
 
 export async function generateState(
@@ -19,16 +20,10 @@ export async function generateState(
 
 	const codeVerifier = generateRandomString(128);
 	const state = generateRandomString(32);
-	const stateCookie = c.context.createAuthCookie("state", {
-		maxAge: 5 * 60 * 1000, // 5 minutes
-	});
-	await c.setSignedCookie(
-		stateCookie.name,
-		state,
-		c.context.secret,
-		stateCookie.attributes,
-	);
-	const data = JSON.stringify({
+	const storeStateStrategy =
+		c.context.oauthConfig?.storeStateStrategy || "cookie";
+
+	const stateData = {
 		callbackURL,
 		codeVerifier,
 		errorURL: c.body?.errorCallbackURL,
@@ -39,11 +34,42 @@ export async function generateState(
 		 */
 		expiresAt: Date.now() + 10 * 60 * 1000,
 		requestSignUp: c.body?.requestSignUp,
+	};
+
+	if (storeStateStrategy === "cookie") {
+		// Store state data in an encrypted cookie
+		const encryptedData = await symmetricEncrypt({
+			key: c.context.secret,
+			data: JSON.stringify(stateData),
+		});
+
+		const stateCookie = c.context.createAuthCookie("oauth_state", {
+			maxAge: 10 * 60 * 1000, // 10 minutes
+		});
+
+		await c.setCookie(stateCookie.name, encryptedData, stateCookie.attributes);
+
+		return {
+			state,
+			codeVerifier,
+		};
+	}
+
+	// Default: database strategy
+	const stateCookie = c.context.createAuthCookie("state", {
+		maxAge: 5 * 60 * 1000, // 5 minutes
 	});
+	await c.setSignedCookie(
+		stateCookie.name,
+		state,
+		c.context.secret,
+		stateCookie.attributes,
+	);
+
 	const expiresAt = new Date();
 	expiresAt.setMinutes(expiresAt.getMinutes() + 10);
 	const verification = await c.context.internalAdapter.createVerificationValue({
-		value: data,
+		value: JSON.stringify(stateData),
 		identifier: state,
 		expiresAt,
 	});
@@ -63,64 +89,111 @@ export async function generateState(
 
 export async function parseState(c: GenericEndpointContext) {
 	const state = c.query.state || c.body.state;
-	const data = await c.context.internalAdapter.findVerificationValue(state);
-	if (!data) {
-		c.context.logger.error("State Mismatch. Verification not found", {
-			state,
-		});
-		const errorURL =
-			c.context.options.onAPIError?.errorURL || `${c.context.baseURL}/error`;
-		throw c.redirect(`${errorURL}?error=please_restart_the_process`);
-	}
+	const storeStateStrategy =
+		c.context.oauthConfig.storeStateStrategy || "cookie";
 
-	const parsedData = z
-		.object({
-			callbackURL: z.string(),
-			codeVerifier: z.string(),
-			errorURL: z.string().optional(),
-			newUserURL: z.string().optional(),
-			expiresAt: z.number(),
-			link: z
-				.object({
-					email: z.string(),
-					userId: z.coerce.string(),
-				})
-				.optional(),
-			requestSignUp: z.boolean().optional(),
-		})
-		.parse(JSON.parse(data.value));
+	const stateDataSchema = z.object({
+		callbackURL: z.string(),
+		codeVerifier: z.string(),
+		errorURL: z.string().optional(),
+		newUserURL: z.string().optional(),
+		expiresAt: z.number(),
+		link: z
+			.object({
+				email: z.string(),
+				userId: z.coerce.string(),
+			})
+			.optional(),
+		requestSignUp: z.boolean().optional(),
+	});
+
+	let parsedData: z.infer<typeof stateDataSchema>;
+
+	if (storeStateStrategy === "cookie") {
+		// Retrieve state data from encrypted cookie
+		const stateCookie = c.context.createAuthCookie("oauth_state");
+		const encryptedData = await c.getCookie(stateCookie.name);
+
+		if (!encryptedData) {
+			c.context.logger.error("State Mismatch. OAuth state cookie not found", {
+				state,
+			});
+			const errorURL =
+				c.context.options.onAPIError?.errorURL || `${c.context.baseURL}/error`;
+			throw c.redirect(`${errorURL}?error=please_restart_the_process`);
+		}
+
+		try {
+			const decryptedData = await symmetricDecrypt({
+				key: c.context.secret,
+				data: encryptedData,
+			});
+
+			parsedData = stateDataSchema.parse(JSON.parse(decryptedData));
+		} catch (error) {
+			c.context.logger.error("Failed to decrypt or parse OAuth state cookie", {
+				error,
+			});
+			const errorURL =
+				c.context.options.onAPIError?.errorURL || `${c.context.baseURL}/error`;
+			throw c.redirect(`${errorURL}?error=please_restart_the_process`);
+		}
+
+		// Clear the cookie after successful parsing
+		c.setCookie(stateCookie.name, "", {
+			maxAge: 0,
+		});
+	} else {
+		// Default: database strategy
+		const data = await c.context.internalAdapter.findVerificationValue(state);
+		if (!data) {
+			c.context.logger.error("State Mismatch. Verification not found", {
+				state,
+			});
+			const errorURL =
+				c.context.options.onAPIError?.errorURL || `${c.context.baseURL}/error`;
+			throw c.redirect(`${errorURL}?error=please_restart_the_process`);
+		}
+
+		parsedData = stateDataSchema.parse(JSON.parse(data.value));
+
+		const stateCookie = c.context.createAuthCookie("state");
+		const stateCookieValue = await c.getSignedCookie(
+			stateCookie.name,
+			c.context.secret,
+		);
+		/**
+		 * This is generally cause security issue and should only be used in
+		 * dev or staging environments. It's currently used by the oauth-proxy
+		 * plugin
+		 */
+		const skipStateCookieCheck = c.context.oauthConfig?.skipStateCookieCheck;
+		if (
+			!skipStateCookieCheck &&
+			(!stateCookieValue || stateCookieValue !== state)
+		) {
+			const errorURL =
+				c.context.options.onAPIError?.errorURL || `${c.context.baseURL}/error`;
+			throw c.redirect(`${errorURL}?error=state_mismatch`);
+		}
+		c.setCookie(stateCookie.name, "", {
+			maxAge: 0,
+		});
+
+		// Delete verification value after retrieval
+		await c.context.internalAdapter.deleteVerificationValue(data.id);
+	}
 
 	if (!parsedData.errorURL) {
 		parsedData.errorURL = `${c.context.baseURL}/error`;
 	}
-	const stateCookie = c.context.createAuthCookie("state");
-	const stateCookieValue = await c.getSignedCookie(
-		stateCookie.name,
-		c.context.secret,
-	);
-	/**
-	 * This is generally cause security issue and should only be used in
-	 * dev or staging environments. It's currently used by the oauth-proxy
-	 * plugin
-	 */
-	const skipStateCookieCheck = c.context.oauthConfig?.skipStateCookieCheck;
-	if (
-		!skipStateCookieCheck &&
-		(!stateCookieValue || stateCookieValue !== state)
-	) {
-		const errorURL =
-			c.context.options.onAPIError?.errorURL || `${c.context.baseURL}/error`;
-		throw c.redirect(`${errorURL}?error=state_mismatch`);
-	}
-	c.setCookie(stateCookie.name, "", {
-		maxAge: 0,
-	});
+
+	// Check expiration
 	if (parsedData.expiresAt < Date.now()) {
-		await c.context.internalAdapter.deleteVerificationValue(data.id);
 		const errorURL =
 			c.context.options.onAPIError?.errorURL || `${c.context.baseURL}/error`;
 		throw c.redirect(`${errorURL}?error=please_restart_the_process`);
 	}
-	await c.context.internalAdapter.deleteVerificationValue(data.id);
+
 	return parsedData;
 }

packages/better-auth/src/plugins/generic-oauth/generic-oauth.test.ts

@@ -892,4 +892,74 @@ describe("oauth2", async () => {
 			userId: session.data?.user.id,
 		});
 	});
+
+	it("should work with cookie-based state storage", async () => {
+		server.service.once("beforeUserinfo", (userInfoResponse) => {
+			userInfoResponse.body = {
+				email: "oauth2-cookie-state@test.com",
+				name: "OAuth2 Cookie State",
+				sub: "oauth2-cookie-state",
+				picture: "https://test.com/picture.png",
+				email_verified: true,
+			};
+			userInfoResponse.statusCode = 200;
+		});
+
+		const { customFetchImpl, auth, cookieSetter } = await getTestInstance({
+			plugins: [
+				genericOAuth({
+					config: [
+						{
+							providerId: "test-cookie",
+							discoveryUrl: `http://localhost:${port}/.well-known/openid-configuration`,
+							clientId: clientId,
+							clientSecret: clientSecret,
+							pkce: true,
+						},
+					],
+				}),
+			],
+			advanced: {
+				oauthConfig: {
+					storeStateStrategy: "cookie",
+				},
+			},
+		});
+		const headers = new Headers();
+		const authClient = createAuthClient({
+			plugins: [genericOAuthClient()],
+			baseURL: "http://localhost:3000",
+			fetchOptions: {
+				customFetchImpl,
+				onSuccess: cookieSetter(headers),
+			},
+		});
+
+		const res = await authClient.signIn.oauth2({
+			providerId: "test-cookie",
+			callbackURL: "http://localhost:3000/dashboard",
+			newUserCallbackURL: "http://localhost:3000/new_user",
+			fetchOptions: {
+				onSuccess: cookieSetter(headers),
+			},
+		});
+		expect(res.data?.url).toContain(`http://localhost:${port}/authorize`);
+
+		const { callbackURL, headers: newHeaders } = await simulateOAuthFlow(
+			res.data?.url || "",
+			headers,
+			customFetchImpl,
+		);
+		expect(callbackURL).toBe("http://localhost:3000/new_user");
+
+		const session = await authClient.getSession({
+			fetchOptions: {
+				headers: newHeaders,
+			},
+		});
+
+		expect(session.data).not.toBeNull();
+		expect(session.data?.user.email).toBe("oauth2-cookie-state@test.com");
+		expect(session.data?.user.name).toBe("OAuth2 Cookie State");
+	});
 });

packages/core/src/types/context.ts

@@ -166,11 +166,20 @@ export type AuthContext<Options extends BetterAuthOptions = BetterAuthOptions> =
 		appName: string;
 		baseURL: string;
 		trustedOrigins: string[];
-		oauthConfig?: {
+		oauthConfig: {
 			/**
 			 * This is dangerous and should only be used in dev or staging environments.
 			 */
 			skipStateCookieCheck?: boolean;
+			/**
+			 * Strategy for storing OAuth state
+			 *
+			 * - "cookie": Store state in an encrypted cookie (stateless)
+			 * - "database": Store state in the database
+			 *
+			 * @default "cookie"
+			 */
+			storeStateStrategy: "database" | "cookie";
 		};
 		/**
 		 * New session that will be set after the request

packages/core/src/types/init-options.ts

@@ -241,6 +241,15 @@ export type BetterAuthAdvancedOptions = {
 		 * @default false
 		 */
 		skipStateCookieCheck?: boolean;
+		/**
+		 * Strategy for storing OAuth state
+		 *
+		 * - "cookie": Store state in an encrypted cookie (stateless)
+		 * - "database": Store state in the database
+		 *
+		 * @default "cookie"
+		 */
+		storeStateStrategy?: "database" | "cookie";
 	};
 };