add x-ray trace id header in lambda-http (awslabs#557)

twu-AWS · web-flow · commit 1fc2ccbca3f4 · 2022-10-27T20:37:06.000-07:00
* fix for Github issue awslabs#552 * Github issue awslabs#552 fmt check * github issue awslabs#552 apply clippy fix * remove unnessecory comment
diff --git a/lambda-http/src/request.rs b/lambda-http/src/request.rs
@@ -16,12 +16,12 @@ use aws_lambda_events::apigw::{ApiGatewayV2httpRequest, ApiGatewayV2httpRequestC
 use aws_lambda_events::apigw::{ApiGatewayWebsocketProxyRequest, ApiGatewayWebsocketProxyRequestContext};
 use aws_lambda_events::{encodings::Body, query_map::QueryMap};
 use http::header::HeaderName;
-use http::HeaderMap;
+use http::{HeaderMap, HeaderValue};
 use serde::Deserialize;
 use serde_json::error::Error as JsonError;
 use std::future::Future;
 use std::pin::Pin;
-use std::{io::Read, mem};
+use std::{env, io::Read, mem};
 use url::Url;
 
 /// Internal representation of an Lambda http event from
@@ -119,8 +119,9 @@ fn into_api_gateway_v2_request(ag: ApiGatewayV2httpRequest) -> http::Request<Bod
         .extension(RequestContext::ApiGatewayV2(ag.request_context));
 
     let mut headers = ag.headers;
+    update_xray_trace_id_header(&mut headers);
     if let Some(cookies) = ag.cookies {
-        if let Ok(header_value) = http::header::HeaderValue::from_str(&cookies.join(";")) {
+        if let Ok(header_value) = HeaderValue::from_str(&cookies.join(";")) {
             headers.append(http::header::COOKIE, header_value);
         }
     }
@@ -142,6 +143,13 @@ fn into_api_gateway_v2_request(ag: ApiGatewayV2httpRequest) -> http::Request<Bod
     req
 }
 
+fn update_xray_trace_id_header(headers: &mut HeaderMap) {
+    if let Ok(xray_trace_id) = env::var("_X_AMZN_TRACE_ID") {
+        if let Ok(header_value) = HeaderValue::from_str(&xray_trace_id) {
+            headers.insert(HeaderName::from_static("x-amzn-trace-id"), header_value);
+        }
+    }
+}
 #[cfg(feature = "apigw_rest")]
 fn into_proxy_request(ag: ApiGatewayProxyRequest) -> http::Request<Body> {
     let http_method = ag.http_method;
@@ -179,6 +187,7 @@ fn into_proxy_request(ag: ApiGatewayProxyRequest) -> http::Request<Body> {
     // multi-value_headers our cannoncial source of request headers
     let mut headers = ag.multi_value_headers;
     headers.extend(ag.headers);
+    update_xray_trace_id_header(&mut headers);
 
     let base64 = ag.is_base64_encoded.unwrap_or_default();
     let mut req = builder
@@ -229,6 +238,7 @@ fn into_alb_request(alb: AlbTargetGroupRequest) -> http::Request<Body> {
     // multi-value_headers our cannoncial source of request headers
     let mut headers = alb.multi_value_headers;
     headers.extend(alb.headers);
+    update_xray_trace_id_header(&mut headers);
 
     let base64 = alb.is_base64_encoded;
 
@@ -291,6 +301,7 @@ fn into_websocket_request(ag: ApiGatewayWebsocketProxyRequest) -> http::Request<
     // multi-value_headers our cannoncial source of request headers
     let mut headers = ag.multi_value_headers;
     headers.extend(ag.headers);
+    update_xray_trace_id_header(&mut headers);
 
     let base64 = ag.is_base64_encoded.unwrap_or_default();
     let mut req = builder
diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs
@@ -134,8 +134,10 @@ where
             let request_id = &ctx.request_id.clone();
 
             let xray_trace_id = &ctx.xray_trace_id.clone();
-            env::set_var("_X_AMZN_TRACE_ID", xray_trace_id);
-
+            match xray_trace_id {
+                Some(trace_id) => env::set_var("_X_AMZN_TRACE_ID", trace_id),
+                None => env::remove_var("_X_AMZN_TRACE_ID"),
+            }
             let body = match serde_json::from_slice(&body) {
                 Ok(body) => body,
                 Err(err) => {
diff --git a/lambda-runtime/src/types.rs b/lambda-runtime/src/types.rs
@@ -90,7 +90,7 @@ pub struct Context {
     /// The ARN of the Lambda function being invoked.
     pub invoked_function_arn: String,
     /// The X-Ray trace ID for the current invocation.
-    pub xray_trace_id: String,
+    pub xray_trace_id: Option<String>,
     /// The client context object sent by the AWS mobile SDK. This field is
     /// empty unless the function is invoked using an AWS mobile SDK.
     pub client_context: Option<ClientContext>,
@@ -139,9 +139,7 @@ impl TryFrom<HeaderMap> for Context {
                 .to_owned(),
             xray_trace_id: headers
                 .get("lambda-runtime-trace-id")
-                .unwrap_or(&HeaderValue::from_static(""))
-                .to_str()?
-                .to_owned(),
+                .map(|v| String::from_utf8_lossy(v.as_bytes()).to_string()),
             client_context,
             identity,
             ..Default::default()
