Try all resolved UDP peers for native auth

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<dyn AsyncWrite + Unpin + Send>;
 
+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(
+        let requested_forwardings = forwarding_requests(
-            &socket,
+            &local_forwards,
-            target_udp_addr,
+            &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(
+            requested_forwardings,
-                &local_forwards,
-                &remote_forwards,
-                &dynamic_forwards,
-                forward_agent,
-            ),
             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) {
+    let udp_addrs = match resolve_addrs(&udp_host, dosh_port) {
-        Ok(addr) => {
+        Ok(addrs) => {
-            println!("[ok] udp resolve: {addr}");
+            println!(
-            addr
+                "[ok] udp resolve: {}",
+                addrs
+                    .iter()
+                    .map(ToString::to_string)
+                    .collect::<Vec<_>>()
+                    .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_any_addr(
-    match try_native_auth_check(
+        &udp_addrs,
-        &socket,
-        udp_addr,
         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 target_udp_addrs = resolve_addrs(&target_udp_host, dosh_port)?;
-    let socket = bind_udp_for_peer(target_udp_addr).await?;
     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(
+    let NativeAuthAttempt {
-        &socket,
+        socket,
-        target_udp_addr,
+        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<SocketAddr> {
-    (host, port)
+    let addrs = resolve_addrs(host, port)?;
+    first_resolved_addr(&addrs, host, port)
+}
+
+fn resolve_addrs(host: &str, port: u16) -> Result<Vec<SocketAddr>> {
+    let addrs = (host, port)
         .to_socket_addrs()
         .with_context(|| format!("resolve UDP target {host}:{port}"))?
-        .next()
+        .collect::<Vec<_>>();
+    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<SocketAddr> {
+    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<u16>,
+    udp_host: &str,
+    udp_port: u16,
+    cli_identity: Option<&Path>,
+    session: &str,
+    mode: &str,
+    cols: u16,
+    rows: u16,
+    requested_forwardings: Vec<ForwardingRequest>,
+    ssh_config_hint: Option<&SshConfig>,
+    requested_env: Vec<EnvVar>,
+    allow_passphrase_prompt: bool,
+) -> Result<NativeAuthAttempt> {
+    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<u16>,
+    cli_identity: Option<&Path>,
+    ssh_config_hint: Option<&SshConfig>,
+) -> Result<NativeAuthCheckOkBody> {
+    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<SocketAddr>,
     first_frame: Option<Frame>,
     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 {