Merge branch 'main' into shreya.malpani/fix-e2e-test-trigger

Author: shreyamalpani

bottlecap/src/appsec/processor/mod.rs

@@ -151,10 +151,18 @@ impl Processor {
     /// Returns the first `aws.lambda` span from the provided trace, if one
     /// exists.
     ///
+    /// Placeholder spans (resource == `INVOCATION_SPAN_RESOURCE`) emitted by
+    /// Go and Java tracers are excluded: they are always dropped by the chunk
+    /// processor before reaching the backend, so tagging them would waste the
+    /// `AppSec` context and trigger a premature context deletion that would leave
+    /// the real, extension-built `aws.lambda` span untagged.
+    ///
     /// # Returns
     /// The span on which security information will be attached.
     pub fn service_entry_span_mut(trace: &mut [Span]) -> Option<&mut Span> {
-        trace.iter_mut().find(|span| span.name == "aws.lambda")
+        trace.iter_mut().find(|span| {
+            span.name == "aws.lambda" && span.resource != crate::traces::INVOCATION_SPAN_RESOURCE
+        })
     }
 
     /// Processes an intercepted [`Span`].
@@ -812,4 +820,41 @@ mod tests {
             result
         );
     }
+
+    #[test]
+    fn service_entry_span_mut_skips_placeholder_lambda_span() {
+        let mut trace = vec![
+            Span {
+                name: "aws.lambda".into(),
+                resource: crate::traces::INVOCATION_SPAN_RESOURCE.into(),
+                span_id: 1,
+                ..Default::default()
+            },
+            Span {
+                name: "aws.lambda".into(),
+                resource: "real.lambda.invocation".into(),
+                span_id: 2,
+                ..Default::default()
+            },
+        ];
+
+        let selected = Processor::service_entry_span_mut(&mut trace)
+            .expect("expected non-placeholder aws.lambda span");
+
+        assert_eq!(selected.name, "aws.lambda");
+        assert_ne!(selected.resource, crate::traces::INVOCATION_SPAN_RESOURCE);
+        assert_eq!(selected.span_id, 2);
+    }
+
+    #[test]
+    fn service_entry_span_mut_returns_none_for_only_placeholder() {
+        let mut trace = vec![Span {
+            name: "aws.lambda".into(),
+            resource: crate::traces::INVOCATION_SPAN_RESOURCE.into(),
+            span_id: 1,
+            ..Default::default()
+        }];
+
+        assert!(Processor::service_entry_span_mut(&mut trace).is_none());
+    }
 }

bottlecap/src/lifecycle/invocation/processor.rs

@@ -615,7 +615,6 @@ impl Processor {
             );
         }
 
-        // todo(duncanista): Add missing metric tags for ASM
         // Add dynamic and trigger tags
         context
             .invocation_span
@@ -626,6 +625,19 @@ impl Processor {
             context.invocation_span.meta.extend(trigger_tags);
         }
 
+        // Ensure _dd.appsec.enabled is present on the invocation span when AAP is enabled.
+        // complete_inferred_spans (called below) propagates this metric from the invocation
+        // span to the inferred trigger span. AppSec's process_span will set it again from the
+        // security context when it runs, but this baseline guarantees the tag is always present
+        // even when the context cannot be found at flush time.
+        if self.config.serverless_appsec_enabled {
+            context
+                .invocation_span
+                .metrics
+                .entry("_dd.appsec.enabled".to_string())
+                .or_insert(1.0);
+        }
+
         self.inferrer
             .complete_inferred_spans(&context.invocation_span);
 
@@ -1405,6 +1417,7 @@ impl Processor {
         execution_id: &str,
         execution_name: &str,
         first_invocation: Option<bool>,
+        execution_status: Option<String>,
     ) {
         if let Err(e) = self
             .durable_context_tx
@@ -1413,6 +1426,7 @@ impl Processor {
                 execution_id: execution_id.to_owned(),
                 execution_name: execution_name.to_owned(),
                 first_invocation,
+                execution_status,
             })
             .await
         {
@@ -2332,4 +2346,103 @@ mod tests {
             "no contexts should be ready to send yet"
         );
     }
