diff --git a/setup/controlplane.ts b/setup/controlplane.ts index 0c6d195..50c774e 100644 --- a/setup/controlplane.ts +++ b/setup/controlplane.ts @@ -15,7 +15,10 @@ import path from 'path'; import pg from 'pg'; -import { ensureControlplaneAgentApiKey } from '../src/controlplane-agent-keys.js'; +import { + ensureControlplaneAgentApiKey, + syncControlplaneAgentApiKeysFromDisk, +} from '../src/controlplane-agent-keys.js'; import { BETTER_AUTH_SECRET, BETTER_AUTH_URL, @@ -85,6 +88,7 @@ export async function run(_args: string[]): Promise { await runSchemaMigration(pool); logger.info('Schema ready'); appendSetupLog('schema migration: ready'); + await syncControlplaneAgentApiKeysFromDisk(pool); // ── 2. Provision default agents ────────────────────────────────── const defaultAgents = getDefaultAgents(TENANT_ID); diff --git a/src/controlplane-agent-keys.test.ts b/src/controlplane-agent-keys.test.ts new file mode 100644 index 0000000..63c1385 --- /dev/null +++ b/src/controlplane-agent-keys.test.ts @@ -0,0 +1,124 @@ +import fs from 'fs'; +import path from 'path'; + +import type { Pool } from 'pg'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const { + tmpDir, + getAgentMock, + getAgentsMock, + setAgentApiKeyMock, +} = vi.hoisted(() => ({ + tmpDir: `${process.cwd()}/tmp/controlplane-agent-keys-test`, + getAgentMock: vi.fn(), + getAgentsMock: vi.fn(), + setAgentApiKeyMock: vi.fn(), +})); + +vi.mock('./config.js', () => ({ + DATA_DIR: tmpDir, + TENANT_ID: '', +})); + +vi.mock('./platform-identity.js', () => ({ + SERVICE_NAME: 'clawdie', +})); + +vi.mock('./controlplane-db.js', async () => { + const actual = await vi.importActual( + './controlplane-db.js', + ); + return { + ...actual, + getAgent: getAgentMock, + getAgents: getAgentsMock, + setAgentApiKey: setAgentApiKeyMock, + }; +}); + +import { + deriveAgentApiKeyLookupHash, + generateAgentApiKey, +} from './controlplane-db.js'; +import { + ensureControlplaneAgentApiKey, + getControlplaneAgentApiKeyPath, + readControlplaneAgentApiKey, + syncControlplaneAgentApiKeysFromDisk, + writeControlplaneAgentApiKey, +} from './controlplane-agent-keys.js'; + +describe('controlplane-agent-keys', () => { + beforeEach(() => { + fs.rmSync(tmpDir, { recursive: true, force: true }); + fs.mkdirSync(tmpDir, { recursive: true }); + getAgentMock.mockReset(); + getAgentsMock.mockReset(); + setAgentApiKeyMock.mockReset(); + }); + + it('writes and reads agent API keys from disk', () => { + writeControlplaneAgentApiKey('sysadmin', 'sk-test-123'); + expect(readControlplaneAgentApiKey('sysadmin')).toBe('sk-test-123'); + expect(fs.existsSync(getControlplaneAgentApiKeyPath('sysadmin'))).toBe(true); + }); + + it('reuses an in-sync on-disk key without rewriting DB state', async () => { + const apiKey = 'sk-test-123'; + writeControlplaneAgentApiKey('sysadmin', apiKey); + getAgentMock.mockResolvedValue({ + id: 'sysadmin', + api_key_hash: 'argon2id$present', + api_key_lookup_hash: deriveAgentApiKeyLookupHash(apiKey), + }); + + await expect( + ensureControlplaneAgentApiKey({} as Pool, 'sysadmin'), + ).resolves.toBe(apiKey); + expect(setAgentApiKeyMock).not.toHaveBeenCalled(); + }); + + it('backfills DB state when an old on-disk key has no lookup hash row', async () => { + const apiKey = 'sk-test-legacy'; + writeControlplaneAgentApiKey('sysadmin', apiKey); + getAgentMock.mockResolvedValue({ + id: 'sysadmin', + api_key_hash: 'argon2id$present', + api_key_lookup_hash: null, + }); + + await ensureControlplaneAgentApiKey({} as Pool, 'sysadmin'); + expect(setAgentApiKeyMock).toHaveBeenCalledWith( + expect.anything(), + 'sysadmin', + apiKey, + ); + }); + + it('syncs only agents that have key files on disk', async () => { + const sysadminKey = generateAgentApiKey(); + writeControlplaneAgentApiKey('sysadmin', sysadminKey); + getAgentsMock.mockResolvedValue([ + { + id: 'sysadmin', + api_key_hash: 'argon2id$present', + api_key_lookup_hash: null, + }, + { + id: 'db-admin', + api_key_hash: 'argon2id$present', + api_key_lookup_hash: null, + }, + ]); + + await syncControlplaneAgentApiKeysFromDisk({} as Pool); + + expect(setAgentApiKeyMock).toHaveBeenCalledTimes(1); + expect(setAgentApiKeyMock).toHaveBeenCalledWith( + expect.anything(), + 'sysadmin', + sysadminKey, + ); + }); +}); diff --git a/src/controlplane-agent-keys.ts b/src/controlplane-agent-keys.ts index e179c98..db7071d 100644 --- a/src/controlplane-agent-keys.ts +++ b/src/controlplane-agent-keys.ts @@ -5,10 +5,11 @@ import type { Pool } from 'pg'; import { DATA_DIR, TENANT_ID } from './config.js'; import { + deriveAgentApiKeyLookupHash, generateAgentApiKey, getAgent, + getAgents, setAgentApiKey, - verifyPassword, } from './controlplane-db.js'; import { SERVICE_NAME } from './platform-identity.js'; @@ -57,6 +58,23 @@ export function writeControlplaneAgentApiKey( fs.chmodSync(target, 0o600); } +export async function syncControlplaneAgentApiKeysFromDisk( + pool: Pool, +): Promise { + const agents = await getAgents(pool); + for (const agent of agents) { + const apiKey = readControlplaneAgentApiKey(agent.id); + if (!apiKey) continue; + const lookupHash = deriveAgentApiKeyLookupHash(apiKey); + if (agent.api_key_hash && agent.api_key_lookup_hash === lookupHash) { + syncedApiKeys.set(agent.id, apiKey); + continue; + } + await setAgentApiKey(pool, agent.id, apiKey); + syncedApiKeys.set(agent.id, apiKey); + } +} + export async function ensureControlplaneAgentApiKey( pool: Pool, agentId: string, @@ -77,8 +95,9 @@ export async function ensureControlplaneAgentApiKey( return apiKey; } + const lookupHash = deriveAgentApiKeyLookupHash(apiKey); const inSync = - !!agent.api_key_hash && (await verifyPassword(apiKey, agent.api_key_hash)); + !!agent.api_key_hash && agent.api_key_lookup_hash === lookupHash; if (!inSync) { await setAgentApiKey(pool, agentId, apiKey); } diff --git a/src/controlplane-api.test.ts b/src/controlplane-api.test.ts index 737b38e..3c46995 100644 --- a/src/controlplane-api.test.ts +++ b/src/controlplane-api.test.ts @@ -14,6 +14,7 @@ import path from 'path'; import type { ControlplaneApiOptions } from './controlplane-api.js'; import { + deriveAgentApiKeyLookupHash, generateAgentApiKey, hashPassword, type Agent, @@ -102,9 +103,10 @@ const validAgentApiKey = generateAgentApiKey(); const validAgentAuth = `Bearer ${validAgentApiKey}`; const hashedTestPass = await hashPassword('testpass'); const hashedAgentApiKey = await hashPassword(validAgentApiKey); +const validAgentApiKeyLookupHash = deriveAgentApiKeyLookupHash(validAgentApiKey); const defaultAgents: Agent[] = [ { id: 'clawdie', role: 'orchestrator', adapter: 'pi-local', heartbeat_enabled: false, heartbeat_interval_sec: null, budget_allocation_pct: 80, api_key_hash: null, created_at: new Date() }, - { id: 'sysadmin', role: 'sysadmin', adapter: 'pi-local', heartbeat_enabled: true, heartbeat_interval_sec: 86400, budget_allocation_pct: 10, api_key_hash: hashedAgentApiKey, created_at: new Date() }, + { id: 'sysadmin', role: 'sysadmin', adapter: 'pi-local', heartbeat_enabled: true, heartbeat_interval_sec: 86400, budget_allocation_pct: 10, api_key_hash: hashedAgentApiKey, api_key_lookup_hash: validAgentApiKeyLookupHash, created_at: new Date() }, ]; interface SetupOptions { @@ -138,7 +140,12 @@ function defaultMockQuery( if (/SELECT.*FROM operators/i.test(sql) && params?.[0] === 'admin') { return { rows: [{ username: 'admin', hashed_password: hashedTestPass, created_at: new Date() }] }; } - if (/FROM agents/i.test(sql)) return { rows: agents }; + if (/FROM agents/i.test(sql)) { + if (/api_key_lookup_hash\s*=\s*\$1/i.test(sql)) { + return { rows: agents.filter((agent) => agent.api_key_lookup_hash === params?.[0]) }; + } + return { rows: agents }; + } if (/FROM agent_budgets/i.test(sql)) { if (params?.[0]) { return { rows: budgets.filter((b) => b.agent_id === params[0]) }; diff --git a/src/controlplane-db.test.ts b/src/controlplane-db.test.ts index 40f7e25..a8f25a3 100644 --- a/src/controlplane-db.test.ts +++ b/src/controlplane-db.test.ts @@ -16,6 +16,7 @@ import os from 'os'; import { getDefaultAgents, DEFAULT_AGENTS, + deriveAgentApiKeyLookupHash, generateAgentApiKey, getAgentByApiKey, hashPassword, @@ -211,6 +212,7 @@ describe('getAgentByApiKey', () => { it('returns the matching agent when a hash matches', async () => { const apiKey = generateAgentApiKey(); const apiKeyHash = await hashPassword(apiKey); + const apiKeyLookupHash = deriveAgentApiKeyLookupHash(apiKey); const pool = { query: vi.fn().mockResolvedValue({ rows: [ @@ -222,6 +224,7 @@ describe('getAgentByApiKey', () => { heartbeat_interval_sec: 86400, budget_allocation_pct: 10, api_key_hash: apiKeyHash, + api_key_lookup_hash: apiKeyLookupHash, created_at: new Date(), }, ], @@ -232,10 +235,15 @@ describe('getAgentByApiKey', () => { id: 'sysadmin', role: 'sysadmin', }); + expect(pool.query).toHaveBeenCalledWith( + expect.stringContaining('WHERE api_key_lookup_hash = $1'), + [apiKeyLookupHash], + ); }); it('returns null when no agent hash matches the token', async () => { const otherApiKeyHash = await hashPassword('different-token'); + const otherApiKeyLookupHash = deriveAgentApiKeyLookupHash('different-token'); const pool = { query: vi.fn().mockResolvedValue({ rows: [ @@ -247,6 +255,7 @@ describe('getAgentByApiKey', () => { heartbeat_interval_sec: 86400, budget_allocation_pct: 10, api_key_hash: otherApiKeyHash, + api_key_lookup_hash: otherApiKeyLookupHash, created_at: new Date(), }, ], diff --git a/src/controlplane-db.ts b/src/controlplane-db.ts index 4cc97f7..e518eb5 100644 --- a/src/controlplane-db.ts +++ b/src/controlplane-db.ts @@ -31,6 +31,7 @@ export interface Agent { heartbeat_interval_sec: number | null; budget_allocation_pct: number; api_key_hash: string | null; + api_key_lookup_hash?: string | null; created_at: Date; } @@ -170,6 +171,7 @@ CREATE TABLE IF NOT EXISTS agents ( heartbeat_interval_sec INTEGER, budget_allocation_pct INTEGER NOT NULL, api_key_hash TEXT, + api_key_lookup_hash TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); @@ -233,6 +235,7 @@ export async function runSchemaMigration(pool: pg.Pool): Promise { await ensureRoleConstraint(pool); await ensureApprovalDecisionColumn(pool); await ensureParentTaskId(pool); + await ensureAgentApiKeyLookupColumn(pool); await ensureModelCatalogSchema(pool); // Seed default agents (idempotent via upsert) @@ -356,6 +359,17 @@ async function ensureParentTaskId(pool: pg.Pool): Promise { ); } +async function ensureAgentApiKeyLookupColumn(pool: pg.Pool): Promise { + await pool.query( + 'ALTER TABLE agents ADD COLUMN IF NOT EXISTS api_key_lookup_hash TEXT', + ); + await pool.query( + `CREATE UNIQUE INDEX IF NOT EXISTS idx_agents_api_key_lookup_hash + ON agents(api_key_lookup_hash) + WHERE api_key_lookup_hash IS NOT NULL`, + ); +} + // ── Agent queries ────────────────────────────────────────────────────────── export async function upsertAgent( @@ -403,16 +417,21 @@ export function generateAgentApiKey(length = 48): string { return generatePassword(length); } +export function deriveAgentApiKeyLookupHash(apiKey: string): string { + return crypto.createHash('sha256').update(apiKey).digest('base64url'); +} + export async function setAgentApiKey( pool: pg.Pool, agentId: string, apiKey: string, ): Promise { const hashed = await hashPassword(apiKey); - await pool.query('UPDATE agents SET api_key_hash = $2 WHERE id = $1', [ - agentId, - hashed, - ]); + const lookupHash = deriveAgentApiKeyLookupHash(apiKey); + await pool.query( + 'UPDATE agents SET api_key_hash = $2, api_key_lookup_hash = $3 WHERE id = $1', + [agentId, hashed, lookupHash], + ); } export async function getAgentByApiKey( @@ -420,16 +439,16 @@ export async function getAgentByApiKey( apiKey: string, ): Promise { if (!apiKey.trim()) return null; + const lookupHash = deriveAgentApiKeyLookupHash(apiKey); const result = await pool.query( - 'SELECT * FROM agents WHERE api_key_hash IS NOT NULL ORDER BY id', + `SELECT * FROM agents + WHERE api_key_lookup_hash = $1 + LIMIT 1`, + [lookupHash], ); - for (const agent of result.rows) { - if (!agent.api_key_hash) continue; - if (await verifyPassword(apiKey, agent.api_key_hash)) { - return agent; - } - } - return null; + const agent = result.rows[0]; + if (!agent?.api_key_hash) return null; + return (await verifyPassword(apiKey, agent.api_key_hash)) ? agent : null; } // ── Budget queries ───────────────────────────────────────────────────────── diff --git a/src/index.ts b/src/index.ts index 1c72003..6d3f87f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -4,6 +4,7 @@ import path from 'path'; import { ensureControlplaneAgentApiKey, getControlplaneOrchestratorId, + syncControlplaneAgentApiKeysFromDisk, } from './controlplane-agent-keys.js'; import { AGENT_ENGINE, @@ -1170,6 +1171,7 @@ async function main(): Promise { // Controlplane uses the memory pool — ensure schema + agents + budgets exist. await runSchemaMigration(getMemoryPool()); + await syncControlplaneAgentApiKeysFromDisk(getMemoryPool()); const sessionDir = getControlplaneSessionDir(TMP_DIR); fs.mkdirSync(sessionDir, { recursive: true });