diff --git a/rama-tcp/src/client/connect.rs b/rama-tcp/src/client/connect.rs index c2e00f7a..8c2708eb 100644 --- a/rama-tcp/src/client/connect.rs +++ b/rama-tcp/src/client/connect.rs @@ -61,10 +61,7 @@ implConnectIpMode{ } } -/* -TODO - Using the wrapper types -in the tcp_connect fn, update the dns type to use the DnsResolveIpMode -*/ + /// Trait used internally by [`tcp_connect`] and the `TcpConnector` /// to actually establish the [`TcpStream`.] @@ -197,8 +194,8 @@ async fn tcp_connect_inner( ctx: &Context, domain: Domain, port: u16, - dns: Dns, - connector: Connector, + dns: DnsResolveIpMode, + connector: ConnectIpMode, ) -> Result<(TcpStream, SocketAddr), OpaqueError> where State: Clone + Send + Sync + 'static, @@ -210,37 +207,41 @@ where let connected = Arc::new(AtomicBool::new(false)); let sem = Arc::new(Semaphore::new(3)); - // IPv6 - let ipv6_tx = tx.clone(); - let ipv6_domain = domain.clone(); - let ipv6_connected = connected.clone(); - let ipv6_sem = sem.clone(); - ctx.spawn(tcp_connect_inner_branch( - dns.clone(), - connector.clone(), - IpKind::Ipv6, - ipv6_domain, - port, - ipv6_tx, - ipv6_connected, - ipv6_sem, - )); - - // IPv4 - let ipv4_tx = tx; - let ipv4_domain = domain.clone(); - let ipv4_connected = connected.clone(); - let ipv4_sem = sem; - ctx.spawn(tcp_connect_inner_branch( - dns, - connector, - IpKind::Ipv4, - ipv4_domain, - port, - ipv4_tx, - ipv4_connected, - ipv4_sem, - )); + match dns.mode { + DnsResolveIpMode::SingleIpV4 | DnsResolveIpMode::DualPreferIpV4 | DnsResovleIpMode::Dual =>{ + //IPV4 + ctx.spawn(tcp_connect_inner_branch( + dns.clone(), + connector.clone(), + IpKind::Ipv4, + domain.clone(), + port, + tx.clone(), + connected.clone(), + sem.clone(), + )); + } + _ => {} + + } + + match dns.mode { + DnsResolveIpMode::SingleIpV6 | DnsResovleIpMode::Dual =>{ + //IPV6 + ctx.spawn(tcp_connect_inner_branch( + dns.clone(), + connector.clone(), + IpKind::Ipv6, + domain.clone(), + port, + tx.clone(), + connected.clone(), + sem.clone(), + )); + } + _ => {} + + } // wait for the first connection to succeed, // ignore the rest of the connections (sorry, but not sorry) @@ -274,39 +275,31 @@ async fn tcp_connect_inner_branch( Dns: DnsResolver> + Clone, Connector: TcpStreamConnector + Send + 'static> + Clone, { - /*TODO - modify the match to use dns.ip_mode and match for - SingleIpv4 or DualPreferIpv4 or Dual - */ + let ip_it = match ip_kind { - IpKind::Ipv4 => match dns.ipv4_lookup(domain).await { + IpKind::Ipv4 if matches!(dns.mode, DnsResolveIpMode::Dual |DnsResolveIpMode::SingleIpV4 |DnsResolveIpMode::DualPreferIpV4 ) =>{ + match dns.ipv4_lookup(domain).await { Ok(ips) => Either::A(ips.into_iter().map(IpAddr::V4)), Err(err) => { let err = OpaqueError::from_boxed(err.into()); tracing::trace!(err = %err, "[{ip_kind:?}] failed to resolve domain to IPv4 addresses"); return; } - }, - /*TODO - modify the match to use dns.ip_mode and match for - SingleIpv6 or Dual - */ - IpKind::Ipv6 => match dns.ipv6_lookup(domain).await { + } + }, + IpKind::Ipv6 if matches!(dns.mode, DnsResolveIpMode::Dual |DnsResolveIpMode::SingleIpV6)=> { + match dns.ipv6_lookup(domain).await { Ok(ips) => Either::B(ips.into_iter().map(IpAddr::V6)), Err(err) => { let err = OpaqueError::from_boxed(err.into()); tracing::trace!(err = ?err, "[{ip_kind:?}] failed to resolve domain to IPv6 addresses"); return; } - }, + } + }, + _ => return, }; - /* - TODO - Not sure here, we need to get the matched ip type and - use it - */ - for (index, ip) in ip_it.enumerate() { let addr = (ip, port).into();