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

/**
 * Security-weighted tests for the Cognito auth layer.
 *
 * The token claims here mirror the shape the 0a spike actually observed from
 * live Cognito (access token: prefixed `scope` string, `client_id`, `token_use`,
 * no native `aud`). The audience-boundary, scope-isolation, TTL-ceiling and
 * revocation cases are the trust-tier guarantees from design.md §2 — they are
 * the reason this file carries the heaviest coverage in the platform.
 */

import { describe, it, expect } from 'vitest';
import { generateKeyPair, SignJWT, type JWTVerifyGetKey } from 'jose';

import {
  CognitoAuthProvider,
  AuthError,
  extractScopes,
  extractBearerToken,
  cognitoIssuer,
  type CognitoAuthConfig,
} from './cognito-auth.js';
import { requireScope } from './auth.js';

type KeyPair = Awaited<ReturnType<typeof generateKeyPair>>;

const ISSUER = cognitoIssuer('us-east-1', 'us-east-1_TESTPOOL');
const OPS_CLIENT = 'ops-app-client-id';
const FIN_CLIENT = 'finance-app-client-id';

// Signing keys for the suite, plus a second pair to forge bad signatures.
const signing: KeyPair = await generateKeyPair('RS256');
const attacker: KeyPair = await generateKeyPair('RS256');

/** A JWKS resolver that returns our test public key (stands in for the pool's JWKS). */
const jwks: JWTVerifyGetKey = async () => signing.publicKey;

function baseConfig(overrides: Partial<CognitoAuthConfig> = {}): CognitoAuthConfig {
  return {
    issuer: ISSUER,
    audience: 'sh-mcp-ops',
    allowedClientIds: [OPS_CLIENT],
    scopePrefix: 'sh-mcp-ops',
    jwks,
    ...overrides,
  };
}

function financeConfig(overrides: Partial<CognitoAuthConfig> = {}): CognitoAuthConfig {
  return baseConfig({
    audience: 'sh-mcp-finance',
    allowedClientIds: [FIN_CLIENT],
    scopePrefix: 'sh-mcp-finance',
    maxTtlSeconds: 900,
    ttlGuardedScopes: ['finance:read', 'finance:admin'],
    ...overrides,
  });
}

interface MintOpts {
  issuer?: string;
  clientId?: string;
  scope?: string;
  tokenUse?: string;
  sub?: string;
  iat?: number;
  ttlSeconds?: number;
  signer?: KeyPair['privateKey'];
}

async function mint(opts: MintOpts = {}): Promise<string> {
  const now = Math.floor(Date.now() / 1000);
  const iat = opts.iat ?? now;
  const ttl = opts.ttlSeconds ?? 3600;
  return new SignJWT({
    token_use: opts.tokenUse ?? 'access',
    client_id: opts.clientId ?? OPS_CLIENT,
    scope: opts.scope ?? 'sh-mcp-ops/ops:read sh-mcp-ops/ops:tasks openid email',
  })
    .setProtectedHeader({ alg: 'RS256' })
    .setSubject(opts.sub ?? 'user-sub-123')
    .setIssuer(opts.issuer ?? ISSUER)
    .setIssuedAt(iat)
    .setExpirationTime(iat + ttl)
    .sign(opts.signer ?? signing.privateKey);
}

