impl multi_value_headers (aws#58)

softprops · davidbarsky · commit 26d0ece3355d · 2018-12-25T23:31:51.000-05:00
* impl multi_value_headers

* remove redundant closure

* impl multi_value_headers in response

* fmt

* unprefix

* missing import in test

* rustfmt

* improve test assertion msg

* rustfmt
diff --git a/lambda-http/src/request.rs b/lambda-http/src/request.rs
@@ -5,7 +5,7 @@ use std::{borrow::Cow, collections::HashMap, fmt, mem};
 
 use http::{
     self,
-    header::{HeaderValue, HOST},
+    header::{HeaderName, HeaderValue, HOST},
     HeaderMap, Method, Request as HttpRequest,
 };
 use serde::{
@@ -31,6 +31,8 @@ pub(crate) struct GatewayRequest<'a> {
     pub(crate) http_method: Method,
     #[serde(deserialize_with = "deserialize_headers")]
     pub(crate) headers: HeaderMap<HeaderValue>,
+    #[serde(default, deserialize_with = "deserialize_multi_value_headers")]
+    pub(crate) multi_value_headers: HeaderMap<HeaderValue>,
     #[serde(deserialize_with = "nullable_default")]
     pub(crate) query_string_parameters: StrMap,
     #[serde(deserialize_with = "nullable_default")]
@@ -78,6 +80,7 @@ pub struct Identity {
     pub user_arn: Option<String>,
 }
 
+/// Deserialize a str into an http::Method
 fn deserialize_method<'de, D>(deserializer: D) -> Result<Method, D::Error>
 where
     D: Deserializer<'de>,
@@ -102,6 +105,49 @@ where
     deserializer.deserialize_str(MethodVisitor)
 }
 
+/// Deserialize a map of Cow<'_, str> => Vec<Cow<'_, str>> into an http::HeaderMap
+fn deserialize_multi_value_headers<'de, D>(deserializer: D) -> Result<HeaderMap<HeaderValue>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    struct HeaderVisitor;
+
+    impl<'de> Visitor<'de> for HeaderVisitor {
+        type Value = HeaderMap<HeaderValue>;
+
+        fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+            write!(formatter, "a multi valued HeaderMap<HeaderValue>")
+        }
+
+        fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
+        where
+            A: MapAccess<'de>,
+        {
+            let mut headers = map
+                .size_hint()
+                .map(HeaderMap::with_capacity)
+                .unwrap_or_else(HeaderMap::new);
+            while let Some((key, values)) = map.next_entry::<Cow<'_, str>, Vec<Cow<'_, str>>>()? {
+                // note the aws docs for multi value headers include an empty key. I'm not sure if this is a doc bug
+                // or not by the http crate doesn't handle it
+                // https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format
+                if !key.is_empty() {
+                    for value in values {
+                        let header_name = key.parse::<HeaderName>().map_err(A::Error::custom)?;
+                        let header_value =
+                            HeaderValue::from_shared(value.into_owned().into()).map_err(A::Error::custom)?;
+                        headers.append(header_name, header_value);
+                    }
+                }
+            }
+            Ok(headers)
+        }
+    }
+
+    deserializer.deserialize_map(HeaderVisitor)
+}
+
+/// Deserialize a map of Cow<'_, str> => Cow<'_, str> into an http::HeaderMap
 fn deserialize_headers<'de, D>(deserializer: D) -> Result<HeaderMap<HeaderValue>, D::Error>
 where
     D: Deserializer<'de>,
@@ -119,11 +165,13 @@ where
         where
             A: MapAccess<'de>,
         {
-            let mut headers = http::HeaderMap::new();
+            let mut headers = map
+                .size_hint()
+                .map(HeaderMap::with_capacity)
+                .unwrap_or_else(HeaderMap::new);
             while let Some((key, value)) = map.next_entry::<Cow<'_, str>, Cow<'_, str>>()? {
-                let header_name = key.parse::<http::header::HeaderName>().map_err(A::Error::custom)?;
-                let header_value =
-                    http::header::HeaderValue::from_shared(value.into_owned().into()).map_err(A::Error::custom)?;
+                let header_name = key.parse::<HeaderName>().map_err(A::Error::custom)?;
+                let header_value = HeaderValue::from_shared(value.into_owned().into()).map_err(A::Error::custom)?;
                 headers.append(header_name, header_value);
             }
             Ok(headers)
@@ -150,6 +198,7 @@ impl<'a> From<GatewayRequest<'a>> for HttpRequest<Body> {
             path,
             http_method,
             headers,
+            mut multi_value_headers,
             query_string_parameters,
             path_parameters,
             stage_variables,
@@ -191,8 +240,20 @@ impl<'a> From<GatewayRequest<'a>> for HttpRequest<Body> {
             })
             .expect("failed to build request");
 
+        // merge headers into multi_value_headers and make
+        // multi_value_headers our cannoncial source of request headers
+        for (key, value) in headers {
+            // see HeaderMap#into_iter() docs for cases when key element may be None
+            if let Some(first_key) = key {
+                // if it contains the key, avoid appending a duplicate value
+                if !multi_value_headers.contains_key(&first_key) {
+                    multi_value_headers.append(first_key, value);
+                }
+            }
+        }
+
         // no builder method that sets headers in batch
-        mem::replace(req.headers_mut(), headers);
+        mem::replace(req.headers_mut(), multi_value_headers);
 
         req
     }
