sh-mcp/packages/shared/src/cognito-auth.ts

/**
 * Concrete `AuthProvider` for Amazon Cognito access tokens.
 *
 * This is the real implementation the 0a spike unblocked. The spike proved the
 * exact token shape we receive (see SPIKE findings / design.md §2): a Cognito
 * **access** token carries `sub`, `client_id`, `scope` (a space-separated,
 * resource-server-prefixed string), `iss`, `token_use`, `iat`, `exp` — and
 * crucially **no native `aud` claim and no `email`**. Two consequences drive
 * this implementation:
 *
 *   1. Audience binding cannot use the JWT `aud` claim. Instead the
 *      **`client_id` allow-list IS the audience boundary** — each trust tier
 *      gets its own Cognito app client, and a server only accepts tokens minted
 *      by app clients on its `allowedClientIds`. An ops token presented to the
 *      finance server is rejected because the ops app-client id is not on
 *      finance's allow-list (design.md §2.5).
 *   2. Scopes arrive prefixed with the resource server ("sh-mcp-ops/ops:read").
 *      We accept only scopes carrying this server's `scopePrefix`, strip the
 *      prefix to the internal `Scope` ("ops:read"), and drop everything else —
 *      so a cross-tier scope can never leak into an AuthContext.
 *
 * The class verifies the JWT signature against the pool's JWKS, validates the
 * issuer, enforces the client_id allow-list, applies the optional finance TTL
 * ceiling and deny-list, and returns a populated `AuthContext`. Scope-per-tool
 * enforcement still happens in each handler via `requireScope()` (see auth.ts).
 *
 * Config (client_id allow-list, audience, scope prefix, TTL rule) is INJECTED,
 * never hardcoded — it comes from each server's SSM/CDK env (design.md §2,
 * memory: "the client/audience matrix is config, not hardcoded into the core").
 */

import { jwtVerify, createRemoteJWKSet, type JWTVerifyGetKey, type JWTPayload } from 'jose';

import type { AuthContext, Scope } from './types.js';
import type { AuthProvider } from './auth.js';

// ---------------------------------------------------------------------------
// AuthError
// ---------------------------------------------------------------------------

/**
 * Reason an inbound token was rejected. Distinct from `ScopeError`
 * (auth.ts): an `AuthError` is an authentication failure (→ HTTP 401),
 * whereas a `ScopeError` is an authorization failure on a valid identity
 * (→ HTTP 403).
 */
export type AuthErrorCode =
  | 'missing_token' // no / malformed Authorization header
  | 'invalid_token' // bad signature, wrong issuer, not an access token, no sub
  | 'client_not_allowed' // client_id absent from this server's allow-list (audience boundary)
  | 'ttl_exceeded' // token lifetime exceeds the ceiling for a guarded scope (finance)
  | 'revoked'; // user is on the deny-list

/** Thrown by `CognitoAuthProvider.authenticate()` on any authentication failure. */
export class AuthError extends Error {
  readonly code: AuthErrorCode;

  constructor(code: AuthErrorCode, message: string) {
    super(message);
    this.name = 'AuthError';
    this.code = code;
    // Maintain proper prototype chain for `instanceof` checks.
    Object.setPrototypeOf(this, new.target.prototype);
  }
}

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

/** Optional immediate-revocation check (design.md §2.3 — Cognito-backed deny-list). */
export interface DenyListChecker {
  /** Resolve `true` if `sub` has been hard-revoked and must be rejected now. */
  isDenied(sub: string): Promise<boolean>;
}

/**
 * Per-server configuration for {@link CognitoAuthProvider}. Injected from each
 * server's environment — never hardcoded into the shared core.
 */
