diff --git a/libs/scrap/src/common/wayland.rs b/libs/scrap/src/common/wayland.rs index 30b5f4d54..7c96f162b 100644 --- a/libs/scrap/src/common/wayland.rs +++ b/libs/scrap/src/common/wayland.rs @@ -17,10 +17,17 @@ pub fn set_map_err(f: fn(err: String) -> io::Error) { } fn map_err(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) } } diff --git a/libs/scrap/src/wayland/pipewire.rs b/libs/scrap/src/wayland/pipewire.rs index 8859d0d3b..3d99a5e55 100644 --- a/libs/scrap/src/wayland/pipewire.rs +++ b/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, pub streams: Vec, @@ -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)); diff --git a/src/server/wayland.rs b/src/server/wayland.rs index 1e0efc0f4..653c86411 100644 --- a/src/server/wayland.rs +++ b/src/server/wayland.rs @@ -133,7 +133,7 @@ pub(super) fn is_inited() -> Option { } } -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> { check_init().await?; let cap_map = CAP_DISPLAY_INFO.read().unwrap(); @@ -230,6 +285,11 @@ pub(super) async fn get_displays() -> ResultType> { 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 { + 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 {}",