OmniRoute/src/server/authz/pipeline.ts

import { jwtVerify, SignJWT } from "jose";
import { NextResponse, type NextRequest } from "next/server";
import { getCachedSettings } from "../../lib/db/readCache";
import { isDraining } from "../../lib/gracefulShutdown";
import { checkBodySize, getBodySizeLimit } from "../../shared/middleware/bodySizeGuard";
import { generateRequestId } from "../../shared/utils/requestId";
import { applyCorsHeaders } from "../cors/origins";
import { validateBrowserMutationOrigin } from "../origin/publicOrigin";
import { classifyRoute } from "./classify";
import { validateDashboardCsrfToken } from "./csrf";
import {
  classifyStampedPeerLocality,
  resolveStampedPeer,
  resolveStampedViaProxy,
} from "./peerStamp";
import { checkRequestIP } from "@omniroute/open-sse/services/ipFilter.ts";
import { clientApiPolicy } from "./policies/clientApi";
import { managementPolicy } from "./policies/management";
import { publicPolicy } from "./policies/public";
import {
  AUTHZ_HEADER_AUTH_ID,
  AUTHZ_HEADER_AUTH_KIND,
  AUTHZ_HEADER_AUTH_LABEL,
  AUTHZ_HEADER_AUTH_SCOPES,
  AUTHZ_HEADER_PEER_LOCALITY,
  AUTHZ_HEADER_REQUEST_ID,
  AUTHZ_HEADER_ROUTE_CLASS,
  AUTHZ_TRUSTED_HEADERS,
  PEER_IP_HEADER,
  VIA_PROXY_HEADER,
} from "./headers";
import type { AuthSubject, RouteClass, RouteClassification } from "./types";
import type { AuthOutcome, RoutePolicy } from "./context";

export interface AuthzPipelineOptions {
  enforce?: boolean;
}

const POLICIES: Record<RouteClass, RoutePolicy> = {
  PUBLIC: publicPolicy,
  CLIENT_API: clientApiPolicy,
  MANAGEMENT: managementPolicy,
};

let staleDashboardJwtWarningEmitted = false;

function isStaleDashboardJwtError(error: unknown): boolean {
  const code =
    typeof error === "object" &&
    error !== null &&
    "code" in error &&
    typeof (error as { code?: unknown }).code === "string"
      ? (error as { code: string }).code
      : "";

  if (
    code === "ERR_JWS_SIGNATURE_VERIFICATION_FAILED" ||
    code === "ERR_JWT_EXPIRED" ||
    code === "ERR_JWS_INVALID" ||
    code === "ERR_JWT_CLAIM_VALIDATION_FAILED"
  ) {
    return true;
  }

  return error instanceof Error && error.message.includes("signature verification failed");
}

function stampSubject(headers: Headers, subject: AuthSubject): void {
  headers.set(AUTHZ_HEADER_AUTH_KIND, subject.kind);
  headers.set(AUTHZ_HEADER_AUTH_ID, subject.id);
  if (subject.label) headers.set(AUTHZ_HEADER_AUTH_LABEL, subject.label);
  if (subject.scopes && subject.scopes.length > 0) {
    headers.set(AUTHZ_HEADER_AUTH_SCOPES, subject.scopes.join(","));
  }
}

function rejectionResponse(
  outcome: Extract<AuthOutcome, { allow: false }>,
  classification: RouteClassification,
  requestId: string
): NextResponse {
  const response = NextResponse.json(
    {
      error: {
        code: outcome.code,
        message: outcome.message,
        correlation_id: requestId,
      },
    },
    { status: outcome.status }
  );
  response.headers.set(AUTHZ_HEADER_REQUEST_ID, requestId);
  response.headers.set(AUTHZ_HEADER_ROUTE_CLASS, classification.routeClass);
  return response;
}

function isDashboardPath(pathname: string): boolean {
  return (
    pathname === "/dashboard" ||
    pathname.startsWith("/dashboard/") ||
    pathname === "/home" ||
    pathname.startsWith("/home/")
  );
}

function isManagementDashboardRoute(
  classification: RouteClassification,
  pathname: string
): boolean {
  return classification.routeClass === "MANAGEMENT" && isDashboardPath(pathname);
}

