wasmtime_wasi_http/

field_map.rs

use crate::{ForbiddenHeaderFilter, WasiHttpHooks};
use http::header::Entry;
use http::{HeaderMap, HeaderName, HeaderValue};
use std::fmt;
use std::ops::Deref;
use std::sync::Arc;
use wasmtime::Result;

/// A wrapper around [`http::HeaderMap`] which implements `wasi:http` semantics.
///
/// The main differences from [`http::HeaderMap`] and this type are:
///
/// * A slimmed down mutability API to just what `wasi:http` needs.
/// * `FieldMap` is cheaply clone-able with the internal `HeaderMap` being
///   behind an `Arc`.
/// * `FieldMap` is either immutable or mutable. Mutations on immutable values
///   are rejected with an error. Mutations on mutable values will never panic
///   unlike `HeaderMap` and additionally require a limit to be set on the size
///   of the map.
///
/// Overall the intention is that this is a slim wrapper around
/// [`http::HeaderMap`] with slightly different ownership, panic, and error
/// semantics.
#[derive(Debug, Clone)]
pub struct FieldMap {
    map: Arc<HeaderMap>,
    limit: Limit,
    size: usize,
}

#[derive(Debug, Clone)]
enum Limit {
    Mutable(usize),
    Immutable,
}

impl Default for FieldMap {
    fn default() -> Self {
        Self {
            map: Arc::new(HeaderMap::new()),
            size: 0,
            limit: Limit::Immutable,
        }
    }
}

impl FieldMap {
    /// Creates a new immutable `FieldMap` from the provided
    /// [`http::HeaderMap`].
    ///
    /// The returned value cannot be mutated and attempting to mutate it will
    /// return an error.
    pub fn new_immutable(hooks: &mut dyn WasiHttpHooks, mut map: HeaderMap) -> Self {
        // Strip out all forbidden headers from `map` to ensure they're never
        // able to enter into a WASI guest.
        remove_forbidden(&mut map, &hooks.forbidden_header_filter());

        let size = Self::content_size(&map);
        Self {
            map: Arc::new(map),
            size,
            limit: Limit::Immutable,
        }
    }

    /// Creates a new, empty, mutable `FieldMap`.
    ///
    /// Mutations are allowed on the returned value and up to `limit` bytes of
    /// memory (roughly) may be consumed by this map.
    pub fn new_mutable(limit: usize) -> Self {
        Self {
            map: Arc::new(HeaderMap::new()),
            size: 0,
            limit: Limit::Mutable(limit),
        }
    }

    /// Calculate the content size of a `HeaderMap`. This is a sum of the size
    /// of all of the keys and all of the values.
    fn content_size(map: &HeaderMap) -> usize {
        let mut sum = 0;
        for key in map.keys() {
            sum += header_name_size(key);
        }
        for value in map.values() {
            sum += header_value_size(value);
        }
        sum
    }

    /// Sets the header `key` to the `values` list provided.
    ///
    /// Removes the previous value, if any.
    ///
    /// If `values` is empty then this removes the header `key`.
    //
    // FIXME(WebAssembly/WASI#900): is this the right behavior?
    pub fn set(
        &mut self,
        hooks: &mut dyn WasiHttpHooks,
        key: String,
        values: Vec<Vec<u8>>,
    ) -> Result<(), FieldMapError> {
        let key = key.parse()?;
        if (hooks.forbidden_header_filter())(&key) {
            return Err(FieldMapError::Forbidden);
        }
        let (map, limit, size) = self.mutable()?;

        // `Content-Length` may only have a single value, see `append`.
        if key == http::header::CONTENT_LENGTH && values.len() > 1 {
            return Err(FieldMapError::InvalidHeaderValue);
        }
        let values = values
            .into_iter()
            .map(|v| parse_header_value(&key, v))
            .collect::<Result<Vec<_>, _>>()?;

        let key_size = header_name_size(&key);
        let prev_values_size: usize = map.get_all(&key).iter().map(header_value_size).sum();
        let mut new_size = *size - prev_values_size;
        new_size += values.iter().map(header_value_size).sum::<usize>();

        let mut values = values.into_iter();
        let Some(first) = values.next() else {
            if map.remove(&key).is_some() {
                *size = new_size - key_size;
            }
            return Ok(());
        };
        if !map.contains_key(&key) {
            new_size += key_size;
        }
        if new_size > limit {
            return Err(FieldMapError::TotalSizeTooBig);
        }

        // `try_entry` reserves a slot even for an existing key, therefore remove it first
        map.remove(&key);
        // `try_entry` reports MaxSizeReached as InvalidHeaderName
        let Entry::Vacant(entry) = map
            .try_entry(key)
            .map_err(|_| FieldMapError::TooManyFields)?
        else {
            // `key` was just removed above
            unreachable!();
        };
        let mut entry = entry.try_insert_entry(first)?;
        for value in values {
            entry.append(value);
        }
        *size = new_size;
        Ok(())
    }

