hpx 2.4.10

//! HTTP Cookies

use std::{convert::TryInto, fmt, sync::Arc, time::SystemTime};

use arc_swap::ArcSwap;
use bytes::Bytes;
use cookie::{Cookie as RawCookie, CookieJar, Expiration, SameSite};
use http::Uri;

use crate::{
    IntoUri,
    error::Error,
    ext::UriExt,
    hash::{HASHER, HashMap},
    header::HeaderValue,
};

/// Cookie header values in two forms.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum Cookies {
    /// All cookies combined into one header (compressed).
    Compressed(HeaderValue),

    /// Each cookie sent as its own header (uncompressed).
    Uncompressed(Vec<HeaderValue>),

    /// No cookies.
    Empty,
}

/// Actions for a persistent cookie store providing session support.
pub trait CookieStore: Send + Sync {
    /// Store a set of Set-Cookie header values received from `uri`
    fn set_cookies(&self, cookie_headers: &mut dyn Iterator<Item = &HeaderValue>, uri: &Uri);

    /// Get any [`Cookies`] in the store for `uri`
    fn cookies(&self, uri: &Uri) -> Cookies;
}

/// Trait for converting types into a shared cookie store ([`Arc<dyn CookieStore>`]).
///
/// Implemented for any [`CookieStore`] type, [`Arc<T>`] where `T: CookieStore`, and [`Arc<dyn
/// CookieStore>`]. Enables ergonomic conversion to a trait object for use in APIs without manual
/// boxing.
pub trait IntoCookieStore {
    /// Converts the implementor into an [`Arc<dyn CookieStore>`].
    ///
    /// This method allows ergonomic conversion of concrete cookie stores, [`Arc<T>`], or
    /// existing [`Arc<dyn CookieStore>`] into a trait object suitable for APIs that expect
    /// a shared cookie store.
    fn into_cookie_store(self) -> Arc<dyn CookieStore>;
}

/// Trait for converting types into an owned cookie ([`Cookie<'static>`]).
pub trait IntoCookie {
    /// Converts the implementor into an optional owned [`Cookie<'static>`].
    fn into_cookie(self) -> Option<Cookie<'static>>;
}

/// A single HTTP cookie.
#[derive(Debug, Clone)]
pub struct Cookie<'a>(RawCookie<'a>);

/// A good default `CookieStore` implementation.
///
/// This is the implementation used when simply calling `cookie_store(true)`.
/// This type is exposed to allow creating one and filling it with some
/// existing cookies more easily, before creating a [`crate::Client`].
///
/// Uses `ArcSwap` for lock-free reads with copy-on-write for mutations.
pub struct Jar {
    compression: bool,
    store: Arc<ArcSwap<HashMap<String, HashMap<String, CookieJar>>>>,
}

// ===== impl CookieStore =====

impl_request_config_value!(Arc<dyn CookieStore>);

// ===== impl IntoCookieStore =====

impl IntoCookieStore for Arc<dyn CookieStore> {
    #[inline]
    fn into_cookie_store(self) -> Arc<dyn CookieStore> {
        self
    }
}

impl<R> IntoCookieStore for Arc<R>
where
    R: CookieStore + 'static,
{
    #[inline]
    fn into_cookie_store(self) -> Arc<dyn CookieStore> {
        self
    }
}

impl<R> IntoCookieStore for R
where
    R: CookieStore + 'static,
{
    #[inline]
    fn into_cookie_store(self) -> Arc<dyn CookieStore> {
        Arc::new(self)
    }
}

// ===== impl IntoCookie =====

impl IntoCookie for Cookie<'_> {
    #[inline]
    fn into_cookie(self) -> Option<Cookie<'static>> {
        Some(self.into_owned())
    }
}

impl IntoCookie for RawCookie<'_> {
    #[inline]
    fn into_cookie(self) -> Option<Cookie<'static>> {
        Some(Cookie(self.into_owned()))
    }
}

impl IntoCookie for &str {
    #[inline]
    fn into_cookie(self) -> Option<Cookie<'static>> {
        RawCookie::parse(self).map(|c| Cookie(c.into_owned())).ok()
    }
}

