Merge branch 'rustdesk:master' into master

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::<reg_display_settings::RegRecovery>(&config_recovery_value)

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: {}",

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<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."),
     }
 }
@@ -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(())
         }
@@ -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<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()?;
@@ -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);
         }