diff --git a/src/privacy_mode.rs b/src/privacy_mode.rs index c3ff04de7..1caa4ee46 100644 --- a/src/privacy_mode.rs +++ b/src/privacy_mode.rs @@ -61,6 +61,8 @@ pub trait PrivacyMode: Sync + Send { fn pre_conn_id(&self) -> i32; + fn get_impl_key(&self) -> &str; + #[inline] fn check_on_conn_id(&self, conn_id: i32) -> ResultType { let pre_conn_id = self.pre_conn_id(); @@ -117,24 +119,21 @@ lazy_static::lazy_static! { } }; - static ref CUR_PRIVACY_MODE_IMPL: Arc> = { + static ref PRIVACY_MODE: Arc>>> = { let mut cur_impl = get_option("privacy-mode-impl-key".to_owned()); if !get_supported_privacy_mode_impl().iter().any(|(k, _)| k == &cur_impl) { cur_impl = DEFAULT_PRIVACY_MODE_IMPL.to_owned(); } - Arc::new(Mutex::new(cur_impl)) - }; - static ref PRIVACY_MODE: Arc>>> = { - let cur_impl = (*CUR_PRIVACY_MODE_IMPL.lock().unwrap()).clone(); + let privacy_mode = match PRIVACY_MODE_CREATOR.lock().unwrap().get(&(&cur_impl as &str)) { - Some(creator) => Some(creator()), + Some(creator) => Some(creator(&cur_impl)), None => None, }; Arc::new(Mutex::new(privacy_mode)) }; } -pub type PrivacyModeCreator = fn() -> Box; +pub type PrivacyModeCreator = fn(impl_key: &str) -> Box; lazy_static::lazy_static! { static ref PRIVACY_MODE_CREATOR: Arc>> = { #[cfg(not(windows))] @@ -144,18 +143,18 @@ lazy_static::lazy_static! { #[cfg(windows)] { if win_exclude_from_capture::is_supported() { - map.insert(win_exclude_from_capture::PRIVACY_MODE_IMPL, || { - Box::new(win_exclude_from_capture::PrivacyModeImpl::default()) + map.insert(win_exclude_from_capture::PRIVACY_MODE_IMPL, |impl_key: &str| { + Box::new(win_exclude_from_capture::PrivacyModeImpl::new(impl_key)) }); } else { - map.insert(win_mag::PRIVACY_MODE_IMPL, || { - Box::new(win_mag::PrivacyModeImpl::default()) + map.insert(win_mag::PRIVACY_MODE_IMPL, |impl_key: &str| { + Box::new(win_mag::PrivacyModeImpl::new(impl_key)) }); } #[cfg(feature = "virtual_display_driver")] - map.insert(win_virtual_display::PRIVACY_MODE_IMPL, || { - Box::new(win_virtual_display::PrivacyModeImpl::default()) + map.insert(win_virtual_display::PRIVACY_MODE_IMPL, |impl_key: &str| { + Box::new(win_virtual_display::PrivacyModeImpl::new(impl_key)) }); } Arc::new(Mutex::new(map)) @@ -174,13 +173,15 @@ pub fn clear() -> Option<()> { #[inline] pub fn switch(impl_key: &str) { - let mut cur_impl_lock = CUR_PRIVACY_MODE_IMPL.lock().unwrap(); - if *cur_impl_lock == impl_key { - return; + let mut privacy_mode_lock = PRIVACY_MODE.lock().unwrap(); + if let Some(privacy_mode) = privacy_mode_lock.as_ref() { + if privacy_mode.get_impl_key() == impl_key { + return; + } } + if let Some(creator) = PRIVACY_MODE_CREATOR.lock().unwrap().get(impl_key) { - *PRIVACY_MODE.lock().unwrap() = Some(creator()); - *cur_impl_lock = impl_key.to_owned(); + *privacy_mode_lock = Some(creator(impl_key)); } } @@ -208,13 +209,15 @@ pub fn turn_on_privacy(impl_key: &str, conn_id: i32) -> Option> // Check or switch privacy mode implementation let impl_key = get_supported_impl(impl_key); - let mut cur_impl_lock = CUR_PRIVACY_MODE_IMPL.lock().unwrap(); + let mut cur_impl_key = "".to_string(); if let Some(privacy_mode) = privacy_mode_lock.as_ref() { + cur_impl_key = privacy_mode.get_impl_key().to_string(); let check_on_conn_id = privacy_mode.check_on_conn_id(conn_id); match check_on_conn_id.as_ref() { Ok(true) => { - if *cur_impl_lock == impl_key { + if cur_impl_key == impl_key { + // Same peer, same implementation. return Some(Ok(true)); } else { // Same peer, switch to new implementation. @@ -225,7 +228,7 @@ pub fn turn_on_privacy(impl_key: &str, conn_id: i32) -> Option> } } - if *cur_impl_lock != impl_key { + if cur_impl_key != impl_key { if let Some(creator) = PRIVACY_MODE_CREATOR .lock() .unwrap() @@ -235,8 +238,7 @@ pub fn turn_on_privacy(impl_key: &str, conn_id: i32) -> Option> privacy_mode.clear(); } - *privacy_mode_lock = Some(creator()); - *cur_impl_lock = impl_key.to_owned(); + *privacy_mode_lock = Some(creator(&impl_key)); } else { return Some(Err(anyhow!("Unsupported privacy mode: {}", impl_key))); } @@ -313,9 +315,23 @@ pub fn get_supported_privacy_mode_impl() -> Vec<(&'static str, &'static str)> { } } +#[inline] +pub fn get_cur_impl_key() -> Option { + PRIVACY_MODE + .lock() + .unwrap() + .as_ref() + .map(|pm| pm.get_impl_key().to_owned()) +} + #[inline] pub fn is_current_privacy_mode_impl(impl_key: &str) -> bool { - *CUR_PRIVACY_MODE_IMPL.lock().unwrap() == impl_key + PRIVACY_MODE + .lock() + .unwrap() + .as_ref() + .map(|pm| pm.get_impl_key() == impl_key) + .unwrap_or(false) } #[inline] @@ -347,3 +363,22 @@ pub fn check_privacy_mode_err( pub fn is_privacy_mode_supported() -> bool { !DEFAULT_PRIVACY_MODE_IMPL.is_empty() } + +#[inline] +pub fn get_privacy_mode_conn_id() -> Option { + PRIVACY_MODE + .lock() + .unwrap() + .as_ref() + .map(|pm| pm.pre_conn_id()) +} + +#[inline] +pub fn is_in_privacy_mode() -> bool { + PRIVACY_MODE + .lock() + .unwrap() + .as_ref() + .map(|pm| pm.pre_conn_id() != INVALID_PRIVACY_MODE_CONN_ID) + .unwrap_or(false) +} diff --git a/src/privacy_mode/win_topmost_window.rs b/src/privacy_mode/win_topmost_window.rs index 6b94d58f1..687f59155 100644 --- a/src/privacy_mode/win_topmost_window.rs +++ b/src/privacy_mode/win_topmost_window.rs @@ -65,24 +65,12 @@ impl WindowHandlers { } pub struct PrivacyModeImpl { + impl_key: String, conn_id: i32, handlers: WindowHandlers, hwnd: u64, } -impl Default for PrivacyModeImpl { - fn default() -> Self { - Self { - conn_id: INVALID_PRIVACY_MODE_CONN_ID, - handlers: WindowHandlers { - hthread: 0, - hprocess: 0, - }, - hwnd: 0, - } - } -} - impl PrivacyMode for PrivacyModeImpl { fn init(&self) -> ResultType<()> { Ok(()) @@ -163,9 +151,26 @@ impl PrivacyMode for PrivacyModeImpl { fn pre_conn_id(&self) -> i32 { self.conn_id } + + #[inline] + fn get_impl_key(&self) -> &str { + &self.impl_key + } } impl PrivacyModeImpl { + pub fn new(impl_key: &str) -> Self { + Self { + impl_key: impl_key.to_owned(), + conn_id: INVALID_PRIVACY_MODE_CONN_ID, + handlers: WindowHandlers { + hthread: 0, + hprocess: 0, + }, + hwnd: 0, + } + } + #[inline] pub fn get_hwnd(&self) -> u64 { self.hwnd @@ -372,4 +377,3 @@ pub(super) fn wait_find_privacy_hwnd(msecs: u128) -> ResultType { std::thread::sleep(Duration::from_millis(100)); } } - diff --git a/src/privacy_mode/win_virtual_display.rs b/src/privacy_mode/win_virtual_display.rs index 1d3ffa30e..8183d3275 100644 --- a/src/privacy_mode/win_virtual_display.rs +++ b/src/privacy_mode/win_virtual_display.rs @@ -35,23 +35,13 @@ struct Display { } pub struct PrivacyModeImpl { + impl_key: String, conn_id: i32, displays: Vec, virtual_displays: Vec, virtual_displays_added: Vec, } -impl Default for PrivacyModeImpl { - fn default() -> Self { - Self { - conn_id: INVALID_PRIVACY_MODE_CONN_ID, - displays: Vec::new(), - virtual_displays: Vec::new(), - virtual_displays_added: Vec::new(), - } - } -} - struct TurnOnGuard<'a> { privacy_mode: &'a mut PrivacyModeImpl, succeeded: bool, @@ -82,6 +72,16 @@ impl<'a> Drop for TurnOnGuard<'a> { } impl PrivacyModeImpl { + pub fn new(impl_key: &str) -> Self { + Self { + impl_key: impl_key.to_owned(), + conn_id: INVALID_PRIVACY_MODE_CONN_ID, + displays: Vec::new(), + virtual_displays: Vec::new(), + virtual_displays_added: Vec::new(), + } + } + // mainly from https://github.com/fufesou/rustdesk/blob/44c3a52ca8502cf53b58b59db130611778d34dbe/libs/scrap/src/dxgi/mod.rs#L365 fn set_displays(&mut self) { self.displays.clear(); @@ -431,6 +431,11 @@ impl PrivacyMode for PrivacyModeImpl { fn pre_conn_id(&self) -> i32 { self.conn_id } + + #[inline] + fn get_impl_key(&self) -> &str { + &self.impl_key + } } impl Drop for PrivacyModeImpl { diff --git a/src/server/connection.rs b/src/server/connection.rs index 419592821..12e52afba 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -508,21 +508,18 @@ impl Connection { ipc::Data::PrivacyModeState((_, state, impl_key)) => { let msg_out = match state { privacy_mode::PrivacyModeState::OffSucceeded => { - video_service::set_privacy_mode_conn_id(0); crate::common::make_privacy_mode_msg( back_notification::PrivacyModeState::PrvOffSucceeded, impl_key, ) } privacy_mode::PrivacyModeState::OffByPeer => { - video_service::set_privacy_mode_conn_id(0); crate::common::make_privacy_mode_msg( back_notification::PrivacyModeState::PrvOffByPeer, impl_key, ) } privacy_mode::PrivacyModeState::OffUnknown => { - video_service::set_privacy_mode_conn_id(0); crate::common::make_privacy_mode_msg( back_notification::PrivacyModeState::PrvOffUnknown, impl_key, @@ -682,10 +679,10 @@ impl Connection { } } - let video_privacy_conn_id = video_service::get_privacy_mode_conn_id(); - if video_privacy_conn_id == id { - video_service::set_privacy_mode_conn_id(0); - let _ = Self::turn_off_privacy_to_msg(id); + if let Some(video_privacy_conn_id) = privacy_mode::get_privacy_mode_conn_id() { + if video_privacy_conn_id == id { + let _ = Self::turn_off_privacy_to_msg(id); + } } #[cfg(all(feature = "flutter", feature = "plugin_framework"))] #[cfg(not(any(target_os = "android", target_os = "ios")))] @@ -880,7 +877,7 @@ impl Connection { } async fn check_privacy_mode_on(&mut self) -> bool { - if video_service::get_privacy_mode_conn_id() > 0 { + if privacy_mode::is_in_privacy_mode() { self.send_login_error("Someone turns on privacy mode, exit") .await; false @@ -2610,7 +2607,23 @@ impl Connection { impl_key, ) } else { - match privacy_mode::turn_on_privacy(&impl_key, self.inner.id) { + let is_pre_privacy_on = privacy_mode::is_in_privacy_mode(); + let pre_impl_key = privacy_mode::get_cur_impl_key(); + let turn_on_res = privacy_mode::turn_on_privacy(&impl_key, self.inner.id); + + if is_pre_privacy_on { + if let Some(pre_impl_key) = pre_impl_key { + if !privacy_mode::is_current_privacy_mode_impl(&pre_impl_key) { + let off_msg = crate::common::make_privacy_mode_msg( + back_notification::PrivacyModeState::PrvOffSucceeded, + pre_impl_key, + ); + self.send(off_msg).await; + } + } + } + + match turn_on_res { Some(Ok(res)) => { if res { let err_msg = privacy_mode::check_privacy_mode_err( @@ -2619,7 +2632,6 @@ impl Connection { 5_000, ); if err_msg.is_empty() { - video_service::set_privacy_mode_conn_id(self.inner.id); crate::common::make_privacy_mode_msg( back_notification::PrivacyModeState::PrvOnSucceeded, impl_key, @@ -2629,7 +2641,6 @@ impl Connection { "Check privacy mode failed: {}, turn off privacy mode.", &err_msg ); - video_service::set_privacy_mode_conn_id(0); let _ = Self::turn_off_privacy_to_msg(self.inner.id); crate::common::make_privacy_mode_msg_with_details( back_notification::PrivacyModeState::PrvOnFailed, @@ -2646,8 +2657,10 @@ impl Connection { } Some(Err(e)) => { log::error!("Failed to turn on privacy mode. {}", e); - if video_service::get_privacy_mode_conn_id() == 0 { - let _ = Self::turn_off_privacy_to_msg(0); + if !privacy_mode::is_in_privacy_mode() { + let _ = Self::turn_off_privacy_to_msg( + privacy_mode::INVALID_PRIVACY_MODE_CONN_ID, + ); } crate::common::make_privacy_mode_msg_with_details( back_notification::PrivacyModeState::PrvOnFailed, @@ -2674,7 +2687,6 @@ impl Connection { impl_key, ) } else { - video_service::set_privacy_mode_conn_id(0); Self::turn_off_privacy_to_msg(self.inner.id) }; self.send(msg_out).await; diff --git a/src/server/video_service.rs b/src/server/video_service.rs index 509db74b0..c36c6b90d 100644 --- a/src/server/video_service.rs +++ b/src/server/video_service.rs @@ -28,6 +28,7 @@ use super::{ use crate::common::SimpleCallOnReturn; #[cfg(target_os = "linux")] use crate::platform::linux::is_x11; +use crate::privacy_mode::{get_privacy_mode_conn_id, INVALID_PRIVACY_MODE_CONN_ID}; #[cfg(windows)] use crate::{ platform::windows::is_process_consent_running, @@ -68,7 +69,6 @@ lazy_static::lazy_static! { let (tx, rx) = unbounded_channel(); (tx, Arc::new(TokioMutex::new(rx))) }; - static ref PRIVACY_MODE_CONN_ID: Mutex = Mutex::new(0); pub static ref VIDEO_QOS: Arc> = Default::default(); pub static ref IS_UAC_RUNNING: Arc> = Default::default(); pub static ref IS_FOREGROUND_WINDOW_ELEVATED: Arc> = Default::default(); @@ -79,16 +79,6 @@ pub fn notify_video_frame_fetched(conn_id: i32, frame_tm: Option) { FRAME_FETCHED_NOTIFIER.0.send((conn_id, frame_tm)).ok(); } -#[inline] -pub fn set_privacy_mode_conn_id(conn_id: i32) { - *PRIVACY_MODE_CONN_ID.lock().unwrap() = conn_id -} - -#[inline] -pub fn get_privacy_mode_conn_id() -> i32 { - *PRIVACY_MODE_CONN_ID.lock().unwrap() -} - struct VideoFrameController { cur: Instant, send_conn_ids: HashSet, @@ -251,7 +241,9 @@ pub fn test_create_capturer( // Note: This function is extremely expensive, do not call it frequently. #[cfg(windows)] fn check_uac_switch(privacy_mode_id: i32, capturer_privacy_mode_id: i32) -> ResultType<()> { - if capturer_privacy_mode_id != 0 && is_current_privacy_mode_impl(PRIVACY_MODE_IMPL_WIN_MAG) { + if capturer_privacy_mode_id != INVALID_PRIVACY_MODE_CONN_ID + && is_current_privacy_mode_impl(PRIVACY_MODE_IMPL_WIN_MAG) + { if !is_installed() { if privacy_mode_id != capturer_privacy_mode_id { if !is_process_consent_running()? { @@ -323,18 +315,19 @@ fn get_capturer(current: usize, portable_service_running: bool) -> ResultType ResultType ResultType<()> { - let privacy_mode_id_2 = *PRIVACY_MODE_CONN_ID.lock().unwrap(); + let privacy_mode_id_2 = get_privacy_mode_conn_id().unwrap_or(INVALID_PRIVACY_MODE_CONN_ID); if privacy_mode_id != privacy_mode_id_2 { - if privacy_mode_id_2 != 0 { + if privacy_mode_id_2 != INVALID_PRIVACY_MODE_CONN_ID { let msg_out = crate::common::make_privacy_mode_msg( back_notification::PrivacyModeState::PrvOnByOther, "".to_owned(),