fix(desktop): open Local Control on the active session

packages/desktop-shell/src-tauri/src/local_control.rs

@@ -52,7 +52,11 @@ pub struct LocalControlSession {
 }
 
 impl LocalControlSession {
-    pub fn start(runtime_url: &Url, runtime_token: &str) -> Result<Self, String> {
+    pub fn start(
+        runtime_url: &Url,
+        runtime_token: &str,
+        current_url: &Url,
+    ) -> Result<Self, String> {
         let target = runtime_socket_addr(runtime_url)?;
         let lan_ip = primary_lan_ipv4()?;
         let listener = TcpListener::bind((lan_ip, 0))
@@ -66,7 +70,7 @@ impl LocalControlSession {
             .port();
         let public_origin = format!("http://{lan_ip}:{port}");
         let pair_token = random_token();
-        let url = format!("{public_origin}/#token={pair_token}");
+        let url = local_control_url(current_url, lan_ip, port, &pair_token)?;
         let qr_svg = QrCode::new(url.as_bytes())
             .map_err(|error| format!("Failed to generate Local Control QR code: {error}"))?
             .render::<svg::Color>()
@@ -303,7 +307,19 @@ fn rewrite_request(
     pair_token: &str,
     runtime_token: &str,
 ) -> Result<Vec<u8>, u16> {
-    let header = std::str::from_utf8(&request[..header_end]).map_err(|_| 400_u16)?;
+    let header_bytes = &request[..header_end];
+    if header_bytes
+        .iter()
+        .enumerate()
+        .any(|(index, &byte)| match byte {
+            b'\r' => header_bytes.get(index + 1) != Some(&b'\n'),
+            b'\n' => index == 0 || header_bytes[index - 1] != b'\r',
+            _ => false,
+        })
+    {
+        return Err(400);
+    }
+    let header = std::str::from_utf8(header_bytes).map_err(|_| 400_u16)?;
     let public_authority = public_origin.strip_prefix("http://").ok_or(500_u16)?;
     let pair_protocol = format!("qwen-bearer.{}", URL_SAFE_NO_PAD.encode(pair_token));
     let runtime_protocol = format!("qwen-bearer.{}", URL_SAFE_NO_PAD.encode(runtime_token));
@@ -402,15 +418,39 @@ fn runtime_socket_addr(url: &Url) -> Result<SocketAddr, String> {
     Ok(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port))
 }
 
+fn local_control_url(
+    current_url: &Url,
+    lan_ip: Ipv4Addr,
+    port: u16,
+    pair_token: &str,
+) -> Result<String, String> {
+    let workspace = current_url
+        .query_pairs()
+        .find(|(name, value)| name == "workspace" && !value.is_empty())
+        .map(|(_, value)| value.into_owned());
+    let mut url = current_url.clone();
+    url.set_host(Some(&lan_ip.to_string()))
+        .map_err(|error| format!("Failed to construct the Local Control URL: {error}"))?;
+    url.set_port(Some(port))
+        .map_err(|_| "Failed to construct the Local Control URL.".to_string())?;
+    url.set_query(None);
+    if let Some(workspace) = workspace {
+        url.query_pairs_mut().append_pair("workspace", &workspace);
+    }
+    url.set_fragment(Some(&format!("token={pair_token}")));
+    Ok(url.into())
+}
+
 fn primary_lan_ipv4() -> Result<Ipv4Addr, String> {
-    let routed = routed_ipv4().ok();
-    let interfaces = match NetworkInterface::show() {
-        Ok(interfaces) => interfaces,
-        Err(_) => {
-            return routed
-                .ok_or_else(|| "Local Control could not find a usable IPv4 network.".to_string())
-        }
-    };
+    select_lan_ipv4(routed_ipv4().ok(), NetworkInterface::show().ok())
+}
+
+fn select_lan_ipv4(
+    routed: Option<Ipv4Addr>,
+    interfaces: Option<Vec<NetworkInterface>>,
+) -> Result<Ipv4Addr, String> {
+    let interfaces =
+        interfaces.ok_or_else(|| "Local Control could not inspect IPv4 networks.".to_string())?;
     let physical = interfaces
         .into_iter()
         .filter(|interface| {
@@ -441,10 +481,10 @@ fn choose_lan_ipv4(
 ) -> Result<Ipv4Addr, String> {
     physical.sort_unstable();
     physical.dedup();
+    physical.retain(|address| address.is_private() || address.is_link_local());
     if let Some(routed) = routed.filter(|address| physical.contains(address)) {
         return Ok(routed);
     }
-    physical.retain(|address| address.is_private() || address.is_link_local());
     match physical.as_slice() {
         [address] => Ok(*address),
         [] => Err("Local Control could not find a usable IPv4 network.".to_string()),
@@ -526,8 +566,8 @@ fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
 #[cfg(test)]
 mod tests {
     use super::{
-        choose_lan_ipv4, find_header_end, rewrite_request, runtime_socket_addr, spawn_proxy,
-        Connections,
+        choose_lan_ipv4, find_header_end, local_control_url, rewrite_request, runtime_socket_addr,
+        select_lan_ipv4, spawn_proxy, Connections,
     };
     use std::collections::HashMap;
     use std::io::{Read, Write};
@@ -541,7 +581,7 @@ mod tests {
     use url::Url;
 
     #[test]
-    fn prefers_the_physical_lan_over_a_vpn_route() {
+    fn selects_only_a_private_physical_lan() {
         assert_eq!(
             choose_lan_ipv4(
                 Some("10.8.0.2".parse::<Ipv4Addr>().expect("VPN address")),
@@ -552,6 +592,63 @@ mod tests {
                 .parse::<Ipv4Addr>()
                 .expect("expected address"),
         );
+        assert!(choose_lan_ipv4(
+            Some("203.0.113.10".parse().expect("public route")),
+            vec!["203.0.113.10".parse().expect("public interface")],
+        )
+        .is_err());
+        let routed = Ipv4Addr::new(192, 168, 1, 20);
+        assert_eq!(
+            choose_lan_ipv4(Some(routed), vec![routed, Ipv4Addr::new(192, 168, 2, 5)],)
+                .expect("routed LAN"),
+            routed,
+        );
+    }
+
+    #[test]
+    fn rejects_unverified_networks_when_interface_enumeration_fails() {
+        let routed = Ipv4Addr::new(192, 168, 1, 20);
+        assert!(select_lan_ipv4(Some(routed), None).is_err());
+    }
+
+    #[test]
+    fn shares_only_the_current_local_session() {
+        let url = local_control_url(
+            &Url::parse(
+                "http://127.0.0.1:4170/session/a%20b?token=runtime&workspace=work%2Ftree&theme=light#old",
+            )
+            .expect("current URL"),
+            "192.168.1.20".parse().expect("LAN address"),
+            49152,
+            "pair-token",
+        )
+        .expect("Local Control URL");
+        assert_eq!(
+            url,
+            "http://192.168.1.20:49152/session/a%20b?workspace=work%2Ftree#token=pair-token"
+        );
+
+        let url = local_control_url(
+            &Url::parse("http://127.0.0.1:4170/").expect("runtime URL"),
+            "192.168.1.20".parse().expect("LAN address"),
+            49152,
+            "pair-token",
+        )
+        .expect("Local Control URL");
+        assert_eq!(url, "http://192.168.1.20:49152/#token=pair-token");
+
+        let url = local_control_url(
+            &Url::parse("http://127.0.0.1:4170/session/x?workspace=")
+                .expect("current URL"),
+            "192.168.1.20".parse().expect("LAN address"),
+            49152,
+            "pair-token",
+        )
+        .expect("Local Control URL");
+        assert_eq!(
+            url,
+            "http://192.168.1.20:49152/session/x#token=pair-token"
+        );
     }
 
     #[test]
@@ -643,6 +740,35 @@ mod tests {
             Err(403),
         );
 
+        let request = b"GET / HTTP/1.1\nHost: 127.0.0.1:4170\n\n\r\n\r\n";
+        assert_eq!(
+            rewrite_request(
+                request,
+                find_header_end(request).expect("header"),
+                "http://192.168.1.10:49152",
+                "http://127.0.0.1:4170",
+                "127.0.0.1:4170",
+                "pair-token",
+                "runtime-token",
+            ),
+            Err(400),
+        );
+
+        let request =
+            b"GET / HTTP/1.1\r\nHost: 192.168.1.10:49152\r\nX-Inject: a\r\r\n\r\n";
+        assert_eq!(
+            rewrite_request(
+                request,
+                find_header_end(request).expect("header"),
+                "http://192.168.1.10:49152",
+                "http://127.0.0.1:4170",
+                "127.0.0.1:4170",
+                "pair-token",
+                "runtime-token",
+            ),
+            Err(400),
+        );
+
         let pair = URL_SAFE_NO_PAD.encode("pair-token");
         let runtime = URL_SAFE_NO_PAD.encode("runtime-token");
         let request = format!(

packages/desktop-shell/src-tauri/src/main.rs

@@ -304,7 +304,12 @@ fn enable_local_control(
         .as_ref()
         .map(|runtime| (runtime.base_url().clone(), runtime.token().to_string()))
         .ok_or_else(|| "Start a Desktop workspace before enabling Local Control.".to_string())?;
-    let session = LocalControlSession::start(&runtime_url, &runtime_token)?;
+    let current_url = app
+        .get_webview_window("main")
+        .and_then(|window| window.url().ok())
+        .filter(|url| is_same_origin(url, &runtime_url))
+        .unwrap_or_else(|| runtime_url.clone());
+    let session = LocalControlSession::start(&runtime_url, &runtime_token, &current_url)?;
     let info = session.info();
     *local_control = Some(session);
     set_local_control_menu_state(&app, true);