Merge dd4381880e into 3c574a4182

libs/scrap/src/common/wayland.rs

@@ -17,10 +17,17 @@ pub fn set_map_err(f: fn(err: String) -> io::Error) {
 }
 
 fn map_err<E: ToString>(err: E) -> io::Error {
+    let err_str = err.to_string();
+    
+    if err_str.contains("SESSION_REVOKED")
+    {
+        return io::Error::new(io::ErrorKind::ConnectionAborted, err_str);
+    }
+
     if let Some(f) = *MAP_ERR.read().unwrap() {
-        f(err.to_string())
+        f(err_str)
     } else {
-        io::Error::new(io::ErrorKind::Other, err.to_string())
+        io::Error::new(io::ErrorKind::Other, err_str)
     }
 }
 

libs/scrap/src/wayland/pipewire.rs

@@ -95,6 +95,14 @@ pub fn try_close_session() {
     }
 }
 
+pub fn force_close_dead_session() {
+    if let Ok(mut rdp_info) = RDP_SESSION_INFO.lock() {
+        *rdp_info = None;
+        clear_wayland_displays_cache();
+        HAS_POSITION_ATTR.store(false, Ordering::SeqCst);
+    }
+}
+
 pub struct RdpSessionInfo {
     pub conn: Arc<SyncConnection>,
     pub streams: Vec<PwStreamInfo>,
@@ -145,6 +153,17 @@ impl std::fmt::Display for GStreamerError {
 
 impl Error for GStreamerError {}
 
+#[derive(Debug)]
+pub struct SessionRevokedError;
+
+impl std::fmt::Display for SessionRevokedError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "SESSION_REVOKED")
+    }
+}
+
+impl Error for SessionRevokedError {}
+
 #[derive(Clone)]
 pub struct PipeWireCapturable {
     // connection needs to be kept alive for recording
@@ -315,7 +334,17 @@ impl PipeWireRecorder {
             "[gstreamer] Setting pipeline {} to PLAYING state...",
             capturable.fd.as_raw_fd()
         );
-        pipeline.set_state(gst::State::Playing)?;
+        if let Err(e) = pipeline.set_state(gst::State::Playing) {
+            let _ = pipeline.set_state(gst::State::Null);
+
+            if is_server_running() {
+                warn!("[gstreamer] Failed to set PLAYING state: {:?}", e);
+                return Err(hbb_common::anyhow::Error::msg(format!("GStreamer pipeline failed to start: {:?}", e)));
+            } else {
+                warn!("[gstreamer] Failed to set PLAYING state, session was likely revoked: {:?}", e);
+                return Err(hbb_common::anyhow::Error::new(SessionRevokedError));
+            }
+        }
 
         // If `is_server_running()` is false, it means using remote_desktop_portal,
         // which does not use multiple streams, so no need to wait for state change.
@@ -332,9 +361,15 @@ impl PipeWireRecorder {
                 }
                 (result, state, pending) => {
                     warn!(
-                    "[gstreamer] Pipeline {} state change incomplete: result={:?}, state={:?}, pending={:?}",
-                    capturable.fd.as_raw_fd(), result, state, pending
-                );
+                        "[gstreamer] Pipeline {} state change incomplete: result={:?}, state={:?}, pending={:?}",
+                        capturable.fd.as_raw_fd(), result, state, pending
+                    );
+
+                    if let Err(_) = result {
+                        warn!("[gstreamer] Async pipeline error detected. Session was likely terminated.");
+                        let _ = pipeline.set_state(gst::State::Null);
+                        return Err(hbb_common::anyhow::Error::new(SessionRevokedError));
+                    }
                 }
             }
             std::thread::sleep(std::time::Duration::from_millis(150));

src/server/wayland.rs

@@ -133,7 +133,7 @@ pub(super) fn is_inited() -> Option<Message> {
     }
 }
 
