Enhance Auth Endpoints To Manage CSRF Tokens

Now that the middleware is set up to verify CSRF tokens, we need to update your Callback Endpoint and Logout Endpoint to create and clean up CSRF tokens.

Update Callback Endpoint

Your Callback Endpoint needs to be updated to create a CSRF token. The CSRF token will be derived from a session-specific CSRF secret. Once the CSRF token is generated, it will be stored in a cookie so that the frontend JavaScript code can read it. In addition, the CSRF secret will be stored in the session cookie so the CSRF middleware can access it to verify the CSRF token.

//
// Location -> src/pages/api/auth/callback.ts
//
import type { NextApiRequest, NextApiResponse } from 'next';
import { CallbackResultType, PageRouterCallbackResult } from '@wristband/nextjs-auth';
import { wristbandAuth } from '@/wristband-auth';
import { getSession } from '@/session/session';
import { createCsrfSecret, setCsrfTokenCookie } from '@/utils/csrf';

// Callback Endpoint
export default async function handleCallback(req: NextApiRequest, res: NextApiResponse) {
  try {
    const callbackResult: PageRouterCallbackResult = await wristbandAuth.pageRouter.callback(req, res);
    const { callbackData, result } = callbackResult;

    if (result === CallbackResultType.REDIRECT_REQUIRED) {
      return;
    }
    
    const session = await getSession(req, res);
    session.isAuthenticated = true;
    session.accessToken = callbackData.accessToken;
    session.expiresAt = Date.now() + callbackData.expiresIn * 1000;
    session.refreshToken = callbackData.refreshToken;
    session.userId = callbackData.userinfo.sub;
    session.tenantId = callbackData.userinfo.tnt_id;
    session.tenantDomainName = callbackData.tenantDomainName;
    session.tenantCustomDomain = callbackData.tenantCustomDomain || undefined;

    /* ***** BEGIN NEW CSRF LOGIC ***** */
    
    const csrfSecret = createCsrfSecret();
    session.csrfSecret = csrfSecret;
    await setCsrfTokenCookie(csrfSecret, res);
    
    /* ***** END NEW CSRF LOGIC ***** */
    
    await session.save();

    return res.redirect(callbackData!.returnUrl || '<replace_with_a_default_return_url>');
  } catch (error) {
    console.error(error);
    return res.status(500).json({ error: "Internal Server Error" });
  }
}

//
// Location -> src/app/api/auth/callback/route.ts
//
import { NextRequest } from 'next/server';
import { AppRouterCallbackResult, CallbackResultType } from '@wristband/nextjs-auth';
import { wristbandAuth } from '@/wristband-auth';
import { getSession } from '@/session/session';
import { createCsrfSecret, setCsrfTokenCookie } from '@/utils/csrf';

// Callback Endpoint
export async function GET(req: NextRequest) {
  try {
    const callbackResult: AppRouterCallbackResult = await wristbandAuth.appRouter.callback(req);
    const { callbackData, redirectResponse, result } = callbackResult;

    if (result === CallbackResultType.REDIRECT_REQUIRED) {
      return redirectResponse;
    }
    
    const appUrl = callbackData!.returnUrl || '<replace_with_a_default_return_url>';
    const callbackResponse = await wristbandAuth.appRouter.createCallbackResponse(req, appUrl);
    
    const session = await getSession();
    session.isAuthenticated = true;
    session.accessToken = callbackData.accessToken;
    session.expiresAt = Date.now() + callbackData.expiresIn * 1000;
    session.refreshToken = callbackData.refreshToken;
    session.userId = callbackData.userinfo.sub;
    session.tenantId = callbackData.userinfo.tnt_id;
    session.tenantDomainName = callbackData.tenantDomainName;
    session.tenantCustomDomain = callbackData.tenantCustomDomain || undefined;
    
    /* ***** BEGIN NEW CSRF LOGIC ***** */
    
    const csrfSecret = createCsrfSecret();
    session.csrfSecret = csrfSecret;
    await setCsrfTokenCookie(csrfSecret, callbackResponse);

    /* ***** END NEW CSRF LOGIC ***** */
    
    await session.save();

    return callbackResponse;
  } catch (error) {
    console.error(error);
    return Response.json({ error: "Internal Server Error" }, { status: 500 });
  }
}

Update Logout Endpoint

Now, when a user logs out, we need to delete their CSRF cookie before redirecting to the Wristband Logout Endpoint.

//
// Location -> src/pages/api/auth/logout.ts
//
import type { NextApiRequest, NextApiResponse } from 'next';
import { wristbandAuth } from '@/wristband-auth';
import { getSession } from '@/session/session';

// Logout endpoint
export default async function handleLogout(req: NextApiRequest, res: NextApiResponse) {
  try {
    const session = await getSession(req, res);
    const { refreshToken, tenantCustomDomain, tenantDomainName } = session;
    const logoutConfig = { refreshToken, tenantCustomDomain, tenantDomainName };
  	res.setHeader('Set-Cookie', [`session=; Max-Age=0; Path=/`]);
  	session.destroy();
    
    /* ***** BEGIN NEW CSRF LOGIC ***** */
    
    res.setHeader('Set-Cookie', [`CSRF-TOKEN=; Max-Age=0; Path=/`]);
    
    /* ***** END NEW CSRF LOGIC ***** */

    return await wristbandAuth.pageRouter.logout(req, res, logoutConfig);
  } catch (error) {
    console.error(error);
    return res.status(500).json({ error: "Internal Server Error" });
  }
}

//
// Location -> src/app/api/auth/logout/route.ts
//
import type { NextRequest } from 'next/server';
import { cookies } from 'next/headers';
import { wristbandAuth } from '@/wristband-auth';
import { getSession } from '@/session/session';

// Logout endpoint
export async function GET(req: NextRequest) {
  try {
    const session = await getSession();
    const { refreshToken, tenantCustomDomain, tenantDomainName } = session;
    const logoutConfig = { refreshToken, tenantCustomDomain, tenantDomainName };
    
    const cookieStore = await cookies();
    cookieStore.delete('session');
    session.destroy();
    
    /* ***** BEGIN NEW CSRF LOGIC ***** */
    
    cookieStore.delete('CSRF-TOKEN');
    
    /* ***** END NEW CSRF LOGIC ***** */
    
    return await wristbandAuth.appRouter.logout(req, logoutConfig);
  } catch (error) {
    console.error(error);
    return Response.json({ error: "Internal Server Error" }, { status: 500 });
  }
}