axum_session 0.19.0

use std::borrow::{Borrow, BorrowMut};

use hmac::{Hmac, Mac};
use sha2::Sha256;

use cookie::{Cookie, CookieJar, Key};

// Keep these in sync, and keep the key len synced with the `signed` docs as
// well as the `KEYS_INFO` const in secure::Key.
pub(crate) const BASE64_DIGEST_LEN: usize = 44;
pub(crate) const KEY_LEN: usize = 32;

use base64::{prelude::BASE64_STANDARD, DecodeError, Engine};
use tokio::task;

/// Encode `input` as the standard base64 with padding.
pub(crate) fn encode<T: AsRef<[u8]>>(input: T) -> String {
    BASE64_STANDARD.encode(input)
}

/// Decode `input` as the standard base64 with padding.
pub(crate) fn decode<T: AsRef<[u8]>>(input: T) -> Result<Vec<u8>, DecodeError> {
    BASE64_STANDARD.decode(input)
}

pub trait CookiesAdditionJar {
    fn message_signed<'a>(&'a self, key: &Key, message: String) -> AdditionalSignedJar<&'a Self>;
    fn message_signed_mut<'a>(
        &'a mut self,
        key: &Key,
        message: String,
    ) -> AdditionalSignedJar<&'a mut Self>;
}

impl CookiesAdditionJar for CookieJar {
    fn message_signed<'a>(&'a self, key: &Key, message: String) -> AdditionalSignedJar<&'a Self> {
        AdditionalSignedJar::new(self, key, message)
    }

    fn message_signed_mut<'a>(
        &'a mut self,
        key: &Key,
        message: String,
    ) -> AdditionalSignedJar<&'a mut Self> {
        AdditionalSignedJar::new(self, key, message)
    }
}

/// A child cookie jar that authenticates its cookies and Adds Additional measures to ensure integrity.
pub struct AdditionalSignedJar<J> {
    parent: J,
    key: [u8; KEY_LEN],
    message: String,
}

impl<J> AdditionalSignedJar<J> {
    /// Creates a new child `AdditionalSignedJar`
    pub(crate) fn new(parent: J, key: &Key, message: String) -> AdditionalSignedJar<J> {
        AdditionalSignedJar {
            parent,
            key: key.signing().try_into().expect("sign key len"),
            message,
        }
    }

    /// Signs the cookie's value and message providing integrity and authenticity.
    async fn sign_cookie(&self, cookie: &mut Cookie<'_>) {
        let key = self.key.to_owned();
        let message = self.message.to_owned();
        let value = cookie.value().to_owned();

        let new_value = task::spawn_blocking(move || {
            // Compute HMAC-SHA256 of the cookie's value.
            let mut mac = match Hmac::<Sha256>::new_from_slice(&key) {
                Ok(v) => v,
                Err(err) => {
                    tracing::error!(err = %err,  "key is invalid." );
                    return None;
                }
            };

            // Add the payload to the message first.
            let message = format!("{}{}", &value, &message);
            mac.update(message.as_bytes());

            // Cookie's new value is [MAC | original-value].
            let mut new_value = encode(mac.finalize().into_bytes());
            new_value.push_str(&value);
            Some(new_value)
        })
        .await;

        if let Ok(Some(value)) = new_value {
            cookie.set_value(value);
        }
    }

    /// Given a signed value `str` where the signature is prepended to `value`,
    /// verifies the signed value and returns it. If there's a problem, returns
    /// an `Err` with a string describing the issue.
    async fn _verify(&self, cookie_value: &str) -> Result<String, &'static str> {
        if !cookie_value.is_char_boundary(BASE64_DIGEST_LEN) {
            return Err("missing or invalid digest");
        }

        // we clone it so we can pass them to the blocking thread for encryption verification needs.
        let key = self.key.to_owned();
        let message = self.message.to_owned();
        let cookie_value = cookie_value.to_owned();

