Use axum-extra's PrivateCookieJar

Cargo.lock

@@ -514,6 +514,24 @@ dependencies = [
  "mime",
 ]
 
+[[package]]
+name = "axum-extra"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a00a7085c512df12d4e07a862b23a3b3bfe5326ecfc4185b49fb15c8850ba406"
+dependencies = [
+ "axum",
+ "bytes 1.1.0",
+ "cookie",
+ "http",
+ "mime",
+ "pin-project-lite",
+ "tower",
+ "tower-http",
+ "tower-layer",
+ "tower-service",
+]
+
 [[package]]
 name = "axum-macros"
 version = "0.2.0"
@@ -1943,9 +1961,9 @@ version = "0.1.0"
 dependencies = [
  "async-trait",
  "axum",
+ "axum-extra",
  "bincode",
  "chrono",
- "cookie",
  "data-encoding",
  "futures-util",
  "headers",
@@ -2080,6 +2098,7 @@ dependencies = [
  "anyhow",
  "argon2",
  "axum",
+ "axum-extra",
  "axum-macros",
  "chrono",
  "crc",

crates/axum-utils/Cargo.toml

@@ -8,9 +8,9 @@ license = "Apache-2.0"
 [dependencies]
 async-trait = "0.1.53"
 axum = { version = "0.5.4", features = ["headers"] }
+axum-extra = { version = "0.3.0", features = ["cookie-private"] }
 bincode = "1.3.3"
 chrono = "0.4.19"
-cookie = { version = "0.16.0", features = ["private", "percent-encode"] }
 data-encoding = "2.3.2"
 futures-util = "0.3.21"
 headers = "0.3.7"

crates/axum-utils/src/cookies.rs

@@ -14,113 +14,10 @@
 
 //! Private (encrypted) cookie jar, based on axum-extra's cookie jar
 
-use std::{convert::Infallible, marker::PhantomData};
-
-use async_trait::async_trait;
-use axum::{
-    extract::{Extension, FromRequest, RequestParts},
-    response::IntoResponseParts,
-};
-pub use cookie::Cookie;
 use data_encoding::BASE64URL_NOPAD;
-use headers::HeaderMap;
-use http::header::{COOKIE, SET_COOKIE};
 use serde::{de::DeserializeOwned, Serialize};
 use thiserror::Error;
 
-pub struct PrivateCookieJar<K = cookie::Key> {
-    jar: cookie::CookieJar,
-    key: cookie::Key,
-    _marker: PhantomData<K>,
-}
-
-impl<K> PrivateCookieJar<K> {
-    pub fn get(&self, name: &str) -> Option<Cookie<'static>> {
-        self.private_jar().get(name)
-    }
-
-    #[must_use]
-    pub fn remove(mut self, cookie: Cookie<'static>) -> Self {
-        self.private_jar_mut().remove(cookie);
-        self
-    }
-
-    #[must_use]
-    #[allow(clippy::should_implement_trait)]
-    pub fn add(mut self, cookie: Cookie<'static>) -> Self {
-        self.private_jar_mut().add(cookie);
-        self
-    }
-
-    pub fn decrypt(&self, cookie: Cookie<'static>) -> Option<Cookie<'static>> {
-        self.private_jar().decrypt(cookie)
-    }
-
-    fn private_jar(&self) -> cookie::PrivateJar<&'_ cookie::CookieJar> {
-        self.jar.private(&self.key)
-    }
-
-    fn private_jar_mut(&mut self) -> cookie::PrivateJar<&'_ mut cookie::CookieJar> {
-        self.jar.private_mut(&self.key)
-    }
-
-    pub fn set_cookies(self, headers: &mut HeaderMap) {
-        for cookie in self.jar.delta() {
-            if let Ok(header_value) = cookie.encoded().to_string().parse() {
-                headers.append(SET_COOKIE, header_value);
-            }
-        }
-    }
-}
-
-#[async_trait]
-impl<B, K> FromRequest<B> for PrivateCookieJar<K>
-where
-    B: Send,
-    K: Into<cookie::Key> + Clone + Send + Sync + 'static,
-{
-    type Rejection = <Extension<K> as FromRequest<B>>::Rejection;
-
-    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
-        let Extension(key): Extension<K> = Extension::from_request(req).await?;
-        let key = key.into();
-
-        let mut jar = cookie::CookieJar::new();
-        let mut private_jar = jar.private_mut(&key);
-
-        let cookies = req
-            .headers()
-            .get_all(COOKIE)
-            .into_iter()
-            .filter_map(|value| value.to_str().ok())
-            .flat_map(|value| value.split(';'))
-            .filter_map(|cookie| Cookie::parse_encoded(cookie.to_owned()).ok());
-
-        for cookie in cookies {
-            if let Some(cookie) = private_jar.decrypt(cookie) {
-                private_jar.add_original(cookie);
-            }
-        }
-
-        Ok(Self {
-            jar,
-            key,
-            _marker: PhantomData,
-        })
-    }
-}
-
-impl<K> IntoResponseParts for PrivateCookieJar<K> {
-    type Error = Infallible;
-    fn into_response_parts(
-        self,
-        mut res: axum::response::ResponseParts,
-    ) -> Result<axum::response::ResponseParts, Self::Error> {
-        self.set_cookies(res.headers_mut());
-        Ok(res)
-    }
-}
-
 #[derive(Debug, Error)]
 #[error("could not decode cookie")]
 pub enum CookieDecodeError {
@@ -138,7 +35,7 @@ pub trait CookieExt {
         T: Serialize;
 }
 
-impl<'a> CookieExt for Cookie<'a> {
+impl<'a> CookieExt for axum_extra::extract::cookie::Cookie<'a> {
     fn decode<T>(&self) -> Result<T, CookieDecodeError>
     where
         T: DeserializeOwned,

crates/axum-utils/src/csrf.rs

@@ -12,14 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use axum_extra::extract::cookie::{Cookie, PrivateCookieJar};
 use chrono::{DateTime, Duration, Utc};
-use cookie::Cookie;
 use data_encoding::{DecodeError, BASE64URL_NOPAD};
 use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, TimestampSeconds};
 use thiserror::Error;
 
-use crate::{cookies::CookieDecodeError, CookieExt, PrivateCookieJar};
+use crate::{cookies::CookieDecodeError, CookieExt};
 
 /// Failed to validate CSRF token
 #[derive(Debug, Error)]

crates/axum-utils/src/lib.rs

@@ -21,7 +21,7 @@ pub mod url_builder;
 pub mod user_authorization;
 
 pub use self::{
-    cookies::{Cookie, CookieExt, PrivateCookieJar},
+    cookies::CookieExt,
     fancy_error::{fancy_error, internal_error, FancyError},
     session::{SessionInfo, SessionInfoExt},
     url_builder::UrlBuilder,

crates/axum-utils/src/session.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use cookie::Cookie;
+use axum_extra::extract::cookie::{Cookie, PrivateCookieJar};
 use mas_data_model::BrowserSession;
 use mas_storage::{
     user::{lookup_active_session, ActiveSessionLookupError},
@@ -21,7 +21,7 @@ use mas_storage::{
 use serde::{Deserialize, Serialize};
 use sqlx::{Executor, Postgres};
 
-use crate::{CookieExt, PrivateCookieJar};
+use crate::CookieExt;
 
 /// An encrypted cookie to save the session ID
 #[derive(Serialize, Deserialize, Debug, Default)]

crates/handlers/Cargo.toml

@@ -22,6 +22,7 @@ tower = "0.4.12"
 tower-http = { version = "0.3.1", features = ["cors"] }
 axum = "0.5.4"
 axum-macros = "0.2.0"
+axum-extra = { version = "0.3.0", features = ["cookie-private"] }
 
 # Emails
 lettre = { version = "0.10.0-rc.5", default-features = false, features = ["builder"] }

crates/handlers/src/oauth2/authorization.rs

@@ -19,12 +19,13 @@ use axum::{
     extract::{Extension, Form, Query},
     response::{Html, IntoResponse, Redirect, Response},
 };
+use axum_extra::extract::PrivateCookieJar;
 use chrono::Duration;
 use hyper::{
     http::uri::{Parts, PathAndQuery, Uri},
     StatusCode,
 };
-use mas_axum_utils::{PrivateCookieJar, SessionInfoExt};
+use mas_axum_utils::SessionInfoExt;
 use mas_config::Encrypter;
 use mas_data_model::{
     Authentication, AuthorizationCode, AuthorizationGrant, AuthorizationGrantStage, BrowserSession,

crates/handlers/src/oauth2/consent.rs

@@ -18,10 +18,11 @@ use axum::{
     http::uri::{Parts, PathAndQuery},
     response::{Html, IntoResponse, Redirect, Response},
 };
+use axum_extra::extract::PrivateCookieJar;
 use hyper::{StatusCode, Uri};
 use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
-    PrivateCookieJar, SessionInfoExt,
+    SessionInfoExt,
 };
 use mas_config::Encrypter;
 use mas_data_model::AuthorizationGrantStage;

crates/handlers/src/views/account/emails.rs

@@ -16,10 +16,11 @@ use axum::{
     extract::{Extension, Form},
     response::{Html, IntoResponse, Redirect, Response},
 };
+use axum_extra::extract::PrivateCookieJar;
 use lettre::{message::Mailbox, Address};
 use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
-    fancy_error, FancyError, PrivateCookieJar, SessionInfoExt, UrlBuilder,
+    fancy_error, FancyError, SessionInfoExt, UrlBuilder,
 };
 use mas_config::Encrypter;
 use mas_data_model::{BrowserSession, User, UserEmail};

crates/handlers/src/views/account/mod.rs

@@ -19,7 +19,8 @@ use axum::{
     extract::Extension,
     response::{Html, IntoResponse, Redirect, Response},
 };
-use mas_axum_utils::{csrf::CsrfExt, fancy_error, FancyError, PrivateCookieJar, SessionInfoExt};
+use axum_extra::extract::PrivateCookieJar;
+use mas_axum_utils::{csrf::CsrfExt, fancy_error, FancyError, SessionInfoExt};
 use mas_config::Encrypter;
 use mas_storage::user::{count_active_sessions, get_user_emails};
 use mas_templates::{AccountContext, TemplateContext, Templates};

crates/handlers/src/views/account/password.rs

@@ -17,9 +17,10 @@ use axum::{
     extract::{Extension, Form},
     response::{Html, IntoResponse, Redirect, Response},
 };
+use axum_extra::extract::PrivateCookieJar;
 use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
-    fancy_error, FancyError, PrivateCookieJar, SessionInfoExt,
+    fancy_error, FancyError, SessionInfoExt,
 };
 use mas_config::Encrypter;
 use mas_data_model::BrowserSession;

crates/handlers/src/views/index.rs

@@ -16,9 +16,8 @@ use axum::{
     extract::Extension,
     response::{Html, IntoResponse},
 };
-use mas_axum_utils::{
-    csrf::CsrfExt, fancy_error, FancyError, PrivateCookieJar, SessionInfoExt, UrlBuilder,
-};
+use axum_extra::extract::PrivateCookieJar;
+use mas_axum_utils::{csrf::CsrfExt, fancy_error, FancyError, SessionInfoExt, UrlBuilder};
 use mas_config::Encrypter;
 use mas_templates::{IndexContext, TemplateContext, Templates};
 use sqlx::PgPool;

crates/handlers/src/views/login.rs

@@ -16,10 +16,11 @@ use axum::{
     extract::{Extension, Form, Query},
     response::{Html, IntoResponse, Redirect, Response},
 };
+use axum_extra::extract::PrivateCookieJar;
 use hyper::http::uri::{Parts, PathAndQuery, Uri};
 use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
-    fancy_error, FancyError, PrivateCookieJar, SessionInfoExt,
+    fancy_error, FancyError, SessionInfoExt,
 };
 use mas_config::Encrypter;
 use mas_data_model::errors::WrapFormError;

crates/handlers/src/views/logout.rs

@@ -16,9 +16,10 @@ use axum::{
     extract::{Extension, Form},
     response::{IntoResponse, Redirect},
 };
+use axum_extra::extract::PrivateCookieJar;
 use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
-    fancy_error, FancyError, PrivateCookieJar, SessionInfoExt,
+    fancy_error, FancyError, SessionInfoExt,
 };
 use mas_config::Encrypter;
 use mas_storage::user::end_session;

crates/handlers/src/views/reauth.rs

@@ -16,13 +16,14 @@ use axum::{
     extract::{Extension, Form, Query},
     response::{Html, IntoResponse, Redirect, Response},
 };
+use axum_extra::extract::PrivateCookieJar;
 use hyper::{
     http::uri::{Parts, PathAndQuery},
     Uri,
 };
 use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
-    fancy_error, FancyError, PrivateCookieJar, SessionInfoExt,
+    fancy_error, FancyError, SessionInfoExt,
 };
 use mas_config::Encrypter;
 use mas_storage::user::authenticate_session;