// ===== impl Cookie =====

impl<'a> Cookie<'a> {
    pub(crate) fn parse(value: &'a HeaderValue) -> crate::Result<Cookie<'a>> {
        std::str::from_utf8(value.as_bytes())
            .map_err(cookie::ParseError::from)
            .and_then(cookie::Cookie::parse)
            .map_err(Error::decode)
            .map(Cookie)
    }

    /// The name of the cookie.
    #[inline]
    pub fn name(&self) -> &str {
        self.0.name()
    }

    /// The value of the cookie.
    #[inline]
    pub fn value(&self) -> &str {
        self.0.value()
    }

    /// Returns true if the 'HttpOnly' directive is enabled.
    #[inline]
    pub fn http_only(&self) -> bool {
        self.0.http_only().unwrap_or(false)
    }

    /// Returns true if the 'Secure' directive is enabled.
    #[inline]
    pub fn secure(&self) -> bool {
        self.0.secure().unwrap_or(false)
    }

    /// Returns true if  'SameSite' directive is 'Lax'.
    #[inline]
    pub fn same_site_lax(&self) -> bool {
        self.0.same_site() == Some(SameSite::Lax)
    }

    /// Returns true if  'SameSite' directive is 'Strict'.
    #[inline]
    pub fn same_site_strict(&self) -> bool {
        self.0.same_site() == Some(SameSite::Strict)
    }

    /// Returns the path directive of the cookie, if set.
    #[inline]
    pub fn path(&self) -> Option<&str> {
        self.0.path()
    }

    /// Returns the domain directive of the cookie, if set.
    #[inline]
    pub fn domain(&self) -> Option<&str> {
        self.0.domain()
    }

    /// Get the Max-Age information.
    #[inline]
    pub fn max_age(&self) -> Option<std::time::Duration> {
        self.0.max_age().and_then(|d| d.try_into().ok())
    }

    /// The cookie expiration time.
    #[inline]
    pub fn expires(&self) -> Option<SystemTime> {
        match self.0.expires() {
            Some(Expiration::DateTime(offset)) => Some(SystemTime::from(offset)),
            None | Some(Expiration::Session) => None,
        }
    }

    /// Converts `self` into a `Cookie` with a static lifetime with as few
    /// allocations as possible.
    #[inline]
    pub fn into_owned(self) -> Cookie<'static> {
        Cookie(self.0.into_owned())
    }
}

impl fmt::Display for Cookie<'_> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl<'c> From<RawCookie<'c>> for Cookie<'c> {
    #[inline]
    fn from(cookie: RawCookie<'c>) -> Cookie<'c> {
        Cookie(cookie)
    }
}

impl<'c> From<Cookie<'c>> for RawCookie<'c> {
    #[inline]
    fn from(cookie: Cookie<'c>) -> RawCookie<'c> {
        cookie.0
    }
}

// ===== impl Jar =====

macro_rules! into_uri {
    ($expr:expr) => {
        match $expr.into_uri() {
            Ok(u) => u,
            Err(_) => return,
        }
    };
}

impl Jar {
    /// Creates a new [`Jar`] with the specified compression setting.
    pub fn new(compression: bool) -> Self {
        Self {
            compression,
            store: Arc::new(ArcSwap::from_pointee(HashMap::with_hasher(HASHER))),
        }
    }

    /// Clone this [`Jar`], sharing storage but enabling compression.
    pub fn compressed(self: &Arc<Self>) -> Arc<Self> {
        Arc::new(Jar {
            compression: true,
            store: self.store.clone(),
        })
    }

    /// Clone this [`Jar`], sharing storage but disabling compression.
    pub fn uncompressed(self: &Arc<Self>) -> Arc<Self> {
        Arc::new(Jar {
            compression: false,
            store: self.store.clone(),
        })
    }