function getCookieValue(request: NextRequest, name: string): string | null {
  const fromCookies = request.cookies.get(name)?.value;
  if (fromCookies) return fromCookies;

  const cookieHeader = request.headers.get("cookie") || request.headers.get("Cookie");
  if (!cookieHeader) return null;

  for (const segment of cookieHeader.split(";")) {
    const [rawKey, ...rawValue] = segment.split("=");
    if (!rawKey || rawValue.length === 0) continue;
    if (rawKey.trim() === name) return rawValue.join("=").trim() || null;
  }

  return null;
}

function getJwtSecret(): Uint8Array | null {
  const secret = process.env.JWT_SECRET?.trim();
  return secret ? new TextEncoder().encode(secret) : null;
}

function shouldUseSecureCookie(request: NextRequest): boolean {
  if (process.env.AUTH_COOKIE_SECURE === "true") return true;
  const forwardedProto = (request.headers.get("x-forwarded-proto") || "")
    .split(",")[0]
    .trim()
    .toLowerCase();
  return forwardedProto === "https" || request.nextUrl.protocol === "https:";
}

async function refreshDashboardSessionIfNeeded(
  response: NextResponse,
  request: NextRequest
): Promise<void> {
  const secret = getJwtSecret();
  if (!secret) return;

  const token = getCookieValue(request, "auth_token");
  if (!token) return;

  try {
    const { payload } = await jwtVerify(token, secret);
    const exp = typeof payload.exp === "number" ? payload.exp : null;
    if (!exp) return;

    const now = Math.floor(Date.now() / 1000);
    const refreshWindowSeconds = 7 * 24 * 60 * 60;
    if (exp - now >= refreshWindowSeconds) return;

    const freshToken = await new SignJWT({ authenticated: true })
      .setProtectedHeader({ alg: "HS256" })
      .setExpirationTime("30d")
      .sign(secret);

    response.cookies.set("auth_token", freshToken, {
      httpOnly: true,
      secure: shouldUseSecureCookie(request),
      sameSite: "lax",
      path: "/",
    });
  } catch (error) {
    if (isStaleDashboardJwtError(error)) {
      response.cookies.delete("auth_token");
      if (!staleDashboardJwtWarningEmitted) {
        staleDashboardJwtWarningEmitted = true;
        console.warn("[Authz] Dropped stale dashboard session cookie during auto-refresh");
      }
      return;
    }

    console.error("[Authz] JWT auto-refresh failed:", error);
  }
}

function dashboardLoginRedirect(request: NextRequest, requestId: string): NextResponse {
  const response = NextResponse.redirect(new URL(`${request.nextUrl.basePath}/login`, request.url));
  response.cookies.delete("auth_token");
  stampRouteResponse(response, requestId, "MANAGEMENT");
  applyCorsHeaders(response, request);
  return response;
}

function drainingResponse(requestId: string): NextResponse {
  const response = NextResponse.json(
    {
      error: {
        code: "SERVICE_UNAVAILABLE",
        message: "Server is shutting down",
        correlation_id: requestId,
      },
    },
    { status: 503 }
  );
  response.headers.set(AUTHZ_HEADER_REQUEST_ID, requestId);
  return response;
}

function invalidOriginResponse(requestId: string): NextResponse {
  const response = NextResponse.json(
    {
      error: {
        code: "INVALID_ORIGIN",
        message:
          "Invalid request origin. Same-origin dashboard writes must include a valid dashboard CSRF token. " +
          "Refresh the dashboard and retry, or set OMNIROUTE_PUBLIC_BASE_URL for non-dashboard browser integrations.",
        correlation_id: requestId,
      },
    },
    { status: 403 }
  );
  response.headers.set(AUTHZ_HEADER_REQUEST_ID, requestId);
  return response;
}

function isUnsafeMutationMethod(method: string): boolean {
  return !["GET", "HEAD", "OPTIONS"].includes(method.toUpperCase());
}

function stampRouteResponse(
  response: Response,
  requestId: string,
  routeClass: RouteClass
): Response {
  response.headers.set(AUTHZ_HEADER_REQUEST_ID, requestId);
  response.headers.set(AUTHZ_HEADER_ROUTE_CLASS, routeClass);
  return response;
}

