BizOne-portal/frontend/middleware.ts

import { NextRequest, NextResponse } from 'next/server';

const API_URL = process.env.NEXT_PUBLIC_API_URL || 'http://localhost:3001/api';
const AUTH_COOKIE = 'wa_session';
const REFRESH_COOKIE = 'wa_refresh';
const SESSION_REFRESH_LEEWAY_SECONDS = 60;

type RefreshPayload = {
  access_token: string;
  refresh_token: string;
  access_token_max_age_seconds?: number;
  refresh_token_max_age_seconds?: number;
};

function isProtectedPath(pathname: string) {
  if (pathname.startsWith('/dashboard')) {
    return true;
  }

  if (!pathname.startsWith('/api')) {
    return false;
  }

  if (pathname.startsWith('/api/auth/password-reset')) {
    return false;
  }

  if (pathname.startsWith('/api/invitations/')) {
    return false;
  }

  return true;
}

function decodeJwtExpiry(token: string) {
  try {
    const parts = token.split('.');
    if (parts.length < 2) {
      return null;
    }

    const normalized = parts[1].replace(/-/g, '+').replace(/_/g, '/');
    const padded = normalized.padEnd(Math.ceil(normalized.length / 4) * 4, '=');
    const payload = JSON.parse(atob(padded)) as { exp?: unknown };
    return typeof payload.exp === 'number' ? payload.exp : null;
  } catch {
    return null;
  }
}

function tokenNeedsRefresh(token: string) {
  const expiry = decodeJwtExpiry(token);
  if (!expiry) {
    return true;
  }

  return expiry * 1000 <= Date.now() + SESSION_REFRESH_LEEWAY_SECONDS * 1000;
}

function buildRequestCookieHeader(request: NextRequest, payload: RefreshPayload) {
  const pairs = request.cookies.getAll().map((cookie) => [cookie.name, cookie.value] as const);
  const nextCookies = new Map<string, string>(pairs);
  nextCookies.set(AUTH_COOKIE, payload.access_token);
  nextCookies.set(REFRESH_COOKIE, payload.refresh_token);

  return Array.from(nextCookies.entries())
    .map(([name, value]) => `${name}=${value}`)
    .join('; ');
}

function applySessionCookies(response: NextResponse, payload: RefreshPayload) {
  const secure = process.env.NODE_ENV === 'production';

  response.cookies.set(AUTH_COOKIE, payload.access_token, {
    httpOnly: true,
    sameSite: 'lax',
    secure,
    path: '/',
    maxAge: payload.access_token_max_age_seconds || 60 * 60 * 24,
  });
  response.cookies.set(REFRESH_COOKIE, payload.refresh_token, {
    httpOnly: true,
    sameSite: 'lax',
    secure,
    path: '/',
    maxAge: payload.refresh_token_max_age_seconds || 60 * 60 * 24 * 30,
  });
}

function clearSessionCookies(response: NextResponse) {
  response.cookies.delete(AUTH_COOKIE);
  response.cookies.delete(REFRESH_COOKIE);
}

function unauthorizedResponse(request: NextRequest) {
  if (request.nextUrl.pathname.startsWith('/api')) {
    return NextResponse.json({ message: 'Unauthorized' }, { status: 401 });
  }

  return NextResponse.redirect(new URL('/login', request.url));
}

async function refreshSession(refreshToken: string) {
  const response = await fetch(`${API_URL}/auth/refresh`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ refreshToken }),
    cache: 'no-store',
  });

  if (!response.ok) {
    return null;
  }

  return (await response.json()) as RefreshPayload;
}

export async function middleware(request: NextRequest) {
  if (!isProtectedPath(request.nextUrl.pathname)) {
    return NextResponse.next();
  }

  const accessToken = request.cookies.get(AUTH_COOKIE)?.value;
  const refreshToken = request.cookies.get(REFRESH_COOKIE)?.value;

  if (accessToken && !tokenNeedsRefresh(accessToken)) {
    return NextResponse.next();
  }

  if (!refreshToken) {
    const response = unauthorizedResponse(request);
    clearSessionCookies(response);
    return response;
  }

  const payload = await refreshSession(refreshToken);
  if (!payload?.access_token || !payload.refresh_token) {
    const response = unauthorizedResponse(request);
    clearSessionCookies(response);
    return response;
  }

  const requestHeaders = new Headers(request.headers);
  requestHeaders.set('cookie', buildRequestCookieHeader(request, payload));

  const response = NextResponse.next({
    request: {
      headers: requestHeaders,
    },
  });

  applySessionCookies(response, payload);
  return response;
}

export const config = {
  matcher: ['/dashboard/:path*', '/api/:path*'],
};