rama_http_headers/common/
sec_websocket_protocol.rs

1use rama_core::extensions::Extension;
2use rama_utils::str::NonEmptyStr;
3
4derive_non_empty_flat_csv_header! {
5    #[header(name = SEC_WEBSOCKET_PROTOCOL, sep = Comma)]
6    /// The `Sec-WebSocket-Protocol` header, containing one or multiple protocols.
7    ///
8    /// Sub protocols are advertised by the client,
9    /// and the server has to match it if defined.
10    #[derive(Clone, Debug, PartialEq, Eq)]
11    pub struct SecWebSocketProtocol(pub NonEmptySmallVec<3, NonEmptyStr>);
12}
13
14#[derive(Debug, Clone, PartialEq, Eq, Extension)]
15#[extension(tags(http, ws))]
16/// Utility type containing the accepted [`SecWebSocketProtocol`].
17pub struct AcceptedWebSocketProtocol(pub NonEmptyStr);
18
19impl AcceptedWebSocketProtocol {
20    #[inline]
21    #[must_use]
22    /// consume this instance as a [`SecWebSocketProtocol`]
23    ///
24    /// Useful for servers to communicate back to clients.
25    pub fn into_header(self) -> SecWebSocketProtocol {
26        self.into()
27    }
28}
29
30impl From<AcceptedWebSocketProtocol> for SecWebSocketProtocol {
31    fn from(value: AcceptedWebSocketProtocol) -> Self {
32        Self::new(value.0)
33    }
34}
35
36impl SecWebSocketProtocol {
37    #[must_use]
38    /// Return the first protocol in this [`SecWebSocketProtocol`] as the [`AcceptedWebSocketProtocol`].
39    pub fn accept_first_protocol(&self) -> AcceptedWebSocketProtocol {
40        // assumption: we always have at least one item
41        AcceptedWebSocketProtocol(self.0[0].clone())
42    }
43
44    /// returns true if the given protocol is found in this [`SecWebSocketProtocol`]
45    pub fn contains(&self, protocol: impl AsRef<str>) -> Option<AcceptedWebSocketProtocol> {
46        let protocol = protocol.as_ref().trim();
47        self.0.iter().find_map(|candidate| {
48            candidate
49                .trim()
50                .eq_ignore_ascii_case(protocol)
51                .then(|| AcceptedWebSocketProtocol(candidate.clone()))
52        })
53    }
54
55    /// returns true if any of the given protocol is found in this [`SecWebSocketProtocol`]
56    ///
57    /// Searched in order.
58    pub fn contains_any(
59        &self,
60        protocols: impl IntoIterator<Item: AsRef<str>>,
61    ) -> Option<AcceptedWebSocketProtocol> {
62        protocols
63            .into_iter()
64            .find_map(|protocol| self.contains(protocol))
65    }
66
67    pub fn iter(&self) -> impl Iterator<Item = &str> {
68        self.0.iter().map(|it| it.as_ref())
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use crate::common::{test_decode, test_encode};
76
77    #[test]
78    fn protocols_reflective_str_single() {
79        fn assert_encode_decode_eq(s: &str, equal: bool) {
80            let header: SecWebSocketProtocol = test_decode(&[s]).unwrap();
81            let headers = test_encode(header);
82            let output = &headers["sec-websocket-protocol"];
83            if equal {
84                assert_eq!(s, output, "input ({s}) != output ({output:?})");
85            } else {
86                assert_ne!(s, output, "input ({s}) == output ({output:?})");
87            }
88        }
89        assert_encode_decode_eq("foo", true);
90        assert_encode_decode_eq(" foo ", false);
91        assert_encode_decode_eq("x-foo-123", true);
92        assert_encode_decode_eq("X-Foo-Bar", true);
93        assert_encode_decode_eq("a b", true);
94    }
95
96    #[test]
97    fn protocols_reflective_str_multiple() {
98        fn assert_encode_decode_eq(s: &[&'static str], equal: bool) {
99            let header: SecWebSocketProtocol = test_decode(s).unwrap();
100            let headers = test_encode(header);
101            let output = &headers["sec-websocket-protocol"];
102            if equal {
103                assert_eq!(
104                    &s.join(", "),
105                    output,
106                    "input ({s:?}) != output ({output:?})"
107                );
108            } else {
109                assert_ne!(
110                    &s.join(", "),
111                    output,
112                    "input ({s:?}) == output ({output:?})"
113                );
114            }
115        }
116        assert_encode_decode_eq(&["foo"], true);
117        assert_encode_decode_eq(&["x-foo-123", "foo"], true);
118        assert_encode_decode_eq(&["a", "b", "c"], true);
119        assert_encode_decode_eq(&["a b", "c d"], true);
120    }
121
122    #[test]
123    fn test_accept_first_protocol() {
124        let header: SecWebSocketProtocol = test_decode(&["a, b"]).unwrap();
125        assert_eq!("a", header.accept_first_protocol().0);
126    }
127
128    #[test]
129    fn test_contains() {
130        for (input, protocol, expected) in [
131            ("a", "b", None),
132            ("a", "a", Some("a")),
133            ("a", " a", Some("a")),
134            ("a", "A ", Some("a")),
135            ("a", " A ", Some("a")),
136            ("a, b", " A ", Some("a")),
137            ("a, b", "b", Some("b")),
138            ("a, b", " B ", Some("b")),
139            ("a, b", " c ", None),
140        ] {
141            let header: SecWebSocketProtocol = test_decode(&[input]).unwrap();
142            assert_eq!(
143                expected,
144                header.contains(protocol).as_ref().map(|p| p.0.as_ref()),
145                "input: '{input}'"
146            );
147        }
148    }
149
150    #[test]
151    fn test_contains_any() {
152        struct Case {
153            input: &'static str,
154            protocols: &'static [&'static str],
155            expected: Option<&'static str>,
156        }
157        impl Case {
158            fn new(
159                input: &'static str,
160                protocols: &'static [&'static str],
161                expected: Option<&'static str>,
162            ) -> Self {
163                Self {
164                    input,
165                    protocols,
166                    expected,
167                }
168            }
169        }
170
171        for case in [
172            Case::new("a", &["b"], None),
173            Case::new("a", &["a"], Some("a")),
174            Case::new("a", &[" a"], Some("a")),
175            Case::new("a", &[" A "], Some("a")),
176            Case::new("a, b", &["b", "a"], Some("b")),
177            Case::new("a, b", &["c", "a", "b", "a"], Some("a")),
178            Case::new("a, b", &["c", "d"], None),
179            Case::new("a", &["c", "d"], None),
180            Case::new("d", &["c", "d"], Some("d")),
181        ] {
182            let header: SecWebSocketProtocol = test_decode(&[case.input]).unwrap();
183            assert_eq!(
184                case.expected,
185                header
186                    .contains_any(case.protocols)
187                    .as_ref()
188                    .map(|p| p.0.as_ref()),
189                "input: '{}'",
190                case.input,
191            );
192        }
193    }
194}
rama_http_headers/common/sec_websocket_protocol.rs

rama_http_headers/common/
sec_websocket_protocol.rs
