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

/**
 * Unit tests for src/services/OIDCKeyService.ts
 * Mocks pg Pool and node-vault; uses a real in-memory Redis stub.
 */

import { Pool } from 'pg';
import { RedisClientType } from 'redis';
import { OIDCKeyService } from '../../../src/services/OIDCKeyService';
import { IOIDCKey, IJWKSKey } from '../../../src/types/oidc';

// ─── Mocks ────────────────────────────────────────────────────────────────────

jest.mock('pg', () => {
  const mQuery = jest.fn();
  const mPool = { query: mQuery };
  return { Pool: jest.fn(() => mPool) };
});

jest.mock('node-vault', () => {
  return jest.fn(() => ({
    write: jest.fn().mockResolvedValue({}),
    read: jest.fn().mockResolvedValue({
      data: { data: { privateKeyPem: 'mock-pem' } },
    }),
    delete: jest.fn().mockResolvedValue({}),
  }));
});

// ─── Helpers ─────────────────────────────────────────────────────────────────

function makeRedis(): RedisClientType {
  const store = new Map<string, string>();
  return {
    get: jest.fn(async (key: string) => store.get(key) ?? null),
    set: jest.fn(async (key: string, value: string, _opts?: unknown) => {
      store.set(key, value);
      return 'OK';
    }),
    del: jest.fn(async (key: string) => {
      store.delete(key);
      return 1;
    }),
  } as unknown as RedisClientType;
}

function makeSampleJwk(kid = 'key-test-001'): IJWKSKey {
  return { kid, kty: 'RSA', use: 'sig', alg: 'RS256', n: 'abc', e: 'AQAB' };
}

function makeSampleRow(overrides: Partial<IOIDCKey> = {}): IOIDCKey {
  return {
    id: 'uuid-1',
    kid: 'key-test-001',
    algorithm: 'RS256',
    public_key_jwk: makeSampleJwk(),
    vault_key_path: 'dev:no-vault',
    is_current: true,
    created_at: new Date('2026-01-01T00:00:00Z'),
    expires_at: new Date(Date.now() + 3600 * 1000),
    ...overrides,
  };
}

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