export interface CognitoAuthConfig {
  /**
   * Expected token issuer — the Cognito user pool URL,
   * e.g. `https://cognito-idp.us-east-1.amazonaws.com/us-east-1_GsDbGe0pa`.
   * Tokens with any other `iss` are rejected. Use {@link cognitoIssuer}.
   */
  issuer: string;
  /**
   * This server's resource-server identifier, written into `AuthContext.aud`
   * after the client_id allow-list passes (e.g. `"sh-mcp-ops"`). Because the
   * token carries no native `aud`, this is the server's asserted audience, not
   * a value read from the token.
   */
  audience: string;
  /**
   * Cognito app-client ids permitted to call THIS server. This list IS the
   * audience boundary (the token has no `aud`): a token minted for another
   * tier's app client is rejected. Each tier has its own app client.
   */
  allowedClientIds: readonly string[];
  /**
   * Resource-server prefix for this tier, e.g. `"sh-mcp-ops"`. Cognito scopes
   * arrive as `"sh-mcp-ops/ops:read"`; only scopes with this prefix are
   * accepted, the prefix is stripped to the internal `Scope`, and all other
   * (cross-tier or standard `openid`/`email`) scopes are dropped.
   */
  scopePrefix: string;
  /**
   * JWKS key resolver used to verify the token signature. In production build
   * it with {@link cognitoJwks}; tests inject a local key set.
   */
  jwks: JWTVerifyGetKey;
  /**
   * If set, a token whose lifetime (`exp - iat`) exceeds this many seconds is
   * rejected when it carries any scope in {@link ttlGuardedScopes}. design.md
   * §2.5 requires `finance:*` tokens to be ≤ 15 min (900s).
   */
  maxTtlSeconds?: number;
  /** Scopes that trigger the {@link maxTtlSeconds} ceiling (e.g. the finance scopes). */
  ttlGuardedScopes?: readonly Scope[];
  /** Optional deny-list for immediate hard revocation of a `sub`. */
  denyList?: DenyListChecker;
  /** Clock skew tolerance in seconds for `exp`/`nbf` (default 5). */
  clockToleranceSeconds?: number;
}

// ---------------------------------------------------------------------------
// JWKS / issuer helpers
// ---------------------------------------------------------------------------

/** The Cognito issuer URL for a pool — use for {@link CognitoAuthConfig.issuer}. */
export function cognitoIssuer(region: string, userPoolId: string): string {
  return `https://cognito-idp.${region}.amazonaws.com/${userPoolId}`;
}

/**
 * Build a cached remote JWKS resolver for a Cognito user pool. The resolver
 * fetches and caches the pool's signing keys, refreshing on unknown `kid`.
 */
export function cognitoJwks(region: string, userPoolId: string): JWTVerifyGetKey {
  return createRemoteJWKSet(new URL(`${cognitoIssuer(region, userPoolId)}/.well-known/jwks.json`));
}

// ---------------------------------------------------------------------------
// Pure helpers (exported for unit testing)
// ---------------------------------------------------------------------------

/** Every internal scope the platform recognizes (cross-checked from types.ts). */
const KNOWN_SCOPES: readonly Scope[] = [
  'ops:read',
  'ops:tasks',
  'gmail:self',
  'calendar:self',
  'finance:read',
  'finance:admin',
];

/**
 * Parse a Cognito `scope` string into validated internal `Scope`s for one tier.
 *
 * Keeps only entries prefixed `"<scopePrefix>/"`, strips the prefix, and admits
 * the result only if it is a {@link KNOWN_SCOPES} value. Standard scopes
 * (`openid`, `email`) and other tiers' scopes are dropped. Order-preserving and
 * de-duplicated.
 */
export function extractScopes(rawScope: unknown, scopePrefix: string): Scope[] {
  if (typeof rawScope !== 'string' || rawScope.length === 0) return [];
  const wanted = `${scopePrefix}/`;
  const out: Scope[] = [];
  for (const entry of rawScope.split(/\s+/)) {
    if (!entry.startsWith(wanted)) continue;
    const bare = entry.slice(wanted.length);
    if ((KNOWN_SCOPES as readonly string[]).includes(bare) && !out.includes(bare as Scope)) {
      out.push(bare as Scope);
    }
  }
  return out;
}

/**
 * Pull the bearer token out of a transport-agnostic request. Accepts the raw
 * Authorization header string, a Fetch `Headers`-like object, or a plain
 * `{ headers: { authorization } }` shape (Node `IncomingMessage`, Lambda event).
 *
 * @throws {AuthError} `missing_token` if no usable Bearer token is present.
 */
export function extractBearerToken(req: unknown): string {
  const header = getAuthorizationHeader(req);
  if (!header) {
    throw new AuthError('missing_token', 'No Authorization header present.');
  }
  // `\s+(\S.*)` — NOT `\s+(.+)`: requiring the capture to start with a
  // non-whitespace char removes the quantifier overlap (both `\s` and `.`
  // match a space), which otherwise allows polynomial backtracking on a
  // crafted all-whitespace header (CodeQL js/polynomial-redos). Linear now.
  const match = /^Bearer\s+(\S.*)$/i.exec(header.trim());
  const token = match?.[1]?.trim();
  if (!token) {
    throw new AuthError('missing_token', 'Authorization header is not a Bearer token.');
  }
  return token;
}