    /// Get a cookie by name for a given Uri.
    ///
    /// Returns the cookie with the specified name for the domain and path
    /// derived from the given Uri, if it exists.
    ///
    /// # Example
    /// ```
    /// use hpx::cookie::Jar;
    /// let jar = Jar::default();
    /// jar.add_cookie_str(
    ///     "foo=bar; Path=/foo; Domain=example.com",
    ///     "http://example.com/foo",
    /// );
    /// let cookie = jar.get("foo", "http://example.com/foo").unwrap();
    /// assert_eq!(cookie.value(), "bar");
    /// ```
    pub fn get<U: IntoUri>(&self, name: &str, uri: U) -> Option<Cookie<'static>> {
        let uri = uri.into_uri().ok()?;
        let store = self.store.load();
        let cookie = store
            .get(uri.host()?)?
            .get(uri.path())?
            .get(name)?
            .clone()
            .into_owned();
        Some(Cookie(cookie))
    }

    /// Get all cookies in this jar.
    ///
    /// Returns an iterator over all cookies currently stored in the jar,
    /// regardless of domain or path.
    ///
    /// # Example
    /// ```
    /// use hpx::cookie::Jar;
    /// let jar = Jar::default();
    /// jar.add_cookie_str("foo=bar; Domain=example.com", "http://example.com");
    /// for cookie in jar.get_all() {
    ///     println!("{}={}", cookie.name(), cookie.value());
    /// }
    /// ```
    pub fn get_all(&self) -> impl Iterator<Item = Cookie<'static>> {
        let store = self.store.load();
        store
            .values()
            .flat_map(|path_map| {
                path_map.values().flat_map(|name_map| {
                    name_map
                        .iter()
                        .map(|cookie| Cookie(cookie.clone().into_owned()))
                })
            })
            .collect::<Vec<_>>()
            .into_iter()
    }

    /// Add a cookie str to this jar.
    ///
    /// # Example
    ///
    /// ```
    /// use hpx::cookie::Jar;
    ///
    /// let cookie = "foo=bar; Domain=yolo.local";
    /// let jar = Jar::default();
    /// jar.add_cookie_str(cookie, "https://yolo.local");
    /// ```
    pub fn add_cookie_str<U: IntoUri>(&self, cookie: &str, uri: U) {
        self.add(cookie, uri);
    }

    /// Add a cookie to this jar.
    ///
    /// The cookie's domain and path attributes are used if present, otherwise
    /// the domain and path are derived from the provided Uri.
    ///
    /// # Example
    /// ```
    /// use cookie::CookieBuilder;
    /// use hpx::cookie::Jar;
    /// let jar = Jar::default();
    /// let cookie = CookieBuilder::new("foo", "bar")
    ///     .domain("example.com")
    ///     .path("/")
    ///     .build();
    /// jar.add(cookie, "http://example.com");
    ///
    /// // You can also use a string directly
    /// jar.add("foo=bar; Domain=example.com; Path=/", "http://example.com");
    /// ```
    pub fn add<C, U>(&self, cookie: C, uri: U)
    where
        C: IntoCookie,
        U: IntoUri,
    {
        if let Some(cookie) = cookie.into_cookie() {
            let cookie: RawCookie<'static> = cookie.into();
            let uri = into_uri!(uri);
            let domain = cookie
                .domain()
                .map(normalize_domain)
                .or_else(|| uri.host())
                .unwrap_or_default();
            let path = cookie.path().unwrap_or_else(|| normalize_path(&uri));

            self.store.rcu(|current| {
                let mut inner = (**current).clone();
                let name_map = inner
                    .entry(domain.to_owned())
                    .or_insert_with(|| HashMap::with_hasher(HASHER))
                    .entry(path.to_owned())
                    .or_default();

                // RFC 6265: If Max-Age=0 or Expires in the past, remove the cookie
                let expired = cookie
                    .expires_datetime()
                    .is_some_and(|dt| dt <= SystemTime::now())
                    || cookie.max_age().is_some_and(|age| age.is_zero());

                if expired {
                    name_map.remove(cookie.clone());
                } else {
                    name_map.add(cookie.clone());
                }
                inner
            });
        }
    }

    /// Add a cookie to this jar (legacy API).
    ///
    /// The cookie's domain and path attributes are used if present, otherwise
    /// the domain and path are derived from the provided Uri.
    ///
    /// # Example
    /// ```
    /// use cookie::CookieBuilder;
    /// use hpx::cookie::Jar;
    /// let jar = Jar::default();
    /// let cookie = CookieBuilder::new("foo", "bar")
    ///     .domain("example.com")
    ///     .path("/")
    ///     .build();
    /// jar.add_cookie(cookie, "http://example.com");
    /// ```
    pub fn add_cookie<C, U>(&self, cookie: C, uri: U)
    where
        C: Into<RawCookie<'static>>,
        U: IntoUri,
    {
        let cookie: RawCookie<'static> = cookie.into();
        let uri = into_uri!(uri);
        let domain = cookie
            .domain()
            .map(normalize_domain)
            .or_else(|| uri.host())
            .unwrap_or_default();
        let path = cookie.path().unwrap_or_else(|| normalize_path(&uri));

        self.store.rcu(|current| {
            let mut inner = (**current).clone();
            let name_map = inner
                .entry(domain.to_owned())
                .or_insert_with(|| HashMap::with_hasher(HASHER))
                .entry(path.to_owned())
                .or_default();

            // RFC 6265: If Max-Age=0 or Expires in the past, remove the cookie
            let expired = cookie
                .expires_datetime()
                .is_some_and(|dt| dt <= SystemTime::now())
                || cookie.max_age().is_some_and(|age| age.is_zero());

            if expired {
                name_map.remove(cookie.clone());
            } else {
                name_map.add(cookie.clone());
            }
            inner
        });
    }

    /// Remove a cookie by name for a given Uri.
    ///
    /// Removes the cookie with the specified name for the domain and path
    /// derived from the given Uri, if it exists.
    ///
    /// # Example
    /// ```
    /// use hpx::cookie::Jar;
    /// let jar = Jar::default();
    /// jar.add_cookie_str(
    ///     "foo=bar; Path=/foo; Domain=example.com",
    ///     "http://example.com/foo",
    /// );
    /// assert!(jar.get("foo", "http://example.com/foo").is_some());
    /// jar.remove("foo", "http://example.com/foo");
    /// assert!(jar.get("foo", "http://example.com/foo").is_none());
    /// ```
    pub fn remove<C, U>(&self, cookie: C, uri: U)
    where
        C: Into<RawCookie<'static>>,
        U: IntoUri,
    {
        let cookie_raw = cookie.into();
        let uri = into_uri!(uri);
        if let Some(host) = uri.host() {
            let path = uri.path().to_owned();
            self.store.rcu(|current| {
                let mut inner = (**current).clone();
                if let Some(path_map) = inner.get_mut(host)
                    && let Some(name_map) = path_map.get_mut(&path)
                {
                    name_map.remove(cookie_raw.clone());
                }
                inner
            });
        }
    }

    /// Clear all cookies from this jar.
    ///
    /// Removes all cookies from the jar, leaving it empty.
    ///
    /// # Example
    /// ```
    /// use hpx::cookie::Jar;
    /// let jar = Jar::default();
    /// jar.add_cookie_str("foo=bar; Domain=example.com", "http://example.com");
    /// assert_eq!(jar.get_all().count(), 1);
    /// jar.clear();
    /// assert_eq!(jar.get_all().count(), 0);
    /// ```
    pub fn clear(&self) {
        self.store.store(Arc::new(HashMap::with_hasher(HASHER)));
    }
}

