Try all resolved UDP peers in SDK client

src/client.rs

@@ -170,7 +170,7 @@ impl DoshClientBuilder {
             .udp_port
             .or(host_config.port)
             .unwrap_or(self.client.config.dosh_port);
-        let peer_addr = resolve_addr(&udp_host, udp_port)?;
+        let peer_addrs = resolve_addrs(&udp_host, udp_port)?;
         let timeout = self.timeout.unwrap_or_else(|| {
             Duration::from_millis(self.client.config.native_auth_timeout_ms.max(1))
         });
@@ -188,108 +188,32 @@ impl DoshClientBuilder {
                 })
             })
             .collect::<Result<Vec<_>>>()?;
-        let socket = bind_udp_for_peer(peer_addr).await?;
-        let (client_secret, client_public) = generate_native_ephemeral();
-        let hello = NativeClientHello {
-            protocol_version: native::NATIVE_PROTOCOL_VERSION,
-            client_random: crypto::random_32(),
-            client_ephemeral_public: client_public,
-            requested_host: self.host.clone(),
-            requested_user,
-            requested_session: session.clone(),
-            requested_mode: "forward-only".to_string(),
-            terminal_size: (80, 24),
-            supported_aead: vec!["chacha20poly1305".to_string()],
-            supported_user_key_algorithms: supported_user_key_algorithms(),
-            cached_host_key_fingerprint: None,
-            attach_ticket_envelope: None,
-            requested_env: self.env,
-        };
-        let packet = protocol::encode_plain(
-            PacketKind::NativeClientHello,
-            [0u8; 16],
-            1,
-            0,
-            &protocol::to_body(&NativeClientHelloBody {
-                hello: hello.clone(),
-            })?,
-        )?;
-        send_udp_retrying_transient(&socket, &packet, peer_addr, timeout).await?;
-
-        let mut buf = vec![0u8; 65535];
-        let (n, _) = recv_udp_retrying_transient(&socket, &mut buf, timeout).await?;
-        let packet = protocol::decode(&buf[..n])?;
-        if packet.header.kind != PacketKind::NativeServerHello {
-            if packet.header.kind == PacketKind::AttachReject {
-                let reject: AttachReject = protocol::from_body(&packet.body)?;
-                bail!("native auth rejected: {}", reject.reason);
-            }
-            bail!("native auth received unexpected server response");
-        }
-        let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?;
-        verify_server_hello(&hello, &server_hello.hello)?;
-        verify_or_trust_host(
-            &self.client.config,
-            &self.host,
-            &server_hello.hello.host_key,
-            self.trust_on_first_use,
-        )?;
-
-        let session_key = derive_native_session_key(
-            &client_secret,
-            server_hello.hello.server_ephemeral_public,
-            &hello,
-            &server_hello.hello,
-        )?;
-        let auth = sign_auth(
-            &self.client.config,
-            &host_config,
-            &hello,
-            &server_hello.hello,
-            requested_forwardings,
-            self.identity_files,
-            self.use_ssh_agent,
-        )?;
-        let mut pending_id = [0u8; 16];
-        pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]);
-        let auth_packet = protocol::encode_encrypted(
-            PacketKind::NativeUserAuth,
-            pending_id,
-            2,
-            1,
-            &session_key,
-            CLIENT_TO_SERVER,
-            &protocol::to_body(&NativeUserAuthBody { auth })?,
-        )?;
-        send_udp_retrying_transient(&socket, &auth_packet, peer_addr, timeout).await?;
-        let (n, _) = recv_udp_retrying_transient(&socket, &mut buf, timeout).await?;
-        let packet = protocol::decode(&buf[..n])?;
-        if packet.header.kind != PacketKind::NativeAuthOk {
-            if packet.header.kind == PacketKind::AttachReject {
-                let reject: AttachReject = protocol::from_body(&packet.body)?;
-                bail!("native auth rejected: {}", reject.reason);
-            }
-            bail!("native auth received unexpected auth response");
-        }
-        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)?;
-        let ok: NativeAuthOkBody = protocol::from_body(&plain)?;
-        let transport = DoshTransport::new_owned(
-            socket,
-            SessionTransportConfig {
-                role: SessionRole::Client,
-                conn_id: ok.ok.client_id,
-                session_key: ok.ok.session_key,
+        let mut errors = Vec::new();
+        for peer_addr in peer_addrs {
+            match connect_sdk_peer(
                 peer_addr,
-                initial_send_seq: 2,
-                initial_ack: ok.ok.initial_seq,
-                stream: TransportConfig::default(),
-            },
-        );
-        Ok(ConnectedDoshClient {
-            host: self.host,
-            session: ok.ok.session,
-            transport,
-        })
+                &self.client.config,
+                &host_config,
+                &self.host,
+                &requested_user,
+                &session,
+                &requested_forwardings,
+                &self.identity_files,
+                self.use_ssh_agent,
+                self.trust_on_first_use,
+                &self.env,
+                timeout,
+            )
+            .await
+            {
+                Ok(client) => return Ok(client),
+                Err(err) => errors.push(format!("{peer_addr}: {err:#}")),
+            }
+        }
+        Err(anyhow!(
+            "native auth failed for all resolved UDP addresses: {}",
+            errors.join("; ")
+        ))
     }
 }
 
