stack/packages/appservice/src/agent-store.ts

import { createHash, randomBytes, timingSafeEqual } from 'node:crypto';

import { agentSlug } from './agent-registry.dto.js';
import type { AgentSummary } from './agent-registry.dto.js';
import type { AppserviceIntent } from './intent.js';

/** account_data type holding the agent registry on the AS sender user. */
export const AGENTS_ACCOUNT_DATA_TYPE = 'org.uscllc.mosaic_as.agents';

const TOKEN_PREFIX = 'magt_';

interface StoredAgent {
  alias: string;
  host: string;
  display_name?: string;
  created_at: string;
  /** sha256hex of each active token. Plaintext tokens are NEVER stored. */
  token_hashes: string[];
  revoked_at?: string;
}

interface AgentRegistry {
  agents: Record<string, StoredAgent>;
}

const sha256hex = (value: string): string => createHash('sha256').update(value).digest('hex');

const mintToken = (): string => `${TOKEN_PREFIX}${randomBytes(32).toString('base64url')}`;

/**
 * Persists scoped/revocable bridge tokens for agent virtual users in Matrix
 * account_data on the AS sender user (no new infra; survives restart).
 *
 * Tokens are stored only as sha256 hashes (the high-entropy `magt_` token makes
 * plain sha256 safe — no salt/KDF needed since brute force is infeasible).
 *
 * KNOWN v1 LIMIT: Synapse caps a single account_data object (default
 * max_account_data_size, ~100KB). Each agent + hash entry is small, so this
 * supports thousands of agents, but a very large fleet would eventually need a
 * dedicated store. Revoked agents with no active tokens are pruned of hashes
 * (kept as tombstones) to bound growth.
 */
export class AgentTokenStore {
  constructor(private readonly intent: AppserviceIntent) {}

  /** Read the registry fresh from account_data (low-frequency ops favor
   * correctness over caching; verifyToken/list also read fresh). */
  private async read(): Promise<AgentRegistry> {
    const data = await this.intent.getSenderAccountData(AGENTS_ACCOUNT_DATA_TYPE);
    const agents = data?.agents;
    if (agents && typeof agents === 'object') {
      return { agents: agents as Record<string, StoredAgent> };
    }
    return { agents: {} };
  }

  private async write(registry: AgentRegistry): Promise<void> {
    await this.intent.setSenderAccountData(AGENTS_ACCOUNT_DATA_TYPE, {
      agents: registry.agents,
    });
  }

  /** Ensure the virtual user exists, mint a fresh token, store its hash, and
   * return the plaintext token ONCE. Clears any prior revocation. */
  async register(opts: {
    alias: string;
    host: string;
    displayName?: string;
  }): Promise<{ agentUserId: string; token: string }> {
    const slug = agentSlug(opts.alias, opts.host);
    const agentUserId = await this.intent.ensureRegistered(slug);
    if (opts.displayName !== undefined) {
      await this.intent.setDisplayName(slug, opts.displayName);
    }

    const token = mintToken();
    const hash = sha256hex(token);

    const registry = await this.read();
    const existing = registry.agents[agentUserId];
    if (existing) {
      // The agent slug `<alias>-<host>` joins with a `-`, which is also a legal
      // slug char, so distinct pairs can collide on one Matrix id (e.g.
      // a/b-c and a-b/c both -> @agent-a-b-c). They ARE the same Matrix user,
      // but silently overwriting the stored alias/host of a different pair
      // would conflate two logical agents into one token bucket. Reject the
      // ambiguous re-registration instead of overwriting.
      if (existing.alias !== opts.alias || existing.host !== opts.host) {
        throw new Error(
          `agent id collision: ${agentUserId} already registered as ` +
            `${existing.alias}/${existing.host}, refusing ${opts.alias}/${opts.host}`,
        );
      }
      if (opts.displayName !== undefined) existing.display_name = opts.displayName;
      existing.token_hashes = [...existing.token_hashes, hash];
      delete existing.revoked_at;
    } else {
      registry.agents[agentUserId] = {
        alias: opts.alias,
        host: opts.host,
        ...(opts.displayName !== undefined ? { display_name: opts.displayName } : {}),
        created_at: new Date().toISOString(),
        token_hashes: [hash],
      };
    }
    await this.write(registry);
    return { agentUserId, token };
  }

  /** Return the agentUserId bound to an active (non-revoked) token, else null.
   * Constant-time hash comparison; no early-out on match. */
  async verifyToken(token: string): Promise<string | null> {
    if (!token.startsWith(TOKEN_PREFIX)) return null;
    const presented = Buffer.from(sha256hex(token), 'hex');

    const registry = await this.read();
    let matched: string | null = null;
    for (const [agentUserId, agent] of Object.entries(registry.agents)) {
      if (agent.revoked_at) continue;
      for (const stored of agent.token_hashes) {
        const candidate = Buffer.from(stored, 'hex');
        if (candidate.length === presented.length && timingSafeEqual(candidate, presented)) {
          // No early break: keep scanning so timing does not reveal match position.
          matched = agentUserId;
        }
      }
    }
    return matched;
  }

  /** Revoke all active tokens for an agent. Idempotent; returns count revoked. */
  async revoke(agentUserId: string): Promise<number> {
    const registry = await this.read();
    const agent = registry.agents[agentUserId];
    if (!agent) return 0;
    const count = agent.token_hashes.length;
    agent.token_hashes = [];
    agent.revoked_at = new Date().toISOString();
    await this.write(registry);
    return count;
  }

  /** List agents with at least one active token (never advertise revoked/phantom). */
  async list(): Promise<AgentSummary[]> {
    const registry = await this.read();
    const out: AgentSummary[] = [];
    for (const [agentUserId, agent] of Object.entries(registry.agents)) {
      if (agent.revoked_at || agent.token_hashes.length === 0) continue;
      out.push({
        agent_user_id: agentUserId,
        alias: agent.alias,
        host: agent.host,
        ...(agent.display_name !== undefined ? { display_name: agent.display_name } : {}),
        created_at: agent.created_at,
        active_token_count: agent.token_hashes.length,
      });
    }
    return out;
  }
}