Skip to content

Commit 53628ad

Browse files
jchrostek-ddclaude
jchrostek-dd
and
claude
committed
fix: align delegated auth proof format with agent implementation
Match the Datadog agent's proof format: header values as arrays, canonical HTTP header casing, no trailing slash on STS URL. Also read DD_ORG_UUID from SERVERLESS_UUID env var and remove trace assertion. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 5e7bff2 commit 53628ad

3 files changed

Lines changed: 118 additions & 177 deletions

File tree

‎bottlecap/src/delegated_auth/auth_proof.rs‎

Lines changed: 40 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ pub fn generate_auth_proof(
6363
} else {
6464
format!("sts.{region}.amazonaws.com")
6565
};
66-
let sts_url = format!("https://{sts_host}/");
66+
let sts_url = format!("https://{sts_host}");
6767

6868
// Get current time for signing
6969
let now = Utc::now();
@@ -136,17 +136,17 @@ pub fn generate_auth_proof(
136136

137137
// Build headers map for the proof
138138
// Using BTreeMap for consistent ordering (important for signature verification)
139-
let mut headers_map: BTreeMap<String, String> = BTreeMap::new();
140-
headers_map.insert("Authorization".to_string(), authorization);
141-
headers_map.insert("Content-Type".to_string(), CONTENT_TYPE.to_string());
142-
headers_map.insert("Host".to_string(), sts_host);
143-
headers_map.insert("x-amz-date".to_string(), amz_date);
144-
headers_map.insert(ORG_ID_HEADER.to_string(), org_uuid.to_string());
139+
let mut headers_map: BTreeMap<String, Vec<String>> = BTreeMap::new();
140+
headers_map.insert("Authorization".to_string(), vec![authorization]);
141+
headers_map.insert("Content-Type".to_string(), vec![CONTENT_TYPE.to_string()]);
142+
headers_map.insert("Host".to_string(), vec![sts_host]);
143+
headers_map.insert("X-Amz-Date".to_string(), vec![amz_date]);
144+
headers_map.insert("X-Ddog-Org-Id".to_string(), vec![org_uuid.to_string()]);
145145

146146
if !aws_credentials.aws_session_token.is_empty() {
147147
headers_map.insert(
148-
"x-amz-security-token".to_string(),
149-
aws_credentials.aws_session_token.clone(),
148+
"X-Amz-Security-Token".to_string(),
149+
vec![aws_credentials.aws_session_token.clone()],
150150
);
151151
}
152152

@@ -248,17 +248,23 @@ mod tests {
248248

249249
// Verify body is base64-encoded GET_CALLER_IDENTITY_BODY
250250
let body = String::from_utf8(
251-
BASE64_STANDARD.decode(parts[0]).expect("Failed to decode base64 body")
252-
).expect("Failed to convert body to UTF-8");
251+
BASE64_STANDARD
252+
.decode(parts[0])
253+
.expect("Failed to decode base64 body"),
254+
)
255+
.expect("Failed to convert body to UTF-8");
253256
assert_eq!(body, GET_CALLER_IDENTITY_BODY);
254257

255258
// Verify method
256259
assert_eq!(parts[2], "POST");
257260

258261
// Verify URL is base64-encoded STS URL
259262
let url = String::from_utf8(
260-
BASE64_STANDARD.decode(parts[3]).expect("Failed to decode base64 URL")
261-
).expect("Failed to convert URL to UTF-8");
263+
BASE64_STANDARD
264+
.decode(parts[3])
265+
.expect("Failed to decode base64 URL"),
266+
)
267+
.expect("Failed to convert URL to UTF-8");
262268
assert!(url.contains("sts.us-east-1.amazonaws.com"));
263269
}
264270

@@ -280,23 +286,28 @@ mod tests {
280286

281287
// Decode and parse headers
282288
let headers_json = String::from_utf8(
283-
BASE64_STANDARD.decode(parts[1]).expect("Failed to decode base64 headers")
284-
).expect("Failed to convert headers to UTF-8");
285-
let headers: BTreeMap<String, String> = serde_json::from_str(&headers_json)
286-
.expect("Failed to parse headers JSON");
289+
BASE64_STANDARD
290+
.decode(parts[1])
291+
.expect("Failed to decode base64 headers"),
292+
)
293+
.expect("Failed to convert headers to UTF-8");
294+
let headers: BTreeMap<String, Vec<String>> =
295+
serde_json::from_str(&headers_json).expect("Failed to parse headers JSON");
287296

