generate a random id with OsRng instead of thread_rng on `SnapSta…

…rt` lambdas (#427)
DataDog · Oct 24, 2024 · 5bf6708 · 5bf6708
1 parent ada1414
commit 5bf6708
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/bottlecap/src/lifecycle/invocation/span_inferrer.rs b/bottlecap/src/lifecycle/invocation/span_inferrer.rs
@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use datadog_trace_protobuf::pb::Span;
-use rand::Rng;
+use rand::{rngs::OsRng, Rng, RngCore};
 use serde_json::Value;
 use tracing::debug;
 
@@ -11,6 +11,7 @@ use crate::lifecycle::invocation::triggers::{
     api_gateway_http_event::APIGatewayHttpEvent, api_gateway_rest_event::APIGatewayRestEvent,
     Trigger,
 };
+use crate::tags::lambda::tags::{INIT_TYPE, SNAP_START_VALUE};
 
 const FUNCTION_TRIGGER_EVENT_SOURCE_TAG: &str = "function_trigger.event_source";
 const FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG: &str = "function_trigger.event_source_arn";
@@ -133,7 +134,10 @@ impl SpanInferrer {
     }
 
     fn generate_span_id() -> u64 {
-        // todo: secure random id with OsRng for SnapStart
+        if std::env::var(INIT_TYPE).map_or(false, |it| it == SNAP_START_VALUE) {
+            return OsRng.next_u64();
+        }
+
         let mut rng = rand::thread_rng();
         rng.gen()
     }

diff --git a/bottlecap/src/tags/lambda/tags.rs b/bottlecap/src/tags/lambda/tags.rs
@@ -10,8 +10,10 @@ use tracing::debug;
 const QUALIFIER_ENV_VAR: &str = "AWS_LAMBDA_FUNCTION_VERSION";
 const RUNTIME_VAR: &str = "AWS_EXECUTION_ENV";
 const MEMORY_SIZE_VAR: &str = "AWS_LAMBDA_FUNCTION_MEMORY_SIZE";
-const INIT_TYPE: &str = "AWS_LAMBDA_INITIALIZATION_TYPE";
+pub const INIT_TYPE: &str = "AWS_LAMBDA_INITIALIZATION_TYPE";
 const INIT_TYPE_KEY: &str = "init_type";
+// Value for INIT_TYPE when the function is using SnapStart
+pub const SNAP_START_VALUE: &str = "snap-start";
 
 // FunctionARNKey is the tag key for a function's arn
 pub const FUNCTION_ARN_KEY: &str = "function_arn";