@@ -224,8 +285,18 @@ mod tests {
     fn deserializes_request_events() {
         // from the docs
         // https://docs.aws.amazon.com/lambda/latest/dg/eventsources.html#eventsources-api-gateway-request
-        let input = include_str!("../tests/data/proxy_request.json");
-        assert!(serde_json::from_str::<GatewayRequest<'_>>(&input).is_ok())
+        let input = include_str!("../tests/data/apigw_proxy_request.json");
+        let result = serde_json::from_str::<GatewayRequest<'_>>(&input);
+        assert!(result.is_ok(), format!("event was not parsed as expected {:?}", result));
+    }
+
+    #[test]
+    fn deserialize_multi_value_events() {
+        // from docs
+        // https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format
+        let input = include_str!("../tests/data/apigw_multi_value_proxy_request.json");
+        let result = serde_json::from_str::<GatewayRequest<'_>>(&input);
+        assert!(result.is_ok(), format!("event was not parsed as expected {:?}", result));
     }
 
     #[test]
diff --git a/lambda-http/src/response.rs b/lambda-http/src/response.rs
@@ -4,7 +4,7 @@
 use std::ops::Not;
 
 use http::{
-    header::{HeaderMap, HeaderValue},
+    header::{HeaderMap, HeaderValue, CONTENT_TYPE},
     Response,
 };
 use serde::{
@@ -22,23 +22,47 @@ pub(crate) struct GatewayResponse {
     pub status_code: u16,
     #[serde(skip_serializing_if = "HeaderMap::is_empty", serialize_with = "serialize_headers")]
     pub headers: HeaderMap<HeaderValue>,
+    #[serde(
+        skip_serializing_if = "HeaderMap::is_empty",
+        serialize_with = "serialize_multi_value_headers"
+    )]
+    pub multi_value_headers: HeaderMap<HeaderValue>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub body: Option<Body>,
     #[serde(skip_serializing_if = "Not::not")]
     pub is_base64_encoded: bool,
 }
 
+#[cfg(test)]
 impl Default for GatewayResponse {
     fn default() -> Self {
         Self {
             status_code: 200,
             headers: Default::default(),
+            multi_value_headers: Default::default(),
             body: Default::default(),
             is_base64_encoded: Default::default(),
         }
     }
 }
 
+/// Serialize a http::HeaderMap into a serde str => str map
+fn serialize_multi_value_headers<S>(headers: &HeaderMap<HeaderValue>, serializer: S) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+{
+    let mut map = serializer.serialize_map(Some(headers.keys_len()))?;
+    for key in headers.keys() {
+        let mut map_values = Vec::new();
+        for value in headers.get_all(key) {
+            map_values.push(value.to_str().map_err(S::Error::custom)?)
+        }
+        map.serialize_entry(key.as_str(), &map_values)?;
+    }
+    map.end()
+}
+
+/// Serialize a http::HeaderMap into a serde str => Vec<str> map
 fn serialize_headers<S>(headers: &HeaderMap<HeaderValue>, serializer: S) -> Result<S::Ok, S::Error>
 where
     S: Serializer,
@@ -65,7 +89,8 @@ where
         GatewayResponse {
             status_code: parts.status.as_u16(),
             body,
-            headers: parts.headers,
+            headers: parts.headers.clone(),
+            multi_value_headers: parts.headers,
             is_base64_encoded,
         }
     }
