trillium_http/
headers.rs

1mod header_name;
2mod header_value;
3mod header_values;
4mod known_header_name;
5mod unknown_header_name;
6
7pub use header_name::HeaderName;
8pub use header_value::HeaderValue;
9pub use header_values::HeaderValues;
10pub use known_header_name::KnownHeaderName;
11
12use header_name::HeaderNameInner;
13use unknown_header_name::UnknownHeaderName;
14
15use hashbrown::{
16    hash_map::{self, Entry},
17    HashMap,
18};
19use smartcow::SmartCow;
20use std::{
21    fmt::{self, Debug, Display, Formatter},
22    hash::{BuildHasherDefault, Hasher},
23};
24
25/// Trillium's header map type
26#[derive(Debug, Clone, PartialEq, Eq)]
27#[must_use]
28pub struct Headers {
29    known: HashMap<KnownHeaderName, HeaderValues, BuildHasherDefault<DirectHasher>>,
30    unknown: HashMap<UnknownHeaderName<'static>, HeaderValues>,
31}
32
33#[cfg(feature = "serde")]
34impl serde::Serialize for Headers {
35    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
36    where
37        S: serde::Serializer,
38    {
39        use serde::ser::SerializeMap;
40        let mut map = serializer.serialize_map(Some(self.len()))?;
41        for (key, values) in self {
42            map.serialize_entry(&key, values)?;
43        }
44        map.end()
45    }
46}
47
48impl Default for Headers {
49    fn default() -> Self {
50        Self::with_capacity(15)
51    }
52}
53
54impl Display for Headers {
55    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
56        for (n, v) in self {
57            for v in v {
58                f.write_fmt(format_args!("{n}: {v}\r\n"))?;
59            }
60        }
61        Ok(())
62    }
63}
64
65impl Headers {
66    /// Construct a new Headers, expecting to see at least this many known headers.
67    pub fn with_capacity(capacity: usize) -> Self {
68        Self {
69            known: HashMap::with_capacity_and_hasher(capacity, BuildHasherDefault::default()),
70            unknown: HashMap::with_capacity(0),
71        }
72    }
73
74    /// Construct a new headers with a default capacity of 15 known headers
75    pub fn new() -> Self {
76        Self::default()
77    }
78
79    /// Extend the capacity of the known headers map by this many
80    pub fn reserve(&mut self, additional: usize) {
81        self.known.reserve(additional);
82    }
83
84    /// Return an iterator over borrowed header names and header
85    /// values. First yields the known headers and then the unknown
86    /// headers, if any.
87    pub fn iter(&self) -> Iter<'_> {
88        self.into()
89    }
90
91    /// Are there zero headers?
92    pub fn is_empty(&self) -> bool {
93        self.known.is_empty() && self.unknown.is_empty()
94    }
95
96    /// How many unique [`HeaderName`] have been added to these [`Headers`]?
97    /// Note that each header name may have more than one [`HeaderValue`].
98    pub fn len(&self) -> usize {
99        self.known.len() + self.unknown.len()
100    }
101
102    /// add the header value or header values into this header map. If
103    /// there is already a header with the same name, the new values
104    /// will be added to the existing ones. To replace any existing
105    /// values, use [`Headers::insert`]
106    pub fn append(&mut self, name: impl Into<HeaderName<'static>>, value: impl Into<HeaderValues>) {
107        let value = value.into();
108        match name.into().0 {
109            HeaderNameInner::KnownHeader(known) => match self.known.entry(known) {
110                Entry::Occupied(mut o) => {
111                    o.get_mut().extend(value);
112                }
113                Entry::Vacant(v) => {
114                    v.insert(value);
115                }
116            },
117
118            HeaderNameInner::UnknownHeader(unknown) => match self.unknown.entry(unknown) {
119                Entry::Occupied(mut o) => {
120                    o.get_mut().extend(value);
121                }
122                Entry::Vacant(v) => {
123                    v.insert(value);
124                }
125            },
126        }
127    }
128
129    /// A slightly more efficient way to combine two [`Headers`] than
130    /// using [`Extend`]
131    pub fn append_all(&mut self, other: Headers) {
132        self.known.reserve(other.known.len());
133        for (name, value) in other.known {
134            match self.known.entry(name) {
135                Entry::Occupied(mut entry) => {
136                    entry.get_mut().extend(value);
137                }
138                Entry::Vacant(entry) => {
139                    entry.insert(value);
140                }
141            }
142        }
143
144        for (name, value) in other.unknown {
145            match self.unknown.entry(name) {
146                Entry::Occupied(mut entry) => {
147                    entry.get_mut().extend(value);
148                }
149                Entry::Vacant(entry) => {
150                    entry.insert(value);
151                }
152            }
153        }
154    }
155
156    /// Combine two [`Headers`], replacing any existing header values
157    pub fn insert_all(&mut self, other: Headers) {
158        self.known.reserve(other.known.len());
159        for (name, value) in other.known {
160            self.known.insert(name, value);
161        }
162
163        for (name, value) in other.unknown {
164            self.unknown.insert(name, value);
165        }
166    }
167
168    /// Add a header value or header values into this header map. If a
169    /// header already exists with the same name, it will be
170    /// replaced. To combine, see [`Headers::append`]
171    pub fn insert(&mut self, name: impl Into<HeaderName<'static>>, value: impl Into<HeaderValues>) {
172        let value = value.into();
173        match name.into().0 {
174            HeaderNameInner::KnownHeader(known) => {
175                self.known.insert(known, value);
176            }
177
178            HeaderNameInner::UnknownHeader(unknown) => {
179                self.unknown.insert(unknown, value);
180            }
181        }
182    }
183
184    /// Add a header value or header values into this header map if
185    /// and only if there is not already a header with the same name.
186    pub fn try_insert(
187        &mut self,
188        name: impl Into<HeaderName<'static>>,
189        value: impl Into<HeaderValues>,
190    ) {
191        let value = value.into();
192        match name.into().0 {
193            HeaderNameInner::KnownHeader(known) => {
194                self.known.entry(known).or_insert(value);
195            }
196
197            HeaderNameInner::UnknownHeader(unknown) => {
198                self.unknown.entry(unknown).or_insert(value);
199            }
200        }
201    }
202
203    /// Retrieves a &str header value if there is at least one header
204    /// in the map with this name. If there are several headers with
205    /// the same name, this follows the behavior defined at
206    /// [`HeaderValues::one`]. Returns None if there is no header with
207    /// the provided header name.
208    pub fn get_str<'a>(&self, name: impl Into<HeaderName<'a>>) -> Option<&str> {
209        self.get_values(name).and_then(HeaderValues::as_str)
210    }
211
212    pub(crate) fn get_lower<'a>(&self, name: impl Into<HeaderName<'a>>) -> Option<SmartCow<'_>> {
213        self.get_values(name).and_then(HeaderValues::as_lower)
214    }
215
216    /// Retrieves a singular header value from this header map. If
217    /// there are several headers with the same name, this follows the
218    /// behavior defined at [`HeaderValues::one`]. Returns None if there is no header with the provided header name
219    pub fn get<'a>(&self, name: impl Into<HeaderName<'a>>) -> Option<&HeaderValue> {
220        self.get_values(name).and_then(HeaderValues::one)
221    }
222
223    /// Takes all headers with the provided header name out of this
224    /// header map and returns them. Returns None if the header did
225    /// not have an entry in this map.
226    pub fn remove<'a>(&mut self, name: impl Into<HeaderName<'a>>) -> Option<HeaderValues> {
227        match name.into().0 {
228            HeaderNameInner::KnownHeader(known) => self.known.remove(&known),
229            HeaderNameInner::UnknownHeader(unknown) => self.unknown.remove(&&unknown),
230        }
231    }
232
233    /// Retrieves a reference to all header values with the provided
234    /// header name. If you expect there to be only one value, use
235    /// [`Headers::get`].
236    pub fn get_values<'a>(&self, name: impl Into<HeaderName<'a>>) -> Option<&HeaderValues> {
237        match name.into().0 {
238            HeaderNameInner::KnownHeader(known) => self.known.get(&known),
239            HeaderNameInner::UnknownHeader(unknown) => self.unknown.get(&&unknown),
240        }
241    }
242
243    /// Predicate function to check whether this header map contains
244    /// the provided header name. If you are using this to
245    /// conditionally insert a value, consider using
246    /// [`Headers::try_insert`] instead.
247    pub fn has_header<'a>(&self, name: impl Into<HeaderName<'a>>) -> bool {
248        match name.into().0 {
249            HeaderNameInner::KnownHeader(known) => self.known.contains_key(&known),
250            HeaderNameInner::UnknownHeader(unknown) => self.unknown.contains_key(&unknown),
251        }
252    }
253
254    /// Convenience function to check whether the value contained in
255    /// this header map for the provided name is
256    /// ascii-case-insensitively equal to the provided comparison
257    /// &str. Returns false if there is no value for the name
258    pub fn eq_ignore_ascii_case<'a>(
259        &'a self,
260        name: impl Into<HeaderName<'a>>,
261        needle: &str,
262    ) -> bool {
263        self.get_str(name)
264            .is_some_and(|v| v.eq_ignore_ascii_case(needle))
265    }
266
267    /// Deprecated because is likely not what you want. It is rarely the case that headers should
268    /// be searched for a matching string instead of carefully parsed according to the appropriate
269    /// header rule. Naive string matching on headers without regard to header structure is a
270    /// possible source of spec noncompliance or occasionally security vulnerability, so trillium
271    /// does not go out of its way to facilitate that.
272    #[deprecated = "Please open an issue if this behavior is important to you. \
273See documentation for deprecation rationale"]
274    pub fn contains_ignore_ascii_case<'a>(
275        &self,
276        name: impl Into<HeaderName<'a>>,
277        needle: &str,
278    ) -> bool {
279        self.get_str(name).is_some_and(|h| {
280            let needle = if needle.chars().all(|c| c.is_ascii_lowercase()) {
281                SmartCow::Borrowed(needle)
282            } else {
283                SmartCow::Owned(needle.chars().map(|c| c.to_ascii_lowercase()).collect())
284            };
285
286            if h.chars().all(|c| c.is_ascii_lowercase()) {
287                h.contains(&*needle)
288            } else {
289                h.to_ascii_lowercase().contains(&*needle)
290            }
291        })
292    }
293
294    /// Chainable method to insert a header
295    pub fn with_inserted_header(
296        mut self,
297        name: impl Into<HeaderName<'static>>,
298        values: impl Into<HeaderValues>,
299    ) -> Self {
300        self.insert(name, values);
301        self
302    }
303
304    /// Chainable method to append a header
305    pub fn with_appended_header(
306        mut self,
307        name: impl Into<HeaderName<'static>>,
308        values: impl Into<HeaderValues>,
309    ) -> Self {
310        self.append(name, values);
311        self
312    }
313
314    /// Chainable method to remove a header
315    pub fn without_header(mut self, name: impl Into<HeaderName<'static>>) -> Self {
316        self.remove(name);
317        self
318    }
319
320    /// Chainable method to remove multiple headers by name
321    pub fn without_headers<I, H>(mut self, names: I) -> Self
322    where
323        I: IntoIterator<Item = H>,
324        H: Into<HeaderName<'static>>,
325    {
326        self.remove_all(names);
327        self
328    }
329
330    /// remove multiple headers by name
331    pub fn remove_all<I, H>(&mut self, names: I)
332    where
333        I: IntoIterator<Item = H>,
334        H: Into<HeaderName<'static>>,
335    {
336        for header in names {
337            self.remove(header.into());
338        }
339    }
340
341    /// if a key does not exist already, execute the provided function and insert a value
342    ///
343    /// this can be useful to avoid calculating an unnecessary header value, or checking for the
344    /// presence of a key before insertion
345    pub fn try_insert_with<F, V>(&mut self, name: impl Into<HeaderName<'static>>, values_fn: F)
346    where
347        F: Fn() -> V,
348        V: Into<HeaderValues>,
349    {
350        match name.into().0 {
351            HeaderNameInner::KnownHeader(known) => {
352                self.known
353                    .entry(known)
354                    .or_insert_with(|| values_fn().into());
355            }
356
357            HeaderNameInner::UnknownHeader(unknown) => {
358                self.unknown
359                    .entry(unknown)
360                    .or_insert_with(|| values_fn().into());
361            }
362        }
363    }
364}
365
366impl<HN, HV> Extend<(HN, HV)> for Headers
367where
368    HN: Into<HeaderName<'static>>,
369    HV: Into<HeaderValues>,
370{
371    fn extend<T: IntoIterator<Item = (HN, HV)>>(&mut self, iter: T) {
372        let iter = iter.into_iter();
373        match iter.size_hint() {
374            (additional, _) if additional > 0 => self.known.reserve(additional),
375            _ => {}
376        };
377
378        for (name, values) in iter {
379            self.append(name, values);
380        }
381    }
382}
383
384impl<HN, HV> FromIterator<(HN, HV)> for Headers
385where
386    HN: Into<HeaderName<'static>>,
387    HV: Into<HeaderValues>,
388{
389    fn from_iter<T: IntoIterator<Item = (HN, HV)>>(iter: T) -> Self {
390        let iter = iter.into_iter();
391        let mut headers = match iter.size_hint() {
392            (0, _) => Self::new(),
393            (n, _) => Self::with_capacity(n),
394        };
395
396        for (name, values) in iter {
397            headers.append(name, values);
398        }
399
400        headers
401    }
402}
403
404#[derive(Default)]
405struct DirectHasher(u8);
406
407impl Hasher for DirectHasher {
408    fn write(&mut self, _: &[u8]) {
409        unreachable!("KnownHeaderName calls write_u64");
410    }
411
412    #[inline]
413    fn write_u8(&mut self, i: u8) {
414        self.0 = i;
415    }
416
417    #[inline]
418    fn finish(&self) -> u64 {
419        u64::from(self.0)
420    }
421}
422
423impl<'a> IntoIterator for &'a Headers {
424    type Item = (HeaderName<'a>, &'a HeaderValues);
425
426    type IntoIter = Iter<'a>;
427
428    fn into_iter(self) -> Self::IntoIter {
429        self.into()
430    }
431}
432
433#[derive(Debug)]
434pub struct IntoIter {
435    known: hash_map::IntoIter<KnownHeaderName, HeaderValues>,
436    unknown: hash_map::IntoIter<UnknownHeaderName<'static>, HeaderValues>,
437}
438
439impl Iterator for IntoIter {
440    type Item = (HeaderName<'static>, HeaderValues);
441
442    fn next(&mut self) -> Option<Self::Item> {
443        let IntoIter { known, unknown } = self;
444        known
445            .next()
446            .map(|(k, v)| (HeaderName::from(k), v))
447            .or_else(|| unknown.next().map(|(k, v)| (HeaderName::from(k), v)))
448    }
449}
450impl From<Headers> for IntoIter {
451    fn from(value: Headers) -> Self {
452        Self {
453            known: value.known.into_iter(),
454            unknown: value.unknown.into_iter(),
455        }
456    }
457}
458
459#[derive(Debug)]
460pub struct Iter<'a> {
461    known: hash_map::Iter<'a, KnownHeaderName, HeaderValues>,
462    unknown: hash_map::Iter<'a, UnknownHeaderName<'static>, HeaderValues>,
463}
464
465impl<'a> From<&'a Headers> for Iter<'a> {
466    fn from(value: &'a Headers) -> Self {
467        Iter {
468            known: value.known.iter(),
469            unknown: value.unknown.iter(),
470        }
471    }
472}
473
474impl<'a> Iterator for Iter<'a> {
475    type Item = (HeaderName<'a>, &'a HeaderValues);
476
477    fn next(&mut self) -> Option<Self::Item> {
478        let Iter { known, unknown } = self;
479        known
480            .next()
481            .map(|(k, v)| (HeaderName::from(*k), v))
482            .or_else(|| unknown.next().map(|(k, v)| (HeaderName::from(&**k), v)))
483    }
484}
485
486impl IntoIterator for Headers {
487    type Item = (HeaderName<'static>, HeaderValues);
488
489    type IntoIter = IntoIter;
490
491    fn into_iter(self) -> Self::IntoIter {
492        self.into()
493    }
494}
495
496#[cfg(test)]
497mod tests {
498    use crate::{Headers, KnownHeaderName};
499
500    #[test]
501    fn header_names_are_case_insensitive_for_access_but_retain_initial_case_in_headers() {
502        let mut headers = Headers::new();
503        headers.insert("my-Header-name", "initial-value");
504        headers.insert("my-Header-NAME", "my-header-value");
505
506        assert_eq!(headers.len(), 1);
507
508        assert_eq!(
509            headers.get_str("My-Header-Name").unwrap(),
510            "my-header-value"
511        );
512
513        headers.append("mY-hEaDer-NaMe", "second-value");
514        assert_eq!(
515            headers.get_values("my-header-name").unwrap(),
516            ["my-header-value", "second-value"].as_slice()
517        );
518
519        assert_eq!(
520            headers.iter().next().unwrap().0.to_string(),
521            "my-Header-name"
522        );
523
524        assert!(headers.remove("my-HEADER-name").is_some());
525        assert!(headers.is_empty());
526    }
527
528    #[test]
529    fn value_case_insensitive_comparison() {
530        let mut headers = Headers::new();
531        headers.insert(KnownHeaderName::Upgrade, "WebSocket");
532        headers.insert(KnownHeaderName::Connection, "upgrade");
533
534        assert!(headers.eq_ignore_ascii_case(KnownHeaderName::Upgrade, "websocket"));
535        assert!(headers.eq_ignore_ascii_case(KnownHeaderName::Connection, "Upgrade"));
536    }
537}