volga 0.9.1

//! Utilities for private cookies

use crate::{
    HttpRequest,
    di::Container,
    error::Error,
    headers::HeaderMap,
    http::{
        Extensions, Parts, Request,
        body::Incoming,
        cookie::get_cookies,
        endpoints::args::{
            FromPayload, FromRawRequest, FromRequestParts, FromRequestRef, Payload, Source,
        },
    },
};
use cookie::{CookieJar, Key, PrivateJar};
use futures_util::future::{Ready, ready};
use std::{fs::File, io::Read, path::Path};

/// Represents a cryptographic pass key for [`PrivateCookies`]
#[derive(Clone)]
pub struct PrivateKey(Key);

/// Represents private HTTP cookies
pub struct PrivateCookies(PrivateKey, CookieJar);

impl Default for PrivateKey {
    #[inline]
    fn default() -> Self {
        Self::from(&[])
    }
}

impl std::fmt::Debug for PrivateCookies {
    #[inline]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("PrivateCookies")
            .field(&"[redacted]")
            .finish()
    }
}

impl std::fmt::Debug for PrivateKey {
    #[inline]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("PrivateKey").field(&"[redacted]").finish()
    }
}

impl PrivateKey {
    /// Creates a new [`PrivateKey`] from a 512-bit cryptographically random string.
    ///
    /// See also [`Key::from`]
    #[inline]
    pub fn from(bytes: &[u8]) -> Self {
        Self(Key::from(bytes))
    }

    /// Creates a new [`PrivateKey`] from a file with 512-bit cryptographically random string.
    ///
    /// > **Note:** It reads the first 64 bytes of the file.
    ///
    /// See also [`Key::from`]
    #[inline]
    pub fn from_file(path: impl AsRef<Path>) -> Self {
        let mut file = File::open(path).expect("File must exists");
        let mut buffer = [0u8; 64];
        file.read_exact(&mut buffer).expect("File must be readable");

        Self(Key::from(&buffer))
    }

    /// Generates signing/encryption keys from a secure, random source.
    /// Keys are generated nondeterministically.
    ///
    /// See also [`Key::generate`]
    #[inline]
    pub fn generate() -> Self {
        Self(Key::generate())
    }
}

impl PrivateCookies {
    /// Creates a new [`PrivateCookies`]
    #[inline]
    pub fn new(key: PrivateKey) -> Self {
        Self(key, CookieJar::default())
    }

    /// Creates a new [`PrivateCookies`] from [`HeaderMap`]
    #[inline]
    pub fn from_headers(key: PrivateKey, headers: &HeaderMap) -> Self {
        let mut jar = CookieJar::new();
        let mut private_jar = jar.private_mut(&key.0);
        for cookie in get_cookies(headers) {
            if let Some(cookie) = private_jar.decrypt(cookie) {
                private_jar.add_original(cookie);
            }
        }
        Self(key, jar)
    }

    /// Unwraps the inner jar and the pass key.
    #[inline]
    pub fn into_parts(self) -> (PrivateKey, CookieJar) {
        (self.0, self.1)
    }

    /// Returns a reference to the cookie inside the private jar by `name`
    /// and authenticates and decrypts the cookie's value, returning a [`Cookie`] with the decrypted value.
    /// If the cookie cannot be found, or the cookie fails to authenticate or decrypt, `None` is returned.
    pub fn get(&self, name: &str) -> Option<cookie::Cookie<'static>> {
        self.private().get(name)
    }

    /// Adds a cookie.
    /// The cookie's value is encrypted with authenticated encryption assuring confidentiality, integrity, and authenticity.
    #[allow(clippy::should_implement_trait)]
    pub fn add<C: Into<cookie::Cookie<'static>>>(mut self, cookie: C) -> Self {
        self.private_mut().add(cookie);
        self
    }

    /// Removes a cookie from the private jar.
    ///
    /// For correct removal, the passed in cookie must contain the same path
    /// and domain as the cookie that was initially set.
    pub fn remove<C: Into<cookie::Cookie<'static>>>(mut self, cookie: C) -> Self {
        self.private_mut().remove(cookie);
        self
    }

    /// Authenticates and decrypts the cookie, returning the plaintext version if decryption succeeds or `None` otherwise.
    /// Authentication and decryption always succeed if a cookie was generated by a [`PrivateCookies`] with the same key as self.
    pub fn decrypt(&self, cookie: cookie::Cookie<'static>) -> Option<cookie::Cookie<'static>> {
        self.private().decrypt(cookie)
    }

    /// Returns an iterator over all the cookies present in this jar.
    pub fn iter(&self) -> impl Iterator<Item = &cookie::Cookie<'static>> + '_ {
        self.1.iter()
    }

    #[inline]
    fn private(&self) -> PrivateJar<&'_ CookieJar> {
        self.1.private(&self.0.0)
    }

    #[inline]
    fn private_mut(&mut self) -> PrivateJar<&'_ mut CookieJar> {
        self.1.private_mut(&self.0.0)
    }
}

