seahaven-ap/packages/api/src/auth/middleware.ts

import { createMiddleware } from "hono/factory";
import { createRemoteJWKSet, jwtVerify, type JWTPayload } from "jose";
import type { AuthUser } from "./upsert-user.js";
import { IdentityConflictError, upsertUserFromIdentity } from "./upsert-user.js";
import type { ApiEnv, UserRole } from "../env.js";
import { isUserRole } from "../env.js";
import type { Db } from "../db/client.js";
import { errorJson } from "../http.js";
import { cookieEventFromHeader, sessionCookieValue } from "./cookies.js";
import { cookieStage, isPublicRoute } from "./oauth.js";

export type AppVariables = {
  user: AuthUser;
};

export type AppBindings = {
  Variables: AppVariables;
};

export type TokenVerifier = (token: string) => Promise<JWTPayload>;

export type AuthDeps = {
  verifyToken?: TokenVerifier;
};

function roleFromClaims(claims: Record<string, unknown>): UserRole | undefined {
  const raw =
    (typeof claims["custom:role"] === "string" && claims["custom:role"]) ||
    (typeof claims.role === "string" && claims.role) ||
    undefined;
  if (raw === undefined) return undefined;
  return isUserRole(raw) ? raw : undefined;
}

function audienceMatches(payload: JWTPayload, expected: string): boolean {
  const claims = payload as JWTPayload & { token_use?: string; client_id?: string };
  if (claims.token_use === "access") {
    return claims.client_id === expected;
  }

  if (typeof payload.aud === "string") {
    return payload.aud === expected;
  }
  if (Array.isArray(payload.aud)) {
    return payload.aud.includes(expected);
  }
  return false;
}

function identityFromPayload(payload: JWTPayload): {
  sub: string;
  email: string;
  name: string;
} | null {
  const sub = typeof payload.sub === "string" ? payload.sub : null;
  const email = typeof payload.email === "string" ? payload.email : null;
  const name =
    (typeof payload.name === "string" && payload.name) ||
    (typeof payload["cognito:username"] === "string" && payload["cognito:username"]) ||
    email;

  if (!sub || !email || !name) {
    return null;
  }
  return { sub, email, name };
}

export function createAuthMiddleware(env: ApiEnv, handle: Db, deps: AuthDeps = {}) {
  const jwks =
    env.cognitoIssuer.length > 0
      ? createRemoteJWKSet(new URL(`${env.cognitoIssuer}/.well-known/jwks.json`))
      : null;

  const verifyToken: TokenVerifier =
    deps.verifyToken ??
    (async (token) => {
      if (!jwks) {
        throw new Error("JWT verification is not configured.");
      }
      const { payload } = await jwtVerify(token, jwks, {
        issuer: env.cognitoIssuer,
      });
      return payload;
    });

  return createMiddleware<AppBindings>(async (c, next) => {
    const path = new URL(c.req.url).pathname;
    if (isPublicRoute(c.req.method, path)) {
      await next();
      return;
    }

    if (env.devAuthBypass) {
      try {
        const user = await upsertUserFromIdentity(handle, {
          cognitoSub: env.devAuthSub,
          email: env.devAuthEmail,
          name: env.devAuthName,
          role: env.devAuthRole,
        });
        c.set("user", user);
      } catch (error) {
        if (error instanceof IdentityConflictError) {
          return errorJson(c, 409, "CONFLICT", error.message);
        }
        throw error;
      }
      await next();
      return;
    }

    const stage = cookieStage(env);
    const idToken = sessionCookieValue(cookieEventFromHeader(c.req.header("cookie")), "id", stage);
    if (!idToken) {
      return errorJson(c, 401, "UNAUTHENTICATED", "Missing id token.");
    }

    let payload: JWTPayload;
    try {
      payload = await verifyToken(idToken);
    } catch {
      return errorJson(c, 401, "UNAUTHENTICATED", "Invalid or expired token.");
    }

    if (!audienceMatches(payload, env.cognitoAudience)) {
      return errorJson(c, 401, "UNAUTHENTICATED", "Token audience does not match this API.");
    }

    const identity = identityFromPayload(payload);
    if (!identity) {
      return errorJson(
        c,
        401,
        "UNAUTHENTICATED",
        "Token is missing required identity claims. Use a Cognito ID token or an access token that includes email.",
      );
    }

    let user: AuthUser;
    try {
      user = await upsertUserFromIdentity(handle, {
        cognitoSub: identity.sub,
        email: identity.email,
        name: identity.name,
        role: roleFromClaims(payload as Record<string, unknown>),
      });
    } catch (error) {
      if (error instanceof IdentityConflictError) {
        return errorJson(c, 409, "CONFLICT", error.message);
      }
      throw error;
    }

    c.set("user", user);
    await next();
  });
}