288-
// Verify required headers
297+
// Verify required headers (canonical casing)
289298
assert!(headers.contains_key("Authorization"));
290299
assert!(headers.contains_key("Content-Type"));
291300
assert!(headers.contains_key("Host"));
292-
assert!(headers.contains_key("x-amz-date"));
293-
assert!(headers.contains_key("x-amz-security-token"));
294-
assert!(headers.contains_key("x-ddog-org-id"));
301+
assert!(headers.contains_key("X-Amz-Date"));
302+
assert!(headers.contains_key("X-Amz-Security-Token"));
303+
assert!(headers.contains_key("X-Ddog-Org-Id"));
295304

296-
// Verify org-id header value
305+
// Verify org-id header value (array format)
297306
assert_eq!(
298-
headers.get("x-ddog-org-id").expect("Missing x-ddog-org-id header"),
299-
"my-org-uuid"
307+
headers
308+
.get("X-Ddog-Org-Id")
309+
.expect("Missing X-Ddog-Org-Id header"),
310+
&vec!["my-org-uuid".to_string()]
300311
);
301312
}
302313

@@ -317,8 +328,11 @@ mod tests {
317328
let proof = result.expect("Failed to generate auth proof");
318329
let parts: Vec<&str> = proof.split('|').collect();
319330
let url = String::from_utf8(
320-
BASE64_STANDARD.decode(parts[3]).expect("Failed to decode base64 URL")
321-
).expect("Failed to convert URL to UTF-8");
331+
BASE64_STANDARD
332+
.decode(parts[3])
333+
.expect("Failed to decode base64 URL"),
334+
)
335+
.expect("Failed to convert URL to UTF-8");
322336
assert!(url.contains("sts.amazonaws.com"));
323337
}
324338
}

