Add Dosh stdio proxy for SSH streams

src/bin/dosh-client.rs

@@ -32,6 +32,10 @@ use dosh::protocol::{
     TicketAttachEnvelope, TicketAttachOkEnvelope,
 };
 use dosh::ssh_agent;
+use dosh::transport::{
+    DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig,
+    TransportEvent,
+};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 use std::collections::BTreeMap;
@@ -258,6 +262,12 @@ async fn main() -> Result<()> {
     if args.server.as_deref() == Some("selftest") {
         return run_selftest_command(&config, &args).await;
     }
+    if args.server.as_deref() == Some("proxy-stdio") {
+        return run_proxy_stdio_command(&config, &args).await;
+    }
+    if args.server.as_deref() == Some("vscode") {
+        return run_vscode_command(&config, &args);
+    }
     if matches!(args.server.as_deref(), Some("recover" | "repair")) {
         return run_recover_command(&config, &args).await;
     }
@@ -512,6 +522,7 @@ async fn main() -> Result<()> {
             ),
             resolved_ssh_config.as_ref(),
             cold_requested_env,
+            true,
         )
         .await
         {
@@ -1349,6 +1360,330 @@ struct ExecServiceClient {
     pending: VecDeque<ExecResponse>,
 }
 
+async fn run_proxy_stdio_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
+    if args.command.len() != 3 {
+        return Err(anyhow!(
+            "usage: dosh proxy-stdio HOST TARGET_HOST TARGET_PORT"
+        ));
+    }
+    anyhow::ensure!(
+        !args.local_auth,
+        "proxy-stdio requires native auth; local auth cannot authenticate SSH ProxyCommand streams"
+    );
+    let requested_server = args.command[0].clone();
+    let target_host = args.command[1].clone();
+    let target_port = args.command[2]
+        .parse::<u16>()
+        .with_context(|| format!("invalid proxy target port {:?}", args.command[2]))?;
+    anyhow::ensure!(
+        valid_forward_host(&target_host),
+        "invalid proxy target host"
+    );
+    anyhow::ensure!(target_port != 0, "proxy target port cannot be 0");
+
+    let hosts = load_hosts_config(None).unwrap_or_default();
+    let host = hosts
+        .hosts
+        .get(&requested_server)
+        .cloned()
+        .unwrap_or_default();
+    let raw_server = host.ssh.clone().unwrap_or_else(|| requested_server.clone());
+    let server = ssh_with_user(&raw_server, host.user.as_deref());
+    let ssh_port = args.ssh_port.or(host.ssh_port).or(config.ssh_port);
+    let resolved_ssh_config = ssh_config(&server, ssh_port).unwrap_or_default();
+    let requested_udp_host = args
+        .dosh_host
+        .clone()
+        .or_else(|| host.dosh_host.clone())
+        .or_else(|| config.dosh_host.clone());
+    let target_udp_host =
+        selected_udp_host(requested_udp_host.as_deref(), &server, &resolved_ssh_config)?;
+    let dosh_port = args.dosh_port.or(host.port).unwrap_or(config.dosh_port);
+    let socket = UdpSocket::bind("0.0.0.0:0").await?;
+    let session = args
+        .session
+        .clone()
+        .unwrap_or_else(protocol::generate_implicit_session_name);
+    let requested_forwardings = vec![ForwardingRequest {
+        kind: ForwardingKind::Local,
+        bind_host: None,
+        listen_port: 0,
+        target_host: Some(target_host.clone()),
+        target_port: Some(target_port),
+    }];
+    let (_frame, cred) = try_native_auth(
+        &socket,
+        config,
+        &requested_server,
+        &server,
+        ssh_port,
+        &target_udp_host,
+        dosh_port,
+        args.ssh_key.as_deref(),
+        &session,
+        "forward-only",
+        80,
+        24,
+        requested_forwardings,
+        Some(&resolved_ssh_config),
+        Vec::new(),
+        false,
+    )
+    .await?;
+
+    let peer_addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
+    let mut transport = DoshTransport::new_owned(
+        socket,
+        SessionTransportConfig {
+            role: SessionRole::Client,
+            conn_id: cred.client_id,
+            session_key: cred.session_key,
+            peer_addr,
+            initial_send_seq: 2,
+            initial_ack: cred.last_rendered_seq,
+            stream: TransportConfig::default(),
+        },
+    );
+    let stream_id = transport.open_target(target_host, target_port).await?;
+    proxy_stdio_loop(transport, stream_id).await
+}
+
+async fn proxy_stdio_loop(mut transport: DoshTransport, stream_id: u64) -> Result<()> {
+    let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
+    std::thread::Builder::new()
+        .name("dosh-proxy-stdin".to_string())
+        .spawn(move || {
+            let mut stdin = std::io::stdin();
+            let mut buf = [0u8; 16 * 1024];
+            loop {
+                match stdin.read(&mut buf) {
+                    Ok(0) => break,
+                    Ok(n) => {
+                        if stdin_tx.send(buf[..n].to_vec()).is_err() {
+                            break;
+                        }
+                    }
+                    Err(_) => break,
+                }
+            }
+        })?;
+    let mut stdout = tokio::io::stdout();
+    let mut maintenance = tokio::time::interval(Duration::from_millis(50));
+    let mut stdin_closed = false;
+    loop {
+        tokio::select! {
+            stdin_msg = stdin_rx.recv(), if !stdin_closed => {
+                match stdin_msg {
+                    Some(bytes) => transport.send(stream_id, bytes).await?,
+                    None => {
+                        stdin_closed = true;
+                        let _ = transport.close(stream_id).await;
+                    }
+                }
+            }
+            event = transport.recv() => {
+                match event? {
+                    SessionEvent::Stream(TransportEvent::OpenOk { stream_id: opened, .. }) if opened == stream_id => {}
+                    SessionEvent::Stream(TransportEvent::OpenReject { stream_id: rejected, reason }) if rejected == stream_id => {
+                        return Err(anyhow!("Dosh proxy stream rejected: {reason}"));
+                    }
+                    SessionEvent::Stream(TransportEvent::Data(data)) if data.stream_id == stream_id => {
+                        for chunk in data.chunks {
+                            stdout.write_all(&chunk).await?;
+                        }
+                    }
+                    SessionEvent::Stream(TransportEvent::Close { stream_id: closed }) if closed == stream_id => {
+                        stdout.flush().await?;
+                        return Ok(());
+                    }
+                    SessionEvent::Stream(_) | SessionEvent::Ping | SessionEvent::Pong | SessionEvent::Ignored => {}
+                }
+            }
+            _ = maintenance.tick() => {
+                transport.maintenance().await?;
+            }
+        }
+    }
+}
+
+fn run_vscode_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
+    let mut tokens = args.command.iter();
+    let first = tokens
+        .next()
+        .ok_or_else(|| anyhow!("usage: dosh vscode [setup|open] HOST [REMOTE_PATH]"))?;
+    let (launch, host, remote_path) = match first.as_str() {
+        "setup" => {
+            let host = tokens
+                .next()
+                .ok_or_else(|| anyhow!("usage: dosh vscode setup HOST"))?;
+            (false, host.as_str(), None)
+        }
+        "open" => {
+            let host = tokens
+                .next()
+                .ok_or_else(|| anyhow!("usage: dosh vscode open HOST [REMOTE_PATH]"))?;
+            (true, host.as_str(), tokens.next().map(String::as_str))
+        }
+        host => (true, host, tokens.next().map(String::as_str)),
+    };
+    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");
+    if launch {
+        launch_vscode_remote(&alias, remote_path)?;
+    } else {
+        println!("Open in VS Code Remote-SSH as: {alias}");
+    }
+    Ok(())
+}
+
+fn ensure_vscode_ssh_host(
+    config: &dosh::config::ClientConfig,
+    args: &Args,
+    requested: &str,
+) -> Result<String> {
+    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
+        .ssh
+        .clone()
+        .unwrap_or_else(|| requested.to_string());
+    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 user = host_config
+        .user
+        .or(ssh_username(&server))
+        .or(ssh_config.user);
+    let alias = format!("dosh-{}", vscode_safe_alias(requested));
+    let exe = std::env::current_exe()
+        .ok()
+        .map(|path| path.display().to_string())
+        .unwrap_or_else(|| "dosh".to_string());
+    let mut proxy_command = shell_word(&exe);
+    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)));
+    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");
+    let main_config = ssh_dir.join("config");
+    ensure_ssh_include(&main_config, "config.dosh")?;
+    let mut block = String::new();
+    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!("    HostKeyAlias {requested}\n"));
+    if let Some(user) = user {
+        block.push_str(&format!("    User {user}\n"));
+    }
+    block.push_str("    ClearAllForwardings yes\n");
+    block.push_str("    ServerAliveInterval 15\n");
+    block.push_str("    ServerAliveCountMax 3\n");
+    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)
+}
+
+fn ensure_ssh_include(path: &Path, include: &str) -> Result<()> {
+    let raw = fs::read_to_string(path).unwrap_or_default();
+    if raw.lines().any(|line| {
+        line.trim()
+            .eq_ignore_ascii_case(&format!("include {include}"))
+    }) {
+        return Ok(());
+    }
+    let mut next = String::new();
+    next.push_str(&format!("Include {include}\n"));
+    if !raw.is_empty() {
+        next.push('\n');
+        next.push_str(&raw);
+    }
+    fs::write(path, next).with_context(|| format!("write {}", path.display()))
+}
+
+fn upsert_managed_block(path: &Path, alias: &str, block: &str) -> Result<()> {
+    let raw = fs::read_to_string(path).unwrap_or_default();
+    let begin = format!("# BEGIN DOSH {alias}");
+    let end = format!("# END DOSH {alias}");
+    let mut out = String::new();
+    let mut skipping = false;
+    let mut replaced = false;
+    for line in raw.lines() {
+        if line.trim() == begin {
+            if !out.ends_with('\n') && !out.is_empty() {
+                out.push('\n');
+            }
+            out.push_str(block);
+            replaced = true;
+            skipping = true;
+            continue;
+        }
+        if skipping {
+            if line.trim() == end {
+                skipping = false;
+            }
+            continue;
+        }
+        out.push_str(line);
+        out.push('\n');
+    }
+    if !replaced {
+        if !out.trim().is_empty() && !out.ends_with("\n\n") {
+            out.push('\n');
+        }
+        out.push_str(block);
+    }
+    fs::write(path, out).with_context(|| format!("write {}", path.display()))
+}
+
+fn vscode_safe_alias(value: &str) -> String {
+    let mut out = String::new();
+    for byte in value.bytes() {
+        if byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_' {
+            out.push(byte as char);
+        } else {
+            out.push('-');
+        }
+    }
+    let out = out.trim_matches('-').to_string();
+    if out.is_empty() {
+        "host".to_string()
+    } else {
+        out
+    }
+}
+
+fn launch_vscode_remote(alias: &str, remote_path: Option<&str>) -> Result<()> {
+    let remote = format!("ssh-remote+{alias}");
+    let mut command = Command::new("code");
+    command.arg("--remote").arg(&remote);
+    if let Some(path) = remote_path {
+        command.arg(path);
+    }
+    match command.status() {
+        Ok(status) if status.success() => Ok(()),
+        Ok(status) => Err(anyhow!("code exited with status {status}")),
+        Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
+            println!(
+                "Run: code --remote {}{}",
+                shell_word(&remote),
+                remote_path
+                    .map(|path| format!(" {}", shell_word(path)))
+                    .unwrap_or_default()
+            );
+            Ok(())
+        }
+        Err(err) => Err(err).context("launch VS Code"),
+    }
+}
+
 impl ExecServiceClient {
     fn connect(port: u16) -> Result<Self> {
         let deadline = Instant::now() + Duration::from_secs(10);
@@ -3174,6 +3509,7 @@ async fn try_native_auth(
     requested_forwardings: Vec<ForwardingRequest>,
     ssh_config_hint: Option<&SshConfig>,
     requested_env: Vec<EnvVar>,
+    allow_passphrase_prompt: bool,
 ) -> Result<(Frame, CachedCredential)> {
     let addr = resolve_addr(udp_host, udp_port)?;
     let owned_ssh_config;
@@ -3267,6 +3603,7 @@ async fn try_native_auth(
         &hello,
         &server_hello.hello,
         requested_forwardings,
+        allow_passphrase_prompt,
     )?;
     let mut pending_id = [0u8; 16];
     pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]);
@@ -3424,6 +3761,7 @@ async fn try_native_auth_check(
         &hello,
         &server_hello.hello,
         Vec::new(),
+        true,
     )?;
     let mut pending_id = [0u8; 16];
     pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]);
