sentryagent-idp/src/services/OAuth2Service.ts

/**
 * OAuth 2.0 Token Service for SentryAgent.ai AgentIdP.
 * Issues, introspects, and revokes RS256 JWT access tokens.
 */

import { TokenRepository } from '../repositories/TokenRepository.js';
import { CredentialRepository } from '../repositories/CredentialRepository.js';
import { AgentRepository } from '../repositories/AgentRepository.js';
import { AuditService } from './AuditService.js';
import { VaultClient } from '../vault/VaultClient.js';
import { IDTokenService } from './IDTokenService.js';
import { EventPublisher } from './EventPublisher.js';
import { EncryptionService } from './EncryptionService.js';
import {
  ITokenPayload,
  ITokenResponse,
  IIntrospectResponse,
  IOAuth2ErrorResponse,
} from '../types/index.js';
import {
  AuthenticationError,
  AuthorizationError,
  FreeTierLimitError,
} from '../utils/errors.js';
import { signToken, verifyToken, decodeToken, getTokenExpiresIn } from '../utils/jwt.js';
import { verifySecret } from '../utils/crypto.js';
import { v4 as uuidv4 } from 'uuid';
import { tokensIssuedTotal } from '../metrics/registry.js';

const FREE_TIER_MAX_MONTHLY_TOKENS = 10000;

/** Result of a token issuance, including either a success response or OAuth2 error. */
export interface IssueTokenResult {
  success: boolean;
  response?: ITokenResponse;
  error?: IOAuth2ErrorResponse;
  httpStatus?: number;
}

/**
 * Service for OAuth 2.0 Client Credentials token issuance, introspection, and revocation.
 */
export class OAuth2Service {
  /**
   * @param tokenRepository - Repository for token revocation and monthly counts.
   * @param credentialRepository - Repository for credential lookup and verification.
   * @param agentRepository - Repository for agent status lookup.
   * @param auditService - The audit log service.
   * @param privateKey - PEM-encoded RSA private key for signing tokens.
   * @param publicKey - PEM-encoded RSA public key for verifying tokens.
   * @param vaultClient - Optional VaultClient for Phase 2 credential verification.
   * @param idTokenService - Optional IDTokenService; when provided and `openid` scope
   *   is requested, an OIDC ID token is appended to the token response.
   * @param eventPublisher - Optional EventPublisher. When provided, token.issued and
   *   token.revoked events are published as webhooks and Kafka messages (fire-and-forget).
   * @param encryptionService - Optional EncryptionService. When provided, encrypted
   *   `secret_hash` values are decrypted before bcrypt verification (SOC 2 CC6.1).
   */
  constructor(
    private readonly tokenRepository: TokenRepository,
    private readonly credentialRepository: CredentialRepository,
    private readonly agentRepository: AgentRepository,
    private readonly auditService: AuditService,
    private readonly privateKey: string,
    private readonly publicKey: string,
    private readonly vaultClient: VaultClient | null = null,
    private readonly idTokenService: IDTokenService | null = null,
    private readonly eventPublisher: EventPublisher | null = null,
    private readonly encryptionService: EncryptionService | null = null,
  ) {}

