Merge pull request #4732 from wasmerio/sock-poll-ready-after-connect

Poll for write-readiness after connecting a blocking socket
wasmerio · May 22, 2024 · dacf004 · dacf004
2 parents 2527d7b + f5d9dab
commit dacf004
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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(()))