fixed windows TFO by using tokio try_write_io

shadowsocks · Jul 7, 2021 · 42b55f1 · 42b55f1
1 parent 0f942ea
commit 42b55f1
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 33 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/shadowsocks-service/src/local/dns/upstream.rs b/crates/shadowsocks-service/src/local/dns/upstream.rs
@@ -3,6 +3,7 @@
 #[cfg(unix)]
 use std::path::Path;
 use std::{
+    cmp::Ordering,
     io::{self, ErrorKind},
     net::SocketAddr,
     sync::Arc,
@@ -170,20 +171,16 @@ impl DnsClient {
                     libc::MSG_PEEK | libc::MSG_DONTWAIT,
                 );
 
-                if ret == 0 {
+                match ret.cmp(&0) {
                     // EOF, connection lost
-                    false
-                } else if ret > 0 {
+                    Ordering::Equal => false,
                     // Data in buffer
-                    true
-                } else {
-                    let err = io::Error::last_os_error();
-                    if err.kind() == ErrorKind::WouldBlock {
+                    Ordering::Greater => true,
+                    Ordering::Less => {
+                        let err = io::Error::last_os_error();
                         // EAGAIN, EWOULDBLOCK
                         // Still connected.
-                        true
-                    } else {
-                        false
+                        err.kind() == ErrorKind::WouldBlock
                     }
                 }
             }
@@ -193,10 +190,10 @@ impl DnsClient {
         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},
+                um::winsock2::{recv, MSG_PEEK, SOCKET},
             };
 
-            let sock = s.as_raw_socket();
+            let sock = s.as_raw_socket() as SOCKET;
 
             unsafe {
                 let mut peek_buf = [0u8; 1];
@@ -208,20 +205,16 @@ impl DnsClient {
                     MSG_PEEK,
                 );
 
-                if ret == 0 {
+                match ret.cmp(&0) {
                     // EOF, connection lost
-                    false
-                } else if ret > 0 {
+                    Ordering::Equal => false,
                     // Data in buffer
-                    true
-                } else {
-                    let err = io::Error::last_os_error();
-                    if err.kind() == ErrorKind::WouldBlock {
+                    Ordering::Greater => true,
+                    Ordering::Less => {
+                        let err = io::Error::last_os_error();
                         // I have to trust the `s` have already set to non-blocking mode
                         // Becuase windows doesn't have MSG_DONTWAIT
-                        true
-                    } else {
-                        false
+                        err.kind() == ErrorKind::WouldBlock
                     }
                 }
             }

diff --git a/crates/shadowsocks/src/net/sys/windows/mod.rs b/crates/shadowsocks/src/net/sys/windows/mod.rs
@@ -283,7 +283,9 @@ impl AsyncWrite for TcpStream {
                 TcpStreamState::FastOpenConnecting(ref mut overlapped) => {
                     let stream = inner.get_mut();
 
-                    let n = ready!(stream.poll_write_io(cx, || {
+                    ready!(stream.poll_write_ready(cx))?;
+
+                    let write_result = stream.try_write_io(|| {
                         unsafe {
                             let sock = stream.as_raw_socket() as SOCKET;
 
@@ -320,11 +322,19 @@ impl AsyncWrite for TcpStream {
                                 Err(io::Error::from_raw_os_error(err))
                             }
                         }
-                    }))?;
+                    });
 
-                    // Connect successfully with fast open
-                    *state = TcpStreamState::Connected;
-                    return Ok(n).into();
+                    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(),
+                    }
                 }
             }
         }