1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220
use super::{cookies_from_request, set_cookies, Cookie, Key};
use axum::{
async_trait,
extract::{FromRequest, RequestParts},
response::{IntoResponse, IntoResponseParts, Response, ResponseParts},
Extension,
};
use cookie_lib::PrivateJar;
use std::{convert::Infallible, fmt, marker::PhantomData};
/// Extractor that grabs private cookies from the request and manages the jar.
///
/// All cookies will be private and encrypted with a [`Key`]. This makes it suitable for storing
/// private data.
///
/// Note that methods like [`PrivateCookieJar::add`], [`PrivateCookieJar::remove`], etc updates the
/// [`PrivateCookieJar`] and returns it. This value _must_ be returned from the handler as part of
/// the response for the changes to be propagated.
///
/// # Example
///
/// ```rust
/// use axum::{
/// Router,
/// Extension,
/// routing::{post, get},
/// extract::TypedHeader,
/// response::{IntoResponse, Redirect},
/// headers::authorization::{Authorization, Bearer},
/// http::StatusCode,
/// };
/// use axum_extra::extract::cookie::{PrivateCookieJar, Cookie, Key};
///
/// async fn set_secret(
/// jar: PrivateCookieJar,
/// ) -> (PrivateCookieJar, Redirect) {
/// let updated_jar = jar.add(Cookie::new("secret", "secret-data"));
/// (updated_jar, Redirect::to("/get"))
/// }
///
/// async fn get_secret(jar: PrivateCookieJar) {
/// if let Some(data) = jar.get("secret") {
/// // ...
/// }
/// }
///
/// // Generate a secure key
/// //
/// // You probably don't wanna generate a new one each time the app starts though
/// let key = Key::generate();
///
/// let app = Router::new()
/// .route("/set", post(set_secret))
/// .route("/get", get(get_secret))
/// // add extension with the key so `PrivateCookieJar` can access it
/// .layer(Extension(key));
/// # let app: Router<axum::body::Body> = app;
/// ```
pub struct PrivateCookieJar<K = Key> {
jar: cookie_lib::CookieJar,
key: Key,
// The key used to extract the key extension. Allows users to use multiple keys for different
// jars. Maybe a library wants its own key.
_marker: PhantomData<K>,
}
impl<K> fmt::Debug for PrivateCookieJar<K> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PrivateCookieJar")
.field("jar", &self.jar)
.field("key", &"REDACTED")
.finish()
}
}
#[async_trait]
impl<B, K> FromRequest<B> for PrivateCookieJar<K>
where
B: Send,
K: Into<Key> + Clone + Send + Sync + 'static,
{
type Rejection = <axum::Extension<K> as FromRequest<B>>::Rejection;
async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
let key = Extension::<K>::from_request(req).await?.0.into();
let mut jar = cookie_lib::CookieJar::new();
let mut private_jar = jar.private_mut(&key);
for cookie in cookies_from_request(req) {
if let Some(cookie) = private_jar.decrypt(cookie) {
private_jar.add_original(cookie);
}
}
Ok(Self {
jar,
key,
_marker: PhantomData,
})
}
}
impl<K> PrivateCookieJar<K> {
/// Get a cookie from the jar.
///
/// If the cookie exists and can be decrypted then it is returned in plaintext.
///
/// # Example
///
/// ```rust
/// use axum_extra::extract::cookie::PrivateCookieJar;
/// use axum::response::IntoResponse;
///
/// async fn handle(jar: PrivateCookieJar) {
/// let value: Option<String> = jar
/// .get("foo")
/// .map(|cookie| cookie.value().to_owned());
/// }
/// ```
pub fn get(&self, name: &str) -> Option<Cookie<'static>> {
self.private_jar().get(name)
}
/// Remove a cookie from the jar.
///
/// # Example
///
/// ```rust
/// use axum_extra::extract::cookie::{PrivateCookieJar, Cookie};
/// use axum::response::IntoResponse;
///
/// async fn handle(jar: PrivateCookieJar) -> PrivateCookieJar {
/// jar.remove(Cookie::named("foo"))
/// }
/// ```
#[must_use]
pub fn remove(mut self, cookie: Cookie<'static>) -> Self {
self.private_jar_mut().remove(cookie);
self
}
/// Add a cookie to the jar.
///
/// The value will automatically be percent-encoded.
///
/// # Example
///
/// ```rust
/// use axum_extra::extract::cookie::{PrivateCookieJar, Cookie};
/// use axum::response::IntoResponse;
///
/// async fn handle(jar: PrivateCookieJar) -> PrivateCookieJar {
/// jar.add(Cookie::new("foo", "bar"))
/// }
/// ```
#[must_use]
#[allow(clippy::should_implement_trait)]
pub fn add(mut self, cookie: Cookie<'static>) -> Self {
self.private_jar_mut().add(cookie);
self
}
/// Authenticates and decrypts `cookie`, returning the plaintext version if decryption succeeds
/// or `None` otherwise.
pub fn decrypt(&self, cookie: Cookie<'static>) -> Option<Cookie<'static>> {
self.private_jar().decrypt(cookie)
}
/// Get an iterator over all cookies in the jar.
///
/// Only cookies with valid authenticity and integrity are yielded by the iterator.
pub fn iter(&self) -> impl Iterator<Item = Cookie<'static>> + '_ {
PrivateCookieJarIter {
jar: self,
iter: self.jar.iter(),
}
}
fn private_jar(&self) -> PrivateJar<&'_ cookie_lib::CookieJar> {
self.jar.private(&self.key)
}
fn private_jar_mut(&mut self) -> PrivateJar<&'_ mut cookie_lib::CookieJar> {
self.jar.private_mut(&self.key)
}
}
impl<K> IntoResponseParts for PrivateCookieJar<K> {
type Error = Infallible;
fn into_response_parts(self, mut res: ResponseParts) -> Result<ResponseParts, Self::Error> {
set_cookies(self.jar, res.headers_mut());
Ok(res)
}
}
impl<K> IntoResponse for PrivateCookieJar<K> {
fn into_response(self) -> Response {
(self, ()).into_response()
}
}
struct PrivateCookieJarIter<'a, K> {
jar: &'a PrivateCookieJar<K>,
iter: cookie_lib::Iter<'a>,
}
impl<'a, K> Iterator for PrivateCookieJarIter<'a, K> {
type Item = Cookie<'static>;
fn next(&mut self) -> Option<Self::Item> {
loop {
let cookie = self.iter.next()?;
if let Some(cookie) = self.jar.get(cookie.name()) {
return Some(cookie);
}
}
}
}