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

/**
 * Integration tests for OIDC endpoints.
 * Tests discovery document, JWKS, id_token issuance, and agent-info endpoint.
 * Uses a real Postgres test DB and Redis test instance.
 */

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

// 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['OIDC_ISSUER'] = 'https://idp.sentryagent.ai';

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

function makeToken(sub: string, scope = 'agents:read agents:write tokens:read'): string {
  return signToken({ sub, client_id: sub, scope, jti: uuidv4() }, privateKey);
}

/**
 * Converts a JWK public key to a PEM string for jwt.verify.
 */
function jwkToPem(jwk: IJWKSKey): string {
  const keyObj = createPublicKey({ key: jwk as unknown as JsonWebKey, format: 'jwk' });
  return keyObj.export({ type: 'spki', format: 'pem' }) as string;
}

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

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

    // Ensure all required tables exist for this test suite
    const migrations = [
      `CREATE TABLE IF NOT EXISTS organizations (
        organization_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
        name VARCHAR(255) NOT NULL,
        slug VARCHAR(128) NOT NULL UNIQUE,
        plan_tier VARCHAR(32) NOT NULL DEFAULT 'free',
        max_agents INTEGER NOT NULL DEFAULT 10,
        max_tokens_per_month INTEGER NOT NULL DEFAULT 10000,
        status VARCHAR(16) NOT NULL DEFAULT 'active',
        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
      )`,
      `CREATE TABLE IF NOT EXISTS agents (
        agent_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
        organization_id UUID NOT NULL DEFAULT 'a0000000-0000-0000-0000-000000000000',
        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',
        did VARCHAR(512),
        did_created_at TIMESTAMPTZ,
        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,
        organization_id UUID NOT NULL DEFAULT 'a0000000-0000-0000-0000-000000000000',
        secret_hash VARCHAR(255) NOT NULL,
        vault_path VARCHAR(512),
        status VARCHAR(16) NOT NULL DEFAULT 'active',
        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(64) NOT NULL,
        outcome VARCHAR(16) NOT NULL,
        ip_address VARCHAR(64) NOT NULL,
        user_agent TEXT NOT NULL,
        metadata JSONB NOT NULL DEFAULT '{}',
        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 token_monthly_counts (
        client_id UUID NOT NULL,
        month_key VARCHAR(7) NOT NULL,
        count INTEGER NOT NULL DEFAULT 0,
        PRIMARY KEY (client_id, month_key)
      )`,
      `CREATE TABLE IF NOT EXISTS oidc_keys (
        id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
        kid VARCHAR(64) 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
      )`,
    ];

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

  afterEach(async () => {
    await pool.query('DELETE FROM audit_events');
    await pool.query('DELETE FROM token_revocations');
    await pool.query('DELETE FROM token_monthly_counts');
    await pool.query('DELETE FROM credentials');
    await pool.query('DELETE FROM agents');
    // Do NOT delete oidc_keys — ensureCurrentKey() runs once at app startup
  });

  afterAll(async () => {
    await pool.query('DELETE FROM oidc_keys');
    await pool.end();
    await closePool();
    await closeRedisClient();
  });

  async function createAgentWithCredentials(): Promise<{
    agentId: string;
    clientSecret: string;
    orgId: string;
  }> {
    const agentId = uuidv4();
    const orgId = 'a0000000-0000-0000-0000-000000000000';
    const token = makeToken(agentId, 'agents:read agents:write tokens:read');

    // Insert agent directly
    await pool.query(
      `INSERT INTO agents
         (agent_id, organization_id, email, agent_type, version, capabilities, owner, deployment_env, status)
       VALUES ($1, $2, $3, 'screener', '1.0.0', '{"agents:read"}', 'test-team', 'production', 'active')`,
      [agentId, orgId, `oidc-test-${agentId}@test.ai`],
    );

    // Generate credential via API
    const credRes = await request(app)
      .post(`/api/v1/agents/${agentId}/credentials`)
      .set('Authorization', `Bearer ${token}`)
      .send({});

    return { agentId, clientSecret: credRes.body.clientSecret, orgId };
  }

  // ── Discovery document ───────────────────────────────────────────────────

  describe('GET /.well-known/openid-configuration', () => {
    it('returns a valid OIDC discovery document', async () => {
      const res = await request(app).get('/.well-known/openid-configuration');

      expect(res.status).toBe(200);
      expect(res.body.issuer).toBe('https://idp.sentryagent.ai');
      expect(res.body.token_endpoint).toContain('/oauth2/token');
      expect(res.body.jwks_uri).toContain('/.well-known/jwks.json');
      expect(res.body.authorization_endpoint).toBeDefined();
    });

    it('includes all required OIDC discovery fields', async () => {
      const res = await request(app).get('/.well-known/openid-configuration');

      const requiredFields = [
        'issuer',
        'authorization_endpoint',
        'token_endpoint',
        'jwks_uri',
        'response_types_supported',
        'subject_types_supported',
        'id_token_signing_alg_values_supported',
        'scopes_supported',
        'claims_supported',
        'grant_types_supported',
      ];

      for (const field of requiredFields) {
        expect(res.body).toHaveProperty(field);
      }
    });

    it('does not require authentication', async () => {
      const res = await request(app).get('/.well-known/openid-configuration');
      expect(res.status).toBe(200);
    });

    it('includes openid in scopes_supported', async () => {
      const res = await request(app).get('/.well-known/openid-configuration');
      expect(res.body.scopes_supported).toContain('openid');
    });

    it('includes RS256 in id_token_signing_alg_values_supported', async () => {
      const res = await request(app).get('/.well-known/openid-configuration');
      expect(res.body.id_token_signing_alg_values_supported).toContain('RS256');
    });
  });

  // ── JWKS endpoint ─────────────────────────────────────────────────────────

  describe('GET /.well-known/jwks.json', () => {
    it('returns JWKS with at least one key', async () => {
      const res = await request(app).get('/.well-known/jwks.json');

      expect(res.status).toBe(200);
      expect(res.body.keys).toBeInstanceOf(Array);
      expect(res.body.keys.length).toBeGreaterThanOrEqual(1);
    });

    it('returns keys with required JWK fields', async () => {
      const res = await request(app).get('/.well-known/jwks.json');
      const key = res.body.keys[0];

      expect(key.kid).toBeDefined();
      expect(key.kty).toBeDefined();
      expect(key.use).toBe('sig');
      expect(key.alg).toBeDefined();
    });

    it('does not require authentication', async () => {
      const res = await request(app).get('/.well-known/jwks.json');
      expect(res.status).toBe(200);
    });

    it('sets Cache-Control: public, max-age=3600', async () => {
      const res = await request(app).get('/.well-known/jwks.json');
      expect(res.headers['cache-control']).toContain('public');
      expect(res.headers['cache-control']).toContain('max-age=3600');
    });
  });

  // ── Token endpoint with openid scope ─────────────────────────────────────

  describe('POST /api/v1/token with openid scope', () => {
    it('returns id_token when openid scope is requested', async () => {
      const { agentId, clientSecret } = await createAgentWithCredentials();

      const res = await request(app)
        .post('/api/v1/token')
        .type('form')
        .send({
          grant_type: 'client_credentials',
          client_id: agentId,
          client_secret: clientSecret,
          scope: 'openid agents:read',
        });

      expect(res.status).toBe(200);
      expect(res.body.id_token).toBeDefined();
      expect(typeof res.body.id_token).toBe('string');
      expect(res.body.id_token.split('.')).toHaveLength(3);
    });

    it('does not return id_token when openid scope is not requested', async () => {
      const { agentId, clientSecret } = await createAgentWithCredentials();

      const res = await request(app)
        .post('/api/v1/token')
        .type('form')
        .send({
          grant_type: 'client_credentials',
          client_id: agentId,
          client_secret: clientSecret,
          scope: 'agents:read',
        });

      expect(res.status).toBe(200);
      expect(res.body.id_token).toBeUndefined();
    });

    it('id_token is verifiable against JWKS from /.well-known/jwks.json', async () => {
      const { agentId, clientSecret } = await createAgentWithCredentials();

      // Issue token with openid scope
      const tokenRes = await request(app)
        .post('/api/v1/token')
        .type('form')
        .send({
          grant_type: 'client_credentials',
          client_id: agentId,
          client_secret: clientSecret,
          scope: 'openid agents:read',
        });

      expect(tokenRes.status).toBe(200);
      const idToken: string = tokenRes.body.id_token;

      // Fetch JWKS
      const jwksRes = await request(app).get('/.well-known/jwks.json');
      const jwks: IJWKSResponse = jwksRes.body;

      // Decode header to get kid
      const decoded = jwt.decode(idToken, { complete: true });
      expect(decoded).not.toBeNull();
      const kid = decoded!.header.kid;

      // Find matching key
      const matchingKey = jwks.keys.find((k) => k.kid === kid);
      expect(matchingKey).toBeDefined();

      // Verify signature
      const publicKeyPem = jwkToPem(matchingKey!);
      const verified = jwt.verify(idToken, publicKeyPem, { algorithms: ['RS256', 'ES256'] });
      expect(verified).toBeDefined();
      const payload = verified as Record<string, unknown>;
      expect(payload['sub']).toBe(agentId);
    });

    it('id_token contains correct agent claims', async () => {
      const { agentId, clientSecret } = await createAgentWithCredentials();

      const res = await request(app)
        .post('/api/v1/token')
        .type('form')
        .send({
          grant_type: 'client_credentials',
          client_id: agentId,
          client_secret: clientSecret,
          scope: 'openid agents:read',
        });

      const idToken: string = res.body.id_token;
      const decoded = jwt.decode(idToken) as Record<string, unknown>;

      expect(decoded['sub']).toBe(agentId);
      expect(decoded['iss']).toBe('https://idp.sentryagent.ai');
      expect(decoded['aud']).toBe(agentId);
      expect(decoded['agent_type']).toBe('screener');
      expect(decoded['deployment_env']).toBe('production');
      expect(decoded['organization_id']).toBeDefined();
    });

    it('id_token header contains kid matching the JWKS', async () => {
      const { agentId, clientSecret } = await createAgentWithCredentials();

      const tokenRes = await request(app)
        .post('/api/v1/token')
        .type('form')
        .send({
          grant_type: 'client_credentials',
          client_id: agentId,
          client_secret: clientSecret,
          scope: 'openid agents:read',
        });

      const jwksRes = await request(app).get('/.well-known/jwks.json');
      const jwks: IJWKSResponse = jwksRes.body;
      const jwksKids = jwks.keys.map((k) => k.kid);

      const decoded = jwt.decode(tokenRes.body.id_token, { complete: true });
      expect(jwksKids).toContain(decoded!.header.kid);
    });
  });

  // ── Agent info endpoint ───────────────────────────────────────────────────

  describe('GET /agent-info', () => {
    it('returns agent identity claims for authenticated caller', async () => {
      const { agentId } = await createAgentWithCredentials();
      const token = makeToken(agentId, 'openid agents:read');

      const res = await request(app)
        .get('/agent-info')
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(200);
      expect(res.body.sub).toBe(agentId);
      expect(res.body.agent_type).toBe('screener');
      expect(res.body.deployment_env).toBe('production');
      expect(res.body.organization_id).toBeDefined();
    });

    it('returns 401 without a Bearer token', async () => {
      const res = await request(app).get('/agent-info');
      expect(res.status).toBe(401);
    });

    it('returns 401 with an invalid Bearer token', async () => {
      const res = await request(app)
        .get('/agent-info')
        .set('Authorization', 'Bearer invalid.token.here');
      expect(res.status).toBe(401);
    });

    it('includes scope in the agent-info response', async () => {
      const { agentId } = await createAgentWithCredentials();
      const token = makeToken(agentId, 'openid agents:read');

      const res = await request(app)
        .get('/agent-info')
        .set('Authorization', `Bearer ${token}`);

      expect(res.status).toBe(200);
      expect(res.body.scope).toContain('openid');
    });

    it('returns 404 for a token referencing a non-existent agent', async () => {
      const unknownAgentId = uuidv4();
      const token = makeToken(unknownAgentId, 'openid agents:read');

      const res = await request(app)
        .get('/agent-info')
        .set('Authorization', `Bearer ${token}`);

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