fix(lambda): intercept raw JSON before deserialization to preserve trace context

lambda_runtime deserializes the event into E before wrap_handler sees it. If the user's
type omits fields like headers, re-serializing with to_value loses them permanently. For
API Gateway triggers this means trace context in event.headers is silently dropped.

wrap_handler now receives LambdaEvent<Value> so lambda_runtime preserves the full event.
The span inferrer runs on the complete JSON, then the library deserializes Value to E
before calling the user handler. The E: Serialize bound is replaced with E: DeserializeOwned,
which is the constraint lambda_runtime itself requires.

Author: Dogbu-cyber

integrations/aws/datadog-lambda-rs/src/invocation.rs

@@ -8,7 +8,7 @@ use lambda_runtime::LambdaEvent;
 use opentelemetry::trace::{FutureExt, SpanKind, TraceContextExt, Tracer, TracerProvider as _};
 use opentelemetry::{Context, KeyValue};
 use opentelemetry_sdk::trace::{SdkTracer, SdkTracerProvider};
-use serde::Serialize;
+use serde_json::Value;
 use std::future::Future;
 use std::sync::atomic::{AtomicBool, Ordering};
 
@@ -40,20 +40,16 @@ pub(crate) struct InvocationScope {
     pub invocation_cx: Context,
 }
 
-pub(crate) fn start_invocation<E>(
-    event: &LambdaEvent<E>,
+pub(crate) fn start_invocation(
+    event: &LambdaEvent<Value>,
     provider: &SdkTracerProvider,
     _config: &Config,
-) -> InvocationScope
-where
-    E: Serialize,
-{
+) -> InvocationScope {
     let is_cold = detect_cold_start();
     let tracer = provider.tracer("datadog-lambda-rs");
-    let payload_value = serde_json::to_value(&event.payload).unwrap_or_default();
 
     let inferrer = SpanInferrer::new(&tracer);
-    let result = inferrer.infer(&payload_value);
+    let result = inferrer.infer(&event.payload);
 
     let lambda_span = LambdaSpan {
         function_name: event.context.env_config.function_name.clone(),

integrations/aws/datadog-lambda-rs/src/lib.rs

@@ -10,7 +10,9 @@ use lambda_runtime::{
     LambdaEvent,
 };
 use opentelemetry_sdk::trace::SdkTracerProvider;
-use serde::Serialize;
+use serde::de::DeserializeOwned;
+use serde_json::Value;
+use std::marker::PhantomData;
 use std::sync::Arc;
 use std::task::{Context as TaskContext, Poll};
 
@@ -25,11 +27,11 @@ pub struct Config {}
 pub fn wrap_handler<F, Fut, E, R>(
     handler: F,
     provider: SdkTracerProvider,
-) -> impl Fn(LambdaEvent<E>) -> BoxFuture<Result<R, HandlerError>>
+) -> impl Fn(LambdaEvent<Value>) -> BoxFuture<Result<R, HandlerError>>
 where
     F: Fn(LambdaEvent<E>) -> Fut + Send + Sync + 'static,
     Fut: std::future::Future<Output = Result<R, HandlerError>> + Send + 'static,
-    E: Serialize + Send + Sync + 'static,
+    E: DeserializeOwned + Send + Sync + 'static,
     R: Send + 'static,
 {
     wrap_handler_with_config(handler, provider, Config::default())
@@ -40,21 +42,26 @@ pub fn wrap_handler_with_config<F, Fut, E, R>(
     handler: F,
     provider: SdkTracerProvider,
     config: Config,
-) -> impl Fn(LambdaEvent<E>) -> BoxFuture<Result<R, HandlerError>>
+) -> impl Fn(LambdaEvent<Value>) -> BoxFuture<Result<R, HandlerError>>
 where
     F: Fn(LambdaEvent<E>) -> Fut + Send + Sync + 'static,
     Fut: std::future::Future<Output = Result<R, HandlerError>> + Send + 'static,
-    E: Serialize + Send + Sync + 'static,
+    E: DeserializeOwned + Send + Sync + 'static,
     R: Send + 'static,
 {
     let handler = Arc::new(handler);
     let config = Arc::new(config);
-    move |event: LambdaEvent<E>| {
+    move |event: LambdaEvent<Value>| {
         let handler = Arc::clone(&handler);
         let provider = provider.clone();
         let config = Arc::clone(&config);
         let scope = start_invocation(&event, &provider, &config);
-        let fut = handler(event);
+        let typed_payload = match serde_json::from_value::<E>(event.payload) {
+            Ok(p) => p,
+            Err(e) => return Box::pin(async move { Err(e.into()) }),
+        };
+        let typed_event = LambdaEvent::new(typed_payload, event.context);
+        let fut = handler(typed_event);
         Box::pin(async move { run_in_invocation_scope(scope, provider, fut).await })
     }
 }
@@ -89,15 +96,17 @@ impl<S> Layer<S> for DatadogLambdaLayer {
             inner,
             provider: self.provider.clone(),
             config: Arc::clone(&self.config),
+            _phantom: PhantomData,
         }
     }
 }
 
 #[doc(hidden)]
-pub struct DatadogLambdaService<S> {
+pub struct DatadogLambdaService<S, E = Value> {
     inner: S,
     provider: SdkTracerProvider,
     config: Arc<Config>,
+    _phantom: PhantomData<fn(LambdaEvent<E>)>,
 }
 
 #[cfg(test)]
@@ -119,27 +128,32 @@ mod tests {
     }
 }
 
-impl<S, E, R, Err> Service<LambdaEvent<E>> for DatadogLambdaService<S>
+impl<S, E> Service<LambdaEvent<Value>> for DatadogLambdaService<S, E>
 where
-    S: Service<LambdaEvent<E>, Response = R, Error = Err> + Send + 'static,
+    S: Service<LambdaEvent<E>> + Send + 'static,
     S::Future: Send + 'static,
-    E: Serialize + Send + Sync + 'static,
-    R: Send + 'static,
-    Err: std::fmt::Display + Send + 'static,
+    S::Response: Send + 'static,
+    S::Error: From<serde_json::Error> + std::fmt::Display + Send + 'static,
+    E: DeserializeOwned + Send + Sync + 'static,
 {
-    type Response = R;
-    type Error = Err;
-    type Future = BoxFuture<Result<R, Err>>;
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = BoxFuture<Result<S::Response, S::Error>>;
 
     fn poll_ready(&mut self, cx: &mut TaskContext<'_>) -> Poll<Result<(), Self::Error>> {
         self.inner.poll_ready(cx)
     }
 
-    fn call(&mut self, event: LambdaEvent<E>) -> Self::Future {
+    fn call(&mut self, event: LambdaEvent<Value>) -> Self::Future {
         let provider = self.provider.clone();
         let config = Arc::clone(&self.config);
         let scope = start_invocation(&event, &provider, &config);
-        let fut = self.inner.call(event);
+        let typed_payload = match serde_json::from_value::<E>(event.payload) {
+            Ok(p) => p,
+            Err(e) => return Box::pin(async move { Err(e.into()) }),
+        };
+        let typed_event = LambdaEvent::new(typed_payload, event.context);
+        let fut = self.inner.call(typed_event);
         Box::pin(async move { run_in_invocation_scope(scope, provider, fut).await })
     }
 }