wasmtime_wasi_http/
field_map.rs

1use crate::WasiHttpHooks;
2use http::header::Entry;
3use http::{HeaderMap, HeaderName, HeaderValue};
4use std::fmt;
5use std::ops::Deref;
6use std::sync::Arc;
7use wasmtime::Result;
8
9/// A wrapper around [`http::HeaderMap`] which implements `wasi:http` semantics.
10///
11/// The main differences from [`http::HeaderMap`] and this type are:
12///
13/// * A slimmed down mutability API to just what `wasi:http` needs.
14/// * `FieldMap` is cheaply clone-able with the internal `HeaderMap` being
15///   behind an `Arc`.
16/// * `FieldMap` is either immutable or mutable. Mutations on immutable values
17///   are rejected with an error. Mutations on mutable values will never panic
18///   unlike `HeaderMap` and additionally require a limit to be set on the size
19///   of the map.
20///
21/// Overall the intention is that this is a slim wrapper around
22/// [`http::HeaderMap`] with slightly different ownership, panic, and error
23/// semantics.
24#[derive(Debug, Clone)]
25pub struct FieldMap {
26    map: Arc<HeaderMap>,
27    limit: Limit,
28    size: usize,
29}
30
31#[derive(Debug, Clone)]
32enum Limit {
33    Mutable(usize),
34    Immutable,
35}
36
37impl Default for FieldMap {
38    fn default() -> Self {
39        Self {
40            map: Arc::new(HeaderMap::new()),
41            size: 0,
42            limit: Limit::Immutable,
43        }
44    }
45}
46
47impl FieldMap {
48    /// Creates a new immutable `FieldMap` from the provided
49    /// [`http::HeaderMap`].
50    ///
51    /// The returned value cannot be mutated and attempting to mutate it will
52    /// return an error.
53    pub fn new_immutable(hooks: &mut dyn WasiHttpHooks, mut map: HeaderMap) -> Self {
54        // Strip out all forbidden headers from `map` to ensure they're never
55        // able to enter into a WASI guest.
56        let forbidden_keys = Vec::from_iter(map.keys().filter_map(|name| {
57            if hooks.is_forbidden_header(name) {
58                Some(name.clone())
59            } else {
60                None
61            }
62        }));
63        for name in forbidden_keys {
64            map.remove(&name);
65        }
66
67        let size = Self::content_size(&map);
68        Self {
69            map: Arc::new(map),
70            size,
71            limit: Limit::Immutable,
72        }
73    }
74
75    /// Creates a new, empty, mutable `FieldMap`.
76    ///
77    /// Mutations are allowed on the returned value and up to `limit` bytes of
78    /// memory (roughly) may be consumed by this map.
79    pub fn new_mutable(limit: usize) -> Self {
80        Self {
81            map: Arc::new(HeaderMap::new()),
82            size: 0,
83            limit: Limit::Mutable(limit),
84        }
85    }
86
87    /// Calculate the content size of a `HeaderMap`. This is a sum of the size
88    /// of all of the keys and all of the values.
89    pub(crate) fn content_size(map: &HeaderMap) -> usize {
90        let mut sum = 0;
91        for key in map.keys() {
92            sum += header_name_size(key);
93        }
94        for value in map.values() {
95            sum += header_value_size(value);
96        }
97        sum
98    }
99
100    /// Sets the header `key` to the `values` list provided.
101    ///
102    /// Removes the previous value, if any.
103    ///
104    /// If `values` is empty then this removes the header `key`.
105    //
106    // FIXME(WebAssembly/WASI#900): is this the right behavior?
107    pub fn set(
108        &mut self,
109        hooks: &mut dyn WasiHttpHooks,
110        key: String,
111        values: Vec<Vec<u8>>,
112    ) -> Result<(), FieldMapError> {
113        let key = key.parse()?;
114        if hooks.is_forbidden_header(&key) {
115            return Err(FieldMapError::Forbidden);
116        }
117        let (map, limit, size) = self.mutable()?;
118        let key_size = header_name_size(&key);
119        let values = values
120            .into_iter()
121            .map(|v| parse_header_value(&key, v))
122            .collect::<Result<Vec<_>, _>>()?;
123        let values_size = values.iter().map(header_value_size).sum::<usize>();
124        let mut values = values.into_iter();
125        let mut entry = match map.try_entry(key)? {
126            Entry::Vacant(e) => match values.next() {
127                Some(v) => {
128                    update_size(size, limit, *size + values_size + key_size)?;
129                    e.try_insert_entry(v)?
130                }
131                None => return Ok(()),
132            },
133            Entry::Occupied(mut e) => {
134                let prev_values_size = e.iter().map(header_value_size).sum::<usize>();
135                let _prev = match values.next() {
136                    Some(v) => {
137                        update_size(size, limit, *size - prev_values_size + values_size)?;
138                        e.insert(v);
139                    }
140                    None => {
141                        update_size(size, limit, *size - prev_values_size - key_size)?;
142                        e.remove();
143                        return Ok(());
144                    }
145                };
146                e
147            }
148        };
149        for value in values {
150            entry.append(value);
151        }
152        Ok(())
153    }
154
155    /// Remove all values associated with a key in a map.
156    ///
157    /// Returns an empty list if the key is not already present within the map.
158    pub fn remove_all(
159        &mut self,
160        hooks: &mut dyn WasiHttpHooks,
161        key: String,
162    ) -> Result<Vec<HeaderValue>, FieldMapError> {
163        let key = key.parse()?;
164        if hooks.is_forbidden_header(&key) {
165            return Err(FieldMapError::Forbidden);
166        }
167        let (map, _limit, size) = self.mutable()?;
168        match map.try_entry(key)? {
169            Entry::Vacant { .. } => Ok(Vec::new()),
170            Entry::Occupied(e) => {
171                let (name, value_drain) = e.remove_entry_mult();
172                let mut removed = header_name_size(&name);
173                let values = value_drain.collect::<Vec<_>>();
174                for v in values.iter() {
175                    removed += header_value_size(v);
176                }
177                *size -= removed;
178                Ok(values)
179            }
180        }
181    }
182
183    fn mutable(&mut self) -> Result<(&mut HeaderMap, usize, &mut usize), FieldMapError> {
184        match self.limit {
185            Limit::Immutable => Err(FieldMapError::Immutable),
186            Limit::Mutable(limit) => Ok((Arc::make_mut(&mut self.map), limit, &mut self.size)),
187        }
188    }
189
190    /// Add a value associated with a key to the map.
191    ///
192    /// If `key` is already present within the map then `value` is appended to
193    /// the list of values it already has.
194    pub fn append(
195        &mut self,
196        hooks: &mut dyn WasiHttpHooks,
197        key: String,
198        value: Vec<u8>,
199    ) -> Result<bool, FieldMapError> {
200        let key = key.parse()?;
201        if hooks.is_forbidden_header(&key) {
202            return Err(FieldMapError::Forbidden);
203        }
204        let value = parse_header_value(&key, value)?;
205        self.append_raw(key, value)
206    }
207
208    /// Add a value associated with a key to the map.
209    ///
210    /// If `key` is already present within the map then `value` is appended to
211    /// the list of values it already has.
212    ///  TODO
213    pub fn append_raw(
214        &mut self,
215        key: HeaderName,
216        value: HeaderValue,
217    ) -> Result<bool, FieldMapError> {
218        let (map, limit, size) = self.mutable()?;
219        let key_size = header_name_size(&key);
220        let val_size = header_value_size(&value);
221        let new_size = if !map.contains_key(&key) {
222            *size + key_size + val_size
223        } else {
224            *size + val_size
225        };
226        update_size(size, limit, new_size)?;
227        let already_present = map.try_append(key, value)?;
228        self.size = new_size;
229        Ok(already_present)
230    }
231
232    /// Flags this map as mutable, allowing mutations which can allocate as much
233    /// as `limit` memory, in bytes, for this entire map (roughly).
234    pub fn set_mutable(&mut self, limit: usize) {
235        self.limit = Limit::Mutable(limit);
236    }
237
238    /// Flags this map as immutable, forbidding all further mutations.
239    pub fn set_immutable(&mut self) {
240        self.limit = Limit::Immutable;
241    }
242}
243
244/// Returns the size, in accounting cost, to consider for `name`.
245///
246/// This includes both the byte length of the `name` itself as well as the size
247/// of the data structure itself as it'll reside within a `HeaderMap`.
248fn header_name_size(name: &HeaderName) -> usize {
249    name.as_str().len() + size_of::<HeaderName>()
250}
251
252/// Same as `header_name_size`, but for values.
253///
254/// This notably includes the size of `HeaderValue` itself to ensure that all
255/// headers have a nonzero size as otherwise this would never limit addition of
256/// an empty header value.
257fn header_value_size(value: &HeaderValue) -> usize {
258    value.len() + size_of::<HeaderValue>()
259}
260
261fn update_size(size: &mut usize, limit: usize, new: usize) -> Result<(), FieldMapError> {
262    if new > limit {
263        Err(FieldMapError::TotalSizeTooBig)
264    } else {
265        *size = new;
266        Ok(())
267    }
268}
269
270// Note that `DerefMut` is specifically omitted here to force all mutations
271// through the `FieldMap` wrapper.
272impl Deref for FieldMap {
273    type Target = HeaderMap;
274
275    fn deref(&self) -> &HeaderMap {
276        &self.map
277    }
278}
279
280impl From<FieldMap> for HeaderMap {
281    fn from(map: FieldMap) -> Self {
282        Arc::unwrap_or_clone(map.map)
283    }
284}
285
286/// Errors that can happen when mutating/operating on a [`FieldMap`].
287#[derive(Debug, PartialEq, Eq, Clone, Copy)]
288pub enum FieldMapError {
289    /// A mutation was attempted when the map is not mutable.
290    Immutable,
291    /// The map has too many fields and is not allowed to add more.
292    ///
293    /// Note that this is currently a limitation inherited from
294    /// [`http::HeaderMap`].
295    TooManyFields,
296    /// The map's total size, of keys and values, is too large.
297    TotalSizeTooBig,
298    /// An invalid header name was attempted to be added.
299    InvalidHeaderName,
300    /// An invalid header value was attempted to be added.
301    InvalidHeaderValue,
302    /// A forbidden header name was used.
303    Forbidden,
304}
305
306impl fmt::Display for FieldMapError {
307    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
308        let s = match self {
309            FieldMapError::Immutable => "cannot mutate an immutable field map",
310            FieldMapError::TooManyFields => "too many fields in the field map",
311            FieldMapError::TotalSizeTooBig => "total size of fields exceeds limit",
312            FieldMapError::InvalidHeaderName => "invalid header name",
313            FieldMapError::InvalidHeaderValue => "invalid header value",
314            FieldMapError::Forbidden => "forbidden header name",
315        };
316        f.write_str(s)
317    }
318}
319
320impl std::error::Error for FieldMapError {}
321
322impl From<http::header::MaxSizeReached> for FieldMapError {
323    fn from(_: http::header::MaxSizeReached) -> Self {
324        Self::TooManyFields
325    }
326}
327
328impl From<http::header::InvalidHeaderName> for FieldMapError {
329    fn from(_: http::header::InvalidHeaderName) -> Self {
330        Self::InvalidHeaderName
331    }
332}
333
334impl From<http::header::InvalidHeaderValue> for FieldMapError {
335    fn from(_: http::header::InvalidHeaderValue) -> Self {
336        Self::InvalidHeaderValue
337    }
338}
339
340fn parse_header_value(
341    name: &http::HeaderName,
342    value: Vec<u8>,
343) -> Result<http::HeaderValue, FieldMapError> {
344    if name == http::header::CONTENT_LENGTH {
345        let s = str::from_utf8(value.as_ref()).or(Err(FieldMapError::InvalidHeaderValue))?;
346        // RFC 9110 defines `Content-Length` as `1*DIGIT`. `u64`'s `FromStr` is
347        // more lenient and also accepts a leading `+`, so reject anything that
348        // isn't a non-empty run of decimal digits.
349        if s.is_empty() || !s.bytes().all(|b| b.is_ascii_digit()) {
350            return Err(FieldMapError::InvalidHeaderValue);
351        }
352        let v: u64 = s.parse().or(Err(FieldMapError::InvalidHeaderValue))?;
353        Ok(v.into())
354    } else {
355        let value = value.try_into()?;
356        Ok(value)
357    }
358}
359
360#[cfg(test)]
361mod tests {
362    use super::{FieldMap, FieldMapError, parse_header_value};
363    use crate::default_hooks;
364    use http::header::{CONTENT_LENGTH, CONTENT_TYPE};
365
366    #[test]
367    fn test_immutable() {
368        let mut map = FieldMap::default();
369        assert_eq!(
370            map.set(default_hooks(), "foo".to_owned(), vec![b"bar".to_vec()]),
371            Err(FieldMapError::Immutable)
372        );
373        assert_eq!(
374            map.append(default_hooks(), "foo".to_owned(), b"bar".to_vec()),
375            Err(FieldMapError::Immutable)
376        );
377        assert_eq!(
378            map.remove_all(default_hooks(), "foo".to_owned()),
379            Err(FieldMapError::Immutable)
380        );
381    }
382
383    #[test]
384    fn test_limits() {
385        let mut map = FieldMap::new_mutable(100);
386        loop {
387            match map.append(default_hooks(), "foo".to_owned(), b"bar".to_vec()) {
388                Ok(_) => {}
389                Err(FieldMapError::TotalSizeTooBig) => break,
390                Err(e) => panic!("unexpected error: {e}"),
391            }
392        }
393
394        map = FieldMap::new_mutable(100);
395        for i in 0.. {
396            match map.set(
397                default_hooks(),
398                "foo".to_owned(),
399                (0..i).map(|j| format!("bar{j}").into_bytes()).collect(),
400            ) {
401                Ok(_) => {}
402                Err(FieldMapError::TotalSizeTooBig) => break,
403                Err(e) => panic!("unexpected error: {e}"),
404            }
405        }
406
407        map = FieldMap::new_mutable(100);
408        for i in 0.. {
409            match map.set(default_hooks(), format!("foo{i}"), vec![b"bar".to_vec()]) {
410                Ok(_) => {}
411                Err(FieldMapError::TotalSizeTooBig) => break,
412                Err(e) => panic!("unexpected error: {e}"),
413            }
414        }
415    }
416
417    #[test]
418    fn test_size() -> Result<(), FieldMapError> {
419        let mut map = FieldMap::new_mutable(2000);
420        let name = "foo".to_owned();
421        let hooks = default_hooks();
422
423        map.append(hooks, name.clone(), b"bar".to_vec())?;
424        assert!(map.size > 0);
425        map.remove_all(hooks, name.clone())?;
426        assert_eq!(map.size, 0);
427
428        map.set(hooks, name.clone(), vec![b"bar".to_vec()])?;
429        assert!(map.size > 0);
430        map.remove_all(hooks, name.clone())?;
431        assert_eq!(map.size, 0);
432
433        map.set(hooks, name.clone(), vec![])?;
434        assert_eq!(map.size, 0);
435        map.set(hooks, name.clone(), vec![b"bar".to_vec()])?;
436        assert!(map.size > 0);
437        map.set(hooks, name.clone(), vec![])?;
438        assert_eq!(map.size, 0);
439
440        map.set(hooks, name.clone(), vec![b"bar".to_vec()])?;
441        assert!(map.size > 0);
442        map.set(hooks, name.clone(), vec![b"bar".to_vec(), b"baz".to_vec()])?;
443        assert!(map.size > 0);
444        map.remove_all(hooks, name.clone())?;
445        assert_eq!(map.size, 0);
446
447        Ok(())
448    }
449
450    #[test]
451    fn content_length_rejects_non_digits() {
452        assert!(parse_header_value(&CONTENT_LENGTH, b"0".to_vec()).is_ok());
453        assert!(parse_header_value(&CONTENT_LENGTH, b"1234".to_vec()).is_ok());
454
455        // `u64::from_str` accepts these but they are not `1*DIGIT` per RFC 9110.
456        assert!(parse_header_value(&CONTENT_LENGTH, b"+5".to_vec()).is_err());
457        assert!(parse_header_value(&CONTENT_LENGTH, b"-5".to_vec()).is_err());
458        assert!(parse_header_value(&CONTENT_LENGTH, b" 5".to_vec()).is_err());
459        assert!(parse_header_value(&CONTENT_LENGTH, b"".to_vec()).is_err());
460
461        // other header names are unaffected
462        assert!(parse_header_value(&CONTENT_TYPE, b"text/plain".to_vec()).is_ok());
463    }
464}
wasmtime_wasi_http/field_map.rs

wasmtime_wasi_http/
field_map.rs
