diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 20d35aa..51a8fbe 100644 --- a/src/bin/dosh-client.rs +++ b/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, } +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::() + .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::>(); + 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 { + 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 { let deadline = Instant::now() + Duration::from_secs(10); @@ -3174,6 +3509,7 @@ async fn try_native_auth( requested_forwardings: Vec, ssh_config_hint: Option<&SshConfig>, requested_env: Vec, + 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, + allow_passphrase_prompt: bool, ) -> Result { 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 { + 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( 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";