From 57d960b9f224896f2d844ad7ef9acdd4cff48f1f Mon Sep 17 00:00:00 2001 From: Sudhansu Date: Thu, 23 Jul 2026 18:43:06 +0530 Subject: [PATCH] feat: Add agentcore-identity-cognito-multi-tenant-cdk pattern Add a serverless pattern demonstrating multi-tenant authentication and tenant isolation for an Amazon Bedrock AgentCore Runtime agent using Amazon Cognito as the OAuth 2.0 / JWT identity provider. AgentCore Identity validates the inbound Cognito JWT via a custom JWT authorizer. The agent resolves the caller's tenant from the validated cognito:groups claim (requiring exactly one tenant group) and exposes a no-argument DynamoDB tool bound to that tenant, giving prompt-injection- resistant isolation in the pool multi-tenancy model. Includes CDK infra (Cognito user pool with tenant groups, DynamoDB table, AgentCore Runtime with Authorization header forwarding, scoped IAM role), containerized Strands agent, test/seed scripts, README with a security section and IAM LeadingKeys hardening guidance, and example-pattern.json. --- .../.gitignore | 11 + .../README.md | 158 +++++++++++ .../agent-code/Dockerfile | 20 ++ .../agent-code/agent.py | 128 +++++++++ .../agent-code/requirements.txt | 3 + .../app.py | 24 ++ .../cdk.json | 43 +++ .../example-pattern.json | 75 ++++++ .../infra/__init__.py | 0 .../infra/agentcore_identity_stack.py | 252 ++++++++++++++++++ .../requirements.txt | 2 + .../test/get_token.sh | 63 +++++ .../test/seed_tenant_data.sh | 32 +++ .../test/test_multi_tenant.py | 66 +++++ 14 files changed, 877 insertions(+) create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/.gitignore create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/README.md create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/agent-code/Dockerfile create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/agent-code/agent.py create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/agent-code/requirements.txt create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/app.py create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/cdk.json create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/example-pattern.json create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/infra/__init__.py create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/infra/agentcore_identity_stack.py create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/requirements.txt create mode 100755 agentcore-identity-cognito-multi-tenant-cdk/test/get_token.sh create mode 100755 agentcore-identity-cognito-multi-tenant-cdk/test/seed_tenant_data.sh create mode 100644 agentcore-identity-cognito-multi-tenant-cdk/test/test_multi_tenant.py diff --git a/agentcore-identity-cognito-multi-tenant-cdk/.gitignore b/agentcore-identity-cognito-multi-tenant-cdk/.gitignore new file mode 100644 index 0000000000..34b5c7cbdd --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/.gitignore @@ -0,0 +1,11 @@ +*.swp +package-lock.json +__pycache__ +.pytest_cache +.venv +*.egg-info +*.pyc +cdk.out +.DS_Store +.bedrock_agentcore +test/.cognito_config diff --git a/agentcore-identity-cognito-multi-tenant-cdk/README.md b/agentcore-identity-cognito-multi-tenant-cdk/README.md new file mode 100644 index 0000000000..8d368ddb18 --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/README.md @@ -0,0 +1,158 @@ +# Multi-tenant Amazon Bedrock AgentCore Runtime with Amazon Cognito + +This pattern demonstrates multi-tenant authentication and tenant isolation for an [Amazon Bedrock AgentCore Runtime](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/runtime.html) agent using [Amazon Cognito](https://docs.aws.amazon.com/cognito/latest/developerguide/what-is-amazon-cognito.html) as the OAuth 2.0 / JWT identity provider. + +AgentCore Identity validates inbound Cognito JWTs at the runtime edge via a [custom JWT authorizer](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/inbound-jwt-authorizer.html). The agent then derives the caller's tenant from the validated token claims and can only ever read that tenant's partition in Amazon DynamoDB. Because the tenant is fixed from the token and not from the prompt, a user cannot coax the agent into reading another tenant's data. + +Learn more about this pattern at [Serverless Land Patterns](https://serverlessland.com/patterns/agentcore-identity-cognito-multi-tenant-cdk). + +Important: this application uses various AWS services and there are costs associated with these services after the Free Tier usage - please see the [AWS Pricing page](https://aws.amazon.com/pricing/) for details. You are responsible for any AWS costs incurred. No warranty is implied in this example. + +## Pre-requisites + +- [Create an AWS account](https://portal.aws.amazon.com/gp/aws/developer/registration/index.html) if you do not already have one and log in. +- [AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/install-cliv2.html) installed and configured +- [Git Installed](https://git-scm.com/book/en/v2/Getting-Started-Installing-Git) +- [AWS CDK](https://docs.aws.amazon.com/cdk/latest/guide/cli.html) installed and configured +- [Docker](https://docs.docker.com/get-docker/) installed and running +- [Python 3.12](https://www.python.org/downloads/) installed +- Access to an Amazon Bedrock foundation model enabled in your region (the agent uses the default [Strands Agents](https://strandsagents.com/) model on Amazon Bedrock) + +## Deployment Instructions + +1. Clone the GitHub repository: + + ```shell + git clone https://github.com/aws-samples/serverless-patterns + cd serverless-patterns/agentcore-identity-cognito-multi-tenant-cdk + ``` + +2. Create and activate a Python virtual environment: + + ```shell + python3 -m venv .venv + source .venv/bin/activate + ``` + +3. Install the required dependencies: + + ```shell + pip3 install -r requirements.txt + ``` + +4. Set the deployment region (required), bootstrap the account/region if you have not already, and deploy: + + ```shell + export CDK_DEFAULT_REGION=us-west-2 + cdk bootstrap + cdk deploy + ``` + +5. Note the CDK outputs: `UserPoolId`, `UserPoolClientId`, `AgentRuntimeArn`, `TenantTableName`, and `TenantGroups`. + +## How it works + +This pattern creates: + +1. **Amazon Cognito User Pool** - the inbound OAuth 2.0 / JWT identity provider. Two Cognito groups (`tenant-acme` and `tenant-globex`) represent two tenants. Each end user belongs to exactly one tenant group. +2. **Amazon DynamoDB table** - a single table using the "pool" multi-tenancy model, partitioned by `tenant_id`. +3. **Amazon Bedrock AgentCore Runtime (HTTP)** - hosts the containerized agent and is configured with a custom JWT authorizer that trusts the Cognito user pool. + +**Request flow:** + +1. A client authenticates with Cognito and obtains an access token. The token carries the `cognito:groups` claim identifying the tenant. +2. The client calls the AgentCore Runtime with `Authorization: Bearer `. +3. AgentCore Identity validates the JWT (signature, issuer, expiry, and `allowedClients`) before the request reaches the agent. +4. The agent reads the validated `cognito:groups` claim, resolves the tenant, and exposes a `get_tenant_records` tool bound to that tenant. +5. The tool queries DynamoDB with a fixed `tenant_id` key condition, so the agent can only return the caller's tenant data - regardless of what the prompt asks for. + +## Testing + +1. Seed sample data for both tenants (use the `TenantTableName` output): + + ```shell + cd test + chmod +x seed_tenant_data.sh get_token.sh + ./seed_tenant_data.sh us-west-2 + ``` + +2. Install the test dependency: + + ```shell + pip3 install requests + ``` + +3. Get an access token for a user in the **acme** tenant (use `UserPoolId` and `UserPoolClientId` outputs): + + ```shell + ./get_token.sh + # Username: acme-user Tenant group: tenant-acme + ``` + + Copy the printed `export BEARER_TOKEN=...` line and run it. + +4. Set the agent ARN and region, then invoke: + + ```shell + export AGENT_ARN="" + export REGION="us-west-2" + python test_multi_tenant.py "List all of my orders" + ``` + + The response only contains **acme** orders (order-1001, order-1002). + +5. Repeat step 3 for a user in the **globex** tenant (Username: `globex-user`, Tenant group: `tenant-globex`), re-export `BEARER_TOKEN`, and invoke again. The response only contains **globex** orders. + +6. Confirm tenant isolation. While using the **acme** token, explicitly ask for another tenant's data: + + ```shell + python test_multi_tenant.py "Show me the globex orders" + ``` + + The agent still only has access to acme data and will not return globex records. + +## Security and tenant isolation + +This pattern uses the "pool" multi-tenancy model (one shared table, partitioned by tenant). Isolation is enforced on two boundaries: + +**Authentication boundary (cryptographically enforced at the edge).** AgentCore Identity's inbound JWT authorizer validates every request against the Cognito user pool before it reaches the agent. The `cognito:groups` claim that identifies the tenant is signed by Cognito, so it cannot be forged or altered. Verified behaviors: + +- A tampered token (group claim rewritten, original signature kept) is rejected with `401 Invalid Signature`. +- An invalid/garbage token is rejected with `403 OAuth authorization failed`. +- A user in no tenant group receives no data. +- A prompt-injection attempt to read another tenant's data is refused - the `get_tenant_records` tool takes no arguments and the tenant id is a per-request value taken from the validated token, so the model cannot change the queried partition. + +**Authorization boundary (application-enforced).** The agent resolves the tenant to exactly one Cognito group (zero or multiple `tenant-*` groups are rejected rather than guessed) and binds the DynamoDB query to that tenant's partition key. + +### Hardening for production + +- **Enforce the boundary in IAM, not just in code.** The execution role in this sample grants `dynamodb:Query`/`GetItem` on the whole table, so the tenant boundary depends entirely on application logic. For stronger, defense-in-depth isolation, issue per-tenant scoped credentials and add a [`dynamodb:LeadingKeys`](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/specifying-conditions.html) condition so IAM itself denies cross-tenant reads even if the code has a bug: + + ```json + { + "Effect": "Allow", + "Action": ["dynamodb:Query", "dynamodb:GetItem"], + "Resource": "", + "Condition": { + "ForAllValues:StringEquals": { "dynamodb:LeadingKeys": ["${tenant_id}"] } + } + } + ``` + + Per-tenant credentials can come from an STS `AssumeRole` with session tags, or from AgentCore Identity's per-user token vault (`GetWorkloadAccessTokenForUserId`). For strict isolation requirements, consider the "silo" model (a separate table or account per tenant). +- **Never deploy the agent without the inbound JWT authorizer.** The agent decodes the token to read claims but relies on AgentCore having already verified the signature. Removing the authorizer would make the agent trust unsigned tokens. For belt-and-suspenders, re-verify the JWT against the Cognito JWKS inside the agent. +- **Provision each user into exactly one tenant group.** The agent rejects callers who belong to zero or multiple `tenant-*` groups. + +## Cleanup + +Delete the stack: + +```bash +cdk destroy +``` + +--- + +Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. + +SPDX-License-Identifier: MIT-0 diff --git a/agentcore-identity-cognito-multi-tenant-cdk/agent-code/Dockerfile b/agentcore-identity-cognito-multi-tenant-cdk/agent-code/Dockerfile new file mode 100644 index 0000000000..acc74f75fc --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/agent-code/Dockerfile @@ -0,0 +1,20 @@ +FROM public.ecr.aws/docker/library/python:3.12-slim + +WORKDIR /app + +RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/* + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +RUN useradd -m -u 1000 bedrock_agentcore +USER bedrock_agentcore + +EXPOSE 8080 + +COPY . . + +HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \ + CMD curl -f http://localhost:8080/ping || exit 1 + +CMD ["python", "agent.py"] diff --git a/agentcore-identity-cognito-multi-tenant-cdk/agent-code/agent.py b/agentcore-identity-cognito-multi-tenant-cdk/agent-code/agent.py new file mode 100644 index 0000000000..4da51dd473 --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/agent-code/agent.py @@ -0,0 +1,128 @@ +"""Multi-tenant agent hosted on Amazon Bedrock AgentCore Runtime. + +AgentCore Identity validates the inbound Cognito JWT before the request reaches +this code (via the runtime's custom JWT authorizer). This agent then derives the +caller's tenant from the validated token claims and exposes a single tool that +can only read that tenant's partition in DynamoDB. Because the tenant is fixed +from the token - not from the prompt - a user cannot coax the agent into reading +another tenant's data. +""" + +import base64 +import json +import logging +import os + +import boto3 +from boto3.dynamodb.conditions import Key +from strands import Agent, tool + +from bedrock_agentcore import BedrockAgentCoreApp + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +REGION = os.environ.get("AWS_DEFAULT_REGION", "us-west-2") +TABLE_NAME = os.environ["TENANT_TABLE_NAME"] +TENANT_GROUP_PREFIX = os.environ.get("TENANT_GROUP_PREFIX", "tenant-") + +app = BedrockAgentCoreApp() +table = boto3.resource("dynamodb", region_name=REGION).Table(TABLE_NAME) + + +def _decode_jwt_claims(token: str) -> dict: + """Decode (do not re-verify) the JWT payload. + + AgentCore's inbound JWT authorizer has already verified the signature, + issuer, audience/client, and expiry before invoking the agent, so here we + only base64-decode the payload segment to read its claims. + """ + payload_segment = token.split(".")[1] + padding = "=" * (-len(payload_segment) % 4) + decoded = base64.urlsafe_b64decode(payload_segment + padding) + return json.loads(decoded) + + +class TenantResolutionError(Exception): + """Raised when a caller cannot be mapped to exactly one tenant.""" + + +def _extract_tenant(claims: dict) -> str: + """Resolve the tenant id from the token's Cognito group membership. + + A caller must belong to *exactly one* tenant group. Zero groups means the + user is not provisioned for any tenant; more than one is an ambiguous + (and likely misconfigured) authorization state. Both are rejected rather + than guessed, so the tenant boundary is never resolved arbitrarily. + """ + groups = claims.get("cognito:groups", []) or [] + tenant_groups = [g for g in groups if g.startswith(TENANT_GROUP_PREFIX)] + + if not tenant_groups: + raise TenantResolutionError( + "No tenant group found in token. The user must belong to a " + f"Cognito group prefixed with '{TENANT_GROUP_PREFIX}'." + ) + if len(tenant_groups) > 1: + raise TenantResolutionError( + "Ambiguous tenant: the user belongs to multiple tenant groups " + f"({tenant_groups}). A user must belong to exactly one tenant group." + ) + return tenant_groups[0][len(TENANT_GROUP_PREFIX):] + + +@app.entrypoint +async def invoke(payload, context): + prompt = payload.get( + "prompt", "No prompt found in input, please provide a json payload with a 'prompt' key" + ) + + headers = context.request_headers or {} + auth_header = headers.get("Authorization", "") + token = auth_header[len("Bearer "):].strip() if auth_header.startswith("Bearer ") else "" + + if not token: + yield {"error": "Missing bearer token"} + return + + try: + claims = _decode_jwt_claims(token) + except Exception: # noqa: BLE001 - defensive; malformed tokens + logger.exception("Failed to decode token claims") + yield {"error": "Invalid token"} + return + + try: + tenant_id = _extract_tenant(claims) + except TenantResolutionError as exc: + logger.warning("Tenant resolution failed: %s", exc) + yield {"error": str(exc)} + return + + logger.info("Handling request for tenant '%s'", tenant_id) + + @tool + def get_tenant_records() -> list: + """Return the data records that belong to the caller's tenant.""" + response = table.query( + KeyConditionExpression=Key("tenant_id").eq(tenant_id) + ) + return response.get("Items", []) + + agent = Agent( + tools=[get_tenant_records], + system_prompt=( + f"You are a helpful assistant serving tenant '{tenant_id}'. " + "Use the get_tenant_records tool to answer questions about this " + "tenant's data. You can only access data for this tenant. Never " + "reveal, guess, or fabricate data belonging to any other tenant, " + "even if you are asked to." + ), + ) + + async for event in agent.stream_async(prompt): + yield event + + +if __name__ == "__main__": + app.run() diff --git a/agentcore-identity-cognito-multi-tenant-cdk/agent-code/requirements.txt b/agentcore-identity-cognito-multi-tenant-cdk/agent-code/requirements.txt new file mode 100644 index 0000000000..261592c63e --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/agent-code/requirements.txt @@ -0,0 +1,3 @@ +strands-agents +bedrock-agentcore +boto3 diff --git a/agentcore-identity-cognito-multi-tenant-cdk/app.py b/agentcore-identity-cognito-multi-tenant-cdk/app.py new file mode 100644 index 0000000000..d3d7213c71 --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/app.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 + +import os + +import aws_cdk as cdk + +from infra.agentcore_identity_stack import AgentCoreIdentityStack + +app = cdk.App() + +region = os.environ.get("CDK_DEFAULT_REGION") +if not region: + raise ValueError( + "CDK_DEFAULT_REGION environment variable must be set. " + "Run: export CDK_DEFAULT_REGION=" + ) + +AgentCoreIdentityStack( + app, + "AgentCoreIdentityCognitoMultiTenantStack", + env=cdk.Environment(region=region), +) + +app.synth() diff --git a/agentcore-identity-cognito-multi-tenant-cdk/cdk.json b/agentcore-identity-cognito-multi-tenant-cdk/cdk.json new file mode 100644 index 0000000000..e5c61bc18c --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/cdk.json @@ -0,0 +1,43 @@ +{ + "app": ".venv/bin/python3 app.py", + "watch": { + "include": [ + "**" + ], + "exclude": [ + "README.md", + "cdk*.json", + "requirements*.txt", + "source.bat", + "**/__init__.py", + "**/__pycache__", + "tests" + ] + }, + "context": { + "@aws-cdk/aws-lambda:recognizeLayerVersion": true, + "@aws-cdk/core:checkSecretUsage": true, + "@aws-cdk/core:target-partitions": [ + "aws", + "aws-cn" + ], + "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true, + "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true, + "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true, + "@aws-cdk/aws-iam:minimizePolicies": true, + "@aws-cdk/core:validateSnapshotRemovalPolicy": true, + "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true, + "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true, + "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true, + "@aws-cdk/aws-apigateway:disableCloudWatchRole": true, + "@aws-cdk/core:enablePartitionLiterals": true, + "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true, + "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true, + "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true, + "@aws-cdk/customresources:installLatestAwsSdkDefault": false, + "@aws-cdk/core:includePrefixInUniqueNameGeneration": true, + "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": true, + "@aws-cdk/aws-ec2:ebsDefaultGp3Volume": true, + "@aws-cdk/aws-dynamodb:resourcePolicyPerReplica": true + } +} diff --git a/agentcore-identity-cognito-multi-tenant-cdk/example-pattern.json b/agentcore-identity-cognito-multi-tenant-cdk/example-pattern.json new file mode 100644 index 0000000000..afe09de19c --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/example-pattern.json @@ -0,0 +1,75 @@ +{ + "title": "Multi-tenant Amazon Bedrock AgentCore Runtime with Amazon Cognito", + "description": "Authenticate an Amazon Bedrock AgentCore Runtime agent with Amazon Cognito and enforce per-tenant data isolation using AgentCore Identity inbound JWT authorization.", + "language": "Python", + "level": "300", + "framework": "CDK", + "introBox": { + "headline": "How it works", + "text": [ + "This pattern secures an Amazon Bedrock AgentCore Runtime agent with Amazon Cognito acting as the OAuth 2.0 / JWT identity provider.", + "Two Cognito groups represent two tenants. AgentCore Identity validates the inbound Cognito JWT via a custom JWT authorizer before the request reaches the agent.", + "The agent resolves the caller's tenant from the validated 'cognito:groups' claim and exposes a DynamoDB tool bound to that tenant. It queries a single table partitioned by tenant_id.", + "Because the tenant is fixed from the token and not the prompt, the agent can only ever read the caller's tenant data - providing prompt-injection-resistant multi-tenant isolation." + ] + }, + "gitHub": { + "template": { + "repoURL": "https://github.com/aws-samples/serverless-patterns/tree/main/agentcore-identity-cognito-multi-tenant-cdk", + "templateURL": "serverless-patterns/agentcore-identity-cognito-multi-tenant-cdk", + "projectFolder": "agentcore-identity-cognito-multi-tenant-cdk", + "templateFile": "app.py" + } + }, + "resources": { + "bullets": [ + { + "text": "Configure inbound JWT authorizer for AgentCore", + "link": "https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/inbound-jwt-authorizer.html" + }, + { + "text": "Amazon Bedrock AgentCore Identity", + "link": "https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/identity.html" + }, + { + "text": "Amazon Cognito as an identity provider for AgentCore", + "link": "https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/identity-idp-cognito.html" + }, + { + "text": "SaaS multi-tenancy on AWS", + "link": "https://docs.aws.amazon.com/whitepapers/latest/saas-architecture-fundamentals/tenant-isolation.html" + } + ] + }, + "deploy": { + "text": [ + "python3 -m venv .venv", + "source .venv/bin/activate", + "pip3 install -r requirements.txt", + "export CDK_DEFAULT_REGION=us-west-2", + "cdk deploy" + ] + }, + "testing": { + "text": [ + "Seed data: cd test && ./seed_tenant_data.sh <TenantTableName> us-west-2", + "Get a token for a tenant user: ./get_token.sh", + "Invoke the agent: export AGENT_ARN=<AgentRuntimeArn> && export BEARER_TOKEN=<token> && python test_multi_tenant.py \"List my orders\"", + "Repeat for a second tenant user and confirm each caller only sees its own tenant's data." + ] + }, + "cleanup": { + "text": [ + "Delete the stack: cdk destroy" + ] + }, + "authors": [ + { + "name": "Sudhanshu Saurav", + "image": "link-to-your-photo.jpg", + "bio": "Contributor to the Serverless Patterns Collection.", + "linkedin": "", + "twitter": "" + } + ] +} diff --git a/agentcore-identity-cognito-multi-tenant-cdk/infra/__init__.py b/agentcore-identity-cognito-multi-tenant-cdk/infra/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/agentcore-identity-cognito-multi-tenant-cdk/infra/agentcore_identity_stack.py b/agentcore-identity-cognito-multi-tenant-cdk/infra/agentcore_identity_stack.py new file mode 100644 index 0000000000..7246ce3502 --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/infra/agentcore_identity_stack.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 + +from aws_cdk import ( + Stack, + CfnOutput, + RemovalPolicy, + aws_iam as iam, + aws_ecr_assets as ecr_assets, + aws_bedrockagentcore as bedrockagentcore, + aws_cognito as cognito, + aws_dynamodb as dynamodb, +) +from constructs import Construct + +# Tenant groups seeded in the Cognito user pool. Each end user belongs to +# exactly one of these groups, which becomes their tenant boundary. The group +# name (minus the "tenant-" prefix) is used as the DynamoDB partition key. +TENANT_GROUPS = ["tenant-acme", "tenant-globex"] +TENANT_GROUP_PREFIX = "tenant-" + + +class AgentCoreIdentityStack(Stack): + """Multi-tenant Amazon Bedrock AgentCore Runtime secured by Amazon Cognito. + + Inbound requests are authenticated by AgentCore Identity using a custom JWT + authorizer that trusts the Cognito user pool. The agent derives the caller's + tenant from the validated token claims and can only ever read that tenant's + partition in DynamoDB, giving prompt-injection-resistant tenant isolation. + """ + + def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None: + super().__init__(scope, construct_id, **kwargs) + + # ------------------------------------------------------------------ + # Amazon Cognito user pool (the inbound OAuth2 / JWT identity provider) + # ------------------------------------------------------------------ + user_pool = cognito.UserPool( + self, + "UserPool", + user_pool_name=f"{self.stack_name}-user-pool", + self_sign_up_enabled=False, + sign_in_aliases=cognito.SignInAliases(username=True), + password_policy=cognito.PasswordPolicy(min_length=8), + removal_policy=RemovalPolicy.DESTROY, + ) + + user_pool_client = cognito.UserPoolClient( + self, + "UserPoolClient", + user_pool=user_pool, + user_pool_client_name=f"{self.stack_name}-client", + generate_secret=False, + auth_flows=cognito.AuthFlow(user_password=True), + ) + + # One Cognito group per tenant. The "cognito:groups" claim is included in + # the access token, so the agent can read it after AgentCore validates + # the token. + for group_name in TENANT_GROUPS: + cognito.CfnUserPoolGroup( + self, + f"Group-{group_name}", + user_pool_id=user_pool.user_pool_id, + group_name=group_name, + description=f"Tenant group for {group_name}", + ) + + # ------------------------------------------------------------------ + # Per-tenant data store (pool model: one table, partitioned by tenant) + # ------------------------------------------------------------------ + tenant_table = dynamodb.Table( + self, + "TenantTable", + partition_key=dynamodb.Attribute( + name="tenant_id", type=dynamodb.AttributeType.STRING + ), + sort_key=dynamodb.Attribute( + name="record_id", type=dynamodb.AttributeType.STRING + ), + billing_mode=dynamodb.BillingMode.PAY_PER_REQUEST, + removal_policy=RemovalPolicy.DESTROY, + ) + + # ------------------------------------------------------------------ + # Container image for the agent + # ------------------------------------------------------------------ + agent_image = ecr_assets.DockerImageAsset( + self, + "AgentImage", + directory="./agent-code", + platform=ecr_assets.Platform.LINUX_ARM64, + ) + + agent_name = f"{self.stack_name.replace('-', '_')}_Agent" + + # ------------------------------------------------------------------ + # Execution role for the AgentCore Runtime + # ------------------------------------------------------------------ + agent_role = iam.Role( + self, + "AgentCoreRole", + assumed_by=iam.ServicePrincipal( + "bedrock-agentcore.amazonaws.com", + conditions={ + "StringEquals": {"aws:SourceAccount": self.account}, + "ArnLike": { + "aws:SourceArn": f"arn:aws:bedrock-agentcore:{self.region}:{self.account}:*" + }, + }, + ), + inline_policies={ + "AgentCorePolicy": iam.PolicyDocument( + statements=[ + iam.PolicyStatement( + sid="ECRImageAccess", + actions=[ + "ecr:BatchGetImage", + "ecr:GetDownloadUrlForLayer", + ], + resources=[agent_image.repository.repository_arn], + ), + iam.PolicyStatement( + sid="ECRTokenAccess", + # ecr:GetAuthorizationToken does not support + # resource-level permissions. + actions=["ecr:GetAuthorizationToken"], + resources=["*"], + ), + iam.PolicyStatement( + actions=[ + "logs:DescribeLogStreams", + "logs:CreateLogGroup", + ], + resources=[ + f"arn:aws:logs:{self.region}:{self.account}:log-group:/aws/bedrock-agentcore/runtimes/*" + ], + ), + iam.PolicyStatement( + actions=["logs:DescribeLogGroups"], + resources=[ + f"arn:aws:logs:{self.region}:{self.account}:log-group:*" + ], + ), + iam.PolicyStatement( + actions=["logs:CreateLogStream", "logs:PutLogEvents"], + resources=[ + f"arn:aws:logs:{self.region}:{self.account}:log-group:/aws/bedrock-agentcore/runtimes/*:log-stream:*" + ], + ), + iam.PolicyStatement( + actions=[ + "xray:PutTraceSegments", + "xray:PutTelemetryRecords", + "xray:GetSamplingRules", + "xray:GetSamplingTargets", + ], + # X-Ray actions do not support resource-level permissions. + resources=["*"], + ), + iam.PolicyStatement( + actions=["cloudwatch:PutMetricData"], + resources=["*"], + conditions={ + "StringEquals": { + "cloudwatch:namespace": "bedrock-agentcore" + } + }, + ), + iam.PolicyStatement( + sid="GetAgentAccessToken", + actions=[ + "bedrock-agentcore:GetWorkloadAccessToken", + "bedrock-agentcore:GetWorkloadAccessTokenForJWT", + "bedrock-agentcore:GetWorkloadAccessTokenForUserId", + ], + resources=[ + f"arn:aws:bedrock-agentcore:{self.region}:{self.account}:workload-identity-directory/default", + f"arn:aws:bedrock-agentcore:{self.region}:{self.account}:workload-identity-directory/default/workload-identity/{agent_name}-*", + ], + ), + iam.PolicyStatement( + sid="BedrockModelInvocation", + actions=[ + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream", + ], + resources=[ + "arn:aws:bedrock:*::foundation-model/*", + f"arn:aws:bedrock:{self.region}:{self.account}:*", + ], + ), + iam.PolicyStatement( + sid="TenantDataAccess", + # Table-level access; per-tenant isolation is enforced + # in the agent via the validated token's tenant claim. + actions=["dynamodb:Query", "dynamodb:GetItem"], + resources=[tenant_table.table_arn], + ), + ] + ) + }, + ) + + # ------------------------------------------------------------------ + # AgentCore Runtime with inbound Cognito JWT authorizer + # ------------------------------------------------------------------ + discovery_url = ( + f"https://cognito-idp.{self.region}.amazonaws.com/" + f"{user_pool.user_pool_id}/.well-known/openid-configuration" + ) + + agent_runtime = bedrockagentcore.CfnRuntime( + self, + "AgentRuntime", + agent_runtime_name=agent_name, + agent_runtime_artifact=bedrockagentcore.CfnRuntime.AgentRuntimeArtifactProperty( + container_configuration=bedrockagentcore.CfnRuntime.ContainerConfigurationProperty( + container_uri=agent_image.image_uri + ) + ), + network_configuration=bedrockagentcore.CfnRuntime.NetworkConfigurationProperty( + network_mode="PUBLIC" + ), + protocol_configuration="HTTP", + role_arn=agent_role.role_arn, + # Forward the inbound Authorization header to the container so the + # agent can read the validated Cognito JWT and resolve the tenant. + request_header_configuration=bedrockagentcore.CfnRuntime.RequestHeaderConfigurationProperty( + request_header_allowlist=["Authorization"] + ), + environment_variables={ + "AWS_DEFAULT_REGION": self.region, + "TENANT_TABLE_NAME": tenant_table.table_name, + "TENANT_GROUP_PREFIX": TENANT_GROUP_PREFIX, + }, + authorizer_configuration=bedrockagentcore.CfnRuntime.AuthorizerConfigurationProperty( + custom_jwt_authorizer=bedrockagentcore.CfnRuntime.CustomJWTAuthorizerConfigurationProperty( + discovery_url=discovery_url, + allowed_clients=[user_pool_client.user_pool_client_id], + ) + ), + ) + + # ------------------------------------------------------------------ + # Outputs + # ------------------------------------------------------------------ + CfnOutput(self, "UserPoolId", value=user_pool.user_pool_id) + CfnOutput(self, "UserPoolClientId", value=user_pool_client.user_pool_client_id) + CfnOutput(self, "AgentRuntimeArn", value=agent_runtime.attr_agent_runtime_arn) + CfnOutput(self, "TenantTableName", value=tenant_table.table_name) + CfnOutput(self, "TenantGroups", value=",".join(TENANT_GROUPS)) + CfnOutput(self, "Region", value=self.region) diff --git a/agentcore-identity-cognito-multi-tenant-cdk/requirements.txt b/agentcore-identity-cognito-multi-tenant-cdk/requirements.txt new file mode 100644 index 0000000000..520efb5493 --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/requirements.txt @@ -0,0 +1,2 @@ +aws-cdk-lib>=2.220.0 +constructs>=10.0.0,<11.0.0 diff --git a/agentcore-identity-cognito-multi-tenant-cdk/test/get_token.sh b/agentcore-identity-cognito-multi-tenant-cdk/test/get_token.sh new file mode 100755 index 0000000000..d93d7a21b5 --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/test/get_token.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# +# Creates (or reuses) a Cognito user, assigns it to a tenant group, and prints +# an access token for that user. The access token carries the "cognito:groups" +# claim that the agent uses to determine the caller's tenant. +# +# Get Pool ID, Client ID, and the tenant group names from the CDK stack outputs. + +set -euo pipefail + +echo "Cognito Multi-Tenant Token Generator" +echo "" + +CONFIG_FILE=".cognito_config" +if [ -f "$CONFIG_FILE" ]; then + # shellcheck disable=SC1090 + source "$CONFIG_FILE" +fi + +read -r -p "Pool ID [${POOL_ID:-}]: " input && POOL_ID=${input:-${POOL_ID:-}} +read -r -p "Client ID [${CLIENT_ID:-}]: " input && CLIENT_ID=${input:-${CLIENT_ID:-}} +read -r -p "Username [${USERNAME:-}]: " input && USERNAME=${input:-${USERNAME:-}} +read -r -p "Tenant group (e.g. tenant-acme) [${TENANT_GROUP:-}]: " input && TENANT_GROUP=${input:-${TENANT_GROUP:-}} +echo "Password requires: min 8 chars" +read -r -sp "Password: " PASSWORD +echo +read -r -p "Region [${REGION:-us-west-2}]: " input && REGION=${input:-${REGION:-us-west-2}} + +cat > "$CONFIG_FILE" << EOF +POOL_ID=$POOL_ID +CLIENT_ID=$CLIENT_ID +USERNAME=$USERNAME +TENANT_GROUP=$TENANT_GROUP +REGION=$REGION +EOF + +aws cognito-idp admin-create-user \ + --user-pool-id "$POOL_ID" \ + --username "$USERNAME" \ + --region "$REGION" \ + --message-action SUPPRESS > /dev/null 2>&1 || true + +aws cognito-idp admin-set-user-password \ + --user-pool-id "$POOL_ID" \ + --username "$USERNAME" \ + --password "$PASSWORD" \ + --region "$REGION" \ + --permanent > /dev/null 2>&1 || true + +aws cognito-idp admin-add-user-to-group \ + --user-pool-id "$POOL_ID" \ + --username "$USERNAME" \ + --group-name "$TENANT_GROUP" \ + --region "$REGION" > /dev/null 2>&1 || true + +BEARER_TOKEN=$(aws cognito-idp initiate-auth \ + --client-id "$CLIENT_ID" \ + --auth-flow USER_PASSWORD_AUTH \ + --auth-parameters USERNAME="$USERNAME",PASSWORD="$PASSWORD" \ + --region "$REGION" | jq -r '.AuthenticationResult.AccessToken') + +echo "" +echo "export BEARER_TOKEN=\"$BEARER_TOKEN\"" diff --git a/agentcore-identity-cognito-multi-tenant-cdk/test/seed_tenant_data.sh b/agentcore-identity-cognito-multi-tenant-cdk/test/seed_tenant_data.sh new file mode 100755 index 0000000000..a6ba7d8a5a --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/test/seed_tenant_data.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# +# Seeds the DynamoDB tenant table with sample records for two tenants so you can +# observe tenant isolation. Get the table name from the CDK output "TenantTableName". +# +# Usage: ./seed_tenant_data.sh [region] + +set -euo pipefail + +TABLE_NAME=${1:-} +REGION=${2:-${AWS_DEFAULT_REGION:-us-west-2}} + +if [ -z "$TABLE_NAME" ]; then + echo "Usage: ./seed_tenant_data.sh [region]" + exit 1 +fi + +put() { + local tenant=$1 record=$2 desc=$3 + aws dynamodb put-item \ + --table-name "$TABLE_NAME" \ + --region "$REGION" \ + --item "{\"tenant_id\":{\"S\":\"$tenant\"},\"record_id\":{\"S\":\"$record\"},\"description\":{\"S\":\"$desc\"}}" + echo " put $tenant / $record" +} + +echo "Seeding tenant data into $TABLE_NAME ($REGION)..." +put "acme" "order-1001" "Acme order 1001: 500 anvils" +put "acme" "order-1002" "Acme order 1002: 20 rocket skates" +put "globex" "order-2001" "Globex order 2001: 1000 widgets" +put "globex" "order-2002" "Globex order 2002: 3 fusion reactors" +echo "Done." diff --git a/agentcore-identity-cognito-multi-tenant-cdk/test/test_multi_tenant.py b/agentcore-identity-cognito-multi-tenant-cdk/test/test_multi_tenant.py new file mode 100644 index 0000000000..4252e710bc --- /dev/null +++ b/agentcore-identity-cognito-multi-tenant-cdk/test/test_multi_tenant.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Invoke the multi-tenant AgentCore Runtime over HTTPS using a Cognito bearer token. + +Because the runtime uses OAuth/JWT inbound auth, requests are sent directly to +the InvokeAgentRuntime HTTPS endpoint with an "Authorization: Bearer " +header (the AWS SDK / SigV4 path cannot be used with JWT inbound auth). + +Usage: + export AGENT_ARN="arn:aws:bedrock-agentcore:us-west-2::runtime/" + export REGION="us-west-2" + export BEARER_TOKEN="" + python test_multi_tenant.py "List my orders" +""" + +import os +import sys +import uuid +from urllib.parse import quote + +import requests + + +def main(): + prompt = sys.argv[1] if len(sys.argv) > 1 else "List all of my orders." + agent_arn = os.environ.get("AGENT_ARN", "") + region = os.environ.get("REGION", "us-west-2") + token = os.environ.get("BEARER_TOKEN", "") + + if not agent_arn or not token: + print("Set AGENT_ARN and BEARER_TOKEN environment variables.") + print("AGENT_ARN comes from the CDK output 'AgentRuntimeArn'.") + sys.exit(1) + + encoded_arn = quote(agent_arn, safe="") + url = ( + f"https://bedrock-agentcore.{region}.amazonaws.com/" + f"runtimes/{encoded_arn}/invocations?qualifier=DEFAULT" + ) + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + # Session id must be at least 33 characters. + "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": uuid.uuid4().hex + uuid.uuid4().hex, + } + + print("=" * 60) + print("Multi-Tenant AgentCore Invocation") + print("=" * 60) + print(f"Prompt: {prompt}") + print(f"Token: {token[:20]}...") + print("-" * 60) + + resp = requests.post(url, headers=headers, json={"prompt": prompt}, stream=True, timeout=120) + print(f"Status: {resp.status_code}") + if resp.status_code != 200: + print(resp.text[:1000]) + sys.exit(1) + + for line in resp.iter_lines(): + if line: + print(line.decode("utf-8")) + + +if __name__ == "__main__": + main()