From 21d6ab4beb57730ab2395098085d26a9328030d7 Mon Sep 17 00:00:00 2001
From: Nicolas SEYS <seys.nicolas@gmail.com>
Date: Wed, 29 Nov 2023 10:12:10 +0100
Subject: [PATCH] Make session and cookie age configurable

---
 warpgate-common/src/config/defaults.rs | 10 ++++++++++
 warpgate-common/src/config/mod.rs      |  8 ++++++++
 warpgate-protocol-http/src/common.rs   |  3 ---
 warpgate-protocol-http/src/lib.rs      | 13 ++++++++-----
 warpgate-protocol-http/src/session.rs  |  6 +++---
 5 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/warpgate-common/src/config/defaults.rs b/warpgate-common/src/config/defaults.rs
index 6299310..491fa7b 100644
--- a/warpgate-common/src/config/defaults.rs
+++ b/warpgate-common/src/config/defaults.rs
@@ -56,6 +56,16 @@ pub(crate) fn _default_retention() -> Duration {
     Duration::SECOND * 60 * 60 * 24 * 7
 }
 
+#[inline]
+pub(crate) fn _default_session_max_age() -> Duration {
+    Duration::SECOND * 60 * 30
+}
+
+#[inline]
+pub(crate) fn _default_cookie_max_age() -> Duration {
+    Duration::SECOND * 60 * 60 * 24
+}
+
 #[inline]
 pub(crate) fn _default_empty_vec<T>() -> Vec<T> {
     vec![]
diff --git a/warpgate-common/src/config/mod.rs b/warpgate-common/src/config/mod.rs
index 191ff3c..11c252c 100644
--- a/warpgate-common/src/config/mod.rs
+++ b/warpgate-common/src/config/mod.rs
@@ -254,6 +254,12 @@ pub struct WarpgateConfigStore {
     #[serde(default = "_default_database_url")]
     pub database_url: Secret<String>,
 
+    #[serde(default = "_default_session_max_age", with = "humantime_serde")]
+    pub session_max_age: Duration,
+
+    #[serde(default = "_default_cookie_max_age", with = "humantime_serde")]
+    pub cookie_max_age: Duration,
+
     #[serde(default)]
     pub ssh: SSHConfig,
 
@@ -280,6 +286,8 @@ impl Default for WarpgateConfigStore {
             recordings: <_>::default(),
             external_host: None,
             database_url: _default_database_url(),
+            session_max_age: _default_session_max_age(),
+            cookie_max_age: _default_cookie_max_age(),
             ssh: <_>::default(),
             http: <_>::default(),
             mysql: <_>::default(),
diff --git a/warpgate-protocol-http/src/common.rs b/warpgate-protocol-http/src/common.rs
index adba601..70a3265 100644
--- a/warpgate-protocol-http/src/common.rs
+++ b/warpgate-protocol-http/src/common.rs
@@ -1,5 +1,4 @@
 use std::sync::Arc;
-use std::time::Duration;
 
 use http::StatusCode;
 use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
@@ -19,8 +18,6 @@ pub const PROTOCOL_NAME: ProtocolName = "HTTP";
 static TARGET_SESSION_KEY: &str = "target_name";
 static AUTH_SESSION_KEY: &str = "auth";
 static AUTH_STATE_ID_SESSION_KEY: &str = "auth_state_id";
-pub static SESSION_MAX_AGE: Duration = Duration::from_secs(60 * 30);
-pub static COOKIE_MAX_AGE: Duration = Duration::from_secs(60 * 60 * 24);
 pub static SESSION_COOKIE_NAME: &str = "warpgate-http-session";
 
 pub trait SessionExt {
diff --git a/warpgate-protocol-http/src/lib.rs b/warpgate-protocol-http/src/lib.rs
index 49b5e18..08cc9bc 100644
--- a/warpgate-protocol-http/src/lib.rs
+++ b/warpgate-protocol-http/src/lib.rs
@@ -36,9 +36,7 @@ use warpgate_common::{
 use warpgate_core::{ProtocolServer, Services, TargetTestError};
 use warpgate_web::Assets;
 
-use crate::common::{
-    endpoint_admin_auth, endpoint_auth, page_auth, COOKIE_MAX_AGE, SESSION_COOKIE_NAME,
-};
+use crate::common::{endpoint_admin_auth, endpoint_auth, page_auth, SESSION_COOKIE_NAME};
 use crate::error::error_page;
 use crate::middleware::{CookieHostMiddleware, TicketMiddleware};
 use crate::session::{SessionStore, SharedSessionStorage};
@@ -86,6 +84,11 @@ impl ProtocolServer for HTTPProtocolServer {
             )
         };
 
+        let (cookie_max_age, session_max_age) = {
+            let config = self.services.config.lock().await;
+            (config.store.cookie_max_age, config.store.session_max_age)
+        };
+
         let app = Route::new()
             .nest(
                 "/@warpgate",
@@ -149,7 +152,7 @@ impl ProtocolServer for HTTPProtocolServer {
             .with(ServerSession::new(
                 CookieConfig::default()
                     .secure(false)
-                    .max_age(COOKIE_MAX_AGE)
+                    .max_age(cookie_max_age)
                     .name(SESSION_COOKIE_NAME),
                 session_storage.clone(),
             ))
@@ -160,7 +163,7 @@ impl ProtocolServer for HTTPProtocolServer {
 
         tokio::spawn(async move {
             loop {
-                session_store.lock().await.vacuum().await;
+                session_store.lock().await.vacuum(session_max_age).await;
                 tokio::time::sleep(Duration::from_secs(60)).await;
             }
         });
diff --git a/warpgate-protocol-http/src/session.rs b/warpgate-protocol-http/src/session.rs
index c5847ec..c7156f0 100644
--- a/warpgate-protocol-http/src/session.rs
+++ b/warpgate-protocol-http/src/session.rs
@@ -12,7 +12,7 @@ use tracing::*;
 use warpgate_common::SessionId;
 use warpgate_core::{Services, SessionStateInit, WarpgateServerHandle};
 
-use crate::common::{PROTOCOL_NAME, SESSION_MAX_AGE};
+use crate::common::PROTOCOL_NAME;
 use crate::session_handle::{
     HttpSessionHandle, SessionHandleCommand, WarpgateServerHandleFromRequest,
 };
@@ -169,11 +169,11 @@ impl SessionStore {
         }
     }
 
-    pub async fn vacuum(&mut self) {
+    pub async fn vacuum(&mut self, session_max_age: Duration) {
         let now = Instant::now();
         let mut to_remove = vec![];
         for (id, timestamp) in self.session_timestamps.iter() {
-            if now.duration_since(*timestamp) > SESSION_MAX_AGE {
+            if now.duration_since(*timestamp) > session_max_age {
                 to_remove.push(*id);
             }
         }