From 42b55f1916aa9bd14bdaea36e8c31b33d78db059 Mon Sep 17 00:00:00 2001
From: zonyitoo <zonyitoo@gmail.com>
Date: Wed, 7 Jul 2021 01:02:21 +0800
Subject: [PATCH] fixed windows TFO by using tokio try_write_io

---
 Cargo.lock                                    | 12 +++---
 .../src/local/dns/upstream.rs                 | 37 ++++++++-----------
 crates/shadowsocks/src/net/sys/windows/mod.rs | 20 +++++++---
 3 files changed, 36 insertions(+), 33 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 519e14fecf58..7d7be3c6a7c8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1028,9 +1028,9 @@ dependencies = [
 
 [[package]]
 name = "ordered-float"
-version = "2.5.1"
+version = "2.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f100fcfb41e5385e0991f74981732049f9b896821542a219420491046baafdc2"
+checksum = "6dea6388d3d5498ec651701f14edbaf463c924b5d8829fb2848ccf0bcc7b3c69"
 dependencies = [
  "num-traits",
 ]
@@ -1814,8 +1814,8 @@ checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c"
 
 [[package]]
 name = "tokio"
-version = "1.8.0"
-source = "git+https://github.com/zonyitoo/tokio.git#00159823901e36085bc22b4db32dcb562181e134"
+version = "1.8.1"
+source = "git+https://github.com/zonyitoo/tokio.git#6fa7de17f01c761acb34fa8b62c46ce5708208f0"
 dependencies = [
  "autocfg",
  "bytes",
@@ -1843,8 +1843,8 @@ dependencies = [
 
 [[package]]
 name = "tokio-macros"
-version = "1.2.0"
-source = "git+https://github.com/zonyitoo/tokio.git#00159823901e36085bc22b4db32dcb562181e134"
+version = "1.3.0"
+source = "git+https://github.com/zonyitoo/tokio.git#6fa7de17f01c761acb34fa8b62c46ce5708208f0"
 dependencies = [
  "proc-macro2",
  "quote",
diff --git a/crates/shadowsocks-service/src/local/dns/upstream.rs b/crates/shadowsocks-service/src/local/dns/upstream.rs
index 89b7d901476c..e8e86427e8fd 100644
--- 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
index 17d05a933606..e0c5ece5000c 100644
--- 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(),
+                    }
                 }
             }
         }