impl TryFrom<&Parts> for PrivateCookies {
    type Error = Error;

    #[inline]
    fn try_from(parts: &Parts) -> Result<Self, Self::Error> {
        Container::try_from(parts)?
            .resolve::<PrivateKey>()
            .map(|key| PrivateCookies::from_headers(key, &parts.headers))
            .map_err(Into::into)
    }
}

impl TryFrom<(&Extensions, &HeaderMap)> for PrivateCookies {
    type Error = Error;

    #[inline]
    fn try_from((extensions, headers): (&Extensions, &HeaderMap)) -> Result<Self, Self::Error> {
        Container::try_from(extensions)?
            .resolve::<PrivateKey>()
            .map(|key| PrivateCookies::from_headers(key, headers))
            .map_err(Into::into)
    }
}

impl FromRequestRef for PrivateCookies {
    #[inline]
    fn from_request(req: &HttpRequest) -> Result<Self, Error> {
        Self::try_from((req.extensions(), req.headers()))
    }
}

impl FromRawRequest for PrivateCookies {
    #[inline]
    fn from_request(req: Request<Incoming>) -> impl Future<Output = Result<Self, Error>> + Send {
        ready(Self::try_from((req.extensions(), req.headers())))
    }
}

impl FromRequestParts for PrivateCookies {
    #[inline]
    fn from_parts(parts: &Parts) -> Result<Self, Error> {
        parts.try_into()
    }
}

impl FromPayload for PrivateCookies {
    type Future = Ready<Result<Self, Error>>;

    const SOURCE: Source = Source::Parts;