        task::spawn_blocking(move || {
            // Split [MAC | original-value] into its two parts.
            let (digest_str, value) = cookie_value.split_at(BASE64_DIGEST_LEN);
            let digest = decode(digest_str).map_err(|_| "bad base64 digest")?;

            // Perform the verification.
            let mut mac = Hmac::<Sha256>::new_from_slice(&key).map_err(|_| "key is invalid.")?;
            // Add message here so we can check if it matches.
            let message = format!("{}{}", value, &message);
            mac.update(message.as_bytes());
            mac.verify_slice(&digest)
                .map(|_| value.to_string())
                .map_err(|_| "value did not verify")
        })
        .await
        .map_err(|_| "thread join in _verify was unsuccessful")?
    }

    /// Verifies the authenticity and integrity of `cookie`, returning the
    /// plaintext version if verification succeeds or `None` otherwise.
    /// Verification _always_ succeeds if `cookie` was generated by a
    /// `AdditionalSignedJar` with the same key and message as `self`.
    ///
    /// # Example
    ///
    /// ```rust ignore
    /// use cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// assert!(jar.message_signed(&key, "".to_owned()).get("name").is_none());
    ///
    /// jar.message_signed_mut(&key, "".to_owned()).add(("name", "value"));
    /// assert_eq!(jar.message_signed(&key, "".to_owned()).get("name").unwrap().value(), "value");
    ///
    /// let plain = jar.get("name").cloned().unwrap();
    /// assert_ne!(plain.value(), "value");
    /// let verified = jar.message_signed(&key, "".to_owned()).verify(plain).unwrap();
    /// assert_eq!(verified.value(), "value");
    ///
    /// let plain = Cookie::new("plaintext", "hello");
    /// assert!(jar.message_signed(&key, "".to_owned()).verify(plain).is_none());
    /// ```
    pub async fn verify(&self, mut cookie: Cookie<'static>) -> Option<Cookie<'static>> {
        match self._verify(cookie.value()).await {
            Ok(value) => {
                cookie.set_value(value);
                Some(cookie)
            }
            Err(err) => {
                tracing::warn!(
                    err = %err,
                    "possibly suspicious activity: Verification failed for Cookie. cookie validation string was: {}",
                    self.message
                );
                None
            }
        }
    }
}

impl<J: Borrow<CookieJar>> AdditionalSignedJar<J> {
    /// Returns a reference to the `Cookie` inside this jar with the name `name`
    /// and verifies the authenticity and integrity of the cookie's value,
    /// returning a `Cookie` with the authenticated value. If the cookie cannot
    /// be found, or the cookie fails to verify, `None` is returned.
    ///
    /// # Example
    ///
    /// ```rust ignore
    /// use cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let jar = CookieJar::new();
    /// assert!(jar.message_signed(&key, "".to_owned()).get("name").is_none());
    ///
    /// let mut jar = jar;
    /// let mut signed_jar = jar.message_signed_mut(&key, "".to_owned());
    /// signed_jar.add(Cookie::new("name", "value"));
    /// assert_eq!(signed_jar.get("name").unwrap().value(), "value");
    /// ```
    pub async fn get(&self, name: &str) -> Option<Cookie<'static>> {
        match self.parent.borrow().get(name) {
            Some(c) => self.verify(c.clone()).await,
            None => None,
        }
    }
}

impl<J: BorrowMut<CookieJar>> AdditionalSignedJar<J> {
    /// Adds `cookie` to the parent jar. The cookie's value is signed assuring
    /// integrity and authenticity.
    ///
    /// # Example
    ///
    /// ```rust ignore
    /// use cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// jar.message_signed_mut(&key, "".to_owned()).add(("name", "value")).await;
    ///
    /// assert_ne!(jar.get("name").unwrap().value(), "value");
    /// assert!(jar.get("name").unwrap().value().contains("value"));
    /// assert_eq!(jar.message_signed(&key, "".to_owned()).get("name").unwrap().value(), "value");
    /// ```
    pub async fn add<C: Into<Cookie<'static>>>(&mut self, cookie: C) {
        let mut cookie = cookie.into();
        self.sign_cookie(&mut cookie).await;
        self.parent.borrow_mut().add(cookie);
    }

    /// Adds an "original" `cookie` to this jar. The cookie's value is signed
    /// assuring integrity and authenticity. Adding an original cookie does not
    /// affect the [`CookieJar::delta()`] computation. This method is intended
    /// to be used to seed the cookie jar with cookies received from a client's
    /// HTTP message.
    ///
    /// For accurate `delta` computations, this method should not be called
    /// after calling `remove`.
    ///
    /// # Example
    ///
    /// ```rust ignore
    /// use cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// jar.message_signed_mut(&key, "".to_owned()).add_original(("name", "value")).await;
    ///
    /// assert_eq!(jar.iter().count(), 1);
    /// assert_eq!(jar.delta().count(), 0);
    /// ```
    pub async fn add_original<C: Into<Cookie<'static>>>(&mut self, cookie: C) {
        let mut cookie = cookie.into();
        self.sign_cookie(&mut cookie).await;
        self.parent.borrow_mut().add_original(cookie);
    }

    /// Removes `cookie` from the parent jar.
    ///
    /// For correct removal, the passed in `cookie` must contain the same `path`
    /// and `domain` as the cookie that was initially set.
    ///
    /// This is identical to [`CookieJar::remove()`]. See the method's
    /// documentation for more details.
    ///
    /// # Example
    ///
    /// ```rust ignore
    /// use cookie::{CookieJar, Cookie, Key};
    ///
    /// let key = Key::generate();
    /// let mut jar = CookieJar::new();
    /// let mut signed_jar = jar.message_signed_mut(&key, "".to_owned());
    ///
    /// signed_jar.add(("name", "value"));
    /// assert!(signed_jar.get("name").is_some());
    ///
    /// signed_jar.remove("name");
    /// assert!(signed_jar.get("name").is_none());
    /// ```
    pub fn remove<C: Into<Cookie<'static>>>(&mut self, cookie: C) {
        self.parent.borrow_mut().remove(cookie.into());
    }
}

