Skip to content

Commit

Permalink
Fix/virtual display do not plug out if not plugged in (rustdesk#9372)
Browse files Browse the repository at this point in the history 
* fix: win VD, do not plug out if not plugged in

Signed-off-by: fufesou <[email protected]>

* Forcibly virtual display on clicking button "-"

Signed-off-by: fufesou <[email protected]>

---------

Signed-off-by: fufesou <[email protected]>
  • Loading branch information
@fufesou
fufesou authored Sep 16, 2024
1 parent f4c038e commit c5038b1
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 21 deletions.
11 changes: 7 additions & 4 deletions src/privacy_mode/win_virtual_display.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}

Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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::<reg_display_settings::RegRecovery>(&config_recovery_value)
Expand Down
2 changes: 1 addition & 1 deletion src/server/connection.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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: {}",
Expand Down
86 changes: 70 additions & 16 deletions src/virtual_display_manager.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -76,17 +77,17 @@ pub fn plug_in_monitor(idx: u32, modes: Vec<virtual_display::MonitorMode>) -> 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."),
}
}
Expand All @@ -102,12 +103,16 @@ pub fn plug_in_peer_request(modes: Vec<Vec<virtual_display::MonitorMode>>) -> 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(())
}
Expand Down Expand Up @@ -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::{
Expand All @@ -405,6 +410,14 @@ pub mod amyuni_idd {
static ref LOCK: Arc<Mutex<()>> = Default::default();
static ref LAST_PLUG_IN_HEADLESS_TIME: Arc<Mutex<Option<Instant>>> = 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<Option<Vec<u8>>> {
let cur_exe = std::env::current_exe()?;
Expand Down Expand Up @@ -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(())
}
Expand All @@ -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(())
}

Expand Down Expand Up @@ -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);
}
Expand Down

0 comments on commit c5038b1

Please sign in to comment.