From e34f76764ac73a924389826c4a24544e15c3ef9e Mon Sep 17 00:00:00 2001
From: Nicolas Moutschen <nicolas.moutschen@gmail.com>
Date: Sun, 21 Nov 2021 19:31:20 +0100
Subject: [PATCH] fix: add support for null/invalid type for headers

---
 lambda-http/src/request.rs | 62 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 61 insertions(+), 1 deletion(-)

diff --git a/lambda-http/src/request.rs b/lambda-http/src/request.rs
index e5e3cd8d..88740dd1 100644
--- a/lambda-http/src/request.rs
+++ b/lambda-http/src/request.rs
@@ -352,7 +352,7 @@ where
         }
     }
 
-    deserializer.deserialize_map(HeaderVisitor)
+    Ok(deserializer.deserialize_map(HeaderVisitor).unwrap_or_default())
 }
 
 /// deserializes (json) null values to their default values
@@ -633,6 +633,17 @@ mod tests {
         let req = result.expect("failed to parse request");
         assert_eq!(req.method(), "GET");
         assert_eq!(req.uri(), "https://xxx.execute-api.us-east-1.amazonaws.com/");
+
+        // Ensure this is an APIGWv2 request
+        let req_context = req.request_context();
+        assert!(
+            match req_context {
+                RequestContext::ApiGatewayV2(_) => true,
+                _ => false,
+            },
+            "expected ApiGatewayV2 context, got {:?}",
+            req_context
+        );
     }
 
     #[test]
@@ -657,6 +668,17 @@ mod tests {
         assert_eq!(req.method(), "POST");
         assert_eq!(req.uri(), "https://id.execute-api.us-east-1.amazonaws.com/my/path?parameter1=value1&parameter1=value2&parameter2=value");
         assert_eq!(cookie_header, Ok("cookie1=value1;cookie2=value2"));
+
+        // Ensure this is an APIGWv2 request
+        let req_context = req.request_context();
+        assert!(
+            match req_context {
+                RequestContext::ApiGatewayV2(_) => true,
+                _ => false,
+            },
+            "expected ApiGatewayV2 context, got {:?}",
+            req_context
+        );
     }
 
     #[test]
@@ -678,6 +700,17 @@ mod tests {
             req.uri(),
             "https://wt6mne2s9k.execute-api.us-west-2.amazonaws.com/test/hello"
         );
+
+        // Ensure this is an APIGW request
+        let req_context = req.request_context();
+        assert!(
+            match req_context {
+                RequestContext::ApiGateway(_) => true,
+                _ => false,
+            },
+            "expected ApiGateway context, got {:?}",
+            req_context
+        );
     }
 
     #[test]
@@ -695,6 +728,17 @@ mod tests {
         let req = result.expect("failed to parse request");
         assert_eq!(req.method(), "GET");
         assert_eq!(req.uri(), "https://lambda-846800462-us-east-2.elb.amazonaws.com/");
+
+        // Ensure this is an ALB request
+        let req_context = req.request_context();
+        assert!(
+            match req_context {
+                RequestContext::Alb(_) => true,
+                _ => false,
+            },
+            "expected Alb context, got {:?}",
+            req_context
+        );
     }
 
     #[test]
@@ -808,4 +852,20 @@ mod tests {
             Test { foo: HashMap::new() }
         )
     }
+
+    #[test]
+    fn deserialize_null_headers() {
+        #[derive(Debug, PartialEq, Deserialize)]
+        struct Test {
+            #[serde(deserialize_with = "deserialize_headers")]
+            headers: http::HeaderMap,
+        }
+
+        assert_eq!(
+            serde_json::from_str::<Test>(r#"{"headers":null}"#).expect("failed to deserialize"),
+            Test {
+                headers: http::HeaderMap::new()
+            }
+        )
+    }
 }