    /// Remove all values associated with a key in a map.
    ///
    /// Returns an empty list if the key is not already present within the map.
    pub fn remove_all(
        &mut self,
        hooks: &mut dyn WasiHttpHooks,
        key: String,
    ) -> Result<Vec<HeaderValue>, FieldMapError> {
        let key = key.parse()?;
        if (hooks.forbidden_header_filter())(&key) {
            return Err(FieldMapError::Forbidden);
        }
        let (map, _limit, size) = self.mutable()?;
        // try_entry misreports a full map as InvalidHeaderName.
        let values = map.get_all(&key).iter().cloned().collect::<Vec<_>>();
        if values.is_empty() {
            return Ok(Vec::new());
        }
        let mut removed = header_name_size(&key);
        for v in values.iter() {
            removed += header_value_size(v);
        }
        let _ = map.remove(&key);
        *size -= removed;
        Ok(values)
    }

    fn mutable(&mut self) -> Result<(&mut HeaderMap, usize, &mut usize), FieldMapError> {
        match self.limit {
            Limit::Immutable => Err(FieldMapError::Immutable),
            Limit::Mutable(limit) => Ok((Arc::make_mut(&mut self.map), limit, &mut self.size)),
        }
    }

    /// Add a value associated with a key to the map.
    ///
    /// If `key` is already present within the map then `value` is appended to
    /// the list of values it already has.
    pub fn append(
        &mut self,
        hooks: &mut dyn WasiHttpHooks,
        key: String,
        value: Vec<u8>,
    ) -> Result<bool, FieldMapError> {
        let key = key.parse()?;
        if (hooks.forbidden_header_filter())(&key) {
            return Err(FieldMapError::Forbidden);
        }
        let value = parse_header_value(&key, value)?;
        // Multiple `Content-Length` values make message framing ambiguous, as
        // different parties may disagree on which one to use, so only a single
        // value is permitted.
        if key == http::header::CONTENT_LENGTH && self.contains_key(&key) {
            return Err(FieldMapError::InvalidHeaderValue);
        }
        self.append_raw(key, value)
    }

    /// Add a value associated with a key to the map.
    ///
    /// If `key` is already present within the map then `value` is appended to
    /// the list of values it already has.
    ///  TODO
    pub fn append_raw(
        &mut self,
        key: HeaderName,
        value: HeaderValue,
    ) -> Result<bool, FieldMapError> {
        let (map, limit, size) = self.mutable()?;
        let key_size = header_name_size(&key);
        let val_size = header_value_size(&value);
        let new_size = if !map.contains_key(&key) {
            *size + key_size + val_size
        } else {
            *size + val_size
        };
        if new_size > limit {
            return Err(FieldMapError::TotalSizeTooBig);
        }
        let already_present = map.try_append(key, value)?;
        *size = new_size;
        Ok(already_present)
    }

    /// Flags this map as mutable, allowing mutations which can allocate as much
    /// as `limit` memory, in bytes, for this entire map (roughly).
    pub fn set_mutable(&mut self, limit: usize) {
        self.limit = Limit::Mutable(limit);
    }

    /// Flags this map as immutable, forbidding all further mutations.
    pub fn set_immutable(&mut self) {
        self.limit = Limit::Immutable;
    }
}

/// Returns the size, in accounting cost, to consider for `name`.
///
/// This includes both the byte length of the `name` itself as well as the size
/// of the data structure itself as it'll reside within a `HeaderMap`.
fn header_name_size(name: &HeaderName) -> usize {
    name.as_str().len() + size_of::<HeaderName>()
}