+
+    fn setup_appsec() -> Processor {
+        let aws_config = Arc::new(AwsConfig {
+            region: "us-east-1".into(),
+            aws_lwa_proxy_lambda_runtime_api: Some("***".into()),
+            function_name: "test-function".into(),
+            sandbox_init_time: Instant::now(),
+            runtime_api: "***".into(),
+            exec_wrapper: None,
+            initialization_type: "on-demand".into(),
+        });
+        let config = Arc::new(config::Config {
+            service: Some("test-service".to_string()),
+            serverless_appsec_enabled: true,
+            ..config::Config::default()
+        });
+        let tags_provider = Arc::new(provider::Provider::new(
+            Arc::clone(&config),
+            LAMBDA_RUNTIME_SLUG.to_string(),
+            &HashMap::from([("function_arn".to_string(), "test-arn".to_string())]),
+        ));
+        let (service, handle) =
+            dogstatsd::aggregator::AggregatorService::new(dogstatsd::metric::EMPTY_TAGS, 1024)
+                .expect("failed to create aggregator service");
+        tokio::spawn(service.run());
+        let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(&config)));
+        let (durable_context_tx, _) = tokio::sync::mpsc::channel(1);
+        Processor::new(
+            tags_provider,
+            config,
+            aws_config,
+            handle,
+            propagator,
+            durable_context_tx,
+        )
+    }
+
+    #[tokio::test]
+    async fn enrich_ctx_sets_appsec_enabled_when_aap_enabled() {
+        let mut p = setup_appsec();
+        let request_id = String::from("req-appsec");
+        p.on_invoke_event(request_id.clone());
+        p.on_platform_start(request_id.clone(), chrono::Utc::now());
+
+        let ctx = p
+            .enrich_ctx_at_platform_done(&request_id, Status::Success)
+            .expect("context must be present");
+
+        assert_eq!(
+            ctx.invocation_span.metrics.get("_dd.appsec.enabled"),
+            Some(&1.0),
+            "_dd.appsec.enabled must be 1.0 when AAP is enabled"
+        );
+    }
+
+    #[tokio::test]
+    async fn enrich_ctx_does_not_set_appsec_enabled_when_aap_disabled() {
+        let mut p = setup();
+        let request_id = String::from("req-no-appsec");
+        p.on_invoke_event(request_id.clone());
+        p.on_platform_start(request_id.clone(), chrono::Utc::now());
+
+        let ctx = p
+            .enrich_ctx_at_platform_done(&request_id, Status::Success)
+            .expect("context must be present");
+
+        assert!(
+            !ctx.invocation_span
+                .metrics
+                .contains_key("_dd.appsec.enabled"),
+            "_dd.appsec.enabled must not be set when AAP is disabled"
+        );
+    }
+
+    #[tokio::test]
+    async fn enrich_ctx_does_not_override_existing_appsec_enabled() {
+        let mut p = setup_appsec();
+        let request_id = String::from("req-appsec-preset");
+        p.on_invoke_event(request_id.clone());
+        p.on_platform_start(request_id.clone(), chrono::Utc::now());
+
+        // Pre-set a different value to verify or_insert does not overwrite it.
+        p.context_buffer
+            .get_mut(&request_id)
+            .expect("context must exist")
+            .invocation_span
+            .metrics
+            .insert("_dd.appsec.enabled".to_string(), 0.0);
+
+        let ctx = p
+            .enrich_ctx_at_platform_done(&request_id, Status::Success)
+            .expect("context must be present");
+
+        assert_eq!(
+            ctx.invocation_span.metrics.get("_dd.appsec.enabled"),
+            Some(&0.0),
+            "pre-existing _dd.appsec.enabled value must not be overwritten"
+        );
+    }
 }

bottlecap/src/lifecycle/invocation/processor_service.rs

@@ -115,6 +115,7 @@ pub enum ProcessorCommand {
         execution_id: String,
         execution_name: String,
         first_invocation: Option<bool>,
+        execution_status: Option<String>,
     },
     OnOutOfMemoryError {
         timestamp: i64,
@@ -391,13 +392,15 @@ impl InvocationProcessorHandle {
         execution_id: String,
         execution_name: String,
         first_invocation: Option<bool>,
+        execution_status: Option<String>,
     ) -> Result<(), mpsc::error::SendError<ProcessorCommand>> {
         self.sender
             .send(ProcessorCommand::ForwardDurableContext {
                 request_id,
                 execution_id,
                 execution_name,
                 first_invocation,
+                execution_status,
             })
             .await
     }
@@ -617,13 +620,15 @@ impl InvocationProcessorService {
                     execution_id,
                     execution_name,
                     first_invocation,
+                    execution_status,
                 } => {
                     self.processor
                         .forward_durable_context(
                             &request_id,
                             &execution_id,
                             &execution_name,
                             first_invocation,
+                            execution_status,
                         )
                         .await;
                 }

bottlecap/src/logs/lambda/mod.rs

@@ -9,6 +9,7 @@ pub struct DurableContextUpdate {
     pub execution_id: String,
     pub execution_name: String,
     pub first_invocation: Option<bool>,
+    pub execution_status: Option<String>,
 }
 
 /// Durable execution context stored per `request_id` in `LambdaProcessor::durable_context_map`.
@@ -17,6 +18,7 @@ pub struct DurableExecutionContext {
     pub execution_id: String,
     pub execution_name: String,
     pub first_invocation: Option<bool>,
+    pub execution_status: Option<String>,
 }
 
 ///
@@ -66,6 +68,11 @@ pub struct Lambda {
         skip_serializing_if = "Option::is_none"
     )]
     pub first_invocation: Option<bool>,
