sentryagent-idp/tests/integration/federation.test.ts

/**
 * Integration tests for Federation endpoints.
 * Uses a real Postgres test DB and Redis test instance.
 *
 * For the verify endpoint, a real RS256 key pair is generated in test setup.
 * A partner is registered with a mocked JWKS endpoint, and a valid JWT is
 * signed with the private key and verified through the full stack.
 */

import crypto, { generateKeyPairSync } from 'crypto';
import request from 'supertest';
import { Application } from 'express';
import { v4 as uuidv4 } from 'uuid';
import { Pool } from 'pg';
import jwt from 'jsonwebtoken';

// Set test environment variables before importing app
const { privateKey, publicKey } = crypto.generateKeyPairSync('rsa', {
  modulusLength: 2048,
  publicKeyEncoding: { type: 'spki', format: 'pem' },
  privateKeyEncoding: { type: 'pkcs8', format: 'pem' },
});

process.env['DATABASE_URL'] =
  process.env['TEST_DATABASE_URL'] ??
  'postgresql://sentryagent:sentryagent@localhost:5432/sentryagent_idp_test';
process.env['REDIS_URL'] = process.env['TEST_REDIS_URL'] ?? 'redis://localhost:6379/1';
process.env['JWT_PRIVATE_KEY'] = privateKey;
process.env['JWT_PUBLIC_KEY'] = publicKey;
process.env['NODE_ENV'] = 'test';
process.env['DEFAULT_ORG_ID'] = 'org_system';

import { createApp } from '../../src/app';
import { signToken } from '../../src/utils/jwt';
import { closePool } from '../../src/db/pool';
import { closeRedisClient } from '../../src/cache/redis';

// ─── Token helpers ────────────────────────────────────────────────────────────

const CALLER_ID = uuidv4();
const ADMIN_SCOPE = 'admin:orgs';
const AGENT_SCOPE = 'agents:read';

function makeToken(sub: string = CALLER_ID, scope: string = ADMIN_SCOPE): string {
  return signToken({ sub, client_id: sub, scope, jti: uuidv4() }, privateKey);
}

// ─── Partner IdP key pair (for federated token signing) ──────────────────────

/**
 * A separate RS256 key pair representing an external IdP that is registered
 * as a federation partner. Tokens signed with this key should be verifiable
 * through POST /api/v1/federation/verify.
 */
const partnerKeys = generateKeyPairSync('rsa', {
  modulusLength: 2048,
  publicKeyEncoding: { type: 'spki', format: 'pem' },
  privateKeyEncoding: { type: 'pkcs8', format: 'pem' },
});

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

const PARTNER_ISSUER = 'https://external-idp.test.sentryagent.ai';
const PARTNER_JWKS_URI = 'https://external-idp.test.sentryagent.ai/.well-known/jwks.json';

/** Signs a JWT with the partner's private key, simulating an external IdP token. */
function makePartnerToken(claims: Record<string, unknown> = {}): string {
  return jwt.sign(
    {
      iss: PARTNER_ISSUER,
      sub: `partner-agent-${uuidv4()}`,
      aud: 'sentryagent-idp',
      ...claims,
    },
    partnerKeys.privateKey,
    {
      algorithm: 'RS256',
      header: { alg: 'RS256', kid: PARTNER_KID, typ: 'JWT' },
    },
  );
}

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

