From 8ac4a8fa7a4a55add899d955b731fe1a1cdcc692 Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Mon, 3 May 2021 11:22:11 -0500
Subject: [PATCH] Check WSAE* errors on windows

---
 vm/src/stdlib/socket.rs | 31 ++++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs
index a93ea3868..a20fde292 100644
--- a/vm/src/stdlib/socket.rs
+++ b/vm/src/stdlib/socket.rs
@@ -29,6 +29,19 @@ type RawSocket = std::os::unix::io::RawFd;
 #[cfg(windows)]
 type RawSocket = std::os::windows::raw::SOCKET;
 
+#[cfg(unix)]
+macro_rules! errcode {
+    ($e:ident) => {
+        c::$e
+    };
+}
+#[cfg(windows)]
+macro_rules! errcode {
+    ($e:ident) => {
+        paste::paste!(c::[<WSA $e>])
+    };
+}
+
 #[cfg(unix)]
 mod c {
     pub use libc::*;
@@ -135,7 +148,10 @@ impl PySocket {
                 Ok(addr) if family == -1 => family = addr.family() as i32,
                 Err(e)
                     if family == -1
-                        || matches!(e.raw_os_error(), Some(libc::ENOTSOCK) | Some(libc::EBADF)) =>
+                        || matches!(
+                            e.raw_os_error(),
+                            Some(errcode!(ENOTSOCK)) | Some(errcode!(EBADF))
+                        ) =>
                 {
                     return Err(e.into_pyexception(vm))
                 }
@@ -356,7 +372,12 @@ impl PySocket {
             vm.check_signals()?;
             self.timeout.load() != 0.0
         } else {
-            self.timeout.load() > 0.0 && err.raw_os_error() == Some(libc::EINPROGRESS)
+            #[cfg(unix)]
+            use c::EINPROGRESS;
+            #[cfg(windows)]
+            use c::WSAEWOULDBLOCK as EINPROGRESS;
+
+            self.timeout.load() > 0.0 && err.raw_os_error() == Some(EINPROGRESS)
         };
 
         if wait_connect {
@@ -1303,13 +1324,13 @@ fn _socket_dup(x: RawSocket, vm: &VirtualMachine) -> PyResult<RawSocket> {
 
 fn _socket_close(x: RawSocket, vm: &VirtualMachine) -> PyResult<()> {
     #[cfg(unix)]
-    use libc::{close, ECONNRESET};
+    use libc::close;
     #[cfg(windows)]
-    use winapi::um::winsock2::{closesocket as close, WSAECONNRESET as ECONNRESET};
+    use winapi::um::winsock2::closesocket as close;
     let ret = unsafe { close(x as _) };
     if ret < 0 {
         let err = super::os::errno();
-        if err.raw_os_error() != Some(ECONNRESET) {
+        if err.raw_os_error() != Some(errcode!(ECONNRESET)) {
             return Err(err.into_pyexception(vm));
         }
     }