From 05997d72e8e32e76234075311e7ed1a95efd6436 Mon Sep 17 00:00:00 2001
From: Bereket Engida <bekacru@gmail.com>
Date: Tue, 19 Nov 2024 19:39:18 +0300
Subject: [PATCH] fix: add existing session metadata on enable and disable 2fa

---
 demo/nextjs/lib/auth.ts                       |  7 ------
 .../better-auth/src/db/internal-adapter.ts    |  4 ++--
 .../src/plugins/two-factor/index.ts           | 24 +++++++++++++++++++
 .../src/plugins/two-factor/totp/index.ts      |  5 ++++
 .../src/plugins/two-factor/two-factor.test.ts | 16 ++++++++++---
 5 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/demo/nextjs/lib/auth.ts b/demo/nextjs/lib/auth.ts
index 2b3f3b7a1c..6406d20a59 100644
--- a/demo/nextjs/lib/auth.ts
+++ b/demo/nextjs/lib/auth.ts
@@ -16,7 +16,6 @@ import { resend } from "./email/resend";
 import { MysqlDialect } from "kysely";
 import { createPool } from "mysql2/promise";
 import { nextCookies } from "better-auth/next-js";
-import * as ac from "./access-control";
 
 const from = process.env.BETTER_AUTH_EMAIL || "delivered@resend.dev";
 const to = process.env.TEST_EMAIL || "";
@@ -112,12 +111,6 @@ export const auth = betterAuth({
 	},
 	plugins: [
 		organization({
-			ac: ac.ac,
-			roles: {
-				admin: ac.admin,
-				owner: ac.owner,
-				member: ac.member,
-			},
 			async sendInvitationEmail(data) {
 				const res = await resend.emails.send({
 					from,
diff --git a/packages/better-auth/src/db/internal-adapter.ts b/packages/better-auth/src/db/internal-adapter.ts
index 2a8c4187e3..5da0b04a96 100644
--- a/packages/better-auth/src/db/internal-adapter.ts
+++ b/packages/better-auth/src/db/internal-adapter.ts
@@ -185,8 +185,6 @@ export const createInternalAdapter = (
 		) => {
 			const headers = request instanceof Request ? request.headers : request;
 			const data: Session = {
-				id: generateId(32),
-				userId,
 				/**
 				 * If the user doesn't want to be remembered
 				 * set the session to expire in 1 day.
@@ -198,6 +196,8 @@ export const createInternalAdapter = (
 				ipAddress: request ? getIp(request, ctx.options) || "" : "",
 				userAgent: headers?.get("user-agent") || "",
 				...override,
+				id: generateId(32),
+				userId,
 			};
 			const res = await createWithHooks(
 				data,
diff --git a/packages/better-auth/src/plugins/two-factor/index.ts b/packages/better-auth/src/plugins/two-factor/index.ts
index f51d7af2b1..2414143078 100644
--- a/packages/better-auth/src/plugins/two-factor/index.ts
+++ b/packages/better-auth/src/plugins/two-factor/index.ts
@@ -90,6 +90,8 @@ export const twoFactor = (options?: TwoFactorOptions) => {
 						const newSession = await ctx.context.internalAdapter.createSession(
 							updatedUser.id,
 							ctx.request,
+							false,
+							ctx.context.session.session,
 						);
 						/**
 						 * Update the session cookie with the new user data
@@ -98,6 +100,11 @@ export const twoFactor = (options?: TwoFactorOptions) => {
 							session: newSession,
 							user,
 						});
+
+						//remove current session
+						await ctx.context.internalAdapter.deleteSession(
+							ctx.context.session.session.id,
+						);
 					}
 					//delete existing two factor
 					await ctx.context.adapter.deleteMany({
@@ -164,6 +171,23 @@ export const twoFactor = (options?: TwoFactorOptions) => {
 							},
 						],
 					});
+					const newSession = await ctx.context.internalAdapter.createSession(
+						user.id,
+						ctx.request,
+						false,
+						ctx.context.session.session,
+					);
+					/**
+					 * Update the session cookie with the new user data
+					 */
+					await setSessionCookie(ctx, {
+						session: newSession,
+						user,
+					});
+					//remove current session
+					await ctx.context.internalAdapter.deleteSession(
+						ctx.context.session.session.id,
+					);
 					return ctx.json({ status: true });
 				},
 			),
diff --git a/packages/better-auth/src/plugins/two-factor/totp/index.ts b/packages/better-auth/src/plugins/two-factor/totp/index.ts
index 8ded2c927f..67eeb6d4ac 100644
--- a/packages/better-auth/src/plugins/two-factor/totp/index.ts
+++ b/packages/better-auth/src/plugins/two-factor/totp/index.ts
@@ -179,6 +179,11 @@ export const totp2fa = (options: TOTPOptions, twoFactorTable: string) => {
 				const newSession = await ctx.context.internalAdapter.createSession(
 					user.id,
 					ctx.request,
+					false,
+					ctx.context.session.session,
+				);
+				await ctx.context.internalAdapter.deleteSession(
+					ctx.context.session.session.id,
 				);
 				await setSessionCookie(ctx, {
 					session: newSession,
diff --git a/packages/better-auth/src/plugins/two-factor/two-factor.test.ts b/packages/better-auth/src/plugins/two-factor/two-factor.test.ts
index 1570bb826b..b03eb9487c 100644
--- a/packages/better-auth/src/plugins/two-factor/two-factor.test.ts
+++ b/packages/better-auth/src/plugins/two-factor/two-factor.test.ts
@@ -108,6 +108,7 @@ describe("two factor", async () => {
 			code,
 			fetchOptions: {
 				headers,
+				onSuccess: sessionSetter(headers),
 			},
 		});
 		expect(res.data?.session).toBeDefined();
@@ -356,10 +357,17 @@ describe("two factor auth api", async () => {
 				password: testUser.password,
 			},
 			headers,
+			asResponse: true,
 		});
-		expect(res.backupCodes.length).toBe(10);
-		expect(res.totpURI).toBeDefined();
+		headers = convertSetCookieToCookie(res.headers);
 
+		const json = (await res.json()) as {
+			status: boolean;
+			backupCodes: string[];
+			totpURI: string;
+		};
+		expect(json.backupCodes.length).toBe(10);
+		expect(json.totpURI).toBeDefined();
 		const session = await auth.api.getSession({
 			headers,
 		});
@@ -428,8 +436,10 @@ describe("two factor auth api", async () => {
 			body: {
 				password: testUser.password,
 			},
+			asResponse: true,
 		});
-		expect(res.status).toBe(true);
+		headers = convertSetCookieToCookie(res.headers);
+		expect(res.status).toBe(200);
 		const session = await auth.api.getSession({
 			headers,
 		});