starred/shadowsocks-rust
1
0
Fork 0
mirror of https://github.com/shadowsocks/shadowsocks-rust.git synced 2026-11-06 16:14:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

migrate TFO to separate tokio-tfo crate

Browse Source
This commit is contained in:
zonyitoo
2021-10-23 11:53:01 +08:00
parent 6915ce9591
commit bf4664e7c2
9 changed files with 358 additions and 813 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+18
View File
@@ -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"
crates/shadowsocks/Cargo.toml
+1
View File
@@ -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/src/net/sys/mod.rs
-80
View File
@@ -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/unix/bsd/freebsd.rs
+61 -129
View File
@@ -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<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(),
}
}
pub fn nodelay(&self) -> io::Result<bool> {
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<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),
}
}
}
crates/shadowsocks/src/net/sys/unix/bsd/macos.rs
+60 -126
View File
@@ -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<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(),
}
}
pub fn nodelay(&self) -> io::Result<bool> {
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<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.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<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),
}
}
}
crates/shadowsocks/src/net/sys/unix/linux/mod.rs
+61 -209
View File
@@ -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<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(),
}
}
pub fn nodelay(&self) -> io::Result<bool> {
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<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.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<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),
}
}
}
crates/shadowsocks/src/net/sys/unix/mod.rs
+46 -1
View File
@@ -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<UdpSocke
UdpSocket::from_std(socket.into())
}
}
pub fn set_common_sockopt_after_connect<S: AsRawFd>(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(())
}
crates/shadowsocks/src/net/sys/windows/mod.rs
+95 -258
View File
@@ -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<LPFN_CONNECTEX> = 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<OVERLAPPED>),
}
// 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<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(),
}
}
pub fn nodelay(&self) -> io::Result<bool> {
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<io::Result<()>> {
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<io::Result<usize>> {
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<OVERLAPPED> = 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<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),
}
}
}
@@ -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<S: AsRawSocket>(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(())
}
crates/shadowsocks/src/net/tcp.rs
+16 -10
View File
@@ -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<SocketAddr> {
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<SocketAddr> {
self.0.peer_addr()
}
/// Gets the value of the `TCP_NODELAY` option on this socket.
pub fn nodelay(&self) -> io::Result<bool> {
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)
}
}