trillium_http/
headers.rs

1//! Header types
2mod entry;
3mod header_name;
4mod header_value;
5mod header_values;
6mod known_header_name;
7mod unknown_header_name;
8
9#[cfg(feature = "unstable")]
10pub mod qpack;
11
12#[cfg(not(feature = "unstable"))]
13pub(crate) mod qpack;
14
15use crate::headers::entry::{OccupiedEntryInner, VacantEntryInner};
16pub use entry::{Entry, OccupiedEntry, VacantEntry};
17use hashbrown::{
18    HashMap,
19    hash_map::{self, Entry as HashbrownEntry},
20};
21pub use header_name::HeaderName;
22use header_name::HeaderNameInner;
23pub use header_value::HeaderValue;
24pub use header_values::HeaderValues;
25pub use known_header_name::KnownHeaderName;
26use std::{
27    collections::{
28        BTreeMap,
29        btree_map::{self, Entry as BTreeEntry},
30    },
31    fmt::{self, Debug, Display, Formatter},
32};
33use unknown_header_name::UnknownHeaderName;
34
35/// Trillium's header map type
36#[derive(Debug, Clone, PartialEq, Eq, Default)]
37#[must_use]
38pub struct Headers {
39    known: BTreeMap<KnownHeaderName, HeaderValues>,
40    unknown: HashMap<UnknownHeaderName<'static>, HeaderValues>,
41}
42
43/// Default Server header
44pub const SERVER_HEADER: HeaderValue =
45    HeaderValue::const_new(concat!("trillium-http/", env!("CARGO_PKG_VERSION")));
46
47#[cfg(feature = "serde")]
48impl serde::Serialize for Headers {
49    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
50    where
51        S: serde::Serializer,
52    {
53        use serde::ser::SerializeMap;
54        let mut map = serializer.serialize_map(Some(self.len()))?;
55        for (key, values) in self {
56            map.serialize_entry(&key, values)?;
57        }
58        map.end()
59    }
60}
61
62impl Display for Headers {
63    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
64        // for testing
65        let mut data = self.iter().collect::<Vec<_>>();
66        data.sort_by(|(a, _), (b, _)| a.cmp(b));
67        for (n, v) in data {
68            for v in v {
69                f.write_fmt(format_args!("{n}: {v}\r\n"))?;
70            }
71        }
72        Ok(())
73    }
74}
75
76#[cfg(feature = "parse")]
77impl Headers {
78    #[doc(hidden)]
79    pub fn extend_parse(&mut self, bytes: &[u8]) -> Result<usize, crate::Error> {
80        use memchr::memmem::Finder;
81
82        let mut new_header_count = 0;
83        let mut last_line = 0;
84        for newline in Finder::new(b"\r\n").find_iter(bytes) {
85            if newline == last_line {
86                continue;
87            }
88
89            let line = &bytes[last_line..newline];
90            let colon = memchr::memchr(b':', line).ok_or(crate::Error::InvalidHead)?;
91
92            let header_name = HeaderName::parse(&line[..colon])?.to_owned();
93
94            let token_end = last_line + colon;
95
96            let mut value_start = token_end + 1;
97
98            while bytes
99                .get(value_start)
100                .is_some_and(|b| matches!(b, b'\t' | b' '))
101            {
102                value_start += 1;
103            }
104
105            let header_value = HeaderValue::parse(bytes[value_start..newline].trim_ascii_end());
106            self.append(header_name, header_value);
107            new_header_count += 1;
108            last_line = newline + 2;
109        }
110        Ok(new_header_count)
111    }
112
113    #[doc(hidden)]
114    pub fn parse(bytes: &[u8]) -> Result<Self, crate::Error> {
115        let mut headers = Headers::new();
116        headers.extend_parse(bytes)?;
117        Ok(headers)
118    }
119}
120
121impl Headers {
122    /// Construct a new headers with a default capacity
123    pub fn new() -> Self {
124        Self::default()
125    }
126
127    /// Return an iterator over borrowed header names and header
128    /// values. First yields the known headers and then the unknown
129    /// headers, if any.
130    pub fn iter(&self) -> Iter<'_> {
131        self.into()
132    }
133
134    /// Are there zero headers?
135    pub fn is_empty(&self) -> bool {
136        self.known.is_empty() && self.unknown.is_empty()
137    }
138
139    /// How many unique [`HeaderName`] have been added to these [`Headers`]?
140    /// Note that each header name may have more than one [`HeaderValue`].
141    pub fn len(&self) -> usize {
142        self.known.len() + self.unknown.len()
143    }
144
145    /// add the header value or header values into this header map. If
146    /// there is already a header with the same name, the new values
147    /// will be added to the existing ones. To replace any existing
148    /// values, use [`Headers::insert`]
149    ///
150    /// Identical to [`headers.entry(name).append(values)`][Entry::append]
151    pub fn append(
152        &mut self,
153        name: impl Into<HeaderName<'static>>,
154        values: impl Into<HeaderValues>,
155    ) -> &mut HeaderValues {
156        self.entry(name).append(values)
157    }
158
159    /// A slightly more efficient way to combine two [`Headers`] than
160    /// using [`Extend`]
161    pub fn append_all(&mut self, other: Headers) {
162        for (name, value) in other.known {
163            match self.known.entry(name) {
164                BTreeEntry::Occupied(mut entry) => {
165                    entry.get_mut().extend(value);
166                }
167                BTreeEntry::Vacant(entry) => {
168                    entry.insert(value);
169                }
170            }
171        }
172
173        for (name, value) in other.unknown {
174            match self.unknown.entry(name) {
175                HashbrownEntry::Occupied(mut entry) => {
176                    entry.get_mut().extend(value);
177                }
178                HashbrownEntry::Vacant(entry) => {
179                    entry.insert(value);
180                }
181            }
182        }
183    }
184
185    /// Combine two [`Headers`], replacing any existing header values
186    pub fn insert_all(&mut self, other: Headers) {
187        for (name, value) in other.known {
188            self.known.insert(name, value);
189        }
190
191        for (name, value) in other.unknown {
192            self.unknown.insert(name, value);
193        }
194    }
195
196    /// Add a header value or header values into this header map. If a
197    /// header already exists with the same name, it will be
198    /// replaced. To combine, see [`Headers::append`]
199    pub fn insert(
200        &mut self,
201        name: impl Into<HeaderName<'static>>,
202        values: impl Into<HeaderValues>,
203    ) -> &mut Self {
204        self.entry(name).insert(values);
205        self
206    }
207
208    /// Add a header value or header values into this header map if
209    /// and only if there is not already a header with the same name.
210    ///
211    /// Identical to [`headers.entry(name).or_insert(default)`][Entry::or_insert]
212    pub fn try_insert(
213        &mut self,
214        name: impl Into<HeaderName<'static>>,
215        values: impl Into<HeaderValues>,
216    ) -> &mut Self {
217        self.entry(name).or_insert(values);
218        self
219    }
220
221    /// if a key does not exist already, execute the provided function and insert a value
222    ///
223    /// Identical to
224    /// [`headers.entry(name).or_insert_with(values)`][Entry::or_insert_with]
225    pub fn try_insert_with<V>(
226        &mut self,
227        name: impl Into<HeaderName<'static>>,
228        values: impl FnOnce() -> V,
229    ) -> &mut HeaderValues
230    where
231        V: Into<HeaderValues>,
232    {
233        self.entry(name).or_insert_with(values)
234    }
235
236    /// Return a view into the entry for this header name, whether or not it is populated.
237    ///
238    /// See also [`Entry`]
239    pub fn entry(&mut self, name: impl Into<HeaderName<'static>>) -> Entry<'_> {
240        match name.into().0 {
241            HeaderNameInner::KnownHeader(known) => match self.known.entry(known) {
242                BTreeEntry::Vacant(vacant) => {
243                    Entry::Vacant(VacantEntry(VacantEntryInner::Known(vacant)))
244                }
245                BTreeEntry::Occupied(occupied) => {
246                    Entry::Occupied(OccupiedEntry(OccupiedEntryInner::Known(occupied)))
247                }
248            },
249
250            HeaderNameInner::UnknownHeader(unknown) => match self.unknown.entry(unknown) {
251                HashbrownEntry::Occupied(occupied) => {
252                    Entry::Occupied(OccupiedEntry(OccupiedEntryInner::Unknown(occupied)))
253                }
254
255                HashbrownEntry::Vacant(vacant) => {
256                    Entry::Vacant(VacantEntry(VacantEntryInner::Unknown(vacant)))
257                }
258            },
259        }
260    }
261
262    /// Retrieves a &str header value if there is at least one header
263    /// in the map with this name. If there are several headers with
264    /// the same name, this follows the behavior defined at
265    /// [`HeaderValues::one`]. Returns None if there is no header with
266    /// the provided header name.
267    pub fn get_str<'a>(&self, name: impl Into<HeaderName<'a>>) -> Option<&str> {
268        self.get_values(name).and_then(HeaderValues::as_str)
269    }
270
271    /// Retrieves a singular header value from this header map. If there are several headers with
272    /// the same name, this follows the behavior defined at [`HeaderValues::one`]. Returns None if
273    /// there is no header with the provided header name
274    pub fn get<'a>(&self, name: impl Into<HeaderName<'a>>) -> Option<&HeaderValue> {
275        self.get_values(name).and_then(HeaderValues::one)
276    }
277
278    /// Takes all headers with the provided header name out of this header map and returns
279    /// them. Returns None if the header did not have an entry in this map.
280    pub fn remove<'a>(&mut self, name: impl Into<HeaderName<'a>>) -> Option<HeaderValues> {
281        match name.into().0 {
282            HeaderNameInner::KnownHeader(known) => self.known.remove(&known),
283            HeaderNameInner::UnknownHeader(unknown) => self.unknown.remove(&&unknown),
284        }
285    }
286
287    /// Retrieves a reference to all header values with the provided
288    /// header name. If you expect there to be only one value, use
289    /// [`Headers::get`].
290    pub fn get_values<'a>(&self, name: impl Into<HeaderName<'a>>) -> Option<&HeaderValues> {
291        match name.into().0 {
292            HeaderNameInner::KnownHeader(known) => self.known.get(&known),
293            HeaderNameInner::UnknownHeader(unknown) => self.unknown.get(&&unknown),
294        }
295    }
296
297    /// Predicate function to check whether this header map contains
298    /// the provided header name. If you are using this to
299    /// conditionally insert a value, consider using
300    /// [`Headers::try_insert`] instead.
301    pub fn has_header<'a>(&self, name: impl Into<HeaderName<'a>>) -> bool {
302        match name.into().0 {
303            HeaderNameInner::KnownHeader(known) => self.known.contains_key(&known),
304            HeaderNameInner::UnknownHeader(unknown) => self.unknown.contains_key(&unknown),
305        }
306    }
307
308    /// Convenience function to check whether the value contained in
309    /// this header map for the provided name is
310    /// ascii-case-insensitively equal to the provided comparison
311    /// &str. Returns false if there is no value for the name
312    pub fn eq_ignore_ascii_case<'a>(
313        &'a self,
314        name: impl Into<HeaderName<'a>>,
315        needle: &str,
316    ) -> bool {
317        self.get_str(name)
318            .is_some_and(|v| v.eq_ignore_ascii_case(needle))
319    }
320
321    /// Chainable method to insert a header
322    pub fn with_inserted_header(
323        mut self,
324        name: impl Into<HeaderName<'static>>,
325        values: impl Into<HeaderValues>,
326    ) -> Self {
327        self.insert(name, values);
328        self
329    }
330
331    /// Chainable method to append a header
332    pub fn with_appended_header(
333        mut self,
334        name: impl Into<HeaderName<'static>>,
335        values: impl Into<HeaderValues>,
336    ) -> Self {
337        self.append(name, values);
338        self
339    }
340
341    /// Chainable method to remove a header
342    pub fn without_header<'a>(mut self, name: impl Into<HeaderName<'a>>) -> Self {
343        self.remove(name);
344        self
345    }
346
347    /// Chainable method to remove multiple headers by name
348    pub fn without_headers<'a, I, H>(mut self, names: I) -> Self
349    where
350        I: IntoIterator<Item = H>,
351        H: Into<HeaderName<'a>>,
352    {
353        self.remove_all(names);
354        self
355    }
356
357    /// remove multiple headers by name
358    pub fn remove_all<'a, I, H>(&mut self, names: I)
359    where
360        I: IntoIterator<Item = H>,
361        H: Into<HeaderName<'a>>,
362    {
363        for name in names {
364            self.remove(name);
365        }
366    }
367}
368
369impl<HN, HV> Extend<(HN, HV)> for Headers
370where
371    HN: Into<HeaderName<'static>>,
372    HV: Into<HeaderValues>,
373{
374    fn extend<T: IntoIterator<Item = (HN, HV)>>(&mut self, iter: T) {
375        for (name, values) in iter {
376            self.append(name, values);
377        }
378    }
379}
380
381impl<HN, HV> FromIterator<(HN, HV)> for Headers
382where
383    HN: Into<HeaderName<'static>>,
384    HV: Into<HeaderValues>,
385{
386    fn from_iter<T: IntoIterator<Item = (HN, HV)>>(iter: T) -> Self {
387        let iter = iter.into_iter();
388        let mut headers = Self::new();
389        for (name, values) in iter {
390            headers.append(name, values);
391        }
392
393        headers
394    }
395}
396
397impl<'a> IntoIterator for &'a Headers {
398    type IntoIter = Iter<'a>;
399    type Item = (HeaderName<'a>, &'a HeaderValues);
400
401    fn into_iter(self) -> Self::IntoIter {
402        self.into()
403    }
404}
405
406/// An owned iterator for Headers
407#[derive(Debug)]
408pub struct IntoIter {
409    known: btree_map::IntoIter<KnownHeaderName, HeaderValues>,
410    unknown: hash_map::IntoIter<UnknownHeaderName<'static>, HeaderValues>,
411}
412
413impl Iterator for IntoIter {
414    type Item = (HeaderName<'static>, HeaderValues);
415
416    fn next(&mut self) -> Option<Self::Item> {
417        let IntoIter { known, unknown } = self;
418        known
419            .next()
420            .map(|(k, v)| (HeaderName::from(k), v))
421            .or_else(|| unknown.next().map(|(k, v)| (HeaderName::from(k), v)))
422    }
423}
424
425impl From<Headers> for IntoIter {
426    fn from(value: Headers) -> Self {
427        Self {
428            known: value.known.into_iter(),
429            unknown: value.unknown.into_iter(),
430        }
431    }
432}
433
434/// A borrowed iterator for Headers
435#[derive(Debug)]
436pub struct Iter<'a> {
437    known: btree_map::Iter<'a, KnownHeaderName, HeaderValues>,
438    unknown: hash_map::Iter<'a, UnknownHeaderName<'static>, HeaderValues>,
439}
440
441impl<'a> From<&'a Headers> for Iter<'a> {
442    fn from(value: &'a Headers) -> Self {
443        Iter {
444            known: value.known.iter(),
445            unknown: value.unknown.iter(),
446        }
447    }
448}
449
450impl<'a> Iterator for Iter<'a> {
451    type Item = (HeaderName<'a>, &'a HeaderValues);
452
453    fn next(&mut self) -> Option<Self::Item> {
454        let Iter { known, unknown } = self;
455        known
456            .next()
457            .map(|(k, v)| (HeaderName::from(*k), v))
458            .or_else(|| unknown.next().map(|(k, v)| (HeaderName::from(&**k), v)))
459    }
460}
461
462impl IntoIterator for Headers {
463    type IntoIter = IntoIter;
464    type Item = (HeaderName<'static>, HeaderValues);
465
466    fn into_iter(self) -> Self::IntoIter {
467        self.into()
468    }
469}
trillium_http/headers.rs

trillium_http/
headers.rs
