diff --git a/axum-extra/src/extract/cookie/mod.rs b/axum-extra/src/extract/cookie/mod.rs index 50fa6031ac..107eedff7e 100644 --- a/axum-extra/src/extract/cookie/mod.rs +++ b/axum-extra/src/extract/cookie/mod.rs @@ -200,6 +200,82 @@ impl CookieJar { pub fn iter(&self) -> impl Iterator> { self.jar.iter() } + + /// Add a cookie with the specified prefix to the jar. + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::{CookieJar, Cookie}; + /// use cookie::prefix::{Host, Secure}; + /// + /// async fn handler(jar: CookieJar) -> CookieJar { + /// // Add a cookie with the "__Host-" prefix + /// let with_host = jar.clone().add_prefixed(Host, Cookie::new("session_id", "value")); + /// + /// // Add a cookie with the "__Secure-" prefix + /// let _with_secure = jar.add_prefixed(Secure, Cookie::new("auth", "token")); + /// + /// with_host + /// } + /// ``` + #[must_use] + pub fn add_prefixed( + mut self, + prefix: P, + cookie: Cookie<'static>, + ) -> Self { + let mut prefixed_jar = self.jar.prefixed_mut(prefix); + prefixed_jar.add(cookie); + self + } + + /// Get a signed cookie with the specified prefix from the jar. + /// + /// If the cookie exists and its signature is valid, it is returned with its original name + /// (without the prefix) and plaintext value. + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::{CookieJar, Cookie}; + /// use cookie::prefix::{Host, Secure}; + /// + /// async fn handler(jar: CookieJar) { + /// if let Some(cookie) = jar.get_prefixed(cookie::prefix::Host, "session_id") { + /// let value = cookie.value(); + /// } + /// } + /// ``` + pub fn get_prefixed( + &self, + prefix: P, + name: &str, + ) -> Option> { + let prefixed_jar = self.jar.prefixed(prefix); + prefixed_jar.get(name) + } + + /// Remove a cookie with the specified prefix from the jar. + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::CookieJar; + /// use cookie::prefix::{Host, Secure}; + /// + /// async fn handler(jar: CookieJar) -> CookieJar { + /// // Remove a cookie with the "__Host-" prefix + /// jar.remove_prefixed(Host, "session_id") + /// } + /// ``` + #[must_use] + pub fn remove_prefixed(mut self, prefix: P, name: S) -> Self + where + P: cookie::prefix::Prefix, + S: Into, + { + let mut prefixed_jar = self.jar.prefixed_mut(prefix); + prefixed_jar.remove(name.into()); + self + } } impl IntoResponseParts for CookieJar { @@ -232,6 +308,7 @@ fn set_cookies(jar: cookie::CookieJar, headers: &mut HeaderMap) { mod tests { use super::*; use axum::{body::Body, extract::FromRef, http::Request, routing::get, Router}; + use cookie::prefix::Host; use http_body_util::BodyExt; use tower::ServiceExt; @@ -269,6 +346,18 @@ mod tests { .unwrap(); let cookie_value = res.headers()["set-cookie"].to_str().unwrap(); + assert!(cookie_value.starts_with("key=")); + + // For signed/private cookies, verify that the plaintext value is not directly visible + // (only for signed and private jars, not for the regular CookieJar) + if std::any::type_name::<$jar>().contains("Private") + || std::any::type_name::<$jar>().contains("Signed") + { + assert!(!cookie_value.contains("key=value")); + } else { + assert!(cookie_value.contains("key=value")); + } + let res = app .clone() .oneshot( @@ -302,17 +391,113 @@ mod tests { }; } + macro_rules! cookie_prefixed_test { + ($name:ident, $jar:ty) => { + #[tokio::test] + async fn $name() { + async fn set_cookie_prefixed(jar: $jar) -> impl IntoResponse { + jar.add_prefixed(Host, Cookie::new("key", "value")) + } + + async fn get_cookie_prefixed(jar: $jar) -> impl IntoResponse { + jar.get_prefixed(Host, "key").unwrap().value().to_owned() + } + + async fn remove_cookie_prefixed(jar: $jar) -> impl IntoResponse { + jar.remove_prefixed(Host, "key") + } + + let state = AppState { + key: Key::generate(), + custom_key: CustomKey(Key::generate()), + }; + + let app = Router::new() + .route("/set", get(set_cookie_prefixed)) + .route("/get", get(get_cookie_prefixed)) + .route("/remove", get(remove_cookie_prefixed)) + .with_state(state); + + let res = app + .clone() + .oneshot(Request::builder().uri("/set").body(Body::empty()).unwrap()) + .await + .unwrap(); + let cookie_value = res.headers()["set-cookie"].to_str().unwrap(); + assert!(cookie_value.contains("__Host-key")); + + // For signed/private cookies, verify that the plaintext value is not directly visible + // (only for signed and private jars, not for the regular CookieJar) + if std::any::type_name::<$jar>().contains("Private") + || std::any::type_name::<$jar>().contains("Signed") + { + assert!(!cookie_value.contains("key=value")); + } else { + assert!(cookie_value.contains("key=value")); + } + + // Extract just the cookie part (before the first semicolon) + // Set-Cookie: __Host-key=value; Secure; Path=/ -> __Host-key=value + let cookie_header_value = cookie_value.split(';').next().unwrap().trim(); + + let res = app + .clone() + .oneshot( + Request::builder() + .uri("/get") + .header("cookie", cookie_header_value) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = body_text(res).await; + assert_eq!(body, "value"); + + let res = app + .clone() + .oneshot( + Request::builder() + .uri("/remove") + .header("cookie", cookie_value) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert!(res.headers()["set-cookie"] + .to_str() + .unwrap() + .contains("__Host-key=;")); + } + }; + } + cookie_test!(plaintext_cookies, CookieJar); #[cfg(feature = "cookie-signed")] cookie_test!(signed_cookies, SignedCookieJar); #[cfg(feature = "cookie-signed")] + cookie_prefixed_test!(signed_cookies_prefixed, SignedCookieJar); + #[cfg(feature = "cookie-signed")] cookie_test!(signed_cookies_with_custom_key, SignedCookieJar); + #[cfg(feature = "cookie-signed")] + cookie_prefixed_test!( + signed_cookies_prefixed_with_custom_key, + SignedCookieJar + ); #[cfg(feature = "cookie-private")] cookie_test!(private_cookies, PrivateCookieJar); #[cfg(feature = "cookie-private")] + cookie_prefixed_test!(private_cookies_prefixed, PrivateCookieJar); + #[cfg(feature = "cookie-private")] cookie_test!(private_cookies_with_custom_key, PrivateCookieJar); + #[cfg(feature = "cookie-private")] + cookie_prefixed_test!( + private_cookies_prefixed_with_custom_key, + PrivateCookieJar + ); #[derive(Clone)] struct AppState { diff --git a/axum-extra/src/extract/cookie/private.rs b/axum-extra/src/extract/cookie/private.rs index f852b8c4ba..7611b2bef2 100644 --- a/axum-extra/src/extract/cookie/private.rs +++ b/axum-extra/src/extract/cookie/private.rs @@ -247,7 +247,7 @@ impl PrivateCookieJar { /// Authenticates and decrypts `cookie`, returning the plaintext version if decryption succeeds /// or `None` otherwise. pub fn decrypt(&self, cookie: Cookie<'static>) -> Option> { - self.private_jar().decrypt(cookie) + self.private_jar().decrypt(cookie.clone()) } /// Get an iterator over all cookies in the jar. @@ -267,6 +267,86 @@ impl PrivateCookieJar { fn private_jar_mut(&mut self) -> PrivateJar<&'_ mut cookie::CookieJar> { self.jar.private_mut(&self.key) } + /// Add a signed cookie with the specified prefix to the jar. + /// + /// The cookie's value will be signed using the jar's key, and the prefix will determine the + /// cookie's name and attributes (e.g., `Secure`, `Path=/` for `__Host-`). + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::{PrivateCookieJar, Cookie}; + /// use cookie::prefix::Host; + /// + /// async fn handler(jar: PrivateCookieJar) -> PrivateCookieJar { + /// jar.add_prefixed(Host, Cookie::new("session_id", "value")) + /// } + /// ``` + #[must_use] + pub fn add_prefixed( + self, + _prefix: P, + cookie: Cookie<'static>, + ) -> Self { + let mut jar = self.jar; + jar.remove(Cookie::new(cookie.name().to_owned(), "")); + + let prefixed_name = format!("{}{}", P::PREFIX, cookie.name()); + let mut new_cookie = cookie; + new_cookie.set_name(prefixed_name); + jar.private_mut(&self.key).add(new_cookie); + + Self { + jar, + key: self.key, + _marker: self._marker, + } + } + /// Get a signed cookie with the specified prefix from the jar. + /// + /// If the cookie exists and its signature is valid, it is returned with its original name + /// (without the prefix) and plaintext value. + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::PrivateCookieJar; + /// + /// async fn handler(jar: PrivateCookieJar) { + /// if let Some(cookie) = jar.get_prefixed(cookie::prefix::Host, "session_id") { + /// let value = cookie.value(); + /// } + /// } + /// ``` + pub fn get_prefixed( + &self, + _prefix: P, + name: &str, + ) -> Option> { + let prefixed_name = format!("{}{name}", P::PREFIX); + self.jar + .get(&prefixed_name) + .and_then(|c| self.decrypt(c.clone())) + } + /// Remove a signed cookie with the specified prefix from the jar. + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::PrivateCookieJar; + /// use cookie::prefix::Host; + /// + /// async fn handler(jar: PrivateCookieJar) -> PrivateCookieJar { + /// jar.remove_prefixed(Host, "session_id") + /// } + /// ``` + #[must_use] + pub fn remove_prefixed(mut self, prefix: P, name: S) -> Self + where + P: cookie::prefix::Prefix, + S: Into, + { + let mut prefixed_jar = self.jar.prefixed_mut(prefix); + prefixed_jar.remove(name.into()); + self + } } impl IntoResponseParts for PrivateCookieJar { diff --git a/axum-extra/src/extract/cookie/signed.rs b/axum-extra/src/extract/cookie/signed.rs index 92bf917145..1ac7a68a1b 100644 --- a/axum-extra/src/extract/cookie/signed.rs +++ b/axum-extra/src/extract/cookie/signed.rs @@ -285,6 +285,103 @@ impl SignedCookieJar { fn signed_jar_mut(&mut self) -> SignedJar<&'_ mut cookie::CookieJar> { self.jar.signed_mut(&self.key) } + /// Add a signed cookie with the specified prefix to the jar. + /// + /// The cookie's value will be signed using the jar's key, and the prefix will determine the + /// cookie's name and attributes (e.g., `Secure`, `Path=/` for `__Host-`). + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::{SignedCookieJar, Cookie}; + /// use cookie::prefix::Host; + /// + /// async fn handler(jar: SignedCookieJar) -> SignedCookieJar { + /// jar.add_prefixed(Host, Cookie::new("session_id", "value")) + /// } + /// ``` + #[must_use] + pub fn add_prefixed( + self, + prefix: P, + cookie: Cookie<'static>, + ) -> Self { + // Step 1: First add the cookie to the jar normally, which signs its value + let jar_with_signed_cookie = self.add(cookie.clone()); + let cookie_name = cookie.name().to_owned(); + let mut modified_jar = jar_with_signed_cookie; + + // Step 2: Retrieve the signed cookie that was just added + if let Some(signed_cookie) = modified_jar.jar.get(&cookie_name) { + // Extract the signed value (the value with signature attached) + let signed_value = signed_cookie.value().to_owned(); + + // Step 3: Remove the original non-prefixed cookie + modified_jar + .jar + .remove(Cookie::new(cookie_name.clone(), "")); + + // Step 4: Create a prefixed jar to handle proper attribute enforcement + // (prefixed cookies require specific attributes like Secure, Path=/, etc.) + let mut prefixed_jar = modified_jar.jar.prefixed_mut(prefix); + + // Step 5: Create a new cookie with the same base name but with the signed value + let mut prefixed_cookie = cookie.clone(); + prefixed_cookie.set_value(signed_value); + + // Step 6: Add the cookie to the prefixed jar, which will: + // - Add the prefix to the name (e.g., __Host- or __Secure-) + // - Set required security attributes based on the prefix + prefixed_jar.add(prefixed_cookie); + } + + modified_jar + } + /// Get a signed cookie with the specified prefix from the jar. + /// + /// If the cookie exists and its signature is valid, it is returned with its original name + /// (without the prefix) and plaintext value. + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::SignedCookieJar; + /// + /// async fn handler(jar: SignedCookieJar) { + /// if let Some(cookie) = jar.get_prefixed(cookie::prefix::Host, "session_id") { + /// let value = cookie.value(); + /// } + /// } + /// ``` + pub fn get_prefixed( + &self, + _prefix: P, + name: &str, + ) -> Option> { + let prefixed_name = format!("{}{name}", P::PREFIX); + self.jar + .get(&prefixed_name) + .and_then(|c| self.verify(c.clone())) + } + /// Remove a signed cookie with the specified prefix from the jar. + /// + /// # Example + /// ```rust + /// use axum_extra::extract::cookie::SignedCookieJar; + /// use cookie::prefix::Host; + /// + /// async fn handler(jar: SignedCookieJar) -> SignedCookieJar { + /// jar.remove_prefixed(Host, "session_id") + /// } + /// ``` + #[must_use] + pub fn remove_prefixed(mut self, prefix: P, name: S) -> Self + where + P: cookie::prefix::Prefix, + S: Into, + { + let mut prefixed_jar = self.jar.prefixed_mut(prefix); + prefixed_jar.remove(name.into()); + self + } } impl IntoResponseParts for SignedCookieJar {