From d4b5df56041c71859a40326d420ee1e0187ea0fa Mon Sep 17 00:00:00 2001
From: Illya Gerasymchuk <hello@illya.sh>
Date: Thu, 9 Jan 2025 15:55:10 +0000
Subject: [PATCH] feat: improve auto-cookie refresh + /api/me/info handling

---
 src/app/api/me/info/route.ts |  18 ++++--
 src/lib/auth/index.ts        |   3 +-
 src/middleware.ts            | 122 +++++++++++++++++++++++++++++++----
 3 files changed, 123 insertions(+), 20 deletions(-)

diff --git a/src/app/api/me/info/route.ts b/src/app/api/me/info/route.ts
index d993e0c..2127e87 100644
--- a/src/app/api/me/info/route.ts
+++ b/src/app/api/me/info/route.ts
@@ -1,8 +1,3 @@
-/**
- * In-development verstion of /api/me/info
- * 
- * Later, this will replace /api/me/info and its usages.
- */
 import { NextResponse } from "next/server";
 import { cookies } from "next/headers";
 import prisma from "@/lib/prisma";
@@ -11,6 +6,7 @@ import { UserService } from "@/services/UserService";
 import logger from "@/logging";
 import { deriveUserId } from "@/lib/user/derive";
 import { ApiResponse } from "@/lib/api-response";
+import { JWSInvalid, JWSSignatureVerificationFailed, JWTClaimValidationFailed, JWTExpired, JWTInvalid } from "jose/errors";
 
 const userService = new UserService(prisma);
 
@@ -42,10 +38,22 @@ export async function GET() {
     return ApiResponse.success(userInfo);
   } catch (error) {
     logger.error("User info error:", error);
+  
+    if (
+      error instanceof JWTInvalid ||
+      error instanceof JWTExpired ||
+      error instanceof JWTClaimValidationFailed ||
+      error instanceof JWSSignatureVerificationFailed ||
+      error instanceof JWSInvalid
+    ) {
+      return ApiResponse.unauthorized("Invalid or expired access token");
+    }
+
 
     if (error instanceof Error && error.message === "Invalid token") {
       return ApiResponse.unauthorized("Unauthorized");
     }
+ 
 
     return ApiResponse.error("Internal server error");
   }
diff --git a/src/lib/auth/index.ts b/src/lib/auth/index.ts
index d71dae8..d1a3c65 100644
--- a/src/lib/auth/index.ts
+++ b/src/lib/auth/index.ts
@@ -12,7 +12,7 @@ export async function getOrCreateUserFromRequest(req: Request): Promise<User | n
     // Get the access token from cookies
     const cookieStore = await cookies();
     const accessToken = cookieStore.get("access_token")?.value;
-
+    
     if (!accessToken) {
       return null;
     }
@@ -22,6 +22,7 @@ export async function getOrCreateUserFromRequest(req: Request): Promise<User | n
 
     // Resolve user from payload
     const user = await userService.findOrCreateUser(payload.authSource);