@@ -114,7 +139,7 @@ where
 impl IntoResponse for serde_json::Value {
     fn into_response(self) -> Response<Body> {
         Response::builder()
-            .header(http::header::CONTENT_TYPE, "application/json")
+            .header(CONTENT_TYPE, "application/json")
             .body(
                 serde_json::to_string(&self)
                     .expect("unable to serialize serde_json::Value")
@@ -127,6 +152,7 @@ impl IntoResponse for serde_json::Value {
 #[cfg(test)]
 mod tests {
     use super::{Body, GatewayResponse, IntoResponse};
+    use http::{header::CONTENT_TYPE, Response};
     use serde_json::{self, json};
 
     #[test]
@@ -139,7 +165,7 @@ mod tests {
         assert_eq!(
             response
                 .headers()
-                .get(http::header::CONTENT_TYPE)
+                .get(CONTENT_TYPE)
                 .map(|h| h.to_str().expect("invalid header")),
             Some("application/json")
         )
@@ -176,4 +202,19 @@ mod tests {
             r#"{"statusCode":200,"body":"foo"}"#
         );
     }
+
+    #[test]
+    fn serialize_multi_value_headers() {
+        let res: GatewayResponse = Response::builder()
+            .header("multi", "a")
+            .header("multi", "b")
+            .body(Body::from(()))
+            .expect("failed to create response")
+            .into();
+        let json = serde_json::to_string(&res).expect("failed to serialize to json");
+        assert_eq!(
+            json,
+            r#"{"statusCode":200,"headers":{"multi":"a"},"multiValueHeaders":{"multi":["a","b"]}}"#
+        )
+    }
 }
diff --git a/lambda-http/tests/data/apigw_multi_value_proxy_request.json b/lambda-http/tests/data/apigw_multi_value_proxy_request.json
@@ -0,0 +1,131 @@
+{
+  "resource": "/{proxy+}",
+  "path": "/hello/world",
+  "httpMethod": "POST",
+  "headers": {
+    "Accept": "*/*",
+    "Accept-Encoding": "gzip, deflate",
+    "cache-control": "no-cache",
+    "CloudFront-Forwarded-Proto": "https",
+    "CloudFront-Is-Desktop-Viewer": "true",
+    "CloudFront-Is-Mobile-Viewer": "false",
+    "CloudFront-Is-SmartTV-Viewer": "false",
+    "CloudFront-Is-Tablet-Viewer": "false",
+    "CloudFront-Viewer-Country": "US",
+    "Content-Type": "application/json",
+    "headerName": "headerValue",
+    "Host": "gy415nuibc.execute-api.us-east-1.amazonaws.com",
+    "Postman-Token": "9f583ef0-ed83-4a38-aef3-eb9ce3f7a57f",
+    "User-Agent": "PostmanRuntime/2.4.5",
+    "Via": "1.1 d98420743a69852491bbdea73f7680bd.cloudfront.net (CloudFront)",
+    "X-Amz-Cf-Id": "pn-PWIJc6thYnZm5P0NMgOUglL1DYtl0gdeJky8tqsg8iS_sgsKD1A==",
+    "X-Forwarded-For": "54.240.196.186, 54.182.214.83",
+    "X-Forwarded-Port": "443",
+    "X-Forwarded-Proto": "https"
+  },
+  "multiValueHeaders":{
+    "Accept":[
+      "*/*"
+    ],
+    "Accept-Encoding":[
+      "gzip, deflate"
+    ],
+    "cache-control":[
+      "no-cache"
+    ],
+    "CloudFront-Forwarded-Proto":[
+      "https"
+    ],
+    "CloudFront-Is-Desktop-Viewer":[
+      "true"
+    ],
+    "CloudFront-Is-Mobile-Viewer":[
+      "false"
+    ],
+    "CloudFront-Is-SmartTV-Viewer":[
+      "false"
+    ],
+    "CloudFront-Is-Tablet-Viewer":[
+      "false"
+    ],
+    "CloudFront-Viewer-Country":[
+      "US"
+    ],
+    "":[
+      ""
+    ],
+    "Content-Type":[
+      "application/json"
+    ],
+    "headerName":[
+      "headerValue"
+    ],
+    "Host":[
+      "gy415nuibc.execute-api.us-east-1.amazonaws.com"
+    ],
+    "Postman-Token":[
+      "9f583ef0-ed83-4a38-aef3-eb9ce3f7a57f"
+    ],
+    "User-Agent":[
+      "PostmanRuntime/2.4.5"
+    ],
+    "Via":[
+      "1.1 d98420743a69852491bbdea73f7680bd.cloudfront.net (CloudFront)"
+    ],
+    "X-Amz-Cf-Id":[
+      "pn-PWIJc6thYnZm5P0NMgOUglL1DYtl0gdeJky8tqsg8iS_sgsKD1A=="
+    ],
+    "X-Forwarded-For":[
+      "54.240.196.186, 54.182.214.83"
+    ],
+    "X-Forwarded-Port":[
+      "443"
+    ],
+    "X-Forwarded-Proto":[
+      "https"
+    ]
+  },
+  "queryStringParameters": {
+    "name": "me",
+    "multivalueName": "me"
+  },
+  "multiValueQueryStringParameters":{
+    "name":[
+      "me"
+    ],
+    "multivalueName":[
+      "you",
+      "me"
+    ]
+  },
+  "pathParameters": {
+    "proxy": "hello/world"
+  },
+  "stageVariables": {
+    "stageVariableName": "stageVariableValue"
+  },
+  "requestContext": {
+    "accountId": "12345678912",
+    "resourceId": "roq9wj",
+    "stage": "testStage",
+    "requestId": "deef4878-7910-11e6-8f14-25afc3e9ae33",
+    "identity": {
+      "cognitoIdentityPoolId": null,
+      "accountId": null,
+      "cognitoIdentityId": null,
+      "caller": null,
+      "apiKey": null,
+      "sourceIp": "192.168.196.186",
+      "cognitoAuthenticationType": null,
+      "cognitoAuthenticationProvider": null,
+      "userArn": null,
+      "userAgent": "PostmanRuntime/2.4.5",
+      "user": null
+    },
+    "resourcePath": "/{proxy+}",
+    "httpMethod": "POST",
+    "apiId": "gy415nuibc"
+  },
+  "body": "{\r\n\t\"a\": 1\r\n}",
+  "isBase64Encoded": false
+}
diff --git a/lambda-http/tests/data/apigw_proxy_request.json b/lambda-http/tests/data/apigw_proxy_request.json
