Start the desktop on a tool connection when the viewer sends capture_desktop (#99)

src/desktop_control_tests.rs

@@ -215,3 +215,47 @@ async fn locked_control_unlock_requests_and_tools_work_without_a_video_frame() {
         drop(lock);
     }).await.expect("locked bootstrap/tool regression exceeded deadline");
 }
+
+#[tokio::test]
+async fn capture_desktop_on_a_port_mapping_peer_promotes_it_once() {
+    tokio::time::timeout(Duration::from_secs(25), async {
+        let input = Arc::new(InputInjector::without_desktop());
+        let (host, mut viewer) = connect(input.clone(), false, None, &["CONTROL_DATA_CHANNEL", "FILE_DATA_CHANNEL"]).await;
+        let (requests, mut requested) = mpsc::unbounded_channel();
+        host.on_desktop_request(move || { let _ = requests.send(()); });
+        receive(&mut viewer.messages, "CONTROL_DATA_CHANNEL", 13).await;
+        assert!(!host.role.is_desktop() && host.role.input().is_none());
+        // The Mac viewer opens the desktop on its existing port-mapping peer (#99).
+        send(&viewer, "CONTROL_DATA_CHANNEL", crate::initialization::envelope(5, 0, 3, &[8, 4])).await;
+        tokio::time::timeout(Duration::from_secs(5), requested.recv()).await.unwrap().unwrap();
+        host.promote(None);
+        assert!(host.role.is_desktop());
+        assert!(host.role.input().is_some_and(|session| session.is_current()));
+        // CONTROL repeats the desktop initialization for the same viewer.
+        receive(&mut viewer.messages, "CONTROL_DATA_CHANNEL", 13).await;
+        // The port-mapping session on this peer still answers.
+        send(&viewer, "FILE_DATA_CHANNEL", mapping(2, b"still mapped")).await;
+        receive(&mut viewer.messages, "FILE_DATA_CHANNEL", 27).await;
+        // A repeated capture_desktop does not ask again, and a later promote is a no-op.
+        send(&viewer, "CONTROL_DATA_CHANNEL", crate::initialization::envelope(6, 0, 3, &[8, 4])).await;
+        send(&viewer, "CONTROL_DATA_CHANNEL", crate::initialization::envelope(7, 0, 3, &[8, 0])).await;
+        receive(&mut viewer.messages, "CONTROL_DATA_CHANNEL", 3).await;
+        assert!(requested.try_recv().is_err());
+        let lease = host.role.input().unwrap();
+        host.promote(None);
+        assert!(lease.is_current());
+        // A newer desktop demotes it; the port mapping stays, and capture_desktop promotes it again.
+        host.demote();
+        assert!(!host.role.is_desktop() && host.role.input().is_none());
+        send(&viewer, "FILE_DATA_CHANNEL", mapping(2, b"after demotion")).await;
+        receive(&mut viewer.messages, "FILE_DATA_CHANNEL", 27).await;
+        send(&viewer, "CONTROL_DATA_CHANNEL", crate::initialization::envelope(8, 0, 3, &[8, 4])).await;
+        tokio::time::timeout(Duration::from_secs(5), requested.recv()).await.unwrap().unwrap();
+        host.promote(None);
+        assert!(host.role.is_desktop() && host.role.input().is_some_and(|session| session.is_current()));
+        receive(&mut viewer.messages, "CONTROL_DATA_CHANNEL", 13).await;
+        host.pc.close().await.unwrap();
+        viewer.pc.close().await.unwrap();
+        viewer.router.lock().await.stop().await.unwrap();
+    }).await.expect("capture_desktop promotion exceeded deadline");
+}

src/main.rs

