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 },