From 62cb9eb51e16d9e289c03b237c43c042e325a760 Mon Sep 17 00:00:00 2001 From: fufesou Date: Tue, 31 May 2022 15:54:21 +0800 Subject: [PATCH] privacy_mode_win_magnifier: more check on privacy mode Signed-off-by: fufesou --- libs/scrap/src/dxgi/mag.rs | 34 ++++++++----- src/server/connection.rs | 44 +++++++++++------ src/server/video_service.rs | 98 ++++++++++++++++++++----------------- src/ui/win_privacy.rs | 4 +- 4 files changed, 106 insertions(+), 74 deletions(-) diff --git a/libs/scrap/src/dxgi/mag.rs b/libs/scrap/src/dxgi/mag.rs index 3e74e960c..0d63088b7 100644 --- a/libs/scrap/src/dxgi/mag.rs +++ b/libs/scrap/src/dxgi/mag.rs @@ -277,6 +277,29 @@ impl CapturerMag { height: usize, use_yuv: bool, ) -> Result { + unsafe { + let x = GetSystemMetrics(SM_XVIRTUALSCREEN); + let y = GetSystemMetrics(SM_YVIRTUALSCREEN); + let w = GetSystemMetrics(SM_CXVIRTUALSCREEN); + let h = GetSystemMetrics(SM_CYVIRTUALSCREEN); + if !(origin.0 == x as _ && origin.1 == y as _ && width == w as _ && height == h as _) { + return Err(Error::new( + ErrorKind::Other, + format!( + "Failed Check screen rect ({}, {}, {} , {}) to ({}, {}, {}, {})", + origin.0, + origin.1, + origin.0 + width as i32, + origin.1 + height as i32, + x, + y, + x + w, + y + h + ), + )); + } + } + let mut s = Self { mag_interface: MagInterface::new()?, host_window: 0 as _, @@ -364,17 +387,6 @@ impl CapturerMag { )); } - let x = GetSystemMetrics(SM_XVIRTUALSCREEN); - let y = GetSystemMetrics(SM_YVIRTUALSCREEN); - let w = GetSystemMetrics(SM_CXVIRTUALSCREEN); - let h = GetSystemMetrics(SM_CYVIRTUALSCREEN); - s.rect = RECT { - left: x as _, - top: y as _, - right: (x + w) as _, - bottom: (y + h) as _, - }; - // Create the magnifier control. s.magnifier_window = CreateWindowExA( 0, diff --git a/src/server/connection.rs b/src/server/connection.rs index 746278fd0..3705d4e65 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -386,9 +386,12 @@ impl Connection { } } - if video_service::get_privacy_mode_conn_id() == id { + 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 _ = privacy_mode::turn_off_privacy(id).await; + let _ = privacy_mode::turn_off_privacy(id); + } else if video_privacy_conn_id == 0 { + let _ = privacy_mode::turn_off_privacy(0); } video_service::notify_video_frame_feched(id, None); video_service::update_test_latency(id, 0); @@ -657,7 +660,8 @@ impl Connection { features: Some(Features { privacy_mode: video_service::is_privacy_mode_supported(), ..Default::default() - }).into(), + }) + .into(), ..Default::default() }; let mut sub_service = false; @@ -1170,18 +1174,30 @@ impl Connection { } else { match privacy_mode::turn_on_privacy(self.inner.id) { Ok(true) => { - video_service::set_privacy_mode_conn_id(self.inner.id); - crate::common::make_privacy_mode_msg( - back_notification::PrivacyModeState::OnSucceeded, - ) - } - Ok(false) => { - crate::common::make_privacy_mode_msg( - back_notification::PrivacyModeState::OnFailedPlugin, - ) + if video_service::test_create_capturer(self.inner.id, 5_000) { + video_service::set_privacy_mode_conn_id(self.inner.id); + crate::common::make_privacy_mode_msg( + back_notification::PrivacyModeState::OnSucceeded, + ) + } else { + log::error!( + "Wait privacy mode timeout, turn off privacy mode" + ); + video_service::set_privacy_mode_conn_id(0); + let _ = privacy_mode::turn_off_privacy(self.inner.id); + crate::common::make_privacy_mode_msg( + back_notification::PrivacyModeState::OnFailed, + ) + } } + Ok(false) => crate::common::make_privacy_mode_msg( + back_notification::PrivacyModeState::OnFailedPlugin, + ), Err(e) => { log::error!("Failed to turn on privacy mode. {}", e); + if video_service::get_privacy_mode_conn_id() == 0 { + let _ = privacy_mode::turn_off_privacy(0); + } crate::common::make_privacy_mode_msg( back_notification::PrivacyModeState::OnFailed, ) @@ -1197,7 +1213,7 @@ impl Connection { ) } else { video_service::set_privacy_mode_conn_id(0); - privacy_mode::turn_off_privacy(self.inner.id).await + privacy_mode::turn_off_privacy(self.inner.id) }; self.send(msg_out).await; } @@ -1354,7 +1370,7 @@ fn try_activate_screen() { mod privacy_mode { use super::*; - pub(super) async fn turn_off_privacy(_conn_id: i32) -> Message { + pub(super) fn turn_off_privacy(_conn_id: i32) -> Message { #[cfg(windows)] { use crate::ui::win_privacy::*; diff --git a/src/server/video_service.rs b/src/server/video_service.rs index 1f344edab..b890e7b3d 100644 --- a/src/server/video_service.rs +++ b/src/server/video_service.rs @@ -19,12 +19,9 @@ // https://slhck.info/video/2017/03/01/rate-control.html use super::*; -use hbb_common::tokio::{ - runtime::Runtime, - sync::{ - mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}, - Mutex as TokioMutex, - }, +use hbb_common::tokio::sync::{ + mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender}, + Mutex as TokioMutex, }; use scrap::{Capturer, Config, Display, EncodeFrame, Encoder, Frame, VideoCodecId, STRIDE_ALIGN}; use std::{ @@ -80,7 +77,6 @@ pub fn is_privacy_mode_supported() -> bool { struct VideoFrameController { cur: Instant, send_conn_ids: HashSet, - rt: Runtime, } impl VideoFrameController { @@ -88,7 +84,6 @@ impl VideoFrameController { Self { cur: Instant::now(), send_conn_ids: HashSet::new(), - rt: Runtime::new().unwrap(), } } @@ -103,46 +98,29 @@ impl VideoFrameController { } } - fn blocking_wait_next(&mut self, timeout_millis: u128) { + #[tokio::main(flavor = "current_thread")] + async fn try_wait_next(&mut self, fetched_conn_ids: &mut HashSet, timeout_millis: u64) { if self.send_conn_ids.is_empty() { return; } - let send_conn_ids = self.send_conn_ids.clone(); - self.rt.block_on(async move { - let mut fetched_conn_ids = HashSet::new(); - let begin = Instant::now(); - while begin.elapsed().as_millis() < timeout_millis { - let timeout_dur = - Duration::from_millis((timeout_millis - begin.elapsed().as_millis()) as u64); - match tokio::time::timeout( - timeout_dur, - FRAME_FETCHED_NOTIFIER.1.lock().await.recv(), - ) - .await - { - Err(_) => { - // break if timeout - // log::error!("blocking wait frame receiving timeout {}", timeout_millis); - break; - } - Ok(Some((id, instant))) => { - if let Some(tm) = instant { - log::trace!("Channel recv latency: {}", tm.elapsed().as_secs_f32()); - } - fetched_conn_ids.insert(id); - - // break if all connections have received current frame - if fetched_conn_ids.len() >= send_conn_ids.len() { - break; - } - } - Ok(None) => { - // this branch would nerver be reached - } - } + let timeout_dur = Duration::from_millis(timeout_millis as u64); + match tokio::time::timeout(timeout_dur, FRAME_FETCHED_NOTIFIER.1.lock().await.recv()).await + { + Err(_) => { + // break if timeout + // log::error!("blocking wait frame receiving timeout {}", timeout_millis); } - }); + Ok(Some((id, instant))) => { + if let Some(tm) = instant { + log::trace!("Channel recv latency: {}", tm.elapsed().as_secs_f32()); + } + fetched_conn_ids.insert(id); + } + Ok(None) => { + // this branch would nerver be reached + } + } } } @@ -271,7 +249,7 @@ fn create_capturer(privacy_mode_id: i32, display: Display) -> ResultType { @@ -286,6 +264,7 @@ fn create_capturer(privacy_mode_id: i32, display: Display) -> ResultType { let c1 = Capturer::new(display, use_yuv).with_context(|| "Failed to create capturer")?; + log::debug!("Create capturer dxgi|gdi"); Box::new(c1) } }; @@ -309,6 +288,19 @@ fn ensure_close_virtual_device() -> ResultType<()> { Ok(()) } +pub fn test_create_capturer(privacy_mode_id: i32, timeout_millis: u64) -> bool { + let test_begin = Instant::now(); + while test_begin.elapsed().as_millis() < timeout_millis as _ { + if let Ok((_, _, display)) = get_current_display() { + if let Ok(_) = create_capturer(privacy_mode_id, display) { + return true; + } + } + std::thread::sleep(Duration::from_millis(300)); + } + false +} + fn run(sp: GenericService) -> ResultType<()> { #[cfg(windows)] ensure_close_virtual_device()?; @@ -330,6 +322,10 @@ fn run(sp: GenericService) -> ResultType<()> { ); let privacy_mode_id = *PRIVACY_MODE_CONN_ID.lock().unwrap(); + log::debug!( + "Try create capturer with privacy mode id {}", + privacy_mode_id, + ); let mut c = create_capturer(privacy_mode_id, display)?; let q = get_image_quality(); @@ -403,6 +399,7 @@ fn run(sp: GenericService) -> ResultType<()> { bail!("SWITCH"); } } + *LAST_ACTIVE.lock().unwrap() = now; frame_controller.reset(); @@ -478,8 +475,17 @@ fn run(sp: GenericService) -> ResultType<()> { _ => {} } - // i love 3, 6, 8 - frame_controller.blocking_wait_next(3_000); + let mut fetched_conn_ids = HashSet::new(); + let timeout_millis = 3_000u64; + let wait_begin = Instant::now(); + while wait_begin.elapsed().as_millis() < timeout_millis as _ { + check_privacy_mode_changed(&sp, privacy_mode_id)?; + frame_controller.try_wait_next(&mut fetched_conn_ids, 300); + // break if all connections have received current frame + if fetched_conn_ids.len() >= frame_controller.send_conn_ids.len() { + break; + } + } let elapsed = now.elapsed(); // may need to enable frame(timeout) diff --git a/src/ui/win_privacy.rs b/src/ui/win_privacy.rs index 71daaf0f0..11d0ff91b 100644 --- a/src/ui/win_privacy.rs +++ b/src/ui/win_privacy.rs @@ -154,7 +154,7 @@ pub fn start() -> ResultType<()> { let hwnd = wait_find_privacy_hwnd(1_000)?; if !hwnd.is_null() { - log::info!("Privacy window is already created"); + log::info!("Privacy window is ready"); return Ok(()); } @@ -320,9 +320,7 @@ async fn set_privacy_mode_state( state: PrivacyModeState, ms_timeout: u64, ) -> ResultType<()> { - println!("set_privacy_mode_state begin"); let mut c = connect(ms_timeout, "_cm").await?; - println!("set_privacy_mode_state connect done"); c.send(&Data::PrivacyModeState((conn_id, state))).await }