From 2d533e05925db1eb872713455a19e4e233624c73 Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Sun, 15 Nov 2020 00:38:35 +0800 Subject: [PATCH] ProxyUpstream sending queries with both TCP & UDP if possible --- Cargo.lock | 39 ++++--- src/relay/dnsrelay/mod.rs | 32 ++---- src/relay/dnsrelay/upstream.rs | 192 +++++++++++++++++++-------------- 3 files changed, 142 insertions(+), 121 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 17dc2b63..81408d51 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -770,7 +770,7 @@ checksum = "3c1ad908cc71012b7bea4d0c53ba96a8cba9962f048fa68d143376143d863b7a" [[package]] name = "hyper" version = "0.14.0-dev" -source = "git+https://github.com/hyperium/hyper.git#b819b428d314f2203642a015545967601b8e518a" +source = "git+https://github.com/hyperium/hyper.git#b3b8010e1c6eadc2769103e13f1c1d8cca521a5b" dependencies = [ "bytes 0.5.6", "futures-channel", @@ -1159,9 +1159,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.5.1" +version = "1.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f53cef67919d7d247eb9a2f128ca9e522789967ef1eb4ccd8c71a95a8aedf596" +checksum = "13bd41f508810a131401606d54ac32a467c97172d74ba7662562ebba5ad07fa0" [[package]] name = "opaque-debug" @@ -1498,9 +1498,9 @@ dependencies = [ [[package]] name = "resolv-conf" -version = "0.6.3" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11834e137f3b14e309437a8276714eed3a80d1ef894869e510f2c0c0b98b9f4a" +checksum = "52e44394d2086d010551b14b53b1f24e31647570cd1deb0379e2c21b329aba00" dependencies = [ "hostname", "quick-error", @@ -1996,9 +1996,18 @@ dependencies = [ [[package]] name = "tinyvec" -version = "0.3.4" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "238ce071d267c5710f9d31451efec16c5ee22de34df17cc05e56cbc92e967117" +checksum = "b78a366903f506d2ad52ca8dc552102ffdd3e937ba8a227f024dc1d1eae28575" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c" [[package]] name = "tokio" @@ -2143,7 +2152,7 @@ checksum = "efd1f82c56340fdf16f2a953d7bda4f8fdffba13d93b00844c25572110b26079" [[package]] name = "trust-dns-https" version = "0.20.0-alpha.3" -source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#d035ee38f4fc05c872d8590a02ff8579cf506af9" +source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#751299675322fb0792a1345838119cf182cceaf0" dependencies = [ "bytes 0.5.6", "cfg-if 1.0.0", @@ -2165,7 +2174,7 @@ dependencies = [ [[package]] name = "trust-dns-proto" version = "0.20.0-alpha.3" -source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#d035ee38f4fc05c872d8590a02ff8579cf506af9" +source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#751299675322fb0792a1345838119cf182cceaf0" dependencies = [ "async-trait", "cfg-if 1.0.0", @@ -2189,7 +2198,7 @@ dependencies = [ [[package]] name = "trust-dns-resolver" version = "0.20.0-alpha.3" -source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#d035ee38f4fc05c872d8590a02ff8579cf506af9" +source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#751299675322fb0792a1345838119cf182cceaf0" dependencies = [ "cfg-if 1.0.0", "futures-util", @@ -2214,7 +2223,7 @@ dependencies = [ [[package]] name = "trust-dns-rustls" version = "0.20.0-alpha.3" -source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#d035ee38f4fc05c872d8590a02ff8579cf506af9" +source = "git+https://github.com/bluejekyll/trust-dns.git?branch=upgrade-to-tokio-0.3#751299675322fb0792a1345838119cf182cceaf0" dependencies = [ "futures-channel", "futures-io", @@ -2265,18 +2274,18 @@ dependencies = [ [[package]] name = "unicode-normalization" -version = "0.1.13" +version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6fb19cf769fa8c6a80a162df694621ebeb4dafb606470b2b2fce0be40a98a977" +checksum = "b7f98e67a4d84f730d343392f9bfff7d21e3fca562b9cb7a43b768350beeddc6" dependencies = [ "tinyvec", ] [[package]] name = "unicode-segmentation" -version = "1.6.0" +version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e83e153d1053cbb5a118eeff7fd5be06ed99153f00dbcd8ae310c5fb2b22edc0" +checksum = "db8716a166f290ff49dabc18b44aa407cb7c6dbe1aa0971b44b8a24b0ca35aae" [[package]] name = "unicode-width" diff --git a/src/relay/dnsrelay/mod.rs b/src/relay/dnsrelay/mod.rs index 9e3df35f..a88920f7 100644 --- a/src/relay/dnsrelay/mod.rs +++ b/src/relay/dnsrelay/mod.rs @@ -17,14 +17,10 @@ use crate::{ acl::AccessControl, config::ConfigType, context::SharedContext, - relay::{ - loadbalancing::server::{PlainPingBalancer, ServerType}, - sys::create_udp_socket, - utils::try_timeout, - }, + relay::{sys::create_udp_socket, utils::try_timeout}, }; -use self::upstream::{ProxyTcpUpstream, ProxyUdpUpstream}; +use self::upstream::ProxyUpstream; pub(crate) mod upstream; @@ -409,27 +405,11 @@ pub async fn run(context: SharedContext) -> io::Result<()> { .clone() .expect("remote query DNS address"); - let udp_dns = { - let balancer = PlainPingBalancer::new(context.clone(), ServerType::Udp).await; + let proxy_upstream = ProxyUpstream::new(context.clone(), remote_dns_addr).await; + let proxy_relay = Arc::new(DnsRelay::new(context.clone(), proxy_upstream)); - let relay = DnsRelay::new( - context.clone(), - ProxyUdpUpstream::new(move || balancer.pick_server(), remote_dns_addr.clone()), - ); - - run_udp(Arc::new(relay), bind_addr) - }; - - let tcp_dns = { - let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await; - - let relay = DnsRelay::new( - context.clone(), - ProxyTcpUpstream::new(context.clone(), move || balancer.pick_server(), remote_dns_addr.clone()), - ); - - run_tcp(Arc::new(relay), bind_addr) - }; + let udp_dns = run_udp(proxy_relay.clone(), bind_addr); + let tcp_dns = run_tcp(proxy_relay, bind_addr); select! { res = udp_dns => res, diff --git a/src/relay/dnsrelay/upstream.rs b/src/relay/dnsrelay/upstream.rs index 00510dd6..139e0ba9 100644 --- a/src/relay/dnsrelay/upstream.rs +++ b/src/relay/dnsrelay/upstream.rs @@ -4,7 +4,6 @@ use std::{ fmt, fmt::{Debug, Formatter}, io, - marker::PhantomData, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}, }; @@ -23,10 +22,10 @@ use trust_dns_proto::{ }; use crate::{ - config::{Config, LocalDnsAddr}, + config::{Config, LocalDnsAddr, Mode, ServerConfig}, context::{Context, SharedContext}, relay::{ - loadbalancing::server::{ServerData, SharedServerStatistic}, + loadbalancing::server::{PlainPingBalancer, ServerType}, socks5::Address, sys::{create_outbound_udp_socket, tcp_stream_connect}, tcprelay::ProxyStream, @@ -212,101 +211,67 @@ impl Upstream for TcpUpstream { } } -pub struct ProxyTcpUpstream -where - S: ServerData, -{ - context: SharedContext, - svr_cfg: F, - ns: Address, - _s: PhantomData, +enum ProxyUpstreamMode { + TcpOnly { + balancer: PlainPingBalancer, + }, + UdpOnly { + balancer: PlainPingBalancer, + }, + TcpAndUdp { + tcp_balancer: PlainPingBalancer, + udp_balancer: PlainPingBalancer, + }, } -impl ProxyTcpUpstream -where - S: ServerData, - F: Fn() -> SharedServerStatistic + Send + Sync, -{ - pub fn new(context: SharedContext, svr_cfg: F, ns: Address) -> ProxyTcpUpstream { - ProxyTcpUpstream { - context, - svr_cfg, - ns, - _s: PhantomData, +impl Debug for ProxyUpstreamMode { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match *self { + ProxyUpstreamMode::TcpOnly { .. } => f.write_str("TcpOnly"), + ProxyUpstreamMode::UdpOnly { .. } => f.write_str("UdpOnly"), + ProxyUpstreamMode::TcpAndUdp { .. } => f.write_str("TcpAndUdp"), } } } -impl Debug for ProxyTcpUpstream -where - S: ServerData, -{ - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - f.debug_struct("ProxyTcpUpstream").field("ns", &self.ns).finish() - } +pub struct ProxyUpstream { + context: SharedContext, + ns: Address, + mode: ProxyUpstreamMode, } -#[async_trait] -impl Upstream for ProxyTcpUpstream -where - S: ServerData, - F: Fn() -> SharedServerStatistic + Send + Sync, -{ - async fn lookup(&self, _context: &Context, query: &Query) -> io::Result { - let svr_data = (self.svr_cfg)(); - let svr_cfg = svr_data.server_config(); +impl ProxyUpstream { + pub async fn new(context: SharedContext, ns: Address) -> ProxyUpstream { + let mode = match context.config().mode { + Mode::TcpOnly => ProxyUpstreamMode::TcpOnly { + balancer: PlainPingBalancer::new(context.clone(), ServerType::Tcp).await, + }, + Mode::UdpOnly => ProxyUpstreamMode::UdpOnly { + balancer: PlainPingBalancer::new(context.clone(), ServerType::Udp).await, + }, + Mode::TcpAndUdp => ProxyUpstreamMode::TcpAndUdp { + tcp_balancer: PlainPingBalancer::new(context.clone(), ServerType::Tcp).await, + udp_balancer: PlainPingBalancer::new(context.clone(), ServerType::Udp).await, + }, + }; + + ProxyUpstream { context, ns, mode } + } + + async fn tcp_lookup(&self, svr_cfg: &ServerConfig, query: &Query) -> io::Result { trace!( "DNS TCP proxied query {:?} via {} to {}", query, svr_cfg.addr(), self.ns ); - let mut stream = ProxyStream::connect_proxied(self.context.clone(), &svr_cfg, &self.ns).await?; + let mut stream = ProxyStream::connect_proxied(self.context.clone(), svr_cfg, &self.ns).await?; stream_lookup(query, &mut stream).await } -} -pub struct ProxyUdpUpstream -where - S: ServerData, -{ - svr_cfg: F, - ns: Address, - _s: PhantomData, -} + async fn udp_lookup(&self, svr_cfg: &ServerConfig, query: &Query) -> io::Result { + let context = &self.context; -impl ProxyUdpUpstream -where - S: ServerData, - F: Fn() -> SharedServerStatistic + Send + Sync, -{ - pub fn new(svr_cfg: F, ns: Address) -> ProxyUdpUpstream { - ProxyUdpUpstream { - svr_cfg, - ns, - _s: PhantomData, - } - } -} - -impl Debug for ProxyUdpUpstream -where - S: ServerData, -{ - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - f.debug_struct("ProxyUdpUpstream").field("ns", &self.ns).finish() - } -} - -#[async_trait] -impl Upstream for ProxyUdpUpstream -where - S: ServerData, - F: Fn() -> SharedServerStatistic + Send + Sync, -{ - async fn lookup(&self, context: &Context, query: &Query) -> io::Result { - let svr_data = (self.svr_cfg)(); - let svr_cfg = svr_data.server_config(); trace!( "DNS UDP proxied query {:?} via {} to {}", query, @@ -326,6 +291,73 @@ where } } +impl Debug for ProxyUpstream { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("ProxyUpstream").field("ns", &self.ns).finish() + } +} + +#[async_trait] +impl Upstream for ProxyUpstream { + async fn lookup(&self, _context: &Context, query: &Query) -> io::Result { + match self.mode { + ProxyUpstreamMode::TcpOnly { ref balancer } => { + let svr_cfg = balancer.pick_server(); + self.tcp_lookup(svr_cfg.server_config(), query).await + } + ProxyUpstreamMode::UdpOnly { ref balancer } => { + let svr_cfg = balancer.pick_server(); + self.udp_lookup(svr_cfg.server_config(), query).await + } + ProxyUpstreamMode::TcpAndUdp { + ref tcp_balancer, + ref udp_balancer, + } => { + let tcp_svr_cfg = tcp_balancer.pick_server(); + let tcp_fut = self.tcp_lookup(tcp_svr_cfg.server_config(), query); + tokio::pin!(tcp_fut); + + let udp_svr_cfg = udp_balancer.pick_server(); + let udp_fut = self.udp_lookup(udp_svr_cfg.server_config(), query); + tokio::pin!(udp_fut); + + match future::select(tcp_fut, udp_fut).await { + Either::Left((tcp_result, udp_fut)) => match tcp_result { + Ok(message) => { + trace!("ProxyUpstream {} TCP query answer, {:?}", self.ns, message); + Ok(message) + } + Err(err) => { + trace!( + "ProxyUpstream {} TCP query failed, error: {}, continue with UDP query, {:?}", + self.ns, + err, + query + ); + udp_fut.await + } + }, + Either::Right((udp_result, tcp_fut)) => match udp_result { + Ok(message) => { + trace!("ProxyUpstream {} UDP query answer, {:?}", self.ns, message); + Ok(message) + } + Err(err) => { + trace!( + "ProxyUpstream {} UDP query failed, error: {}, continue with TCP query, {:?}", + self.ns, + err, + query + ); + tcp_fut.await + } + }, + } + } + } + } +} + #[cfg(unix)] #[derive(Debug)] pub struct UnixSocketUpstream {