/// Same as `header_name_size`, but for values.
///
/// This notably includes the size of `HeaderValue` itself to ensure that all
/// headers have a nonzero size as otherwise this would never limit addition of
/// an empty header value.
fn header_value_size(value: &HeaderValue) -> usize {
    value.len() + size_of::<HeaderValue>()
}

// Note that `DerefMut` is specifically omitted here to force all mutations
// through the `FieldMap` wrapper.
impl Deref for FieldMap {
    type Target = HeaderMap;

    fn deref(&self) -> &HeaderMap {
        &self.map
    }
}

impl From<FieldMap> for HeaderMap {
    fn from(map: FieldMap) -> Self {
        Arc::unwrap_or_clone(map.map)
    }
}

/// Errors that can happen when mutating/operating on a [`FieldMap`].
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum FieldMapError {
    /// A mutation was attempted when the map is not mutable.
    Immutable,
    /// The map has too many fields and is not allowed to add more.
    ///
    /// Note that this is currently a limitation inherited from
    /// [`http::HeaderMap`].
    TooManyFields,
    /// The map's total size, of keys and values, is too large.
    TotalSizeTooBig,
    /// An invalid header name was attempted to be added.
    InvalidHeaderName,
    /// An invalid header value was attempted to be added.
    InvalidHeaderValue,
    /// A forbidden header name was used.
    Forbidden,
}

impl fmt::Display for FieldMapError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s = match self {
            FieldMapError::Immutable => "cannot mutate an immutable field map",
            FieldMapError::TooManyFields => "too many fields in the field map",
            FieldMapError::TotalSizeTooBig => "total size of fields exceeds limit",
            FieldMapError::InvalidHeaderName => "invalid header name",
            FieldMapError::InvalidHeaderValue => "invalid header value",
            FieldMapError::Forbidden => "forbidden header name",
        };
        f.write_str(s)
    }
}

impl std::error::Error for FieldMapError {}

impl From<http::header::MaxSizeReached> for FieldMapError {
    fn from(_: http::header::MaxSizeReached) -> Self {
        Self::TooManyFields
    }
}

impl From<http::header::InvalidHeaderName> for FieldMapError {
    fn from(_: http::header::InvalidHeaderName) -> Self {
        Self::InvalidHeaderName
    }
}

impl From<http::header::InvalidHeaderValue> for FieldMapError {
    fn from(_: http::header::InvalidHeaderValue) -> Self {
        Self::InvalidHeaderValue
    }
}

fn parse_header_value(
    name: &http::HeaderName,
    value: Vec<u8>,
) -> Result<http::HeaderValue, FieldMapError> {
    if name == http::header::CONTENT_LENGTH {
        let s = str::from_utf8(value.as_ref()).or(Err(FieldMapError::InvalidHeaderValue))?;
        // RFC 9110 defines `Content-Length` as `1*DIGIT`. `u64`'s `FromStr` is
        // more lenient and also accepts a leading `+`, so reject anything that
        // isn't a non-empty run of decimal digits.
        if s.is_empty() || !s.bytes().all(|b| b.is_ascii_digit()) {
            return Err(FieldMapError::InvalidHeaderValue);
        }
        let v: u64 = s.parse().or(Err(FieldMapError::InvalidHeaderValue))?;
        Ok(v.into())
    } else {
        let value = value.try_into()?;
        Ok(value)
    }
}

#[cfg(test)]
mod tests {
    use super::{FieldMap, FieldMapError, parse_header_value};
    use crate::default_hooks;
    use http::header::{CONTENT_LENGTH, CONTENT_TYPE};
    use http::{HeaderName, HeaderValue};

    #[test]
    fn test_immutable() {
        let mut map = FieldMap::default();
        assert_eq!(
            map.set(default_hooks(), "foo".to_owned(), vec![b"bar".to_vec()]),
            Err(FieldMapError::Immutable)
        );
        assert_eq!(
            map.append(default_hooks(), "foo".to_owned(), b"bar".to_vec()),
            Err(FieldMapError::Immutable)
        );
        assert_eq!(
            map.remove_all(default_hooks(), "foo".to_owned()),
            Err(FieldMapError::Immutable)
        );
    }