function getAuthorizationHeader(req: unknown): string | undefined {
  if (typeof req === 'string') return req;
  if (req === null || typeof req !== 'object') return undefined;

  const headers = (req as { headers?: unknown }).headers;
  if (!headers || typeof headers !== 'object') return undefined;

  // Fetch `Headers`-like (has a .get method).
  const get = (headers as { get?: unknown }).get;
  if (typeof get === 'function') {
    const v = (headers as Headers).get('authorization');
    return v ?? undefined;
  }

  // Plain object headers (case-insensitive lookup, array-valued allowed).
  const h = headers as Record<string, unknown>;
  const v = h['authorization'] ?? h['Authorization'];
  if (typeof v === 'string') return v;
  if (Array.isArray(v) && typeof v[0] === 'string') return v[0];
  return undefined;
}

// ---------------------------------------------------------------------------
// CognitoAuthProvider
// ---------------------------------------------------------------------------

/**
 * Verifies a Cognito access token and resolves it to an {@link AuthContext}.
 *
 * Each server constructs one with its own config and calls `authenticate(req)`
 * once per inbound request, before any tool handler runs.
 */
export class CognitoAuthProvider implements AuthProvider {
  constructor(private readonly config: CognitoAuthConfig) {}

  async authenticate(req: unknown): Promise<AuthContext> {
    const token = extractBearerToken(req);

    let payload: JWTPayload & {
      token_use?: unknown;
      client_id?: unknown;
      scope?: unknown;
    };
    try {
      const result = await jwtVerify(token, this.config.jwks, {
        issuer: this.config.issuer,
        clockTolerance: this.config.clockToleranceSeconds ?? 5,
      });
      payload = result.payload;
    } catch (err) {
      throw new AuthError('invalid_token', `Token verification failed: ${(err as Error).message}`);
    }

    // Must be an access token — id tokens carry different claims and are not
    // the credential the agent callout presents.
    if (payload.token_use !== 'access') {
      throw new AuthError(
        'invalid_token',
        `Expected token_use "access", got "${String(payload.token_use)}".`,
      );
    }

    // client_id allow-list == audience boundary (token has no native aud).
    const clientId = typeof payload.client_id === 'string' ? payload.client_id : undefined;
    if (!clientId || !this.config.allowedClientIds.includes(clientId)) {
      throw new AuthError(
        'client_not_allowed',
        `client_id "${clientId ?? '(none)'}" is not permitted for audience "${this.config.audience}".`,
      );
    }

    const sub = typeof payload.sub === 'string' ? payload.sub : undefined;
    if (!sub) {
      throw new AuthError('invalid_token', 'Token has no "sub" claim.');
    }

    const scopes = extractScopes(payload.scope, this.config.scopePrefix);

    // Finance TTL ceiling (design.md §2.5): a token bearing a guarded scope must
    // be short-lived. exp/iat are validated numbers here (jwtVerify checked exp).
    if (this.config.maxTtlSeconds != null && this.config.ttlGuardedScopes?.length) {
      const carriesGuarded = scopes.some((s) => this.config.ttlGuardedScopes!.includes(s));
      if (carriesGuarded) {
        const iat = typeof payload.iat === 'number' ? payload.iat : undefined;
        const exp = typeof payload.exp === 'number' ? payload.exp : undefined;
        if (iat == null || exp == null || exp - iat > this.config.maxTtlSeconds) {
          throw new AuthError(
            'ttl_exceeded',
            `Token lifetime exceeds the ${this.config.maxTtlSeconds}s ceiling required for ` +
              `${this.config.ttlGuardedScopes!.join('/')} scopes.`,
          );
        }
      }
    }

    // Immediate hard revocation (checked last — most expensive, may hit DynamoDB).
    if (this.config.denyList && (await this.config.denyList.isDenied(sub))) {
      throw new AuthError('revoked', `User "${sub}" is on the deny-list (revoked).`);
    }

    return { sub, scopes, aud: this.config.audience };
  }
}