  /**
   * Issues a signed RS256 JWT access token via the OAuth 2.0 Client Credentials grant.
   * Validates client credentials, checks agent status, enforces 10k monthly limit,
   * and writes an async fire-and-forget audit event.
   *
   * @param clientId - The agent UUID acting as client_id.
   * @param clientSecret - The plain-text client secret.
   * @param scope - Space-separated OAuth 2.0 scopes requested.
   * @param ipAddress - Client IP for audit logging.
   * @param userAgent - Client User-Agent for audit logging.
   * @returns The token response with access_token, token_type, expires_in, scope.
   * @throws AuthenticationError if the client credentials are invalid.
   * @throws AuthorizationError if the agent is suspended or decommissioned.
   * @throws FreeTierLimitError if the monthly token limit is reached.
   */
  async issueToken(
    clientId: string,
    clientSecret: string,
    scope: string,
    ipAddress: string,
    userAgent: string,
  ): Promise<ITokenResponse> {
    // Look up the agent
    const agent = await this.agentRepository.findById(clientId);
    if (!agent) {
      void this.auditService.logEvent(
        clientId,
        'auth.failed',
        'failure',
        ipAddress,
        userAgent,
        { reason: 'agent_not_found', clientId },
      );
      throw new AuthenticationError('Client authentication failed. Invalid client_id or client_secret.');
    }

    // Find active credentials for the agent and verify secret
    const { credentials } = await this.credentialRepository.findByAgentId(clientId, {
      status: 'active',
      page: 1,
      limit: 100,
    });

    let credentialVerified = false;
    for (const cred of credentials) {
      const credRow = await this.credentialRepository.findById(cred.credentialId);
      if (credRow) {
        // Check expiry before attempting secret verification
        if (credRow.expiresAt !== null && credRow.expiresAt < new Date()) {
          continue;
        }

        let matches: boolean;
        if (credRow.vaultPath !== null && this.vaultClient !== null) {
          // Phase 2: verify against Vault-stored secret
          // vault_path may be encrypted — decryption is not needed here since
          // verifySecret uses agent/credential IDs to locate the Vault entry.
          matches = await this.vaultClient.verifySecret(
            clientId,
            credRow.credentialId,
            clientSecret,
          );
        } else {
          // Phase 1: verify against bcrypt hash.
          // Decrypt the stored hash if EncryptionService is configured and the
          // value appears to be encrypted (backward-compat for pre-encryption rows).
          let secretHash = credRow.secretHash;
          if (
            this.encryptionService !== null &&
            this.encryptionService.isEncrypted(secretHash)
          ) {
            secretHash = await this.encryptionService.decryptColumn(secretHash);
          }
          matches = await verifySecret(clientSecret, secretHash);
        }

        if (matches) {
          credentialVerified = true;
          break;
        }
      }
    }

    if (!credentialVerified) {
      void this.auditService.logEvent(
        clientId,
        'auth.failed',
        'failure',
        ipAddress,
        userAgent,
        { reason: 'invalid_client_secret', clientId },
      );
      throw new AuthenticationError('Client authentication failed. Invalid client_id or client_secret.');
    }

    // Check agent status
    if (agent.status === 'suspended') {
      void this.auditService.logEvent(
        clientId,
        'auth.failed',
        'failure',
        ipAddress,
        userAgent,
        { reason: 'agent_suspended', clientId },
      );
      throw new AuthorizationError('Agent is currently suspended and cannot obtain tokens.');
    }

    if (agent.status === 'decommissioned') {
      void this.auditService.logEvent(
        clientId,
        'auth.failed',
        'failure',
        ipAddress,
        userAgent,
        { reason: 'agent_decommissioned', clientId },
      );
      throw new AuthorizationError('Agent is decommissioned and cannot obtain tokens.');
    }

    // Check monthly token limit
    const monthlyCount = await this.tokenRepository.getMonthlyCount(clientId);
    if (monthlyCount >= FREE_TIER_MAX_MONTHLY_TOKENS) {
      throw new FreeTierLimitError(
        'Free tier monthly token limit of 10,000 requests has been reached.',
        { limit: FREE_TIER_MAX_MONTHLY_TOKENS, current: monthlyCount },
      );
    }

    // Issue the token
    const jti = uuidv4();
    const now = Math.floor(Date.now() / 1000);
    const expiresIn = getTokenExpiresIn();

    const payload: Omit<ITokenPayload, 'iat' | 'exp'> = {
      sub: clientId,
      client_id: clientId,
      scope,
      jti,
      organization_id: agent.organizationId ?? 'org_system',
    };

    const accessToken = signToken(payload, this.privateKey);

    // Increment monthly count (fire-and-forget)
    void this.tokenRepository.incrementMonthlyCount(clientId);

    // Audit event (fire-and-forget — do not await for latency)
    const expiresAtDate = new Date((now + expiresIn) * 1000);
    void this.auditService.logEvent(
      clientId,
      'token.issued',
      'success',
      ipAddress,
      userAgent,
      { scope, expiresAt: expiresAtDate.toISOString() },
    );

    // Instrument: count successful token issuances
    tokensIssuedTotal.inc({ scope });

    // Publish event (fire-and-forget)
    void this.eventPublisher?.publishEvent(
      agent.organizationId ?? 'org_system',
      'token.issued',
      { agentId: clientId, scope, jti },
    );

    const tokenResponse: ITokenResponse = {
      access_token: accessToken,
      token_type: 'Bearer',
      expires_in: expiresIn,
      scope,
    };

    // OIDC: append id_token when the `openid` scope was requested and IDTokenService is wired
    const scopeList = scope.split(' ');
    if (scopeList.includes('openid') && this.idTokenService !== null) {
      const claims = await this.idTokenService.buildIDTokenClaims(agent, clientId, scope);
      tokenResponse.id_token = await this.idTokenService.signIDToken(claims);
    }

    return tokenResponse;
  }

