Skip to content

Commit 917ed63

Browse files
calaveracalavera
authored
Fix request query strings (#395)
- 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 <[email protected]>
1 parent d7b7aa7 commit 917ed63

File tree

2 files changed

+89
-6
lines changed

2 files changed

+89
-6
lines changed

‎lambda-http/src/request.rs‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -441,7 +441,7 @@ impl<'a> From<LambdaRequest<'a>> for http::Request<Body> {
441441
.method(http_method)
442442
.uri({
443443
let host = headers.get(http::header::HOST).and_then(|val| val.to_str().ok());
444-
match host {
444+
let mut uri = match host {
445445
Some(host) => {
446446
format!(
447447
"{}://{}{}",
@@ -454,7 +454,17 @@ impl<'a> From<LambdaRequest<'a>> for http::Request<Body> {
454454
)
455455
}
456456
None => path.to_string(),
457+
};
458+
459+
if !multi_value_query_string_parameters.is_empty() {
460+
uri.push('?');
461+
uri.push_str(multi_value_query_string_parameters.to_query_string().as_str());
462+
} else if !query_string_parameters.is_empty() {
463+
uri.push('?');
464+
uri.push_str(query_string_parameters.to_query_string().as_str());
457465
}
466+
467+
uri
458468
})
459469
// multi-valued query string parameters are always a super
460470
// set of singly valued query string parameters,
@@ -507,7 +517,7 @@ impl<'a> From<LambdaRequest<'a>> for http::Request<Body> {
507517
.method(http_method)
508518
.uri({
509519
let host = headers.get(http::header::HOST).and_then(|val| val.to_str().ok());
510-
match host {
520+
let mut uri = match host {
511521
Some(host) => {
512522
format!(
513523
"{}://{}{}",
@@ -520,7 +530,17 @@ impl<'a> From<LambdaRequest<'a>> for http::Request<Body> {
520530
)
521531
}
522532
None => path.to_string(),
533+
};
534+
535+
if !multi_value_query_string_parameters.is_empty() {
536+
uri.push('?');
537+
uri.push_str(multi_value_query_string_parameters.to_query_string().as_str());
538+
} else if !query_string_parameters.is_empty() {
539+
uri.push('?');
540+
uri.push_str(query_string_parameters.to_query_string().as_str());
523541
}
542+
543+
uri
524544
})
525545
// multi valued query string parameters are always a super
526546
// set of singly valued query string parameters,
@@ -698,7 +718,7 @@ mod tests {
698718
assert_eq!(req.method(), "GET");
699719
assert_eq!(
700720
req.uri(),
701-
"https://wt6mne2s9k.execute-api.us-west-2.amazonaws.com/test/hello"
721+
"https://wt6mne2s9k.execute-api.us-west-2.amazonaws.com/test/hello?name=me"
702722
);
703723

704724
// Ensure this is an APIGW request
@@ -727,7 +747,10 @@ mod tests {
727747
);
728748
let req = result.expect("failed to parse request");
729749
assert_eq!(req.method(), "GET");
730-
assert_eq!(req.uri(), "https://lambda-846800462-us-east-2.elb.amazonaws.com/");
750+
assert_eq!(
751+
req.uri(),
752+
"https://lambda-846800462-us-east-2.elb.amazonaws.com/?myKey=val2"
753+
);
731754

732755
// Ensure this is an ALB request
733756
let req_context = req.request_context();
@@ -822,7 +845,7 @@ mod tests {
822845
);
823846
let req = result.expect("failed to parse request");
824847
assert_eq!(req.method(), "GET");
825-
assert_eq!(req.uri(), "/test/hello");
848+
assert_eq!(req.uri(), "/test/hello?name=me");
826849
}
827850

828851
#[test]

‎lambda-http/src/strmap.rs‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,20 @@ impl StrMap {
4141
StrMapIter {
4242
data: self,
4343
keys: self.0.keys(),
44+
current: None,
45+
next_idx: 0,
46+
}
47+
}
48+
49+
/// Return the URI query representation for this map
50+
pub fn to_query_string(&self) -> String {
51+
if self.is_empty() {
52+
"".into()
53+
} else {
54+
self.iter()
55+
.map(|(k, v)| format!("{}={}", k, v))
56+
.collect::<Vec<_>>()
57+
.join("&")
4458
}
4559
}
4660
}
@@ -62,14 +76,40 @@ impl From<HashMap<String, Vec<String>>> for StrMap {
6276
pub struct StrMapIter<'a> {
6377
data: &'a StrMap,
6478
keys: Keys<'a, String, Vec<String>>,
79+
current: Option<(&'a String, Vec<&'a str>)>,
80+
next_idx: usize,
6581
}
6682

6783
impl<'a> Iterator for StrMapIter<'a> {
6884
type Item = (&'a str, &'a str);
6985

7086
#[inline]
7187
fn next(&mut self) -> Option<(&'a str, &'a str)> {
72-
self.keys.next().and_then(|k| self.data.get(k).map(|v| (k.as_str(), v)))
88+
if self.current.is_none() {
89+
self.current = self.keys.next().map(|k| (k, self.data.get_all(k).unwrap_or_default()));
90+
};
91+
92+
let mut reset = false;
93+
let ret = if let Some((key, values)) = &self.current {
94+
let value = values[self.next_idx];
95+
96+
if self.next_idx + 1 < values.len() {
97+
self.next_idx += 1;
98+
} else {
99+
reset = true;
100+
}
101+
102+
Some((key.as_str(), value))
103+
} else {
104+
None
105+
};
106+
107+
if reset {
108+
self.current = None;
109+
self.next_idx = 0;
110+
}
111+
112+
ret
73113
}
74114
}
75115

@@ -158,4 +198,24 @@ mod tests {
158198
values.sort();
159199
assert_eq!(values, vec!["bar", "boom"]);
160200
}
201+
202+
#[test]
203+
fn test_empty_str_map_to_query_string() {
204+
let data = HashMap::new();
205+
let strmap = StrMap(data.into());
206+
let query = strmap.to_query_string();
207+
assert_eq!("", &query);
208+
}
209+
210+
#[test]
211+
fn test_str_map_to_query_string() {
212+
let mut data = HashMap::new();
213+
data.insert("foo".into(), vec!["bar".into(), "qux".into()]);
214+
data.insert("baz".into(), vec!["quux".into()]);
215+
216+
let strmap = StrMap(data.into());
217+
let query = strmap.to_query_string();
218+
assert!(query.contains("foo=bar&foo=qux"));
219+
assert!(query.contains("baz=quux"));
220+
}
161221
}

0 commit comments

Comments
 (0)