@@ -305,6 +229,125 @@ impl ConnectedDoshClient {
     }
 }
 
+#[allow(clippy::too_many_arguments)]
+async fn connect_sdk_peer(
+    peer_addr: SocketAddr,
+    config: &ClientConfig,
+    host_config: &HostConfig,
+    host: &str,
+    requested_user: &str,
+    session: &str,
+    requested_forwardings: &[ForwardingRequest],
+    identity_files: &[PathBuf],
+    use_ssh_agent: Option<bool>,
+    trust_on_first_use: Option<bool>,
+    requested_env: &[EnvVar],
+    timeout: Duration,
+) -> Result<ConnectedDoshClient> {
+    let socket = bind_udp_for_peer(peer_addr).await?;
+    let (client_secret, client_public) = generate_native_ephemeral();
+    let hello = NativeClientHello {
+        protocol_version: native::NATIVE_PROTOCOL_VERSION,
+        client_random: crypto::random_32(),
+        client_ephemeral_public: client_public,
+        requested_host: host.to_string(),
+        requested_user: requested_user.to_string(),
+        requested_session: session.to_string(),
+        requested_mode: "forward-only".to_string(),
+        terminal_size: (80, 24),
+        supported_aead: vec!["chacha20poly1305".to_string()],
+        supported_user_key_algorithms: supported_user_key_algorithms(),
+        cached_host_key_fingerprint: None,
+        attach_ticket_envelope: None,
+        requested_env: requested_env.to_vec(),
+    };
+    let packet = protocol::encode_plain(
+        PacketKind::NativeClientHello,
+        [0u8; 16],
+        1,
+        0,
+        &protocol::to_body(&NativeClientHelloBody {
+            hello: hello.clone(),
+        })?,
+    )?;
+    send_udp_retrying_transient(&socket, &packet, peer_addr, timeout).await?;
+
+    let mut buf = vec![0u8; 65535];
+    let (n, _) = recv_udp_retrying_transient(&socket, &mut buf, timeout).await?;
+    let packet = protocol::decode(&buf[..n])?;
+    if packet.header.kind != PacketKind::NativeServerHello {
+        if packet.header.kind == PacketKind::AttachReject {
+            let reject: AttachReject = protocol::from_body(&packet.body)?;
+            bail!("native auth rejected: {}", reject.reason);
+        }
+        bail!("native auth received unexpected server response");
+    }
+    let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?;
+    verify_server_hello(&hello, &server_hello.hello)?;
+    verify_or_trust_host(
+        config,
+        host,
+        &server_hello.hello.host_key,
+        trust_on_first_use,
+    )?;
+
+    let session_key = derive_native_session_key(
+        &client_secret,
+        server_hello.hello.server_ephemeral_public,
+        &hello,
+        &server_hello.hello,
+    )?;
+    let auth = sign_auth(
+        config,
+        host_config,
+        &hello,
+        &server_hello.hello,
+        requested_forwardings.to_vec(),
+        identity_files.to_vec(),
+        use_ssh_agent,
+    )?;
+    let mut pending_id = [0u8; 16];
+    pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]);
+    let auth_packet = protocol::encode_encrypted(
+        PacketKind::NativeUserAuth,
+        pending_id,
+        2,
+        1,
+        &session_key,
+        CLIENT_TO_SERVER,
+        &protocol::to_body(&NativeUserAuthBody { auth })?,
+    )?;
+    send_udp_retrying_transient(&socket, &auth_packet, peer_addr, timeout).await?;
+    let (n, _) = recv_udp_retrying_transient(&socket, &mut buf, timeout).await?;
+    let packet = protocol::decode(&buf[..n])?;
+    if packet.header.kind != PacketKind::NativeAuthOk {
+        if packet.header.kind == PacketKind::AttachReject {
+            let reject: AttachReject = protocol::from_body(&packet.body)?;
+            bail!("native auth rejected: {}", reject.reason);
+        }
+        bail!("native auth received unexpected auth response");
+    }
+    let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)?;
+    let ok: NativeAuthOkBody = protocol::from_body(&plain)?;
+    let transport = DoshTransport::new_owned(
+        socket,
+        SessionTransportConfig {
+            role: SessionRole::Client,
+            conn_id: ok.ok.client_id,
+            session_key: ok.ok.session_key,
+            peer_addr,
+            initial_send_seq: 2,
+            initial_ack: ok.ok.initial_seq,
+            stream: TransportConfig::default(),
+        },
+    );
+    Ok(ConnectedDoshClient {
+        host: host.to_string(),
+        session: ok.ok.session,
+        transport,
+    })
+}
+
 fn verify_or_trust_host(
     config: &ClientConfig,
     host: &str,
@@ -383,11 +426,22 @@ fn default_identity_paths() -> Vec<PathBuf> {
         .collect()
 }
 
-fn resolve_addr(host: &str, port: u16) -> Result<SocketAddr> {
-    (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)
+}
+
+#[cfg(test)]
+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}"))
 }
 
@@ -471,6 +525,24 @@ mod tests {
         );
     }
 
+    #[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 default_identity_paths_are_expanded() {
         assert!(