From 6979c091f22f89fb6fb8ffb9a3c85a89b0678cb2 Mon Sep 17 00:00:00 2001
From: Luke Vella <me@lukevella.com>
Date: Thu, 13 Feb 2025 10:58:25 +0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20Check=20for=20new=20session=20fi?=
 =?UTF-8?q?rst=20(#1563)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/web/src/auth/edge/with-auth.ts           | 13 +++--
 .../auth/legacy/next-auth-cookie-migration.ts | 51 +++++++++++++------
 apps/web/src/next-auth.ts                     |  4 +-
 3 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/apps/web/src/auth/edge/with-auth.ts b/apps/web/src/auth/edge/with-auth.ts
index 226c7f5d1..2ca2a771f 100644
--- a/apps/web/src/auth/edge/with-auth.ts
+++ b/apps/web/src/auth/edge/with-auth.ts
@@ -15,19 +15,22 @@ export const withAuth = (
   middleware: (request: NextAuthRequest) => Promise<NextResponse>,
 ) => {
   return async (request: NextAuthRequest) => {
-    let legacySession: Session | null = null;
+    let session: Session | null = null;
 
     try {
-      legacySession = await getLegacySession();
+      session = await auth();
     } catch (e) {
       console.error(e);
     }
 
-    let session = legacySession;
+    let isLegacySession = false;
 
     if (!session) {
       try {
-        session = await auth();
+        session = await getLegacySession();
+        if (session) {
+          isLegacySession = true;
+        }
       } catch (e) {
         console.error(e);
       }
@@ -50,7 +53,7 @@ export const withAuth = (
 
     const middlewareRes = await middleware(request);
 
-    if (legacySession) {
+    if (isLegacySession) {
       try {
         await migrateLegacyJWT(middlewareRes);
       } catch (e) {
diff --git a/apps/web/src/auth/legacy/next-auth-cookie-migration.ts b/apps/web/src/auth/legacy/next-auth-cookie-migration.ts
index 0faac5613..0f41b4451 100644
--- a/apps/web/src/auth/legacy/next-auth-cookie-migration.ts
+++ b/apps/web/src/auth/legacy/next-auth-cookie-migration.ts
@@ -2,6 +2,7 @@ import { absoluteUrl } from "@rallly/utils/absolute-url";
 import { cookies } from "next/headers";
 import type { NextResponse } from "next/server";
 import type { Session } from "next-auth";
+import type { JWT } from "next-auth/jwt";
 import { encode } from "next-auth/jwt";
 
 import { decodeLegacyJWT } from "./helpers/jwt";
@@ -16,7 +17,7 @@ const newCookieName = prefix + "authjs.session-token";
 export async function getLegacySession(): Promise<Session | null> {
   const cookieStore = cookies();
   const legacySessionCookie = cookieStore.get(oldCookieName);
-  if (legacySessionCookie) {
+  if (legacySessionCookie && legacySessionCookie.value) {
     const decodedCookie = await decodeLegacyJWT(legacySessionCookie.value);
 
     if (decodedCookie?.sub) {
@@ -45,6 +46,38 @@ async function getLegacyJWT() {
   return null;
 }
 
+function deleteLegacyCookie(res: NextResponse) {
+  const cookieStore = cookies();
+  const oldCookie = cookieStore.get(oldCookieName);
+  if (oldCookie) {
+    // Delete the old cookie
+    res.cookies.set(oldCookieName, oldCookie.value, {
+      httpOnly: true,
+      secure: isSecureCookie,
+      expires: new Date(0),
+      sameSite: "lax",
+      path: "/",
+    });
+  }
+}
+
+async function setNewSessionCookie(res: NextResponse, jwt: JWT) {
+  const newJWT = await encode({
+    token: jwt,
+    secret: process.env.SECRET_PASSWORD,
+    salt: newCookieName,
+  });
+
+  // Set new session cookie
+  res.cookies.set(newCookieName, newJWT, {
+    httpOnly: true,
+    secure: isSecureCookie,
+    expires: new Date(Date.now() + 1000 * 60 * 60 * 24 * 7),
+    sameSite: "lax",
+    path: "/",
+  });
+}
+
 /**
  * Replace the old legacy cookie with the new one
  */
@@ -52,19 +85,7 @@ export async function migrateLegacyJWT(res: NextResponse) {
   const legacyJWT = await getLegacyJWT();
 
   if (legacyJWT) {
-    const newJWT = await encode({
-      token: legacyJWT,
-      secret: process.env.SECRET_PASSWORD,
-      salt: newCookieName,
-    });
-
-    res.cookies.set(newCookieName, newJWT, {
-      httpOnly: true,
-      secure: isSecureCookie,
-      expires: new Date(Date.now() + 1000 * 60 * 60 * 24 * 7),
-      sameSite: "lax",
-      path: "/",
-    });
-    res.cookies.delete(oldCookieName);
+    await setNewSessionCookie(res, legacyJWT);
+    deleteLegacyCookie(res);
   }
 }
diff --git a/apps/web/src/next-auth.ts b/apps/web/src/next-auth.ts
index 9f34fa13b..60b17daf4 100644
--- a/apps/web/src/next-auth.ts
+++ b/apps/web/src/next-auth.ts
@@ -177,12 +177,12 @@ const {
 });
 
 const auth = async () => {
-  const session = await getLegacySession();
+  const session = await originalAuth();
   if (session) {
     return session;
   }
 
-  return originalAuth();
+  return getLegacySession();
 };
 
 export { auth, handlers, signIn, signOut };