impl CookieStore for Jar {
    fn set_cookies(&self, cookie_headers: &mut dyn Iterator<Item = &HeaderValue>, uri: &Uri) {
        let cookies = cookie_headers
            .map(Cookie::parse)
            .filter_map(Result::ok)
            .map(|cookie| cookie.0.into_owned());

        for cookie in cookies {
            self.add_cookie(cookie, uri);
        }
    }

    fn cookies(&self, uri: &Uri) -> Cookies {
        let host = match uri.host() {
            Some(h) => h,
            None => return Cookies::Empty,
        };

        let store = self.store.load();
        let iter = store
            .iter()
            .filter(|(domain, _)| domain_match(host, domain))
            .flat_map(|(_, path_map)| {
                path_map
                    .iter()
                    .filter(|(path, _)| path_match(uri.path(), path))
                    .flat_map(|(_, name_map)| {
                        name_map.iter().filter(|cookie| {
                            if cookie.secure() == Some(true) && uri.is_http() {
                                return false;
                            }

                            if cookie
                                .expires_datetime()
                                .is_some_and(|dt| dt <= SystemTime::now())
                            {
                                return false;
                            }

                            true
                        })
                    })
            });

        if self.compression {
            let cookies = iter.fold(String::new(), |mut cookies, cookie| {
                if !cookies.is_empty() {
                    cookies.push_str("; ");
                }
                cookies.push_str(cookie.name());
                cookies.push('=');
                cookies.push_str(cookie.value());
                cookies
            });

            if cookies.is_empty() {
                return Cookies::Empty;
            }

            HeaderValue::from_maybe_shared(Bytes::from(cookies))
                .map(Cookies::Compressed)
                .unwrap_or(Cookies::Empty)
        } else {
            let cookies = iter
                .map(|cookie| {
                    let name = cookie.name();
                    let value = cookie.value();

                    let mut cookie_str = String::with_capacity(name.len() + 1 + value.len());
                    cookie_str.push_str(name);
                    cookie_str.push('=');
                    cookie_str.push_str(value);

                    HeaderValue::from_maybe_shared(Bytes::from(cookie_str))
                })
                .filter_map(Result::ok)
                .collect();

            Cookies::Uncompressed(cookies)
        }
    }
}

