check cached DnsClient connection state before query

- ref #567

Author: zonyitoo

Cargo.lock

@@ -114,6 +114,9 @@ shadowsocks = { version = "1.11.2", path = "../shadowsocks" }
 [target.'cfg(any(target_os = "macos", target_os = "ios", target_os = "freebsd", target_os = "netbsd", target_os = "openbsd"))'.dependencies]
 nix = "0.21"
 
+[target.'cfg(windows)'.dependencies]
+winapi = { version = "0.3", features = ["mswsock", "winsock2"] }
+
 [dev-dependencies]
 byteorder = "1.3"
 env_logger = "0.8"

crates/shadowsocks-service/Cargo.toml

@@ -10,7 +10,7 @@ use std::{
     time::Duration,
 };
 
-use log::trace;
+use log::{debug, trace};
 use shadowsocks::{config::ServerConfig, net::ConnectOpts, relay::socks5::Address};
 use tokio::sync::Mutex;
 use trust_dns_resolver::proto::{error::ProtoError, op::Message};
@@ -123,6 +123,13 @@ impl DnsClientCache {
         let mut last_err = None;
 
         for _ in 0..self.retry_count {
+            // UNIX stream won't keep connection alive
+            //
+            // https://github.com/shadowsocks/shadowsocks-rust/pull/567
+            //
+            // 1. The cost of recreating UNIX stream sockets are very low
+            // 2. This feature is only used by shadowsocks-android, and it doesn't support connection reuse
+
             let mut client = match DnsClient::connect_unix_stream(ns).await {
                 Ok(client) => client,
                 Err(err) => {
@@ -245,8 +252,14 @@ impl DnsClientCache {
     {
         // Check if there already is a cached client
         if let Some(q) = self.cache.lock().await.get_mut(key) {
-            if let Some(c) = q.pop_front() {
+            while let Some(mut c) = q.pop_front() {
                 trace!("take cached DNS client for {:?}", key);
+
+                if !c.check_connected().await {
+                    debug!("cached DNS client for {:?} is lost", key);
+                    continue;
+                }
+
                 return Ok(c);
             }
         }

crates/shadowsocks-service/src/local/dns/client_cache.rs

@@ -2,7 +2,12 @@
 
 #[cfg(unix)]
 use std::path::Path;
-use std::{io, net::SocketAddr, sync::Arc, time::Duration};
+use std::{
+    io::{self, ErrorKind},
+    net::SocketAddr,
+    sync::Arc,
+    time::Duration,
+};
 
 use byteorder::{BigEndian, ByteOrder};
 use bytes::{BufMut, BytesMut};
@@ -145,6 +150,92 @@ impl DnsClient {
             }
         }
     }
+
+    /// Check if the underlying connection is still connecting
+    ///
+    /// This will only work for TCP and UNIX Stream connections.
+    /// UDP clients will always return `true`.
+    pub async fn check_connected(&mut self) -> bool {
+        #[cfg(unix)]
+        fn check_peekable<F: std::os::unix::io::AsRawFd>(fd: &mut F) -> bool {
+            let fd = fd.as_raw_fd();
+
+            unsafe {
+                let mut peek_buf = [0u8; 1];
+
+                let ret = libc::recv(
+                    fd,
+                    peek_buf.as_mut_ptr() as *mut libc::c_void,
+                    peek_buf.len(),
+                    libc::MSG_PEEK | libc::MSG_DONTWAIT,
+                );
+
+                if ret == 0 {
+                    // EOF, connection lost
+                    false
+                } else if ret > 0 {
+                    // Data in buffer
+                    true
+                } else {
+                    let err = io::Error::last_os_error();
+                    if err.kind() == ErrorKind::WouldBlock {
+                        // EAGAIN, EWOULDBLOCK
+                        // Still connected.
+                        true
+                    } else {
+                        false
+                    }
+                }
+            }
+        }
+
+        #[cfg(windows)]
+        fn check_peekable<F: std::os::windows::io::AsRawSocket>(s: &mut F) -> bool {
+            use winapi::{
+                ctypes::{c_char, c_int},
+                um::winsock2::{recv, MSG_PEEK},
+            };
+
+            let sock = s.as_raw_socket();
+
+            unsafe {
+                let mut peek_buf = [0u8; 1];
+
+                let ret = recv(
+                    sock,
+                    peek_buf.as_mut_ptr() as *mut c_char,
+                    peek_buf.len() as c_int,
+                    MSG_PEEK,
+                );
+
+                if ret == 0 {
+                    // EOF, connection lost
+                    false
+                } else if ret > 0 {
+                    // Data in buffer
+                    true
+                } else {
+                    let err = io::Error::last_os_error();
+                    if err.kind() == ErrorKind::WouldBlock {
+                        // I have to trust the `s` have already set to non-blocking mode
+                        // Becuase windows doesn't have MSG_DONTWAIT
+                        true
+                    } else {
+                        false
+                    }
+                }
+            }
+        }
+
+        match *self {
+            DnsClient::TcpLocal { ref mut stream } => check_peekable(stream),
+            DnsClient::UdpLocal { .. } => true,
+            #[cfg(unix)]
+            DnsClient::UnixStream { ref mut stream } => check_peekable(stream),
+            DnsClient::TcpRemote { ref mut stream } => check_peekable(stream.get_mut().get_mut()),
+            DnsClient::UdpRemote { .. } => true,
+        }
+    }
 }
 
 pub async fn stream_query<S>(stream: &mut S, r: &Message) -> Result<Message, ProtoError>

crates/shadowsocks-service/src/local/dns/upstream.rs

@@ -30,6 +30,11 @@ impl<S> MonProxyStream<S> {
     pub fn get_ref(&self) -> &S {
         &self.stream
     }
+
+    #[inline]
+    pub fn get_mut(&mut self) -> &mut S {
+        &mut self.stream
+    }
 }
 
 impl<S> AsyncRead for MonProxyStream<S>

crates/shadowsocks-service/src/net/mon_stream.rs

@@ -281,9 +281,11 @@ impl AsyncWrite for TcpStream {
                 }
 
                 TcpStreamState::FastOpenConnecting(ref mut overlapped) => {
-                    let n = ready!(inner.poll_write_io(cx, || {
+                    let stream = inner.get_mut();
+
+                    let n = ready!(stream.poll_write_io(cx, || {
                         unsafe {
-                            let sock = inner.as_raw_socket() as SOCKET;
+                            let sock = stream.as_raw_socket() as SOCKET;
 
                             let mut bytes_sent: DWORD = 0;
                             let mut flags: DWORD = 0;

crates/shadowsocks/src/net/sys/windows/mod.rs

@@ -1,9 +1,9 @@
 //! TcpStream wrappers that supports connecting with options
 
 #[cfg(unix)]
-use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
+use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
 #[cfg(windows)]
-use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket};
+use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket};
 use std::{
     io::{self, ErrorKind},
     net::SocketAddr,
@@ -335,3 +335,17 @@ fn setsockopt_with_opt(f: &tokio::net::TcpStream, opts: &AcceptOpts) -> io::Resu
     f.set_nodelay(opts.tcp.nodelay)?;
     Ok(())
 }
+
+#[cfg(unix)]
+impl AsRawFd for TcpStream {
+    fn as_raw_fd(&self) -> RawFd {
+        self.0.as_raw_fd()
+    }
+}
+
+#[cfg(windows)]
+impl AsRawSocket for TcpStream {
+    fn as_raw_socket(&self) -> RawSocket {
+        self.0.as_raw_socket()
+    }
+}