From a4a97fc2141fb71314fc7742d7dcf17c2486a1c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?jordan=20gonz=C3=A1lez?=
 <30836115+duncanista@users.noreply.github.com>
Date: Thu, 24 Oct 2024 16:46:39 -0400
Subject: [PATCH] generate a random id with `OsRng` instead of `thread_rng` on
 `SnapStart` lambdas

---
 bottlecap/src/lifecycle/invocation/span_inferrer.rs | 8 ++++++--
 bottlecap/src/tags/lambda/tags.rs                   | 4 +++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/bottlecap/src/lifecycle/invocation/span_inferrer.rs b/bottlecap/src/lifecycle/invocation/span_inferrer.rs
index 6141d8dd..70f86d7d 100644
--- 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
index 7e6ce54e..ad1fb0c8 100644
--- 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";