From caf0cf9cf7d3e18dbe88ad34a0f13f9cad6f230c Mon Sep 17 00:00:00 2001
From: Bereket Engida <86073083+Bekacru@users.noreply.github.com>
Date: Mon, 16 Dec 2024 16:53:01 +0300
Subject: [PATCH] fix: refersh cookie cache properly when it's expired (#909)

---
 demo/nextjs/lib/auth.ts                       | 16 +++++
 .../src/api/routes/session-api.test.ts        | 31 +++++++++-
 .../better-auth/src/api/routes/session.ts     |  8 ++-
 packages/better-auth/src/cookies/index.ts     | 61 +++++++++++--------
 4 files changed, 87 insertions(+), 29 deletions(-)

diff --git a/demo/nextjs/lib/auth.ts b/demo/nextjs/lib/auth.ts
index c23fbdf6..83c95db0 100644
--- a/demo/nextjs/lib/auth.ts
+++ b/demo/nextjs/lib/auth.ts
@@ -42,6 +42,22 @@ export const auth = betterAuth({
 		dialect,
 		type: process.env.USE_MYSQL ? "mysql" : "sqlite",
 	},
+	databaseHooks: {
+		user: {
+			update: {
+				async before(user) {
+					if (user.emailVerified) {
+						return {
+							data: {
+								...user,
+								emailVerifiedAt: new Date().toISOString(),
+							},
+						};
+					}
+				},
+			},
+		},
+	},
 	emailVerification: {
 		async sendVerificationEmail({ user, url }) {
 			const res = await resend.emails.send({
diff --git a/packages/better-auth/src/api/routes/session-api.test.ts b/packages/better-auth/src/api/routes/session-api.test.ts
index b778843f..530fb039 100644
--- a/packages/better-auth/src/api/routes/session-api.test.ts
+++ b/packages/better-auth/src/api/routes/session-api.test.ts
@@ -435,7 +435,7 @@ describe("cookie cache", async () => {
 	};
 	const adapter = memoryAdapter(database);
 
-	const { client, testUser, auth } = await getTestInstance({
+	const { client, testUser, auth, cookieSetter } = await getTestInstance({
 		database: adapter,
 		session: {
 			cookieCache: {
@@ -498,4 +498,33 @@ describe("cookie cache", async () => {
 		expect(session.data).not.toBeNull();
 		expect(fn).toHaveBeenCalledTimes(3);
 	});
+
+	it("should reset cache when expires", async () => {
+		expect(fn).toHaveBeenCalledTimes(3);
+		await client.getSession({
+			fetchOptions: {
+				headers,
+			},
+		});
+		vi.useFakeTimers();
+		await vi.advanceTimersByTimeAsync(1000 * 60 * 10); // 10 minutes
+		await client.getSession({
+			fetchOptions: {
+				headers,
+				onSuccess(context) {
+					cookieSetter(headers)(context);
+				},
+			},
+		});
+		expect(fn).toHaveBeenCalledTimes(5);
+		await client.getSession({
+			fetchOptions: {
+				headers,
+				onSuccess(context) {
+					cookieSetter(headers)(context);
+				},
+			},
+		});
+		expect(fn).toHaveBeenCalledTimes(5);
+	});
 });
diff --git a/packages/better-auth/src/api/routes/session.ts b/packages/better-auth/src/api/routes/session.ts
index 6fab5a76..0feb34bd 100644
--- a/packages/better-auth/src/api/routes/session.ts
+++ b/packages/better-auth/src/api/routes/session.ts
@@ -1,7 +1,11 @@
 import { APIError } from "better-call";
 import { createAuthEndpoint, createAuthMiddleware } from "../call";
 import { getDate } from "../../utils/date";
-import { deleteSessionCookie, setSessionCookie } from "../../cookies";
+import {
+	deleteSessionCookie,
+	setCookieCache,
+	setSessionCookie,
+} from "../../cookies";
 import { z } from "zod";
 import type {
 	BetterAuthOptions,
@@ -229,7 +233,7 @@ export const getSession = <Option extends BetterAuthOptions>() =>
 						user: InferUser<Option>;
 					});
 				}
-
+				await setCookieCache(ctx, session);
 				return ctx.json(
 					session as unknown as {
 						session: InferSession<Option>;
diff --git a/packages/better-auth/src/cookies/index.ts b/packages/better-auth/src/cookies/index.ts
index e4219690..30be42bf 100644
--- a/packages/better-auth/src/cookies/index.ts
+++ b/packages/better-auth/src/cookies/index.ts
@@ -92,38 +92,13 @@ export function getCookies(options: BetterAuthOptions) {
 
 export type BetterAuthCookies = ReturnType<typeof getCookies>;
 
-export async function setSessionCookie(
+export async function setCookieCache(
 	ctx: GenericEndpointContext,
 	session: {
 		session: Session & Record<string, any>;
 		user: User;
 	},
-	dontRememberMe?: boolean,
-	overrides?: Partial<CookieOptions>,
 ) {
-	const options = ctx.context.authCookies.sessionToken.options;
-	const maxAge = dontRememberMe
-		? undefined
-		: ctx.context.sessionConfig.expiresIn;
-	await ctx.setSignedCookie(
-		ctx.context.authCookies.sessionToken.name,
-		session.session.token,
-		ctx.context.secret,
-		{
-			...options,
-			maxAge,
-			...overrides,
-		},
-	);
-
-	if (dontRememberMe) {
-		await ctx.setSignedCookie(
-			ctx.context.authCookies.dontRememberToken.name,
-			"true",
-			ctx.context.secret,
-			ctx.context.authCookies.dontRememberToken.options,
-		);
-	}
 	const shouldStoreSessionDataInCookie =
 		ctx.context.options.session?.cookieCache?.enabled;
 
@@ -157,7 +132,41 @@ export async function setSessionCookie(
 			ctx.context.authCookies.sessionData.options,
 		);
 	}
+}
 
+export async function setSessionCookie(
+	ctx: GenericEndpointContext,
+	session: {
+		session: Session & Record<string, any>;
+		user: User;
+	},
+	dontRememberMe?: boolean,
+	overrides?: Partial<CookieOptions>,
+) {
+	const options = ctx.context.authCookies.sessionToken.options;
+	const maxAge = dontRememberMe
+		? undefined
+		: ctx.context.sessionConfig.expiresIn;
+	await ctx.setSignedCookie(
+		ctx.context.authCookies.sessionToken.name,
+		session.session.token,
+		ctx.context.secret,
+		{
+			...options,
+			maxAge,
+			...overrides,
+		},
+	);
+
+	if (dontRememberMe) {
+		await ctx.setSignedCookie(
+			ctx.context.authCookies.dontRememberToken.name,
+			"true",
+			ctx.context.secret,
+			ctx.context.authCookies.dontRememberToken.options,
+		);
+	}
+	await setCookieCache(ctx, session);
 	ctx.context.setNewSession(session);
 	/**
 	 * If secondary storage is enabled, store the session data in the secondary storage