From 4dc57400f4e7e63f15a3da36fcb9da2b3cd1c290 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:59:50 -0400 Subject: [PATCH] Try all resolved UDP peers in SDK client --- src/client.rs | 282 +++++++++++++++++++++++++++++++------------------- 1 file changed, 177 insertions(+), 105 deletions(-) diff --git a/src/client.rs b/src/client.rs index bc23084..a85dfd7 100644 --- a/src/client.rs +++ b/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::>>()?; - 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, + trust_on_first_use: Option, + requested_env: &[EnvVar], + timeout: Duration, +) -> Result { + 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 { .collect() } -fn resolve_addr(host: &str, port: u16) -> Result { - (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) +} + +#[cfg(test)] +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}")) } @@ -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!(