Fix Dosh stdio proxy for VS Code

src/bin/dosh-client.rs

@@ -1491,6 +1491,7 @@ async fn proxy_stdio_loop(mut transport: DoshTransport, stream_id: u64) -> Resul
                         for chunk in data.chunks {
                             stdout.write_all(&chunk).await?;
                         }
+                        stdout.flush().await?;
                     }
                     SessionEvent::Stream(TransportEvent::Close { stream_id: closed }) if closed == stream_id => {
                         stdout.flush().await?;
@@ -1529,22 +1530,29 @@ fn run_vscode_command(config: &dosh::config::ClientConfig, args: &Args) -> Resul
     if tokens.next().is_some() {
         return Err(anyhow!("dosh vscode accepts at most one remote path"));
     }
-    let alias = ensure_vscode_ssh_host(config, args, host)?;
-    println!("[ok] VS Code SSH host: {alias}");
-    println!("[ok] Dosh proxy: dosh proxy-stdio {host} 127.0.0.1 22");
+    let generated = ensure_vscode_ssh_host(config, args, host)?;
+    println!("[ok] VS Code SSH host: {}", generated.alias);
+    println!("[ok] SSH target: 127.0.0.1:{}", generated.target_ssh_port);
+    println!("[ok] Dosh proxy: {}", generated.proxy_command);
     if launch {
-        launch_vscode_remote(&alias, remote_path)?;
+        launch_vscode_remote(&generated.alias, remote_path)?;
     } else {
-        println!("Open in VS Code Remote-SSH as: {alias}");
+        println!("Open in VS Code Remote-SSH as: {}", generated.alias);
     }
     Ok(())
 }
 
+struct GeneratedVscodeHost {
+    alias: String,
+    proxy_command: String,
+    target_ssh_port: u16,
+}
+
 fn ensure_vscode_ssh_host(
     config: &dosh::config::ClientConfig,
     args: &Args,
     requested: &str,
-) -> Result<String> {
+) -> Result<GeneratedVscodeHost> {
     let hosts = load_hosts_config(None).unwrap_or_default();
     let host_config = hosts.hosts.get(requested).cloned().unwrap_or_default();
     let raw_server = host_config
@@ -1554,6 +1562,7 @@ fn ensure_vscode_ssh_host(
     let server = ssh_with_user(&raw_server, host_config.user.as_deref());
     let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port);
     let ssh_config = ssh_config(&server, ssh_port).unwrap_or_default();
+    let target_ssh_port = ssh_config.port.or(ssh_port).unwrap_or(22);
     let user = host_config
         .user
         .or(ssh_username(&server))
@@ -1563,11 +1572,22 @@ fn ensure_vscode_ssh_host(
         .ok()
         .map(|path| path.display().to_string())
         .unwrap_or_else(|| "dosh".to_string());
-    let mut proxy_command = shell_word(&exe);
+    let mut proxy_command = ssh_config_word(&exe);
+    if let Some(host) = args
+        .dosh_host
+        .clone()
+        .or_else(|| host_config.dosh_host.clone())
+        .or_else(|| config.dosh_host.clone())
+    {
+        proxy_command.push_str(&format!(" --dosh-host {}", ssh_config_word(&host)));
+    }
     if let Some(port) = args.dosh_port.or(host_config.port) {
         proxy_command.push_str(&format!(" --dosh-port {port}"));
     }
-    proxy_command.push_str(&format!(" proxy-stdio {} %h %p", shell_word(requested)));
+    proxy_command.push_str(&format!(
+        " proxy-stdio {} %h %p",
+        ssh_config_word(requested)
+    ));
     let ssh_dir = expand_tilde("~/.ssh");
     fs::create_dir_all(&ssh_dir).with_context(|| format!("create {}", ssh_dir.display()))?;
     let include_path = ssh_dir.join("config.dosh");
@@ -1577,7 +1597,7 @@ fn ensure_vscode_ssh_host(
     block.push_str(&format!("# BEGIN DOSH {alias}\n"));
     block.push_str(&format!("Host {alias}\n"));
     block.push_str("    HostName 127.0.0.1\n");
-    block.push_str("    Port 22\n");
+    block.push_str(&format!("    Port {target_ssh_port}\n"));
     block.push_str(&format!("    HostKeyAlias {requested}\n"));
     if let Some(user) = user {
         block.push_str(&format!("    User {user}\n"));
@@ -1588,7 +1608,11 @@ fn ensure_vscode_ssh_host(
     block.push_str(&format!("    ProxyCommand {proxy_command}\n"));
     block.push_str(&format!("# END DOSH {alias}\n"));
     upsert_managed_block(&include_path, &alias, &block)?;
-    Ok(alias)
+    Ok(GeneratedVscodeHost {
+        alias,
+        proxy_command,
+        target_ssh_port,
+    })
 }
 
 fn ensure_ssh_include(path: &Path, include: &str) -> Result<()> {
@@ -3145,6 +3169,14 @@ fn shell_word(value: &str) -> String {
     format!("'{}'", value.replace('\'', "'\\''"))
 }
 
+fn ssh_config_word(value: &str) -> String {
+    if cfg!(windows) {
+        format!("\"{}\"", value.replace('"', "\\\""))
+    } else {
+        shell_word(value)
+    }
+}
+
 fn local_bootstrap(
     session: &str,
     mode: &str,