fixed windows TFO by using tokio try_write_io

shadowsocks · Jul 7, 2021 · 250aa95 · 250aa95
1 parent 0f942ea
commit 250aa95
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/crates/shadowsocks-service/src/local/dns/upstream.rs b/crates/shadowsocks-service/src/local/dns/upstream.rs
@@ -193,10 +193,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];

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(),
+                    }
                 }
             }
         }