@@ -3462,6 +3800,7 @@ fn sign_native_user_auth(
     hello: &dosh::native::NativeClientHello,
     server_hello: &dosh::native::NativeServerHello,
     requested_forwardings: Vec<ForwardingRequest>,
+    allow_passphrase_prompt: bool,
 ) -> Result<dosh::native::NativeUserAuth> {
     let mut errors = Vec::new();
     if config.use_ssh_agent && !ssh_config.identities_only {
@@ -3477,7 +3816,12 @@ fn sign_native_user_auth(
         errors.push("ssh-agent: skipped because SSH config sets IdentitiesOnly=yes".to_string());
     }
 
-    match load_first_native_identity(config, cli_identity, ssh_config) {
+    let identity = if allow_passphrase_prompt {
+        load_first_native_identity(config, cli_identity, ssh_config)
+    } else {
+        load_first_native_identity_noninteractive(config, cli_identity, ssh_config)
+    };
+    match identity {
         Ok(identity) => {
             sign_user_auth_with_private_key(&identity, hello, server_hello, requested_forwardings)
         }
@@ -3504,6 +3848,19 @@ fn load_first_native_identity(
     )
 }
 
+fn load_first_native_identity_noninteractive(
+    config: &dosh::config::ClientConfig,
+    cli_identity: Option<&Path>,
+    ssh_config: &SshConfig,
+) -> Result<ssh_key::PrivateKey> {
+    load_first_native_identity_with_prompt(config, cli_identity, ssh_config, |path| {
+        Err(anyhow!(
+            "{} is encrypted; unlock it in ssh-agent or use an unencrypted key for proxy-stdio",
+            path.display()
+        ))
+    })
+}
+
 fn load_first_native_identity_with_prompt<F>(
     config: &dosh::config::ClientConfig,
     cli_identity: Option<&Path>,
@@ -6629,7 +6986,8 @@ mod tests {
         server_version_mismatch, should_hold_during_startup_gate, should_hold_post_submit_input,
         split_after_command_submit, ssh_destination_host, ssh_username, ssh_with_user,
         startup_command, status_ssh_target, strip_stale_mouse_reports, toml_bare_key_or_quoted,
-        update_check_requested, update_version_status, valid_forward_host,
+        update_check_requested, update_version_status, upsert_managed_block, valid_forward_host,
+        vscode_safe_alias,
     };
     use dosh::config::{ClientConfig, CommandExtension, HostConfig};
     use dosh::native::EnvVar;
@@ -6674,6 +7032,35 @@ mod tests {
         }
     }
 
+    #[test]
+    fn vscode_alias_sanitizes_host_names() {
+        assert_eq!(vscode_safe_alias("palav"), "palav");
+        assert_eq!(vscode_safe_alias("user@example.com"), "user-example-com");
+        assert_eq!(vscode_safe_alias(":///"), "host");
+    }
+
+    #[test]
+    fn managed_ssh_block_is_replaced_not_duplicated() {
+        let dir = tempfile::tempdir().unwrap();
+        let path = dir.path().join("config.dosh");
+        upsert_managed_block(
+            &path,
+            "dosh-palav",
+            "# BEGIN DOSH dosh-palav\nHost dosh-palav\n    HostName 127.0.0.1\n# END DOSH dosh-palav\n",
+        )
+        .unwrap();
+        upsert_managed_block(
+            &path,
+            "dosh-palav",
+            "# BEGIN DOSH dosh-palav\nHost dosh-palav\n    HostName localhost\n# END DOSH dosh-palav\n",
+        )
+        .unwrap();
+        let raw = fs::read_to_string(path).unwrap();
+        assert_eq!(raw.matches("# BEGIN DOSH dosh-palav").count(), 1);
+        assert!(raw.contains("HostName localhost"));
+        assert!(!raw.contains("HostName 127.0.0.1"));
+    }
+
     #[test]
     fn parses_ssh_config_hostname() {
         let raw = "user alice\nhostname server.example.com\nport 22\n";