From e5a8d62296f04c27bf67911c6c180694806f2627 Mon Sep 17 00:00:00 2001 From: David Calavera Date: Sun, 9 Jan 2022 13:19:52 -0800 Subject: [PATCH] Fix request query strings - Add query strings to the request uri. This ensures that we don't mangle information coming into lambda. - Fix bug in StrMapIter that ignored multi-value query parameters. Until now, this iterator only looped over the first element of each key, even if the key had multiple values assigned. Signed-off-by: David Calavera --- lambda-http/src/request.rs | 33 +++++++++++++++++--- lambda-http/src/strmap.rs | 62 +++++++++++++++++++++++++++++++++++++- 2 files changed, 89 insertions(+), 6 deletions(-) diff --git a/lambda-http/src/request.rs b/lambda-http/src/request.rs index de61f726..4b230e51 100644 --- a/lambda-http/src/request.rs +++ b/lambda-http/src/request.rs @@ -441,7 +441,7 @@ impl<'a> From> for http::Request { .method(http_method) .uri({ let host = headers.get(http::header::HOST).and_then(|val| val.to_str().ok()); - match host { + let mut uri = match host { Some(host) => { format!( "{}://{}{}", @@ -454,7 +454,17 @@ impl<'a> From> for http::Request { ) } None => path.to_string(), + }; + + if !multi_value_query_string_parameters.is_empty() { + uri.push('?'); + uri.push_str(multi_value_query_string_parameters.to_query_string().as_str()); + } else if !query_string_parameters.is_empty() { + uri.push('?'); + uri.push_str(query_string_parameters.to_query_string().as_str()); } + + uri }) // multi-valued query string parameters are always a super // set of singly valued query string parameters, @@ -507,7 +517,7 @@ impl<'a> From> for http::Request { .method(http_method) .uri({ let host = headers.get(http::header::HOST).and_then(|val| val.to_str().ok()); - match host { + let mut uri = match host { Some(host) => { format!( "{}://{}{}", @@ -520,7 +530,17 @@ impl<'a> From> for http::Request { ) } None => path.to_string(), + }; + + if !multi_value_query_string_parameters.is_empty() { + uri.push('?'); + uri.push_str(multi_value_query_string_parameters.to_query_string().as_str()); + } else if !query_string_parameters.is_empty() { + uri.push('?'); + uri.push_str(query_string_parameters.to_query_string().as_str()); } + + uri }) // multi valued query string parameters are always a super // set of singly valued query string parameters, @@ -698,7 +718,7 @@ mod tests { assert_eq!(req.method(), "GET"); assert_eq!( req.uri(), - "https://wt6mne2s9k.execute-api.us-west-2.amazonaws.com/test/hello" + "https://wt6mne2s9k.execute-api.us-west-2.amazonaws.com/test/hello?name=me" ); // Ensure this is an APIGW request @@ -727,7 +747,10 @@ 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/"); + assert_eq!( + req.uri(), + "https://lambda-846800462-us-east-2.elb.amazonaws.com/?myKey=val2" + ); // Ensure this is an ALB request let req_context = req.request_context(); @@ -822,7 +845,7 @@ mod tests { ); let req = result.expect("failed to parse request"); assert_eq!(req.method(), "GET"); - assert_eq!(req.uri(), "/test/hello"); + assert_eq!(req.uri(), "/test/hello?name=me"); } #[test] diff --git a/lambda-http/src/strmap.rs b/lambda-http/src/strmap.rs index 281a1101..066c575a 100644 --- a/lambda-http/src/strmap.rs +++ b/lambda-http/src/strmap.rs @@ -41,6 +41,20 @@ impl StrMap { StrMapIter { data: self, keys: self.0.keys(), + current: None, + next_idx: 0, + } + } + + /// Return the URI query representation for this map + pub fn to_query_string(&self) -> String { + if self.is_empty() { + "".into() + } else { + self.iter() + .map(|(k, v)| format!("{}={}", k, v)) + .collect::>() + .join("&") } } } @@ -62,6 +76,8 @@ impl From>> for StrMap { pub struct StrMapIter<'a> { data: &'a StrMap, keys: Keys<'a, String, Vec>, + current: Option<(&'a String, Vec<&'a str>)>, + next_idx: usize, } impl<'a> Iterator for StrMapIter<'a> { @@ -69,7 +85,31 @@ impl<'a> Iterator for StrMapIter<'a> { #[inline] fn next(&mut self) -> Option<(&'a str, &'a str)> { - self.keys.next().and_then(|k| self.data.get(k).map(|v| (k.as_str(), v))) + if self.current.is_none() { + self.current = self.keys.next().map(|k| (k, self.data.get_all(k).unwrap_or_default())); + }; + + let mut reset = false; + let ret = if let Some((key, values)) = &self.current { + let value = values[self.next_idx]; + + if self.next_idx + 1 < values.len() { + self.next_idx += 1; + } else { + reset = true; + } + + Some((key.as_str(), value)) + } else { + None + }; + + if reset { + self.current = None; + self.next_idx = 0; + } + + ret } } @@ -158,4 +198,24 @@ mod tests { values.sort(); assert_eq!(values, vec!["bar", "boom"]); } + + #[test] + fn test_empty_str_map_to_query_string() { + let data = HashMap::new(); + let strmap = StrMap(data.into()); + let query = strmap.to_query_string(); + assert_eq!("", &query); + } + + #[test] + fn test_str_map_to_query_string() { + let mut data = HashMap::new(); + data.insert("foo".into(), vec!["bar".into(), "qux".into()]); + data.insert("baz".into(), vec!["quux".into()]); + + let strmap = StrMap(data.into()); + let query = strmap.to_query_string(); + assert!(query.contains("foo=bar&foo=qux")); + assert!(query.contains("baz=quux")); + } }