From 66e04c9fe75de81964e304951335561838ea2691 Mon Sep 17 00:00:00 2001 From: 21pages Date: Mon, 18 Jul 2022 14:03:52 +0800 Subject: [PATCH] keep session when network interrupt Signed-off-by: 21pages --- libs/hbb_common/protos/message.proto | 1 + src/client.rs | 4 ++ src/server/connection.rs | 94 +++++++++++++++++++++++----- 3 files changed, 83 insertions(+), 16 deletions(-) diff --git a/libs/hbb_common/protos/message.proto b/libs/hbb_common/protos/message.proto index 0538f1aef..2044388f8 100644 --- a/libs/hbb_common/protos/message.proto +++ b/libs/hbb_common/protos/message.proto @@ -63,6 +63,7 @@ message LoginRequest { PortForward port_forward = 8; } bool video_ack_required = 9; + uint64 session_id = 10; } message ChatMessage { string text = 1; } diff --git a/src/client.rs b/src/client.rs index 2350d243a..28101896e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -28,6 +28,7 @@ use hbb_common::{ log, message_proto::{option_message::BoolOption, *}, protobuf::Message as _, + rand, rendezvous_proto::*, socket_client, sodiumoxide::crypto::{box_, secretbox, sign}, @@ -782,6 +783,7 @@ pub struct LoginConfigHandler { pub version: i64, pub conn_id: i32, features: Option, + session_id: u64, } impl Deref for LoginConfigHandler { @@ -805,6 +807,7 @@ impl LoginConfigHandler { let config = self.load_config(); self.remember = !config.password.is_empty(); self.config = config; + self.session_id = rand::random(); } pub fn should_auto_login(&self) -> String { @@ -1140,6 +1143,7 @@ impl LoginConfigHandler { my_id, my_name: crate::username(), option: self.get_option_message(true).into(), + session_id: self.session_id, ..Default::default() }; if self.is_file_transfer { diff --git a/src/server/connection.rs b/src/server/connection.rs index 9c7025efc..04b070cfe 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -36,6 +36,7 @@ pub type Sender = mpsc::UnboundedSender<(Instant, Arc)>; lazy_static::lazy_static! { static ref LOGIN_FAILURES: Arc::>> = Default::default(); + static ref SESSIONS: Arc::>> = Default::default(); } pub static CLICK_TIME: AtomicI64 = AtomicI64::new(0); pub static MOUSE_MOVE_TIME: AtomicI64 = AtomicI64::new(0); @@ -54,6 +55,14 @@ enum MessageInput { BlockOff, } +#[derive(Clone, Debug)] +struct Session { + name: String, + session_id: u64, + last_recv_time: Arc>, + random_password: String, +} + pub struct Connection { inner: ConnInner, stream: super::Stream, @@ -81,6 +90,8 @@ pub struct Connection { video_ack_required: bool, peer_info: (String, String), api_server: String, + lr: LoginRequest, + last_recv_time: Arc>, } impl Subscriber for ConnInner { @@ -112,6 +123,7 @@ const H1: Duration = Duration::from_secs(3600); const MILLI1: Duration = Duration::from_millis(1); const SEND_TIMEOUT_VIDEO: u64 = 12_000; const SEND_TIMEOUT_OTHER: u64 = SEND_TIMEOUT_VIDEO * 10; +const SESSION_TIMEOUT: Duration = Duration::from_secs(30); impl Connection { pub async fn start( @@ -165,6 +177,8 @@ impl Connection { video_ack_required: false, peer_info: Default::default(), api_server: "".to_owned(), + lr: Default::default(), + last_recv_time: Arc::new(Mutex::new(Instant::now())), }; #[cfg(not(any(target_os = "android", target_os = "ios")))] tokio::spawn(async move { @@ -224,6 +238,7 @@ impl Connection { msg_out.set_misc(misc); conn.send(msg_out).await; conn.on_close("Close requested from connection manager", false); + SESSIONS.lock().unwrap().remove(&conn.lr.my_id); break; } ipc::Data::ChatMessage{text} => { @@ -317,6 +332,7 @@ impl Connection { }, Ok(bytes) => { last_recv_time = Instant::now(); + *conn.last_recv_time.lock().unwrap() = Instant::now(); if let Ok(msg_in) = Message::parse_from_bytes(&bytes) { if !conn.on_message(msg_in).await { break; @@ -780,38 +796,77 @@ impl Connection { self.tx_input.send(MessageInput::Key((msg, press))).ok(); } - fn validate_password(&mut self, lr_password: Vec) -> bool { - let validate = |password: String| { - if password.len() == 0 { - return false; - } - let mut hasher = Sha256::new(); - hasher.update(password); - hasher.update(&self.hash.salt); - let mut hasher2 = Sha256::new(); - hasher2.update(&hasher.finalize()[..]); - hasher2.update(&self.hash.challenge); - hasher2.finalize()[..] == lr_password[..] - }; + fn validate_one_password(&self, password: String) -> bool { + if password.len() == 0 { + return false; + } + let mut hasher = Sha256::new(); + hasher.update(password); + hasher.update(&self.hash.salt); + let mut hasher2 = Sha256::new(); + hasher2.update(&hasher.finalize()[..]); + hasher2.update(&self.hash.challenge); + hasher2.finalize()[..] == self.lr.password[..] + } + + fn validate_password(&mut self) -> bool { if password::security_enabled() { - if validate(Config::get_security_password()) { + if self.validate_one_password(Config::get_security_password()) { return true; } } if password::random_password_valid() { - if validate(password::random_password()) { + let password = password::random_password(); + if self.validate_one_password(password.clone()) { if password::onetime_password_activated() { password::set_onetime_password_activated(false); } + SESSIONS.lock().unwrap().insert( + self.lr.my_id.clone(), + Session { + name: self.lr.my_name.clone(), + session_id: self.lr.session_id, + last_recv_time: self.last_recv_time.clone(), + random_password: password, + }, + ); return true; } } + false + } + fn is_of_recent_session(&mut self) -> bool { + let session = SESSIONS + .lock() + .unwrap() + .get(&self.lr.my_id) + .map(|s| s.to_owned()); + if let Some(session) = session { + if session.name == self.lr.my_name + && session.session_id == self.lr.session_id + && !self.lr.password.is_empty() + && self.validate_one_password(session.random_password.clone()) + && session.last_recv_time.lock().unwrap().elapsed() < SESSION_TIMEOUT + { + SESSIONS.lock().unwrap().insert( + self.lr.my_id.clone(), + Session { + name: self.lr.my_name.clone(), + session_id: self.lr.session_id, + last_recv_time: self.last_recv_time.clone(), + random_password: session.random_password, + }, + ); + return true; + } + } false } async fn on_message(&mut self, msg: Message) -> bool { if let Some(message::Union::LoginRequest(lr)) = msg.union { + self.lr = lr.clone(); if let Some(o) = lr.option.as_ref() { self.update_option(o).await; if let Some(q) = o.video_codec_state.clone().take() { @@ -882,6 +937,12 @@ impl Connection { } if !crate::is_ip(&lr.username) && lr.username != Config::get_id() { self.send_login_error("Offline").await; + } else if self.is_of_recent_session() { + self.try_start_cm(lr.my_id, lr.my_name, true); + self.send_logon_response().await; + if self.port_forward_socket.is_some() { + return false; + } } else if lr.password.is_empty() { self.try_start_cm(lr.my_id, lr.my_name, false); } else { @@ -901,7 +962,7 @@ impl Connection { .await; } else if time == failure.0 && failure.1 > 6 { self.send_login_error("Please try 1 minute later").await; - } else if !self.validate_password(lr.password.clone()) { + } else if !self.validate_password() { if failure.0 == time { failure.1 += 1; failure.2 += 1; @@ -1135,6 +1196,7 @@ impl Connection { } Some(misc::Union::CloseReason(_)) => { self.on_close("Peer close", true); + SESSIONS.lock().unwrap().remove(&self.lr.my_id); return false; } _ => {}