‎integration-tests/lib/stacks/delegated-auth.ts‎

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import * as cdk from 'aws-cdk-lib';
2+
import * as iam from 'aws-cdk-lib/aws-iam';
23
import * as lambda from 'aws-cdk-lib/aws-lambda';
34
import { Construct } from 'constructs';
45
import {
@@ -20,35 +21,44 @@ export class DelegatedAuthStack extends cdk.Stack {
2021

2122
const extensionLayer = getExtensionLayer(this);
2223

23-
// Happy Path Function - Uses delegated auth only
24-
const happyPathFunctionName = `${id}-happy-path`;
25-
const happyPathFunction = new lambda.Function(this, happyPathFunctionName, {
24+
const functionName = id;
25+
const roleName = `${id}-role`;
26+
const role = new iam.Role(this, 'ExecutionRole', {
27+
roleName,
28+
assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
29+
managedPolicies: [
30+
iam.ManagedPolicy.fromAwsManagedPolicyName('service-role/AWSLambdaBasicExecutionRole'),
31+
],
32+
});
33+
34+
const fn = new lambda.Function(this, functionName, {
35+
role,
2636
runtime: defaultNodeRuntime,
2737
architecture: lambda.Architecture.ARM_64,
2838
handler: 'index.handler',
2939
code: lambda.Code.fromAsset('./lambda/delegated-auth'),
30-
functionName: happyPathFunctionName,
40+
functionName,
3141
timeout: cdk.Duration.seconds(30),
3242
memorySize: 256,
3343
environment: {
3444
DD_SITE: 'datadoghq.com',
3545
DD_ENV: 'integration',
3646
DD_VERSION: '1.0.0',
37-
DD_SERVICE: happyPathFunctionName,
47+
DD_SERVICE: functionName,
3848
DD_SERVERLESS_FLUSH_STRATEGY: 'end',
3949
DD_SERVERLESS_LOGS_ENABLED: 'true',
4050
DD_LOG_LEVEL: 'debug',
4151
// Delegated auth config
42-
DD_ORG_UUID: '447397',
52+
DD_ORG_UUID: process.env.SERVERLESS_UUID || '',
4353
DD_DELEGATED_AUTH_ENABLED: 'true',
4454
TS: Date.now().toString(),
4555
},
46-
logGroup: createLogGroup(this, happyPathFunctionName),
56+
logGroup: createLogGroup(this, functionName),
4757
});
48-
happyPathFunction.addLayers(extensionLayer);
58+
fn.addLayers(extensionLayer);
4959

50-
new cdk.CfnOutput(this, 'HappyPathRoleArn', {
51-
value: happyPathFunction.role!.roleArn,
60+
new cdk.CfnOutput(this, 'RoleArn', {
61+
value: fn.role!.roleArn,
5262
description: 'IAM Role ARN - configure in intake mapping',
5363
});
5464
}

‎integration-tests/tests/delegated-auth.test.ts‎

Lines changed: 58 additions & 141 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,8 @@ import { getIdentifier } from '../config';
55
const identifier = getIdentifier();
66
const stackName = `integ-${identifier}-delegated-auth`;
77

8-
// Function names from CDK stack
9-
const HAPPY_PATH_FUNCTION = `${stackName}-happy-path`;
10-
const FALLBACK_FUNCTION = `${stackName}-fallback`;
8+
// Function name matches the CDK stack id
9+
const FUNCTION_NAME = stackName;
1110

1211
// Default wait time for Datadog to index logs and traces after Lambda invocation
1312
const DEFAULT_DATADOG_INDEXING_WAIT_MS = 5 * 60 * 1000; // 5 minutes
@@ -33,146 +32,64 @@ async function getDatadogTelemetryByRequestId(
3332
}
3433

3534
describe('Delegated Authentication Integration Tests', () => {
35+
let invocationResult: { requestId: string; statusCode?: number };
36+
let telemetry: DatadogTelemetry;
37+
let logs: string[];
3638

37-
describe('Happy Path - Delegated Auth Success', () => {
38-
let invocationResult: { requestId: string; statusCode?: number };
39-
let telemetry: DatadogTelemetry;
40-
let logs: string[];
41-
42-
beforeAll(async () => {
43-
console.log(`Testing happy path function: ${HAPPY_PATH_FUNCTION}`);
44-
45-
// Force cold start to ensure extension initializes fresh
46-
await forceColdStart(HAPPY_PATH_FUNCTION);
47-
48-
// Invoke the function
49-
invocationResult = await invokeLambda(HAPPY_PATH_FUNCTION, {});
50-
51-
console.log(`Invocation completed, requestId: ${invocationResult.requestId}`);
52-
53-
// Wait for telemetry to be indexed in Datadog
54-
console.log(`Waiting ${DEFAULT_DATADOG_INDEXING_WAIT_MS / 1000}s for Datadog indexing...`);
55-
await sleep(DEFAULT_DATADOG_INDEXING_WAIT_MS);
56-
57-
// Collect telemetry from Datadog
58-
telemetry = await getDatadogTelemetryByRequestId(
59-
HAPPY_PATH_FUNCTION,
60-
invocationResult.requestId
61-
);
62-
logs = telemetry.logs.map((log: DatadogLog) => log.message);
63-
64-
console.log(`Collected ${telemetry.logs.length} logs and ${telemetry.traces.length} traces`);
65-
}, 600000); // 10 minute timeout
66-
67-
it('should invoke Lambda successfully', () => {
68-
expect(invocationResult).toBeDefined();
69-
expect(invocationResult.statusCode).toBe(200);
70-
});
71-
72-
it('should have function log output', () => {
73-
expect(telemetry).toBeDefined();
74-
expect(telemetry.logs).toBeDefined();
75-
expect(telemetry.logs.length).toBeGreaterThan(0);
76-
});
77-
78-
it('should show delegated auth API key obtained successfully', () => {
79-
// Look for log message indicating delegated auth succeeded
80-
const delegatedAuthLog = logs.find((log: string) =>
81-
log.includes('Delegated auth') &&
82-
(log.includes('API key obtained') || log.includes('success'))
83-
);
84-
expect(delegatedAuthLog).toBeDefined();
85-
});
86-
87-
it('should NOT show fallback to static API key', () => {
88-
// Ensure no fallback occurred
89-
const fallbackLog = logs.find((log: string) =>
90-
log.includes('fallback') || log.includes('Falling back')
91-
);
92-
expect(fallbackLog).toBeUndefined();
93-
});
94-
95-
it('should send telemetry to Datadog (validates API key works)', () => {
96-
// If we have logs in Datadog, the obtained API key is working
97-
expect(telemetry.logs).toBeDefined();
98-
expect(telemetry.logs.length).toBeGreaterThan(0);
99-
});
100-
101-
it('should send at least one trace to Datadog', () => {
102-
// Traces indicate the extension is functioning correctly
103-
expect(telemetry.traces).toBeDefined();
104-
expect(telemetry.traces.length).toBeGreaterThanOrEqual(1);
105-
});
39+
beforeAll(async () => {
40+
console.log(`Testing delegated auth function: ${FUNCTION_NAME}`);
41+
42+
// Force cold start to ensure extension initializes fresh
43+
await forceColdStart(FUNCTION_NAME);
44+
45+
// Invoke the function
46+
invocationResult = await invokeLambda(FUNCTION_NAME, {});
47+
48+
console.log(`Invocation completed, requestId: ${invocationResult.requestId}`);
49+
50+
// Wait for telemetry to be indexed in Datadog
51+
console.log(`Waiting ${DEFAULT_DATADOG_INDEXING_WAIT_MS / 1000}s for Datadog indexing...`);
52+
await sleep(DEFAULT_DATADOG_INDEXING_WAIT_MS);
53+
54+
// Collect telemetry from Datadog
55+
telemetry = await getDatadogTelemetryByRequestId(
56+
FUNCTION_NAME,
57+
invocationResult.requestId
58+
);
59+
logs = telemetry.logs.map((log: DatadogLog) => log.message);
60+
61+
console.log(`Collected ${telemetry.logs.length} logs and ${telemetry.traces.length} traces`);
62+
}, 600000); // 10 minute timeout
63+
64+
it('should invoke Lambda successfully', () => {
65+
expect(invocationResult).toBeDefined();
66+
expect(invocationResult.statusCode).toBe(200);
10667
});
10768

108-
describe('Fallback Path - Invalid Org UUID Falls Back to Static Key', () => {
109-
let invocationResult: { requestId: string; statusCode?: number };
110-
let telemetry: DatadogTelemetry;
111-
let logs: string[];
112-
113-
beforeAll(async () => {
114-
console.log(`Testing fallback function: ${FALLBACK_FUNCTION}`);
115-
116-
// Force cold start to ensure extension initializes fresh
117-
await forceColdStart(FALLBACK_FUNCTION);
118-
119-
// Invoke the function
120-
invocationResult = await invokeLambda(FALLBACK_FUNCTION, {});
121-
122-
console.log(`Invocation completed, requestId: ${invocationResult.requestId}`);
123-
124-
// Wait for telemetry to be indexed in Datadog
125-
console.log(`Waiting ${DEFAULT_DATADOG_INDEXING_WAIT_MS / 1000}s for Datadog indexing...`);
126-
await sleep(DEFAULT_DATADOG_INDEXING_WAIT_MS);
127-
128-
// Collect telemetry from Datadog
129-
telemetry = await getDatadogTelemetryByRequestId(
130-
FALLBACK_FUNCTION,
131-
invocationResult.requestId
132-
);
133-
logs = telemetry.logs.map((log: DatadogLog) => log.message);
134-
135-
console.log(`Collected ${telemetry.logs.length} logs and ${telemetry.traces.length} traces`);
136-
}, 600000); // 10 minute timeout
137-
138-
it('should invoke Lambda successfully', () => {
139-
expect(invocationResult).toBeDefined();
140-
expect(invocationResult.statusCode).toBe(200);
141-
});
142-
143-
it('should have function log output', () => {
144-
expect(telemetry).toBeDefined();
145-
expect(telemetry.logs).toBeDefined();
146-
expect(telemetry.logs.length).toBeGreaterThan(0);
147-
});
148-
149-
it('should show delegated auth failure', () => {
150-
// Look for log message indicating delegated auth failed
151-
const failureLog = logs.find((log: string) =>
152-
(log.includes('Delegated auth') || log.includes('delegated auth')) &&
153-
(log.includes('fail') || log.includes('error') || log.includes('Error'))
154-
);
155-
expect(failureLog).toBeDefined();
156-
});
157-
158-
it('should show fallback to static API key', () => {
159-
// Look for log message indicating fallback occurred
160-
const fallbackLog = logs.find((log: string) =>
161-
log.includes('fallback') || log.includes('Falling back') || log.includes('using static')
162-
);
163-
expect(fallbackLog).toBeDefined();
164-
});
165-
166-
it('should still send telemetry to Datadog (via fallback key)', () => {
167-
// Even with delegated auth failure, telemetry should work via fallback
168-
expect(telemetry.logs).toBeDefined();
169-
expect(telemetry.logs.length).toBeGreaterThan(0);
170-
});
171-
172-
it('should still send traces to Datadog (via fallback key)', () => {
173-
// Traces should still work via the fallback static API key
174-
expect(telemetry.traces).toBeDefined();
175-
expect(telemetry.traces.length).toBeGreaterThanOrEqual(1);
176-
});
69+
it('should have function log output', () => {
70+
expect(telemetry).toBeDefined();
71+
expect(telemetry.logs).toBeDefined();
72+
expect(telemetry.logs.length).toBeGreaterThan(0);
17773
});
74+
75+
it('should show delegated auth API key obtained successfully', () => {
76+
const delegatedAuthLog = logs.find((log: string) =>
77+
log.includes('Delegated auth') &&
78+
(log.includes('API key obtained') || log.includes('success'))
79+
);
80+
expect(delegatedAuthLog).toBeDefined();
81+
});
82+
83+
it('should NOT show fallback to static API key', () => {
84+
const fallbackLog = logs.find((log: string) =>
85+
log.includes('fallback') || log.includes('Falling back')
86+
);
87+
expect(fallbackLog).toBeUndefined();
88+
});
89+
90+
it('should send telemetry to Datadog (validates API key works)', () => {
91+
expect(telemetry.logs).toBeDefined();
92+
expect(telemetry.logs.length).toBeGreaterThan(0);
93+
});
94+
17895
});

0 commit comments

Comments
 (0)