warp-tech · Eugeny · Aug 7, 2023 · Aug 7, 2023
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/warpgate-common/src/auth/state.rs b/warpgate-common/src/auth/state.rs
@@ -1,8 +1,11 @@
 use std::collections::HashSet;
 
+use chrono::{DateTime, Utc};
+use rand::Rng;
 use uuid::Uuid;
 
 use super::{AuthCredential, CredentialKind, CredentialPolicy, CredentialPolicyResponse};
+use crate::SessionId;
 
 #[derive(Debug, Clone)]
 pub enum AuthResult {
@@ -13,34 +16,54 @@ pub enum AuthResult {
 
 pub struct AuthState {
     id: Uuid,
+    session_id: Option<Uuid>,
     username: String,
     protocol: String,
     force_rejected: bool,
     policy: Box<dyn CredentialPolicy + Sync + Send>,
     valid_credentials: Vec<AuthCredential>,
+    started: DateTime<Utc>,
+    identification_string: String,
+}
+
+fn generate_identification_string() -> String {
+    let mut s = String::new();
+    let mut rng = rand::thread_rng();
+    for _ in 0..4 {
+        s.push_str(&format!("{:X}", rng.gen_range(0..16)));
+    }
+    s
 }
 
 impl AuthState {
     pub fn new(
         id: Uuid,
+        session_id: Option<SessionId>,
         username: String,
         protocol: String,
         policy: Box<dyn CredentialPolicy + Sync + Send>,
     ) -> Self {
         Self {
             id,
+            session_id,
             username,
             protocol,
             force_rejected: false,
             policy,
             valid_credentials: vec![],
+            started: Utc::now(),
+            identification_string: generate_identification_string(),
         }
     }
 
     pub fn id(&self) -> &Uuid {
         &self.id
     }
 
+    pub fn session_id(&self) -> &Option<SessionId> {
+        &self.session_id
+    }
+
     pub fn username(&self) -> &str {
         &self.username
     }
@@ -49,6 +72,14 @@ impl AuthState {
         &self.protocol
     }
 
+    pub fn started(&self) -> &DateTime<Utc> {
+        &self.started
+    }
+
+    pub fn identification_string(&self) -> &str {
+        &self.identification_string
+    }
+
     pub fn add_valid_credential(&mut self, credential: AuthCredential) {
         self.valid_credentials.push(credential);
     }

diff --git a/warpgate-core/src/auth_state_store.rs b/warpgate-core/src/auth_state_store.rs
@@ -6,7 +6,7 @@ use once_cell::sync::Lazy;
 use tokio::sync::{broadcast, Mutex};
 use uuid::Uuid;
 use warpgate_common::auth::{AuthResult, AuthState};
-use warpgate_common::WarpgateError;
+use warpgate_common::{WarpgateError, SessionId};
 
 use crate::ConfigProvider;
 
@@ -49,6 +49,7 @@ impl AuthStateStore {
 
     pub async fn create(
         &mut self,
+        session_id: Option<&SessionId>,
         username: &str,
         protocol: &str,
     ) -> Result<(Uuid, Arc<Mutex<AuthState>>), WarpgateError> {
@@ -63,7 +64,13 @@ impl AuthStateStore {
             return Err(WarpgateError::UserNotFound)
         };
 
-        let state = AuthState::new(id, username.to_string(), protocol.to_string(), policy);
+        let state = AuthState::new(
+            id,
+            session_id.copied(),
+            username.to_string(),
+            protocol.to_string(),
+            policy,
+        );
         self.store
             .insert(id, (Arc::new(Mutex::new(state)), Instant::now()));
 

diff --git a/warpgate-protocol-http/src/api/auth.rs b/warpgate-protocol-http/src/api/auth.rs
@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use chrono::{DateTime, Utc};
 use poem::session::Session;
 use poem::web::Data;
 use poem::Request;
@@ -66,7 +67,10 @@ enum LogoutResponse {
 #[derive(Object)]
 struct AuthStateResponseInternal {
     pub protocol: String,
+    pub address: Option<String>,
+    pub started: DateTime<Utc>,
     pub state: ApiAuthState,
+    pub identification_string: String,
 }
 
 #[derive(ApiResponse)]
@@ -214,7 +218,7 @@ impl Api {
         let Some(state_arc) = store.get(&state_id.0) else {
             return Ok(AuthStateResponse::NotFound);
         };
-        serialize_auth_state_inner(state_arc).await
+        serialize_auth_state_inner(state_arc, *services).await
     }
 
     #[oai(
@@ -237,7 +241,7 @@ impl Api {
         state_arc.lock().await.reject();
         store.complete(&state_id.0).await;
         session.clear_auth_state();
-        serialize_auth_state_inner(state_arc).await
+        serialize_auth_state_inner(state_arc, *services).await
     }
 
     #[oai(
@@ -256,7 +260,7 @@ impl Api {
         let Some(state_arc) = state_arc else {
             return Ok(AuthStateResponse::NotFound);
         };
-        serialize_auth_state_inner(state_arc).await
+        serialize_auth_state_inner(state_arc, *services).await
     }
 
     #[oai(
@@ -284,7 +288,7 @@ impl Api {
         if let AuthResult::Accepted { .. } = auth_result {
             services.auth_state_store.lock().await.complete(&id).await;
         }
-        serialize_auth_state_inner(state_arc).await
+        serialize_auth_state_inner(state_arc, *services).await
     }
 
     #[oai(
@@ -304,7 +308,7 @@ impl Api {
         };
         state_arc.lock().await.reject();
         services.auth_state_store.lock().await.complete(&id).await;
-        serialize_auth_state_inner(state_arc).await
+        serialize_auth_state_inner(state_arc, *services).await
     }
 }
 
@@ -339,10 +343,25 @@ async fn get_auth_state(
 
 async fn serialize_auth_state_inner(
     state_arc: Arc<Mutex<AuthState>>,
+    services: &Services,
 ) -> poem::Result<AuthStateResponse> {
     let state = state_arc.lock().await;
+
+    let session_state_store = services.state.lock().await;
+    let session_state = state
+        .session_id()
+        .and_then(|session_id| session_state_store.sessions.get(&session_id));
+
+    let peer_addr = match session_state {
+        Some(x) => x.lock().await.remote_address,
+        None => None,
+    };
+
     Ok(AuthStateResponse::Ok(Json(AuthStateResponseInternal {
         protocol: state.protocol().to_string(),
+        address: peer_addr.map(|x| x.ip().to_string()),
+        started: state.started().clone(),
         state: state.verify().into(),
+        identification_string: state.identification_string().to_owned(),
     })))
 }
diff --git a/warpgate-protocol-http/src/common.rs b/warpgate-protocol-http/src/common.rs
@@ -198,7 +198,9 @@ pub async fn get_auth_state_for_request(
     match session.get_auth_state_id() {
         Some(id) => Ok(store.get(&id.0).ok_or(WarpgateError::InconsistentState)?),
         None => {
-            let (id, state) = store.create(username, crate::common::PROTOCOL_NAME).await?;
+            let (id, state) = store
+                .create(None, username, crate::common::PROTOCOL_NAME)
+                .await?;
             session.set(AUTH_STATE_ID_SESSION_KEY, AuthStateId(id));
             Ok(state)
         }

diff --git a/warpgate-protocol-mysql/src/session.rs b/warpgate-protocol-mysql/src/session.rs
@@ -183,7 +183,11 @@ impl MySqlSession {
                     .auth_state_store
                     .lock()
                     .await
-                    .create(&username, crate::common::PROTOCOL_NAME)
+                    .create(
+                        Some(&self.server_handle.lock().await.id()),
+                        &username,
+                        crate::common::PROTOCOL_NAME,
+                    )
                     .await?
                     .1;
                 let mut state = state_arc.lock().await;

diff --git a/warpgate-protocol-ssh/src/server/session.rs b/warpgate-protocol-ssh/src/server/session.rs
@@ -228,7 +228,7 @@ impl ServerSession {
                 .auth_state_store
                 .lock()
                 .await
-                .create(username, crate::PROTOCOL_NAME)
+                .create(Some(&self.id), username, crate::PROTOCOL_NAME)
                 .await?
                 .1;
             self.auth_state = Some(state);
@@ -1279,6 +1279,8 @@ impl ServerSession {
                     let Some(auth_state) = self.auth_state.as_ref() else {
                         return russh::server::Auth::Reject { proceed_with_methods: None};
                     };
+                    let identification_string =
+                        auth_state.lock().await.identification_string().to_owned();
                     let auth_state_id = *auth_state.lock().await.id();
                     let event = self
                         .services
@@ -1311,11 +1313,19 @@ impl ServerSession {
                     russh::server::Auth::Partial {
                         name: Cow::Owned(format!(
                             concat!(
-                            "----------------------------------------------------------------\n",
-                            "Warpgate authentication: please open {} in your browser\n",
-                            "----------------------------------------------------------------\n"
+                            "-----------------------------------------------------------------------\n",
+                            "Warpgate authentication: please open the following URL in your browser:\n",
+                            "{}\n\n",
+                            "Make sure you're seeing this security key: {}\n",
+                            "-----------------------------------------------------------------------\n"
                         ),
-                            login_url
+                            login_url,
+                            identification_string
+                                .chars()
+                                .into_iter()
+                                .map(|x| x.to_string())
+                                .collect::<Vec<_>>()
+                                .join(" ")
                         )),
                         instructions: Cow::Borrowed(""),
                         prompts: Cow::Owned(vec![(Cow::Borrowed("Press Enter when done: "), true)]),

diff --git a/warpgate-web/src/admin/lib/openapi-schema.json b/warpgate-web/src/admin/lib/openapi-schema.json
@@ -2,7 +2,7 @@
   "openapi": "3.0.0",
   "info": {
     "title": "Warpgate Web Admin",
-    "version": "0.7.1"
+    "version": "0.7.4"
   },
   "servers": [
     {