Skip to content

Commit da0ac91

Browse files
ashakirinashakirin
committed
Added policy to deny specific tools
1 parent 127ef06 commit da0ac91

10 files changed

Lines changed: 409 additions & 37 deletions

File tree

‎apps/java-spring-ai-agents/aiagent/src/main/java/com/example/agent/ChatService.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,14 @@
2020
import org.springframework.beans.factory.annotation.Qualifier;
2121
import org.springframework.beans.factory.annotation.Value;
2222
import org.springframework.core.io.ByteArrayResource;
23+
import org.springframework.http.HttpHeaders;
2324
import org.springframework.http.MediaType;
2425
import org.springframework.http.MediaTypeFactory;
2526
import org.springframework.stereotype.Service;
2627
import org.springframework.util.MimeType;
2728
import org.springframework.util.MimeTypeUtils;
29+
import org.springframework.web.context.request.RequestAttributes;
30+
import org.springframework.web.context.request.RequestContextHolder;
2831
import reactor.core.publisher.Flux;
2932
import tools.jackson.databind.json.JsonMapper;
3033

@@ -40,7 +43,6 @@ public boolean hasFile() {
4043

4144
@Service
4245
public class ChatService {
43-
4446
private static final Logger logger = LoggerFactory.getLogger(ChatService.class);
4547

4648
private final ChatClient chatClient;
@@ -127,6 +129,10 @@ public ChatService(AgentCoreMemory agentCoreMemory,
127129

128130
@AgentCoreInvocation
129131
public Flux<String> chat(ChatRequest request, AgentCoreContext context) {
132+
String authorization = context.getHeader(HttpHeaders.AUTHORIZATION);
133+
RequestContextHolder.currentRequestAttributes()
134+
.setAttribute(HttpHeaders.AUTHORIZATION, authorization, RequestAttributes.SCOPE_REQUEST);
135+
130136
if (request.hasFile()) {
131137
return processDocument(request.prompt(), request.fileBase64(), request.fileName())
132138
.collectList()

‎apps/java-spring-ai-agents/aiagent/src/main/java/com/example/agent/OAuthMcpConfig.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package com.example.agent;
2+
3+
import io.micrometer.context.ContextRegistry;
4+
import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer;
5+
import org.slf4j.Logger;
6+
import org.slf4j.LoggerFactory;
7+
import org.springframework.context.annotation.Bean;
8+
import org.springframework.context.annotation.Configuration;
9+
import org.springframework.http.HttpHeaders;
10+
import org.springframework.web.context.request.RequestAttributes;
11+
import org.springframework.web.context.request.RequestAttributesThreadLocalAccessor;
12+
import org.springframework.web.context.request.RequestContextHolder;
13+
14+
@Configuration
15+
public class OAuthMcpConfig {
16+
private static final Logger logger = LoggerFactory.getLogger(OAuthMcpConfig.class);
17+
18+
static {
19+
ContextRegistry.getInstance().registerThreadLocalAccessor(new RequestAttributesThreadLocalAccessor());
20+
}
21+
22+
@Bean
23+
McpSyncHttpClientRequestCustomizer oauthRequestCustomizer() {
24+
logger.info("OAuth token injection configured");
25+
26+
return (builder, method, endpoint, body, context) -> {
27+
String auth = getAuthFromRequestContext();
28+
if (auth != null) {
29+
logger.info("Authorization header propagated to MCP calls");
30+
builder.setHeader(HttpHeaders.AUTHORIZATION, auth);
31+
}
32+
};
33+
}
34+
35+
private String getAuthFromRequestContext() {
36+
try {
37+
return (String) RequestContextHolder.currentRequestAttributes()
38+
.getAttribute(HttpHeaders.AUTHORIZATION, RequestAttributes.SCOPE_REQUEST);
39+
} catch (IllegalStateException e) {
40+
logger.warn("Authorization header cannot be retrieved from local context: " + e.getMessage(), e);
41+
return null;
42+
}
43+
}
44+
}

‎apps/java-spring-ai-agents/aiagent/src/main/java/com/example/agent/SigV4McpConfig.java‎

Lines changed: 37 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -15,43 +15,44 @@
1515
import software.amazon.awssdk.http.SdkHttpMethod;
1616
import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
1717

18+
// Deactivated in favor to OAuthMcpConfig, because policy can evaluate only JWT principle
1819
@Configuration
1920
public class SigV4McpConfig {
2021

21-
private static final Logger log = LoggerFactory.getLogger(SigV4McpConfig.class);
22-
private static final Set<String> RESTRICTED_HEADERS = Set.of("content-length", "host", "expect");
23-
24-
@Bean
25-
McpSyncHttpClientRequestCustomizer sigV4RequestCustomizer() {
26-
var signer = Aws4Signer.create();
27-
var credentialsProvider = DefaultCredentialsProvider.create();
28-
var region = new DefaultAwsRegionProviderChain().getRegion();
29-
log.info("SigV4 MCP request customizer: region={}, service=bedrock-agentcore", region);
30-
31-
return (builder, method, endpoint, body, context) -> {
32-
byte[] bodyBytes = (body != null) ? body.getBytes(java.nio.charset.StandardCharsets.UTF_8) : null;
33-
34-
var sdkRequestBuilder = SdkHttpFullRequest.builder();
35-
sdkRequestBuilder.uri(endpoint);
36-
sdkRequestBuilder.method(SdkHttpMethod.valueOf(method));
37-
38-
if (bodyBytes != null && bodyBytes.length > 0) {
39-
sdkRequestBuilder.contentStreamProvider(() -> new ByteArrayInputStream(bodyBytes));
40-
sdkRequestBuilder.putHeader("Content-Length", String.valueOf(bodyBytes.length));
41-
}
42-
sdkRequestBuilder.putHeader("Content-Type", "application/json");
43-
44-
var signedRequest = signer.sign(sdkRequestBuilder.build(), Aws4SignerParams.builder()
45-
.signingName("bedrock-agentcore")
46-
.signingRegion(region)
47-
.awsCredentials(credentialsProvider.resolveCredentials())
48-
.build());
49-
50-
signedRequest.headers().forEach((name, values) -> {
51-
if (!RESTRICTED_HEADERS.contains(name.toLowerCase())) {
52-
values.forEach(value -> builder.setHeader(name, value));
53-
}
54-
});
55-
};
56-
}
22+
// private static final Logger log = LoggerFactory.getLogger(SigV4McpConfig.class);
23+
// private static final Set<String> RESTRICTED_HEADERS = Set.of("content-length", "host", "expect");
24+
//
25+
// @Bean
26+
// McpSyncHttpClientRequestCustomizer sigV4RequestCustomizer() {
27+
// var signer = Aws4Signer.create();
28+
// var credentialsProvider = DefaultCredentialsProvider.create();
29+
// var region = new DefaultAwsRegionProviderChain().getRegion();
30+
// log.info("SigV4 MCP request customizer: region={}, service=bedrock-agentcore", region);
31+
//
32+
// return (builder, method, endpoint, body, context) -> {
33+
// byte[] bodyBytes = (body != null) ? body.getBytes(java.nio.charset.StandardCharsets.UTF_8) : null;
34+
//
35+
// var sdkRequestBuilder = SdkHttpFullRequest.builder();
36+
// sdkRequestBuilder.uri(endpoint);
37+
// sdkRequestBuilder.method(SdkHttpMethod.valueOf(method));
38+
//
39+
// if (bodyBytes != null && bodyBytes.length > 0) {
40+
// sdkRequestBuilder.contentStreamProvider(() -> new ByteArrayInputStream(bodyBytes));
41+
// sdkRequestBuilder.putHeader("Content-Length", String.valueOf(bodyBytes.length));
42+
// }
43+
// sdkRequestBuilder.putHeader("Content-Type", "application/json");
44+
//
45+
// var signedRequest = signer.sign(sdkRequestBuilder.build(), Aws4SignerParams.builder()
46+
// .signingName("bedrock-agentcore")
47+
// .signingRegion(region)
48+
// .awsCredentials(credentialsProvider.resolveCredentials())
49+
// .build());
50+
//
51+
// signedRequest.headers().forEach((name, values) -> {
52+
// if (!RESTRICTED_HEADERS.contains(name.toLowerCase())) {
53+
// values.forEach(value -> builder.setHeader(name, value));
54+
// }
55+
// });
56+
// };
57+
// }
5758
}

‎apps/java-spring-ai-agents/scripts/policy/01-create-policy-engine.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Step 1: Create Policy Engine
4+
5+
Run: python 01-create-policy-engine.py
6+
Auto-updates ENGINE_ID in .env
7+
"""
8+
9+
from config import update_env
10+
from policy_commands import create_policy_engine, list_policy_engines
11+
12+
ENGINE_NAME = "TravelPolicyEngine"
13+
14+
if __name__ == "__main__":
15+
print("=== Creating Policy Engine ===")
16+
engine = create_policy_engine(ENGINE_NAME)
17+
18+
# Auto-update .env
19+
update_env("ENGINE_ID", engine['policyEngineId'])
20+
21+
print("\n=== All Policy Engines ===")
22+
list_policy_engines()

‎apps/java-spring-ai-agents/scripts/policy/02-create-policy.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Step 2: Create Policy
4+
5+
Run: python 02-create-policy.py
6+
Requires: ENGINE_ID set in config.py
7+
"""
8+
9+
from config import ENGINE_ID, GATEWAY_ARN, TARGET_NAME
10+
from policy_commands import create_policy, list_policies, delete_all_policies
11+
import time
12+
13+
# Policy: Permit all tools for alice, EXCEPT searchFlights
14+
POLICY = f'''permit(
15+
principal is AgentCore::OAuthUser,
16+
action,
17+
resource == AgentCore::Gateway::"{GATEWAY_ARN}"
18+
) when {{
19+
principal.hasTag("username") &&
20+
principal.getTag("username") == "alice"
21+
}} unless {{
22+
action == AgentCore::Action::"{TARGET_NAME}___searchFlights"
23+
}};'''
24+
25+
if __name__ == "__main__":
26+
print(f"Using ENGINE_ID: {ENGINE_ID}")
27+
print(f"Using GATEWAY_ARN: {GATEWAY_ARN}")
28+
29+
print("\n=== Deleting existing policies ===")
30+
delete_all_policies(ENGINE_ID)
31+
time.sleep(3)
32+
33+
print("\n=== Creating policy ===")
34+
create_policy(ENGINE_ID, "PermitAllExceptFlights", POLICY)
35+
36+
time.sleep(5)
37+
38+
print("\n=== Policy status ===")
39+
list_policies(ENGINE_ID)

‎apps/java-spring-ai-agents/scripts/policy/03-attach-policy-engine.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Step 3: Attach Policy Engine to Gateway
4+
5+
Run: python 03-attach-policy-engine.py
6+
Requires: ENGINE_ID and GATEWAY_ID set in config.py
7+
"""
8+
9+
from config import GATEWAY_ID, ENGINE_ARN
10+
from policy_commands import attach_policy_engine, get_gateway_policy_config
11+
12+
if __name__ == "__main__":
13+
import time
14+
15+
print(f"Using GATEWAY_ID: {GATEWAY_ID}")
16+
print(f"Using ENGINE_ARN: {ENGINE_ARN}")
17+
18+
print("\n=== Attaching Policy Engine ===")
19+
attach_policy_engine(GATEWAY_ID, ENGINE_ARN, mode="ENFORCE")
20+
21+
print("\n=== Waiting for attachment ===")
22+
for i in range(10):
23+
time.sleep(2)
24+
config = get_gateway_policy_config(GATEWAY_ID)
25+
if config:
26+
print("✓ Policy engine attached")
27+
break
28+
print(f" Attempt {i+1}/10...")
29+
else:
30+
print("✗ Attachment not confirmed after 10 attempts")

‎apps/java-spring-ai-agents/scripts/policy/README.md‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# AgentCore Policy-Based Access Control
2+
3+
Demonstrates Cedar policy-based access control for MCP tools through AgentCore Gateway.
4+
5+
## Overview
6+
7+
```
8+
chat-agent (JWT) → Gateway → Policy Engine → MCP Runtime
9+
↓
10+
Cedar policy checks
11+
user's username tag
12+
```
13+
14+
## Key Finding: Use `unless` for Denying Tools
15+
16+
AgentCore rejects standalone `forbid` policies as "Overly Restrictive".
17+
18+
**Solution:** Use `permit ... unless` to deny specific tools:
19+
20+
```cedar
21+
permit(
22+
principal is AgentCore::OAuthUser,
23+
action,
24+
resource == AgentCore::Gateway::"..."
25+
) when {
26+
principal.hasTag("username") &&
27+
principal.getTag("username") == "alice"
28+
} unless {
29+
action == AgentCore::Action::"travel-mcp___searchFlights"
30+
};
31+
```
32+
33+
## Setup Steps
34+
35+
```bash
36+
cd /Users/shakirin/Projects/agentcore/samples/policy/scripts/policy
37+
source .venv/bin/activate
38+
39+
# 1. Create policy engine (one-time)
40+
python 01-create-policy-engine.py
41+
42+
# 2. Create policy (update ENGINE_ID in script first)
43+
python 02-create-policy.py
44+
45+
# 3. Attach to gateway (update IDs in script first)
46+
python 03-attach-policy-engine.py
47+
```
48+
49+
## Test Results
50+
51+
| Tool | Policy | Result |
52+
|------|--------|--------|
53+
| `searchHotels` | Permitted | ✅ Returns hotel data |
54+
| `searchFlights` | Denied via `unless` | ❌ Tool not available |
55+
56+
## JWT Token Requirements
57+
58+
The policy uses `username` tag from Cognito user tokens:
59+
60+
```json
61+
{
62+
"username": "alice",
63+
"client_id": "...",
64+
"token_use": "access"
65+
}
66+
```
67+
68+
Get user token:
69+
```bash
70+
TOKEN=$(aws cognito-idp initiate-auth \
71+
--client-id $CLIENT_ID \
72+
--auth-flow USER_PASSWORD_AUTH \
73+
--auth-parameters "USERNAME=alice,PASSWORD=$PASSWORD,SECRET_HASH=$SECRET_HASH" \
74+
--region us-east-1 \
75+
--query 'AuthenticationResult.AccessToken' --output text)
76+
```
77+
78+
## Files
79+
80+
- `policy.cedar` - Working Cedar policy with `unless` clause
81+
- `policy_commands.py` - Helper functions for policy management
82+
- `01-create-policy-engine.py` - Create policy engine
83+
- `02-create-policy.py` - Create/update policy
84+
- `03-attach-policy-engine.py` - Attach engine to gateway

‎apps/java-spring-ai-agents/scripts/policy/config.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import os
2+
from dotenv import load_dotenv
3+
4+
load_dotenv()
5+
6+
REGION = os.getenv("REGION", "us-east-1")
7+
ACCOUNT_ID = os.getenv("ACCOUNT_ID")
8+
GATEWAY_ID = os.getenv("GATEWAY_ID")
9+
ENGINE_ID = os.getenv("ENGINE_ID")
10+
TARGET_NAME = os.getenv("TARGET_NAME", "travel-mcp")
11+
12+
# Derived ARNs
13+
GATEWAY_ARN = f"arn:aws:bedrock-agentcore:{REGION}:{ACCOUNT_ID}:gateway/{GATEWAY_ID}"
14+
ENGINE_ARN = f"arn:aws:bedrock-agentcore:{REGION}:{ACCOUNT_ID}:policy-engine/{ENGINE_ID}" if ENGINE_ID else None
15+
16+
17+
def update_env(key: str, value: str):
18+
"""Update a value in .env file."""
19+
env_path = os.path.join(os.path.dirname(__file__), ".env")
20+
21+
with open(env_path, "r") as f:
22+
lines = f.readlines()
23+
24+
updated = False
25+
for i, line in enumerate(lines):
26+
if line.startswith(f"{key}="):
27+
lines[i] = f"{key}={value}\n"
28+
updated = True
29+
break
30+
31+
if not updated:
32+
lines.append(f"{key}={value}\n")
33+
34+
with open(env_path, "w") as f:
35+
f.writelines(lines)
36+
37+
# Update current process
38+
os.environ[key] = value
39+
print(f"Updated .env: {key}={value}")

‎apps/java-spring-ai-agents/scripts/policy/deny-tools-policy.cedar‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# Cedar policy: Permit all tools EXCEPT cancelTrip and deleteExpense for user alice
2+
#
3+
# Uses 'unless' clause to deny specific tools while permitting all others.
4+
# This approach works around AgentCore's "Overly Restrictive" safety check
5+
# that rejects standalone 'forbid' policies.
6+
7+
permit(
8+
principal is AgentCore::OAuthUser,
9+
action,
10+
resource == AgentCore::Gateway::"arn:aws:bedrock-agentcore:us-east-1:724772082315:gateway/policy-demo-gateway-wwh6rjluyl"
11+
) when {
12+
principal.hasTag("username") &&
13+
principal.getTag("username") == "alice"
14+
} unless {
15+
action == AgentCore::Action::"backoffice___cancelTrip" ||
16+
action == AgentCore::Action::"backoffice___deleteExpense"
17+
};

0 commit comments

Comments
 (0)