sentryagent-idp/tests/unit/services/IDTokenService.test.ts

/**
 * Unit tests for src/services/IDTokenService.ts
 * Mocks OIDCKeyService; uses real RSA key pairs for signing/verification.
 */

import crypto from 'crypto';
import jwt from 'jsonwebtoken';
import { IDTokenService } from '../../../src/services/IDTokenService';
import { OIDCKeyService } from '../../../src/services/OIDCKeyService';
import { IAgent } from '../../../src/types/index';
import { IOIDCKey, IJWKSKey, IJWKSResponse } from '../../../src/types/oidc';

// ─── Real RSA key pair for signing tests ─────────────────────────────────────

const { privateKey: rsaPrivateKey, publicKey: rsaPublicKey } = crypto.generateKeyPairSync('rsa', {
  modulusLength: 2048,
});

const privateKeyPem = rsaPrivateKey.export({ format: 'pem', type: 'pkcs8' }) as string;
const publicKeyJwkRaw = rsaPublicKey.export({ format: 'jwk' }) as Record<string, string>;

const TEST_KID = 'key-test-rsa-001';

const testPublicJwk: IJWKSKey = {
  kid: TEST_KID,
  kty: publicKeyJwkRaw['kty'] ?? 'RSA',
  use: 'sig',
  alg: 'RS256',
  n: publicKeyJwkRaw['n'],
  e: publicKeyJwkRaw['e'],
};

const mockCurrentKey: IOIDCKey = {
  id: 'uuid-test',
  kid: TEST_KID,
  algorithm: 'RS256',
  public_key_jwk: testPublicJwk,
  vault_key_path: 'dev:no-vault',
  is_current: true,
  created_at: new Date(),
  expires_at: new Date(Date.now() + 3600 * 1000),
};

const mockJwks: IJWKSResponse = { keys: [testPublicJwk] };

// ─── Mock OIDCKeyService ──────────────────────────────────────────────────────

jest.mock('../../../src/services/OIDCKeyService');

const MockOIDCKeyService = OIDCKeyService as jest.MockedClass<typeof OIDCKeyService>;

// ─── Mock agent ───────────────────────────────────────────────────────────────

const MOCK_AGENT: IAgent = {
  agentId: 'agent-uuid-001',
  organizationId: 'org-uuid-001',
  email: 'agent@sentryagent.ai',
  agentType: 'screener',
  version: '1.0.0',
  capabilities: ['agents:read'],
  owner: 'team-alpha',
  deploymentEnv: 'production',
  status: 'active',
  createdAt: new Date(),
  updatedAt: new Date(),
};

// ─── Tests ───────────────────────────────────────────────────────────────────

