1 + | /*
|
2 + | * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
3 + | * SPDX-License-Identifier: Apache-2.0
|
4 + | */
|
5 + |
|
6 + | use crate::query::fmt_string as percent_encode_query;
|
7 + | use http_02x::uri::InvalidUri;
|
8 + | use http_02x::Uri;
|
9 + |
|
10 + | /// Utility for updating the query string in a [`Uri`].
|
11 + | #[allow(missing_debug_implementations)]
|
12 + | pub struct QueryWriter {
|
13 + | base_uri: Uri,
|
14 + | new_path_and_query: String,
|
15 + | prefix: Option<char>,
|
16 + | }
|
17 + |
|
18 + | impl QueryWriter {
|
19 + | /// Creates a new `QueryWriter` from a string
|
20 + | pub fn new_from_string(uri: &str) -> Result<Self, InvalidUri> {
|
21 + | Ok(Self::new(&Uri::try_from(uri)?))
|
22 + | }
|
23 + |
|
24 + | /// Creates a new `QueryWriter` based off the given `uri`.
|
25 + | pub fn new(uri: &Uri) -> Self {
|
26 + | let new_path_and_query = uri
|
27 + | .path_and_query()
|
28 + | .map(|pq| pq.to_string())
|
29 + | .unwrap_or_default();
|
30 + | let prefix = if uri.query().is_none() {
|
31 + | Some('?')
|
32 + | } else if !uri.query().unwrap_or_default().is_empty() {
|
33 + | Some('&')
|
34 + | } else {
|
35 + | None
|
36 + | };
|
37 + | QueryWriter {
|
38 + | base_uri: uri.clone(),
|
39 + | new_path_and_query,
|
40 + | prefix,
|
41 + | }
|
42 + | }
|
43 + |
|
44 + | /// Clears all query parameters.
|
45 + | pub fn clear_params(&mut self) {
|
46 + | if let Some(index) = self.new_path_and_query.find('?') {
|
47 + | self.new_path_and_query.truncate(index);
|
48 + | self.prefix = Some('?');
|
49 + | }
|
50 + | }
|
51 + |
|
52 + | /// Inserts a new query parameter. The key and value are percent encoded
|
53 + | /// by `QueryWriter`. Passing in percent encoded values will result in double encoding.
|
54 + | pub fn insert(&mut self, k: &str, v: &str) {
|
55 + | self.insert_encoded(&percent_encode_query(k), &percent_encode_query(v));
|
56 + | }
|
57 + |
|
58 + | /// Inserts a new already encoded query parameter. The key and value will be inserted
|
59 + | /// as is.
|
60 + | pub fn insert_encoded(&mut self, encoded_k: &str, encoded_v: &str) {
|
61 + | if let Some(prefix) = self.prefix {
|
62 + | self.new_path_and_query.push(prefix);
|
63 + | }
|
64 + | self.prefix = Some('&');
|
65 + | self.new_path_and_query.push_str(encoded_k);
|
66 + | self.new_path_and_query.push('=');
|
67 + | self.new_path_and_query.push_str(encoded_v)
|
68 + | }
|
69 + |
|
70 + | /// Returns just the built query string.
|
71 + | pub fn build_query(self) -> String {
|
72 + | self.build_uri().query().unwrap_or_default().to_string()
|
73 + | }
|
74 + |
|
75 + | /// Returns a full [`Uri`] with the query string updated.
|
76 + | pub fn build_uri(self) -> Uri {
|
77 + | let mut parts = self.base_uri.into_parts();
|
78 + | parts.path_and_query = Some(
|
79 + | self.new_path_and_query
|
80 + | .parse()
|
81 + | .expect("adding query should not invalidate URI"),
|
82 + | );
|
83 + | Uri::from_parts(parts).expect("a valid URL in should always produce a valid URL out")
|
84 + | }
|
85 + | }
|
86 + |
|
87 + | #[cfg(test)]
|
88 + | mod test {
|
89 + | use super::QueryWriter;
|
90 + | use http_02x::Uri;
|
91 + |
|
92 + | #[test]
|
93 + | fn empty_uri() {
|
94 + | let uri = Uri::from_static("http://www.example.com");
|
95 + | let mut query_writer = QueryWriter::new(&uri);
|
96 + | query_writer.insert("key", "val%ue");
|
97 + | query_writer.insert("another", "value");
|
98 + | assert_eq!(
|
99 + | query_writer.build_uri(),
|
100 + | Uri::from_static("http://www.example.com?key=val%25ue&another=value")
|
101 + | );
|
102 + | }
|
103 + |
|
104 + | #[test]
|
105 + | fn uri_with_path() {
|
106 + | let uri = Uri::from_static("http://www.example.com/path");
|
107 + | let mut query_writer = QueryWriter::new(&uri);
|
108 + | query_writer.insert("key", "val%ue");
|
109 + | query_writer.insert("another", "value");
|
110 + | assert_eq!(
|
111 + | query_writer.build_uri(),
|
112 + | Uri::from_static("http://www.example.com/path?key=val%25ue&another=value")
|
113 + | );
|
114 + | }
|
115 + |
|
116 + | #[test]
|
117 + | fn uri_with_path_and_query() {
|
118 + | let uri = Uri::from_static("http://www.example.com/path?original=here");
|
119 + | let mut query_writer = QueryWriter::new(&uri);
|
120 + | query_writer.insert("key", "val%ue");
|
121 + | query_writer.insert("another", "value");
|
122 + | assert_eq!(
|
123 + | query_writer.build_uri(),
|
124 + | Uri::from_static(
|
125 + | "http://www.example.com/path?original=here&key=val%25ue&another=value"
|
126 + | )
|
127 + | );
|
128 + | }
|
129 + |
|
130 + | #[test]
|
131 + | fn build_query() {
|
132 + | let uri = Uri::from_static("http://www.example.com");
|
133 + | let mut query_writer = QueryWriter::new(&uri);
|
134 + | query_writer.insert("key", "val%ue");
|
135 + | query_writer.insert("ano%ther", "value");
|
136 + | assert_eq!("key=val%25ue&ano%25ther=value", query_writer.build_query());
|
137 + | }
|
138 + |
|
139 + | #[test]
|
140 + | // This test ensures that the percent encoding applied to queries always produces a valid URI if
|
141 + | // the starting URI is valid
|
142 + | fn doesnt_panic_when_adding_query_to_valid_uri() {
|
143 + | let uri = Uri::from_static("http://www.example.com");
|
144 + |
|
145 + | let mut problematic_chars = Vec::new();
|
146 + |
|
147 + | for byte in u8::MIN..=u8::MAX {
|
148 + | match std::str::from_utf8(&[byte]) {
|
149 + | // If we can't make a str from the byte then we certainly can't make a URL from it
|
150 + | Err(_) => {
|
151 + | continue;
|
152 + | }
|
153 + | Ok(value) => {
|
154 + | let mut query_writer = QueryWriter::new(&uri);
|
155 + | query_writer.insert("key", value);
|
156 + |
|
157 + | if std::panic::catch_unwind(|| query_writer.build_uri()).is_err() {
|
158 + | problematic_chars.push(char::from(byte));
|
159 + | };
|
160 + | }
|
161 + | }
|
162 + | }
|
163 + |
|
164 + | if !problematic_chars.is_empty() {
|
165 + | panic!("we got some bad bytes here: {:#?}", problematic_chars)
|
166 + | }
|
167 + | }
|
168 + |
|
169 + | #[test]
|
170 + | fn clear_params() {
|
171 + | let uri = Uri::from_static("http://www.example.com/path?original=here&foo=1");
|
172 + | let mut query_writer = QueryWriter::new(&uri);
|
173 + | query_writer.clear_params();
|
174 + | query_writer.insert("new", "value");
|
175 + | assert_eq!("new=value", query_writer.build_query());
|
176 + | }
|
177 + | }
|