Skip to main content

fiberplane_models/
query_data.rs

1use std::borrow::Cow;
2
3const MIME_TYPE_PREFIX: &str = "application/x-www-form-urlencoded,";
4
5/// Returns the value of a field in the query data.
6///
7/// Returns an empty string if the field has no value.
8pub fn get_query_field<'a>(query_data: &'a str, field_name: &str) -> Cow<'a, str> {
9    if let Some(data) = query_data.strip_prefix(MIME_TYPE_PREFIX) {
10        for (key, value) in form_urlencoded::parse(data.as_bytes()) {
11            if key == field_name {
12                return value;
13            }
14        }
15    }
16
17    Cow::Borrowed("")
18}
19
20/// Returns whether the query data string contains any query data that we
21/// understand.
22pub fn has_query_data(query_data: impl AsRef<str>) -> bool {
23    if let Some(data) = query_data.as_ref().strip_prefix(MIME_TYPE_PREFIX) {
24        !data.is_empty()
25    } else {
26        false
27    }
28}
29
30/// Sets the value of a field in the query data.
31///
32/// Returns the new query data.
33///
34/// This functions maintains an alphabetical ordering of the keys in order to
35/// guarantee a consistent result when separate fields are set out of order.
36/// This is to maintain convergence for our OT algorithm.
37pub fn set_query_field(
38    query_data: impl AsRef<str>,
39    field_name: impl AsRef<str>,
40    value: impl AsRef<str>,
41) -> String {
42    let mut new_query_data = MIME_TYPE_PREFIX.to_owned();
43    if let Some(data) = query_data.as_ref().strip_prefix(MIME_TYPE_PREFIX) {
44        let (before, after): (Vec<_>, Vec<_>) = form_urlencoded::parse(data.as_bytes())
45            .filter(|(key, _)| key.as_ref() != field_name.as_ref())
46            .partition(|(key, _)| key.as_ref() < field_name.as_ref());
47        for (key, value) in before
48            .iter()
49            .chain(&[(
50                Cow::Borrowed(field_name.as_ref()),
51                Cow::Borrowed(value.as_ref()),
52            )])
53            .chain(after.iter())
54        {
55            append_query_field(&mut new_query_data, key, value);
56        }
57    } else {
58        append_query_field(&mut new_query_data, field_name, value);
59    }
60    new_query_data
61}
62
63/// Removes a field from the query data.
64///
65/// Returns the new query data.
66pub fn unset_query_field(query_data: impl AsRef<str>, field_name: impl AsRef<str>) -> String {
67    let mut new_query_data = MIME_TYPE_PREFIX.to_owned();
68    if let Some(data) = query_data.as_ref().strip_prefix(MIME_TYPE_PREFIX) {
69        for (key, value) in form_urlencoded::parse(data.as_bytes()) {
70            if key != field_name.as_ref() {
71                append_query_field(&mut new_query_data, &key, &value);
72            }
73        }
74    }
75    new_query_data
76}
77
78fn append_query_field(
79    query_data: &mut String,
80    field_name: impl AsRef<str>,
81    value: impl AsRef<str>,
82) {
83    if query_data.len() > MIME_TYPE_PREFIX.len() {
84        query_data.push('&');
85    }
86
87    query_data.extend(form_urlencoded::byte_serialize(
88        field_name.as_ref().as_bytes(),
89    ));
90    query_data.push('=');
91    query_data.extend(form_urlencoded::byte_serialize(value.as_ref().as_bytes()));
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn test_get_query_field() {
100        assert_eq!(
101            get_query_field("application/x-www-form-urlencoded,trace_id=123", "trace_id"),
102            "123"
103        );
104        assert_eq!(
105            get_query_field("application/x-www-form-urlencoded,trace_id=123", "id"),
106            ""
107        );
108        assert_eq!(get_query_field("trace_id=123", "trace_id"), "");
109        assert_eq!(get_query_field("", "trace_id"), "");
110
111        assert_eq!(
112            get_query_field("application/x-www-form-urlencoded,hi+=%26there", "hi "),
113            "&there"
114        );
115    }
116
117    #[test]
118    fn test_set_query_field() {
119        assert_eq!(
120            &set_query_field(
121                "application/x-www-form-urlencoded,trace_id=123",
122                "trace_id",
123                "456"
124            ),
125            "application/x-www-form-urlencoded,trace_id=456"
126        );
127        assert_eq!(
128            &set_query_field(
129                "application/x-www-form-urlencoded,trace_id=123",
130                "id",
131                "456"
132            ),
133            "application/x-www-form-urlencoded,id=456&trace_id=123"
134        );
135        assert_eq!(
136            &set_query_field("trace_id=123", "trace_id", "456"),
137            "application/x-www-form-urlencoded,trace_id=456"
138        );
139        assert_eq!(
140            &set_query_field("", "trace_id", "456"),
141            "application/x-www-form-urlencoded,trace_id=456"
142        );
143
144        assert_eq!(
145            &set_query_field(
146                "application/x-www-form-urlencoded,hi+=%26there",
147                "hi!",
148                "-_.!~*'()#"
149            ),
150            "application/x-www-form-urlencoded,hi+=%26there&hi%21=-_.%21%7E*%27%28%29%23"
151        );
152    }
153
154    #[test]
155    fn test_unset_query_field() {
156        assert_eq!(
157            &unset_query_field("application/x-www-form-urlencoded,trace_id=123", "trace_id"),
158            "application/x-www-form-urlencoded,"
159        );
160        assert_eq!(
161            &unset_query_field("application/x-www-form-urlencoded,trace_id=123", "id"),
162            "application/x-www-form-urlencoded,trace_id=123"
163        );
164        assert_eq!(
165            &unset_query_field("trace_id=123", "trace_id"),
166            "application/x-www-form-urlencoded,"
167        );
168        assert_eq!(
169            &unset_query_field("", "trace_id"),
170            "application/x-www-form-urlencoded,"
171        );
172
173        assert_eq!(
174            &unset_query_field("application/x-www-form-urlencoded,hi+=%26there", "hi "),
175            "application/x-www-form-urlencoded,"
176        );
177    }
178}