+
     return user;
   } catch (error) {
     logger.error("Error getting user from request:", error);
diff --git a/src/middleware.ts b/src/middleware.ts
index 4adfcd3..315f95a 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -4,6 +4,7 @@ import { verifyToken } from "./lib/auth/jwt";
 import { AppError } from "./lib/errors";
 import logger from "./logging";
 import { JWTPayload } from "jose";
+import { cookies } from "next/headers";
 
 const getBaseUrl = () => process.env.NEXT_APP_URL;
 
@@ -144,7 +145,7 @@ async function handleTokenRefresh(refreshToken: string, cookieManager: CookieMan
 
     // Get all Set-Cookie headers
     const cookies = refreshResponse.headers.getSetCookie();
-    
+
     // Parse the cookies to get the new tokens
     const newAccessToken = cookies
       .find(cookie => cookie.startsWith('access_token='))
@@ -167,7 +168,7 @@ async function handleTokenRefresh(refreshToken: string, cookieManager: CookieMan
 }
 
 // Update authenticateRequest to use CookieManager
-async function authenticateRequest(cookieManager: CookieManager): Promise<AuthResult> {
+async function authenticateRequest(cookieManager: CookieManager, request: NextRequest): Promise<AuthResult> {
   const { accessToken, refreshToken } = cookieManager.getRequestTokens();
 
   // No tokens available
@@ -197,6 +198,10 @@ async function authenticateRequest(cookieManager: CookieManager): Promise<AuthRe
       if (refreshSuccessful) {
         const { accessToken: newAccessToken } = cookieManager.getResponseTokens();
         if (newAccessToken) {
+          // We don't need to set the cookie on the request anymore since it's handled by cookieManager
+          // and will be propagated through the response
+          const cookieStore = await cookies();
+
           const isAdmin = await checkAdminAccess(newAccessToken);
           return {
             isAuthenticated: true,
@@ -212,31 +217,120 @@ async function authenticateRequest(cookieManager: CookieManager): Promise<AuthRe
   return { isAuthenticated: false, error: "Authentication failed" };
 }
 
-// Update middleware function to use CookieManager
+// Update middleware function to validate tokens first
 export async function middleware(request: NextRequest) {
   const path = request.nextUrl.pathname;
   const routeType = getRouteType(path);
 
   logger.debug("[Middleware] Processing request:", { path, routeType });
 
-  const response = NextResponse.next();
-  const cookieManager = new CookieManager(request, response);
+  // Skip token validation for public and auth paths
+  if (isPublicPath(path)) {
+    return NextResponse.next();
+  }
 
-  // For authentication requests, just authenticate
   if (isAuthPath(path)) {
-    return response;
+    return NextResponse.next();
   }
 
-  // Not auth request: authenticate and check permissions
-  const authResult = await authenticateRequest(cookieManager);
+  // Get tokens from request
+  const accessToken = request.cookies.get("access_token")?.value;
+  const refreshToken = request.cookies.get("refresh_token")?.value;
 
-  // Skip authentication for public paths
-  if (isPublicPath(path)) {
-    return response;
+  logger.error(`1. accessToken: ${accessToken} refreshToken: ${refreshToken}`);
+
+  // No tokens available
+  if (!accessToken && !refreshToken) {
+    return generateUnauthorizedResponse(routeType, request);
+  }
+
+  // Try to validate access token
+  if (accessToken) {
+    try {
+      await verifyToken(accessToken);
+      const isAdmin = await checkAdminAccess(accessToken);
+      
+      // Token is valid, create response and proceed
+      const response = NextResponse.next();
+      
+      // For admin routes, check permissions
+      if (routeType === 'admin' && !isAdmin) {
+        return generateAdminUnauthorizedResponse(routeType, request);
+      }
+
+      return response;
+    } catch (error) {
+      logger.debug("[Middleware] Access token invalid, will attempt refresh");
+      // Continue to refresh token logic
+    }
+  }
+
+  // Try refresh token if access token is invalid or missing
+  if (refreshToken) {
+    try {
+      const baseUrl = getBaseUrl();
+      const refreshResponse = await fetch(`${baseUrl}/api/auth/refresh`, {
+        method: "POST",
+        headers: {
+          Cookie: `refresh_token=${refreshToken}`,
+        },
+      });
+
+      logger.error(`2. refreshResponse: ${refreshResponse}`);
+
+      if (!refreshResponse.ok) {
+        return generateUnauthorizedResponse(routeType, request);
+      }
+
+      // Get the new tokens from the refresh response
+      const cookies = refreshResponse.headers.getSetCookie();
+      const newAccessToken = cookies
+        .find(cookie => cookie.startsWith('access_token='))
+        ?.split(';')[0]
+        .split('=')[1];
+
+      if (!newAccessToken) {
+        return generateUnauthorizedResponse(routeType, request);
+      }
+
+      logger.error(`3. newAccessToken: ${newAccessToken}`);
+
+      const requestHeaders = new Headers(request.headers);
+      const requestCookies = requestHeaders.get('Cookie') || '';
+      const cookieArray = requestCookies.split('; ');
+      const updatedCookies = cookieArray.filter(cookie => !cookie.startsWith('access_token='));
+      updatedCookies.push(`access_token=${newAccessToken}`);
+      
+      requestHeaders.set('Cookie', updatedCookies.join('; '));
+      
+      const modifiedRequest = new NextRequest(request, {
+        headers: requestHeaders,
+      });
+
+      // Verify admin status with new token
+      const isAdmin = await checkAdminAccess(newAccessToken);
+
+      // Create response with the new tokens
+      const response = NextResponse.next({request: modifiedRequest});
+      
+      // Set each cookie as a separate Set-Cookie header as per HTTP/1.1 spec
+      cookies.forEach(cookie => {
+        response.headers.append('Set-Cookie', cookie);
+      });
+
+      // For admin routes, check permissions
+      if (routeType === 'admin' && !isAdmin) {
+        return generateAdminUnauthorizedResponse(routeType, request);
+      }
+
+      return response;
+    } catch (error) {
+      logger.error("[Middleware] Refresh failed:", error);
+      return generateUnauthorizedResponse(routeType, request);
+    }
   }
 
-  // Handle response based on authentication and permissions
-  return handleAuthResponse(request, routeType, authResult, response);
+  return generateUnauthorizedResponse(routeType, request);
 }
 
 /**