use super::{cookies_from_request, set_cookies};
use axum::{
async_trait,
extract::{FromRef, FromRequestParts},
response::{IntoResponse, IntoResponseParts, Response, ResponseParts},
};
use cookie::SignedJar;
use cookie::{Cookie, Key};
use http::{request::Parts, HeaderMap};
use std::{convert::Infallible, fmt, marker::PhantomData};
pub struct SignedCookieJar<K = Key> {
jar: cookie::CookieJar,
key: Key,
_marker: PhantomData<K>,
}
impl<K> fmt::Debug for SignedCookieJar<K> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SignedCookieJar")
.field("jar", &self.jar)
.field("key", &"REDACTED")
.finish()
}
}
#[async_trait]
impl<S, K> FromRequestParts<S> for SignedCookieJar<K>
where
S: Send + Sync,
K: FromRef<S> + Into<Key>,
{
type Rejection = Infallible;
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
let k = K::from_ref(state);
let key = k.into();
let SignedCookieJar {
jar,
key,
_marker: _,
} = SignedCookieJar::from_headers(&parts.headers, key);
Ok(SignedCookieJar {
jar,
key,
_marker: PhantomData,
})
}
}
impl SignedCookieJar {
pub fn from_headers(headers: &HeaderMap, key: Key) -> Self {
let mut jar = cookie::CookieJar::new();
let mut signed_jar = jar.signed_mut(&key);
for cookie in cookies_from_request(headers) {
if let Some(cookie) = signed_jar.verify(cookie) {
signed_jar.add_original(cookie);
}
}
Self {
jar,
key,
_marker: PhantomData,
}
}
pub fn new(key: Key) -> Self {
Self {
jar: Default::default(),
key,
_marker: PhantomData,
}
}
}
impl<K> SignedCookieJar<K> {
pub fn get(&self, name: &str) -> Option<Cookie<'static>> {
self.signed_jar().get(name)
}
#[must_use]
pub fn remove(mut self, cookie: Cookie<'static>) -> Self {
self.signed_jar_mut().remove(cookie);
self
}
#[must_use]
#[allow(clippy::should_implement_trait)]
pub fn add(mut self, cookie: Cookie<'static>) -> Self {
self.signed_jar_mut().add(cookie);
self
}
pub fn verify(&self, cookie: Cookie<'static>) -> Option<Cookie<'static>> {
self.signed_jar().verify(cookie)
}
pub fn iter(&self) -> impl Iterator<Item = Cookie<'static>> + '_ {
SignedCookieJarIter {
jar: self,
iter: self.jar.iter(),
}
}
fn signed_jar(&self) -> SignedJar<&'_ cookie::CookieJar> {
self.jar.signed(&self.key)
}
fn signed_jar_mut(&mut self) -> SignedJar<&'_ mut cookie::CookieJar> {
self.jar.signed_mut(&self.key)
}
}
impl<K> IntoResponseParts for SignedCookieJar<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 SignedCookieJar<K> {
fn into_response(self) -> Response {
(self, ()).into_response()
}
}
struct SignedCookieJarIter<'a, K> {
jar: &'a SignedCookieJar<K>,
iter: cookie::Iter<'a>,
}
impl<'a, K> Iterator for SignedCookieJarIter<'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);
}
}
}
}