crates/handlers/src/views/register.rs

@@ -19,10 +19,11 @@ use axum::{
     extract::{Extension, Form, Query},
     response::{Html, IntoResponse, Redirect, Response},
 };
+use axum_extra::extract::PrivateCookieJar;
 use hyper::http::uri::{Parts, PathAndQuery, Uri};
 use mas_axum_utils::{
     csrf::{CsrfExt, ProtectedForm},
-    fancy_error, FancyError, PrivateCookieJar, SessionInfoExt,
+    fancy_error, FancyError, SessionInfoExt,
 };
 use mas_config::Encrypter;
 use mas_storage::user::{register_user, start_session};

crates/handlers/src/views/verify.rs

@@ -16,8 +16,9 @@ use axum::{
     extract::{Extension, Path},
     response::{Html, IntoResponse},
 };
+use axum_extra::extract::PrivateCookieJar;
 use chrono::Duration;
-use mas_axum_utils::{csrf::CsrfExt, fancy_error, FancyError, PrivateCookieJar, SessionInfoExt};
+use mas_axum_utils::{csrf::CsrfExt, fancy_error, FancyError, SessionInfoExt};
 use mas_config::Encrypter;
 use mas_storage::user::{
     consume_email_verification, lookup_user_email_verification_code, mark_user_email_as_verified,