Skip to content

Commit c2ac1e0

Browse files
authored
[APPSEC-60218] Fix AWS API Gateway endpoints correlation HTTP span tags (#967)
## Overview - API Gateway REST/HTTP inferred spans now emit `span.type` as `web` - droped `operation_name` and `apiname` (in api gateway REST, HTTP and websocket) - Inferred spans now carry the API Gateway ARN in `dd_resource_key` - Appsec enablement and the json data from the service entry span are propagated into inferred and wrapped spans ## Testing - Tested if dd_resource_key are correctly set for api gateway (HTTP/REST) - Tested if inferred spans correctly have the appsec data propagated (appsec enabled via entry span or via config) - Updated tests to remove `operation_name` occurence Co-authored-by: flavien.darche <flavien.darche@datadoghq.com>
1 parent 948a938 commit c2ac1e0

5 files changed

Lines changed: 214 additions & 16 deletions

File tree

bottlecap/src/lifecycle/invocation/span_inferrer.rs

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,12 @@ impl SpanInferrer {
227227
);
228228
}
229229

230+
if let Some(dd_resource_key) = t.get_dd_resource_key(&aws_config.region) {
231+
inferred_span
232+
.meta
233+
.insert("dd_resource_key".to_string(), dd_resource_key);
234+
}
235+
230236
self.wrapped_inferred_span = wrapped_inferred_span;
231237
self.span_pointers = span_pointers;
232238

@@ -278,12 +284,15 @@ impl SpanInferrer {
278284
invocation_span.service.clone(),
279285
);
280286
s.meta.insert("span.kind".to_string(), "server".to_string());
287+
let appsec_enabled = self.config.serverless_appsec_enabled;
288+
propagate_appsec(appsec_enabled, invocation_span, s);
281289

282290
if let Some(ws) = &mut self.wrapped_inferred_span {
283291
ws.trace_id = invocation_span.trace_id;
284292
ws.error = invocation_span.error;
285293
ws.meta
286294
.insert(String::from("peer.service"), s.service.clone());
295+
propagate_appsec(appsec_enabled, invocation_span, ws);
287296

288297
// The wrapper span should be the parent of the inferred span,
289298
// therefore the `parent_id` of the inferred span should be the
@@ -325,6 +334,34 @@ impl SpanInferrer {
325334
}
326335
}
327336

