From c161bfdc4b3d01de69fa3ba29f4143f3cb9a778c Mon Sep 17 00:00:00 2001 From: Oliver Gibbs Date: Wed, 16 Sep 2026 06:57:18 +0000 Subject: [PATCH] fix(vender): refuse agent credentials for resources outside the caller organisation The worker forwards the organisation carried on its dispatch payload to the credential vender and refuses to request credentials without one. The vender now requires that organisation, resolves the owning organisation of the agent and of every declared datastore and integration, and rejects the whole request before any role is created or assumed when an id belongs to another organisation or cannot be resolved. The policy manager refuses to add an agent principal to a datastore or integration role trust policy, so the boundary that previously held only by construction is now enforced and pinned by test. The vender Lambda gains read access to the agent-config, datastores and integrations tables. --- .../test_credential_vender_org_forwarding.py | 98 +++++++ arbiter/workerWrapper/index.py | 32 ++- backend/bin/app.ts | 5 + backend/lib/arbiter-stack.ts | 27 ++ backend/lib/backend-stack.ts | 9 + ...gent-credential-vender-org-scoping.test.ts | 246 ++++++++++++++++++ .../__tests__/agent-credential-vender.test.ts | 158 ++++++++--- backend/src/lambda/agent-credential-vender.ts | 151 ++++++++++- ...licy-manager.agent-trust-rejection.test.ts | 164 ++++++++++++ backend/src/utils/policy-manager.ts | 21 ++ 10 files changed, 861 insertions(+), 50 deletions(-) create mode 100644 arbiter/workerWrapper/__tests__/test_credential_vender_org_forwarding.py create mode 100644 backend/src/lambda/__tests__/agent-credential-vender-org-scoping.test.ts create mode 100644 backend/src/utils/__tests__/policy-manager.agent-trust-rejection.test.ts diff --git a/arbiter/workerWrapper/__tests__/test_credential_vender_org_forwarding.py b/arbiter/workerWrapper/__tests__/test_credential_vender_org_forwarding.py new file mode 100644 index 00000000..3bf82c44 --- /dev/null +++ b/arbiter/workerWrapper/__tests__/test_credential_vender_org_forwarding.py @@ -0,0 +1,98 @@ +"""Wave 2b (branch fix/vender-org-scoping) — workerWrapper forwards orgId +from the Supervisor dispatch payload to the credential vender. + +The Supervisor's ``process_agent_call`` dispatch payload already carries a +server-derived, non-empty ``orgId`` (see supervisor/index.py ~L1024). This +slice makes ``process_event`` read that field off the SQS event and forward +it as ``org`` on the credential-vender invoke payload built by +``get_scoped_credentials``. Fail CLOSED — no credential request at all, +plus a clear log — when the dispatch payload lacks orgId, rather than +silently vending without an org (which is exactly the gap wave 2b closes). +""" + +import sys +import os +import json +from unittest.mock import patch, MagicMock + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +os.environ.setdefault("AGENT_CONFIG_TABLE", "fake-table") +os.environ.setdefault("AGENT_BUCKET_NAME", "fake-bucket") +os.environ.setdefault("COMPLETION_BUS_NAME", "fake-bus") +os.environ.setdefault("CREDENTIAL_VENDER_FUNCTION", "fake-vender-fn") + +import index # noqa: E402 + + +def _base_agent_config(): + return { + "config": json.dumps({ + "tools": [], + "filename": "agent.zip", + "description": "test agent", + "requiredPermissions": {"dataStores": ["ds-1"]}, + }) + } + + +def _run_process_event(event, agent_config=None): + agent_config = agent_config or _base_agent_config() + with patch.object(index, "load_config_from_dynamodb", return_value=agent_config), \ + patch.object(index, "load_file_from_s3_into_tmp"), \ + patch.object(index, "run_agent_in_subprocess", return_value="ok"), \ + patch.object(index, "post_task_complete"), \ + patch.object(index, "CREDENTIAL_VENDER_FUNCTION", "fake-vender-fn"), \ + patch.object(index, "TOOLS_CONFIG_TABLE", None): + index.process_event(event, {}) + + +class TestOrgIdForwardedToVender: + def test_org_id_forwarded_on_vender_invoke_payload(self): + event = { + "orchestration_id": "orch-1", + "agent_use_id": "use-1", + "agent_input": {"x": 1}, + "node": "agent1", + "orgId": "org-abc", + } + mock_lambda_client = MagicMock() + mock_lambda_client.invoke.return_value = { + "Payload": MagicMock(read=lambda: json.dumps({"credentials": {}}).encode()) + } + with patch.object(index, "_get_lambda_client", return_value=mock_lambda_client): + _run_process_event(event) + + assert mock_lambda_client.invoke.called + payload = json.loads(mock_lambda_client.invoke.call_args.kwargs["Payload"]) + assert payload["org"] == "org-abc" + + def test_missing_org_id_fails_closed_no_vender_call(self): + """No orgId on the dispatch payload -> get_scoped_credentials must + refuse to invoke the vender at all (fail closed), not silently vend + org-less credentials.""" + event = { + "orchestration_id": "orch-2", + "agent_use_id": "use-2", + "agent_input": {"x": 1}, + "node": "agent1", + } + mock_lambda_client = MagicMock() + with patch.object(index, "_get_lambda_client", return_value=mock_lambda_client): + _run_process_event(event) + + assert not mock_lambda_client.invoke.called + + def test_empty_string_org_id_fails_closed(self): + event = { + "orchestration_id": "orch-3", + "agent_use_id": "use-3", + "agent_input": {"x": 1}, + "node": "agent1", + "orgId": "", + } + mock_lambda_client = MagicMock() + with patch.object(index, "_get_lambda_client", return_value=mock_lambda_client): + _run_process_event(event) + + assert not mock_lambda_client.invoke.called diff --git a/arbiter/workerWrapper/index.py b/arbiter/workerWrapper/index.py index 9237b1e5..ce2da1be 100644 --- a/arbiter/workerWrapper/index.py +++ b/arbiter/workerWrapper/index.py @@ -471,12 +471,18 @@ def _merge_required_permissions(agent_permissions: dict | None, tool_bindings: d return merged if merged else None -def get_scoped_credentials(agent_name: str, required_permissions: dict, app_id: str | None = None) -> dict | None: +def get_scoped_credentials(agent_name: str, required_permissions: dict, app_id: str | None = None, org_id: str | None = None) -> dict | None: """ Invoke the credential vender Lambda to get scoped IAM credentials for this agent based on its declared permissions. When app_id is provided, the credential vender uses the app-scoped IAM role (citadel-agent-{appId}) instead of the agent-level role (Req 4 AC 5). + + ``org_id`` is the SERVER-DERIVED org carried on the Supervisor's worker + dispatch payload (never trusted from agent/subprocess output). Wave 2b + (branch fix/vender-org-scoping): fail CLOSED — issue no credential + request at all, with a clear log — when org_id is missing/empty, rather + than letting the vender silently grant org-blind AssumeRole access. """ if not CREDENTIAL_VENDER_FUNCTION: print("CREDENTIAL_VENDER_FUNCTION not set, skipping credential vending") @@ -485,10 +491,22 @@ def get_scoped_credentials(agent_name: str, required_permissions: dict, app_id: if not required_permissions: return None + if not isinstance(org_id, str) or not org_id: + print(json.dumps({ + 'level': 'ERROR', + 'component': 'WorkerWrapper', + 'action': 'credential_vend_refused_no_org', + 'agentId': agent_name, + 'error': 'org_id missing/empty on dispatch payload — refusing to ' + 'request scoped credentials (fail closed, no vender call).', + })) + return None + try: payload_data = { 'agentId': agent_name, 'requiredPermissions': required_permissions, + 'org': org_id, } if app_id: payload_data['appId'] = app_id @@ -1271,7 +1289,10 @@ def _process_workflow_node(event, message_attributes=None): ) required_permissions = config.get('requiredPermissions') - scoped_credentials = get_scoped_credentials(msg.agent_id, required_permissions) + scoped_credentials = get_scoped_credentials( + msg.agent_id, required_permissions, + org_id=_resolve_execution_org_id(msg.execution_id), + ) fileName = config['filename'] load_file_from_s3_into_tmp(os.environ["AGENT_BUCKET_NAME"], fileName) @@ -1420,6 +1441,11 @@ def process_event(event, context, message_attributes=None): model_override = event.get('modelOverride') system_prompt_addition = event.get('systemPromptAddition') app_id = event.get('appId') # App-scoped credential vending (Req 4 AC 5) + # Wave 2b (fix/vender-org-scoping): the Supervisor's process_agent_call + # dispatch payload already carries a server-derived, non-empty orgId + # (supervisor/index.py process_agent_call). Forward it verbatim to the + # credential vender — never re-derive or default it here. + org_id = event.get('orgId') # CIT-102 Pass B: frozen contract keys (Pass A dispatch payload — # supervisor.process_agent_call). Absent-tolerant reads: a non-eval @@ -1553,7 +1579,7 @@ def process_event(event, context, message_attributes=None): # Vend scoped credentials based on merged permissions # When appId is present, use app-scoped IAM role (Req 4 AC 5) # Eventual consistency: binding updates are picked up on next invocation (Req 10.8) - scoped_credentials = get_scoped_credentials(agent_name, required_permissions, app_id=app_id) + scoped_credentials = get_scoped_credentials(agent_name, required_permissions, app_id=app_id, org_id=org_id) fileName = config['filename'] print("loading file from s3...") diff --git a/backend/bin/app.ts b/backend/bin/app.ts index 056dfe4e..bc7b0465 100644 --- a/backend/bin/app.ts +++ b/backend/bin/app.ts @@ -248,6 +248,11 @@ const arbiterStack = new ArbiterStack(app, `citadel-arbiter-${environment}`, { // forward-compatible wiring for the fabricator's // design-assessment precondition gate. agentDesignAssessmentsTable: backendStack.agentDesignAssessmentsTable, + // Wave 2b (fix/vender-org-scoping): the AgentCredentialVender Lambda + // resolves declared dataStore/integration ids to their owning org + // against these tables before granting any AssumeRole policy. + dataStoresTable: backendStack.dataStoresTable, + integrationsTable: backendStack.integrationsTable, registryArn: backendStack.registryArn, registryId: backendStack.registryId, // Governance UI Wave 1: the new governance-ui-resolver lives in diff --git a/backend/lib/arbiter-stack.ts b/backend/lib/arbiter-stack.ts index cbfdbd92..016d196e 100644 --- a/backend/lib/arbiter-stack.ts +++ b/backend/lib/arbiter-stack.ts @@ -154,6 +154,14 @@ interface ArbiterStackProps extends cdk.StackProps { // gate is forward-compatible -- when the table/prop is absent the // gate's env-var fallback simply no-ops. agentDesignAssessmentsTable?: dynamodb.Table; + // Wave 2b (fix/vender-org-scoping): optional read handles so the + // AgentCredentialVender Lambda can resolve declared dataStore/integration + // ids to their owning org before granting any AssumeRole policy. Optional + // because some test paths construct ArbiterStack without them; when + // absent the vender's lookups resolve nothing and every org-scoped + // request that declares a dataStore/integration id fails closed. + dataStoresTable?: dynamodb.Table; + integrationsTable?: dynamodb.Table; registryArn?: string; registryId?: string; // Governance UI Wave 1: optional AppSync API handle so the new @@ -543,6 +551,17 @@ export class ArbiterStack extends cdk.Stack { memorySize: 256, environment: { ENVIRONMENT: props.environment, + AGENT_CONFIG_TABLE: props.agentConfigTable.tableName, + // Wave 2b (fix/vender-org-scoping): omitted entirely when the + // stack is constructed without these optional tables (test paths); + // the vender's org-resolution lookups then find nothing and + // every request declaring a dataStore/integration id fails closed. + ...(props.dataStoresTable && { + DATASTORES_TABLE: props.dataStoresTable.tableName, + }), + ...(props.integrationsTable && { + INTEGRATIONS_TABLE: props.integrationsTable.tableName, + }), }, initialPolicy: [ new PolicyStatement({ @@ -571,6 +590,14 @@ export class ArbiterStack extends cdk.Stack { }, ); + // Wave 2b (fix/vender-org-scoping): read-only access to the agent + // config table (agent's own orgId) and, when provisioned, the + // datastore/integration tables the vender resolves declared ids + // against before granting any AssumeRole policy. + props.agentConfigTable.grantReadData(credentialVenderLambda); + props.dataStoresTable?.grantReadData(credentialVenderLambda); + props.integrationsTable?.grantReadData(credentialVenderLambda); + // Tool-call idempotency ledger (PR1). Org-scoped, TTL'd operational // dedupe table — NOT an audit artifact (distinct from the 90-day // governance ledger). PK = orgId#executionId, SK = diff --git a/backend/lib/backend-stack.ts b/backend/lib/backend-stack.ts index 3953650d..b8ef0b3a 100644 --- a/backend/lib/backend-stack.ts +++ b/backend/lib/backend-stack.ts @@ -47,6 +47,13 @@ export class BackendStack extends cdk.Stack { public readonly accessLogsBucket: Bucket; public readonly workflowsTable: dynamodb.Table; public readonly appsTable: dynamodb.Table; + /** + * Wave 2b (fix/vender-org-scoping): exposed publicly so ArbiterStack's + * AgentCredentialVender Lambda can be granted read-only access and + * resolve declared dataStore/integration ids to their owning org. + */ + public readonly dataStoresTable: dynamodb.Table; + public readonly integrationsTable: dynamodb.Table; /** * Model-config table — exposed publicly (CIT-026 replay package, * design §4) so TelemetryStack can grant read-only access for the @@ -357,6 +364,7 @@ export class BackendStack extends cdk.Stack { }, projectionType: dynamodb.ProjectionType.ALL, }); + this.integrationsTable = integrationsTable; // Workflows Table this.workflowsTable = new dynamodb.Table(this, "WorkflowsTable", { @@ -2745,6 +2753,7 @@ export class BackendStack extends cdk.Stack { removalPolicy: cdk.RemovalPolicy.DESTROY, pointInTimeRecoverySpecification: { pointInTimeRecoveryEnabled: true }, }); + this.dataStoresTable = dataStoresTable; dataStoresTable.addGlobalSecondaryIndex({ indexName: "OrgIndex", diff --git a/backend/src/lambda/__tests__/agent-credential-vender-org-scoping.test.ts b/backend/src/lambda/__tests__/agent-credential-vender-org-scoping.test.ts new file mode 100644 index 00000000..54578384 --- /dev/null +++ b/backend/src/lambda/__tests__/agent-credential-vender-org-scoping.test.ts @@ -0,0 +1,246 @@ +/** + * TDD Tests for agent-credential-vender org scoping + * (wave 2b, branch fix/vender-org-scoping) + * + * Covers: + * - missing orgId on the request -> fail closed (no policy/role calls) + * - agent's own orgId mismatch vs request org -> rejected + * - declared datastore/integration id belonging to a different org -> whole + * request rejected (not silently dropped) + * - declared id that cannot be resolved -> rejected + * - same-org ids -> allowed, policies computed normally + */ + +const mockGetAccountContext = jest.fn(); +const mockEnsureRole = jest.fn(); +const mockAssumeScopedRole = jest.fn(); +const mockDynamoSend = jest.fn(); + +jest.mock("../../utils/policy-manager", () => { + const MockPolicyManager = jest.fn().mockImplementation(() => ({ + getAccountContext: mockGetAccountContext, + ensureRole: mockEnsureRole, + assumeScopedRole: mockAssumeScopedRole, + })); + ( + MockPolicyManager as unknown as { + getRoleName: (id: string, scope: string) => string; + } + ).getRoleName = (id: string, scope: string) => { + const prefixes: Record = { + datastore: "citadel-ds-", + integration: "citadel-int-", + agent: "citadel-agent-", + }; + return `${prefixes[scope] || "citadel-ds-"}${id}`; + }; + ( + MockPolicyManager as unknown as { + buildPolicyDocument: ( + policies: Array<{ actions: string[]; resources: string[] }>, + ) => unknown; + } + ).buildPolicyDocument = ( + policies: Array<{ actions: string[]; resources: string[] }>, + ) => ({ + Version: "2012-10-17", + Statement: policies.map((p) => ({ + Effect: "Allow", + Action: p.actions, + Resource: p.resources, + })), + }); + return { PolicyManager: MockPolicyManager }; +}); + +jest.mock("@aws-sdk/lib-dynamodb", () => { + const actual = jest.requireActual("@aws-sdk/lib-dynamodb"); + return { + ...actual, + DynamoDBDocumentClient: { + from: jest.fn(() => ({ send: mockDynamoSend })), + }, + }; +}); + +process.env.AGENT_CONFIG_TABLE = "agent-config-table"; +process.env.DATASTORES_TABLE = "datastores-table"; +process.env.INTEGRATIONS_TABLE = "integrations-table"; + +import { handler } from "../agent-credential-vender"; + +const AGENT_ORG_A = { agentId: "agent-1", orgId: "org-a" }; +const DS_ORG_A = { dataStoreId: "ds-a", orgId: "org-a" }; +const DS_ORG_B = { dataStoreId: "ds-b", orgId: "org-b" }; +const INT_ORG_A = { integrationId: "int-a", orgId: "org-a" }; +const INT_ORG_B = { integrationId: "int-b", orgId: "org-b" }; + +function mockDynamoResolvers(opts: { + agent?: Record | null; + dataStores?: Record | null>; + integrations?: Record | null>; +}) { + mockDynamoSend.mockImplementation((command: unknown) => { + const input = (command as { input?: Record }).input || {}; + const tableName = input.TableName as string | undefined; + if (tableName === "agent-config-table") { + return Promise.resolve({ + Item: opts.agent === undefined ? AGENT_ORG_A : opts.agent, + }); + } + if (tableName === "datastores-table") { + const key = input.Key as { dataStoreId: string } | undefined; + const id = key?.dataStoreId as string; + const row = opts.dataStores?.[id]; + return Promise.resolve({ Item: row === undefined ? undefined : row }); + } + if (tableName === "integrations-table") { + const values = input.ExpressionAttributeValues as + Record | undefined; + const id = values?.[":id"] as string; + const row = opts.integrations?.[id]; + return Promise.resolve({ Items: row ? [row] : [] }); + } + return Promise.resolve({}); + }); +} + +describe("agent-credential-vender org scoping", () => { + beforeEach(() => { + mockGetAccountContext.mockClear(); + mockEnsureRole.mockClear(); + mockAssumeScopedRole.mockClear(); + mockDynamoSend.mockClear(); + + mockGetAccountContext.mockResolvedValue({ + accountId: "123456789012", + region: "us-west-2", + }); + mockEnsureRole.mockResolvedValue(undefined); + mockAssumeScopedRole.mockResolvedValue({ + accessKeyId: "AKIA_SCOPED", + secretAccessKey: "SECRET_SCOPED", + sessionToken: "TOKEN_SCOPED", + }); + + mockDynamoResolvers({ + agent: AGENT_ORG_A, + dataStores: { "ds-a": DS_ORG_A, "ds-b": DS_ORG_B }, + integrations: { "int-a": INT_ORG_A, "int-b": INT_ORG_B }, + }); + }); + + test("rejects request missing org (fail closed, never defaults org)", async () => { + const result = await handler({ + agentId: "agent-1", + requiredPermissions: { dataStores: ["ds-a"] }, + } as never); + + expect(result.credentials).toBeNull(); + expect(result.error).toBeDefined(); + expect(mockEnsureRole).not.toHaveBeenCalled(); + expect(mockAssumeScopedRole).not.toHaveBeenCalled(); + }); + + test("rejects when agent record org does not match request org", async () => { + mockDynamoResolvers({ + agent: AGENT_ORG_A, // agent belongs to org-a + dataStores: { "ds-a": DS_ORG_A }, + }); + + const result = await handler({ + agentId: "agent-1", + org: "org-b", // request claims a different org + requiredPermissions: { dataStores: ["ds-a"] }, + } as never); + + expect(result.credentials).toBeNull(); + expect(result.error).toBeDefined(); + expect(mockEnsureRole).not.toHaveBeenCalled(); + }); + + test("rejects the whole request when a declared datastore id belongs to a different org", async () => { + const result = await handler({ + agentId: "agent-1", + org: "org-a", + requiredPermissions: { dataStores: ["ds-a", "ds-b"] }, + } as never); + + expect(result.credentials).toBeNull(); + expect(result.error).toBeDefined(); + // Cross-org id must not silently drop — the whole vend is refused. + expect(mockEnsureRole).not.toHaveBeenCalled(); + expect(mockAssumeScopedRole).not.toHaveBeenCalled(); + }); + + test("rejects the whole request when a declared integration id belongs to a different org", async () => { + const result = await handler({ + agentId: "agent-1", + org: "org-a", + requiredPermissions: { integrations: ["int-a", "int-b"] }, + } as never); + + expect(result.credentials).toBeNull(); + expect(result.error).toBeDefined(); + expect(mockEnsureRole).not.toHaveBeenCalled(); + }); + + test("rejects when a declared datastore id cannot be resolved", async () => { + mockDynamoResolvers({ + agent: AGENT_ORG_A, + dataStores: { "ds-a": DS_ORG_A }, // ds-missing absent + }); + + const result = await handler({ + agentId: "agent-1", + org: "org-a", + requiredPermissions: { dataStores: ["ds-a", "ds-missing"] }, + } as never); + + expect(result.credentials).toBeNull(); + expect(result.error).toBeDefined(); + expect(mockEnsureRole).not.toHaveBeenCalled(); + }); + + test("rejects when a declared integration id cannot be resolved", async () => { + mockDynamoResolvers({ + agent: AGENT_ORG_A, + integrations: {}, // int-missing absent + }); + + const result = await handler({ + agentId: "agent-1", + org: "org-a", + requiredPermissions: { integrations: ["int-missing"] }, + } as never); + + expect(result.credentials).toBeNull(); + expect(result.error).toBeDefined(); + expect(mockEnsureRole).not.toHaveBeenCalled(); + }); + + test("allows the vend when all declared ids and the agent belong to the request org", async () => { + const result = await handler({ + agentId: "agent-1", + org: "org-a", + requiredPermissions: { dataStores: ["ds-a"], integrations: ["int-a"] }, + } as never); + + expect(result.error).toBeUndefined(); + expect(result.credentials).not.toBeNull(); + expect(mockEnsureRole).toHaveBeenCalledTimes(1); + expect(mockAssumeScopedRole).toHaveBeenCalledTimes(1); + }); + + test("allows a vend with only model permissions (no datastore/integration lookups needed) for a same-org agent", async () => { + const result = await handler({ + agentId: "agent-1", + org: "org-a", + requiredPermissions: { models: ["anthropic.claude-sonnet-4-20250514"] }, + } as never); + + expect(result.error).toBeUndefined(); + expect(result.credentials).not.toBeNull(); + expect(mockEnsureRole).toHaveBeenCalledTimes(1); + }); +}); diff --git a/backend/src/lambda/__tests__/agent-credential-vender.test.ts b/backend/src/lambda/__tests__/agent-credential-vender.test.ts index 13c69309..07d7e199 100644 --- a/backend/src/lambda/__tests__/agent-credential-vender.test.ts +++ b/backend/src/lambda/__tests__/agent-credential-vender.test.ts @@ -1,72 +1,153 @@ /** * TDD Tests for agent-credential-vender Lambda + * + * Updated for wave 2b (branch fix/vender-org-scoping): the vender now + * requires `org` on every request and resolves the agent record plus every + * declared dataStore/integration id to that org before calling + * ensureRole/assumeScopedRole. These tests supply a same-org agent and + * same-org declared ids so the pre-existing (non-org) assertions below + * still exercise the exact same policy-computation path. */ const mockGetAccountContext = jest.fn(); const mockEnsureRole = jest.fn(); const mockAssumeScopedRole = jest.fn(); +const mockDynamoSend = jest.fn(); -jest.mock('../../utils/policy-manager', () => { +jest.mock("../../utils/policy-manager", () => { const MockPolicyManager = jest.fn().mockImplementation(() => ({ getAccountContext: mockGetAccountContext, ensureRole: mockEnsureRole, assumeScopedRole: mockAssumeScopedRole, })); // Static methods - (MockPolicyManager as unknown as { getRoleName: (id: string, scope: string) => string }).getRoleName = (id: string, scope: string) => { - const prefixes: Record = { datastore: 'citadel-ds-', integration: 'citadel-int-', agent: 'citadel-agent-' }; - return `${prefixes[scope] || 'citadel-ds-'}${id}`; + ( + MockPolicyManager as unknown as { + getRoleName: (id: string, scope: string) => string; + } + ).getRoleName = (id: string, scope: string) => { + const prefixes: Record = { + datastore: "citadel-ds-", + integration: "citadel-int-", + agent: "citadel-agent-", + }; + return `${prefixes[scope] || "citadel-ds-"}${id}`; }; - (MockPolicyManager as unknown as { - buildPolicyDocument: (policies: Array<{ actions: string[]; resources: string[] }>) => unknown; - }).buildPolicyDocument = (policies: Array<{ actions: string[]; resources: string[] }>) => ({ - Version: '2012-10-17', - Statement: policies.map(p => ({ Effect: 'Allow', Action: p.actions, Resource: p.resources })), + ( + MockPolicyManager as unknown as { + buildPolicyDocument: ( + policies: Array<{ actions: string[]; resources: string[] }>, + ) => unknown; + } + ).buildPolicyDocument = ( + policies: Array<{ actions: string[]; resources: string[] }>, + ) => ({ + Version: "2012-10-17", + Statement: policies.map((p) => ({ + Effect: "Allow", + Action: p.actions, + Resource: p.resources, + })), }); return { PolicyManager: MockPolicyManager }; }); -import { handler } from '../agent-credential-vender'; +jest.mock("@aws-sdk/lib-dynamodb", () => { + const actual = jest.requireActual("@aws-sdk/lib-dynamodb"); + return { + ...actual, + DynamoDBDocumentClient: { + from: jest.fn(() => ({ send: mockDynamoSend })), + }, + }; +}); + +process.env.AGENT_CONFIG_TABLE = "agent-config-table"; +process.env.DATASTORES_TABLE = "datastores-table"; +process.env.INTEGRATIONS_TABLE = "integrations-table"; + +import { handler } from "../agent-credential-vender"; + +const ORG = "org-legacy"; -describe('agent-credential-vender', () => { +describe("agent-credential-vender", () => { beforeEach(() => { mockGetAccountContext.mockClear(); mockEnsureRole.mockClear(); mockAssumeScopedRole.mockClear(); + mockDynamoSend.mockClear(); - mockGetAccountContext.mockResolvedValue({ accountId: '123456789012', region: 'us-west-2' }); + mockGetAccountContext.mockResolvedValue({ + accountId: "123456789012", + region: "us-west-2", + }); mockEnsureRole.mockResolvedValue(undefined); mockAssumeScopedRole.mockResolvedValue({ - accessKeyId: 'AKIA_SCOPED', - secretAccessKey: 'SECRET_SCOPED', - sessionToken: 'TOKEN_SCOPED', + accessKeyId: "AKIA_SCOPED", + secretAccessKey: "SECRET_SCOPED", + sessionToken: "TOKEN_SCOPED", + }); + + // Every agent/dataStore/integration id used below resolves to the same + // org as the request, so the org-scoping gate is a no-op pass-through + // and these tests still exercise only the policy-computation path. + mockDynamoSend.mockImplementation((command: unknown) => { + const input = + (command as { input?: Record }).input || {}; + const tableName = input.TableName as string | undefined; + if (tableName === "agent-config-table") { + return Promise.resolve({ + Item: { + agentId: (input.Key as { agentId: string }).agentId, + orgId: ORG, + }, + }); + } + if (tableName === "datastores-table") { + const key = input.Key as { dataStoreId: string }; + return Promise.resolve({ + Item: { dataStoreId: key.dataStoreId, orgId: ORG }, + }); + } + if (tableName === "integrations-table") { + const values = input.ExpressionAttributeValues as Record< + string, + string + >; + return Promise.resolve({ + Items: [{ integrationId: values[":id"], orgId: ORG }], + }); + } + return Promise.resolve({}); }); }); - test('returns scoped credentials for an agent with model permissions', async () => { + test("returns scoped credentials for an agent with model permissions", async () => { const result = await handler({ - agentId: 'agent-1', - requiredPermissions: { models: ['anthropic.claude-sonnet-4-20250514'] }, + agentId: "agent-1", + org: ORG, + requiredPermissions: { models: ["anthropic.claude-sonnet-4-20250514"] }, }); expect(result.credentials).toBeDefined(); - expect(result.credentials!.accessKeyId).toBe('AKIA_SCOPED'); + expect(result.credentials!.accessKeyId).toBe("AKIA_SCOPED"); expect(mockEnsureRole).toHaveBeenCalledTimes(1); - expect(mockEnsureRole.mock.calls[0][0]).toBe('agent-1'); - expect(mockEnsureRole.mock.calls[0][3]).toBe('agent'); + expect(mockEnsureRole.mock.calls[0][0]).toBe("agent-1"); + expect(mockEnsureRole.mock.calls[0][3]).toBe("agent"); const policies = mockEnsureRole.mock.calls[0][1]; - expect(policies[0].actions).toContain('bedrock:InvokeModel'); + expect(policies[0].actions).toContain("bedrock:InvokeModel"); }); - test('returns scoped credentials with datastore and integration access', async () => { + test("returns scoped credentials with datastore and integration access", async () => { const result = await handler({ - agentId: 'agent-2', + agentId: "agent-2", + org: ORG, requiredPermissions: { - models: ['anthropic.claude-sonnet-4-20250514'], - dataStores: ['ds-abc'], - integrations: ['int-xyz'], + models: ["anthropic.claude-sonnet-4-20250514"], + dataStores: ["ds-abc"], + integrations: ["int-xyz"], }, }); @@ -78,26 +159,31 @@ describe('agent-credential-vender', () => { expect(policies).toHaveLength(3); }); - test('returns null credentials when no permissions declared', async () => { - const result = await handler({ agentId: 'agent-3', requiredPermissions: {} }); + test("returns null credentials when no permissions declared", async () => { + const result = await handler({ + agentId: "agent-3", + org: ORG, + requiredPermissions: {}, + }); expect(result.credentials).toBeNull(); expect(mockEnsureRole).not.toHaveBeenCalled(); }); - test('returns null credentials when requiredPermissions is missing', async () => { - const result = await handler({ agentId: 'agent-4' }); + test("returns null credentials when requiredPermissions is missing", async () => { + const result = await handler({ agentId: "agent-4", org: ORG }); expect(result.credentials).toBeNull(); }); - test('returns error when PolicyManager fails', async () => { - mockEnsureRole.mockRejectedValueOnce(new Error('IAM failure')); + test("returns error when PolicyManager fails", async () => { + mockEnsureRole.mockRejectedValueOnce(new Error("IAM failure")); const result = await handler({ - agentId: 'agent-5', - requiredPermissions: { models: ['anthropic.claude-sonnet-4-20250514'] }, + agentId: "agent-5", + org: ORG, + requiredPermissions: { models: ["anthropic.claude-sonnet-4-20250514"] }, }); - expect(result.error).toContain('IAM failure'); + expect(result.error).toContain("IAM failure"); expect(result.credentials).toBeNull(); }); }); diff --git a/backend/src/lambda/agent-credential-vender.ts b/backend/src/lambda/agent-credential-vender.ts index 1b79f566..f6568b59 100644 --- a/backend/src/lambda/agent-credential-vender.ts +++ b/backend/src/lambda/agent-credential-vender.ts @@ -6,19 +6,43 @@ * a per-agent IAM role with only the permissions the agent declares. * * Input: - * { agentId: string, requiredPermissions: { models?, dataStores?, integrations? } } + * { agentId: string, org: string, requiredPermissions: { models?, dataStores?, integrations? } } * * Output: * { credentials: { accessKeyId, secretAccessKey, sessionToken } | null, error?: string } + * + * Wave 2b (branch fix/vender-org-scoping): `org` is REQUIRED and never + * defaulted — fail closed with a clear error when absent. Before computing + * any policy or creating/assuming any role, every declared dataStore/ + * integration id is resolved to its OWNING org and the agent's own record + * org is checked against the request org. Any mismatch or unresolvable id + * REJECTS the whole request (never silently dropped) — this is the only + * gate that prevents a cross-org agent from being granted AssumeRole on + * another org's citadel-ds-{id} / citadel-int-{id} role. */ -import { PolicyManager } from '../utils/policy-manager'; -import { computeAgentPolicies, AgentPermissions } from '../utils/policy-helpers'; +import { DynamoDBClient } from "@aws-sdk/client-dynamodb"; +import { + DynamoDBDocumentClient, + GetCommand, + QueryCommand, +} from "@aws-sdk/lib-dynamodb"; +import { PolicyManager } from "../utils/policy-manager"; +import { + computeAgentPolicies, + AgentPermissions, +} from "../utils/policy-helpers"; const policyManager = new PolicyManager(); +const dynamodb = DynamoDBDocumentClient.from(new DynamoDBClient({})); + +const AGENT_CONFIG_TABLE = process.env.AGENT_CONFIG_TABLE || ""; +const DATASTORES_TABLE = process.env.DATASTORES_TABLE || ""; +const INTEGRATIONS_TABLE = process.env.INTEGRATIONS_TABLE || ""; interface VendCredentialsEvent { agentId: string; + org?: string; requiredPermissions?: AgentPermissions; } @@ -31,8 +55,96 @@ interface VendCredentialsResult { error?: string; } -export async function handler(event: VendCredentialsEvent): Promise { - const { agentId, requiredPermissions } = event; +/** Thrown internally to short-circuit to the fail-closed error response. */ +class OrgScopeError extends Error {} + +async function getAgentOrgId(agentId: string): Promise { + if (!AGENT_CONFIG_TABLE) return null; + const result = await dynamodb.send( + new GetCommand({ TableName: AGENT_CONFIG_TABLE, Key: { agentId } }), + ); + const orgId = result.Item?.orgId; + return typeof orgId === "string" && orgId ? orgId : null; +} + +async function getDataStoreOrgId(dataStoreId: string): Promise { + if (!DATASTORES_TABLE) return null; + const result = await dynamodb.send( + new GetCommand({ TableName: DATASTORES_TABLE, Key: { dataStoreId } }), + ); + const orgId = result.Item?.orgId; + return typeof orgId === "string" && orgId ? orgId : null; +} + +async function getIntegrationOrgId( + integrationId: string, +): Promise { + if (!INTEGRATIONS_TABLE) return null; + const result = await dynamodb.send( + new QueryCommand({ + TableName: INTEGRATIONS_TABLE, + IndexName: "IntegrationIdIndex", + KeyConditionExpression: "integrationId = :id", + ExpressionAttributeValues: { ":id": integrationId }, + }), + ); + const orgId = result.Items?.[0]?.orgId; + return typeof orgId === "string" && orgId ? orgId : null; +} + +/** + * Resolves every declared dataStore/integration id to its owning orgId and + * verifies the agent's own record org, all against `requestOrg`. Throws + * `OrgScopeError` (fail closed, whole request rejected) on the first + * unresolvable id or org mismatch — never silently drops an id and never + * proceeds partially scoped. + */ +async function assertSameOrgOrThrow( + agentId: string, + requestOrg: string, + permissions: AgentPermissions, +): Promise { + const agentOrgId = await getAgentOrgId(agentId); + if (!agentOrgId || agentOrgId !== requestOrg) { + throw new OrgScopeError( + `Agent ${agentId} org (${agentOrgId ?? "unresolved"}) does not match request org ${requestOrg}`, + ); + } + + for (const dsId of permissions.dataStores ?? []) { + const dsOrgId = await getDataStoreOrgId(dsId); + if (!dsOrgId || dsOrgId !== requestOrg) { + throw new OrgScopeError( + `DataStore ${dsId} could not be resolved to request org ${requestOrg}`, + ); + } + } + + for (const intId of permissions.integrations ?? []) { + const intOrgId = await getIntegrationOrgId(intId); + if (!intOrgId || intOrgId !== requestOrg) { + throw new OrgScopeError( + `Integration ${intId} could not be resolved to request org ${requestOrg}`, + ); + } + } +} + +export async function handler( + event: VendCredentialsEvent, +): Promise { + const { agentId, org, requiredPermissions } = event; + + if (!org || typeof org !== "string") { + console.error("Failed to vend agent credentials: missing org on request", { + agentId, + }); + return { + credentials: null, + error: + 'Vend refused: request is missing a required "org" — never defaults to a fallback organization.', + }; + } if (!requiredPermissions) { return { credentials: null }; @@ -41,19 +153,32 @@ export async function handler(event: VendCredentialsEvent): Promise 0) || - (requiredPermissions.dataStores && requiredPermissions.dataStores.length > 0) || - (requiredPermissions.integrations && requiredPermissions.integrations.length > 0); + (requiredPermissions.dataStores && + requiredPermissions.dataStores.length > 0) || + (requiredPermissions.integrations && + requiredPermissions.integrations.length > 0); if (!hasPermissions) { return { credentials: null }; } try { + await assertSameOrgOrThrow(agentId, org, requiredPermissions); + const { accountId, region } = await policyManager.getAccountContext(); - const policies = computeAgentPolicies(agentId, requiredPermissions, accountId, region); + const policies = computeAgentPolicies( + agentId, + requiredPermissions, + accountId, + region, + ); - await policyManager.ensureRole(agentId, policies, accountId, 'agent'); - const credentials = await policyManager.assumeScopedRole(agentId, accountId, 'agent'); + await policyManager.ensureRole(agentId, policies, accountId, "agent"); + const credentials = await policyManager.assumeScopedRole( + agentId, + accountId, + "agent", + ); return { credentials: { @@ -63,7 +188,11 @@ export async function handler(event: VendCredentialsEvent): Promise { + let manager: PolicyManager; + + beforeEach(() => { + iamMock.reset(); + stsMock.reset(); + stsMock.on(GetCallerIdentityCommand).resolves({ + Account: "123456789012", + Arn: "arn:aws:sts::123456789012:assumed-role/LambdaRole/session", + }); + iamMock.on(CreateRoleCommand).resolves({}); + iamMock.on(PutRolePolicyCommand).resolves({}); + manager = new PolicyManager(); + }); + + test("trust principal is exactly the creating Lambda role when no extra principals supplied", async () => { + await manager.ensureRole( + "ds-1", + [{ actions: ["s3:GetObject"], resources: ["*"] }], + "123456789012", + "datastore", + ); + + const createCalls = iamMock.commandCalls(CreateRoleCommand); + const trustDoc = JSON.parse( + createCalls[0].args[0].input.AssumeRolePolicyDocument as string, + ); + expect(trustDoc.Statement[0].Principal.AWS).toBe( + "arn:aws:iam::123456789012:role/LambdaRole", + ); + }); + + test("trust principal includes an explicitly supplied crossArn", async () => { + await manager.ensureRole( + "int-1", + [{ actions: ["s3:GetObject"], resources: ["*"] }], + "123456789012", + "integration", + "arn:aws:iam::999999999999:role/CrossAccountRole", + ); + + const createCalls = iamMock.commandCalls(CreateRoleCommand); + const trustDoc = JSON.parse( + createCalls[0].args[0].input.AssumeRolePolicyDocument as string, + ); + expect(trustDoc.Statement[0].Principal.AWS).toEqual([ + "arn:aws:iam::123456789012:role/LambdaRole", + "arn:aws:iam::999999999999:role/CrossAccountRole", + ]); + }); + + test("trust principal includes an explicitly supplied non-agent additionalTrustedPrincipal", async () => { + await manager.ensureRole( + "ds-2", + [{ actions: ["s3:GetObject"], resources: ["*"] }], + "123456789012", + "datastore", + undefined, + ["arn:aws:iam::123456789012:role/HealthMonitorRole"], + ); + + const createCalls = iamMock.commandCalls(CreateRoleCommand); + const trustDoc = JSON.parse( + createCalls[0].args[0].input.AssumeRolePolicyDocument as string, + ); + expect(trustDoc.Statement[0].Principal.AWS).toEqual([ + "arn:aws:iam::123456789012:role/LambdaRole", + "arn:aws:iam::123456789012:role/HealthMonitorRole", + ]); + }); + + test("rejects a citadel-agent-* principal passed via additionalTrustedPrincipals for a datastore role", async () => { + await expect( + manager.ensureRole( + "ds-3", + [{ actions: ["s3:GetObject"], resources: ["*"] }], + "123456789012", + "datastore", + undefined, + ["arn:aws:iam::123456789012:role/citadel-agent-evil"], + ), + ).rejects.toThrow(/citadel-agent-/); + + // Never reaches CreateRole -- the rejection happens before any IAM call. + expect(iamMock.commandCalls(CreateRoleCommand)).toHaveLength(0); + }); + + test("rejects a citadel-agent-* principal passed via additionalTrustedPrincipals for an integration role", async () => { + await expect( + manager.ensureRole( + "int-2", + [{ actions: ["s3:GetObject"], resources: ["*"] }], + "123456789012", + "integration", + undefined, + ["arn:aws:iam::123456789012:role/citadel-agent-evil"], + ), + ).rejects.toThrow(/citadel-agent-/); + + expect(iamMock.commandCalls(CreateRoleCommand)).toHaveLength(0); + }); + + test("rejects when a citadel-agent-* principal is mixed in among otherwise-valid principals", async () => { + await expect( + manager.ensureRole( + "ds-4", + [{ actions: ["s3:GetObject"], resources: ["*"] }], + "123456789012", + "datastore", + undefined, + [ + "arn:aws:iam::123456789012:role/HealthMonitorRole", + "arn:aws:iam::123456789012:role/citadel-agent-sneaky", + ], + ), + ).rejects.toThrow(/citadel-agent-/); + + expect(iamMock.commandCalls(CreateRoleCommand)).toHaveLength(0); + }); + + test("does not reject citadel-agent-* trust when creating an agent-scoped role itself", async () => { + // The rejection is specific to ds/int scopes trusting an agent + // principal; an agent role legitimately being created is unaffected. + await expect( + manager.ensureRole( + "agent-9", + [{ actions: ["bedrock:InvokeModel"], resources: ["*"] }], + "123456789012", + "agent", + undefined, + ["arn:aws:iam::123456789012:role/citadel-agent-other"], + ), + ).resolves.not.toThrow(); + }); +}); diff --git a/backend/src/utils/policy-manager.ts b/backend/src/utils/policy-manager.ts index cfce7d25..2ea9eb15 100644 --- a/backend/src/utils/policy-manager.ts +++ b/backend/src/utils/policy-manager.ts @@ -115,6 +115,27 @@ export class PolicyManager { } } + // Wave 2b (fix/vender-org-scoping): a datastore/integration-scoped role + // must NEVER trust an agent role directly -- agents reach ds/int + // credentials exclusively through the credential vender's own + // sts:AssumeRole GRANT (computeAgentPolicies), never through the + // ds/int role's OWN trust policy. Reject even when a citadel-agent-* + // principal arrives via additionalTrustedPrincipals (e.g. a + // misconfigured caller), before any IAM call is made. Agent-scoped + // roles are exempt -- an agent role legitimately trusting another + // agent role (or itself) is not the invariant being enforced here. + if (scope !== "agent") { + const agentPrincipal = principals.find((p) => /citadel-agent-/.test(p)); + if (agentPrincipal) { + throw new PermissionError( + `Refusing to create ${scope} role ${roleName}: trust policy would ` + + `include a citadel-agent-* principal (${agentPrincipal}), which ` + + `is never permitted to be trusted directly by a datastore/` + + `integration role.`, + ); + } + } + const trustStatement: Record = { Effect: "Allow", Principal: { AWS: principals.length === 1 ? principals[0] : principals },