describe('CognitoAuthProvider.authenticate', () => {
  it('accepts a valid access token and extracts sub + this tier’s scopes', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    const ctx = await provider.authenticate(`Bearer ${await mint()}`);
    expect(ctx.sub).toBe('user-sub-123');
    expect(ctx.aud).toBe('sh-mcp-ops');
    expect(ctx.scopes).toEqual(['ops:read', 'ops:tasks']);
  });

  it('drops cross-tier and standard (openid/email) scopes', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    const token = await mint({
      scope: 'sh-mcp-ops/ops:read sh-mcp-finance/finance:read openid email',
    });
    const ctx = await provider.authenticate(`Bearer ${token}`);
    expect(ctx.scopes).toEqual(['ops:read']);
  });

  it('rejects a token from the wrong issuer', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    const token = await mint({ issuer: 'https://evil.example.com/pool' });
    await expect(provider.authenticate(`Bearer ${token}`)).rejects.toMatchObject({
      code: 'invalid_token',
    });
  });

  it('AUDIENCE BOUNDARY: rejects an ops token presented to the finance server', async () => {
    const finance = new CognitoAuthProvider(financeConfig());
    const opsToken = await mint({ clientId: OPS_CLIENT, scope: 'sh-mcp-ops/ops:read' });
    await expect(finance.authenticate(`Bearer ${opsToken}`)).rejects.toMatchObject({
      code: 'client_not_allowed',
    });
  });

  it('rejects an id token (token_use !== "access")', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    const token = await mint({ tokenUse: 'id' });
    await expect(provider.authenticate(`Bearer ${token}`)).rejects.toMatchObject({
      code: 'invalid_token',
    });
  });

  it('rejects an expired token', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    const now = Math.floor(Date.now() / 1000);
    const token = await mint({ iat: now - 7200, ttlSeconds: 3600 }); // expired ~1h ago
    await expect(provider.authenticate(`Bearer ${token}`)).rejects.toMatchObject({
      code: 'invalid_token',
    });
  });

  it('rejects a token signed by an unknown key (forged signature)', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    const token = await mint({ signer: attacker.privateKey });
    await expect(provider.authenticate(`Bearer ${token}`)).rejects.toMatchObject({
      code: 'invalid_token',
    });
  });

  it('rejects a missing Authorization header', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    await expect(provider.authenticate({ headers: {} })).rejects.toMatchObject({
      code: 'missing_token',
    });
  });

  it('FINANCE TTL: rejects a finance-scoped token whose lifetime exceeds the ceiling', async () => {
    const finance = new CognitoAuthProvider(financeConfig());
    const longToken = await mint({
      clientId: FIN_CLIENT,
      scope: 'sh-mcp-finance/finance:read',
      ttlSeconds: 3600,
    });
    await expect(finance.authenticate(`Bearer ${longToken}`)).rejects.toMatchObject({
      code: 'ttl_exceeded',
    });
  });

  it('FINANCE TTL: accepts a finance-scoped token within the ceiling', async () => {
    const finance = new CognitoAuthProvider(financeConfig());
    const shortToken = await mint({
      clientId: FIN_CLIENT,
      scope: 'sh-mcp-finance/finance:read',
      ttlSeconds: 600,
    });
    const ctx = await finance.authenticate(`Bearer ${shortToken}`);
    expect(ctx.scopes).toEqual(['finance:read']);
  });

  it('does not apply the TTL ceiling to non-guarded scopes', async () => {
    const finance = new CognitoAuthProvider(financeConfig());
    // ops:read carried under the finance prefix is known but not TTL-guarded.
    const longToken = await mint({
      clientId: FIN_CLIENT,
      scope: 'sh-mcp-finance/ops:read',
      ttlSeconds: 3600,
    });
    const ctx = await finance.authenticate(`Bearer ${longToken}`);
    expect(ctx.scopes).toEqual(['ops:read']);
  });

  it('rejects a revoked (deny-listed) user', async () => {
    const provider = new CognitoAuthProvider(
      baseConfig({ denyList: { isDenied: async (sub) => sub === 'revoked-user' } }),
    );
    const token = await mint({ sub: 'revoked-user' });
    await expect(provider.authenticate(`Bearer ${token}`)).rejects.toMatchObject({
      code: 'revoked',
    });
  });

  it('returns a context that satisfies requireScope for granted scopes only', async () => {
    const provider = new CognitoAuthProvider(baseConfig());
    const ctx = await provider.authenticate(`Bearer ${await mint()}`);
    expect(() => requireScope(ctx, 'ops:read')).not.toThrow();
    expect(() => requireScope(ctx, 'finance:read')).toThrow();
  });
});

describe('extractScopes', () => {
  it('strips the tier prefix and keeps known scopes in order', () => {
    expect(extractScopes('sh-mcp-ops/ops:read sh-mcp-ops/ops:tasks', 'sh-mcp-ops')).toEqual([
      'ops:read',
      'ops:tasks',
    ]);
  });

  it('drops other tiers and standard scopes', () => {
    expect(
      extractScopes('sh-mcp-ops/ops:read sh-mcp-finance/finance:read openid email', 'sh-mcp-ops'),
    ).toEqual(['ops:read']);
  });

  it('drops prefixed-but-unknown scopes', () => {
    expect(extractScopes('sh-mcp-ops/bogus:scope', 'sh-mcp-ops')).toEqual([]);
  });

  it('de-duplicates', () => {
    expect(extractScopes('sh-mcp-ops/ops:read sh-mcp-ops/ops:read', 'sh-mcp-ops')).toEqual([
      'ops:read',
    ]);
  });

  it('handles empty / non-string input', () => {
    expect(extractScopes('', 'sh-mcp-ops')).toEqual([]);
    expect(extractScopes(undefined, 'sh-mcp-ops')).toEqual([]);
    expect(extractScopes(null, 'sh-mcp-ops')).toEqual([]);
  });
});

describe('extractBearerToken', () => {
  it('reads a raw Authorization header string', () => {
    expect(extractBearerToken('Bearer abc.def.ghi')).toBe('abc.def.ghi');
  });

  it('is case-insensitive on the scheme', () => {
    expect(extractBearerToken('bearer abc')).toBe('abc');
  });

  it('reads a plain headers object (either header casing)', () => {
    expect(extractBearerToken({ headers: { authorization: 'Bearer xyz' } })).toBe('xyz');
    expect(extractBearerToken({ headers: { Authorization: 'Bearer XYZ' } })).toBe('XYZ');
  });

  it('reads a Fetch Headers-like object', () => {
    const headers = new Headers({ authorization: 'Bearer fetchtoken' });
    expect(extractBearerToken({ headers })).toBe('fetchtoken');
  });

  it('throws AuthError on a missing header', () => {
    expect(() => extractBearerToken({ headers: {} })).toThrow(AuthError);
  });

  it('throws AuthError on a non-Bearer header', () => {
    expect(() => extractBearerToken('Basic abc')).toThrow(AuthError);
  });

  // ReDoS guard (CodeQL js/polynomial-redos): the matcher is /\s+(\S.*)/, not
  // the ambiguous /\s+(.+)/. These pin the behavior that the `\S` fix preserves.
  it('still captures a token that follows multiple separating spaces', () => {
    expect(extractBearerToken('Bearer    abc.def')).toBe('abc.def');
  });

  it('rejects a "Bearer" header with no token after the whitespace', () => {
    expect(() => extractBearerToken(`Bearer ${' '.repeat(5_000)}`)).toThrow(AuthError);
  });
});