fix handleCallback not setting cookie

Author: frederikprijck

src/server/auth-client.ts

@@ -413,14 +413,12 @@ export class AuthClient {
     }
   }
 
-    /**
+  /**
    * Proxy handler that routes to the appropriate proxy handler method
    * @param nextRequest The incoming NextRequest instance.
    * @returns A Promise that resolves to a NextResponse or undefined.
    */
-  async proxyHandler(
-    nextRequest: NextRequest,
-  ): Promise<NextResponse | void> {
+  async proxyHandler(nextRequest: NextRequest): Promise<NextResponse | void> {
     // Convert the incomming request to Auth0Request
     const auth0Req = new Auth0NextRequest(nextRequest as NextRequest);
     // Convert the incoming response to Auth0Response
@@ -485,7 +483,7 @@ export class AuthClient {
     nextRequest: NextRequest | NextApiRequest,
     fallback?: (req: Auth0Request, res: Auth0Response) => Promise<Auth0Response>,
     nextResponse?: NextApiResponse
-  ): Promise<NextResponse | undefined> {
+  ): Promise<NextResponse | void> {
     // Convert the incomming request to Auth0Request
     const auth0Req = isRequest(nextRequest)
       ? new Auth0NextRequest(nextRequest as NextRequest)
@@ -554,14 +552,27 @@ export class AuthClient {
   }
 
   /**
-   * Unwraps an Auth0Response by extracting the underlying NextResponse.
+   * Unwraps an Auth0Response by extracting the underlying NextResponse, or calling `.end()` on the underlying NextApiResponse.
    * This utility simplifies the pattern of awaiting handler calls and accessing .res
+   * 
+   * @param handler A function that returns a Promise resolving to an Auth0Response.
+   * @returns A Promise that resolves to a NextResponse or void (in case of Pages Router usage).
    */
   async #unwrapHandler(
     handler: () => Promise<Auth0Response>
-  ): Promise<NextResponse> {
+  ): Promise<NextResponse | void> {
     const auth0Response = await handler();
-    return auth0Response.res;
+    const response = auth0Response.res;
+    const canEndRequest =
+      response && "end" in response && typeof response.end === "function";
+
+    // When the underlying response supports .end(), call it to finalize the response without returning the NextApiResponse instance.
+    // NextResponse does not have .end(), instead we return the NextResponse instance.
+    if (canEndRequest) {
+      response.end();
+    } else {
+      return response as NextResponse;
+    }
   }
 
   async startInteractiveLogin(
@@ -806,14 +817,25 @@ export class AuthClient {
       auth0Req.getCookies(),
       state
     );
+
+    const baseOnCallbackCtx: OnCallbackContext = {
+      request: auth0Req.req,
+      response: isRequest(auth0Req.req) ? undefined : auth0Res.res
+    };
+
     if (!transactionStateCookie) {
-      const errorRes = await this.onCallback(new InvalidStateError(), {}, null);
+      const errorRes = await this.onCallback(
+        new InvalidStateError(),
+        baseOnCallbackCtx,
+        null
+      );
       auth0Res.setResponse(errorRes);
       return auth0Res;
     }
 
     const transactionState = transactionStateCookie.payload;
     const onCallbackCtx: OnCallbackContext = {
+      ...baseOnCallbackCtx,
       responseType: transactionState.responseType,
       returnTo: transactionState.returnTo
     };
@@ -881,7 +903,7 @@ export class AuthClient {
         session
       );
 
-      auth0Res.setResponse(res);
+      auth0Res.setResponse(res ?? onCallbackCtx.response);
       await this.transactionStore.delete(auth0Res.getCookies(), state);
 
       return auth0Res;
@@ -1028,7 +1050,7 @@ export class AuthClient {
     // if not then filter id_token claims with default rules
     session = await this.finalizeSession(session, oidcRes.id_token);
 
-    auth0Res.setResponse(res);
+    auth0Res.setResponse(res ?? onCallbackCtx.response);
     await this.sessionStore.set(
       auth0Req.getCookies(),
       auth0Res.getCookies(),
@@ -1537,17 +1559,32 @@ export class AuthClient {
     error: SdkError | null,
     ctx: OnCallbackContext
   ) {
-    if (error) {
-      return new NextResponse(error.message, {
-        status: 500
-      });
-    }
+    const redirectUrl = createRouteUrl(ctx.returnTo || "/", this.appBaseUrl);
 
-    const res = NextResponse.redirect(
-      createRouteUrl(ctx.returnTo || "/", this.appBaseUrl)
-    );
+    if (ctx.request && isRequest(ctx.request)) {
+      if (error) {
+        return new NextResponse(error.message, {
+          status: 500
+        });
+      }
+
+      const res = NextResponse.redirect(redirectUrl);
 
-    return res;
+      return res;
+    } else {
+      if (error) {
+        ctx.response?.status(500).send(error.message);
+      } else {
+        // We do not use ctx.response.redirect(), as that would call `res.end()` immediately.
+        // Even though that should be okay conceptually,
+        // it would require changing code to not set cookies after calling `onCallback()`,
+        // which is not in scope.
+
+        // This also means that anyone using a custom onCallback should not end the response themselves.
+        // If they would, the Set-Cookie headers set after onCallback would be ignored and the transaction cookie will not exist.
+        ctx.response?.status(307).setHeader("Location", redirectUrl.toString());
+      }
+    }
   }
 
   /**
@@ -1561,7 +1598,7 @@ export class AuthClient {
   ): Promise<Auth0Response> {
     const response = await this.onCallback(error, ctx, null);
 
-    auth0Res.setResponse(response);
+    auth0Res.setResponse(response ?? ctx.response);
 
     // Clean up the transaction cookie on error to prevent accumulation
     if (state) {
@@ -2953,4 +2990,4 @@ export async function buildConnectAccountErrorResponse(
       null
     ];
   }
-}
+}

src/server/http/auth0-next-api-response.ts

@@ -26,7 +26,10 @@ export class Auth0NextApiResponse extends Auth0Response<NextApiResponse> {
    * @param url The URL to redirect to.
    */
   redirect(url: string) {
-    this.res = this.res.redirect(url);
+    // We do not use .redirect(), as that would call `res.end()` immediately.
+    // Even though that should be okay conceptually,
+    // it would require changing code to not set cookies after calling `redirect()`.
+    this.res = this.res.status(307).setHeader("Location", url);
     return this;
   }