  /**
   * Introspects a token per RFC 7662.
   * Always returns 200; check the `active` field for validity.
   * Scope enforcement (`tokens:read`) is handled upstream by OPA middleware.
   *
   * @param token - The JWT string to introspect.
   * @param callerPayload - The decoded payload of the calling agent's token.
   * @param ipAddress - Client IP for audit logging.
   * @param userAgent - Client User-Agent for audit logging.
   * @returns The introspection response.
   */
  async introspectToken(
    token: string,
    callerPayload: ITokenPayload,
    ipAddress: string,
    userAgent: string,
  ): Promise<IIntrospectResponse> {
    try {
      const payload = verifyToken(token, this.publicKey);
      const revoked = await this.tokenRepository.isRevoked(payload.jti);

      if (revoked) {
        void this.auditService.logEvent(
          callerPayload.sub,
          'token.introspected',
          'success',
          ipAddress,
          userAgent,
          { targetJti: payload.jti, active: false },
        );
        return { active: false };
      }

      void this.auditService.logEvent(
        callerPayload.sub,
        'token.introspected',
        'success',
        ipAddress,
        userAgent,
        { targetJti: payload.jti, active: true },
      );

      return {
        active: true,
        sub: payload.sub,
        client_id: payload.client_id,
        scope: payload.scope,
        token_type: 'Bearer',
        iat: payload.iat,
        exp: payload.exp,
      };
    } catch {
      // Token is invalid or expired — return inactive per RFC 7662
      return { active: false };
    }
  }

  /**
   * Revokes a token per RFC 7009.
   * Idempotent — revoking an already-revoked or expired token returns success.
   * An agent may only revoke its own tokens.
   *
   * @param token - The JWT string to revoke.
   * @param callerPayload - The decoded payload of the calling agent's token.
   * @param ipAddress - Client IP for audit logging.
   * @param userAgent - Client User-Agent for audit logging.
   * @throws AuthorizationError if the caller tries to revoke another agent's token.
   */
  async revokeToken(
    token: string,
    callerPayload: ITokenPayload,
    ipAddress: string,
    userAgent: string,
  ): Promise<void> {
    // Decode the token without verification to extract claims
    const decoded = decodeToken(token);

    if (decoded !== null) {
      // Only the token owner can revoke their own token
      if (decoded.sub !== callerPayload.sub) {
        throw new AuthorizationError('You do not have permission to revoke this token.');
      }

      // Add to revocation list
      const expiresAt = new Date(decoded.exp * 1000);
      await this.tokenRepository.addToRevocationList(decoded.jti, expiresAt);

      void this.auditService.logEvent(
        callerPayload.sub,
        'token.revoked',
        'success',
        ipAddress,
        userAgent,
        { jti: decoded.jti },
      );

      // Publish event (fire-and-forget)
      void this.eventPublisher?.publishEvent(
        callerPayload.organization_id ?? 'org_system',
        'token.revoked',
        { jti: decoded.jti },
      );
    }
    // If token is malformed/undecoded, per RFC 7009 we still return success
  }
}