From f5d9dab5de855cfdb8c1e0c8aa065cda944582a8 Mon Sep 17 00:00:00 2001
From: Arshia Ghafoori <arshia001@live.com>
Date: Wed, 22 May 2024 13:58:00 +0330
Subject: [PATCH] Poll for write-readiness after connecting a blocking socket

---
 lib/wasix/src/net/socket.rs                  |  4 ++++
 lib/wasix/src/syscalls/mod.rs                |  4 ++--
 lib/wasix/src/syscalls/wasix/sock_bind.rs    |  2 +-
 lib/wasix/src/syscalls/wasix/sock_connect.rs | 12 +++++++++++-
 lib/wasix/src/syscalls/wasix/sock_listen.rs  |  2 +-
 5 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/lib/wasix/src/net/socket.rs b/lib/wasix/src/net/socket.rs
index bbf6c8dd9e9..f83d12859d3 100644
--- a/lib/wasix/src/net/socket.rs
+++ b/lib/wasix/src/net/socket.rs
@@ -536,6 +536,7 @@ impl InodeSocket {
         net: &dyn VirtualNetworking,
         peer: SocketAddr,
         timeout: Option<std::time::Duration>,
+        nonblocking: bool,
     ) -> Result<Option<InodeSocket>, Errno> {
         let new_write_timeout;
         let new_read_timeout;
@@ -576,6 +577,9 @@ impl InodeSocket {
                                 if let Some(dont_route) = dont_route {
                                     ret.set_dontroute(dont_route).ok();
                                 }
+                                if !nonblocking {
+                                    futures::future::poll_fn(|cx| ret.poll_write_ready(cx)).await?;
+                                }
                                 Ok(ret)
                             })
                         }
diff --git a/lib/wasix/src/syscalls/mod.rs b/lib/wasix/src/syscalls/mod.rs
index 676c280c98f..f4b1ee10911 100644
--- a/lib/wasix/src/syscalls/mod.rs
+++ b/lib/wasix/src/syscalls/mod.rs
@@ -760,7 +760,7 @@ pub(crate) fn __sock_upgrade<'a, F, Fut>(
     actor: F,
 ) -> Result<(), Errno>
 where
-    F: FnOnce(crate::net::socket::InodeSocket) -> Fut,
+    F: FnOnce(crate::net::socket::InodeSocket, Fdflags) -> Fut,
     Fut: std::future::Future<Output = Result<Option<crate::net::socket::InodeSocket>, Errno>> + 'a,
 {
     let env = ctx.data();
@@ -786,7 +786,7 @@ where
                 drop(guard);
 
                 // Start the work using the socket
-                let work = actor(socket);
+                let work = actor(socket, fd_entry.flags);
 
                 // Block on the work and process it
                 let res = InlineWaker::block_on(work);
diff --git a/lib/wasix/src/syscalls/wasix/sock_bind.rs b/lib/wasix/src/syscalls/wasix/sock_bind.rs
index ca030f41eda..bdea8822e3f 100644
--- a/lib/wasix/src/syscalls/wasix/sock_bind.rs
+++ b/lib/wasix/src/syscalls/wasix/sock_bind.rs
@@ -48,7 +48,7 @@ pub(crate) fn sock_bind_internal(
         ctx,
         sock,
         Rights::SOCK_BIND,
-        move |socket| async move { socket.bind(tasks.deref(), net.deref(), addr).await }
+        move |socket, _| async move { socket.bind(tasks.deref(), net.deref(), addr).await }
     ));
 
     Ok(Ok(()))
diff --git a/lib/wasix/src/syscalls/wasix/sock_connect.rs b/lib/wasix/src/syscalls/wasix/sock_connect.rs
index e935bec8627..3942d3dbcb4 100644
--- a/lib/wasix/src/syscalls/wasix/sock_connect.rs
+++ b/lib/wasix/src/syscalls/wasix/sock_connect.rs
@@ -58,7 +58,17 @@ pub(crate) fn sock_connect_internal(
         ctx,
         sock,
         Rights::SOCK_CONNECT,
-        move |mut socket| async move { socket.connect(tasks.deref(), net.deref(), addr, None).await }
+        move |mut socket, flags| async move {
+            socket
+                .connect(
+                    tasks.deref(),
+                    net.deref(),
+                    addr,
+                    None,
+                    flags.contains(Fdflags::NONBLOCK),
+                )
+                .await
+        }
     ));
 
     Ok(Ok(()))
diff --git a/lib/wasix/src/syscalls/wasix/sock_listen.rs b/lib/wasix/src/syscalls/wasix/sock_listen.rs
index 1bc969ebc28..c706bd8e2d1 100644
--- a/lib/wasix/src/syscalls/wasix/sock_listen.rs
+++ b/lib/wasix/src/syscalls/wasix/sock_listen.rs
@@ -49,7 +49,7 @@ pub(crate) fn sock_listen_internal(
         ctx,
         sock,
         Rights::SOCK_LISTEN,
-        |socket| async move { socket.listen(tasks.deref(), net.deref(), backlog).await }
+        |socket, _| async move { socket.listen(tasks.deref(), net.deref(), backlog).await }
     ));
 
     Ok(Ok(()))