-pub(super) async fn check_init() -> ResultType<()> {
+async fn check_init_once() -> ResultType<()> {
     if !is_x11() {
         if CAP_DISPLAY_INFO.read().unwrap().is_empty() {
             if crate::input_service::wayland_use_uinput() {
@@ -172,7 +172,14 @@ pub(super) async fn check_init() -> ResultType<()> {
                 }
                 log::debug!("Attempting to fix logical size with try_fix_logical_size()");
                 try_fix_logical_size(&mut all);
-                *PIPEWIRE_INITIALIZED.write().unwrap() = true;
+
+                // Bail early if no displays were found: the portal session was likely revoked (e.g., after screen lock).
+                if all.is_empty() {
+                    log::warn!("check_init: no displays from PipeWire portal, session revoked.");
+                    scrap::wayland::pipewire::close_session();
+                    bail!("No displays returned by PipeWire portal. Try reconnecting to request a new screen-sharing session.");
+                }
+
                 let num = all.len();
                 let primary = super::display_service::get_primary_2(&all);
                 super::display_service::check_update_displays(&all);
@@ -196,6 +203,8 @@ pub(super) async fn check_init() -> ResultType<()> {
                 );
 
                 // Create individual CapDisplayInfo for each display with its own capturer
+                let init_result: ResultType<()> = (|| {
+
                 for (idx, display) in all.into_iter().enumerate() {
                     let capturer =
                         Box::into_raw(Box::new(Capturer::new(display).with_context(|| {
@@ -214,12 +223,58 @@ pub(super) async fn check_init() -> ResultType<()> {
 
                     lock.insert(idx, cap_display_info as u64);
                 }
+
+                Ok(())
+                })();
+
+                if let Err(e) = init_result {
+                    log::error!("check_init: capturer loop failed, cleaning up partial state: {:?}", e);
+                    for (_, addr) in lock.iter() {
+                        let cap_display_info: *mut CapDisplayInfo = *addr as _;
+                        unsafe {
+                            let _box_capturer = Box::from_raw((*cap_display_info).capturer.0);
+                            let _box_cap_display_info = Box::from_raw(cap_display_info);
+                        }
+                    }
+                    lock.clear();
+
+                    let err_str = format!("{:?}", e);
+                    if err_str.contains("SESSION_REVOKED") {
+                        log::warn!("check_init: Detected Wayland session death. Forcing hard reset");
+                        *PIPEWIRE_INITIALIZED.write().unwrap() = false;
+                        scrap::wayland::pipewire::force_close_dead_session();
+                    } else {
+                        scrap::wayland::pipewire::close_session();
+                    }
+
+                    return Err(e);
+                }
+
+                // Only mark as initialized after the entire loop succeeds.
+                *PIPEWIRE_INITIALIZED.write().unwrap() = true;
             }
         }
     }
     Ok(())
 }
 
+pub(super) async fn check_init() -> ResultType<()> {
+    const MAX_RETRIES: usize = 1;
+    let mut retry_count = 0;
+    loop {
+        let result = check_init_once().await;
+        if let Err(ref e) = result {
+            if format!("{:?}", e).contains("SESSION_REVOKED") && retry_count < MAX_RETRIES {
+                retry_count += 1;
+                // Brief pause before re-requesting the portal permission dialog, to avoid back-to-back prompts firing immediately.
+                tokio::time::sleep(std::time::Duration::from_millis(500)).await;
+                continue;
+            }
+        }
+        return result;
+    }
+}
+
 pub(super) async fn get_displays() -> ResultType<Vec<DisplayInfo>> {
     check_init().await?;
     let cap_map = CAP_DISPLAY_INFO.read().unwrap();
@@ -230,6 +285,11 @@ pub(super) async fn get_displays() -> ResultType<Vec<DisplayInfo>> {
             Ok(cap_display_info.displays.clone())
         }
     } else {
+        drop(cap_map);
+        log::warn!(
+            "get_displays: map empty after check_init(); resetting PIPEWIRE_INITIALIZED to allow retry."
+        );
+        *PIPEWIRE_INITIALIZED.write().unwrap() = false;
         bail!("Failed to get capturer display info");
     }
 }
@@ -271,23 +331,64 @@ pub(super) fn get_capturer_for_display(
     if is_x11() {
         bail!("Do not call this function if not wayland");
     }
-    let cap_map = CAP_DISPLAY_INFO.read().unwrap();
-    if let Some(addr) = cap_map.get(&display_idx) {
-        let cap_display_info: *const CapDisplayInfo = *addr as _;
-        unsafe {
-            let cap_display_info = &*cap_display_info;
-            let rect = cap_display_info.rects[cap_display_info.current];
-            Ok(super::video_service::CapturerInfo {
-                origin: rect.0,
-                width: rect.1,
-                height: rect.2,
-                ndisplay: cap_display_info.num,
-                current: cap_display_info.current,
-                privacy_mode_id: 0,
-                _capturer_privacy_mode_id: 0,
-                capturer: Box::new(cap_display_info.capturer.clone()),
-            })
+
+    let build_capturer_info =
+        |addr: u64| -> ResultType<super::video_service::CapturerInfo> {
+            let cap_display_info: *const CapDisplayInfo = addr as _;
+            unsafe {
+                let cap_display_info = &*cap_display_info;
+                let rect = cap_display_info.rects[cap_display_info.current];
+                Ok(super::video_service::CapturerInfo {
+                    origin: rect.0,
+                    width: rect.1,
+                    height: rect.2,
+                    ndisplay: cap_display_info.num,
+                    current: cap_display_info.current,
+                    privacy_mode_id: 0,
+                    _capturer_privacy_mode_id: 0,
+                    capturer: Box::new(cap_display_info.capturer.clone()),
+                })
+            }
+        };
+
+    {
+        let cap_map = CAP_DISPLAY_INFO.read().unwrap();
+        if let Some(&addr) = cap_map.get(&display_idx) {
+            return build_capturer_info(addr);
         }
+    }
+
+    log::warn!(
+        "get_capturer_for_display: display {} not found in CAP_DISPLAY_INFO.",
+        display_idx
+    );
+
+    // Wait until all active capturers have exited before reinitializing.
+    let active_count = *ACTIVE_DISPLAY_COUNT.read().unwrap();
+    if active_count > 1 {
+        bail!(
+            "Display {} not found in CAP_DISPLAY_INFO, but {} active capturer(s) are still running. Skipping reinitialization now.",
+            display_idx, active_count
+        );
+    }
+
+    log::info!(
+        "get_capturer_for_display: no active capturers, reinitializing PipeWire session for display {}.",
+        display_idx
+    );
+    
+    // to-do: Potential use-after-free: clone() aliases raw pointer and clear() frees it.
+    // Requires ownership refactor to make lifetime sound.
+    clear();
+    ensure_inited()?;
+
+    let cap_map = CAP_DISPLAY_INFO.read().unwrap();
+    if let Some(&addr) = cap_map.get(&display_idx) {
+        log::info!(
+            "get_capturer_for_display: re-initialization succeeded for display {}.",
+            display_idx
+        );
+        build_capturer_info(addr)
     } else {
         bail!(
             "Failed to get capturer display info for display {}",