async function getBodySizeSettings(): Promise<Record<string, unknown> | undefined> {
  try {
    return await getCachedSettings();
  } catch (error) {
    console.warn(
      "[Authz] Failed to load request body limit settings:",
      error instanceof Error ? error.message : error
    );
    return undefined;
  }
}

export async function runAuthzPipeline(
  request: NextRequest,
  options: AuthzPipelineOptions = {}
): Promise<Response> {
  const { pathname } = request.nextUrl;
  const method = request.method;

  const requestId = generateRequestId();

  if (pathname === "/") {
    const response = NextResponse.redirect(
      new URL(`${request.nextUrl.basePath}/dashboard`, request.url)
    );
    return stampRouteResponse(response, requestId, "MANAGEMENT");
  }

  const classification = classifyRoute(pathname, method);
  const guardedPathname = classification.normalizedPath;
  const managementDashboardRoute = isManagementDashboardRoute(classification, pathname);

  // Relax the CORS origin fallback ONLY for the token-authenticated API
  // surface (CLIENT_API: /v1/*, /v1beta/*, codex/responses aliases) and
  // read-only PUBLIC endpoints. These authenticate via Authorization /
  // x-api-key headers that browsers never auto-attach, so echoing the caller's
  // Origin (or `*`) there carries no credentialed-session / CSRF risk — it just
  // lets browser/Electron clients (issue #5242) read responses they are already
  // entitled to. MANAGEMENT (cookie-authed dashboard) and non-read-only PUBLIC
  // routes (e.g. /api/cloud/, which sets Allow-Credentials in its own handler)
  // stay exactly fail-closed.
  const corsRelaxOrigin =
    classification.routeClass === "CLIENT_API" ||
    (classification.routeClass === "PUBLIC" && classification.reason === "public_readonly_prefix");

  if (guardedPathname.startsWith("/api/") && isDraining()) {
    const response = drainingResponse(requestId);
    stampRouteResponse(response, requestId, classification.routeClass);
    applyCorsHeaders(response, request, corsRelaxOrigin);
    return response;
  }

  if (guardedPathname.startsWith("/api/") && method !== "GET" && method !== "OPTIONS") {
    const bodySizeSettings = await getBodySizeSettings();
    const bodySizeRejection = checkBodySize(
      request,
      getBodySizeLimit(guardedPathname, bodySizeSettings)
    );
    if (bodySizeRejection) {
      stampRouteResponse(bodySizeRejection, requestId, classification.routeClass);
      applyCorsHeaders(bodySizeRejection, request, corsRelaxOrigin);
      return bodySizeRejection;
    }
  }

  const requestHeaders = new Headers(request.headers);
  for (const trusted of AUTHZ_TRUSTED_HEADERS) {
    requestHeaders.delete(trusted);
  }
  // The trusted peer-IP + via-proxy stamps are read by the policy from the
  // ORIGINAL request (above); strip them from the forwarded headers so the
  // per-process token never reaches route handlers or upstream providers.
  requestHeaders.delete(PEER_IP_HEADER);
  requestHeaders.delete(VIA_PROXY_HEADER);

  requestHeaders.set(AUTHZ_HEADER_ROUTE_CLASS, classification.routeClass);
  requestHeaders.set(AUTHZ_HEADER_REQUEST_ID, requestId);
  // Stamp a trusted, non-secret locality verdict derived from the real stamped
  // peer IP AND the via-proxy marker. Route handlers (e.g. cliTokenAuth) read
  // this instead of re-deriving locality from the spoofable Host header. The
  // client-supplied values (if any) were already removed by the
  // AUTHZ_TRUSTED_HEADERS strip above. When the via-proxy marker is set, a
  // loopback socket is the proxy hop, not the end-user — verdict is downgraded
  // to "remote" so the LOCAL_ONLY gate is not bypassed by a request arriving
  // through an external reverse proxy (nginx / Caddy / Cloudflare Tunnel).
  // See peerStamp.ts and the upstream da667836 reference for the full rationale.
  const peerLocality = classifyStampedPeerLocality(
    request.headers.get(PEER_IP_HEADER),
    request.headers.get(VIA_PROXY_HEADER),
    process.env.OMNIROUTE_PEER_STAMP_TOKEN
  );
  requestHeaders.set(AUTHZ_HEADER_PEER_LOCALITY, peerLocality);

  if (method === "OPTIONS") {
    const preflight = new NextResponse(null, { status: 204 });
    preflight.headers.set(AUTHZ_HEADER_REQUEST_ID, requestId);
    preflight.headers.set(AUTHZ_HEADER_ROUTE_CLASS, classification.routeClass);
    applyCorsHeaders(preflight, request, corsRelaxOrigin);
    return preflight;
  }

  if (!options.enforce) {
    const response = NextResponse.next({ request: { headers: requestHeaders } });
    response.headers.set(AUTHZ_HEADER_REQUEST_ID, requestId);
    response.headers.set(AUTHZ_HEADER_ROUTE_CLASS, classification.routeClass);
    applyCorsHeaders(response, request, corsRelaxOrigin);
    return response;
  }

  // IP filter (#6131): enforce the operator's IP blacklist/whitelist on the
  // external surface. Loopback is exempt so the local operator can never lock
  // themselves out of the dashboard (they can always fix the list from
  // localhost). checkIP is a no-op when the filter is disabled.
  //
  // D1 (#9033): on a direct connection the proxy runtime has no socket, so
  // checkRequestIP reads only forwarding headers + undefined request.ip and
  // falls to "unknown", never blocking the blacklisted client. Resolve the
  // trusted peer IP from the authenticated stamp and pass it to checkRequestIP,
  // but only when NOT behind a reverse proxy (the via-proxy marker means the
  // peer IP is the proxy hop, e.g. 127.0.0.1, and the real client is in XFF).
  if (peerLocality !== "loopback") {
    const trustedPeerIp = resolveStampedPeer(
      request.headers.get(PEER_IP_HEADER),
      process.env.OMNIROUTE_PEER_STAMP_TOKEN
    );
    const viaProxy = resolveStampedViaProxy(
      request.headers.get(VIA_PROXY_HEADER),
      process.env.OMNIROUTE_PEER_STAMP_TOKEN
    );
    const ipVerdict = checkRequestIP(request, viaProxy ? null : trustedPeerIp);
    if (!ipVerdict.allowed) {
      const blocked = NextResponse.json(
        { error: ipVerdict.reason || "Access denied" },
        { status: 403 }
      );
      stampRouteResponse(blocked, requestId, classification.routeClass);
      applyCorsHeaders(blocked, request, corsRelaxOrigin);
      return blocked;
    }
  }

  const policy = POLICIES[classification.routeClass];
  const outcome = await policy.evaluate({ request, classification, requestId });

  if (!outcome.allow) {
    if (managementDashboardRoute) {
      return dashboardLoginRedirect(request, requestId);
    }

    const rejection = rejectionResponse(outcome, classification, requestId);
    applyCorsHeaders(rejection, request, corsRelaxOrigin);
    return rejection;
  }

  if (
    classification.routeClass === "MANAGEMENT" &&
    outcome.subject.kind === "dashboard_session" &&
    isUnsafeMutationMethod(method)
  ) {
    const originVerdict = validateBrowserMutationOrigin(request);
    const csrfOriginFallback =
      originVerdict.reason === "invalid-origin" && validateDashboardCsrfToken(request);
    if (!originVerdict.ok && !csrfOriginFallback) {
      const rejection = invalidOriginResponse(requestId);
      rejection.headers.set(AUTHZ_HEADER_ROUTE_CLASS, classification.routeClass);
      applyCorsHeaders(rejection, request);
      return rejection;
    }
  }

  stampSubject(requestHeaders, outcome.subject);

  const response = NextResponse.next({ request: { headers: requestHeaders } });
  response.headers.set(AUTHZ_HEADER_REQUEST_ID, requestId);
  response.headers.set(AUTHZ_HEADER_ROUTE_CLASS, classification.routeClass);
  applyCorsHeaders(response, request, corsRelaxOrigin);
  if (managementDashboardRoute) {
    await refreshDashboardSessionIfNeeded(response, request);
  }
  return response;
}