@@ -737,6 +737,8 @@ async fn serve(mut capture: Arc<ScreenCapture>, input: Arc<InputInjector>, shutd
     let mut capture_types: HashMap<String, (u64,crate::session_ice::SessionIce)> = HashMap::new();
     let mut mac_by_ice: HashMap<String, bool> = HashMap::new();
     let mut next_mac = false;
+    // Tool peers ask here to become the desktop (capture_desktop on an existing connection, #99).
+    let (desktop_requests, mut desktop_requested) = mpsc::unbounded_channel::<String>();
 
     loop {
         if *shutdown.borrow() {
@@ -821,6 +823,7 @@ async fn serve(mut capture: Arc<ScreenCapture>, input: Arc<InputInjector>, shutd
                 _ = shutdown.changed() => Some(SignalingEvent::Stop),
                 command = commands.recv(), if assist.is_some() => Some(SignalingEvent::AssistCommand(command.unwrap_or(assist::Command::Stop))),
                 _ = assist_tick.tick(), if assist.is_some() => Some(SignalingEvent::AssistTick),
+                Some(ice) = desktop_requested.recv() => Some(SignalingEvent::DesktopRequested(ice)),
                 event = rx_event.recv() => event,
             };
             let Some(evt) = event else { break; };
@@ -855,8 +858,8 @@ async fn serve(mut capture: Arc<ScreenCapture>, input: Arc<InputInjector>, shutd
                         ::webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState::Closed |
                         ::webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState::Failed)).map(|(id, _)| id.clone()).collect();
                     let mut desktop_ended = false;
-                    for id in ended { if let Some(s) = sessions.remove(&id) { desktop_ended |= s.desktop; let _ = s.pc.close().await; } }
-                    if desktop_ended && !sessions.values().any(|s| s.desktop) { if let Some(s) = assist.as_mut() { s.session_ended()?; } }
+                    for id in ended { if let Some(s) = sessions.remove(&id) { desktop_ended |= s.role.is_desktop(); let _ = s.pc.close().await; } }
+                    if desktop_ended && !sessions.values().any(|s| s.role.is_desktop()) { if let Some(s) = assist.as_mut() { s.session_ended()?; } }
                     if let Some(s) = assist.as_mut() { s.tick()?; }
                 }
                 SignalingEvent::Connected => {
@@ -943,8 +946,8 @@ async fn serve(mut capture: Arc<ScreenCapture>, input: Arc<InputInjector>, shutd
                         ::webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState::Closed |
                         ::webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState::Failed)).map(|(k, _)| k.clone()).collect();
                     let mut desktop_ended = false;
-                    for k in ended { if let Some(old) = sessions.remove(&k) { desktop_ended |= old.desktop; let _ = old.pc.close().await; } }
-                    if desktop_ended && !sessions.values().any(|s| s.desktop) { if let Some(s) = assist.as_mut() { s.session_ended()?; } }
+                    for k in ended { if let Some(old) = sessions.remove(&k) { desktop_ended |= old.role.is_desktop(); let _ = old.pc.close().await; } }
+                    if desktop_ended && !sessions.values().any(|s| s.role.is_desktop()) { if let Some(s) = assist.as_mut() { s.session_ended()?; } }
                     if let Some(s) = assist.as_mut() { s.tick()?; }
 
                     if soac.msg_type == "offer" {
@@ -960,6 +963,8 @@ async fn serve(mut capture: Arc<ScreenCapture>, input: Arc<InputInjector>, shutd
                         for id in superseded_peer_ids(&sessions, &ice, desktop) {
                             if let Some(old) = sessions.remove(&id) { let _ = old.pc.close().await; }
                         }
+                        // A tool peer promoted by capture_desktop stops its desktop but keeps its tools.
+                        if desktop { for old in sessions.values() { old.demote(); } }
                         if desktop && desktop_backend::headless() {
                             log::info!("Desktop connection ignored: headless mode (terminal, file transfer and port mapping remain available)");
                             continue;
@@ -981,7 +986,12 @@ async fn serve(mut capture: Arc<ScreenCapture>, input: Arc<InputInjector>, shutd
                         if desktop { if let Some(portal) = &portal_desktop { portal.begin(); } }
                         match WebRtcSession::new(session_capture, input.clone(), signaling_client.clone(), desktop, transport,
                             if desktop { portal_desktop.clone() } else { None }).await {
-                            Ok(sess) => {
+                            Ok(mut sess) => {
+                                sess.controller_mac = controller_mac;
+                                if !desktop {
+                                    let (requests, peer) = (desktop_requests.clone(), ice.clone());
+                                    sess.on_desktop_request(move || { let _ = requests.send(peer.clone()); });
+                                }
                                 // Register before awaiting SDP: shutdown can always close
                                 // a peer whose answer may already be queued to the viewer.
                                 sessions.insert(ice.clone(), sess);
@@ -1026,13 +1036,50 @@ async fn serve(mut capture: Arc<ScreenCapture>, input: Arc<InputInjector>, shutd
                     let mut desktop_ended = false;
                     for k in targets {
                         if let Some(old) = sessions.remove(&k) {
-                            desktop_ended |= old.desktop;
+                            desktop_ended |= old.role.is_desktop();
                             if let Err(e) = old.pc.close().await { log::warn!("Close released peer failed: {}", e); }
                         }
                     }
-                    if desktop_ended && !sessions.values().any(|s| s.desktop) { if let Some(s) = assist.as_mut() { s.session_ended()?; } }
+                    if desktop_ended && !sessions.values().any(|s| s.role.is_desktop()) { if let Some(s) = assist.as_mut() { s.session_ended()?; } }
                     if let Some(s) = assist.as_mut() { s.tick()?; }
                 }
+                SignalingEvent::DesktopRequested(ice) => {
+                    let mut sessions = webrtc_sessions.lock().await;
+                    let Some(peer) = sessions.get(&ice).filter(|peer| !peer.role.is_desktop()) else { continue; };
+                    if desktop_backend::headless() {
+                        log::info!("Desktop request on a tool connection ignored: headless mode");
+                        continue;
+                    }
+                    if portal_desktop.is_some() {
+                        log::info!("Desktop request on a tool connection not supported with the Portal backend; open the desktop before port mapping");
+                        continue;
+                    }
+                    let controller_mac = peer.controller_mac;
+                    let current = peer.desktop_capture();
+                    // Same recovery as a desktop offer: the display may have come up or been resized after this peer.
+                    let resized = current.init_error.is_none() && current.backend == desktop_backend::DesktopBackend::X11
+                        && tokio::task::spawn_blocking(crate::capture::x11_size).await.ok().flatten().is_some_and(|size| size != (current.width, current.height));
+                    if resized { log::info!("Desktop geometry changed since this connection opened; reinitializing capture"); }
+                    let screen = if current.init_error.is_some() || resized {
+                        match tokio::task::spawn_blocking(ScreenCapture::new).await {
+                            Ok(Ok(screen)) => { capture = Arc::new(screen); Some(capture.clone()) },
+                            Ok(Err(error)) => { log::info!("Desktop request on a tool connection ignored: {error}"); continue; },
+                            Err(_) => { log::warn!("Desktop initialization worker failed"); continue; },
+                        }
+                    } else { None };
+                    video::warm_auto_probe(); video::check_ffmpeg_usable();
+                    input.set_mac_controller(controller_mac);
+                    input.bind_portal(screen.as_ref().map_or_else(|| sessions[&ice].capture.portal.clone(), |screen| screen.portal.clone()));
+                    // One desktop at a time, as for a new desktop offer. Peers offered as desktops
+                    // close; a promoted tool peer only stops its desktop and keeps its tools.
+                    let superseded: Vec<String> = sessions.iter().filter(|(id, old)| id.as_str() != ice && old.role.is_desktop())
+                        .map(|(id, _)| id.clone()).collect();
+                    for id in superseded {
+                        if sessions[&id].desktop { if let Some(old) = sessions.remove(&id) { let _ = old.pc.close().await; } }
+                        else { sessions[&id].demote(); }
+                    }
+                    if let Some(peer) = sessions.get(&ice) { peer.promote(screen); }
+                }
                 SignalingEvent::Disconnected => {
                     capture_types.clear();
                     next_capture_type=None;

src/signaling.rs

@@ -46,6 +46,8 @@ pub enum SignalingEvent {
     Released(Value),
     Left(Value),
     Disconnected,
+    /// capture_desktop on this tool peer (ice_id); raised by the peer, not by signaling.
+    DesktopRequested(String),
 }
 
 pub struct SignalingClient {

src/system_actions.rs

@@ -12,6 +12,8 @@ pub const FF_SIMPLE_ACTION: u8 = 2;
 pub const FF_SYSTEM_METRICS: u8 = 2;
 
 pub const SHOW_DESKTOP: u64 = 2;
+/// Opens the desktop on an existing connection (the viewer's port mapping reuses its peer, #99).
+pub const CAPTURE_DESKTOP: u64 = 4;
 pub const REBOOT_SYSTEM: u64 = 9;
 pub const LOCK_SYSTEM: u64 = 10;
 pub const SHUTDOWN_SYSTEM: u64 = 12;