🐛 Fix infinite loop when trying to migrate legacy cookie

apps/web/.env.test

@@ -1,5 +1,6 @@
 PORT=3002
 NEXT_PUBLIC_BASE_URL=http://localhost:3002
+AUTH_URL=http://localhost:3002
 SECRET_PASSWORD=abcdef1234567890abcdef1234567890
 DATABASE_URL=postgres://postgres:postgres@localhost:5450/rallly
 [email protected]

apps/web/src/auth/legacy/helpers/jwt.ts

@@ -1,7 +1,6 @@
 import hkdf from "@panva/hkdf";
 import { jwtDecrypt } from "jose";
-
-import type { JWT } from "./types";
+import type { JWT } from "next-auth/jwt";
 
 /** Decodes a NextAuth.js issued JWT. */
 export async function decodeLegacyJWT(token: string): Promise<JWT | null> {

apps/web/src/auth/legacy/next-auth-cookie-migration.ts

@@ -1,65 +1,68 @@
-import type { NextRequest } from "next/server";
-import { NextResponse } from "next/server";
+import { absoluteUrl } from "@rallly/utils/absolute-url";
+import { cookies } from "next/headers";
+import type { NextRequest, NextResponse } from "next/server";
+import type { Session } from "next-auth";
 import { encode } from "next-auth/jwt";
 
 import { decodeLegacyJWT } from "./helpers/jwt";
 
-const isSecureCookie =
-  process.env.NEXT_PUBLIC_BASE_URL?.startsWith("https://") ?? false;
+const isSecureCookie = absoluteUrl().startsWith("https://");
 
 const prefix = isSecureCookie ? "__Secure-" : "";
 
 const oldCookieName = prefix + "next-auth.session-token";
 const newCookieName = prefix + "authjs.session-token";
 
-/**
- * Migrates the next-auth cookies to the new authjs cookie names
- * This is needed for next-auth v5 which renamed the cookie prefix from 'next-auth' to 'authjs'
- */
-export function withAuthMigration(
-  middleware: (req: NextRequest) => void | Response | Promise<void | Response>,
-) {
-  return async (req: NextRequest) => {
-    const oldCookie = req.cookies.get(oldCookieName);
+export async function getLegacySession(): Promise<Session | null> {
+  const cookieStore = cookies();
+  const legacySessionCookie = cookieStore.get(oldCookieName);
+  if (legacySessionCookie) {
+    const decodedCookie = await decodeLegacyJWT(legacySessionCookie.value);
 
-    // If the old cookie doesn't exist, return the middleware
-    if (!oldCookie) {
-      return middleware(req);
+    if (decodedCookie?.sub) {
+      const { sub: id, ...rest } = decodedCookie;
+      return {
+        user: { id, ...rest },
+        expires: new Date(decodedCookie.exp ?? 0).toISOString(),
+      };
     }
+  }
 
-    const response = NextResponse.redirect(req.url);
-    response.cookies.delete(oldCookieName);
+  return null;
+}
 
-    // If the new cookie exists, delete the old cookie first and rerun middleware
-    if (req.cookies.get(newCookieName)) {
-      return response;
+async function getLegacyJWT() {
+  const cookieStore = cookies();
+  const legacySessionCookie = cookieStore.get(oldCookieName);
+  if (legacySessionCookie) {
+    const decodedCookie = await decodeLegacyJWT(legacySessionCookie.value);
+    if (decodedCookie) {
+      return decodedCookie;
     }
+  }
+  return null;
+}
 
-    const decodedCookie = await decodeLegacyJWT(oldCookie.value);
-
-    // If old cookie is invalid, delete the old cookie first and rerun middleware
-    if (!decodedCookie) {
-      return response;
-    }
+/**
+ * Replace the old legacy cookie with the new one
+ */
+export async function migrateLegacyJWT(res: NextResponse) {
+  const legacyJWT = await getLegacyJWT();
 
-    // Set the new cookie
-    const encodedCookie = await encode({
-      token: decodedCookie,
+  if (legacyJWT) {
+    const newJWT = await encode({
+      token: legacyJWT,
       secret: process.env.SECRET_PASSWORD,
       salt: newCookieName,
     });
 
-    // Set the new cookie with the same value and attributes
-    response.cookies.set(newCookieName, encodedCookie, {
-      path: "/",
+    res.cookies.set(newCookieName, newJWT, {
+      httpOnly: true,
       secure: isSecureCookie,
+      expires: new Date(Date.now() + 1000 * 60 * 60 * 24 * 7),
       sameSite: "lax",
-      httpOnly: true,
+      path: "/",
     });
-
-    // Delete the old cookie
-    response.cookies.delete(oldCookieName);
-
-    return response;
-  };
+    res.cookies.delete(oldCookieName);
+  }
 }

apps/web/src/auth/middleware.ts

@@ -1,15 +1,42 @@
-import type { NextRequest } from "next/server";
-import type { NextAuthRequest } from "next-auth";
+import { cookies } from "next/headers";
+import type { NextAuthRequest, Session } from "next-auth";
 import NextAuth from "next-auth";
 
 import { nextAuthConfig } from "@/next-auth.config";
 
+import {
+  getLegacySession,
+  migrateLegacyJWT,
+} from "./legacy/next-auth-cookie-migration";
+import type { NextResponse } from "next/server";
+
 const { auth } = NextAuth(nextAuthConfig);
 
-export function withAuth(
-  middleware: (
-    req: NextAuthRequest,
-  ) => void | Response | Promise<void | Response>,
-): (req: NextRequest) => void | Response | Promise<void | Response> {
-  return (req: NextRequest) => auth(middleware)(req, undefined as never);
-}
+export const withAuth = (
+  middleware: (request: NextAuthRequest) => Promise<NextResponse>,
+) => {
+  return async (request: NextAuthRequest) => {
+    const legacySession = await getLegacySession();
+
+    const session = legacySession || (await auth());
+
+    const res = await nextAuthConfig.callbacks.authorized({
+      request,
+      auth: session,
+    });
+
+    request.auth = session;
+
+    if (res !== true) {
+      return res;
+    }
+
+    const middlewareRes = await middleware(request);
+
+    if (legacySession) {
+      await migrateLegacyJWT(middlewareRes);
+    }
+
+    return middlewareRes;
+  };
+};

apps/web/src/middleware.ts

@@ -3,44 +3,41 @@ import { withPostHog } from "@rallly/posthog/next/middleware";
 import { NextResponse } from "next/server";
 
 import { getLocaleFromHeader } from "@/app/guest";
-import { withAuthMigration } from "@/auth/legacy/next-auth-cookie-migration";
 import { withAuth } from "@/auth/middleware";
 
 const supportedLocales = Object.keys(languages);
 
-export const middleware = withAuthMigration(
-  withAuth(async (req) => {
-    const { nextUrl } = req;
-    const newUrl = nextUrl.clone();
-
-    const isLoggedIn = req.auth?.user?.email;
-
-    // if the user is already logged in, don't let them access the login page
-    if (/^\/(login)/.test(newUrl.pathname) && isLoggedIn) {
-      newUrl.pathname = "/";
-      return NextResponse.redirect(newUrl);
-    }
-
-    // Check if locale is specified in cookie
-    let locale = req.auth?.user?.locale;
-    if (locale && supportedLocales.includes(locale)) {
-      newUrl.pathname = `/${locale}${newUrl.pathname}`;
-    } else {
-      // Check if locale is specified in header
-      locale = await getLocaleFromHeader(req);
-      newUrl.pathname = `/${locale}${newUrl.pathname}`;
-    }
-
-    const res = NextResponse.rewrite(newUrl);
-    res.headers.set("x-pathname", newUrl.pathname);
-
-    if (req.auth?.user?.id) {
-      await withPostHog(res, { distinctID: req.auth.user.id });
-    }
-
-    return res;
-  }),
-);
+export const middleware = withAuth(async (req) => {
+  const { nextUrl } = req;
+  const newUrl = nextUrl.clone();
+
+  const isLoggedIn = req.auth?.user?.email;
+
+  // if the user is already logged in, don't let them access the login page
+  if (/^\/(login)/.test(newUrl.pathname) && isLoggedIn) {
+    newUrl.pathname = "/";
+    return NextResponse.redirect(newUrl);
+  }
+
+  // Check if locale is specified in cookie
+  let locale = req.auth?.user?.locale;
+  if (locale && supportedLocales.includes(locale)) {
+    newUrl.pathname = `/${locale}${newUrl.pathname}`;
+  } else {
+    // Check if locale is specified in header
+    locale = await getLocaleFromHeader(req);
+    newUrl.pathname = `/${locale}${newUrl.pathname}`;
+  }
+
+  const res = NextResponse.rewrite(newUrl);
+  res.headers.set("x-pathname", newUrl.pathname);
+
+  if (req.auth?.user?.id) {
+    await withPostHog(res, { distinctID: req.auth.user.id });
+  }
+
+  return res;
+});
 
 export const config = {
   matcher: ["/((?!api|_next/static|_next/image|static|.*\\.).*)"],

apps/web/src/next-auth.ts

@@ -7,6 +7,7 @@ import z from "zod";
 import { CustomPrismaAdapter } from "./auth/adapters/prisma";
 import { isEmailBlocked } from "./auth/helpers/is-email-blocked";
 import { mergeGuestsIntoUser } from "./auth/helpers/merge-user";
+import { getLegacySession } from "./auth/legacy/next-auth-cookie-migration";
 import { EmailProvider } from "./auth/providers/email";
 import { GoogleProvider } from "./auth/providers/google";
 import { GuestProvider } from "./auth/providers/guest";
@@ -22,7 +23,12 @@ const sessionUpdateSchema = z.object({
   weekStart: z.number().nullish(),
 });
 
-export const { auth, handlers, signIn, signOut } = NextAuth({
+const {
+  auth: originalAuth,
+  handlers,
+  signIn,
+  signOut,
+} = NextAuth({
   ...nextAuthConfig,
   adapter: CustomPrismaAdapter({
     migrateData: async (userId) => {
@@ -169,3 +175,14 @@ export const { auth, handlers, signIn, signOut } = NextAuth({
     },
   },
 });
+
+const auth = async () => {
+  const session = await getLegacySession();
+  if (session) {
+    return session;
+  }
+
+  return originalAuth();
+};
+
+export { auth, handlers, signIn, signOut };

apps/web/src/utils/session/session-config.ts

@@ -1,8 +1,10 @@
+import { absoluteUrl } from "@rallly/utils/absolute-url";
+
 export const sessionConfig = {
   password: process.env.SECRET_PASSWORD ?? "",
   cookieName: "rallly-session",
   cookieOptions: {
-    secure: process.env.NEXT_PUBLIC_BASE_URL?.startsWith("https://") ?? false,
+    secure: absoluteUrl().startsWith("https://") ?? false,
   },
   ttl: 60 * 60 * 24 * 30, // 30 days
 };

apps/web/tests/helpers/next-auth-v4.ts

@@ -0,0 +1,36 @@
+import hkdf from "@panva/hkdf";
+import { EncryptJWT } from "jose";
+import type { JWT } from "next-auth/jwt";
+
+const now = () => (Date.now() / 1000) | 0;
+export async function getDerivedEncryptionKey(
+  keyMaterial: string | Buffer,
+  salt: string,
+) {
+  return await hkdf(
+    "sha256",
+    keyMaterial,
+    salt,
+    `NextAuth.js Generated Encryption Key${salt ? ` (${salt})` : ""}`,
+    32,
+  );
+}
+
+interface JWTEncodeParams {
+  token?: JWT;
+  salt?: string;
+  secret: string | Buffer;
+  maxAge?: number;
+}
+
+export async function encode(params: JWTEncodeParams) {
+  /** @note empty `salt` means a session token. See {@link JWTEncodeParams.salt}. */
+  const { token = {}, secret, maxAge = 30 * 24 * 60 * 60, salt = "" } = params;
+  const encryptionSecret = await getDerivedEncryptionKey(secret, salt);
+  return await new EncryptJWT(token)
+    .setProtectedHeader({ alg: "dir", enc: "A256GCM" })
+    .setIssuedAt()
+    .setExpirationTime(now() + maxAge)
+    .setJti("some-random-id")
+    .encrypt(encryptionSecret);
+}