From c5038b1a78fe6d0ef6f4ef46aae180896859fe7f Mon Sep 17 00:00:00 2001 From: fufesou <13586388+fufesou@users.noreply.github.com> Date: Mon, 16 Sep 2024 15:57:40 +0800 Subject: [PATCH] Fix/virtual display do not plug out if not plugged in (#9372) * fix: win VD, do not plug out if not plugged in Signed-off-by: fufesou * Forcibly virtual display on clicking button "-" Signed-off-by: fufesou --------- Signed-off-by: fufesou --- src/privacy_mode/win_virtual_display.rs | 11 ++-- src/server/connection.rs | 2 +- src/virtual_display_manager.rs | 86 ++++++++++++++++++++----- 3 files changed, 78 insertions(+), 21 deletions(-) diff --git a/src/privacy_mode/win_virtual_display.rs b/src/privacy_mode/win_virtual_display.rs index 25997f036..782d7ed75 100644 --- a/src/privacy_mode/win_virtual_display.rs +++ b/src/privacy_mode/win_virtual_display.rs @@ -150,8 +150,11 @@ impl PrivacyModeImpl { } fn restore_plug_out_monitor(&mut self) { - let _ = - virtual_display_manager::plug_out_monitor_indices(&self.virtual_displays_added, true); + let _ = virtual_display_manager::plug_out_monitor_indices( + &self.virtual_displays_added, + true, + false, + ); self.virtual_displays_added.clear(); } @@ -312,7 +315,7 @@ impl PrivacyModeImpl { // No physical displays, no need to use the privacy mode. if self.displays.is_empty() { - virtual_display_manager::plug_out_monitor_indices(&displays, false)?; + virtual_display_manager::plug_out_monitor_indices(&displays, false, false)?; bail!(NO_PHYSICAL_DISPLAYS); } @@ -509,7 +512,7 @@ pub fn restore_reg_connectivity(plug_out_monitors: bool) { return; } if plug_out_monitors { - let _ = virtual_display_manager::plug_out_monitor(-1, true); + let _ = virtual_display_manager::plug_out_monitor(-1, true, false); } if let Ok(reg_recovery) = serde_json::from_str::(&config_recovery_value) diff --git a/src/server/connection.rs b/src/server/connection.rs index 7181478b3..e351b0d50 100644 --- a/src/server/connection.rs +++ b/src/server/connection.rs @@ -2713,7 +2713,7 @@ impl Connection { } } } else { - if let Err(e) = virtual_display_manager::plug_out_monitor(t.display, false) { + if let Err(e) = virtual_display_manager::plug_out_monitor(t.display, false, true) { log::error!("Failed to plug out virtual display {}: {}", t.display, e); self.send(make_msg(format!( "Failed to plug out virtual displays: {}", diff --git a/src/virtual_display_manager.rs b/src/virtual_display_manager.rs index 138087c75..41e5b3fc8 100644 --- a/src/virtual_display_manager.rs +++ b/src/virtual_display_manager.rs @@ -8,6 +8,7 @@ pub const AMYUNI_IDD_DEVICE_STRING: &'static str = "USB Mobile Monitor Virtual D const IDD_IMPL: &str = IDD_IMPL_AMYUNI; const IDD_IMPL_RUSTDESK: &str = "rustdesk_idd"; const IDD_IMPL_AMYUNI: &str = "amyuni_idd"; +const IDD_PLUG_OUT_ALL_INDEX: i32 = -1; pub fn is_amyuni_idd() -> bool { IDD_IMPL == IDD_IMPL_AMYUNI @@ -76,17 +77,17 @@ pub fn plug_in_monitor(idx: u32, modes: Vec) -> Re } } -pub fn plug_out_monitor(index: i32, force_all: bool) -> ResultType<()> { +pub fn plug_out_monitor(index: i32, force_all: bool, force_one: bool) -> ResultType<()> { match IDD_IMPL { IDD_IMPL_RUSTDESK => { - let indices = if index == -1 { + let indices = if index == IDD_PLUG_OUT_ALL_INDEX { rustdesk_idd::get_virtual_displays() } else { vec![index as _] }; rustdesk_idd::plug_out_peer_request(&indices) } - IDD_IMPL_AMYUNI => amyuni_idd::plug_out_monitor(index, force_all), + IDD_IMPL_AMYUNI => amyuni_idd::plug_out_monitor(index, force_all, force_one), _ => bail!("Unsupported virtual display implementation."), } } @@ -102,12 +103,16 @@ pub fn plug_in_peer_request(modes: Vec>) -> Re } } -pub fn plug_out_monitor_indices(indices: &[u32], force_all: bool) -> ResultType<()> { +pub fn plug_out_monitor_indices( + indices: &[u32], + force_all: bool, + force_one: bool, +) -> ResultType<()> { match IDD_IMPL { IDD_IMPL_RUSTDESK => rustdesk_idd::plug_out_peer_request(indices), IDD_IMPL_AMYUNI => { for _idx in indices.iter() { - amyuni_idd::plug_out_monitor(0, force_all)?; + amyuni_idd::plug_out_monitor(0, force_all, force_one)?; } Ok(()) } @@ -382,7 +387,7 @@ pub mod amyuni_idd { use hbb_common::{bail, lazy_static, log, tokio::time::Instant, ResultType}; use std::{ ptr::null_mut, - sync::{Arc, Mutex}, + sync::{atomic, Arc, Mutex}, time::Duration, }; use winapi::{ @@ -405,6 +410,14 @@ pub mod amyuni_idd { static ref LOCK: Arc> = Default::default(); static ref LAST_PLUG_IN_HEADLESS_TIME: Arc>> = Arc::new(Mutex::new(None)); } + const VIRTUAL_DISPLAY_MAX_COUNT: usize = 4; + // The count of virtual displays plugged in. + // This count is not accurate, because: + // 1. The virtual display driver may also be controlled by other processes. + // 2. RustDesk may crash and restart, but the virtual displays are kept. + // + // to-do: Maybe a better way is to add an option asking the user if plug out all virtual displays on disconnect. + static VIRTUAL_DISPLAY_COUNT: atomic::AtomicUsize = atomic::AtomicUsize::new(0); fn get_deviceinstaller64_work_dir() -> ResultType>> { let cur_exe = std::env::current_exe()?; @@ -510,7 +523,7 @@ pub mod amyuni_idd { pub fn reset_all() -> ResultType<()> { let _ = crate::privacy_mode::turn_off_privacy(0, None); - let _ = plug_out_monitor(-1, true); + let _ = plug_out_monitor(super::IDD_PLUG_OUT_ALL_INDEX, true, false); *LAST_PLUG_IN_HEADLESS_TIME.lock().unwrap() = None; Ok(()) } @@ -522,6 +535,18 @@ pub mod amyuni_idd { unsafe { win_device::device_io_control(&INTERFACE_GUID, PLUG_MONITOR_IO_CONTROL_CDOE, &cmd, 0)?; } + // No need to consider concurrency here. + if add { + // If the monitor is plugged in, increase the count. + // Though there's already a check of `VIRTUAL_DISPLAY_MAX_COUNT`, it's still better to check here for double ensure. + if VIRTUAL_DISPLAY_COUNT.load(atomic::Ordering::SeqCst) < VIRTUAL_DISPLAY_MAX_COUNT { + VIRTUAL_DISPLAY_COUNT.fetch_add(1, atomic::Ordering::SeqCst); + } + } else { + if VIRTUAL_DISPLAY_COUNT.load(atomic::Ordering::SeqCst) > 0 { + VIRTUAL_DISPLAY_COUNT.fetch_sub(1, atomic::Ordering::SeqCst); + } + } Ok(()) } @@ -607,44 +632,73 @@ pub mod amyuni_idd { bail!("Failed to install driver."); } - if get_monitor_count() == 4 { - bail!("There are already 4 monitors plugged in."); + if get_monitor_count() == VIRTUAL_DISPLAY_MAX_COUNT { + bail!("There are already {VIRTUAL_DISPLAY_MAX_COUNT} monitors plugged in."); } plug_in_monitor_(true, is_async) } - pub fn plug_out_monitor(index: i32, force_all: bool) -> ResultType<()> { - let all_count = windows::get_device_names(None).len(); + // `index` the display index to plug out. -1 means plug out all. + // `force_all` is used to forcibly plug out all virtual displays. + // `force_one` is used to forcibly plug out one virtual display managed by other processes + // if there're no virtual displays managed by RustDesk. + pub fn plug_out_monitor(index: i32, force_all: bool, force_one: bool) -> ResultType<()> { + let plug_out_all = index == super::IDD_PLUG_OUT_ALL_INDEX; + // If `plug_out_all and force_all` is true, forcibly plug out all virtual displays. + // Though the driver may be controlled by other processes, + // we still forcibly plug out all virtual displays. + // + // 1. RustDesk plug in 2 virtual displays. (RustDesk) + // 2. Other process plug out all virtual displays. (User mannually) + // 3. Other process plug in 1 virtual display. (User mannually) + // 4. RustDesk plug out all virtual displays in this call. (RustDesk disconnect) + // + // This is not a normal scenario, RustDesk will plug out virtual display unexpectedly. + let mut plug_in_count = VIRTUAL_DISPLAY_COUNT.load(atomic::Ordering::Relaxed); let amyuni_count = get_monitor_count(); + if !plug_out_all { + if plug_in_count == 0 && amyuni_count > 0 { + if force_one { + plug_in_count = 1; + } else { + bail!("The virtual display is managed by other processes."); + } + } + } else { + // Ignore the message if trying to plug out all virtual displays. + } + + let all_count = windows::get_device_names(None).len(); let mut to_plug_out_count = match all_count { 0 => return Ok(()), 1 => { - if amyuni_count == 0 { + if plug_in_count == 0 { bail!("No virtual displays to plug out.") } else { if force_all { 1 } else { - bail!("This only virtual display cannot be pulled out.") + bail!("This only virtual display cannot be plugged out.") } } } _ => { - if all_count == amyuni_count { + if all_count == plug_in_count { if force_all { all_count } else { all_count - 1 } } else { - amyuni_count + plug_in_count } } }; - if to_plug_out_count != 0 && index != -1 { + if to_plug_out_count != 0 && !plug_out_all { to_plug_out_count = 1; } + for _i in 0..to_plug_out_count { let _ = plug_monitor_(false); }