Handle cookies better by setting the right flags & expiration

crates/handlers/src/upstream_oauth2/authorize.rs

@@ -16,10 +16,8 @@ use axum::{
     extract::{Path, Query, State},
     response::{IntoResponse, Redirect},
 };
-use axum_extra::extract::PrivateCookieJar;
 use hyper::StatusCode;
-use mas_axum_utils::http_client_factory::HttpClientFactory;
-use mas_keystore::Encrypter;
+use mas_axum_utils::{cookies::CookieJar, http_client_factory::HttpClientFactory};
 use mas_oidc_client::requests::authorization_code::AuthorizationRequestData;
 use mas_router::UrlBuilder;
 use mas_storage::{
@@ -68,7 +66,7 @@ pub(crate) async fn get(
     State(http_client_factory): State<HttpClientFactory>,
     mut repo: BoxRepository,
     State(url_builder): State<UrlBuilder>,
-    cookie_jar: PrivateCookieJar<Encrypter>,
+    cookie_jar: CookieJar,
     Path(provider_id): Path<Ulid>,
     Query(query): Query<OptionalPostAuthAction>,
 ) -> Result<impl IntoResponse, RouteError> {

crates/handlers/src/upstream_oauth2/callback.rs

@@ -16,9 +16,8 @@ use axum::{
     extract::{Path, Query, State},
     response::IntoResponse,
 };
-use axum_extra::extract::PrivateCookieJar;
 use hyper::StatusCode;
-use mas_axum_utils::http_client_factory::HttpClientFactory;
+use mas_axum_utils::{cookies::CookieJar, http_client_factory::HttpClientFactory};
 use mas_jose::claims::ClaimError;
 use mas_keystore::{Encrypter, Keystore};
 use mas_oidc_client::requests::{
@@ -133,7 +132,7 @@ pub(crate) async fn get(
     State(url_builder): State<UrlBuilder>,
     State(encrypter): State<Encrypter>,
     State(keystore): State<Keystore>,
-    cookie_jar: PrivateCookieJar<Encrypter>,
+    cookie_jar: CookieJar,
     Path(provider_id): Path<Ulid>,
     Query(params): Query<QueryParams>,
 ) -> Result<impl IntoResponse, RouteError> {

crates/handlers/src/upstream_oauth2/cookie.rs

@@ -14,14 +14,12 @@
 
 // TODO: move that to a standalone cookie manager
 
-use axum_extra::extract::{cookie::Cookie, PrivateCookieJar};
 use chrono::{DateTime, Duration, NaiveDateTime, Utc};
-use mas_axum_utils::CookieExt;
+use mas_axum_utils::cookies::CookieJar;
 use mas_router::PostAuthAction;
 use mas_storage::Clock;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
-use time::OffsetDateTime;
 use ulid::Ulid;
 
 /// Name of the cookie
@@ -62,30 +60,24 @@ pub struct UpstreamSessionNotFound;
 
 impl UpstreamSessions {
     /// Load the upstreams sessions cookie
-    pub fn load<K>(cookie_jar: &PrivateCookieJar<K>) -> Self {
-        cookie_jar
-            .get(COOKIE_NAME)
-            .and_then(|c| c.decode().ok())
-            .unwrap_or_default()
+    pub fn load(cookie_jar: &CookieJar) -> Self {
+        match cookie_jar.load(COOKIE_NAME) {
+            Ok(Some(sessions)) => sessions,
+            Ok(None) => Self::default(),
+            Err(e) => {
+                tracing::warn!("Invalid upstream sessions cookie: {}", e);
+                Self::default()
+            }
+        }
     }
 
     /// Save the upstreams sessions to the cookie jar
-    pub fn save<K, C>(self, cookie_jar: PrivateCookieJar<K>, clock: &C) -> PrivateCookieJar<K>
+    pub fn save<C>(self, cookie_jar: CookieJar, clock: &C) -> CookieJar
     where
         C: Clock,
     {
-        let now = clock.now();
-        let this = self.expire(now);
-        let mut cookie = Cookie::named(COOKIE_NAME).encode(&this);
-        cookie.set_path("/");
-        cookie.set_http_only(true);
-
-        let expiration = now + Duration::seconds(SESSION_MAX_TIME_SECS);
-        let expiration = OffsetDateTime::from_unix_timestamp(expiration.timestamp())
-            .expect("invalid unix timestamp");
-        cookie.set_expires(expiration);
-
-        cookie_jar.add(cookie)
+        let this = self.expire(clock.now());
+        cookie_jar.save(COOKIE_NAME, &this, false)
     }
 
     fn expire(mut self, now: DateTime<Utc>) -> Self {

crates/handlers/src/upstream_oauth2/link.rs

@@ -17,15 +17,14 @@ use axum::{
     response::{Html, IntoResponse},
     Form,
 };
-use axum_extra::extract::PrivateCookieJar;
 use hyper::StatusCode;
 use mas_axum_utils::{
+    cookies::CookieJar,
     csrf::{CsrfExt, ProtectedForm},
     SessionInfoExt,
 };
 use mas_data_model::{UpstreamOAuthProviderImportPreference, User};
 use mas_jose::jwt::Jwt;
-use mas_keystore::Encrypter;
 use mas_storage::{
     job::{JobRepositoryExt, ProvisionUserJob},
     upstream_oauth2::{UpstreamOAuthLinkRepository, UpstreamOAuthSessionRepository},
@@ -170,7 +169,7 @@ pub(crate) async fn get(
     clock: BoxClock,
     mut repo: BoxRepository,
     State(templates): State<Templates>,
-    cookie_jar: PrivateCookieJar<Encrypter>,
+    cookie_jar: CookieJar,
     Path(link_id): Path<Ulid>,
 ) -> Result<impl IntoResponse, RouteError> {
     let sessions_cookie = UpstreamSessionsCookie::load(&cookie_jar);
@@ -350,7 +349,7 @@ pub(crate) async fn post(
     mut rng: BoxRng,
     clock: BoxClock,
     mut repo: BoxRepository,
-    cookie_jar: PrivateCookieJar<Encrypter>,
+    cookie_jar: CookieJar,
     Path(link_id): Path<Ulid>,
     Form(form): Form<ProtectedForm<FormData>>,
 ) -> Result<impl IntoResponse, RouteError> {