aws_smithy_http_server/instrumentation/sensitivity/uri/
query.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! A wrapper around a query string [`&str`](str) to allow for sensitivity.
7
8use std::fmt::{Debug, Display, Error, Formatter};
9
10use crate::instrumentation::{sensitivity::Sensitive, MakeFmt};
11
12/// Marks the sensitive data of a query string pair.
13#[derive(Debug, Default, PartialEq, Eq)]
14pub struct QueryMarker {
15    /// Set to `true` to mark the key as sensitive.
16    pub key: bool,
17    /// Set to `true` to mark the value as sensitive.
18    pub value: bool,
19}
20
21/// A wrapper around a query string [`&str`](str) which modifies the behavior of [`Display`]. Specific query string
22/// values are marked as sensitive by providing predicate over the keys. This accommodates the [httpQuery trait] and
23/// the [httpQueryParams trait].
24///
25/// The [`Display`] implementation will respect the `unredacted-logging` flag.
26///
27/// # Example
28///
29/// ```
30/// # use aws_smithy_http_server::instrumentation::sensitivity::uri::{Query, QueryMarker};
31/// # let uri = "";
32/// // Query string value with key "name" is redacted
33/// let uri = Query::new(&uri, |x| QueryMarker { key: false, value: x == "name" } );
34/// println!("{uri}");
35/// ```
36///
37/// [httpQuery trait]: https://smithy.io/2.0/spec/http-bindings.html#httpquery-trait
38/// [httpQueryParams trait]: https://smithy.io/2.0/spec/http-bindings.html#httpqueryparams-trait
39#[allow(missing_debug_implementations)]
40pub struct Query<'a, F> {
41    query: &'a str,
42    marker: F,
43}
44
45impl<'a, F> Query<'a, F> {
46    /// Constructs a new [`Query`].
47    pub fn new(query: &'a str, marker: F) -> Self {
48        Self { query, marker }
49    }
50}
51
52#[inline]
53fn write_pair<'a, F>(section: &'a str, marker: F, f: &mut Formatter<'_>) -> Result<(), Error>
54where
55    F: Fn(&'a str) -> QueryMarker,
56{
57    if let Some((key, value)) = section.split_once('=') {
58        match (marker)(key) {
59            QueryMarker { key: true, value: true } => write!(f, "{}={}", Sensitive(key), Sensitive(value)),
60            QueryMarker {
61                key: true,
62                value: false,
63            } => write!(f, "{}={value}", Sensitive(key)),
64            QueryMarker {
65                key: false,
66                value: true,
67            } => write!(f, "{key}={}", Sensitive(value)),
68            QueryMarker {
69                key: false,
70                value: false,
71            } => write!(f, "{key}={value}"),
72        }
73    } else {
74        write!(f, "{section}")
75    }
76}
77
78impl<'a, F> Display for Query<'a, F>
79where
80    F: Fn(&'a str) -> QueryMarker,
81{
82    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
83        let mut it = self.query.split('&');
84
85        if let Some(section) = it.next() {
86            write_pair(section, &self.marker, f)?;
87        }
88
89        for section in it {
90            write!(f, "&")?;
91            write_pair(section, &self.marker, f)?;
92        }
93
94        Ok(())
95    }
96}
97
98/// A [`MakeFmt`] producing [`Query`].
99#[derive(Clone)]
100pub struct MakeQuery<F>(pub(crate) F);
101
102impl<F> Debug for MakeQuery<F> {
103    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
104        f.debug_tuple("MakeQuery").field(&"...").finish()
105    }
106}
107impl<'a, F> MakeFmt<&'a str> for MakeQuery<F>
108where
109    F: Clone,
110{
111    type Target = Query<'a, F>;
112
113    fn make(&self, path: &'a str) -> Self::Target {
114        Query::new(path, self.0.clone())
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use http::Uri;
121
122    use crate::instrumentation::sensitivity::uri::tests::{
123        ALL_KEYS_QUERY_STRING_EXAMPLES, ALL_PAIRS_QUERY_STRING_EXAMPLES, ALL_VALUES_QUERY_STRING_EXAMPLES, EXAMPLES,
124        QUERY_STRING_EXAMPLES, X_QUERY_STRING_EXAMPLES,
125    };
126
127    use super::*;
128
129    #[test]
130    fn mark_none() {
131        let originals = EXAMPLES.into_iter().chain(QUERY_STRING_EXAMPLES).map(Uri::from_static);
132        for original in originals {
133            if let Some(query) = original.query() {
134                let output = Query::new(query, |_| QueryMarker::default()).to_string();
135                assert_eq!(output, query, "original = {original}");
136            }
137        }
138    }
139
140    #[test]
141    fn mark_all_keys() {
142        let originals = QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
143        let expecteds = ALL_KEYS_QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
144        for (original, expected) in originals.zip(expecteds) {
145            let output = Query::new(original.query().unwrap(), |_| QueryMarker {
146                key: true,
147                value: false,
148            })
149            .to_string();
150            assert_eq!(output, expected.query().unwrap(), "original = {original}");
151        }
152    }
153
154    #[test]
155    fn mark_all_values() {
156        let originals = QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
157        let expecteds = ALL_VALUES_QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
158        for (original, expected) in originals.zip(expecteds) {
159            let output = Query::new(original.query().unwrap(), |_| QueryMarker {
160                key: false,
161                value: true,
162            })
163            .to_string();
164            assert_eq!(output, expected.query().unwrap(), "original = {original}");
165        }
166    }
167
168    #[test]
169    fn mark_all_pairs() {
170        let originals = QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
171        let expecteds = ALL_PAIRS_QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
172        for (original, expected) in originals.zip(expecteds) {
173            let output = Query::new(original.query().unwrap(), |_| QueryMarker { key: true, value: true }).to_string();
174            assert_eq!(output, expected.query().unwrap(), "original = {original}");
175        }
176    }
177
178    #[test]
179    fn mark_x() {
180        let originals = QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
181        let expecteds = X_QUERY_STRING_EXAMPLES.into_iter().map(Uri::from_static);
182        for (original, expected) in originals.zip(expecteds) {
183            let output = Query::new(original.query().unwrap(), |key| QueryMarker {
184                key: false,
185                value: key == "x",
186            })
187            .to_string();
188            assert_eq!(output, expected.query().unwrap(), "original = {original}");
189        }
190    }
191}
aws_smithy_http_server/instrumentation/sensitivity/uri/query.rs

aws_smithy_http_server/instrumentation/sensitivity/uri/
query.rs
