Panic when X-Ray header not present

lambda-runtime-client/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lambda_runtime_client"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["Stefano Buliani", "David Barsky"]
 edition = "2018"
 description = "Client SDK for AWS Lambda's runtime APIs"
@@ -24,4 +24,7 @@ serde_json = "^1"
 serde_derive = "^1"
 log = "0.4"
 lambda_runtime_errors = { path = "../lambda-runtime-errors", version = "^0.1" }
-failure = "^0.1"
+failure = "^0.1"
+
+[dev-dependencies]
+chrono = "^0.4"

lambda-runtime-client/src/client.rs

@@ -109,7 +109,7 @@ pub struct EventContext {
     /// The AWS request ID generated by the Lambda service.
     pub aws_request_id: String,
     /// The X-Ray trace ID for the current invocation.
-    pub xray_trace_id: String,
+    pub xray_trace_id: Option<String>,
     /// The execution deadline for the current invocation in milliseconds.
     pub deadline: i64,
     /// The client context object sent by the AWS mobile SDK. This field is
@@ -404,7 +404,17 @@ impl<'ev> RuntimeClient {
             headers.get(LambdaHeaders::FunctionArn.as_str()),
             &LambdaHeaders::FunctionArn,
         )?;
-        let xray_trace_id = header_string(headers.get(LambdaHeaders::TraceId.as_str()), &LambdaHeaders::TraceId)?;
+        let xray_trace_id = match headers.get(LambdaHeaders::TraceId.as_str()) {
+            Some(trace_id) => match trace_id.to_str() {
+                Ok(trace_str) => Some(trace_str.to_owned()),
+                Err(e) => {
+                    // we do not fail on this error.
+                    error!("Could not parse X-Ray trace id as string: {}", e);
+                    None
+                }
+            },
+            None => None
+        };
         let deadline = header_string(headers.get(LambdaHeaders::Deadline.as_str()), &LambdaHeaders::Deadline)?
             .parse::<i64>()
             .context(ApiErrorKind::Recoverable(
@@ -461,3 +471,42 @@ fn header_string(value: Option<&HeaderValue>, header_type: &LambdaHeaders) -> Re
         }
     }
 }
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use super::*;
+    use chrono::{Utc, Duration};
+
+    fn get_headers() -> HeaderMap<HeaderValue> {
+        let mut headers: HeaderMap<HeaderValue> = HeaderMap::new();
+        headers.insert(LambdaHeaders::RequestId.as_str(), HeaderValue::from_str("req_id").unwrap());
+        headers.insert(LambdaHeaders::FunctionArn.as_str(), HeaderValue::from_str("func_arn").unwrap());
+        headers.insert(LambdaHeaders::TraceId.as_str(), HeaderValue::from_str("trace").unwrap());
+        let deadline = Utc::now() + Duration::seconds(10);
+        headers.insert(LambdaHeaders::Deadline.as_str(), HeaderValue::from_str(&deadline.timestamp_millis().to_string()).unwrap());
+        //headers.insert(LambdaHeaders::ClientContext.as_str(), HeaderValue::from_str("{}").unwrap());
+        //headers.insert(LambdaHeaders::CognitoIdentity.as_str(), HeaderValue::from_str("{}").unwrap());
+        headers
+    }
+
+    #[test]
+    fn get_event_context_with_empty_trace_id() {
+        let client = RuntimeClient::new("localhost:8081", None, None).expect("Could not initialize runtime client");
+        let mut headers = get_headers();
+        headers.remove(LambdaHeaders::TraceId.as_str());
+        let headers_result = client.get_event_context(&headers);
+        assert_eq!(false, headers_result.is_err());
+        let ok_result = headers_result.unwrap();
+        assert_eq!(None, ok_result.xray_trace_id);
+        assert_eq!("req_id", ok_result.aws_request_id);
+    }
+
+    #[test]
+    fn get_event_context_populates_trace_id_when_present() {
+        let client = RuntimeClient::new("localhost:8081", None, None).expect("Could not initialize runtime client");
+        let headers = get_headers();
+        let headers_result = client.get_event_context(&headers);
+        assert_eq!(false, headers_result.is_err());
+        assert_eq!(Some("trace".to_owned()), headers_result.unwrap().xray_trace_id);
+    }
+}

lambda-runtime-core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lambda_runtime_core"
-version = "0.1.0"
+version = "0.1.1"
 authors = ["Stefano Buliani", "David Barsky"]
 description = "Rust runtime for AWS Lambda"
 keywords = ["AWS", "Lambda", "Runtime", "Rust"]

lambda-runtime-core/src/context.rs

@@ -32,7 +32,7 @@ pub struct Context {
     /// by the Lambda Runtime APIs as a header. Developers can use this value
     /// with the AWS SDK to create new, custom sub-segments to the current
     /// invocation.
-    pub xray_trace_id: String,
+    pub xray_trace_id: Option<String>,
     /// The name of the CloudWatch log stream for the current execution
     /// environment. This value is extracted from the `AWS_LAMBDA_LOG_STREAM_NAME`
     /// environment variable set by the Lambda service.
@@ -72,7 +72,7 @@ impl Context {
     /// A new, populated `Context` object.
     pub(super) fn new(local_settings: lambda_env::FunctionSettings) -> Context {
         Context {
-            xray_trace_id: String::from(""),
+            xray_trace_id: None,
             memory_limit_in_mb: local_settings.memory_size,
             function_name: local_settings.function_name,
             function_version: local_settings.version,
@@ -107,7 +107,7 @@ pub(crate) mod tests {
             function_version: "$LATEST".to_string(),
             invoked_function_arn: "arn:aws:lambda".to_string(),
             aws_request_id: "123".to_string(),
-            xray_trace_id: "123".to_string(),
+            xray_trace_id: Some("123".to_string()),
             log_stream_name: "logStream".to_string(),
             log_group_name: "logGroup".to_string(),
             client_context: Option::default(),

lambda-runtime/src/lib.rs

@@ -142,7 +142,7 @@ pub(crate) mod tests {
             function_version: "$LATEST".to_string(),
             invoked_function_arn: "arn:aws:lambda".to_string(),
             aws_request_id: "123".to_string(),
-            xray_trace_id: "123".to_string(),
+            xray_trace_id: Some("123".to_string()),
             log_stream_name: "logStream".to_string(),
             log_group_name: "logGroup".to_string(),
             client_context: Option::default(),