1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
use std::borrow::Cow;

const MIME_TYPE_PREFIX: &str = "application/x-www-form-urlencoded,";

/// Returns the value of a field in the query data.
///
/// Returns an empty string if the field has no value.
pub fn get_query_field<'a>(query_data: &'a str, field_name: &str) -> Cow<'a, str> {
    if let Some(data) = query_data.strip_prefix(MIME_TYPE_PREFIX) {
        for (key, value) in form_urlencoded::parse(data.as_bytes()) {
            if key == field_name {
                return value;
            }
        }
    }

    Cow::Borrowed("")
}

/// Returns whether the query data string contains any query data that we
/// understand.
pub fn has_query_data(query_data: impl AsRef<str>) -> bool {
    if let Some(data) = query_data.as_ref().strip_prefix(MIME_TYPE_PREFIX) {
        !data.is_empty()
    } else {
        false
    }
}

/// Sets the value of a field in the query data.
///
/// Returns the new query data.
///
/// This functions maintains an alphabetical ordering of the keys in order to
/// guarantee a consistent result when separate fields are set out of order.
/// This is to maintain convergence for our OT algorithm.
pub fn set_query_field(
    query_data: impl AsRef<str>,
    field_name: impl AsRef<str>,
    value: impl AsRef<str>,
) -> String {
    let mut new_query_data = MIME_TYPE_PREFIX.to_owned();
    if let Some(data) = query_data.as_ref().strip_prefix(MIME_TYPE_PREFIX) {
        let (before, after): (Vec<_>, Vec<_>) = form_urlencoded::parse(data.as_bytes())
            .filter(|(key, _)| key.as_ref() != field_name.as_ref())
            .partition(|(key, _)| key.as_ref() < field_name.as_ref());
        for (key, value) in before
            .iter()
            .chain(&[(
                Cow::Borrowed(field_name.as_ref()),
                Cow::Borrowed(value.as_ref()),
            )])
            .chain(after.iter())
        {
            append_query_field(&mut new_query_data, key, value);
        }
    } else {
        append_query_field(&mut new_query_data, field_name, value);
    }
    new_query_data
}

/// Removes a field from the query data.
///
/// Returns the new query data.
pub fn unset_query_field(query_data: impl AsRef<str>, field_name: impl AsRef<str>) -> String {
    let mut new_query_data = MIME_TYPE_PREFIX.to_owned();
    if let Some(data) = query_data.as_ref().strip_prefix(MIME_TYPE_PREFIX) {
        for (key, value) in form_urlencoded::parse(data.as_bytes()) {
            if key != field_name.as_ref() {
                append_query_field(&mut new_query_data, &key, &value);
            }
        }
    }
    new_query_data
}

fn append_query_field(
    query_data: &mut String,
    field_name: impl AsRef<str>,
    value: impl AsRef<str>,
) {
    if query_data.len() > MIME_TYPE_PREFIX.len() {
        query_data.push('&');
    }

    query_data.extend(form_urlencoded::byte_serialize(
        field_name.as_ref().as_bytes(),
    ));
    query_data.push('=');
    query_data.extend(form_urlencoded::byte_serialize(value.as_ref().as_bytes()));
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_get_query_field() {
        assert_eq!(
            get_query_field("application/x-www-form-urlencoded,trace_id=123", "trace_id"),
            "123"
        );
        assert_eq!(
            get_query_field("application/x-www-form-urlencoded,trace_id=123", "id"),
            ""
        );
        assert_eq!(get_query_field("trace_id=123", "trace_id"), "");
        assert_eq!(get_query_field("", "trace_id"), "");

        assert_eq!(
            get_query_field("application/x-www-form-urlencoded,hi+=%26there", "hi "),
            "&there"
        );
    }

    #[test]
    fn test_set_query_field() {
        assert_eq!(
            &set_query_field(
                "application/x-www-form-urlencoded,trace_id=123",
                "trace_id",
                "456"
            ),
            "application/x-www-form-urlencoded,trace_id=456"
        );
        assert_eq!(
            &set_query_field(
                "application/x-www-form-urlencoded,trace_id=123",
                "id",
                "456"
            ),
            "application/x-www-form-urlencoded,id=456&trace_id=123"
        );
        assert_eq!(
            &set_query_field("trace_id=123", "trace_id", "456"),
            "application/x-www-form-urlencoded,trace_id=456"
        );
        assert_eq!(
            &set_query_field("", "trace_id", "456"),
            "application/x-www-form-urlencoded,trace_id=456"
        );

        assert_eq!(
            &set_query_field(
                "application/x-www-form-urlencoded,hi+=%26there",
                "hi!",
                "-_.!~*'()#"
            ),
            "application/x-www-form-urlencoded,hi+=%26there&hi%21=-_.%21%7E*%27%28%29%23"
        );
    }

    #[test]
    fn test_unset_query_field() {
        assert_eq!(
            &unset_query_field("application/x-www-form-urlencoded,trace_id=123", "trace_id"),
            "application/x-www-form-urlencoded,"
        );
        assert_eq!(
            &unset_query_field("application/x-www-form-urlencoded,trace_id=123", "id"),
            "application/x-www-form-urlencoded,trace_id=123"
        );
        assert_eq!(
            &unset_query_field("trace_id=123", "trace_id"),
            "application/x-www-form-urlencoded,"
        );
        assert_eq!(
            &unset_query_field("", "trace_id"),
            "application/x-www-form-urlencoded,"
        );

        assert_eq!(
            &unset_query_field("application/x-www-form-urlencoded,hi+=%26there", "hi "),
            "application/x-www-form-urlencoded,"
        );
    }
}