pub(crate) async fn sign_header(
    value: &str,
    key: &Key,
    message: &str,
) -> Result<String, &'static str> {
    let value = value.to_owned();
    let message = message.to_owned();
    let key = key.to_owned();

    task::spawn_blocking(move || {
        // Compute HMAC-SHA256 of the cookie's value.
        let mut mac =
            Hmac::<Sha256>::new_from_slice(key.signing()).map_err(|_| "Key was invalid.")?;
        // Add the payload to the message first.
        let message = format!("{value}{message}");
        mac.update(message.as_bytes());

        // Cookie's new value is [MAC | original-value].
        let mut new_value = encode(mac.finalize().into_bytes());
        new_value.push_str(&value);
        Ok(new_value)
    })
    .await
    .map_err(|_| "thread join in _verify was unsuccessful")?
}

/// Given a signed value `str` where the signature is prepended to `value`,
/// verifies the signed value and returns it. If there's a problem, returns
/// an `Err` with a string describing the issue.
pub(crate) async fn verify_header(
    header_value: &str,
    key: &Key,
    message: &str,
) -> Result<String, &'static str> {
    let header_value = header_value.to_owned();
    let message = message.to_owned();
    let key = key.to_owned();

    task::spawn_blocking(move || {
        if !header_value.is_char_boundary(BASE64_DIGEST_LEN) {
            return Err("missing or invalid digest");
        }

        // Split [MAC | original-value] into its two parts.
        let (digest_str, value) = header_value.split_at(BASE64_DIGEST_LEN);
        let digest = decode(digest_str).map_err(|_| "bad base64 digest")?;

        // Perform the verification.
        let mut mac =
            Hmac::<Sha256>::new_from_slice(key.signing()).map_err(|_| "Key was invalid.")?;
        // Add message here so we can check if it matches.
        let message = format!("{value}{message}");
        mac.update(message.as_bytes());
        mac.verify_slice(&digest)
            .map(|_| value.to_string())
            .map_err(|_| "value did not verify")
    })
    .await
    .map_err(|_| "thread join in _verify was unsuccessful")?
}

#[cfg(test)]
mod test {
    use crate::sec::signed::CookiesAdditionJar;
    use cookie::{Cookie, CookieJar, Key};

    #[tokio::test]
    async fn roundtrip() {
        // Secret is SHA-256 hash of 'Super secret!' passed through HKDF-SHA256.
        let key = Key::from(&[
            89, 202, 200, 125, 230, 90, 197, 245, 166, 249, 34, 169, 135, 31, 20, 197, 94, 154,
            254, 79, 60, 26, 8, 143, 254, 24, 116, 138, 92, 225, 159, 60, 157, 41, 135, 129, 31,
            226, 196, 16, 198, 168, 134, 4, 42, 1, 196, 24, 57, 103, 241, 147, 201, 185, 233, 10,
            180, 170, 187, 89, 252, 137, 110, 107,
        ]);

        let mut jar = CookieJar::new();
        jar.add(Cookie::new(
            "signed_with_ring014",
            "3tdHXEQ2kf6fxC7dWzBGmpSLMtJenXLKrZ9cHkSsl1w=Tamper-proof",
        ));
        jar.add(Cookie::new(
            "signed_with_ring016",
            "3tdHXEQ2kf6fxC7dWzBGmpSLMtJenXLKrZ9cHkSsl1w=Tamper-proof",
        ));

        let signed = jar.message_signed(&key, "".to_owned());
        assert_eq!(
            signed.get("signed_with_ring014").await.unwrap().value(),
            "Tamper-proof"
        );
        assert_eq!(
            signed.get("signed_with_ring016").await.unwrap().value(),
            "Tamper-proof"
        );
    }
}