http: set cookies for all subdomains

warpgate-protocol-http/src/api/auth.rs

@@ -1,5 +1,5 @@
 use crate::common::SessionExt;
-use crate::session::SessionMiddleware;
+use crate::session::SessionStore;
 use anyhow::Context;
 use poem::session::Session;
 use poem::web::Data;
@@ -54,7 +54,7 @@ impl Api {
         req: &Request,
         session: &Session,
         services: Data<&Services>,
-        session_middleware: Data<&Arc<Mutex<SessionMiddleware>>>,
+        session_middleware: Data<&Arc<Mutex<SessionStore>>>,
         body: Json<LoginRequest>,
     ) -> poem::Result<LoginResponse> {
         let mut credentials = vec![AuthCredential::Password(Secret::new(body.password.clone()))];
@@ -103,7 +103,7 @@ impl Api {
     async fn api_auth_logout(
         &self,
         session: &Session,
-        session_middleware: Data<&Arc<Mutex<SessionMiddleware>>>,
+        session_middleware: Data<&Arc<Mutex<SessionStore>>>,
     ) -> poem::Result<LogoutResponse> {
         session_middleware.lock().await.remove_session(session);
         session.clear();

warpgate-protocol-http/src/lib.rs

@@ -21,19 +21,22 @@ use logging::{log_request_result, span_for_request};
 use poem::endpoint::{EmbeddedFileEndpoint, EmbeddedFilesEndpoint};
 use poem::listener::{Listener, RustlsConfig, TcpListener};
 use poem::middleware::SetHeader;
-use poem::session::{CookieConfig, MemoryStorage, ServerSession};
+use poem::session::MemoryStorage;
 use poem::web::Data;
 use poem::{Endpoint, EndpointExt, FromRequest, IntoEndpoint, IntoResponse, Route, Server};
 use poem_openapi::OpenApiService;
 use tokio::sync::Mutex;
 use tracing::*;
 use warpgate_admin::admin_api_app;
-use warpgate_common::{ProtocolServer, Services, Target, TargetTestError, TlsCertificateAndPrivateKey, TlsCertificateBundle, TlsPrivateKey};
+use warpgate_common::{
+    ProtocolServer, Services, Target, TargetTestError, TlsCertificateAndPrivateKey,
+    TlsCertificateBundle, TlsPrivateKey,
+};
 use warpgate_web::Assets;
 
-use crate::common::{endpoint_admin_auth, endpoint_auth, page_auth, COOKIE_MAX_AGE};
+use crate::common::{endpoint_admin_auth, endpoint_auth, page_auth};
 use crate::error::error_page;
-use crate::session::{SessionMiddleware, SharedSessionStorage};
+use crate::session::{SessionMiddleware, SessionStore, SharedSessionStorage};
 
 pub struct HTTPProtocolServer {
     services: Services,
@@ -62,7 +65,7 @@ impl ProtocolServer for HTTPProtocolServer {
 
         let session_storage =
             SharedSessionStorage(Arc::new(Mutex::new(Box::new(MemoryStorage::default()))));
-        let session_middleware = SessionMiddleware::new();
+        let session_store = SessionStore::new();
 
         let app = Route::new()
             .nest(
@@ -104,7 +107,7 @@ impl ProtocolServer for HTTPProtocolServer {
                 }),
             )
             .around(move |ep, req| async move {
-                let sm = Data::<&Arc<Mutex<SessionMiddleware>>>::from_request_without_body(&req)
+                let sm = Data::<&Arc<Mutex<SessionStore>>>::from_request_without_body(&req)
                     .await?
                     .clone();
 
@@ -118,20 +121,14 @@ impl ProtocolServer for HTTPProtocolServer {
                 SetHeader::new()
                     .overriding(http::header::STRICT_TRANSPORT_SECURITY, "max-age=31536000"),
             )
-            .with(ServerSession::new(
-                CookieConfig::default()
-                    .secure(false)
-                    .max_age(COOKIE_MAX_AGE)
-                    .name("warpgate-http-session"),
-                session_storage.clone(),
-            ))
+            .with(SessionMiddleware::new(session_storage.clone()))
             .data(self.services.clone())
-            .data(session_middleware.clone())
+            .data(session_store.clone())
             .data(session_storage);
 
         tokio::spawn(async move {
             loop {
-                session_middleware.lock().await.vacuum().await;
+                session_store.lock().await.vacuum().await;
                 tokio::time::sleep(Duration::from_secs(60)).await;
             }
         });
@@ -144,9 +141,11 @@ impl ProtocolServer for HTTPProtocolServer {
             let key_path = config.paths_relative_to.join(&config.store.http.key);
 
             TlsCertificateAndPrivateKey {
-                certificate: TlsCertificateBundle::from_file(&certificate_path).await.with_context(|| {
-                    format!("reading TLS private key from '{}'", key_path.display())
-                })?,
+                certificate: TlsCertificateBundle::from_file(&certificate_path)
+                    .await
+                    .with_context(|| {
+                        format!("reading TLS private key from '{}'", key_path.display())
+                    })?,
                 private_key: TlsPrivateKey::from_file(&key_path).await.with_context(|| {
                     format!(
                         "reading TLS certificate from '{}'",
@@ -157,9 +156,10 @@ impl ProtocolServer for HTTPProtocolServer {
         };
 
         info!(?address, "Listening");
-        Server::new(TcpListener::bind(address).rustls(
-            RustlsConfig::new().fallback(certificate_and_key.into()),
-        ))
+        Server::new(
+            TcpListener::bind(address)
+                .rustls(RustlsConfig::new().fallback(certificate_and_key.into())),
+        )
         .run(app)
         .await?;
 

warpgate-protocol-http/src/session.rs

@@ -2,15 +2,22 @@ use std::collections::{BTreeMap, HashMap};
 use std::sync::{Arc, Weak};
 use std::time::{Duration, Instant};
 
-use poem::session::{Session, SessionStorage};
+use async_trait::async_trait;
+use http::header::Entry;
+use poem::middleware::CookieJarManagerEndpoint;
+use poem::session::{
+    CookieConfig, ServerSession as PoemSessionMiddleware, ServerSessionEndpoint, Session,
+    SessionStorage,
+};
+use poem::web::cookie::Cookie;
 use poem::web::{Data, RemoteAddr};
-use poem::{FromRequest, Request};
+use poem::{Endpoint, FromRequest, IntoResponse, Middleware, Request, Response};
 use serde_json::Value;
 use tokio::sync::Mutex;
 use tracing::*;
 use warpgate_common::{Services, SessionId, SessionStateInit, WarpgateServerHandle};
 
-use crate::common::{PROTOCOL_NAME, SESSION_MAX_AGE};
+use crate::common::{COOKIE_MAX_AGE, PROTOCOL_NAME, SESSION_MAX_AGE};
 use crate::session_handle::{
     HttpSessionHandle, SessionHandleCommand, WarpgateServerHandleFromRequest,
 };
@@ -20,7 +27,7 @@ pub struct SharedSessionStorage(pub Arc<Mutex<Box<dyn SessionStorage>>>);
 
 static POEM_SESSION_ID_SESSION_KEY: &str = "poem_session_id";
 
-#[async_trait::async_trait]
+#[async_trait]
 impl SessionStorage for SharedSessionStorage {
     async fn load_session(
         &self,
@@ -55,16 +62,16 @@ impl SessionStorage for SharedSessionStorage {
     }
 }
 
-pub struct SessionMiddleware {
+pub struct SessionStore {
     session_handles: HashMap<SessionId, Arc<Mutex<WarpgateServerHandle>>>,
     session_timestamps: HashMap<SessionId, Instant>,
-    this: Weak<Mutex<SessionMiddleware>>,
+    this: Weak<Mutex<SessionStore>>,
 }
 
 static SESSION_ID_SESSION_KEY: &str = "session_id";
 static REQUEST_COUNTER_SESSION_KEY: &str = "request_counter";
 
-impl SessionMiddleware {
+impl SessionStore {
     pub fn new() -> Arc<Mutex<Self>> {
         Arc::new_cyclic(|me| {
             Mutex::new(Self {
@@ -183,3 +190,60 @@ impl SessionMiddleware {
         }
     }
 }
+
+pub struct SessionMiddleware {
+    inner: PoemSessionMiddleware<SharedSessionStorage>,
+}
+
+impl SessionMiddleware {
+    pub fn new(session_storage: SharedSessionStorage) -> Self {
+        Self {
+            inner: PoemSessionMiddleware::new(
+                CookieConfig::default()
+                    .secure(false)
+                    .max_age(COOKIE_MAX_AGE)
+                    .name("warpgate-http-session"),
+                session_storage,
+            ),
+        }
+    }
+}
+
+pub struct SessionMiddlewareEndpoint<E: Endpoint> {
+    inner: E,
+}
+
+impl<E: Endpoint> Middleware<E> for SessionMiddleware {
+    type Output = SessionMiddlewareEndpoint<
+        CookieJarManagerEndpoint<ServerSessionEndpoint<SharedSessionStorage, E>>,
+    >;
+
+    fn transform(&self, ep: E) -> Self::Output {
+        SessionMiddlewareEndpoint {
+            inner: self.inner.transform(ep),
+        }
+    }
+}
+
+#[async_trait]
+impl<E: Endpoint> Endpoint for SessionMiddlewareEndpoint<E> {
+    type Output = Response;
+
+    async fn call(&self, req: Request) -> poem::Result<Self::Output> {
+        let host = req.original_uri().host().map(|x| x.to_string());
+        let mut resp = self.inner.call(req).await?.into_response();
+        if let Some(host) = host {
+            if let Entry::Occupied(mut entry) = resp.headers_mut().entry(http::header::SET_COOKIE) {
+                if let Ok(cookie_str) = entry.get().to_str() {
+                    if let Ok(mut cookie) = Cookie::parse(cookie_str) {
+                        cookie.set_domain(host);
+                        if let Ok(value) = cookie.to_string().parse() {
+                            entry.insert(value);
+                        }
+                    }
+                }
+            }
+        }
+        Ok(resp)
+    }
+}

warpgate-protocol-http/src/session_handle.rs

@@ -8,7 +8,7 @@ use poem::{FromRequest, Request, RequestBody};
 use tokio::sync::{mpsc, Mutex};
 use warpgate_common::{SessionHandle, WarpgateServerHandle};
 
-use crate::session::SessionMiddleware;
+use crate::session::SessionStore;
 
 #[derive(Clone, Debug, PartialEq)]
 pub enum SessionHandleCommand {
@@ -46,7 +46,7 @@ impl std::ops::Deref for WarpgateServerHandleFromRequest {
 #[async_trait::async_trait]
 impl<'a> FromRequest<'a> for WarpgateServerHandleFromRequest {
     async fn from_request(req: &'a Request, _: &mut RequestBody) -> poem::Result<Self> {
-        let sm = Data::<&Arc<Mutex<SessionMiddleware>>>::from_request_without_body(req).await?;
+        let sm = Data::<&Arc<Mutex<SessionStore>>>::from_request_without_body(req).await?;
         let session: &Session = <_>::from_request_without_body(&req).await?;
         Ok(sm
             .lock()

warpgate-protocol-ssh/src/server/session.rs

@@ -38,6 +38,7 @@ use crate::{
 };
 
 #[derive(Clone)]
+#[allow(clippy::large_enum_variant)]
 enum TargetSelection {
     None,
     NotFound(String),

warpgate-web/src/gateway/TargetList.svelte

@@ -84,7 +84,6 @@ init()
             targetName={selectedTarget?.name}
             username={$serverInfo?.username}
             targetKind={selectedTarget?.kind ?? TargetKind.Ssh}
-            targetURL={selectedTarget?.url}
         />
     </ModalBody>
 </Modal>