    #[test]
    fn test_limits() {
        let mut map = FieldMap::new_mutable(100);
        loop {
            match map.append(default_hooks(), "foo".to_owned(), b"bar".to_vec()) {
                Ok(_) => {}
                Err(FieldMapError::TotalSizeTooBig) => break,
                Err(e) => panic!("unexpected error: {e}"),
            }
        }

        map = FieldMap::new_mutable(100);
        for i in 0.. {
            match map.set(
                default_hooks(),
                "foo".to_owned(),
                (0..i).map(|j| format!("bar{j}").into_bytes()).collect(),
            ) {
                Ok(_) => {}
                Err(FieldMapError::TotalSizeTooBig) => break,
                Err(e) => panic!("unexpected error: {e}"),
            }
        }

        map = FieldMap::new_mutable(100);
        for i in 0.. {
            match map.set(default_hooks(), format!("foo{i}"), vec![b"bar".to_vec()]) {
                Ok(_) => {}
                Err(FieldMapError::TotalSizeTooBig) => break,
                Err(e) => panic!("unexpected error: {e}"),
            }
        }
    }

    #[test]
    fn test_size() -> Result<(), FieldMapError> {
        let mut map = FieldMap::new_mutable(2000);
        let name = "foo".to_owned();
        let hooks = default_hooks();

        map.append(hooks, name.clone(), b"bar".to_vec())?;
        assert!(map.size > 0);
        map.remove_all(hooks, name.clone())?;
        assert_eq!(map.size, 0);

        map.set(hooks, name.clone(), vec![b"bar".to_vec()])?;
        assert!(map.size > 0);
        map.remove_all(hooks, name.clone())?;
        assert_eq!(map.size, 0);

        map.set(hooks, name.clone(), vec![])?;
        assert_eq!(map.size, 0);
        map.set(hooks, name.clone(), vec![b"bar".to_vec()])?;
        assert!(map.size > 0);
        map.set(hooks, name.clone(), vec![])?;
        assert_eq!(map.size, 0);

        map.set(hooks, name.clone(), vec![b"bar".to_vec()])?;
        assert!(map.size > 0);
        map.set(hooks, name.clone(), vec![b"bar".to_vec(), b"baz".to_vec()])?;
        assert!(map.size > 0);
        map.remove_all(hooks, name.clone())?;
        assert_eq!(map.size, 0);

        Ok(())
    }

    #[test]
    fn set_empty_over_limit() -> Result<(), FieldMapError> {
        let mut map = FieldMap::new_mutable(usize::MAX);
        let hooks = default_hooks();
        map.set(hooks, "foo".to_owned(), vec![b"bar".to_vec()])?;
        map.set_mutable(0);

        map.set(hooks, "new".to_owned(), vec![])?;
        map.set(hooks, "foo".to_owned(), vec![])?;
        assert_eq!(map.size, 0);

        Ok(())
    }

    #[test]
    fn append_raw_too_many_fields_keeps_size() {
        let mut map = FieldMap::new_mutable(usize::MAX);
        let mut i = 0;
        let err = loop {
            let name = HeaderName::try_from(format!("foo{i}")).unwrap();
            match map.append_raw(name, HeaderValue::from_static("bar")) {
                Ok(_) => i += 1,
                Err(e) => break e,
            }
        };
        assert_eq!(err, FieldMapError::TooManyFields);
        assert_eq!(map.size, FieldMap::content_size(&map));
    }

    #[test]
    fn content_length_rejects_non_digits() {
        assert!(parse_header_value(&CONTENT_LENGTH, b"0".to_vec()).is_ok());
        assert!(parse_header_value(&CONTENT_LENGTH, b"1234".to_vec()).is_ok());

        // `u64::from_str` accepts these but they are not `1*DIGIT` per RFC 9110.
        assert!(parse_header_value(&CONTENT_LENGTH, b"+5".to_vec()).is_err());
        assert!(parse_header_value(&CONTENT_LENGTH, b"-5".to_vec()).is_err());
        assert!(parse_header_value(&CONTENT_LENGTH, b" 5".to_vec()).is_err());
        assert!(parse_header_value(&CONTENT_LENGTH, b"".to_vec()).is_err());

        // other header names are unaffected
        assert!(parse_header_value(&CONTENT_TYPE, b"text/plain".to_vec()).is_ok());
    }

