write dns client

Author: xerik

Cargo.lock

@@ -27,8 +27,6 @@ tokio-net = "0.2.0-alpha.6"
 radix_trie = "0.1.*"
 structopt = "^0"
 treebitmap = "0.3.1"
-trust-dns = "0.14.0"
-trust-dns-resolver = "*"
-futures01 = { package = "futures", version = "0.1" }
+trust-dns = { version = "^0.17", default-features = false }
 futures-preview =  { version = "0.3.0-alpha.16", features = ["compat"] }
 tokio = "0.2.0-alpha.6"

Cargo.toml

@@ -26,11 +26,7 @@ pub struct Socks5Datagram {
 impl Socks5Datagram {
     /// Creates a UDP socket bound to the specified address which will have its
     /// traffic routed through the specified proxy.
-    pub async fn bind_with_timeout(
-        proxy: SocketAddr,
-        local: SocketAddr,
-        time: Option<Duration>,
-    ) -> Result<Socks5Datagram, SocksError> {
+    pub async fn bind(proxy: SocketAddr, local: SocketAddr) -> Result<Socks5Datagram, SocksError> {
         // we don't know what our IP is from the perspective of the proxy, so
         // don't try to pass `addr` in here.
         let dst = Address::SocketAddress(SocketAddr::V4(SocketAddrV4::new(
@@ -51,10 +47,6 @@ impl Socks5Datagram {
             }
         };
         let socket = net::UdpSocket::bind(&local)?;
-        if time.is_some() {
-            socket.set_read_timeout(time)?;
-        }
-
         let socket = UdpSocket::from_std(socket, &Handle::default())?;
 
         Ok(Socks5Datagram {

asocks5/src/client/udp.rs

@@ -1,5 +1,5 @@
 mod decision_tree;
-mod main;
+pub(crate) mod main;
 mod prefix_match;
 mod util;
 

src/conf/mod.rs

@@ -2,7 +2,6 @@ extern crate byteorder;
 extern crate bytes;
 #[macro_use]
 extern crate failure;
-extern crate futures01;
 extern crate httparse;
 #[macro_use]
 extern crate log;
@@ -15,8 +14,6 @@ extern crate env_logger;
 extern crate tokio;
 extern crate tokio_io;
 extern crate treebitmap;
-extern crate trust_dns;
-extern crate trust_dns_resolver;
 
 use structopt::StructOpt;
 
@@ -31,7 +28,6 @@ use crate::relay::run_with_conf;
 
 use futures::task::Context;
 
-
 use futures::Future;
 use std::pin::Pin;
 use std::task::Poll;
@@ -61,11 +57,7 @@ pub fn run() -> Result<(), i32> {
         for r in conf.relays {
             info!("Starting {}", r);
 
-            if let Err(e) = run_with_conf(
-                r,
-                conf.domain_matcher.clone(),
-                conf.ip_matcher.clone(),
-            ) {
+            if let Err(e) = run_with_conf(r, conf.domain_matcher.clone(), conf.ip_matcher.clone()) {
                 error!("Relay error: {:?}", e);
             }
         }

src/lib.rs

@@ -1,7 +1,6 @@
 use std::io;
 use std::time;
 
-
 use futures::task::Context;
 use futures::{Future, Poll as Poll1};
 use futures_timer::Delay;

src/relay/forwarding/tcp/copy/mod.rs

@@ -91,7 +91,7 @@ async fn carry_out(
         .map_err(|e| format_err!("Error sending {:?} header bytes to {:?}: {}", &pr, a, e))?;
     let (ur, uw) = split(s);
     let (cr, cw) = split(client_stream);
-    run_copy(ur, cw, a, pr.clone(), r.clone(), true, );
+    run_copy(ur, cw, a, pr.clone(), r.clone(), true);
     run_copy(cr, uw, a, pr, r, false);
     Ok(())
 }

src/relay/forwarding/tcp/mod.rs

@@ -4,8 +4,6 @@ use bytes::BytesMut;
 use failure::Error;
 use tokio::net::TcpStream;
 
-
-
 use super::InspectedTcp;
 use crate::relay::inspect::parse::guess_bytes;
 use tokio::prelude::*;

src/relay/inspect/codec.rs

@@ -1,9 +1,8 @@
 use failure::Error;
-use futures::compat::Future01CompatExt;
+
 use std::net::SocketAddr;
 use std::sync::Arc;
 use tokio;
-use trust_dns_resolver::AsyncResolver;
 
 use crate::relay::forwarding::handle_incoming_tcp;
 use crate::relay::TcpRouter;
@@ -14,11 +13,9 @@ use asocks5::socks::TcpRequestHeader;
 use asocks5::Command;
 use tokio::net::TcpStream;
 
-
 use futures::future::ready;
 
-use futures::task::SpawnExt;
-
+use crate::resolver::AsyncResolver;
 use tokio::prelude::*;
 
 pub async fn listen_socks(
@@ -68,13 +65,12 @@ async fn read_address(
     match head.address {
         Address::SocketAddress(a) => Ok(a),
         Address::DomainNameAddress(domain, port) => {
-            let lookup = resolver
-                .lookup_ip(&*domain)
-                .compat()
+            let ips = resolver
+                .resolve(&*domain)
                 .await
                 .map_err(|e| format_err!("Error resolving {}: {}", domain, e))?;
-            let ip = lookup
-                .iter()
+            let ip = ips
+                .into_iter()
                 .next()
                 .ok_or_else(|| format_err!("No address found for domain {}", domain))?;
             Ok(SocketAddr::new(ip, port))

src/relay/listen/socks.rs

@@ -11,23 +11,17 @@ pub use self::route::TcpRouter;
 use crate::conf::Relay;
 use crate::conf::RelayProto;
 use crate::conf::{DomainMatcher, IpMatcher};
-use crate::resolver::create_resolver;
-
-
+use crate::resolver::AsyncResolver;
 
 /// Start a relay
-pub fn run_with_conf(
-    conf: Relay,
-    d: Arc<DomainMatcher>,
-    i: Arc<IpMatcher>,
-) -> Result<(), Error> {
+pub fn run_with_conf(conf: Relay, d: Arc<DomainMatcher>, i: Arc<IpMatcher>) -> Result<(), Error> {
     let rule = conf.rule.val().clone();
     // FIXME support proxy
     let ns = conf.nameserver_or_default();
     if ns.egress.is_some() {
         error!("Resolver config in relay doesn't support proxy yet");
     }
-    let resolver = Arc::new(create_resolver(ns.remote));
+    let resolver = Arc::new(AsyncResolver::new(&ns));
     let router = TcpRouter::new(d, i, rule);
     match conf.listen {
         RelayProto::Socks5(a) => {

src/relay/mod.rs

@@ -3,10 +3,8 @@ use std::io;
 use std::net::SocketAddr;
 
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
-use std::time::Duration;
-use tokio::net::TcpStream;
 
-use super::TIMEOUT;
+use tokio::net::TcpStream;
 
 use crate::conf::NameServerRemote;
 use asocks5::socks::Address;
@@ -44,9 +42,7 @@ impl SockGetterAsync {
     pub async fn get_udp(&self, data: Vec<u8>) -> Result<Vec<u8>, SocksError> {
         use std::str::FromStr;
         let la = SocketAddr::from_str("0.0.0.0:0").unwrap();
-        let socks5 =
-            Socks5Datagram::bind_with_timeout(self.proxy, la, Some(Duration::from_secs(TIMEOUT)))
-                .await?;
+        let socks5 = Socks5Datagram::bind(self.proxy, la).await?;
         let addr = ns_sock_addr(&self.addr);
         let socks5 = socks5.send_to(&data, Address::SocketAddress(addr)).await?;
 

src/resolver/client/socks/mod.rs

@@ -1,4 +1,3 @@
-
 use std::io;
 use std::net::SocketAddr;
 use std::net::UdpSocket as UdpSocketStd;

src/resolver/client/udp.rs

@@ -1,25 +1,93 @@
-use crate::conf::NameServerRemote;
-use futures::compat::Future01CompatExt;
+use super::dnsclient::DnsClient;
+use crate::conf::NameServer;
 
-use futures::task::SpawnExt;
-use trust_dns_resolver::config::{NameServerConfig, Protocol};
-use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
-use trust_dns_resolver::AsyncResolver;
+use failure::Error;
+use std::net::IpAddr;
+use std::str::FromStr;
+use trust_dns::op::Message;
+use trust_dns::op::Query;
+use trust_dns::proto::rr::Name;
+use trust_dns::rr::RecordType;
+use trust_dns::serialize::binary::{BinDecodable, BinDecoder};
 
-pub fn create_resolver(rm: NameServerRemote) -> AsyncResolver {
-    let (addr, p) = match rm {
-        NameServerRemote::Udp(a) => (a, Protocol::Udp),
-        NameServerRemote::Tcp(a) => (a, Protocol::Tcp),
-    };
-    let nc = NameServerConfig {
-        socket_addr: addr,
-        protocol: p,
-        tls_dns_name: None,
-    };
-    let conf = ResolverConfig::from_parts(None, vec![], vec![nc]);
-    let (resolver, fut) = AsyncResolver::new(conf, ResolverOpts::default());
-    tokio::spawn(async move {
-        fut.compat().await.expect("AsyncResolver");
-    });
-    resolver
+use failure::_core::time::Duration;
+use tokio::future::FutureExt;
+
+pub struct AsyncResolver {
+    client: DnsClient,
+}
+impl AsyncResolver {
+    pub fn new(rm: &NameServer) -> AsyncResolver {
+        let client = DnsClient::new(rm);
+        AsyncResolver { client }
+    }
+
+    pub async fn resolve(&self, name: &str) -> Result<Vec<IpAddr>, Error> {
+        let result = self
+            .resolve_eternal(name)
+            .timeout(Duration::from_secs(10))
+            .await?;
+        result
+    }
+    async fn resolve_eternal(&self, name: &str) -> Result<Vec<IpAddr>, Error> {
+        let mut msg = Message::new();
+        let name = Name::from_str(name).unwrap();
+        let query = Query::query(name, RecordType::A);
+        msg.add_query(query);
+        let res = self.client.resolve(msg.to_vec()?).await?;
+        let mut decoder = BinDecoder::new(&res);
+        let message = Message::read(&mut decoder).expect("msg deco err");
+        let ips = message
+            .answers()
+            .into_iter()
+            .filter_map(|rec| {
+                let rd = rec.rdata();
+                let ip = rd.to_ip_addr();
+                ip
+            })
+            .collect();
+        Ok(ips)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::conf::{main::RefVal, Egress, EgressAddr};
+    use crate::conf::{NameServer, NameServerRemote};
+    use crate::resolver::AsyncResolver;
+    use bytes::Bytes;
+    use std::net::IpAddr;
+    use std::net::SocketAddr;
+    use std::str::FromStr;
+
+    #[test]
+    fn udp_test() {
+        let ip = [127, 8, 8, 8];
+        let remote = NameServerRemote::Udp(SocketAddr::new(IpAddr::from(ip), 53001));
+        let ns = NameServer {
+            remote,
+            egress: None,
+        };
+        let resolver = AsyncResolver::new(&ns);
+        let rt = tokio::runtime::Runtime::new().unwrap();
+        let response = rt.block_on(async move { resolver.resolve("www.example.com.").await });
+        assert!(response.is_err());
+    }
+
+    #[test]
+    fn socks_test() {
+        let ip = [8, 8, 8, 8];
+        let remote = NameServerRemote::Tcp(SocketAddr::new(IpAddr::from(ip), 53));
+        let ns = NameServer {
+            remote,
+            egress: Some(RefVal::Val(Egress {
+                name: Bytes::new(),
+                addr: EgressAddr::Socks5(SocketAddr::from_str("1.1.1.1:3128").unwrap()),
+            })),
+        };
+        let resolver = AsyncResolver::new(&ns);
+        let rt = tokio::runtime::Runtime::new().unwrap();
+        let response = rt.block_on(async move { resolver.resolve("www.example.com").await });
+        assert!(response.is_err());
+    }
 }

src/resolver/lookup.rs

@@ -6,5 +6,5 @@ mod handler;
 mod lookup;
 mod serve;
 
-pub use self::lookup::create_resolver;
+pub use self::lookup::AsyncResolver;
 pub use self::serve::serve;

src/resolver/mod.rs

@@ -10,16 +10,11 @@ use crate::conf::DomainMatcher;
 use crate::resolver::handler;
 use crate::resolver::handler::SmartResolver;
 
-
-
 use std::net::UdpSocket as UdpSocketStd;
 use tokio::net::UdpSocket;
 use tokio_net::driver::Handle;
 
-pub fn serve(
-    conf: DnsProxy,
-    matcher: Arc<DomainMatcher>,
-) -> Result<(), Error> {
+pub fn serve(conf: DnsProxy, matcher: Arc<DomainMatcher>) -> Result<(), Error> {
     let handler = handler::SmartResolver::new(matcher, &conf)?;
     let handler = Arc::new(handler);
     let addr = conf.listen;