hpx 2.4.18

//! HTTP Cookies

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

use bytes::Bytes;
use cookie::{Cookie as RawCookie, CookieJar, Expiration, SameSite};
use http::Uri;
use scc::HashMap as SccHashMap;

use crate::{IntoUri, error::Error, ext::UriExt, 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 `scc::HashMap` for lock-free reads with fine-grained entry-level
/// mutations, avoiding the full-map clone that `ArcSwap` requires on every
/// write.
pub struct Jar {
    compression: bool,
    store: Arc<SccHashMap<String, SccHashMap<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 =====

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

    /// 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 host = uri.host()?;
        let path = uri.path();
        let path_entry = self.store.get_sync(host)?;
        let cookie_jar = path_entry.get_sync(path)?;
        cookie_jar
            .get()
            .get(name)
            .map(|c| Cookie(c.clone().into_owned()))
    }

    /// 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) -> Vec<Cookie<'static>> {
        let mut cookies = Vec::new();
        self.store.iter_sync(|_, path_map| {
            path_map.iter_sync(|_, cookie_jar| {
                for cookie in cookie_jar.iter() {
                    cookies.push(Cookie(cookie.clone().into_owned()));
                }
                true
            });
            true
        });
        cookies
    }

    /// 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");
    /// ```
    #[allow(clippy::unwrap_or_default)]
    pub fn add<C, U>(&self, cookie: C, uri: U)
    where
        C: IntoCookie,
        U: IntoUri,
    {
        let Some(cookie) = cookie.into_cookie() else {
            return;
        };
        let cookie: RawCookie<'static> = cookie.into();
        let Ok(uri) = uri.into_uri() else { return };
        let domain = cookie
            .domain()
            .map(normalize_domain)
            .or_else(|| uri.host())
            .unwrap_or_default();
        let path = cookie.path().unwrap_or_else(|| normalize_path(&uri));

        let expired = cookie
            .expires_datetime()
            .is_some_and(|dt| dt <= SystemTime::now())
            || cookie.max_age().is_some_and(|age| age.is_zero());

        self.store
            .entry_sync(domain.to_owned())
            .or_insert_with(SccHashMap::new)
            .get()
            .entry_sync(path.to_owned())
            .or_insert_with(CookieJar::default)
            .get_mut()
            .remove(cookie.name().to_owned());

        if !expired {
            self.store
                .entry_sync(domain.to_owned())
                .or_insert_with(SccHashMap::new)
                .get()
                .entry_sync(path.to_owned())
                .or_insert_with(CookieJar::default)
                .get_mut()
                .add(cookie);
        }
    }

    /// 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");
    /// ```
    #[allow(clippy::unwrap_or_default)]
    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 Ok(uri) = uri.into_uri() else { return };
        let domain = cookie
            .domain()
            .map(normalize_domain)
            .or_else(|| uri.host())
            .unwrap_or_default();
        let path = cookie.path().unwrap_or_else(|| normalize_path(&uri));

        let expired = cookie
            .expires_datetime()
            .is_some_and(|dt| dt <= SystemTime::now())
            || cookie.max_age().is_some_and(|age| age.is_zero());

        self.store
            .entry_sync(domain.to_owned())
            .or_insert_with(SccHashMap::new)
            .get()
            .entry_sync(path.to_owned())
            .or_insert_with(CookieJar::default)
            .get_mut()
            .remove(cookie.name().to_owned());

        if !expired {
            self.store
                .entry_sync(domain.to_owned())
                .or_insert_with(SccHashMap::new)
                .get()
                .entry_sync(path.to_owned())
                .or_insert_with(CookieJar::default)
                .get_mut()
                .add(cookie);
        }
    }

    /// 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());
    /// ```
    #[allow(clippy::unwrap_or_default)]
    pub fn remove<C, U>(&self, cookie: C, uri: U)
    where
        C: Into<RawCookie<'static>>,
        U: IntoUri,
    {
        let cookie_raw = cookie.into();
        let Ok(uri) = uri.into_uri() else { return };
        let Some(host) = uri.host() else { return };
        let path = uri.path().to_owned();

        self.store
            .entry_sync(host.to_owned())
            .or_insert_with(SccHashMap::new)
            .get()
            .entry_sync(path)
            .or_insert_with(CookieJar::default)
            .get_mut()
            .remove(cookie_raw.name().to_owned());
    }

    /// 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.clear_sync();
    }
}

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,
        };

        if self.compression {
            let mut cookies = String::new();
            self.store.iter_sync(|domain, path_map| {
                if !domain_match(host, domain) {
                    return true;
                }
                path_map.iter_sync(|path, cookie_jar| {
                    if !path_match(uri.path(), path) {
                        return true;
                    }
                    for cookie in cookie_jar.iter() {
                        if cookie.secure() == Some(true) && uri.is_http() {
                            continue;
                        }
                        if cookie
                            .expires_datetime()
                            .is_some_and(|dt| dt <= SystemTime::now())
                        {
                            continue;
                        }
                        if !cookies.is_empty() {
                            cookies.push_str("; ");
                        }
                        cookies.push_str(cookie.name());
                        cookies.push('=');
                        cookies.push_str(cookie.value());
                    }
                    true
                });
                true
            });

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

            HeaderValue::from_maybe_shared(Bytes::from(cookies))
                .map(Cookies::Compressed)
                .unwrap_or(Cookies::Empty)
        } else {
            let mut cookie_strs = Vec::new();
            self.store.iter_sync(|domain, path_map| {
                if !domain_match(host, domain) {
                    return true;
                }
                path_map.iter_sync(|path, cookie_jar| {
                    if !path_match(uri.path(), path) {
                        return true;
                    }
                    for cookie in cookie_jar.iter() {
                        if cookie.secure() == Some(true) && uri.is_http() {
                            continue;
                        }
                        if cookie
                            .expires_datetime()
                            .is_some_and(|dt| dt <= SystemTime::now())
                        {
                            continue;
                        }
                        let name = cookie.name();
                        let value = cookie.value();
                        let mut s = String::with_capacity(name.len() + 1 + value.len());
                        s.push_str(name);
                        s.push('=');
                        s.push_str(value);
                        cookie_strs.push(s);
                    }
                    true
                });
                true
            });

            let cookies = cookie_strs
                .into_iter()
                .filter_map(|s| HeaderValue::from_maybe_shared(Bytes::from(s)).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
}