+    #[serde(
+        rename = "durable_function.execution_status",
+        skip_serializing_if = "Option::is_none"
+    )]
+    pub durable_execution_status: Option<String>,
 }
 
 impl Message {

bottlecap/src/logs/lambda/processor.rs

@@ -546,10 +546,11 @@ impl LambdaProcessor {
     /// `request_id`.
     pub fn insert_to_durable_context_map(
         &mut self,
-        request_id: &str,               // key
-        execution_id: &str,             // value
-        execution_name: &str,           // value
-        first_invocation: Option<bool>, // value
+        request_id: &str,                 // key
+        execution_id: &str,               // value
+        execution_name: &str,             // value
+        first_invocation: Option<bool>,   // value
+        execution_status: Option<String>, // value
     ) {
         if self.durable_context_map.contains_key(request_id) {
             error!("LOGS | insert_to_durable_context_map: request_id={request_id} already in map");
@@ -567,6 +568,7 @@ impl LambdaProcessor {
                 execution_id: execution_id.to_string(),
                 execution_name: execution_name.to_string(),
                 first_invocation,
+                execution_status,
             },
         );
         self.drain_held_logs_for_request_id(request_id);
@@ -660,6 +662,12 @@ impl LambdaProcessor {
         if is_platform_log(&log.message.message) {
             log.message.lambda.first_invocation = ctx.first_invocation;
         }
+        if log.message.message.starts_with("END RequestId:") {
+            log.message
+                .lambda
+                .durable_execution_status
+                .clone_from(&ctx.execution_status);
+        }
         if let Ok(s) = serde_json::to_string(&log) {
             // explicitly drop log so we don't accidentally re-use it and push
             // duplicate logs to the aggregator
@@ -2578,4 +2586,41 @@ mod tests {
         let batches = aggregator_handle.get_batches().await.unwrap();
         assert!(batches.is_empty());
     }
+
+    #[tokio::test]
+    async fn test_execution_status_on_end_log() {
+        for (execution_status, expected) in [
+            (Some("SUCCEEDED"), serde_json::json!("SUCCEEDED")),
+            (None, serde_json::Value::Null),
+        ] {
+            let mut processor = make_processor_for_durable_arn_tests();
+            processor.is_durable_function = Some(true);
+            processor.invocation_context.request_id = "req-end".to_string();
+            processor.insert_to_durable_context_map(
+                "req-end",
+                "exec-id-123",
+                "exec-name-abc",
+                Some(false),
+                execution_status.map(str::to_string),
+            );
+            let event = TelemetryEvent {
+                time: Utc.with_ymd_and_hms(2023, 1, 7, 3, 23, 47).unwrap(),
+                record: TelemetryRecord::PlatformRuntimeDone {
+                    request_id: "req-end".to_string(),
+                    status: Status::Success,
+                    error_type: None,
+                    metrics: None,
+                },
+            };
+            let (aggregator_service, aggregator_handle) = AggregatorService::default();
+            tokio::spawn(async move { aggregator_service.run().await });
+            processor.process(event, &aggregator_handle).await;
+            let batches = aggregator_handle.get_batches().await.unwrap();
+            let logs: Vec<serde_json::Value> = serde_json::from_slice(&batches[0]).unwrap();
+            assert_eq!(
+                logs[0]["message"]["lambda"]["durable_function.execution_status"],
+                expected
+            );
+        }
+    }
 }

bottlecap/src/logs/processor.rs

@@ -53,6 +53,7 @@ impl LogsProcessor {
             &update.execution_id,
             &update.execution_name,
             update.first_invocation,
+            update.execution_status,
         );
         let ready_logs = self.take_ready_logs();
         if !ready_logs.is_empty()
@@ -68,6 +69,7 @@ impl LogsProcessor {
         execution_id: &str,
         execution_name: &str,
         first_invocation: Option<bool>,
+        execution_status: Option<String>,
     ) {
         match self {
             LogsProcessor::Lambda(p) => {
@@ -76,6 +78,7 @@ impl LogsProcessor {
                     execution_id,
                     execution_name,
                     first_invocation,
+                    execution_status,
                 );
             }
         }