fix stale primary display selection (#15460)

* fix stale primary display selection

Signed-off-by: 21pages <sunboeasy@gmail.com>

* fix stale display selection during login and switching

  - resolve the primary display from the refreshed login snapshot
  - defer display enumeration until authentication succeeds
  - read Wayland displays and primary index from the same cache snapshot
  - reject stale monitor and camera indices during display switching

Signed-off-by: 21pages <sunboeasy@gmail.com>

* fix inconsistent display snapshots during login

  - return displays from the same enumeration used to select the primary
  - avoid re-reading the shared display cache after updating it
  - use the same converted snapshot during Wayland initialization

Signed-off-by: 21pages <sunboeasy@gmail.com>

* avoid cloning unchanged display snapshots

Signed-off-by: 21pages <sunboeasy@gmail.com>

* fix invalid display subset handling

Signed-off-by: 21pages <sunboeasy@gmail.com>

* minimize code churn in switch_display_to

Signed-off-by: 21pages <sunboeasy@gmail.com>

---------

Signed-off-by: 21pages <sunboeasy@gmail.com>

src/server.rs

@@ -357,15 +357,13 @@ impl Server {
         }
     }
 
-    pub fn try_add_primay_video_service(&mut self) {
-        let primary_video_service_name = video_service::get_service_name(
-            VideoSource::Monitor,
-            *display_service::PRIMARY_DISPLAY_IDX,
-        );
-        if !self.contains(&primary_video_service_name) {
+    pub fn try_add_monitor_service(&mut self, display_idx: usize) {
+        let monitor_service_name =
+            video_service::get_service_name(VideoSource::Monitor, display_idx);
+        if !self.contains(&monitor_service_name) {
             self.add_service(Box::new(video_service::new(
                 VideoSource::Monitor,
-                *display_service::PRIMARY_DISPLAY_IDX,
+                display_idx,
             )));
         }
     }
@@ -381,14 +379,17 @@ impl Server {
         self.connections.insert(conn.id(), conn);
     }
 
-    pub fn add_connection(&mut self, conn: ConnInner, noperms: &Vec<&'static str>) {
-        let primary_video_service_name = video_service::get_service_name(
-            VideoSource::Monitor,
-            *display_service::PRIMARY_DISPLAY_IDX,
-        );
+    pub fn add_monitor_connection(
+        &mut self,
+        conn: ConnInner,
+        noperms: &Vec<&'static str>,
+        display_idx: usize,
+    ) {
+        let monitor_service_name =
+            video_service::get_service_name(VideoSource::Monitor, display_idx);
         for s in self.services.values() {
             let name = s.name();
-            if Self::is_video_service_name(&name) && name != primary_video_service_name {
+            if Self::is_video_service_name(&name) && name != monitor_service_name {
                 continue;
             }
             if !noperms.contains(&(&name as _)) {
@@ -783,8 +784,7 @@ async fn sync_and_watch_config_dir(sync_done_tx: Option<tokio::sync::oneshot::Se
                 loop {
                     sleep(CONFIG_SYNC_INTERVAL_SECS).await;
                     let cfg = (Config::get(), Config2::get());
-                    let should_sync =
-                        cfg != cfg0 || (is_root_config_empty && !cfg.0.is_empty());
+                    let should_sync = cfg != cfg0 || (is_root_config_empty && !cfg.0.is_empty());
                     if should_sync {
                         if is_root_config_empty {
                             log::info!("root config is empty, sync our config to root");

src/server/connection.rs

@@ -503,7 +503,9 @@ impl Connection {
                 tx_video: Some(tx_video),
             },
             require_2fa: crate::auth_2fa::get_2fa(None),
-            display_idx: *display_service::PRIMARY_DISPLAY_IDX,
+            // Defer display enumeration until login succeeds. Monitor login replaces this
+            // with the primary index returned with the refreshed display snapshot.
+            display_idx: 0,
             stream,
             server,
             hash,
@@ -1891,13 +1893,15 @@ impl Connection {
                 Err(err) => {
                     res.set_error(format!("{}", err));
                 }
-                Ok(displays) => {
+                Ok((displays, primary_display_idx)) => {
                     // For compatibility with old versions, we need to send the displays to the peer.
                     // But the displays may be updated later, before creating the video capturer.
                     #[cfg(target_os = "macos")]
                     {
                         self.retina.set_displays(&displays);
                     }
+                    // A separate primary lookup here could race with display hot-plug.
+                    self.display_idx = primary_display_idx;
                     pi.displays = displays;
                     pi.current_display = self.display_idx as _;
                     #[cfg(not(any(target_os = "android", target_os = "ios")))]
@@ -2006,8 +2010,8 @@ impl Connection {
                 #[cfg(not(any(target_os = "android", target_os = "ios")))]
                 let _h = try_start_record_cursor_pos();
                 self.auto_disconnect_timer = Self::get_auto_disconenct_timer();
-                s.try_add_primay_video_service();
-                s.add_connection(self.inner.clone(), &noperms);
+                s.try_add_monitor_service(self.display_idx);
+                s.add_monitor_connection(self.inner.clone(), &noperms, self.display_idx);
             }
         }
     }
@@ -4150,7 +4154,9 @@ impl Connection {
         let display_idx = s.display as usize;
         if self.display_idx != display_idx {
             if let Some(server) = self.server.upgrade() {
-                self.switch_display_to(display_idx, server.clone());
+                if !self.switch_display_to(display_idx, server.clone()) {
+                    return;
+                }
 
                 #[cfg(not(any(target_os = "android", target_os = "ios")))]
                 if !self.view_camera && s.width != 0 && s.height != 0 {
@@ -4177,6 +4183,13 @@ impl Connection {
         }
     }
 
+    fn video_source_count(video_source: VideoSource) -> usize {
+        match video_source {
+            VideoSource::Monitor => display_service::get_sync_displays().len(),
+            VideoSource::Camera => camera::Cameras::get_sync_cameras().len(),
+        }
+    }
+
     fn video_source(&self) -> VideoSource {
         if self.view_camera {
             VideoSource::Camera
@@ -4185,18 +4198,28 @@ impl Connection {
         }
     }
 
-    fn switch_display_to(&mut self, display_idx: usize, server: Arc<RwLock<Server>>) {
+    fn switch_display_to(&mut self, display_idx: usize, server: Arc<RwLock<Server>>) -> bool {
+        let source_count = Self::video_source_count(self.video_source());
+        if display_idx >= source_count {
+            // Do not remap an explicit switch: its resolution belongs to the requested source.
+            log::warn!(
+                "Ignore switch to invalid {:?} index {}, available source count: {}",
+                self.video_source(),
+                display_idx,
+                source_count
+            );
+            return false;
+        }
+
         let new_service_name = video_service::get_service_name(self.video_source(), display_idx);
         let old_service_name =
             video_service::get_service_name(self.video_source(), self.display_idx);
         let mut lock = server.write().unwrap();
-        if display_idx != *display_service::PRIMARY_DISPLAY_IDX {
-            if !lock.contains(&new_service_name) {
-                lock.add_service(Box::new(video_service::new(
-                    self.video_source(),
-                    display_idx,
-                )));
-            }
+        if !lock.contains(&new_service_name) {
+            lock.add_service(Box::new(video_service::new(
+                self.video_source(),
+                display_idx,
+            )));
         }
         // For versions greater than 1.2.4, a `CaptureDisplays` message will be sent immediately.
         // Unnecessary capturers will be removed then.
@@ -4205,6 +4228,7 @@ impl Connection {
         }
         lock.subscribe(&new_service_name, self.inner.clone(), true);
         self.display_idx = display_idx;
+        true
     }
 
     #[cfg(windows)]
@@ -4231,26 +4255,61 @@ impl Connection {
 
     async fn capture_displays(&mut self, add: &[usize], sub: &[usize], set: &[usize]) {
         let video_source = self.video_source();
-        if let Some(sever) = self.server.upgrade() {
-            let mut lock = sever.write().unwrap();
-            for display in add.iter() {
+        let source_count = Self::video_source_count(video_source);
+        // Only add/set can create services; sub only narrows existing subscriptions.
+        let valid_add = add
+            .iter()
+            .copied()
+            .filter(|display| *display < source_count)
+            .collect::<Vec<_>>();
+        let valid_sub = sub
+            .iter()
+            .copied()
+            .filter(|display| *display < source_count)
+            .collect::<Vec<_>>();
+        let valid_set = set
+            .iter()
+            .copied()
+            .filter(|display| *display < source_count)
+            .collect::<Vec<_>>();
+        let invalid_count =
+            add.len() + sub.len() + set.len() - valid_add.len() - valid_sub.len() - valid_set.len();
+        if invalid_count != 0 {
+            log::warn!(
+                "Ignore {} invalid {:?} indices, available source count: {}",
+                invalid_count,
+                video_source,
+                source_count
+            );
+        }
+        // Passing an invalid sub request as an empty exclude list would unsubscribe all services.
+        if (!add.is_empty() && valid_add.is_empty())
+            || (add.is_empty() && !sub.is_empty() && valid_sub.is_empty())
+            || (add.is_empty() && sub.is_empty() && !set.is_empty() && valid_set.is_empty())
+        {
+            return;
+        }
+
+        if let Some(server) = self.server.upgrade() {
+            let mut lock = server.write().unwrap();
+            for display in valid_add.iter() {
                 let service_name = video_service::get_service_name(video_source, *display);
                 if !lock.contains(&service_name) {
                     lock.add_service(Box::new(video_service::new(video_source, *display)));
                 }
             }
-            for display in set.iter() {
+            for display in valid_set.iter() {
                 let service_name = video_service::get_service_name(video_source, *display);
                 if !lock.contains(&service_name) {
                     lock.add_service(Box::new(video_service::new(video_source, *display)));
                 }
             }
             if !add.is_empty() {
-                lock.capture_displays(self.inner.clone(), video_source, add, true, false);
+                lock.capture_displays(self.inner.clone(), video_source, &valid_add, true, false);
             } else if !sub.is_empty() {
-                lock.capture_displays(self.inner.clone(), video_source, sub, false, true);
+                lock.capture_displays(self.inner.clone(), video_source, &valid_sub, false, true);
             } else {
-                lock.capture_displays(self.inner.clone(), video_source, set, true, true);
+                lock.capture_displays(self.inner.clone(), video_source, &valid_set, true, true);
             }
             self.multi_ui_session = lock.get_subbed_displays_count(self.inner.id()) > 1;
             if self.follow_remote_window {

src/server/display_service.rs

@@ -25,9 +25,6 @@ struct ChangedResolution {
 lazy_static::lazy_static! {
     static ref IS_CAPTURER_MAGNIFIER_SUPPORTED: bool = is_capturer_mag_supported();
     static ref CHANGED_RESOLUTIONS: Arc<RwLock<HashMap<String, ChangedResolution>>> = Default::default();
-    // Initial primary display index.
-    // It should not be updated when displays changed.
-    pub static ref PRIMARY_DISPLAY_IDX: usize = get_primary();
     static ref SYNC_DISPLAYS: Arc<Mutex<SyncDisplaysInfo>> = Default::default();
 }
 
@@ -41,22 +38,14 @@ struct SyncDisplaysInfo {
 }
 
 impl SyncDisplaysInfo {
-    fn check_changed(&mut self, displays: Vec<DisplayInfo>) {
-        if self.displays.len() != displays.len() {
-            self.displays = displays;
-            if !TEMP_IGNORE_DISPLAYS_CHANGED.load(Ordering::Relaxed) {
-                self.is_synced = false;
-            }
+    fn check_changed(&mut self, displays: &[DisplayInfo]) {
+        if self.displays.as_slice() == displays {
             return;
         }
-        for (i, d) in displays.iter().enumerate() {
-            if d != &self.displays[i] {
-                self.displays = displays;
-                if !TEMP_IGNORE_DISPLAYS_CHANGED.load(Ordering::Relaxed) {
-                    self.is_synced = false;
-                }
-                return;
-            }
+
+        self.displays = displays.to_vec();
+        if !TEMP_IGNORE_DISPLAYS_CHANGED.load(Ordering::Relaxed) {
+            self.is_synced = false;
         }
     }
 
@@ -304,6 +293,11 @@ pub(super) fn get_display_info(idx: usize) -> Option<DisplayInfo> {
 // Display to DisplayInfo
 // The DisplayInfo is be sent to the peer.
 pub(super) fn check_update_displays(all: &Vec<Display>) {
+    let _ = update_sync_displays(all);
+}
+
+// Return the converted input snapshot while updating the shared display cache.
+pub(super) fn update_sync_displays(all: &Vec<Display>) -> Vec<DisplayInfo> {
     // For compatibility: if only one display, scale remains 1.0 and we use the physical size for `uinput`.
     // If there are multiple displays, we use the logical size for `uinput` by setting scale to d.scale().
     #[cfg(target_os = "linux")]
@@ -346,7 +340,8 @@ pub(super) fn check_update_displays(all: &Vec<Display>) {
             }
         })
         .collect::<Vec<DisplayInfo>>();
-    SYNC_DISPLAYS.lock().unwrap().check_changed(displays);
+    SYNC_DISPLAYS.lock().unwrap().check_changed(&displays);
+    displays
 }
 
 pub fn is_inited_msg() -> Option<Message> {
@@ -357,34 +352,38 @@ pub fn is_inited_msg() -> Option<Message> {
     None
 }
 
-pub async fn update_get_sync_displays_on_login() -> ResultType<Vec<DisplayInfo>> {
+// Return the primary index with the refreshed list so login cannot mix display snapshots.
+pub async fn update_get_sync_displays_on_login() -> ResultType<(Vec<DisplayInfo>, usize)> {
     #[cfg(target_os = "linux")]
     {
         if !is_x11() {
-            return super::wayland::get_displays().await;
+            let (displays, primary_display_idx) =
+                super::wayland::get_displays_and_primary().await?;
+            let primary_display_idx =
+                normalize_primary_display_idx(primary_display_idx, displays.len());
+            return Ok((displays, primary_display_idx));
         }
     }
     #[cfg(not(windows))]
     let displays = display_service::try_get_displays();
     #[cfg(windows)]
     let displays = display_service::try_get_displays_add_amyuni_headless();
-    check_update_displays(&displays?);
-    Ok(SYNC_DISPLAYS.lock().unwrap().displays.clone())
+    let displays = displays?;
+    let primary_display_idx = get_primary_2(&displays);
+    let sync_displays = update_sync_displays(&displays);
+    let primary_display_idx =
+        normalize_primary_display_idx(primary_display_idx, sync_displays.len());
+    Ok((sync_displays, primary_display_idx))
 }
 
 #[inline]
-pub fn get_primary() -> usize {
-    #[cfg(target_os = "linux")]
-    {
-        if !is_x11() {
-            return match super::wayland::get_primary() {
-                Ok(n) => n,
-                Err(_) => 0,
-            };
-        }
+fn normalize_primary_display_idx(primary_display_idx: usize, display_len: usize) -> usize {
+    // Zero is the protocol fallback when the list is empty or its primary index is stale.
+    if primary_display_idx < display_len {
+        primary_display_idx
+    } else {
+        0
     }
-
-    try_get_displays().map(|d| get_primary_2(&d)).unwrap_or(0)
 }
 
 #[inline]
@@ -486,3 +485,16 @@ pub fn try_get_displays_(add_amyuni_headless: bool) -> ResultType<Vec<Display>>
     }
     Ok(displays)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::normalize_primary_display_idx;
+
+    #[test]
+    fn normalize_primary_display_idx_bounds() {
+        assert_eq!(normalize_primary_display_idx(0, 0), 0);
+        assert_eq!(normalize_primary_display_idx(0, 2), 0);
+        assert_eq!(normalize_primary_display_idx(1, 2), 1);
+        assert_eq!(normalize_primary_display_idx(2, 2), 0);
+    }
+}

src/server/wayland.rs

@@ -175,8 +175,7 @@ pub(super) async fn check_init() -> ResultType<()> {
                 *PIPEWIRE_INITIALIZED.write().unwrap() = true;
                 let num = all.len();
                 let primary = super::display_service::get_primary_2(&all);
-                super::display_service::check_update_displays(&all);
-                let mut displays = super::display_service::get_sync_displays();
+                let mut displays = super::display_service::update_sync_displays(&all);
                 for display in displays.iter_mut() {
                     display.cursor_embedded = is_cursor_embedded();
                 }
@@ -220,27 +219,15 @@ pub(super) async fn check_init() -> ResultType<()> {
     Ok(())
 }
 
-pub(super) async fn get_displays() -> ResultType<Vec<DisplayInfo>> {
+pub(super) async fn get_displays_and_primary() -> ResultType<(Vec<DisplayInfo>, usize)> {
     check_init().await?;
+    // Keep one read guard so clear/reinitialization cannot split these across cache snapshots.
     let cap_map = CAP_DISPLAY_INFO.read().unwrap();
     if let Some(addr) = cap_map.values().next() {
         let cap_display_info: *const CapDisplayInfo = *addr as _;
         unsafe {
             let cap_display_info = &*cap_display_info;
-            Ok(cap_display_info.displays.clone())
-        }
-    } else {
-        bail!("Failed to get capturer display info");
-    }
-}
-
-pub(super) fn get_primary() -> ResultType<usize> {
-    let cap_map = CAP_DISPLAY_INFO.read().unwrap();
-    if let Some(addr) = cap_map.values().next() {
-        let cap_display_info: *const CapDisplayInfo = *addr as _;
-        unsafe {
-            let cap_display_info = &*cap_display_info;
-            Ok(cap_display_info.primary)
+            Ok((cap_display_info.displays.clone(), cap_display_info.primary))
         }
     } else {
         bail!("Failed to get capturer display info");