From 45d7211859bd6b60d6dd2115116743eca4d3e8c2 Mon Sep 17 00:00:00 2001
From: Bailey Eaton <xYoshify@gmail.com>
Date: Wed, 18 Dec 2024 22:29:48 +1000
Subject: [PATCH] feat(auto-refresh): override callback result headers/cookies
 if the result is a response

In the instance where a user is using withAuth with the callback pattern, we can't just return the result of their callback directly - we need to be returning our response with our (possibly refreshed) tokens.

An initial workaround was just to invoke onSuccess and ignore whatever the user returned and return our response. This could be breaking if the user returned their own response.

This implementation checks to see if the user has returned their own response from the callback - if so, we merge our responses cookies and headers into it, then return it.

This way we:
- Preserve the user response
- Still pass down our refreshed tokens immediately

Note: If we have conflicting headers or cookie names, this does technically override them favouring ours - which is fine, I think.
---
 src/authMiddleware/authMiddleware.ts | 29 +++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/authMiddleware/authMiddleware.ts b/src/authMiddleware/authMiddleware.ts
index a9107e4..6e8fecb 100644
--- a/src/authMiddleware/authMiddleware.ts
+++ b/src/authMiddleware/authMiddleware.ts
@@ -183,7 +183,7 @@ const handleMiddleware = async (req, options, onSuccess) => {
     if(config.isDebugMode) {
       console.log('authMiddleware: invoking onSuccess callback')
     }
-    return await onSuccess({
+    const callbackResult = await onSuccess({
       token: accessTokenValue,
       user: {
         family_name: idTokenValue.family_name,
@@ -193,6 +193,33 @@ const handleMiddleware = async (req, options, onSuccess) => {
         picture: idTokenValue.picture,
       },
     });
+
+    // If a user returned a response from their onSuccess callback, copy our refreshed tokens to it
+    if (callbackResult instanceof NextResponse) {
+        if(config.isDebugMode) {
+          console.log('authMiddleware: onSuccess callback returned a response, copying our cookies to it')
+        }
+        // Copy our cookies to their response
+        resp.cookies.getAll().forEach(cookie => {
+            callbackResult.cookies.set(cookie.name, cookie.value, {
+              ...cookie
+            });
+        });
+
+        // Copy any headers we set (if any) to their response
+        resp.headers.forEach((value, key) => {
+            callbackResult.headers.set(key, value);
+        });
+
+        return callbackResult;
+    }
+
+    // If they didn't return a response, return our response with the refreshed tokens
+    if(config.isDebugMode) {
+      console.log('authMiddleware: onSuccess callback did not return a response, returning our response')
+    }
+
+    return resp;
   }
 
   if (customValidationValid) {