    #[test]
    fn content_length_rejects_multiple_values() {
        let hooks = default_hooks();
        let cl = || CONTENT_LENGTH.as_str().to_owned();

        let mut map = FieldMap::new_mutable(2000);
        map.append(hooks, cl(), b"5".to_vec()).unwrap();
        assert_eq!(
            map.append(hooks, cl(), b"6".to_vec()),
            Err(FieldMapError::InvalidHeaderValue)
        );
        assert_eq!(
            map.append(hooks, cl(), b"5".to_vec()),
            Err(FieldMapError::InvalidHeaderValue)
        );
        assert_eq!(
            map.append(hooks, "Content-Length".to_owned(), b"6".to_vec()),
            Err(FieldMapError::InvalidHeaderValue)
        );
        assert_eq!(map.get_all(CONTENT_LENGTH).iter().count(), 1);
        assert_eq!(map.size, FieldMap::content_size(&map.map));

        assert_eq!(
            map.set(hooks, cl(), vec![b"5".to_vec(), b"6".to_vec()]),
            Err(FieldMapError::InvalidHeaderValue)
        );
        assert_eq!(
            map.set(hooks, cl(), vec![b"5".to_vec(), b"5".to_vec()]),
            Err(FieldMapError::InvalidHeaderValue)
        );
        assert_eq!(map[CONTENT_LENGTH], "5");

        // Replacing the single value is still fine.
        map.set(hooks, cl(), vec![b"6".to_vec()]).unwrap();
        assert_eq!(map[CONTENT_LENGTH], "6");
        assert_eq!(map.get_all(CONTENT_LENGTH).iter().count(), 1);
        assert_eq!(map.size, FieldMap::content_size(&map.map));
    }

    #[test]
    fn full_map_is_not_invalid_syntax() {
        // try_entry reports MaxSizeReached as InvalidHeaderName (http map.rs:1149),
        // so a full map must report valid names as too many fields instead.
        let mut map = FieldMap::new_mutable(usize::MAX);
        let mut i = 0;
        let err = loop {
            assert_eq!(map.size, FieldMap::content_size(&map.map));
            match map.set(default_hooks(), format!("k{i}"), vec![b"v".to_vec()]) {
                Ok(()) => i += 1,
                Err(e) => break e,
            }
        };
        assert_eq!(err, FieldMapError::TooManyFields);
        assert_eq!(map.size, FieldMap::content_size(&map.map));

        assert_eq!(
            map.set(
                default_hooks(),
                "new".to_owned(),
                vec![b"a".to_vec(), b"b".to_vec()],
            ),
            Err(FieldMapError::TooManyFields)
        );
        assert!(!map.contains_key("new"));
        assert_eq!(map.size, FieldMap::content_size(&map.map));

        assert_eq!(
            map.set(
                default_hooks(),
                "k0".to_owned(),
                vec![b"foo".to_vec(), b"bar".to_vec()]
            ),
            Ok(())
        );
        assert_eq!(map.get("k0"), Some(&HeaderValue::from_static("foo")));
        assert_eq!(
            map.get_all("k0").iter().collect::<Vec<_>>(),
            [
                HeaderValue::from_static("foo"),
                HeaderValue::from_static("bar")
            ]
        );
        assert_eq!(map.size, FieldMap::content_size(&map.map));

        assert_eq!(map.set(default_hooks(), "k0".to_owned(), vec![]), Ok(()));
        assert!(!map.contains_key("k0"));
        assert_eq!(map.size, FieldMap::content_size(&map.map));

        assert_eq!(
            map.set(default_hooks(), "k0".to_owned(), vec![b"v".to_vec()]),
            Ok(())
        );
        assert_eq!(map.get("k0"), Some(&HeaderValue::from_static("v")));
        assert_eq!(map.size, FieldMap::content_size(&map.map));

        // Removing a field never needs new capacity.
        assert_eq!(
            map.remove_all(default_hooks(), "k0".to_owned())
                .map(|v| v.len()),
            Ok(1)
        );
        assert_eq!(map.size, FieldMap::content_size(&map.map));
    }
}

/// Removes all fields from `map` whose name `is_forbidden`.
pub(crate) fn remove_forbidden(map: &mut HeaderMap, is_forbidden: &ForbiddenHeaderFilter) {
    let forbidden_keys = Vec::from_iter(map.keys().filter(|name| is_forbidden(name)).cloned());
    for name in forbidden_keys {
        map.remove(&name);
    }
}