From 1e86063574b7e50743f30b6333c27e3c5747e150 Mon Sep 17 00:00:00 2001
From: Dave Bakker <github@davebakker.io>
Date: Tue, 3 Oct 2023 18:59:56 +0200
Subject: [PATCH] wasi-sockets: Implement initial listen backlog (#7034)

* Allow backlog size to be set before initial listen

* Move set_listen_backlog_size into example_body

* Format code

* Let cap_net_ext handle the default value.

* retrigger checks

* Got lost while pulling in changes from main.
---
 .../wasi-sockets-tests/src/lib.rs             |  2 +
 crates/wasi/src/preview2/host/tcp.rs          | 41 +++++++++++++++----
 crates/wasi/src/preview2/tcp.rs               |  4 ++
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/crates/test-programs/wasi-sockets-tests/src/lib.rs b/crates/test-programs/wasi-sockets-tests/src/lib.rs
index 56c32d3e87..c7a49203bb 100644
--- a/crates/test-programs/wasi-sockets-tests/src/lib.rs
+++ b/crates/test-programs/wasi-sockets-tests/src/lib.rs
@@ -44,6 +44,8 @@ pub fn example_body(net: tcp::Network, sock: tcp::TcpSocket, family: network::Ip
 
     let sub = sock.subscribe();
 
+    sock.set_listen_backlog_size(32).unwrap();
+
     sock.start_listen().unwrap();
     poll::poll_one(&sub);
     sock.finish_listen().unwrap();
diff --git a/crates/wasi/src/preview2/host/tcp.rs b/crates/wasi/src/preview2/host/tcp.rs
index a022f4d1ab..479ae44dc6 100644
--- a/crates/wasi/src/preview2/host/tcp.rs
+++ b/crates/wasi/src/preview2/host/tcp.rs
@@ -159,7 +159,7 @@ impl<T: WasiView> crate::preview2::host::tcp::tcp::HostTcpSocket for T {
         socket
             .tcp_socket()
             .as_socketlike_view::<TcpListener>()
-            .listen(None)?;
+            .listen(socket.listen_backlog_size)?;
 
         socket.tcp_state = TcpState::ListenStarted;
 
@@ -317,16 +317,41 @@ impl<T: WasiView> crate::preview2::host::tcp::tcp::HostTcpSocket for T {
         this: Resource<tcp::TcpSocket>,
         value: u64,
     ) -> Result<(), network::Error> {
-        let table = self.table();
-        let socket = table.get_resource(&this)?;
+        const MIN_BACKLOG: i32 = 1;
+        const MAX_BACKLOG: i32 = i32::MAX; // OS'es will most likely limit it down even further.
+
+        let table = self.table_mut();
+        let socket = table.get_resource_mut(&this)?;
+
+        // Silently clamp backlog size. This is OK for us to do, because operating systems do this too.
+        let value = value
+            .try_into()
+            .unwrap_or(i32::MAX)
+            .clamp(MIN_BACKLOG, MAX_BACKLOG);
 
         match socket.tcp_state {
-            TcpState::Listening => {}
-            _ => return Err(ErrorCode::NotInProgress.into()),
-        }
+            TcpState::Default | TcpState::BindStarted | TcpState::Bound => {
+                // Socket not listening yet. Stash value for first invocation to `listen`.
+                socket.listen_backlog_size = Some(value);
 
-        let value = value.try_into().map_err(|_| ErrorCode::OutOfMemory)?;
-        Ok(rustix::net::listen(socket.tcp_socket(), value)?)
+                Ok(())
+            }
+            TcpState::Listening => {
+                // Try to update the backlog by calling `listen` again.
+                // Not all platforms support this. We'll only update our own value if the OS supports changing the backlog size after the fact.
+
+                rustix::net::listen(socket.tcp_socket(), value)
+                    .map_err(|_| ErrorCode::AlreadyListening)?;
+
+                socket.listen_backlog_size = Some(value);
+
+                Ok(())
+            }
+            TcpState::Connected => Err(ErrorCode::AlreadyConnected.into()),
+            TcpState::Connecting | TcpState::ConnectReady | TcpState::ListenStarted => {
+                Err(ErrorCode::ConcurrencyConflict.into())
+            }
+        }
     }
 
     fn keep_alive(&mut self, this: Resource<tcp::TcpSocket>) -> Result<bool, network::Error> {
diff --git a/crates/wasi/src/preview2/tcp.rs b/crates/wasi/src/preview2/tcp.rs
index 61bef55dd2..3d5925a686 100644
--- a/crates/wasi/src/preview2/tcp.rs
+++ b/crates/wasi/src/preview2/tcp.rs
@@ -53,6 +53,9 @@ pub struct TcpSocket {
 
     /// The current state in the bind/listen/accept/connect progression.
     pub(crate) tcp_state: TcpState,
+
+    /// The desired listen queue size. Set to None to use the system's default.
+    pub(crate) listen_backlog_size: Option<i32>,
 }
 
 pub(crate) struct TcpReadStream {
@@ -264,6 +267,7 @@ impl TcpSocket {
         Ok(Self {
             inner: Arc::new(stream),
             tcp_state: TcpState::Default,
+            listen_backlog_size: None,
         })
     }