    #[inline]
    fn from_payload(payload: Payload<'_>) -> Self::Future {
        let Payload::Parts(parts) = payload else {
            unreachable!()
        };
        ready(Self::from_parts(parts))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::di::ContainerBuilder;
    use crate::headers::{COOKIE, SET_COOKIE};
    use crate::http::{HttpBody, HttpRequest, Request, cookie::set_cookies};

    #[test]
    fn it_creates_cookies_from_empty_headers() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key);
        assert_eq!(cookies.iter().count(), 0);
    }

    #[test]
    fn it_creates_cookies() {
        let key = PrivateKey::from(
            b"f3d9e2a44c6b172a1ea9b9d05e5fe1bcaa8679d032ccae271c503af9618bb2ef7c4e51452dbfcd96f6e9c9d09166a3de77e");
        let cookies = PrivateCookies::new(key.clone());
        let cookies = cookies.add(("session", "abc123"));

        let mut headers = HeaderMap::new();
        set_cookies_for_request(cookies.1, &mut headers);

        let cookies = PrivateCookies::from_headers(key, &headers);
        let cookie = cookies.get("session").expect("Cookie should exist");

        assert_eq!(cookie.value(), "abc123");
    }

    #[test]
    fn it_creates_from_multiple_cookies() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key.clone());
        let cookies = cookies
            .add(("session", "abc123"))
            .add(("user", "john"))
            .add(("theme", "dark"));

        let mut headers = HeaderMap::new();
        set_cookies_for_request(cookies.1, &mut headers);

        let cookies = PrivateCookies::from_headers(key, &headers);

        assert_eq!(cookies.get("session").unwrap().value(), "abc123");
        assert_eq!(cookies.get("user").unwrap().value(), "john");
        assert_eq!(cookies.get("theme").unwrap().value(), "dark");
    }

    #[test]
    fn it_adds_and_removes_cookies() {
        let key = PrivateKey::generate();
        let mut cookies = PrivateCookies::new(key);

        // Add a new cookie
        cookies = cookies.add(cookie::Cookie::new("test", "value"));
        assert_eq!(cookies.get("test").unwrap().value(), "value");

        // Remove a cookie
        cookies = cookies.remove(cookie::Cookie::new("test", ""));
        assert!(cookies.get("test").is_none());
    }

    #[test]
    fn it_sets_cookies_to_headers() {
        let key = PrivateKey::generate();
        let mut cookies = PrivateCookies::new(key);
        cookies = cookies.add(cookie::Cookie::new("session", "xyz789"));

        let mut headers = HeaderMap::new();
        set_cookies(cookies.1, &mut headers);

        let cookie_header = headers
            .get(SET_COOKIE)
            .expect("Cookie header should be set");
        assert!(cookie_header.to_str().unwrap().contains("session"));
    }

    #[tokio::test]
    async fn it_extracts_from_payload() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key.clone());
        let cookies = cookies.add(("test", "value"));

        let mut headers = HeaderMap::new();
        set_cookies_for_request(cookies.1, &mut headers);

        let mut container = ContainerBuilder::new();
        container.register_singleton(key);
        let container = container.build();

        let mut request = Request::builder()
            .extension(container.create_scope())
            .body(())
            .unwrap();

        request.headers_mut().extend(headers);

        let (parts, _) = request.into_parts();
        let payload = Payload::Parts(&parts);

        let cookies = PrivateCookies::from_payload(payload).await.unwrap();

        assert_eq!(cookies.get("test").unwrap().value(), "value");
    }

    #[test]
    fn it_extracts_from_request_ref() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key.clone());
        let cookies = cookies.add(("test", "value"));

        let mut headers = HeaderMap::new();
        set_cookies_for_request(cookies.1, &mut headers);

        let mut container = ContainerBuilder::new();
        container.register_singleton(key);
        let container = container.build();

        let mut request = Request::builder()
            .extension(container.create_scope())
            .body(HttpBody::empty())
            .unwrap();

        request.headers_mut().extend(headers);

        let (parts, body) = request.into_parts();
        let req = HttpRequest::from_parts(parts, body);

        let cookies = <PrivateCookies as FromRequestRef>::from_request(&req).unwrap();

        assert_eq!(cookies.get("test").unwrap().value(), "value");
    }

    #[test]
    fn it_extracts_from_parts() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key.clone());
        let cookies = cookies.add(("test", "value"));

        let mut headers = HeaderMap::new();
        set_cookies_for_request(cookies.1, &mut headers);

        let mut container = ContainerBuilder::new();
        container.register_singleton(key);
        let container = container.build();

        let mut request = Request::builder()
            .extension(container.create_scope())
            .body(())
            .unwrap();

        request.headers_mut().extend(headers);

        let (parts, _) = request.into_parts();

        let cookies = PrivateCookies::from_parts(&parts).unwrap();

        assert_eq!(cookies.get("test").unwrap().value(), "value");
    }

    #[test]
    fn it_tries_extracts_from_parts() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key.clone());
        let cookies = cookies.add(("test", "value"));

        let mut headers = HeaderMap::new();
        set_cookies_for_request(cookies.1, &mut headers);

        let mut container = ContainerBuilder::new();
        container.register_singleton(key);
        let container = container.build();

        let mut request = Request::builder()
            .extension(container.create_scope())
            .body(())
            .unwrap();

        request.headers_mut().extend(headers);

        let (parts, _) = request.into_parts();

        let cookies = PrivateCookies::try_from(&parts).unwrap();

        assert_eq!(cookies.get("test").unwrap().value(), "value");
    }

    #[test]
    fn it_tries_extracts_from_extensions_and_headers() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key.clone());
        let cookies = cookies.add(("test", "value"));

        let mut headers = HeaderMap::new();
        set_cookies_for_request(cookies.1, &mut headers);

        let mut container = ContainerBuilder::new();
        container.register_singleton(key);
        let container = container.build();

        let mut request = Request::builder()
            .extension(container.create_scope())
            .body(())
            .unwrap();

        request.headers_mut().extend(headers);

        let cookies = PrivateCookies::try_from((request.extensions(), request.headers())).unwrap();

        assert_eq!(cookies.get("test").unwrap().value(), "value");
    }

    #[test]
    fn if_return_parts_source() {
        assert_eq!(PrivateCookies::SOURCE, Source::Parts);
    }

    #[tokio::test]
    async fn it_reads_signed_key_from_bytes() {
        let temp_file = crate::test::TempFile::new(
            "f3d9e2a44c6b172a1ea9b9d05e5fe1bcaa8679d032ccae271c503af9618bb2ef7c4e51452dbfcd96f6e9c9d09166a3de77e"
        ).await;
        let _ = PrivateKey::from_file(&temp_file.path);
    }

    #[test]
    fn it_debugs() {
        let key = PrivateKey::generate();
        let cookies = PrivateCookies::new(key.clone());

        assert_eq!(format!("{key:?}"), r#"PrivateKey("[redacted]")"#);
        assert_eq!(format!("{cookies:?}"), r#"PrivateCookies("[redacted]")"#);
    }

    fn set_cookies_for_request(jar: CookieJar, headers: &mut HeaderMap) {
        for cookie in jar.delta() {
            if let Ok(header_value) = cookie.encoded().to_string().parse() {
                headers.append(COOKIE, header_value);
            }
        }
    }
}