describe('Federation Endpoints Integration Tests', () => {
  let app: Application;
  let pool: Pool;

  beforeAll(async () => {
    app = await createApp();
    pool = new Pool({ connectionString: process.env['DATABASE_URL'] });

    const migrations: string[] = [
      `CREATE TABLE IF NOT EXISTS schema_migrations (
        name VARCHAR(255) PRIMARY KEY,
        applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
      )`,

      `CREATE TABLE IF NOT EXISTS organizations (
        organization_id   VARCHAR(40)  PRIMARY KEY,
        name              VARCHAR(100) NOT NULL,
        slug              VARCHAR(50)  NOT NULL UNIQUE,
        plan_tier         VARCHAR(20)  NOT NULL DEFAULT 'free',
        max_agents        INTEGER      NOT NULL DEFAULT 100,
        max_tokens_per_month INTEGER   NOT NULL DEFAULT 10000,
        status            VARCHAR(20)  NOT NULL DEFAULT 'active',
        created_at        TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
        updated_at        TIMESTAMPTZ  NOT NULL DEFAULT NOW()
      )`,

      `INSERT INTO organizations
         (organization_id, name, slug, plan_tier, max_agents, max_tokens_per_month, status)
       VALUES
         ('org_system', 'System', 'system', 'enterprise', 999999, 999999999, 'active')
       ON CONFLICT (organization_id) DO NOTHING`,

      `CREATE TABLE IF NOT EXISTS agents (
        agent_id       UUID         PRIMARY KEY DEFAULT gen_random_uuid(),
        email          VARCHAR(255) NOT NULL UNIQUE,
        agent_type     VARCHAR(32)  NOT NULL,
        version        VARCHAR(64)  NOT NULL,
        capabilities   TEXT[]       NOT NULL DEFAULT '{}',
        owner          VARCHAR(128) NOT NULL,
        deployment_env VARCHAR(16)  NOT NULL,
        status         VARCHAR(24)  NOT NULL DEFAULT 'active',
        organization_id VARCHAR(40) NOT NULL DEFAULT 'org_system',
        created_at     TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
        updated_at     TIMESTAMPTZ  NOT NULL DEFAULT NOW()
      )`,

      `CREATE TABLE IF NOT EXISTS credentials (
        credential_id UUID         PRIMARY KEY DEFAULT gen_random_uuid(),
        client_id     UUID         NOT NULL,
        secret_hash   VARCHAR(255) NOT NULL,
        status        VARCHAR(16)  NOT NULL DEFAULT 'active',
        organization_id VARCHAR(40) NOT NULL DEFAULT 'org_system',
        created_at    TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
        expires_at    TIMESTAMPTZ,
        revoked_at    TIMESTAMPTZ
      )`,

      `CREATE TABLE IF NOT EXISTS audit_events (
        event_id   UUID         PRIMARY KEY DEFAULT gen_random_uuid(),
        agent_id   UUID         NOT NULL,
        action     VARCHAR(32)  NOT NULL,
        outcome    VARCHAR(16)  NOT NULL,
        ip_address VARCHAR(64)  NOT NULL,
        user_agent TEXT         NOT NULL,
        metadata   JSONB        NOT NULL DEFAULT '{}',
        organization_id VARCHAR(40) NOT NULL DEFAULT 'org_system',
        timestamp  TIMESTAMPTZ  NOT NULL DEFAULT NOW()
      )`,

      `CREATE TABLE IF NOT EXISTS token_revocations (
        jti        UUID         PRIMARY KEY,
        expires_at TIMESTAMPTZ  NOT NULL,
        revoked_at TIMESTAMPTZ  NOT NULL DEFAULT NOW()
      )`,

      `CREATE TABLE IF NOT EXISTS oidc_keys (
        id              UUID         PRIMARY KEY DEFAULT gen_random_uuid(),
        kid             VARCHAR(128) NOT NULL UNIQUE,
        algorithm       VARCHAR(16)  NOT NULL,
        public_key_jwk  JSONB        NOT NULL,
        vault_key_path  VARCHAR(512) NOT NULL,
        is_current      BOOLEAN      NOT NULL DEFAULT FALSE,
        created_at      TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
        expires_at      TIMESTAMPTZ  NOT NULL
      )`,

      `CREATE TABLE IF NOT EXISTS federation_partners (
        id                    UUID         PRIMARY KEY DEFAULT gen_random_uuid(),
        name                  VARCHAR(255) NOT NULL,
        issuer                VARCHAR(512) NOT NULL UNIQUE,
        jwks_uri              VARCHAR(512) NOT NULL,
        allowed_organizations JSONB        NOT NULL DEFAULT '[]',
        status                VARCHAR(32)  NOT NULL DEFAULT 'active',
        created_at            TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
        updated_at            TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
        expires_at            TIMESTAMPTZ
      )`,
    ];

    for (const sql of migrations) {
      await pool.query(sql);
    }
  });

  afterEach(async () => {
    await pool.query('DELETE FROM federation_partners');
  });

  afterAll(async () => {
    await pool.end();
    await closePool();
    await closeRedisClient();
  });

  // ─── Helper: register a partner, mocking the JWKS fetch ───────────────────

  async function registerPartner(overrides: Record<string, unknown> = {}): Promise<{ body: Record<string, unknown>; status: number }> {
    // Mock global fetch to return the partner's JWKS when called during registration
    global.fetch = jest.fn().mockResolvedValue({
      ok: true,
      json: jest.fn().mockResolvedValue({ keys: [partnerJwk] }),
    } as unknown as Response);

    const token = makeToken();
    const res = await request(app)
      .post('/api/v1/federation/trust')
      .set('Authorization', `Bearer ${token}`)
      .send({
        name: 'Test External IdP',
        issuer: PARTNER_ISSUER,
        jwks_uri: PARTNER_JWKS_URI,
        ...overrides,
      });

    return { body: res.body as Record<string, unknown>, status: res.status };
  }

  // ─── POST /api/v1/federation/trust ────────────────────────────────────────

  describe('POST /api/v1/federation/trust', () => {
    it('registers a partner and returns 201', async () => {
      const { status, body } = await registerPartner();

      expect(status).toBe(201);
      expect(body['id']).toBeDefined();
      expect(body['issuer']).toBe(PARTNER_ISSUER);
      expect(body['status']).toBe('active');
    });

    it('returns 401 without a token', async () => {
      const res = await request(app)
        .post('/api/v1/federation/trust')
        .send({ name: 'x', issuer: 'https://x.example.com', jwks_uri: 'https://x.example.com/jwks' });

      expect(res.status).toBe(401);
    });

    it('returns 403 without admin:orgs scope', async () => {
      global.fetch = jest.fn().mockResolvedValue({
        ok: true,
        json: jest.fn().mockResolvedValue({ keys: [partnerJwk] }),
      } as unknown as Response);

      const token = makeToken(CALLER_ID, AGENT_SCOPE);
      const res = await request(app)
        .post('/api/v1/federation/trust')
        .set('Authorization', `Bearer ${token}`)
        .send({ name: 'x', issuer: 'https://x.example.com', jwks_uri: 'https://x.example.com/jwks' });

      expect(res.status).toBe(403);
    });

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

      const token = makeToken();
      const res = await request(app)
        .post('/api/v1/federation/trust')
        .set('Authorization', `Bearer ${token}`)
        .send({
          name: 'Bad Partner',
          issuer: 'https://bad.example.com',
          jwks_uri: 'https://bad.example.com/.well-known/jwks.json',
        });

      expect(res.status).toBe(400);
    });
  });

  // ─── GET /api/v1/federation/partners ──────────────────────────────────────

  describe('GET /api/v1/federation/partners', () => {
    it('returns list of registered partners (admin:orgs scope)', async () => {
      await registerPartner();

      const token = makeToken();
      const res = await request(app)
        .get('/api/v1/federation/partners')
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(200);
      expect(Array.isArray(res.body)).toBe(true);
      expect(res.body).toHaveLength(1);
      expect((res.body as Array<Record<string, unknown>>)[0]['issuer']).toBe(PARTNER_ISSUER);
    });

    it('returns 403 without admin:orgs scope', async () => {
      const token = makeToken(CALLER_ID, AGENT_SCOPE);
      const res = await request(app)
        .get('/api/v1/federation/partners')
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(403);
    });

    it('returns 401 without a token', async () => {
      const res = await request(app).get('/api/v1/federation/partners');
      expect(res.status).toBe(401);
    });
  });

  // ─── GET /api/v1/federation/partners/:id ──────────────────────────────────

  describe('GET /api/v1/federation/partners/:id', () => {
    it('returns the specific partner by ID', async () => {
      const { body: created } = await registerPartner();
      const partnerId = created['id'] as string;

      const token = makeToken();
      const res = await request(app)
        .get(`/api/v1/federation/partners/${partnerId}`)
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(200);
      expect((res.body as Record<string, unknown>)['id']).toBe(partnerId);
    });

    it('returns 404 for unknown partner id', async () => {
      const token = makeToken();
      const res = await request(app)
        .get(`/api/v1/federation/partners/${uuidv4()}`)
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(404);
    });
  });

  // ─── PATCH /api/v1/federation/partners/:id ────────────────────────────────

  describe('PATCH /api/v1/federation/partners/:id', () => {
    it('updates the partner name and returns 200', async () => {
      const { body: created } = await registerPartner();
      const partnerId = created['id'] as string;

      const token = makeToken();
      const res = await request(app)
        .patch(`/api/v1/federation/partners/${partnerId}`)
        .set('Authorization', `Bearer ${token}`)
        .send({ name: 'Updated Partner Name' });

      expect(res.status).toBe(200);
      expect((res.body as Record<string, unknown>)['name']).toBe('Updated Partner Name');
    });

    it('returns 404 for unknown partner id', async () => {
      const token = makeToken();
      const res = await request(app)
        .patch(`/api/v1/federation/partners/${uuidv4()}`)
        .set('Authorization', `Bearer ${token}`)
        .send({ name: 'Ghost' });

      expect(res.status).toBe(404);
    });
  });

  // ─── DELETE /api/v1/federation/partners/:id ───────────────────────────────

  describe('DELETE /api/v1/federation/partners/:id', () => {
    it('deletes the partner and returns 204', async () => {
      const { body: created } = await registerPartner();
      const partnerId = created['id'] as string;

      const token = makeToken();
      const res = await request(app)
        .delete(`/api/v1/federation/partners/${partnerId}`)
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(204);
    });

    it('returns 404 when deleting a non-existent partner', async () => {
      const token = makeToken();
      const res = await request(app)
        .delete(`/api/v1/federation/partners/${uuidv4()}`)
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(404);
    });
  });

  // ─── POST /api/v1/federation/verify ───────────────────────────────────────

  describe('POST /api/v1/federation/verify', () => {
    it('verifies a valid token from a registered partner and returns 200', async () => {
      // Register the partner (JWKS cached from registration)
      await registerPartner();

      // After registration, the JWKS is cached in Redis from the registration fetch.
      // Now mock fetch to return the JWKS again for the verify path (cache miss fallback).
      global.fetch = jest.fn().mockResolvedValue({
        ok: true,
        json: jest.fn().mockResolvedValue({ keys: [partnerJwk] }),
      } as unknown as Response);

      const federatedToken = makePartnerToken();
      const token = makeToken(CALLER_ID, AGENT_SCOPE);

      const res = await request(app)
        .post('/api/v1/federation/verify')
        .set('Authorization', `Bearer ${token}`)
        .send({ token: federatedToken });

      expect(res.status).toBe(200);
      expect((res.body as Record<string, unknown>)['valid']).toBe(true);
      expect((res.body as Record<string, unknown>)['issuer']).toBe(PARTNER_ISSUER);
    });

    it('rejects a token from an unknown issuer with 401', async () => {
      const unknownToken = jwt.sign(
        { iss: 'https://completely-unknown.example.com', sub: 'x', aud: 'a' },
        partnerKeys.privateKey,
        { algorithm: 'RS256' },
      );

      const token = makeToken(CALLER_ID, AGENT_SCOPE);
      const res = await request(app)
        .post('/api/v1/federation/verify')
        .set('Authorization', `Bearer ${token}`)
        .send({ token: unknownToken });

      expect(res.status).toBe(401);
      expect((res.body as Record<string, unknown>)['code']).toBe('FEDERATION_VERIFICATION_ERROR');
    });

    it('rejects an alg:none token with 401', async () => {
      const header = Buffer.from(JSON.stringify({ alg: 'none', typ: 'JWT' })).toString('base64url');
      const payload = Buffer.from(JSON.stringify({
        iss: PARTNER_ISSUER,
        sub: 'x',
        aud: 'a',
        iat: 1,
        exp: 9999999999,
      })).toString('base64url');
      const noneToken = `${header}.${payload}.`;

      const token = makeToken(CALLER_ID, AGENT_SCOPE);
      const res = await request(app)
        .post('/api/v1/federation/verify')
        .set('Authorization', `Bearer ${token}`)
        .send({ token: noneToken });

      expect(res.status).toBe(401);
    });

    it('returns 401 without a bearer token', async () => {
      const res = await request(app)
        .post('/api/v1/federation/verify')
        .send({ token: 'some-token' });

      expect(res.status).toBe(401);
    });
  });
});