From 943f2e6da80062ba2dc0eb228213b9723f3b2c6e Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:54:54 -0400 Subject: [PATCH] Try all resolved UDP peers for native auth --- src/bin/dosh-client.rs | 227 +++++++++++++++++++++++++++++++++++------ 1 file changed, 197 insertions(+), 30 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 7149da2..dc9015e 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -90,6 +90,13 @@ const AGENT_STREAM_SENTINEL: &str = "@dosh-agent"; const WINDOWS_OPENSSH_AGENT_PIPE: &str = r"\\.\pipe\openssh-ssh-agent"; type AgentWriter = Box; +struct NativeAuthAttempt { + socket: UdpSocket, + addr: SocketAddr, + frame: Frame, + cred: CachedCredential, +} + /// Current terminal size, with a sane fallback. /// /// `crossterm::size()` returns `Err` when stdout is not a TTY (piped), but it @@ -451,7 +458,8 @@ async fn main() -> Result<()> { }; let cache_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref()); - let target_udp_addr = resolve_addr(&target_udp_host, dosh_port)?; + let target_udp_addrs = resolve_addrs(&target_udp_host, dosh_port)?; + let target_udp_addr = first_resolved_addr(&target_udp_addrs, &target_udp_host, dosh_port)?; let socket = bind_udp_for_peer(target_udp_addr).await?; if let Some(mut cached) = credential.clone() { @@ -472,6 +480,7 @@ async fn main() -> Result<()> { return run_terminal( socket, cred, + Some(target_udp_addr), Some(frame), predict, predict_mode, @@ -510,6 +519,7 @@ async fn main() -> Result<()> { return run_terminal( socket, cred, + Some(target_udp_addr), Some(frame), predict, predict_mode, @@ -540,9 +550,14 @@ async fn main() -> Result<()> { } let cold_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref()); let native_start = Instant::now(); - match try_native_auth( - &socket, - target_udp_addr, + let requested_forwardings = forwarding_requests( + &local_forwards, + &remote_forwards, + &dynamic_forwards, + forward_agent, + ); + match try_native_auth_any_addr( + &target_udp_addrs, &config, &requested_server, &server, @@ -554,19 +569,19 @@ async fn main() -> Result<()> { &mode, cols, rows, - forwarding_requests( - &local_forwards, - &remote_forwards, - &dynamic_forwards, - forward_agent, - ), + requested_forwardings, resolved_ssh_config.as_ref(), cold_requested_env, true, ) .await { - Ok((frame, cred)) => { + Ok(NativeAuthAttempt { + socket, + addr, + frame, + cred, + }) => { log_timing(args.verbose, "native_auth", native_start.elapsed()); if allow_cache { save_cache(&cache_path, &cred)?; @@ -580,6 +595,7 @@ async fn main() -> Result<()> { return run_terminal( socket, cred, + Some(addr), Some(frame), predict, predict_mode, @@ -658,6 +674,7 @@ async fn main() -> Result<()> { run_terminal( socket, cred, + Some(target_udp_addr), Some(first), predict, predict_mode, @@ -1949,10 +1966,17 @@ async fn run_doctor_for_host( ) { println!("[warn] native endpoint: {warning}"); } - let udp_addr = match resolve_addr(&udp_host, dosh_port) { - Ok(addr) => { - println!("[ok] udp resolve: {addr}"); - addr + let udp_addrs = match resolve_addrs(&udp_host, dosh_port) { + Ok(addrs) => { + println!( + "[ok] udp resolve: {}", + addrs + .iter() + .map(ToString::to_string) + .collect::>() + .join(", ") + ); + addrs } Err(err) => { println!("[fail] udp resolve: {err:#}"); @@ -1963,10 +1987,8 @@ async fn run_doctor_for_host( "[info] forwarding requested by config: agent={} default_session={}", config.forward_agent, config.default_session ); - let socket = bind_udp_for_peer(udp_addr).await?; - match try_native_auth_check( - &socket, - udp_addr, + match try_native_auth_check_any_addr( + &udp_addrs, config, requested, &server, @@ -2400,8 +2422,7 @@ async fn run_proxy_stdio_command(config: &dosh::config::ClientConfig, args: &Arg 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 target_udp_addr = resolve_addr(&target_udp_host, dosh_port)?; - let socket = bind_udp_for_peer(target_udp_addr).await?; + let target_udp_addrs = resolve_addrs(&target_udp_host, dosh_port)?; let session = args .session .clone() @@ -2413,9 +2434,13 @@ async fn run_proxy_stdio_command(config: &dosh::config::ClientConfig, args: &Arg target_host: Some(target_host.clone()), target_port: Some(target_port), }]; - let (_frame, cred) = try_native_auth( - &socket, - target_udp_addr, + let NativeAuthAttempt { + socket, + addr, + frame: _frame, + cred, + } = try_native_auth_any_addr( + &target_udp_addrs, config, &requested_server, &server, @@ -2434,14 +2459,13 @@ async fn run_proxy_stdio_command(config: &dosh::config::ClientConfig, args: &Arg ) .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, + peer_addr: addr, initial_send_seq: 2, initial_ack: cred.last_rendered_seq, stream: TransportConfig::default(), @@ -5310,10 +5334,25 @@ fn glob_matches(pattern: &str, value: &str) -> bool { } fn resolve_addr(host: &str, port: u16) -> Result { - (host, port) + let addrs = resolve_addrs(host, port)?; + first_resolved_addr(&addrs, host, port) +} + +fn resolve_addrs(host: &str, port: u16) -> Result> { + let addrs = (host, port) .to_socket_addrs() .with_context(|| format!("resolve UDP target {host}:{port}"))? - .next() + .collect::>(); + if addrs.is_empty() { + return Err(anyhow!("no UDP address resolved for {host}:{port}")); + } + Ok(addrs) +} + +fn first_resolved_addr(addrs: &[SocketAddr], host: &str, port: u16) -> Result { + addrs + .first() + .copied() .ok_or_else(|| anyhow!("no UDP address resolved for {host}:{port}")) } @@ -5373,6 +5412,72 @@ async fn send_terminal_udp(socket: &UdpSocket, packet: &[u8], addr: SocketAddr) } } +#[allow(clippy::too_many_arguments)] +async fn try_native_auth_any_addr( + addrs: &[SocketAddr], + config: &dosh::config::ClientConfig, + requested_server: &str, + server: &str, + ssh_port: Option, + udp_host: &str, + udp_port: u16, + cli_identity: Option<&Path>, + session: &str, + mode: &str, + cols: u16, + rows: u16, + requested_forwardings: Vec, + ssh_config_hint: Option<&SshConfig>, + requested_env: Vec, + allow_passphrase_prompt: bool, +) -> Result { + let mut errors = Vec::new(); + for addr in addrs { + let socket = match bind_udp_for_peer(*addr).await { + Ok(socket) => socket, + Err(err) => { + errors.push(format!("{addr}: bind UDP socket: {err:#}")); + continue; + } + }; + match try_native_auth( + &socket, + *addr, + config, + requested_server, + server, + ssh_port, + udp_host, + udp_port, + cli_identity, + session, + mode, + cols, + rows, + requested_forwardings.clone(), + ssh_config_hint, + requested_env.clone(), + allow_passphrase_prompt, + ) + .await + { + Ok((frame, cred)) => { + return Ok(NativeAuthAttempt { + socket, + addr: *addr, + frame, + cred, + }); + } + Err(err) => errors.push(format!("{addr}: {err:#}")), + } + } + Err(anyhow!( + "native auth failed for all resolved UDP addresses: {}", + errors.join("; ") + )) +} + #[allow(clippy::too_many_arguments)] async fn try_native_auth( socket: &UdpSocket, @@ -5658,6 +5763,46 @@ async fn try_native_auth_check( protocol::from_body(&plain) } +async fn try_native_auth_check_any_addr( + addrs: &[SocketAddr], + config: &dosh::config::ClientConfig, + requested_server: &str, + server: &str, + ssh_port: Option, + cli_identity: Option<&Path>, + ssh_config_hint: Option<&SshConfig>, +) -> Result { + let mut errors = Vec::new(); + for addr in addrs { + let socket = match bind_udp_for_peer(*addr).await { + Ok(socket) => socket, + Err(err) => { + errors.push(format!("{addr}: bind UDP socket: {err:#}")); + continue; + } + }; + match try_native_auth_check( + &socket, + *addr, + config, + requested_server, + server, + ssh_port, + cli_identity, + ssh_config_hint, + ) + .await + { + Ok(check) => return Ok(check), + Err(err) => errors.push(format!("{addr}: {err:#}")), + } + } + Err(anyhow!( + "native auth check failed for all resolved UDP addresses: {}", + errors.join("; ") + )) +} + fn sign_native_user_auth( config: &dosh::config::ClientConfig, cli_identity: Option<&Path>, @@ -6207,6 +6352,7 @@ where async fn run_terminal( socket: UdpSocket, mut cred: CachedCredential, + initial_addr: Option, first_frame: Option, predict: bool, predict_mode: PredictMode, @@ -6223,7 +6369,10 @@ async fn run_terminal( } else { Some(RawMode::enter()?) }; - let mut addr = resolve_addr(&cred.udp_host, cred.udp_port)?; + let mut addr = match initial_addr { + Some(addr) => addr, + None => resolve_addr(&cred.udp_host, cred.udp_port)?, + }; let mut send_seq = 2u64; let mut last_packet_at = Instant::now(); let mut status_tick = tokio::time::interval(Duration::from_secs(1)); @@ -10411,7 +10560,7 @@ mod tests { SshPathTokenContext, StartupGateMode, StatusAction, TERMINAL_SNAPSHOT_RESET, UpdateOptions, UpdateRole, auth_allows, cache_key, cache_server_prefix, cleanup_stream_state, clear_cached_credentials, effective_update_artifact_tag, ensure_tui_safe_status_overlay, - expand_ssh_path_tokens, input_contains_focus_in, input_matches_escape, + expand_ssh_path_tokens, first_resolved_addr, input_contains_focus_in, input_matches_escape, is_local_status_target, is_resume_response_for_client, latest_release_download_url, load_first_native_identity_with_prompt, local_symlink_target_is_dir, local_username_from_env, native_proxy_udp_warning, newest_client_trace_path_from, @@ -10937,6 +11086,24 @@ mod tests { assert_eq!(ssh_destination_host("alice@example.com"), "example.com"); } + #[test] + fn first_resolved_addr_uses_dns_order_and_reports_empty_results() { + let addrs = [ + "127.0.0.1:50000".parse().unwrap(), + "[::1]:50000".parse().unwrap(), + ]; + assert_eq!( + first_resolved_addr(&addrs, "example.test", 50000).unwrap(), + addrs[0] + ); + assert!( + first_resolved_addr(&[], "example.test", 50000) + .unwrap_err() + .to_string() + .contains("no UDP address resolved for example.test:50000") + ); + } + #[test] fn selected_udp_host_supports_ssh_auto_and_localhost() { let ssh = SshConfig {