describe('IDTokenService', () => {
  let mockKeyService: jest.Mocked<OIDCKeyService>;
  let service: IDTokenService;

  beforeEach(() => {
    jest.clearAllMocks();
    delete process.env['OIDC_ISSUER'];
    delete process.env['OIDC_ID_TOKEN_TTL_SECONDS'];

    mockKeyService = new MockOIDCKeyService(
      null as unknown as import('pg').Pool,
      null as unknown as import('redis').RedisClientType,
    ) as jest.Mocked<OIDCKeyService>;

    mockKeyService.getCurrentKey.mockResolvedValue(mockCurrentKey);
    mockKeyService.getPrivateKeyPem.mockResolvedValue(privateKeyPem);
    mockKeyService.getPublicJWKS.mockResolvedValue(mockJwks);

    service = new IDTokenService(mockKeyService);
  });

  // ── buildIDTokenClaims ───────────────────────────────────────────────────

  describe('buildIDTokenClaims()', () => {
    it('includes all required OIDC claims', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid agents:read');

      expect(claims.iss).toBe('https://idp.sentryagent.ai');
      expect(claims.sub).toBe(MOCK_AGENT.agentId);
      expect(claims.aud).toBe('client-abc');
      expect(claims.iat).toBeDefined();
      expect(claims.exp).toBeDefined();
      expect(claims.exp).toBeGreaterThan(claims.iat);
    });

    it('includes agent-specific claims', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');

      expect(claims.agent_type).toBe(MOCK_AGENT.agentType);
      expect(claims.deployment_env).toBe(MOCK_AGENT.deploymentEnv);
      expect(claims.organization_id).toBe(MOCK_AGENT.organizationId);
    });

    it('includes nonce when provided', async () => {
      const claims = await service.buildIDTokenClaims(
        MOCK_AGENT,
        'client-abc',
        'openid',
        'test-nonce-xyz',
      );
      expect(claims.nonce).toBe('test-nonce-xyz');
    });

    it('omits nonce when not provided', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      expect(claims.nonce).toBeUndefined();
    });

    it('includes did when the agent has a DID', async () => {
      const agentWithDID: IAgent = {
        ...MOCK_AGENT,
        did: 'did:web:idp.sentryagent.ai:agents:agent-uuid-001',
      };
      const claims = await service.buildIDTokenClaims(agentWithDID, 'client-abc', 'openid');
      expect(claims.did).toBe(agentWithDID.did);
    });

    it('omits did when the agent does not have a DID', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      expect(claims.did).toBeUndefined();
    });

    it('uses OIDC_ISSUER env var when set', async () => {
      process.env['OIDC_ISSUER'] = 'https://my-custom-issuer.example.com';
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      expect(claims.iss).toBe('https://my-custom-issuer.example.com');
    });

    it('uses OIDC_ID_TOKEN_TTL_SECONDS for expiry', async () => {
      process.env['OIDC_ID_TOKEN_TTL_SECONDS'] = '7200';
      const before = Math.floor(Date.now() / 1000);
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      const after = Math.floor(Date.now() / 1000);
      expect(claims.exp - claims.iat).toBeGreaterThanOrEqual(7200 - 1);
      expect(claims.exp).toBeGreaterThanOrEqual(before + 7200);
      expect(claims.exp).toBeLessThanOrEqual(after + 7200 + 1);
    });
  });

  // ── signIDToken ──────────────────────────────────────────────────────────

  describe('signIDToken()', () => {
    it('produces a valid JWT string', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      const token = await service.signIDToken(claims);

      expect(typeof token).toBe('string');
      expect(token.split('.')).toHaveLength(3);
    });

    it('includes the kid in the JWT header', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      const token = await service.signIDToken(claims);

      const decoded = jwt.decode(token, { complete: true });
      expect(decoded).not.toBeNull();
      expect(decoded!.header.kid).toBe(TEST_KID);
    });

    it('signs with RS256 algorithm by default', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      const token = await service.signIDToken(claims);

      const decoded = jwt.decode(token, { complete: true });
      expect(decoded!.header.alg).toBe('RS256');
    });
  });

  // ── verifyIDToken ────────────────────────────────────────────────────────

  describe('verifyIDToken()', () => {
    it('verifies a valid ID token and returns claims', async () => {
      const claims = await service.buildIDTokenClaims(MOCK_AGENT, 'client-abc', 'openid');
      const token = await service.signIDToken(claims);

      const verified = await service.verifyIDToken(token);
      expect(verified.sub).toBe(MOCK_AGENT.agentId);
      expect(verified.iss).toBe('https://idp.sentryagent.ai');
      expect(verified.aud).toBe('client-abc');
    });

    it('rejects alg:none tokens', async () => {
      // Craft a token with alg:none manually
      const header = Buffer.from(JSON.stringify({ alg: 'none', typ: 'JWT' })).toString('base64url');
      const payload = Buffer.from(JSON.stringify({ sub: 'attacker', iss: 'evil', aud: 'client', iat: 0, exp: 9999999999 })).toString('base64url');
      const noneToken = `${header}.${payload}.`;

      await expect(service.verifyIDToken(noneToken)).rejects.toMatchObject({
        code: 'ID_TOKEN_INVALID',
      });
    });

    it('rejects expired tokens', async () => {
      // Sign a token that was issued in the past and has already expired
      const now = Math.floor(Date.now() / 1000);
      const expiredClaims = {
        iss: 'https://idp.sentryagent.ai',
        sub: MOCK_AGENT.agentId,
        aud: 'client-abc',
        iat: now - 7200,
        exp: now - 3600,
        agent_type: 'screener',
        deployment_env: 'production',
        organization_id: 'org-uuid-001',
      };

      // Sign with the real private key (same as what mockKeyService returns)
      const expiredToken = jwt.sign(expiredClaims, privateKeyPem, {
        algorithm: 'RS256',
        header: { alg: 'RS256', kid: TEST_KID, typ: 'JWT' },
      });

      await expect(service.verifyIDToken(expiredToken)).rejects.toMatchObject({
        code: 'ID_TOKEN_INVALID',
      });
    });

    it('rejects tokens whose kid is not in the JWKS', async () => {
      const unknownKidHeader = Buffer.from(JSON.stringify({ alg: 'RS256', kid: 'unknown-kid', typ: 'JWT' })).toString('base64url');
      const payload = Buffer.from(JSON.stringify({ sub: 'x' })).toString('base64url');
      const fakeToken = `${unknownKidHeader}.${payload}.fakesig`;

      await expect(service.verifyIDToken(fakeToken)).rejects.toMatchObject({
        code: 'ID_TOKEN_INVALID',
      });
    });

    it('rejects malformed tokens', async () => {
      await expect(service.verifyIDToken('not.a.jwt')).rejects.toMatchObject({
        code: 'ID_TOKEN_INVALID',
      });
    });

    it('rejects tokens with missing kid header', async () => {
      // Craft a token with no kid in the header
      const header = Buffer.from(JSON.stringify({ alg: 'RS256', typ: 'JWT' })).toString('base64url');
      const payload = Buffer.from(JSON.stringify({ sub: 'x', exp: 9999999999 })).toString('base64url');
      const noKidToken = `${header}.${payload}.fakesig`;

      await expect(service.verifyIDToken(noKidToken)).rejects.toMatchObject({
        code: 'ID_TOKEN_INVALID',
      });
    });

    it('rejects tokens with unsupported algorithm (HS256)', async () => {
      // Craft a token header claiming HS256
      const header = Buffer.from(JSON.stringify({ alg: 'HS256', kid: TEST_KID, typ: 'JWT' })).toString('base64url');
      const payload = Buffer.from(JSON.stringify({ sub: 'x', exp: 9999999999 })).toString('base64url');
      const fakeToken = `${header}.${payload}.fakesig`;

      await expect(service.verifyIDToken(fakeToken)).rejects.toMatchObject({
        code: 'ID_TOKEN_INVALID',
      });
    });
  });
});