fix: refresh access tokens flow

apps/dashboard/app/(app)/api/auth/refresh/route.ts

@@ -1,34 +1,146 @@
-import { setCookie } from "@/lib/auth/cookies";
+import { getCookie, setCookie } from "@/lib/auth/cookies";
 import { auth } from "@/lib/auth/server";
-import { UNKEY_SESSION_COOKIE } from "@/lib/auth/types";
+import { tokenManager } from "@/lib/auth/token-management-service";
+import {
+  UNKEY_ACCESS_MAX_AGE,
+  UNKEY_ACCESS_TOKEN,
+  UNKEY_REFRESH_MAX_AGE,
+  UNKEY_REFRESH_TOKEN,
+  UNKEY_SESSION_COOKIE,
+  UNKEY_USER_IDENTITY_COOKIE,
+  UNKEY_USER_IDENTITY_MAX_AGE,
+} from "@/lib/auth/types";
+import { type NextRequest, NextResponse } from "next/server";
 
-export async function POST(request: Request) {
+export async function POST(req: NextRequest) {
   try {
-    // Get the current token from the request
-    const currentToken = request.headers.get("x-current-token");
-    if (!currentToken) {
-      console.error("Session refresh failed: no current token");
-      return Response.json({ success: false, error: "Failed to refresh session" }, { status: 401 });
+    // Get the refresh token from the cookie
+    const refreshToken = await getCookie(UNKEY_REFRESH_TOKEN);
+
+    // Get the user identity from the request headers and cookies
+    const requestUserIdentity = req.headers.get("x-user-identity");
+    const cookieUserIdentity = await getCookie(UNKEY_USER_IDENTITY_COOKIE);
+
+    // If we don't have a refresh token, return error
+    if (!refreshToken) {
+      return NextResponse.json({ error: "No refresh token available" }, { status: 401 });
+    }
+
+    // Check user identity consistency
+    // 1. If we have a user identity in the cookie, it should match the request
+    // 2. If we don't have one in the cookie but have one in the request, accept it
+    if (cookieUserIdentity && requestUserIdentity && cookieUserIdentity !== requestUserIdentity) {
+      console.warn("User identity mismatch during refresh");
+      return NextResponse.json({ error: "Invalid user identity" }, { status: 403 });
+    }
+
+    // Use either the cookie identity or request identity
+    const userIdentity = cookieUserIdentity || requestUserIdentity;
+
+    // If we have a user identity, verify token ownership
+    if (userIdentity) {
+      const isValidOwner = tokenManager.verifyTokenOwnership({
+        refreshToken,
+        userIdentity,
+      });
+
+      if (!isValidOwner) {
+        console.warn("Refresh token ownership verification failed");
+        return NextResponse.json({ error: "Invalid refresh token ownership" }, { status: 403 });
+      }
+
+      // If the user identity is only in the request, set it as a cookie
+      if (!cookieUserIdentity && requestUserIdentity) {
+        await setCookie({
+          name: UNKEY_USER_IDENTITY_COOKIE,
+          value: requestUserIdentity,
+          options: {
+            httpOnly: true,
+            secure: true,
+            sameSite: "strict",
+            maxAge: UNKEY_USER_IDENTITY_MAX_AGE,
+            path: "/",
+          },
+        });
+      }
+    }
+
+    // continue with token refresh
+    const result = await auth.refreshAccessToken(refreshToken);
+
+    if (!result || !result.session) {
+      // Remove the invalid token from our ownership mapping
+      tokenManager.removeToken(refreshToken);
+
+      return NextResponse.json({ error: "Failed to refresh session" }, { status: 401 });
     }
-    // Call refreshSession logic here and get new token
-    const { newToken, expiresAt } = await auth.refreshSession(currentToken);
 
-    // Set the new cookie
+    // Calculate remaining time for session in seconds
+    const sessionMaxAge = Math.floor((result.expiresAt.getTime() - Date.now()) / 1000);
+
+    // Update session cookie
     await setCookie({
       name: UNKEY_SESSION_COOKIE,
-      value: newToken,
+      value: result.sessionToken,
       options: {
         httpOnly: true,
         secure: true,
-        sameSite: "lax",
+        sameSite: "strict",
+        maxAge: sessionMaxAge,
         path: "/",
-        maxAge: Math.floor((expiresAt.getTime() - Date.now()) / 1000), // Convert to seconds
       },
     });
 
-    return Response.json({ success: true });
+    // Set access token cookie if available
+    if (result.accessToken) {
+      await setCookie({
+        name: UNKEY_ACCESS_TOKEN,
+        value: result.accessToken,
+        options: {
+          httpOnly: true,
+          secure: true,
+          sameSite: "strict",
+          maxAge: UNKEY_ACCESS_MAX_AGE,
+          path: "/",
+        },
+      });
+    }
+
+    // Update refresh token if available
+    if (result.refreshToken && result.refreshToken !== refreshToken) {
+      await setCookie({
+        name: UNKEY_REFRESH_TOKEN,
+        value: result.refreshToken,
+        options: {
+          httpOnly: true,
+          secure: true,
+          sameSite: "strict",
+          maxAge: UNKEY_REFRESH_MAX_AGE,
+          path: "/",
+        },
+      });
+
+      // Update token ownership mapping
+      if (userIdentity) {
+        tokenManager.updateTokenOwnership({
+          oldToken: refreshToken,
+          newToken: result.refreshToken,
+          userIdentity: userIdentity,
+        });
+      }
+    }
+
+    // Return success response with access token
+    return NextResponse.json({
+      success: true,
+      userId: result.session.userId,
+      orgId: result.session.orgId,
+      role: result.session.role,
+      accessToken: result.accessToken,
+      expiresAt: result.expiresAt,
+    });
   } catch (error) {
-    console.error("Session refresh failed:", error);
-    return Response.json({ success: false, error: "Failed to refresh session" }, { status: 401 });
+    console.error("Refresh error:", error);
+    return NextResponse.json({ error: "Server error during refresh" }, { status: 500 });
   }
 }

apps/dashboard/app/(app)/api/auth/session/route.ts

@@ -0,0 +1,18 @@
+import { getAuth } from "@/lib/auth/get-auth";
+
+export const dynamic = "force-dynamic";
+export async function GET() {
+  try {
+    const { userId } = await getAuth();
+
+    if (!userId) {
+      return new Response(null, { status: 401 });
+    }
+
+    // Just return 200 OK if authenticated
+    return new Response(null, { status: 200 });
+  } catch (error) {
+    console.error("Error checking session:", error);
+    return new Response(null, { status: 401 });
+  }
+}

apps/dashboard/app/(app)/layout.tsx

@@ -1,8 +1,10 @@
+import { randomUUID } from "node:crypto";
 import { AppSidebar } from "@/components/navigation/sidebar/app-sidebar";
 import { SidebarMobile } from "@/components/navigation/sidebar/sidebar-mobile";
 import { SidebarProvider } from "@/components/ui/sidebar";
 import { getIsImpersonator, getOrgId } from "@/lib/auth";
 import { db } from "@/lib/db";
+import { AuthProvider } from "@/providers/AuthProvider";
 import { Empty } from "@unkey/ui";
 import Link from "next/link";
 import { redirect } from "next/navigation";
@@ -29,57 +31,63 @@ export default async function Layout({ children }: LayoutProps) {
     return redirect("/new");
   }
 
+  // Generate a secure random UUID on the server
+  // this will be a fall-back if there isn't already a cookie
+  const serverGeneratedIdentity = randomUUID();
+
   return (
-    <div className="h-[100dvh] relative flex flex-col overflow-hidden bg-white dark:bg-base-12 lg:flex-row">
-      <SidebarProvider>
-        <div className="flex flex-1 overflow-hidden">
-          {/* Desktop Sidebar */}
-          <AppSidebar
-            workspace={{ ...workspace, quotas: workspace.quotas! }}
-            className="bg-gray-1 border-grayA-4"
-          />
+    <AuthProvider requireAuth={true} serverGeneratedIdentity={serverGeneratedIdentity}>
+      <div className="h-[100dvh] relative flex flex-col overflow-hidden bg-white dark:bg-base-12 lg:flex-row">
+        <SidebarProvider>
+          <div className="flex flex-1 overflow-hidden">
+            {/* Desktop Sidebar */}
+            <AppSidebar
+              workspace={{ ...workspace, quotas: workspace.quotas! }}
+              className="bg-gray-1 border-grayA-4"
+            />
 
-          {/* Main content area */}
-          <div className="flex-1 overflow-auto">
-            <div
-              className="isolate bg-base-12 w-full overflow-x-auto flex flex-col items-center"
-              id="layout-wrapper"
-            >
-              {/* Mobile sidebar at the top of content */}
-              <SidebarMobile workspace={workspace} />
+            {/* Main content area */}
+            <div className="flex-1 overflow-auto">
+              <div
+                className="isolate bg-base-12 w-full overflow-x-auto flex flex-col items-center"
+                id="layout-wrapper"
+              >
+                {/* Mobile sidebar at the top of content */}
+                <SidebarMobile workspace={workspace} />
 
-              <div className="w-full">
-                {workspace.enabled ? (
-                  <QueryTimeProvider>{children}</QueryTimeProvider>
-                ) : (
-                  <div className="flex items-center justify-center w-full h-full">
-                    <Empty>
-                      <Empty.Icon />
-                      <Empty.Title>This workspace is disabled</Empty.Title>
-                      <Empty.Description>
-                        Contact{" "}
-                        <Link
-                          href={`mailto:support@unkey.dev?body=workspaceId: ${workspace.id}`}
-                          className="underline"
-                        >
-                          support@unkey.dev
-                        </Link>
-                      </Empty.Description>
-                    </Empty>
-                  </div>
-                )}
-              </div>
-            </div>
-            {isImpersonator ? (
-              <div className="fixed top-0 inset-x-0 z-50 flex justify-center  border-t-2 border-error-9">
-                <div className="bg-error-9  flex -mt-1 font-mono items-center gap-2 text-white text-xs rounded-b overflow-hidden shadow-lg select-none pointer-events-none px-1.5 py-0.5">
-                  Impersonation Mode. Do not change anything and log out after you are done.
+                <div className="w-full">
+                  {workspace.enabled ? (
+                    <QueryTimeProvider>{children}</QueryTimeProvider>
+                  ) : (
+                    <div className="flex items-center justify-center w-full h-full">
+                      <Empty>
+                        <Empty.Icon />
+                        <Empty.Title>This workspace is disabled</Empty.Title>
+                        <Empty.Description>
+                          Contact{" "}
+                          <Link
+                            href={`mailto:support@unkey.dev?body=workspaceId: ${workspace.id}`}
+                            className="underline"
+                          >
+                            support@unkey.dev
+                          </Link>
+                        </Empty.Description>
+                      </Empty>
+                    </div>
+                  )}
                 </div>
               </div>
-            ) : null}
+              {isImpersonator ? (
+                <div className="fixed top-0 inset-x-0 z-50 flex justify-center  border-t-2 border-error-9">
+                  <div className="bg-error-9  flex -mt-1 font-mono items-center gap-2 text-white text-xs rounded-b overflow-hidden shadow-lg select-none pointer-events-none px-1.5 py-0.5">
+                    Impersonation Mode. Do not change anything and log out after you are done.
+                  </div>
+                </div>
+              ) : null}
+            </div>
           </div>
-        </div>
-      </SidebarProvider>
-    </div>
+        </SidebarProvider>
+      </div>
+    </AuthProvider>
   );
 }

apps/dashboard/app/auth/actions.ts

@@ -185,9 +185,9 @@ export async function completeOrgSelection(
 // Used in route handlers, like join
 export async function switchOrg(orgId: string): Promise<{ success: boolean; error?: string }> {
   try {
-    const { newToken, expiresAt } = await auth.switchOrg(orgId);
+    const { sessionToken, expiresAt } = await auth.switchOrg(orgId);
 
-    await SetSessionCookie({ token: newToken, expiresAt });
+    await SetSessionCookie({ token: sessionToken, expiresAt });
 
     return { success: true };
   } catch (error) {

apps/dashboard/app/auth/sso-callback/[[...sso-callback]]/route.ts

@@ -4,7 +4,6 @@ import { AuthErrorCode, SIGN_IN_URL } from "@/lib/auth/types";
 import { type NextRequest, NextResponse } from "next/server";
 export async function GET(request: NextRequest) {
   const authResult = await auth.completeOAuthSignIn(request);
-
   if (!authResult.success) {
     if (
       (authResult.code === AuthErrorCode.ORGANIZATION_SELECTION_REQUIRED ||
@@ -25,7 +24,6 @@ export async function GET(request: NextRequest) {
       }
 
       const response = NextResponse.redirect(url);
-
       return await setCookiesOnResponse(response, authResult.cookies);
     }
 

apps/dashboard/app/new/create-workspace.tsx

@@ -42,7 +42,7 @@ export const CreateWorkspace: React.FC = () => {
         options: {
           httpOnly: true,
           secure: true,
-          sameSite: "lax",
+          sameSite: "strict",
           path: "/",
           maxAge: Math.floor((sessionData.expiresAt!.getTime() - Date.now()) / 1000),
         },

apps/dashboard/lib/auth.ts

@@ -1,14 +1,9 @@
-import { getAuth as noCacheGetAuth } from "@/lib/auth/get-auth";
+import { getAuth as baseGetAuth } from "@/lib/auth/get-auth";
 import { auth } from "@/lib/auth/server";
-import type { User } from "@/lib/auth/types";
+import type { AuthResult, User } from "@/lib/auth/types";
 import { redirect } from "next/navigation";
 import { cache } from "react";
 
-type GetAuthResult = {
-  userId: string | null;
-  orgId: string | null;
-};
-
 export async function getIsImpersonator(): Promise<boolean> {
   const user = await getCurrentUser();
   return user.impersonator !== undefined;
@@ -26,8 +21,8 @@ export async function getIsImpersonator(): Promise<boolean> {
  * @returns Authentication result containing userId and orgId if authenticated, null values otherwise
  * @throws Redirects to sign-in or organization/workspace creation pages if requirements aren't met
  */
-export const getAuth = cache(async (_req?: Request): Promise<GetAuthResult> => {
-  const authResult = await noCacheGetAuth();
+export const getAuth = cache(async (_req?: Request): Promise<AuthResult> => {
+  const authResult = await baseGetAuth();
   if (!authResult.userId) {
     redirect("/auth/sign-in");
   }