describe('OIDCKeyService', () => {
  let pool: Pool;
  let poolQuery: jest.Mock;
  let redis: RedisClientType;
  let service: OIDCKeyService;

  beforeEach(() => {
    jest.clearAllMocks();
    delete process.env['VAULT_ADDR'];
    delete process.env['VAULT_TOKEN'];
    delete process.env['VAULT_MOUNT'];
    delete process.env['OIDC_KEY_ALGORITHM'];
    delete process.env['OIDC_ID_TOKEN_TTL_SECONDS'];

    pool = new Pool();
    poolQuery = pool.query as jest.Mock;
    redis = makeRedis();
    service = new OIDCKeyService(pool, redis);
  });

  // ── ensureCurrentKey ──────────────────────────────────────────────────────

  describe('ensureCurrentKey()', () => {
    it('generates a key when no current key exists', async () => {
      // COUNT returns 0 → no current key
      poolQuery
        .mockResolvedValueOnce({ rows: [{ count: '0' }] }) // ensureCurrentKey COUNT
        .mockResolvedValueOnce({ rows: [] })               // UPDATE demote old key
        .mockResolvedValueOnce({                           // INSERT new key
          rows: [
            {
              id: 'uuid-1',
              kid: 'key-abc',
              algorithm: 'RS256',
              public_key_jwk: makeSampleJwk('key-abc'),
              vault_key_path: 'dev:no-vault',
              is_current: true,
              created_at: new Date(),
              expires_at: new Date(Date.now() + 3600 * 1000),
            },
          ],
        });

      await service.ensureCurrentKey();

      // Should have called INSERT (i.e. generateSigningKeyPair was invoked)
      expect(poolQuery).toHaveBeenCalledTimes(3);
    });

    it('is idempotent — does not generate when a current key already exists', async () => {
      // COUNT returns 1 → key exists
      poolQuery.mockResolvedValueOnce({ rows: [{ count: '1' }] });

      await service.ensureCurrentKey();
      // Only the COUNT query should have been executed
      expect(poolQuery).toHaveBeenCalledTimes(1);
    });

    it('called twice only generates one key', async () => {
      // First call: no key → generates
      poolQuery
        .mockResolvedValueOnce({ rows: [{ count: '0' }] })
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({
          rows: [
            {
              id: 'uuid-1',
              kid: 'key-abc',
              algorithm: 'RS256',
              public_key_jwk: makeSampleJwk('key-abc'),
              vault_key_path: 'dev:no-vault',
              is_current: true,
              created_at: new Date(),
              expires_at: new Date(Date.now() + 3600 * 1000),
            },
          ],
        });

      // Second call: key now exists → no generation
      poolQuery.mockResolvedValueOnce({ rows: [{ count: '1' }] });

      await service.ensureCurrentKey();
      await service.ensureCurrentKey();

      // First call: COUNT + UPDATE + INSERT = 3; second call: COUNT = 1 → total 4
      expect(poolQuery).toHaveBeenCalledTimes(4);
    });
  });

  // ── generateSigningKeyPair ────────────────────────────────────────────────

  describe('generateSigningKeyPair()', () => {
    it('generates an RSA key pair and inserts into the database', async () => {
      const row = {
        id: 'uuid-2',
        kid: 'key-test-002',
        algorithm: 'RS256',
        public_key_jwk: makeSampleJwk('key-test-002'),
        vault_key_path: 'dev:no-vault',
        is_current: true,
        created_at: new Date(),
        expires_at: new Date(Date.now() + 3600 * 1000),
      };

      poolQuery
        .mockResolvedValueOnce({ rows: [] }) // UPDATE demote
        .mockResolvedValueOnce({ rows: [row] }); // INSERT

      const result = await service.generateSigningKeyPair();

      expect(result.algorithm).toBe('RS256');
      expect(result.is_current).toBe(true);
      expect(result.vault_key_path).toBe('dev:no-vault');
      expect(result.public_key_jwk.use).toBe('sig');
    });

    it('throws when OIDC_KEY_ALGORITHM is unsupported', async () => {
      process.env['OIDC_KEY_ALGORITHM'] = 'RS512';
      await expect(service.generateSigningKeyPair()).rejects.toThrow(
        'Unsupported OIDC_KEY_ALGORITHM',
      );
    });

    it('generates an EC P-256 key pair when OIDC_KEY_ALGORITHM=ES256', async () => {
      process.env['OIDC_KEY_ALGORITHM'] = 'ES256';

      const row = {
        id: 'uuid-3',
        kid: 'key-test-003',
        algorithm: 'ES256',
        public_key_jwk: { kid: 'key-test-003', kty: 'EC', use: 'sig', alg: 'ES256', crv: 'P-256', x: 'x', y: 'y' },
        vault_key_path: 'dev:no-vault',
        is_current: true,
        created_at: new Date(),
        expires_at: new Date(Date.now() + 3600 * 1000),
      };

      poolQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [row] });

      const result = await service.generateSigningKeyPair();
      expect(result.algorithm).toBe('ES256');
    });

    it('stores private key as dev:no-vault in dev mode (no Vault env vars)', async () => {
      const row = {
        id: 'uuid-4',
        kid: 'key-test-004',
        algorithm: 'RS256',
        public_key_jwk: makeSampleJwk('key-test-004'),
        vault_key_path: 'dev:no-vault',
        is_current: true,
        created_at: new Date(),
        expires_at: new Date(Date.now() + 3600 * 1000),
      };

      poolQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [row] });

      const result = await service.generateSigningKeyPair();
      expect(result.vault_key_path).toBe('dev:no-vault');
    });

    it('invalidates JWKS Redis cache after generating a new key', async () => {
      const row = {
        id: 'uuid-5',
        kid: 'key-test-005',
        algorithm: 'RS256',
        public_key_jwk: makeSampleJwk('key-test-005'),
        vault_key_path: 'dev:no-vault',
        is_current: true,
        created_at: new Date(),
        expires_at: new Date(Date.now() + 3600 * 1000),
      };

      poolQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [row] });

      await service.generateSigningKeyPair();
      expect(redis.del).toHaveBeenCalledWith('oidc:jwks');
    });
  });

  // ── getCurrentKey ────────────────────────────────────────────────────────

  describe('getCurrentKey()', () => {
    it('returns the current key from the database', async () => {
      const row = makeSampleRow();
      poolQuery.mockResolvedValueOnce({ rows: [row] });

      const key = await service.getCurrentKey();
      expect(key.is_current).toBe(true);
      expect(key.kid).toBe('key-test-001');
    });

    it('throws OIDCKeyNotFoundError when no current key exists', async () => {
      poolQuery.mockResolvedValueOnce({ rows: [] });

      await expect(service.getCurrentKey()).rejects.toMatchObject({
        code: 'OIDC_KEY_NOT_FOUND',
        httpStatus: 500,
      });
    });
  });

  // ── rotateKey ────────────────────────────────────────────────────────────

  describe('rotateKey()', () => {
    it('promotes a new key and the old key remains queryable', async () => {
      const oldRow = makeSampleRow({ kid: 'key-old', is_current: false });
      const newRow = makeSampleRow({ kid: 'key-new', is_current: true });

      // generateSigningKeyPair: UPDATE + INSERT
      poolQuery
        .mockResolvedValueOnce({ rows: [] }) // UPDATE demote
        .mockResolvedValueOnce({ rows: [newRow] }); // INSERT new

      const result = await service.rotateKey();
      expect(result.kid).toBe('key-new');
      expect(result.is_current).toBe(true);

      // Simulate getPublicJWKS returning both (old still in expires_at window)
      poolQuery.mockResolvedValueOnce({
        rows: [oldRow, newRow],
      });

      const jwks = await service.getPublicJWKS();
      expect(jwks.keys).toHaveLength(2);
    });

    it('invalidates JWKS cache on rotation', async () => {
      const newRow = makeSampleRow({ kid: 'key-new', is_current: true });
      poolQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [newRow] });

      await service.rotateKey();
      expect(redis.del).toHaveBeenCalledWith('oidc:jwks');
    });
  });

  // ── getPublicJWKS ────────────────────────────────────────────────────────

  describe('getPublicJWKS()', () => {
    it('returns only non-expired keys', async () => {
      const nonExpiredRow = makeSampleRow();
      poolQuery.mockResolvedValueOnce({ rows: [nonExpiredRow] });

      const jwks = await service.getPublicJWKS();
      expect(jwks.keys).toHaveLength(1);
      expect(jwks.keys[0].kid).toBe('key-test-001');
    });

    it('returns cached JWKS when available', async () => {
      const cachedJwks = JSON.stringify({ keys: [makeSampleJwk('cached-key')] });
      (redis.get as jest.Mock).mockResolvedValueOnce(cachedJwks);

      const jwks = await service.getPublicJWKS();
      expect(jwks.keys[0].kid).toBe('cached-key');
      // Pool should NOT have been queried
      expect(poolQuery).not.toHaveBeenCalled();
    });

    it('caches the JWKS result in Redis', async () => {
      const row = makeSampleRow();
      poolQuery.mockResolvedValueOnce({ rows: [row] });

      await service.getPublicJWKS();
      expect(redis.set).toHaveBeenCalledWith(
        'oidc:jwks',
        expect.any(String),
        expect.objectContaining({ EX: expect.any(Number) }),
      );
    });
  });

  // ── pruneExpiredKeys ─────────────────────────────────────────────────────

  describe('pruneExpiredKeys()', () => {
    it('deletes past-expires_at keys from the database', async () => {
      const expiredRow = makeSampleRow({
        kid: 'key-expired',
        vault_key_path: 'dev:no-vault',
        expires_at: new Date(Date.now() - 1000),
      });

      poolQuery.mockResolvedValueOnce({ rows: [expiredRow] });

      await service.pruneExpiredKeys();
      // DELETE should have been called as the first (and only) DB call
      expect(poolQuery).toHaveBeenCalledTimes(1);
      expect(poolQuery).toHaveBeenCalledWith(
        expect.stringContaining('DELETE FROM oidc_keys'),
      );
    });

    it('invalidates JWKS cache when keys are pruned', async () => {
      const expiredRow = makeSampleRow({
        kid: 'key-expired',
        vault_key_path: 'dev:no-vault',
        expires_at: new Date(Date.now() - 1000),
      });
      poolQuery.mockResolvedValueOnce({ rows: [expiredRow] });

      await service.pruneExpiredKeys();
      expect(redis.del).toHaveBeenCalledWith('oidc:jwks');
    });

    it('does not invalidate cache when no keys are pruned', async () => {
      poolQuery.mockResolvedValueOnce({ rows: [] });
      await service.pruneExpiredKeys();
      expect(redis.del).not.toHaveBeenCalled();
    });
  });

  // ── getPrivateKeyPem ─────────────────────────────────────────────────────

  describe('getPrivateKeyPem()', () => {
    it('returns private key from dev in-memory store for dev:no-vault path', async () => {
      // generateSigningKeyPair runs real crypto; the real kid is dynamic — capture it
      // from the INSERT call args by spying on poolQuery.
      let capturedKid: string | undefined;

      poolQuery
        .mockResolvedValueOnce({ rows: [] }) // UPDATE demote
        .mockImplementationOnce((_sql: string, params: unknown[]) => {
          // INSERT — capture the kid from the query params
          capturedKid = params[0] as string;
          return Promise.resolve({
            rows: [
              {
                id: 'uuid-dev',
                kid: capturedKid,
                algorithm: 'RS256',
                public_key_jwk: makeSampleJwk(capturedKid),
                vault_key_path: 'dev:no-vault',
                is_current: true,
                created_at: new Date(),
                expires_at: new Date(Date.now() + 3600 * 1000),
              },
            ],
          });
        });

      await service.generateSigningKeyPair();

      expect(capturedKid).toBeDefined();
      const pem = await service.getPrivateKeyPem(capturedKid!, 'dev:no-vault');
      expect(pem).toBeDefined();
      expect(typeof pem).toBe('string');
      expect(pem.length).toBeGreaterThan(0);
    });

    it('throws when dev key is not found in memory', async () => {
      await expect(
        service.getPrivateKeyPem('nonexistent-key-xyz-999', 'dev:no-vault'),
      ).rejects.toThrow('not found in memory');
    });

    it('reads private key from Vault when VAULT_ADDR and VAULT_TOKEN are set', async () => {
      process.env['VAULT_ADDR'] = 'http://vault:8200';
      process.env['VAULT_TOKEN'] = 'test-token';

      const pem = await service.getPrivateKeyPem('key-vault-001', 'secret/data/agentidp/oidc/keys/key-vault-001');
      expect(pem).toBe('mock-pem');
    });

    it('throws when VAULT_ADDR and VAULT_TOKEN are missing for a non-dev path', async () => {
      // No VAULT_ADDR/VAULT_TOKEN set (cleared in beforeEach)
      await expect(
        service.getPrivateKeyPem('key-vault-001', 'secret/data/agentidp/oidc/keys/key-vault-001'),
      ).rejects.toThrow('VAULT_ADDR and VAULT_TOKEN are required');
    });
  });

  // ── Vault path — storePrivateKey ─────────────────────────────────────────

  describe('generateSigningKeyPair() — Vault mode', () => {
    it('stores private key in Vault when VAULT_ADDR and VAULT_TOKEN are set', async () => {
      process.env['VAULT_ADDR'] = 'http://vault:8200';
      process.env['VAULT_TOKEN'] = 'test-token';

      const vaultPath = 'secret/data/agentidp/oidc/keys/key-vault-002';
      const row = {
        id: 'uuid-vault',
        kid: 'key-vault-002',
        algorithm: 'RS256',
        public_key_jwk: makeSampleJwk('key-vault-002'),
        vault_key_path: vaultPath,
        is_current: true,
        created_at: new Date(),
        expires_at: new Date(Date.now() + 3600 * 1000),
      };

      poolQuery
        .mockResolvedValueOnce({ rows: [] }) // UPDATE demote
        .mockResolvedValueOnce({ rows: [row] }); // INSERT

      const result = await service.generateSigningKeyPair();
      // vault_key_path returned by DB mock confirms Vault path was used
      expect(result.vault_key_path).toBe(vaultPath);
    });

    it('uses custom VAULT_MOUNT when set', async () => {
      process.env['VAULT_ADDR'] = 'http://vault:8200';
      process.env['VAULT_TOKEN'] = 'test-token';
      process.env['VAULT_MOUNT'] = 'kv';

      const vaultPath = 'kv/data/agentidp/oidc/keys/key-vault-003';
      const row = {
        id: 'uuid-vault-2',
        kid: 'key-vault-003',
        algorithm: 'RS256',
        public_key_jwk: makeSampleJwk('key-vault-003'),
        vault_key_path: vaultPath,
        is_current: true,
        created_at: new Date(),
        expires_at: new Date(Date.now() + 3600 * 1000),
      };

      poolQuery
        .mockResolvedValueOnce({ rows: [] })
        .mockResolvedValueOnce({ rows: [row] });

      const result = await service.generateSigningKeyPair();
      expect(result.vault_key_path).toBe(vaultPath);
    });
  });

  // ── pruneExpiredKeys — Vault path ────────────────────────────────────────

  describe('pruneExpiredKeys() — Vault mode', () => {
    it('deletes Vault key when VAULT_ADDR and VAULT_TOKEN are set', async () => {
      process.env['VAULT_ADDR'] = 'http://vault:8200';
      process.env['VAULT_TOKEN'] = 'test-token';

      const expiredRow = makeSampleRow({
        kid: 'key-vault-expired',
        vault_key_path: 'secret/data/agentidp/oidc/keys/key-vault-expired',
        expires_at: new Date(Date.now() - 1000),
      });

      poolQuery.mockResolvedValueOnce({ rows: [expiredRow] });

      await service.pruneExpiredKeys();
      expect(poolQuery).toHaveBeenCalledTimes(1);
      expect(redis.del).toHaveBeenCalledWith('oidc:jwks');
    });

    it('handles Vault delete errors gracefully (best-effort)', async () => {
      process.env['VAULT_ADDR'] = 'http://vault:8200';
      process.env['VAULT_TOKEN'] = 'test-token';

      // Make vault.delete throw
      const nodeVaultMock = jest.requireMock('node-vault') as jest.Mock;
      nodeVaultMock.mockReturnValueOnce({
        write: jest.fn().mockResolvedValue({}),
        read: jest.fn().mockResolvedValue({ data: { data: { privateKeyPem: 'mock-pem' } } }),
        delete: jest.fn().mockRejectedValue(new Error('Vault unreachable')),
      });

      const expiredRow = makeSampleRow({
        kid: 'key-vault-fail',
        vault_key_path: 'secret/data/agentidp/oidc/keys/key-vault-fail',
        expires_at: new Date(Date.now() - 1000),
      });

      poolQuery.mockResolvedValueOnce({ rows: [expiredRow] });

      // Should not throw — vault delete failure is best-effort
      await expect(service.pruneExpiredKeys()).resolves.not.toThrow();
      expect(redis.del).toHaveBeenCalledWith('oidc:jwks');
    });
  });
});