migrate TFO to separate tokio-tfo crate

Author: zonyitoo

Cargo.lock

@@ -60,6 +60,7 @@ mio = "0.7"
 socket2 = { version = "0.4", features = ["all"] }
 tokio = { version = "1.9.0", features = ["io-util", "macros", "net", "parking_lot", "process", "rt", "sync", "time"] }
 tokio-io-timeout = "1.1"
+tokio-tfo = "0.1.3"
 
 trust-dns-resolver = { version = "0.20", optional = true }
 arc-swap = { version = "1.3", optional = true }

crates/shadowsocks/Cargo.toml

@@ -45,86 +45,6 @@ fn set_common_sockopt_for_connect(addr: SocketAddr, socket: &TcpSocket, opts: &C
     Ok(())
 }
 
-fn set_common_sockopt_after_connect(stream: &tokio::net::TcpStream, opts: &ConnectOpts) -> io::Result<()> {
-    stream.set_nodelay(opts.tcp.nodelay)?;
-    set_common_sockopt_after_connect_sys(stream, opts)?;
-
-    Ok(())
-}
-
-#[cfg(unix)]
-#[inline]
-fn set_common_sockopt_after_connect_sys(stream: &tokio::net::TcpStream, opts: &ConnectOpts) -> io::Result<()> {
-    use socket2::{Socket, TcpKeepalive};
-    use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
-
-    let socket = unsafe { Socket::from_raw_fd(stream.as_raw_fd()) };
-
-    macro_rules! try_sockopt {
-        ($socket:ident . $func:ident ($($arg:expr),*)) => {
-            match $socket . $func ($($arg),*) {
-                Ok(e) => e,
-                Err(err) => {
-                    let _ = socket.into_raw_fd();
-                    return Err(err);
-                }
-            }
-        };
-    }
-
-    if let Some(keepalive_duration) = opts.tcp.keepalive {
-        #[allow(unused_mut)]
-        let mut keepalive = TcpKeepalive::new().with_time(keepalive_duration);
-
-        #[cfg(any(
-            target_os = "freebsd",
-            target_os = "fuchsia",
-            target_os = "linux",
-            target_os = "netbsd",
-            target_vendor = "apple",
-        ))]
-        {
-            keepalive = keepalive.with_interval(keepalive_duration);
-        }
-
-        try_sockopt!(socket.set_tcp_keepalive(&keepalive));
-    }
-
-    let _ = socket.into_raw_fd();
-    Ok(())
-}
-
-#[cfg(windows)]
-#[inline]
-fn set_common_sockopt_after_connect_sys(stream: &tokio::net::TcpStream, opts: &ConnectOpts) -> io::Result<()> {
-    use socket2::{Socket, TcpKeepalive};
-    use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket};
-
-    let socket = unsafe { Socket::from_raw_socket(stream.as_raw_socket()) };
-
-    macro_rules! try_sockopt {
-        ($socket:ident . $func:ident ($($arg:expr),*)) => {
-            match $socket . $func ($($arg),*) {
-                Ok(e) => e,
-                Err(err) => {
-                    let _ = socket.into_raw_socket();
-                    return Err(err);
-                }
-            }
-        };
-    }
-
-    if let Some(keepalive_duration) = opts.tcp.keepalive {
-        let keepalive = TcpKeepalive::new()
-            .with_time(keepalive_duration)
-            .with_interval(keepalive_duration);
-        try_sockopt!(socket.set_tcp_keepalive(&keepalive));
-    }
-
-    let _ = socket.into_raw_socket();
-    Ok(())
-}
-
 #[cfg(all(not(windows), not(unix)))]
 #[inline]
 fn set_common_sockopt_after_connect_sys(_: &tokio::net::TcpStream, _: &ConnectOpts) -> io::Result<()> {

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

@@ -1,39 +1,31 @@
 use std::{
-    io::{self, ErrorKind},
+    io,
     mem,
-    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpStream as StdTcpStream},
-    ops::{Deref, DerefMut},
-    os::unix::io::{AsRawFd, FromRawFd, IntoRawFd},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
+    os::unix::io::{AsRawFd, RawFd},
     pin::Pin,
     task::{self, Poll},
 };
 