337+
fn propagate_appsec(
338+
serverless_appsec_enabled: bool,
339+
invocation_span: &Span,
340+
target_span: &mut Span,
341+
) {
342+
let has_appsec = invocation_span
343+
.metrics
344+
.get("_dd.appsec.enabled")
345+
.copied()
346+
.or(if serverless_appsec_enabled {
347+
Some(1.0)
348+
} else {
349+
None
350+
});
351+
352+
if let Some(enabled) = has_appsec {
353+
target_span
354+
.metrics
355+
.insert("_dd.appsec.enabled".to_string(), enabled);
356+
}
357+
358+
if let Some(json) = invocation_span.meta.get("_dd.appsec.json") {
359+
target_span
360+
.meta
361+
.insert("_dd.appsec.json".to_string(), json.clone());
362+
}
363+
}
364+
328365
pub fn extract_span_context(
329366
payload_value: &Value,
330367
propagator: Arc<impl Propagator>,
@@ -368,6 +405,7 @@ pub fn extract_generated_span_context(
368405
#[cfg(test)]
369406
mod tests {
370407
use super::*;
408+
use crate::lifecycle::invocation::triggers::test_utils::read_json_file;
371409
use crate::traces::propagation::text_map_propagator::DatadogHeaderPropagator;
372410
use serde_json::json;
373411
use std::sync::Arc;
@@ -571,4 +609,117 @@ mod tests {
571609
"Should have SQS as event source"
572610
);
573611
}
612+
613+
fn api_gateway_rest_payload() -> serde_json::Value {
614+
let json = read_json_file("api_gateway_rest_event.json");
615+
serde_json::from_str(&json).expect("Failed to deserialize API Gateway REST payload")
616+
}
617+
618+
fn aws_config(region: &str) -> Arc<AwsConfig> {
619+
Arc::new(AwsConfig {
620+
region: region.to_string(),
621+
aws_lwa_proxy_lambda_runtime_api: Some(String::new()),
622+
runtime_api: String::new(),
623+
function_name: String::new(),
624+
sandbox_init_time: Instant::now(),
625+
exec_wrapper: None,
626+
initialization_type: "on-demand".into(),
627+
})
628+
}
629+
630+
#[test]
631+
fn test_complete_inferred_spans_propagates_appsec_from_invocation() {
632+
let payload = api_gateway_rest_payload();
633+
let aws_config = aws_config("us-east-1");
634+
let mut inferrer = SpanInferrer::new(Arc::new(Config::default()));
635+
636+
inferrer.infer_span(&payload, &aws_config);
637+
638+
let mut invocation_span = Span {
639+
trace_id: 42,
640+
span_id: 100,
641+
service: "lambda-service".to_string(),
642+
..Span::default()
643+
};
644+
if let Some(inferred_span) = &inferrer.inferred_span {
645+
invocation_span.start = inferred_span.start;
646+
}
647+
invocation_span.duration = 1;
648+
invocation_span
649+
.metrics
650+
.insert("_dd.appsec.enabled".to_string(), 1.0);
651+
invocation_span.meta.insert(
652+
"_dd.appsec.json".to_string(),
653+
r#"{"triggers":["rule"]}"#.to_string(),
654+
);
655+
656+
inferrer.complete_inferred_spans(&invocation_span);
657+
658+
let inferred_span = inferrer
659+
.inferred_span
660+
.as_ref()
661+
.expect("Inferred span should still be present");
662+
663+
let appsec_enabled = inferred_span
664+
.metrics
665+
.get("_dd.appsec.enabled")
666+
.copied()
667+
.unwrap_or_default();
668+
assert!(
669+
(appsec_enabled - 1.0).abs() < f64::EPSILON,
670+
"Expected appsec enabled metric to be 1.0"
671+
);
672+
assert_eq!(
673+
inferred_span
674+
.meta
675+
.get("_dd.appsec.json")
676+
.cloned()
677+
.unwrap_or_default(),
678+
r#"{"triggers":["rule"]}"#
679+
);
680+
}
681+
682+
#[test]
683+
fn test_complete_inferred_spans_sets_appsec_when_enabled_in_config() {
684+
let config = Config {
685+
serverless_appsec_enabled: true,
686+
..Config::default()
687+
};
688+
let mut inferrer = SpanInferrer::new(Arc::new(config));
689+
690+
let payload = api_gateway_rest_payload();
691+
let aws_config = aws_config("us-east-1");
692+
inferrer.infer_span(&payload, &aws_config);
693+
694+
let mut invocation_span = Span {
695+
trace_id: 7,
696+
service: "lambda-service".to_string(),
697+
..Span::default()
698+
};
699+
if let Some(inferred_span) = &inferrer.inferred_span {
700+
invocation_span.start = inferred_span.start;
701+
}
702+
invocation_span.duration = 1;
703+
704+
inferrer.complete_inferred_spans(&invocation_span);
705+
706+
let inferred_span = inferrer
707+
.inferred_span
708+
.as_ref()
709+
.expect("Inferred span should still be present");
710+
711+
let appsec_enabled = inferred_span
712+
.metrics
713+
.get("_dd.appsec.enabled")
714+
.copied()
715+
.unwrap_or_default();
716+
assert!(
717+
(appsec_enabled - 1.0).abs() < f64::EPSILON,
718+
"Expected appsec enabled metric to be 1.0"
719+
);
720+
assert!(
721+
!inferred_span.meta.contains_key("_dd.appsec.json"),
722+
"AppSec JSON should not be added when invocation span has none"
723+
);
724+
}
574725
}

bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ impl Trigger for APIGatewayHttpEvent {
107107
span.name = "aws.httpapi".to_string();
108108
span.service = service_name;
109109
span.resource.clone_from(&resource);
110-
span.r#type = "http".to_string();
110+
span.r#type = "web".to_string();
111111
span.start = start_time;
112112
span.meta.extend(HashMap::from([
113113
(
@@ -131,7 +131,6 @@ impl Trigger for APIGatewayHttpEvent {
131131
"http.user_agent".to_string(),
132132
self.request_context.http.user_agent.clone(),
133133
),
134-
("operation_name".to_string(), "aws.httpapi".to_string()),
135134
(
136135
"request_id".to_string(),
137136
self.request_context.request_id.clone(),
@@ -200,6 +199,20 @@ impl Trigger for APIGatewayHttpEvent {
200199
)
201200
}
202201

202+
fn get_dd_resource_key(&self, region: &str) -> Option<String> {
203+
if self.request_context.api_id.is_empty() {
204+
return None;
205+
}
206+
207+
let partition = get_aws_partition_by_region(region);
208+
Some(format!(
209+
"arn:{partition}:apigateway:{region}::/apis/{api_id}",
210+
partition = partition,
211+
region = region,
212+
api_id = self.request_context.api_id
213+
))
214+
}
215+
203216
fn is_async(&self) -> bool {
204217
self.headers
205218
.get("x-amz-invocation-type")
@@ -220,6 +233,7 @@ impl ServiceNameResolver for APIGatewayHttpEvent {
220233
"lambda_api_gateway"
221234
}
222235
}
236+
223237
#[cfg(test)]
224238
mod tests {
225239
use super::*;
@@ -309,7 +323,7 @@ mod tests {
309323
"x02yirxc7a.execute-api.sa-east-1.amazonaws.com"
310324
);
311325
assert_eq!(span.resource, "GET /httpapi/get");
312-
assert_eq!(span.r#type, "http");
326+
assert_eq!(span.r#type, "web");
313327
assert_eq!(
314328
span.meta,
315329
HashMap::from([
@@ -323,7 +337,6 @@ mod tests {
323337
("http.protocol".to_string(), "HTTP/1.1".to_string()),
324338
("http.source_ip".to_string(), "38.122.226.210".to_string()),
325339
("http.user_agent".to_string(), "curl/7.64.1".to_string()),
326-
("operation_name".to_string(), "aws.httpapi".to_string()),
327340
("request_id".to_string(), "FaHnXjKCGjQEJ7A=".to_string()),
328341
])
329342
);
@@ -373,7 +386,7 @@ mod tests {
373386
"9vj54we5ih.execute-api.sa-east-1.amazonaws.com"
374387
);
375388
assert_eq!(span.resource, "GET /user/{user_id}");
376-
assert_eq!(span.r#type, "http");
389+
assert_eq!(span.r#type, "web");
377390
assert_eq!(
378391
span.meta,
379392
HashMap::from([
@@ -386,7 +399,6 @@ mod tests {
386399
("http.protocol".to_string(), "HTTP/1.1".to_string()),
387400
("http.source_ip".to_string(), "76.115.124.192".to_string()),
388401
("http.user_agent".to_string(), "curl/8.1.2".to_string()),
389-
("operation_name".to_string(), "aws.httpapi".to_string()),
390402
("request_id".to_string(), "Ur2JtjEfGjQEPOg=".to_string()),
391403
])
392404
);
@@ -429,6 +441,18 @@ mod tests {
429441
);
430442
}
431443

444+
#[test]
445+
fn test_get_dd_resource_key() {
446+
let json = read_json_file("api_gateway_http_event.json");
447+
let payload = serde_json::from_str(&json).expect("Failed to deserialize into Value");
448+
let event =
449+
APIGatewayHttpEvent::new(payload).expect("Failed to deserialize APIGatewayHttpEvent");
450+
assert_eq!(
451+
event.get_dd_resource_key("sa-east-1"),
452+
Some("arn:aws:apigateway:sa-east-1::/apis/x02yirxc7a".to_string())
453+
);
454+
}
455+
432456
#[test]
433457
fn test_resolve_service_name_with_representation_enabled() {
434458
let json = read_json_file("api_gateway_http_event.json");

bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ impl Trigger for APIGatewayRestEvent {
104104
span.name = "aws.apigateway".to_string();
105105
span.service = service_name;
106106
span.resource = resource;
107-
span.r#type = "http".to_string();
107+
span.r#type = "web".to_string();
108108
span.start = start_time;
109109
span.meta.extend(HashMap::from([
110110
("endpoint".to_string(), self.request_context.path.clone()),
@@ -125,7 +125,6 @@ impl Trigger for APIGatewayRestEvent {
125125
"http.user_agent".to_string(),
126126
self.request_context.identity.user_agent.clone(),
127127
),
128-
("operation_name".to_string(), "aws.apigateway".to_string()),
129128
(
130129
"request_id".to_string(),
131130
self.request_context.request_id.clone(),
@@ -187,6 +186,20 @@ impl Trigger for APIGatewayRestEvent {
187186
)
188187
}
189188

189+
fn get_dd_resource_key(&self, region: &str) -> Option<String> {
190+
if self.request_context.api_id.is_empty() {
191+
return None;
192+
}
193+
194+
let partition = get_aws_partition_by_region(region);
195+
Some(format!(
196+
"arn:{partition}:apigateway:{region}::/restapis/{api_id}",
197+
partition = partition,
198+
region = region,
199+
api_id = self.request_context.api_id
200+
))
201+
}
202+
190203
fn is_async(&self) -> bool {
191204
self.headers
192205
.get("x-amz-invocation-type")
@@ -327,7 +340,7 @@ mod tests {
327340
assert_eq!(span.name, "aws.apigateway");
328341
assert_eq!(span.service, "id.execute-api.us-east-1.amazonaws.com");
329342
assert_eq!(span.resource, "GET /my/path");
330-
assert_eq!(span.r#type, "http");
343+
assert_eq!(span.r#type, "web");
331344

332345
assert_eq!(
333346
span.meta,
@@ -342,7 +355,6 @@ mod tests {
342355
("http.source_ip".to_string(), "IP".to_string()),
343356
("http.user_agent".to_string(), "user-agent".to_string()),
344357
("http.route".to_string(), "/path".to_string()),
345-
("operation_name".to_string(), "aws.apigateway".to_string()),
346358
("request_id".to_string(), "id=".to_string()),
347359
])
348360
);
@@ -389,7 +401,7 @@ mod tests {
389401
"mcwkra0ya4.execute-api.sa-east-1.amazonaws.com"
390402
);
391403
assert_eq!(span.resource, "GET /dev/user/{user_id}/id/{id}");
392-
assert_eq!(span.r#type, "http");
404+
assert_eq!(span.r#type, "web");
393405
let expected = HashMap::from([
394406
("endpoint".to_string(), "/dev/user/42/id/50".to_string()),
395407
(
@@ -402,7 +414,6 @@ mod tests {
402414
("http.source_ip".to_string(), "76.115.124.192".to_string()),
403415
("http.user_agent".to_string(), "curl/8.1.2".to_string()),
404416
("http.route".to_string(), "/user/{id}".to_string()),
405-
("operation_name".to_string(), "aws.apigateway".to_string()),
406417
(
407418
"request_id".to_string(),
408419
"e16399f7-e984-463a-9931-745ba021a27f".to_string(),
@@ -454,6 +465,18 @@ mod tests {
454465
);
455466
}
456467

468+
#[test]
469+
fn test_get_dd_resource_key() {
470+
let json = read_json_file("api_gateway_rest_event.json");
471+
let payload = serde_json::from_str(&json).expect("Failed to deserialize into Value");
472+
let event =
473+
APIGatewayRestEvent::new(payload).expect("Failed to deserialize APIGatewayRestEvent");
474+
assert_eq!(
475+
event.get_dd_resource_key("us-east-1"),
476+
Some("arn:aws:apigateway:us-east-1::/restapis/id".to_string())
477+
);
478+
}
479+
457480
#[test]
458481
fn test_resolve_service_name_with_representation_enabled() {
459482
let json = read_json_file("api_gateway_rest_event.json");

0 commit comments

Comments
 (0)