impl Default for Jar {
    fn default() -> Self {
        Self::new(true)
    }
}

// Jar needs fmt::Debug since the store is ArcSwap now
impl fmt::Debug for Jar {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Jar")
            .field("compression", &self.compression)
            .finish()
    }
}

const DEFAULT_PATH: &str = "/";

/// Determines if the given `host` matches the cookie `domain` according to
/// [RFC 6265 section 5.1.3](https://datatracker.ietf.org/doc/html/rfc6265#section-5.1.3).
///
/// - Returns true if the host and domain are identical.
/// - Returns true if the host is a subdomain of the domain (host ends with ".domain").
/// - Returns false otherwise.
fn domain_match(host: &str, domain: &str) -> bool {
    if domain.is_empty() {
        return false;
    }
    if host == domain {
        return true;
    }
    host.len() > domain.len()
        && host.as_bytes()[host.len() - domain.len() - 1] == b'.'
        && host.ends_with(domain)
}

/// Determines if the request path matches the cookie path according to
/// [RFC 6265 section 5.1.4](https://datatracker.ietf.org/doc/html/rfc6265#section-5.1.4).
///
/// - Returns true if the request path and cookie path are identical.
/// - Returns true if the request path starts with the cookie path, and
///   - the cookie path ends with '/', or
///   - the next character in the request path after the cookie path is '/'.
/// - Returns false otherwise.
fn path_match(req_path: &str, cookie_path: &str) -> bool {
    req_path == cookie_path
        || req_path.starts_with(cookie_path)
            && (cookie_path.ends_with(DEFAULT_PATH)
                || req_path[cookie_path.len()..].starts_with(DEFAULT_PATH))
}

/// Normalizes a domain by stripping any port information.
///
/// According to [RFC 6265 section 5.2.3](https://datatracker.ietf.org/doc/html/rfc6265#section-5.2.3),
/// the domain attribute of a cookie must not include a port. If a port is present (non-standard),
/// it will be ignored for domain matching purposes.
fn normalize_domain(domain: &str) -> &str {
    domain.split(':').next().unwrap_or(domain)
}

/// Computes the normalized default path for a cookie as specified in
/// [RFC 6265 section 5.1.4](https://datatracker.ietf.org/doc/html/rfc6265#section-5.1.4).
///
/// This function normalizes the path for a cookie, ensuring it matches
/// browser and server expectations for default cookie scope.
fn normalize_path(uri: &Uri) -> &str {
    let path = uri.path();
    if !path.starts_with(DEFAULT_PATH) {
        return DEFAULT_PATH;
    }
    if let Some(pos) = path.rfind(DEFAULT_PATH) {
        if pos == 0 {
            return DEFAULT_PATH;
        }
        return &path[..pos];
    }
    DEFAULT_PATH
}