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

/**
 * Unit tests for src/services/FederationService.ts
 * Mocks pg Pool, Redis, and global fetch.
 */

import crypto, { generateKeyPairSync } from 'crypto';
import jwt from 'jsonwebtoken';
import { Pool } from 'pg';
import { RedisClientType } from 'redis';
import { FederationService, FederationPartnerError, FederationPartnerNotFoundError, FederationVerificationError } from '../../../src/services/FederationService';
import { IFederationPartner } from '../../../src/types/federation';
import { IJWKSKey } from '../../../src/types/oidc';

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

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

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

function makeRedis(): { redis: RedisClientType; store: Map<string, string> } {
  const store = new Map<string, string>();
  const redis = {
    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;
  return { redis, store };
}

function makePartnerRow(overrides: Partial<IFederationPartner> = {}): IFederationPartner {
  return {
    id: 'partner-uuid-1',
    name: 'Test Partner',
    issuer: 'https://external-idp.example.com',
    jwks_uri: 'https://external-idp.example.com/.well-known/jwks.json',
    allowed_organizations: [],
    status: 'active',
    created_at: new Date('2026-01-01T00:00:00Z'),
    updated_at: new Date('2026-01-01T00:00:00Z'),
    expires_at: null,
    ...overrides,
  };
}

/** Builds a mock fetch response returning a valid JWKS. */
function mockFetchJWKS(keys: IJWKSKey[]): void {
  global.fetch = jest.fn().mockResolvedValue({
    ok: true,
    json: jest.fn().mockResolvedValue({ keys }),
  } as unknown as Response);
}

/** Generates a real RS256 key pair for signing test tokens. */
function generateRSAKeyPair(): { privateKeyPem: string; publicKeyPem: string; jwk: IJWKSKey } {
  const { privateKey, publicKey } = generateKeyPairSync('rsa', {
    modulusLength: 2048,
    publicKeyEncoding: { type: 'spki', format: 'pem' },
    privateKeyEncoding: { type: 'pkcs8', format: 'pem' },
  });

  const jwkObj = crypto.createPublicKey(publicKey).export({ format: 'jwk' }) as Record<string, string>;
  const jwk: IJWKSKey = {
    kid: 'test-key-001',
    kty: 'RSA',
    use: 'sig',
    alg: 'RS256',
    n: jwkObj['n'],
    e: jwkObj['e'],
  };

  return { privateKeyPem: privateKey, publicKeyPem: publicKey, jwk };
}

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

describe('FederationService', () => {
  let pool: Pool;
  let poolQuery: jest.Mock;
  let redis: RedisClientType;
  let store: Map<string, string>;
  let service: FederationService;

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

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

  // ─────────────────────────────────────────────────────────────────────────
  // registerPartner
  // ─────────────────────────────────────────────────────────────────────────

  describe('registerPartner()', () => {
    it('fetches JWKS, inserts partner, and caches JWKS on success', async () => {
      const { jwk } = generateRSAKeyPair();
      mockFetchJWKS([jwk]);

      const partnerRow = makePartnerRow();
      poolQuery.mockResolvedValueOnce({ rows: [partnerRow] });

      const result = await service.registerPartner({
        name: 'Test Partner',
        issuer: 'https://external-idp.example.com',
        jwks_uri: 'https://external-idp.example.com/.well-known/jwks.json',
      });

      expect(result.id).toBe('partner-uuid-1');
      expect(poolQuery).toHaveBeenCalledWith(
        expect.stringContaining('INSERT INTO federation_partners'),
        expect.any(Array),
      );

      // JWKS should be cached in Redis
      const cacheKey = `federation:jwks:${partnerRow.issuer}`;
      expect(store.has(cacheKey)).toBe(true);
      const cached = JSON.parse(store.get(cacheKey)!);
      expect(cached).toEqual([jwk]);
    });

    it('throws FederationPartnerError (400) when JWKS endpoint is unreachable', async () => {
      global.fetch = jest.fn().mockRejectedValue(new Error('Network error'));

      await expect(
        service.registerPartner({
          name: 'Bad Partner',
          issuer: 'https://bad-idp.example.com',
          jwks_uri: 'https://bad-idp.example.com/.well-known/jwks.json',
        }),
      ).rejects.toMatchObject({
        code: 'FEDERATION_PARTNER_ERROR',
        httpStatus: 400,
      });
    });

    it('throws FederationPartnerError (400) when JWKS response is missing keys array', async () => {
      global.fetch = jest.fn().mockResolvedValue({
        ok: true,
        json: jest.fn().mockResolvedValue({ not_keys: [] }),
      } as unknown as Response);

      await expect(
        service.registerPartner({
          name: 'Bad JWKS Partner',
          issuer: 'https://bad-jwks.example.com',
          jwks_uri: 'https://bad-jwks.example.com/.well-known/jwks.json',
        }),
      ).rejects.toThrow(FederationPartnerError);
    });

    it('throws FederationPartnerError (400) when endpoint returns non-200', async () => {
      global.fetch = jest.fn().mockResolvedValue({
        ok: false,
        status: 404,
        json: jest.fn().mockResolvedValue({}),
      } as unknown as Response);

      await expect(
        service.registerPartner({
          name: 'Gone Partner',
          issuer: 'https://gone.example.com',
          jwks_uri: 'https://gone.example.com/.well-known/jwks.json',
        }),
      ).rejects.toThrow(FederationPartnerError);
    });
  });

  // ─────────────────────────────────────────────────────────────────────────
  // listPartners
  // ─────────────────────────────────────────────────────────────────────────

  describe('listPartners()', () => {
    it('marks expired partners (past expires_at) and returns list', async () => {
      // listPartners makes two DB calls:
      //   1. UPDATE federation_partners SET status = 'expired' ...
      //   2. SELECT * FROM federation_partners ...
      poolQuery.mockResolvedValueOnce({ rows: [], rowCount: 0 }); // UPDATE
      const expiredPartner = makePartnerRow({
        status: 'expired',
        expires_at: new Date(Date.now() - 1000),
      });
      poolQuery.mockResolvedValueOnce({ rows: [expiredPartner] }); // SELECT

      const partners = await service.listPartners();

      expect(poolQuery).toHaveBeenCalledWith(
        expect.stringContaining("SET status = 'expired'"),
      );
      expect(partners).toHaveLength(1);
      expect(partners[0].status).toBe('expired');
    });
  });

  // ─────────────────────────────────────────────────────────────────────────
  // getPartner
  // ─────────────────────────────────────────────────────────────────────────

  describe('getPartner()', () => {
    it('throws FederationPartnerNotFoundError (404) for unknown id', async () => {
      // getPartner makes two DB calls: UPDATE expiry + SELECT
      poolQuery.mockResolvedValueOnce({ rows: [], rowCount: 0 }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [] });              // SELECT → empty

      const err = await service.getPartner('non-existent-id').catch((e: unknown) => e);
      expect(err).toBeInstanceOf(FederationPartnerNotFoundError);
      expect((err as FederationPartnerNotFoundError).httpStatus).toBe(404);
    });

    it('returns the partner when found', async () => {
      const row = makePartnerRow();
      poolQuery.mockResolvedValueOnce({ rows: [], rowCount: 0 }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [row] });           // SELECT

      const result = await service.getPartner('partner-uuid-1');
      expect(result.id).toBe('partner-uuid-1');
    });
  });

  // ─────────────────────────────────────────────────────────────────────────
  // updatePartner
  // ─────────────────────────────────────────────────────────────────────────

  describe('updatePartner()', () => {
    it('updates name and status without invalidating JWKS cache', async () => {
      const existing = makePartnerRow();
      const updated = makePartnerRow({ name: 'Renamed Partner', status: 'suspended' });

      // getPartner: UPDATE expiry + SELECT
      poolQuery.mockResolvedValueOnce({ rows: [] }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [existing] }); // SELECT
      // UPDATE
      poolQuery.mockResolvedValueOnce({ rows: [updated] });

      const result = await service.updatePartner('partner-uuid-1', {
        name: 'Renamed Partner',
        status: 'suspended',
      });

      expect(result.name).toBe('Renamed Partner');
      expect(result.status).toBe('suspended');
      expect(redis.del).not.toHaveBeenCalled();
    });

    it('invalidates JWKS cache when jwks_uri changes', async () => {
      const existing = makePartnerRow({
        jwks_uri: 'https://external-idp.example.com/.well-known/jwks.json',
      });
      const newJwksUri = 'https://external-idp.example.com/.well-known/jwks-new.json';
      const updated = makePartnerRow({ jwks_uri: newJwksUri });
      const cacheKey = `federation:jwks:${existing.issuer}`;
      store.set(cacheKey, JSON.stringify([{ kid: 'old-key', kty: 'RSA', use: 'sig', alg: 'RS256' }]));

      // getPartner: UPDATE expiry + SELECT
      poolQuery.mockResolvedValueOnce({ rows: [] }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [existing] }); // SELECT
      // UPDATE
      poolQuery.mockResolvedValueOnce({ rows: [updated] });

      await service.updatePartner('partner-uuid-1', { jwks_uri: newJwksUri });

      expect(store.has(cacheKey)).toBe(false);
      expect(redis.del).toHaveBeenCalledWith(cacheKey);
    });

    it('does NOT invalidate JWKS cache when jwks_uri is unchanged', async () => {
      const existing = makePartnerRow();
      const updated = makePartnerRow({ name: 'Same JWKS URI' });
      const cacheKey = `federation:jwks:${existing.issuer}`;
      store.set(cacheKey, JSON.stringify([{ kid: 'test', kty: 'RSA', use: 'sig', alg: 'RS256' }]));

      // getPartner: UPDATE expiry + SELECT
      poolQuery.mockResolvedValueOnce({ rows: [] }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [existing] }); // SELECT
      // UPDATE — same jwks_uri as existing
      poolQuery.mockResolvedValueOnce({ rows: [updated] });

      await service.updatePartner('partner-uuid-1', {
        jwks_uri: existing.jwks_uri, // same value → no cache invalidation
        name: 'Same JWKS URI',
      });

      expect(store.has(cacheKey)).toBe(true);
      expect(redis.del).not.toHaveBeenCalled();
    });

    it('throws FederationPartnerNotFoundError when UPDATE returns no rows', async () => {
      const existing = makePartnerRow();

      // getPartner: UPDATE expiry + SELECT
      poolQuery.mockResolvedValueOnce({ rows: [] }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [existing] }); // SELECT
      // UPDATE returns empty (concurrent delete)
      poolQuery.mockResolvedValueOnce({ rows: [] });

      await expect(
        service.updatePartner('partner-uuid-1', { name: 'Ghost' }),
      ).rejects.toBeInstanceOf(FederationPartnerNotFoundError);
    });
  });

  // ─────────────────────────────────────────────────────────────────────────
  // deletePartner
  // ─────────────────────────────────────────────────────────────────────────

  describe('deletePartner()', () => {
    it('deletes the partner and invalidates Redis JWKS cache', async () => {
      const row = makePartnerRow();
      const cacheKey = `federation:jwks:${row.issuer}`;
      store.set(cacheKey, JSON.stringify([{ kid: 'test', kty: 'RSA', use: 'sig', alg: 'RS256' }]));

      // getPartner calls: UPDATE expiry + SELECT
      poolQuery.mockResolvedValueOnce({ rows: [] }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [row] }); // SELECT (inside getPartner)
      // DELETE
      poolQuery.mockResolvedValueOnce({ rows: [{ id: row.id }] });

      await service.deletePartner('partner-uuid-1');

      expect(store.has(cacheKey)).toBe(false);
      expect(redis.del).toHaveBeenCalledWith(cacheKey);
    });

    it('throws FederationPartnerNotFoundError when DELETE returns no rows', async () => {
      const row = makePartnerRow();

      // getPartner calls: UPDATE expiry + SELECT
      poolQuery.mockResolvedValueOnce({ rows: [] }); // UPDATE expiry
      poolQuery.mockResolvedValueOnce({ rows: [row] }); // SELECT (inside getPartner)
      // DELETE returns empty (already deleted)
      poolQuery.mockResolvedValueOnce({ rows: [] });

      await expect(service.deletePartner('partner-uuid-1')).rejects.toBeInstanceOf(
        FederationPartnerNotFoundError,
      );
    });
  });

  // ─────────────────────────────────────────────────────────────────────────
  // verifyFederatedToken
  // ─────────────────────────────────────────────────────────────────────────

  describe('verifyFederatedToken()', () => {
    it('succeeds with a valid RS256 token from a known active partner', async () => {
      const { privateKeyPem, jwk } = generateRSAKeyPair();
      const issuer = 'https://external-idp.example.com';

      const token = jwt.sign(
        { iss: issuer, sub: 'agent-abc', aud: 'sentryagent', organization_id: 'org-1' },
        privateKeyPem,
        { algorithm: 'RS256', header: { alg: 'RS256', kid: 'test-key-001', typ: 'JWT' } },
      );

      const partnerRow = makePartnerRow({ issuer, allowed_organizations: [] });
      poolQuery.mockResolvedValueOnce({ rows: [partnerRow] }); // lookup by issuer

      // Pre-populate JWKS cache
      store.set(`federation:jwks:${issuer}`, JSON.stringify([jwk]));

      const result = await service.verifyFederatedToken({ token });

      expect(result.valid).toBe(true);
      expect(result.issuer).toBe(issuer);
      expect(result.subject).toBe('agent-abc');
    });

    it('rejects alg:none tokens', async () => {
      // Build a token with alg:none by hand (jwt.sign won't allow it, so we craft manually)
      const header = Buffer.from(JSON.stringify({ alg: 'none', typ: 'JWT' })).toString('base64url');
      const payload = Buffer.from(JSON.stringify({ iss: 'https://issuer.example.com', sub: 'x', aud: 'a', iat: 1, exp: 9999999999 })).toString('base64url');
      const noneToken = `${header}.${payload}.`;

      await expect(service.verifyFederatedToken({ token: noneToken })).rejects.toThrow(
        FederationVerificationError,
      );
      await expect(service.verifyFederatedToken({ token: noneToken })).rejects.toMatchObject({
        message: 'alg:none tokens are not accepted',
        httpStatus: 401,
      });
    });

    it('rejects token from unknown issuer', async () => {
      const { privateKeyPem } = generateRSAKeyPair();
      const token = jwt.sign(
        { iss: 'https://unknown.example.com', sub: 'x', aud: 'a' },
        privateKeyPem,
        { algorithm: 'RS256' },
      );

      poolQuery.mockResolvedValueOnce({ rows: [] }); // no partner found

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        code: 'FEDERATION_VERIFICATION_ERROR',
        httpStatus: 401,
      });
    });

    it('rejects token from a suspended partner', async () => {
      const { privateKeyPem } = generateRSAKeyPair();
      const issuer = 'https://suspended.example.com';
      const token = jwt.sign(
        { iss: issuer, sub: 'x', aud: 'a' },
        privateKeyPem,
        { algorithm: 'RS256' },
      );

      const suspendedPartner = makePartnerRow({ issuer, status: 'suspended' });
      poolQuery.mockResolvedValueOnce({ rows: [suspendedPartner] });

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        message: 'Federation partner is suspended',
        httpStatus: 401,
      });
    });

    it('rejects token from an expired partner', async () => {
      const { privateKeyPem } = generateRSAKeyPair();
      const issuer = 'https://expired.example.com';
      const token = jwt.sign(
        { iss: issuer, sub: 'x', aud: 'a' },
        privateKeyPem,
        { algorithm: 'RS256' },
      );

      const expiredPartner = makePartnerRow({
        issuer,
        status: 'active',
        expires_at: new Date(Date.now() - 60_000), // 1 minute ago
      });
      poolQuery.mockResolvedValueOnce({ rows: [expiredPartner] });

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        message: 'Federation partner has expired',
        httpStatus: 401,
      });
    });

    it('rejects token whose organization_id is not in the partner allowlist', async () => {
      const { privateKeyPem, jwk } = generateRSAKeyPair();
      const issuer = 'https://restricted.example.com';

      const token = jwt.sign(
        { iss: issuer, sub: 'agent-abc', aud: 'sentryagent', organization_id: 'org-not-allowed' },
        privateKeyPem,
        { algorithm: 'RS256', header: { alg: 'RS256', kid: 'test-key-001', typ: 'JWT' } },
      );

      const restrictedPartner = makePartnerRow({
        issuer,
        allowed_organizations: ['org-allowed-only'],
      });
      poolQuery.mockResolvedValueOnce({ rows: [restrictedPartner] });
      store.set(`federation:jwks:${issuer}`, JSON.stringify([jwk]));

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        message: "Token organization_id is not in the partner's allowed list",
        httpStatus: 401,
      });
    });

    it('rejects token when expected_issuer does not match token iss', async () => {
      const { privateKeyPem } = generateRSAKeyPair();
      const token = jwt.sign(
        { iss: 'https://actual-issuer.example.com', sub: 'x', aud: 'a' },
        privateKeyPem,
        { algorithm: 'RS256' },
      );

      await expect(
        service.verifyFederatedToken({
          token,
          expected_issuer: 'https://expected-issuer.example.com',
        }),
      ).rejects.toMatchObject({
        message: 'Token issuer does not match expected issuer',
        httpStatus: 401,
      });
    });

    it('rejects a malformed token (not a valid JWT)', async () => {
      await expect(
        service.verifyFederatedToken({ token: 'not.a.jwt' }),
      ).rejects.toMatchObject({
        code: 'FEDERATION_VERIFICATION_ERROR',
        httpStatus: 401,
      });
    });

    it('rejects a token missing the iss claim', async () => {
      const { privateKeyPem } = generateRSAKeyPair();
      // jwt.sign without iss — omit it explicitly
      const token = jwt.sign(
        { sub: 'no-issuer', aud: 'a' },
        privateKeyPem,
        { algorithm: 'RS256' },
      );

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        message: 'Token is missing the iss claim',
        httpStatus: 401,
      });
    });

    it('rejects a token when no key matches the kid in the JWKS', async () => {
      const { privateKeyPem: signingKey, jwk } = generateRSAKeyPair();
      const issuer = 'https://no-kid.example.com';

      // Sign with kid that does NOT appear in the partner JWKS
      const token = jwt.sign(
        { iss: issuer, sub: 'agent', aud: 'sentryagent' },
        signingKey,
        { algorithm: 'RS256', header: { alg: 'RS256', kid: 'missing-key-id', typ: 'JWT' } },
      );

      const partnerRow = makePartnerRow({ issuer, allowed_organizations: [] });
      poolQuery.mockResolvedValueOnce({ rows: [partnerRow] });
      // Cache has a key with a different kid
      store.set(`federation:jwks:${issuer}`, JSON.stringify([{ ...jwk, kid: 'other-key-id' }]));

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        message: expect.stringContaining('No matching key in partner JWKS'),
        httpStatus: 401,
      });
    });

    it('rejects a token with an invalid signature (jwt.verify throws)', async () => {
      const { jwk } = generateRSAKeyPair();
      const { privateKeyPem: wrongPrivateKey } = generateRSAKeyPair();
      const issuer = 'https://badsig.example.com';

      // Sign with the wrong private key — but the JWKS contains the correct public key
      const token = jwt.sign(
        { iss: issuer, sub: 'agent', aud: 'sentryagent' },
        wrongPrivateKey,
        { algorithm: 'RS256', header: { alg: 'RS256', kid: jwk.kid, typ: 'JWT' } },
      );

      const partnerRow = makePartnerRow({ issuer, allowed_organizations: [] });
      poolQuery.mockResolvedValueOnce({ rows: [partnerRow] });
      store.set(`federation:jwks:${issuer}`, JSON.stringify([jwk]));

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        code: 'FEDERATION_VERIFICATION_ERROR',
        httpStatus: 401,
      });
    });

    it('rejects token when JWKS endpoint returns non-JSON on a cache miss', async () => {
      const { privateKeyPem } = generateRSAKeyPair();
      const issuer = 'https://nonjson-jwks.example.com';

      const token = jwt.sign(
        { iss: issuer, sub: 'agent', aud: 'sentryagent' },
        privateKeyPem,
        { algorithm: 'RS256' },
      );

      const partnerRow = makePartnerRow({ issuer });
      poolQuery.mockResolvedValueOnce({ rows: [partnerRow] });

      // Simulate non-JSON response from JWKS endpoint (cache is empty)
      global.fetch = jest.fn().mockResolvedValue({
        ok: true,
        json: jest.fn().mockRejectedValue(new SyntaxError('Unexpected token')),
      } as unknown as Response);

      await expect(service.verifyFederatedToken({ token })).rejects.toMatchObject({
        code: 'FEDERATION_PARTNER_ERROR',
        httpStatus: 400,
      });
    });

    it('caches JWKS on first verify and uses cache on second verify', async () => {
      const { privateKeyPem, jwk } = generateRSAKeyPair();
      const issuer = 'https://caching.example.com';

      const makeToken = (): string =>
        jwt.sign(
          { iss: issuer, sub: 'agent', aud: 'sentryagent' },
          privateKeyPem,
          { algorithm: 'RS256', header: { alg: 'RS256', kid: 'test-key-001', typ: 'JWT' } },
        );

      const partnerRow = makePartnerRow({ issuer });

      // First call: cache miss → fetch JWKS
      poolQuery.mockResolvedValueOnce({ rows: [partnerRow] });
      mockFetchJWKS([jwk]);
      await service.verifyFederatedToken({ token: makeToken() });

      const cacheKey = `federation:jwks:${issuer}`;
      expect(store.has(cacheKey)).toBe(true);

      // Second call: cache hit → fetch should NOT be called again
      poolQuery.mockResolvedValueOnce({ rows: [partnerRow] });
      const fetchMock = global.fetch as jest.Mock;
      fetchMock.mockClear();

      await service.verifyFederatedToken({ token: makeToken() });

      expect(fetchMock).not.toHaveBeenCalled();
    });
  });
});