diff --git a/Cargo.lock b/Cargo.lock index 2a6a396d..a26936d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1574,6 +1574,7 @@ dependencies = [ "thiserror", "tokio", "tokio-io-timeout", + "tokio-tfo", "trust-dns-resolver", "url", "winapi", @@ -1936,6 +1937,23 @@ dependencies = [ "webpki", ] +[[package]] +name = "tokio-tfo" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b22b79bdd90c351d2f03bfd0d4297e215a2b5a395bb999c69fc7041d1ee5847a" +dependencies = [ + "cfg-if", + "futures", + "libc", + "log", + "once_cell", + "pin-project", + "socket2 0.4.2", + "tokio", + "winapi", +] + [[package]] name = "tokio-util" version = "0.6.8" diff --git a/crates/shadowsocks/Cargo.toml b/crates/shadowsocks/Cargo.toml index 5af6466b..c8d8a250 100644 --- a/crates/shadowsocks/Cargo.toml +++ b/crates/shadowsocks/Cargo.toml @@ -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 } diff --git a/crates/shadowsocks/src/net/sys/mod.rs b/crates/shadowsocks/src/net/sys/mod.rs index e29bae69..6d2fba87 100644 --- a/crates/shadowsocks/src/net/sys/mod.rs +++ b/crates/shadowsocks/src/net/sys/mod.rs @@ -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<()> { diff --git a/crates/shadowsocks/src/net/sys/unix/bsd/freebsd.rs b/crates/shadowsocks/src/net/sys/unix/bsd/freebsd.rs index ff8aa467..423f84c0 100644 --- a/crates/shadowsocks/src/net/sys/unix/bsd/freebsd.rs +++ b/crates/shadowsocks/src/net/sys/unix/bsd/freebsd.rs @@ -1,21 +1,19 @@ 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}, @@ -23,17 +21,11 @@ use crate::net::{ 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 { + match *self { + TcpStream::Standard(ref s) => s.local_addr(), + TcpStream::FastOpen(ref s) => s.local_addr(), + } + } + + pub fn peer_addr(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.peer_addr(), + TcpStream::FastOpen(ref s) => s.peer_addr(), + } + } + + pub fn nodelay(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.nodelay(), + TcpStream::FastOpen(ref s) => s.nodelay(), + } + } + + 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 Deref for TcpStream { - type Target = TokioTcpStream; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -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> { - 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> { - 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: - // - - 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> { + 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> { - 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> { - self.project().inner.poll_shutdown(cx) + match self.project() { + TcpStreamProj::Standard(s) => s.poll_shutdown(cx), + TcpStreamProj::FastOpen(s) => s.poll_shutdown(cx), + } } } diff --git a/crates/shadowsocks/src/net/sys/unix/bsd/macos.rs b/crates/shadowsocks/src/net/sys/unix/bsd/macos.rs index e29a50b7..2be50eb5 100644 --- a/crates/shadowsocks/src/net/sys/unix/bsd/macos.rs +++ b/crates/shadowsocks/src/net/sys/unix/bsd/macos.rs @@ -1,22 +1,20 @@ use std::{ io::{self, ErrorKind}, 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, ptr, 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}, @@ -24,17 +22,11 @@ use crate::net::{ ConnectOpts, }; -enum TcpStreamState { - Connected, - FastOpenWrite, -} - /// 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 { @@ -55,140 +47,82 @@ impl TcpStream { // If TFO is not enabled, it just works like a normal 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)); } - // TFO in macos uses connectx - - unsafe { - let raddr = SockAddr::from(addr); - - let mut endpoints: libc::sa_endpoints_t = mem::zeroed(); - endpoints.sae_dstaddr = raddr.as_ptr(); - endpoints.sae_dstaddrlen = raddr.len(); - - let ret = libc::connectx( - socket.as_raw_fd(), - &endpoints as *const _, - libc::SAE_ASSOCID_ANY, - libc::CONNECT_DATA_IDEMPOTENT /* Enable TFO */ | libc::CONNECT_RESUME_ON_READ_WRITE, /* Send SYN with subsequence send/recv */ - ptr::null(), - 0, - ptr::null_mut(), - ptr::null_mut(), - ); - - if ret != 0 { - return Err(io::Error::last_os_error()); - } - } - - 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::FastOpenWrite, - }) + Ok(TcpStream::FastOpen(stream)) + } + + pub fn local_addr(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.local_addr(), + TcpStream::FastOpen(ref s) => s.local_addr(), + } + } + + pub fn peer_addr(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.peer_addr(), + TcpStream::FastOpen(ref s) => s.peer_addr(), + } + } + + pub fn nodelay(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.nodelay(), + TcpStream::FastOpen(ref s) => s.nodelay(), + } + } + + 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 Deref for TcpStream { - type Target = TokioTcpStream; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -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> { - 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> { - loop { - let TcpStreamProj { inner, state } = self.as_mut().project(); - - match *state { - TcpStreamState::Connected => return inner.poll_write(cx, buf), - - TcpStreamState::FastOpenWrite => { - // `CONNECT_RESUME_ON_READ_WRITE` is set when calling `connectx`, - // so the first call of `send` will perform the actual SYN with TFO cookie. - // - // (NOT SURE) If remote server doesn't support TFO or this is the first connection, - // it may return EINPROGRESS just like other platforms (Linux, FreeBSD). - - 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::send(stream.as_raw_fd(), buf.as_ptr() as *const libc::c_void, buf.len(), 0); - if ret >= 0 { - Ok(ret as usize) - } else { - let err = io::Error::last_os_error(); - // EAGAIN and EWOULDBLOCK should have been handled by tokio - // - // 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 - 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> { + 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> { - 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> { - self.project().inner.poll_shutdown(cx) + match self.project() { + TcpStreamProj::Standard(s) => s.poll_shutdown(cx), + TcpStreamProj::FastOpen(s) => s.poll_shutdown(cx), + } } } diff --git a/crates/shadowsocks/src/net/sys/unix/linux/mod.rs b/crates/shadowsocks/src/net/sys/unix/linux/mod.rs index 010c95e0..bba3d71e 100644 --- a/crates/shadowsocks/src/net/sys/unix/linux/mod.rs +++ b/crates/shadowsocks/src/net/sys/unix/linux/mod.rs @@ -1,23 +1,20 @@ 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, - sync::atomic::{AtomicBool, Ordering}, task::{self, Poll}, }; use cfg_if::cfg_if; -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}, @@ -25,18 +22,11 @@ use crate::net::{ ConnectOpts, }; -enum TcpStreamState { - Connected, - FastOpenConnect(SocketAddr), - FastOpenWrite, -} - /// 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 { @@ -95,220 +85,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)); } - let mut connected = false; - - // TFO in Linux was supported since 3.7 - // - // But TCP_FASTOPEN_CONNECT was supported since 4.1, so we have to be compatible with it - static SUPPORT_TCP_FASTOPEN_CONNECT: AtomicBool = AtomicBool::new(true); - if SUPPORT_TCP_FASTOPEN_CONNECT.load(Ordering::Relaxed) { - unsafe { - let enable: libc::c_int = 1; - - let ret = libc::setsockopt( - socket.as_raw_fd(), - libc::IPPROTO_TCP, - libc::TCP_FASTOPEN_CONNECT, - &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(); - if let Some(libc::ENOPROTOOPT) = err.raw_os_error() { - // `TCP_FASTOPEN_CONNECT` is not supported, maybe kernel version < 4.11 - // Fallback to `sendto` with `MSG_FASTOPEN` (Supported after 3.7) - SUPPORT_TCP_FASTOPEN_CONNECT.store(false, Ordering::Relaxed); - } else { - error!("set TCP_FASTOPEN_CONNECT error: {}", err); - return Err(err); - } - } else { - connected = true; - } - } - } - - let stream = if connected { - // call connect() if TCP_FASTOPEN_CONNECT is set - socket.connect(addr).await? - } else { - // call sendto() with MSG_FASTOPEN in poll_read - 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: if connected { - TcpStreamState::FastOpenWrite - } else { - TcpStreamState::FastOpenConnect(addr) - }, - }) + Ok(TcpStream::FastOpen(stream)) + } + + pub fn local_addr(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.local_addr(), + TcpStream::FastOpen(ref s) => s.local_addr(), + } + } + + pub fn peer_addr(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.peer_addr(), + TcpStream::FastOpen(ref s) => s.peer_addr(), + } + } + + pub fn nodelay(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.nodelay(), + TcpStream::FastOpen(ref s) => s.nodelay(), + } + } + + 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 Deref for TcpStream { - type Target = TokioTcpStream; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -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> { - 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> { - loop { - let TcpStreamProj { inner, state } = self.as_mut().project(); - - match *state { - TcpStreamState::Connected => return inner.poll_write(cx, buf), - - TcpStreamState::FastOpenConnect(addr) => { - // Fallback mode. Must be kernal < 4.11 - // - // Uses sendto as BSD-like systems - - 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(), - libc::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 `try_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(), - } - } - - TcpStreamState::FastOpenWrite => { - // First `write` after `TCP_FASTOPEN_CONNECT` - // Kernel >= 4.11 - - 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::send(stream.as_raw_fd(), buf.as_ptr() as *const libc::c_void, buf.len(), 0); - - if ret >= 0 { - Ok(ret as usize) - } else { - 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> { + 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> { - 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> { - self.project().inner.poll_shutdown(cx) + match self.project() { + TcpStreamProj::Standard(s) => s.poll_shutdown(cx), + TcpStreamProj::FastOpen(s) => s.poll_shutdown(cx), + } } } diff --git a/crates/shadowsocks/src/net/sys/unix/mod.rs b/crates/shadowsocks/src/net/sys/unix/mod.rs index 1563ed5e..b55adfec 100644 --- a/crates/shadowsocks/src/net/sys/unix/mod.rs +++ b/crates/shadowsocks/src/net/sys/unix/mod.rs @@ -1,13 +1,16 @@ use std::{ io::{self, ErrorKind}, net::SocketAddr, + os::unix::io::{AsRawFd, FromRawFd, IntoRawFd}, }; use cfg_if::cfg_if; use log::{debug, warn}; -use socket2::{Domain, Protocol, SockAddr, Socket, Type}; +use socket2::{Domain, Protocol, SockAddr, Socket, TcpKeepalive, Type}; use tokio::net::UdpSocket; +use crate::net::ConnectOpts; + cfg_if! { if #[cfg(any(target_os = "linux", target_os = "android"))] { mod linux; @@ -75,3 +78,45 @@ pub async fn create_inbound_udp_socket(addr: &SocketAddr) -> io::Result(stream: &S, opts: &ConnectOpts) -> io::Result<()> { + 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 opts.tcp.nodelay { + try_sockopt!(socket.set_nodelay(true)); + } + + 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(()) +} diff --git a/crates/shadowsocks/src/net/sys/windows/mod.rs b/crates/shadowsocks/src/net/sys/windows/mod.rs index fe13a2a4..574f3747 100644 --- a/crates/shadowsocks/src/net/sys/windows/mod.rs +++ b/crates/shadowsocks/src/net/sys/windows/mod.rs @@ -1,113 +1,45 @@ use std::{ io::{self, ErrorKind}, mem, - net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpStream as StdTcpStream}, - ops::{Deref, DerefMut}, - os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket}, + net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}, + os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket}, pin::Pin, ptr, task::{self, Poll}, }; -use futures::ready; use log::{debug, error, warn}; -use once_cell::sync::Lazy; use pin_project::pin_project; -use socket2::{Domain, Protocol, SockAddr, Socket, Type}; +use socket2::{Domain, Protocol, SockAddr, Socket, TcpKeepalive, Type}; use tokio::{ - io::{AsyncRead, AsyncWrite, Interest, ReadBuf}, + io::{AsyncRead, AsyncWrite, ReadBuf}, net::{TcpSocket, TcpStream as TokioTcpStream, UdpSocket}, }; +use tokio_tfo::TfoStream; use winapi::{ ctypes::{c_char, c_int}, shared::{ - minwindef::{BOOL, DWORD, FALSE, LPDWORD, LPVOID, TRUE}, - winerror::ERROR_IO_PENDING, - ws2def::{AF_INET, IPPROTO_TCP, SIO_GET_EXTENSION_FUNCTION_POINTER}, + minwindef::{BOOL, DWORD, FALSE, LPDWORD, LPVOID}, + ws2def::IPPROTO_TCP, }, um::{ - minwinbase::{LPOVERLAPPED, OVERLAPPED}, - mswsock::{LPFN_CONNECTEX, SIO_UDP_CONNRESET, SO_UPDATE_CONNECT_CONTEXT, WSAID_CONNECTEX}, - winnt::PVOID, - winsock2::{ - closesocket, - setsockopt, - socket, - WSAGetLastError, - WSAGetOverlappedResult, - WSAIoctl, - INVALID_SOCKET, - SOCKET, - SOCKET_ERROR, - SOCK_STREAM, - SOL_SOCKET, - WSA_IO_INCOMPLETE, - }, + mswsock::SIO_UDP_CONNRESET, + winsock2::{setsockopt, WSAGetLastError, WSAIoctl, SOCKET, SOCKET_ERROR}, }, }; -use crate::net::{ - sys::{set_common_sockopt_after_connect, set_common_sockopt_for_connect}, - AddrFamily, - ConnectOpts, -}; +use crate::net::{sys::set_common_sockopt_for_connect, AddrFamily, ConnectOpts}; // ws2ipdef.h // FIXME: Use winapi's definition if issue resolved // https://github.com/retep998/winapi-rs/issues/856 const TCP_FASTOPEN: DWORD = 15; -static PFN_CONNECTEX_OPT: Lazy = Lazy::new(|| unsafe { - let socket = socket(AF_INET, SOCK_STREAM, 0); - if socket == INVALID_SOCKET { - return None; - } - - let mut guid = WSAID_CONNECTEX; - let mut num_bytes: DWORD = 0; - - let mut connectex: LPFN_CONNECTEX = None; - - let ret = WSAIoctl( - socket, - SIO_GET_EXTENSION_FUNCTION_POINTER, - &mut guid as *mut _ as LPVOID, - mem::size_of_val(&guid) as DWORD, - &mut connectex as *mut _ as LPVOID, - mem::size_of_val(&connectex) as DWORD, - &mut num_bytes as *mut _, - ptr::null_mut(), - None, - ); - - if ret != 0 { - let err = WSAGetLastError(); - let e = io::Error::from_raw_os_error(err); - - warn!("Failed to get ConnectEx function from WSA extension, error: {}", e); - } - - let _ = closesocket(socket); - - connectex -}); - -enum TcpStreamState { - Connected, - FastOpenConnect(SocketAddr), - FastOpenConnecting(Box), -} - -// unsafe: OVERLAPPED can be sent between threads -unsafe impl Send for TcpStreamState {} -unsafe impl Sync for TcpStreamState {} - /// 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 { @@ -124,207 +56,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)); } - let sock = socket.as_raw_socket() as SOCKET; - - unsafe { - // TCP_FASTOPEN was supported since Windows 10 - - // Enable TCP_FASTOPEN option - - let enable: DWORD = 1; - - let ret = setsockopt( - sock, - IPPROTO_TCP as c_int, - TCP_FASTOPEN as c_int, - &enable as *const _ as *const c_char, - mem::size_of_val(&enable) as c_int, - ); - - if ret == SOCKET_ERROR { - let err = io::Error::from_raw_os_error(WSAGetLastError()); - error!("set TCP_FASTOPEN error: {}", err); - return Err(err); - } - - if opts.bind_local_addr.is_none() { - // Bind to a dummy address (required for TFO socket) - match addr.ip() { - IpAddr::V4(..) => socket.bind(SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0))?, - IpAddr::V6(..) => socket.bind(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 0))?, - } - } - } - - let stream = TokioTcpStream::from_std(unsafe { StdTcpStream::from_raw_socket(socket.into_raw_socket()) })?; + 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 { + match *self { + TcpStream::Standard(ref s) => s.local_addr(), + TcpStream::FastOpen(ref s) => s.local_addr(), + } + } + + pub fn peer_addr(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.peer_addr(), + TcpStream::FastOpen(ref s) => s.peer_addr(), + } + } + + pub fn nodelay(&self) -> io::Result { + match *self { + TcpStream::Standard(ref s) => s.nodelay(), + TcpStream::FastOpen(ref s) => s.nodelay(), + } + } + + 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 Deref for TcpStream { - type Target = TokioTcpStream; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -impl DerefMut for TcpStream { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.inner +impl AsRawSocket for TcpStream { + fn as_raw_socket(&self) -> RawSocket { + match *self { + TcpStream::Standard(ref s) => s.as_raw_socket(), + TcpStream::FastOpen(ref s) => s.as_raw_socket(), + } } } impl AsyncRead for TcpStream { fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { - self.project().inner.poll_read(cx, buf) - } -} - -fn set_update_connect_context(sock: SOCKET) -> io::Result<()> { - unsafe { - // Make getpeername work - // https://docs.microsoft.com/en-us/windows/win32/api/mswsock/nc-mswsock-lpfn_connectex - let ret = setsockopt(sock, SOL_SOCKET, SO_UPDATE_CONNECT_CONTEXT, ptr::null(), 0); - if ret == SOCKET_ERROR { - let err = WSAGetLastError(); - return Err(io::Error::from_raw_os_error(err)); + match self.project() { + TcpStreamProj::Standard(s) => s.poll_read(cx, buf), + TcpStreamProj::FastOpen(s) => s.poll_read(cx, buf), } } - - Ok(()) } impl AsyncWrite for TcpStream { - fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll> { - loop { - let TcpStreamProj { inner, state } = self.as_mut().project(); - - match *state { - TcpStreamState::Connected => return inner.poll_write(cx, buf), - - TcpStreamState::FastOpenConnect(addr) => { - let saddr = SockAddr::from(addr); - - unsafe { - // https://docs.microsoft.com/en-us/windows/win32/api/mswsock/nc-mswsock-lpfn_connectex - let connect_ex = PFN_CONNECTEX_OPT - .expect("LPFN_CONNECTEX function doesn't exist. It is only supported after Windows 10"); - - let sock = inner.as_raw_socket() as SOCKET; - - let mut overlapped: Box = Box::new(mem::zeroed()); - - let mut bytes_sent: DWORD = 0; - let ret: BOOL = connect_ex( - sock, - saddr.as_ptr(), - saddr.len() as c_int, - buf.as_ptr() as PVOID, - buf.len() as DWORD, - &mut bytes_sent as LPDWORD, - overlapped.as_mut() as LPOVERLAPPED, - ); - - if ret == TRUE { - // Connected successfully. - - // Make getpeername() works - set_update_connect_context(sock)?; - - debug_assert!(bytes_sent as usize <= buf.len()); - - *state = TcpStreamState::Connected; - return Ok(bytes_sent as usize).into(); - } - - let err = WSAGetLastError(); - if err != ERROR_IO_PENDING as c_int { - return Err(io::Error::from_raw_os_error(err)).into(); - } - - // ConnectEx pending (ERROR_IO_PENDING), check later in FastOpenConnecting - *state = TcpStreamState::FastOpenConnecting(overlapped); - } - } - - TcpStreamState::FastOpenConnecting(ref mut overlapped) => { - let stream = inner.get_mut(); - - ready!(stream.poll_write_ready(cx))?; - - let write_result = stream.try_io(Interest::WRITABLE, || { - unsafe { - let sock = stream.as_raw_socket() as SOCKET; - - let mut bytes_sent: DWORD = 0; - let mut flags: DWORD = 0; - - // Fetch ConnectEx's result in a non-blocking way. - let ret: BOOL = WSAGetOverlappedResult( - sock, - overlapped.as_mut() as LPOVERLAPPED, - &mut bytes_sent as LPDWORD, - FALSE, // fWait = false, non-blocking, returns WSA_IO_INCOMPLETE - &mut flags as LPDWORD, - ); - - if ret == TRUE { - // Get ConnectEx's result successfully. Socket is connected - - // Make getpeername() works - set_update_connect_context(sock)?; - - debug_assert!(bytes_sent as usize <= buf.len()); - - return Ok(bytes_sent as usize); - } - - let err = WSAGetLastError(); - if err == WSA_IO_INCOMPLETE { - // ConnectEx is still not connected. Wait for the next round - // - // Let `try_io` clears the write readiness. - Err(ErrorKind::WouldBlock.into()) - } else { - Err(io::Error::from_raw_os_error(err)) - } - } - }); - - match write_result { - Ok(n) => { - // Connect successfully with fast open - *state = TcpStreamState::Connected; - return Ok(n).into(); - } - Err(ref err) if err.kind() == ErrorKind::WouldBlock => { - // Wait again for writable event. - } - Err(err) => return Err(err).into(), - } - } - } + fn poll_write(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll> { + 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> { - 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> { - self.project().inner.poll_shutdown(cx) + match self.project() { + TcpStreamProj::Standard(s) => s.poll_shutdown(cx), + TcpStreamProj::FastOpen(s) => s.poll_shutdown(cx), + } } } @@ -466,3 +273,33 @@ pub async fn create_outbound_udp_socket(af: AddrFamily, opts: &ConnectOpts) -> i Ok(socket) } + +pub fn set_common_sockopt_after_connect(stream: &S, opts: &ConnectOpts) -> io::Result<()> { + 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 opts.tcp.nodelay { + try_sockopt!(socket.set_nodelay(true)); + } + + 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(()) +} diff --git a/crates/shadowsocks/src/net/tcp.rs b/crates/shadowsocks/src/net/tcp.rs index e0a834c1..1cb11227 100644 --- a/crates/shadowsocks/src/net/tcp.rs +++ b/crates/shadowsocks/src/net/tcp.rs @@ -77,19 +77,25 @@ impl TcpStream { Ok(TcpStream(stream)) } -} -impl Deref for TcpStream { - type Target = TokioTcpStream; - - fn deref(&self) -> &Self::Target { - &self.0 + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result { + self.0.local_addr() } -} -impl DerefMut for TcpStream { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.0 + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result { + self.0.peer_addr() + } + + /// Gets the value of the `TCP_NODELAY` option on this socket. + pub fn nodelay(&self) -> io::Result { + self.0.nodelay() + } + + /// Sets the value of the `TCP_NODELAY` option on this socket. + pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> { + self.0.set_nodelay(nodelay) } }