-use futures::ready;
 use log::error;
 use pin_project::pin_project;
-use socket2::SockAddr;
 use tokio::{
-    io::{AsyncRead, AsyncWrite, Interest, ReadBuf},
+    io::{AsyncRead, AsyncWrite, ReadBuf},
     net::{TcpSocket, TcpStream as TokioTcpStream, UdpSocket},
 };
+use tokio_tfo::TfoStream;
 
 use crate::net::{
     sys::{set_common_sockopt_after_connect, set_common_sockopt_for_connect},
     AddrFamily,
     ConnectOpts,
 };
 
-enum TcpStreamState {
-    Connected,
-    FastOpenConnect(SocketAddr),
-}
-
 /// A `TcpStream` that supports TFO (TCP Fast Open)
 #[pin_project(project = TcpStreamProj)]
-pub struct TcpStream {
-    #[pin]
-    inner: TokioTcpStream,
-    state: TcpStreamState,
+pub enum TcpStream {
+    Standard(#[pin] TokioTcpStream),
+    FastOpen(#[pin] TfoStream),
 }
 
 impl TcpStream {
@@ -50,142 +42,82 @@ impl TcpStream {
             let stream = socket.connect(addr).await?;
             set_common_sockopt_after_connect(&stream, opts)?;
 
-            return Ok(TcpStream {
-                inner: stream,
-                state: TcpStreamState::Connected,
-            });
+            return Ok(TcpStream::Standard(stream));
         }
 
-        unsafe {
-            let enable: libc::c_int = 1;
-
-            let ret = libc::setsockopt(
-                socket.as_raw_fd(),
-                libc::IPPROTO_TCP,
-                libc::TCP_FASTOPEN,
-                &enable as *const _ as *const libc::c_void,
-                mem::size_of_val(&enable) as libc::socklen_t,
-            );
-
-            if ret != 0 {
-                let err = io::Error::last_os_error();
-                error!("set TCP_FASTOPEN error: {}", err);
-                return Err(err);
-            }
-        }
-
-        let stream = TokioTcpStream::from_std(unsafe { StdTcpStream::from_raw_fd(socket.into_raw_fd()) })?;
+        let stream = TfoStream::connect_with_socket(socket, addr).await?;
         set_common_sockopt_after_connect(&stream, opts)?;
 
-        Ok(TcpStream {
-            inner: stream,
-            state: TcpStreamState::FastOpenConnect(addr),
-        })
+        Ok(TcpStream::FastOpen(stream))
+    }
+
+    pub fn local_addr(&self) -> io::Result<SocketAddr> {
+        match *self {
+            TcpStream::Standard(ref s) => s.local_addr(),
+            TcpStream::FastOpen(ref s) => s.local_addr(),
+        }
+    }
+
+    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
+        match *self {
+            TcpStream::Standard(ref s) => s.peer_addr(),
+            TcpStream::FastOpen(ref s) => s.peer_addr(),
+        }
     }
-}
 
-impl Deref for TcpStream {
-    type Target = TokioTcpStream;
+    pub fn nodelay(&self) -> io::Result<bool> {
+        match *self {
+            TcpStream::Standard(ref s) => s.nodelay(),
+            TcpStream::FastOpen(ref s) => s.nodelay(),
+        }
+    }
 
-    fn deref(&self) -> &Self::Target {
-        &self.inner
+    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
+        match *self {
+            TcpStream::Standard(ref s) => s.set_nodelay(nodelay),
+            TcpStream::FastOpen(ref s) => s.set_nodelay(nodelay),
+        }
     }
 }
 
-impl DerefMut for TcpStream {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.inner
+impl AsRawFd for TcpStream {
+    fn as_raw_fd(&self) -> RawFd {
+        match *self {
+            TcpStream::Standard(ref s) => s.as_raw_fd(),
+            TcpStream::FastOpen(ref s) => s.as_raw_fd(),
+        }
     }
 }
 
 impl AsyncRead for TcpStream {
     fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
-        self.project().inner.poll_read(cx, buf)
+        match self.project() {
+            TcpStreamProj::Standard(s) => s.poll_read(cx, buf),
+            TcpStreamProj::FastOpen(s) => s.poll_read(cx, buf),
+        }
     }
 }
 
 impl AsyncWrite for TcpStream {
-    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
-        loop {
-            let TcpStreamProj { inner, state } = self.project();
-
-            match *state {
-                TcpStreamState::Connected => return inner.poll_write(cx, buf),
-
-                TcpStreamState::FastOpenConnect(addr) => {
-                    // TCP_FASTOPEN was supported since FreeBSD 12.0
-                    //
-                    // Example program:
-                    // <https://people.freebsd.org/~pkelsey/tfo-tools/tfo-client.c>
-
-                    let saddr = SockAddr::from(addr);
-
-                    let stream = inner.get_mut();
-
-                    // Ensure socket is writable
-                    ready!(stream.poll_write_ready(cx))?;
-
-                    let mut connecting = false;
-                    let send_result = stream.try_io(Interest::WRITABLE, || {
-                        unsafe {
-                            let ret = libc::sendto(
-                                stream.as_raw_fd(),
-                                buf.as_ptr() as *const libc::c_void,
-                                buf.len(),
-                                0, // Yes, BSD doesn't need MSG_FASTOPEN
-                                saddr.as_ptr(),
-                                saddr.len(),
-                            );
-
-                            if ret >= 0 {
-                                Ok(ret as usize)
-                            } else {
-                                // Error occurs
-                                let err = io::Error::last_os_error();
-
-                                // EINPROGRESS
-                                if let Some(libc::EINPROGRESS) = err.raw_os_error() {
-                                    // For non-blocking socket, it returns the number of bytes queued (and transmitted in the SYN-data packet) if cookie is available.
-                                    // If cookie is not available, it transmits a data-less SYN packet with Fast Open cookie request option and returns -EINPROGRESS like connect().
-                                    //
-                                    // So in this state. We have to loop again to call `poll_write` for sending the first packet.
-                                    connecting = true;
-
-                                    // Let `poll_write_io` clears the write readiness.
-                                    Err(ErrorKind::WouldBlock.into())
-                                } else {
-                                    // Other errors, including EAGAIN, EWOULDBLOCK
-                                    Err(err)
-                                }
-                            }
-                        }
-                    });
-
-                    match send_result {
-                        Ok(n) => {
-                            // Connected successfully with fast open
-                            *state = TcpStreamState::Connected;
-                            return Ok(n).into();
-                        }
-                        Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
-                            if connecting {
-                                // Connecting with normal TCP handshakes, write the first packet after connected
-                                *state = TcpStreamState::Connected;
-                            }
-                        }
-                        Err(err) => return Err(err).into(),
-                    }
-                }
-            }
+    fn poll_write(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
+        match self.project() {
+            TcpStreamProj::Standard(s) => s.poll_write(cx, buf),
+            TcpStreamProj::FastOpen(s) => s.poll_write(cx, buf),
         }
     }
 
     fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
-        self.project().inner.poll_flush(cx)
+        match self.project() {
+            TcpStreamProj::Standard(s) => s.poll_flush(cx),
+            TcpStreamProj::FastOpen(s) => s.poll_flush(cx),
+        }
     }
 
     fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
-        self.project().inner.poll_shutdown(cx)
+        match self.project() {
+            TcpStreamProj::Standard(s) => s.poll_shutdown(cx),
+            TcpStreamProj::FastOpen(s) => s.poll_shutdown(cx),
+        }
     }
 }