♻️ Fix cookie migration script

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/helpers/types.ts

@@ -1,46 +0,0 @@
-export interface DefaultJWT extends Record<string, unknown> {
-  name?: string | null;
-  email?: string | null;
-  picture?: string | null;
-  sub?: string;
-}
-
-/**
- * Returned by the `jwt` callback and `getToken`, when using JWT sessions
- *
- * [`jwt` callback](https://next-auth.js.org/configuration/callbacks#jwt-callback) | [`getToken`](https://next-auth.js.org/tutorials/securing-pages-and-api-routes#using-gettoken)
- */
-export interface JWT extends Record<string, unknown>, DefaultJWT {}
-
-export interface JWTEncodeParams {
-  /** The JWT payload. */
-  token?: JWT;
-  /**
-   * Used in combination with `secret` when deriving the encryption secret for the various NextAuth.js-issued JWTs.
-   * @note When no `salt` is passed, we assume this is a session token.
-   * This is for backwards-compatibility with currently active sessions, so they won't be invalidated when upgrading the package.
-   */
-  salt?: string;
-  /** The key material used to encode the NextAuth.js issued JWTs. Defaults to `NEXTAUTH_SECRET`. */
-  secret: string | Buffer;
-  /**
-   * The maximum age of the NextAuth.js issued JWT in seconds.
-   * @default 30 * 24 * 60 * 60 // 30 days
-   */
-  maxAge?: number;
-}
-
-export interface JWTDecodeParams {
-  /** The NextAuth.js issued JWT to be decoded */
-  token?: string;
-  /**
-   * Used in combination with `secret` when deriving the encryption secret for the various NextAuth.js-issued JWTs.
-   * @note When no `salt` is passed, we assume this is a session token.
-   * This is for backwards-compatibility with currently active sessions, so they won't be invalidated when upgrading the package.
-   */
-  salt?: string;
-  /** The key material used to decode the NextAuth.js issued JWTs. Defaults to `NEXTAUTH_SECRET`. */
-  secret: string | Buffer;
-}
-
-export type Secret = string | Buffer;

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

@@ -1,5 +1,7 @@
 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";
@@ -11,6 +13,60 @@ const prefix = isSecureCookie ? "__Secure-" : "";
 const oldCookieName = prefix + "next-auth.session-token";
 const newCookieName = prefix + "authjs.session-token";
 
+export async function getLegacySession(): Promise<Session | null> {
+  const cookieStore = cookies();
+  const legacySessionCookie = cookieStore.get(oldCookieName);
+  if (legacySessionCookie) {
+    const decodedCookie = await decodeLegacyJWT(legacySessionCookie.value);
+
+    if (decodedCookie?.sub) {
+      const { sub: id, ...rest } = decodedCookie;
+      return {
+        user: { id, ...rest },
+        expires: new Date(decodedCookie.exp ?? 0).toISOString(),
+      };
+    }
+  }
+
+  return null;
+}
+
+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;
+}
+
+/**
+ * Replace the old legacy cookie with the new one
+ */
+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);
+  }
+}
+
 /**
  * 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'

apps/web/src/auth/middleware.ts

@@ -1,25 +1,42 @@
-import type { NextRequest } from "next/server";
-import { NextResponse } 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) => Promise<NextResponse>,
-): (req: NextRequest) => Promise<NextResponse> {
-  return async (req: NextRequest) => {
-    const res = await auth(middleware)(req, undefined as never);
-    if (res) {
-      return new NextResponse(res.body, {
-        status: res.status,
-        headers: res.headers,
-        url: res.url,
-        statusText: res.statusText,
-      });
+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;
     }
-    return NextResponse.next();
+
+    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();
+export const middleware = withAuth(async (req) => {
+  const { nextUrl } = req;
+  const newUrl = nextUrl.clone();
 
-    const isLoggedIn = req.auth?.user?.email;
+  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);
-    }
+  // 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}`;
-    }
+  // 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);
+  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 });
-    }
+  if (req.auth?.user?.id) {
+    await withPostHog(res, { distinctID: req.auth.user.id });
+  }
 
-    return res;
-  }),
-);
+  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/tests/next-auth-migration.spec.ts

@@ -48,7 +48,8 @@ test.describe.serial(() => {
       },
     ]);
 
-    await page.goto("/");
+    // For some reason it doesn't work unless we need to redirect
+    await page.goto("/login");
 
     // Check if the poll title exists in the page content
     await expect(page.getByText("Test Poll")).toBeVisible();