From d2a13693c2858b6b0fe1feb2dfff604bc2b0aa30 Mon Sep 17 00:00:00 2001
From: Harald Hoyer <harald@profian.com>
Date: Fri, 21 Jan 2022 14:19:48 +0100
Subject: [PATCH] wasi: enable TcpListener and TcpStream

With the addition of `sock_accept()` to snapshot1, simple networking via
a passed `TcpListener` is possible. This patch implements the basics to
make a simple server work.

Signed-off-by: Harald Hoyer <harald@profian.com>
---
 library/std/src/os/wasi/net/mod.rs | 20 +++++++++
 library/std/src/sys/wasi/fd.rs     |  4 ++
 library/std/src/sys/wasi/net.rs    | 66 ++++++++++++++++++++++++------
 3 files changed, 77 insertions(+), 13 deletions(-)

diff --git a/library/std/src/os/wasi/net/mod.rs b/library/std/src/os/wasi/net/mod.rs
index e6bcf87887f0..73c097d4a50a 100644
--- a/library/std/src/os/wasi/net/mod.rs
+++ b/library/std/src/os/wasi/net/mod.rs
@@ -1,3 +1,23 @@
 //! WASI-specific networking functionality
 
 #![unstable(feature = "wasi_ext", issue = "71213")]
+
+use crate::io;
+use crate::net;
+use crate::sys_common::AsInner;
+
+/// WASI-specific extensions to [`std::net::TcpListener`].
+///
+/// [`std::net::TcpListener`]: crate::net::TcpListener
+pub trait TcpListenerExt {
+    /// Accept a socket.
+    ///
+    /// This corresponds to the `sock_accept` syscall.
+    fn sock_accept(&self, flags: u16) -> io::Result<u32>;
+}
+
+impl TcpListenerExt for net::TcpListener {
+    fn sock_accept(&self, flags: u16) -> io::Result<u32> {
+        self.as_inner().as_inner().as_inner().sock_accept(flags)
+    }
+}
diff --git a/library/std/src/sys/wasi/fd.rs b/library/std/src/sys/wasi/fd.rs
index e4f4456611cd..0b9c8e61db84 100644
--- a/library/std/src/sys/wasi/fd.rs
+++ b/library/std/src/sys/wasi/fd.rs
@@ -228,6 +228,10 @@ impl WasiFd {
         unsafe { wasi::path_remove_directory(self.as_raw_fd() as wasi::Fd, path).map_err(err2io) }
     }
 
+    pub fn sock_accept(&self, flags: wasi::Fdflags) -> io::Result<wasi::Fd> {
+        unsafe { wasi::sock_accept(self.as_raw_fd() as wasi::Fd, flags).map_err(err2io) }
+    }
+
     pub fn sock_recv(
         &self,
         ri_data: &mut [IoSliceMut<'_>],
diff --git a/library/std/src/sys/wasi/net.rs b/library/std/src/sys/wasi/net.rs
index a4dbb225376e..c66e0e4d328a 100644
--- a/library/std/src/sys/wasi/net.rs
+++ b/library/std/src/sys/wasi/net.rs
@@ -1,5 +1,6 @@
 #![deny(unsafe_op_in_unsafe_fn)]
 
+use super::err2io;
 use super::fd::WasiFd;
 use crate::convert::TryFrom;
 use crate::fmt;
@@ -87,24 +88,24 @@ impl TcpStream {
         unsupported()
     }
 
-    pub fn read(&self, _: &mut [u8]) -> io::Result<usize> {
-        unsupported()
+    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.read_vectored(&mut [IoSliceMut::new(buf)])
     }
 
-    pub fn read_vectored(&self, _: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
-        unsupported()
+    pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        self.socket().as_inner().read(bufs)
     }
 
     pub fn is_read_vectored(&self) -> bool {
         true
     }
 
-    pub fn write(&self, _: &[u8]) -> io::Result<usize> {
-        unsupported()
+    pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
+        self.write_vectored(&[IoSlice::new(buf)])
     }
 
-    pub fn write_vectored(&self, _: &[IoSlice<'_>]) -> io::Result<usize> {
-        unsupported()
+    pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        self.socket().as_inner().write(bufs)
     }
 
     pub fn is_write_vectored(&self) -> bool {
@@ -155,8 +156,23 @@ impl TcpStream {
         unsupported()
     }
 
-    pub fn set_nonblocking(&self, _: bool) -> io::Result<()> {
-        unsupported()
+    pub fn set_nonblocking(&self, state: bool) -> io::Result<()> {
+        let fdstat = unsafe {
+            wasi::fd_fdstat_get(self.socket().as_inner().as_raw_fd() as wasi::Fd).map_err(err2io)?
+        };
+
+        let mut flags = fdstat.fs_flags;
+
+        if state {
+            flags |= wasi::FDFLAGS_NONBLOCK;
+        } else {
+            flags &= !wasi::FDFLAGS_NONBLOCK;
+        }
+
+        unsafe {
+            wasi::fd_fdstat_set_flags(self.socket().as_inner().as_raw_fd() as wasi::Fd, flags)
+                .map_err(err2io)
+        }
     }
 
     pub fn socket(&self) -> &Socket {
@@ -194,7 +210,16 @@ impl TcpListener {
     }
 
     pub fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> {
-        unsupported()
+        let fd = unsafe {
+            wasi::sock_accept(self.as_inner().as_inner().as_raw_fd() as _, 0).map_err(err2io)?
+        };
+
+        Ok((
+            TcpStream::from_inner(unsafe { Socket::from_raw_fd(fd as _) }),
+            // WASI has no concept of SocketAddr yet
+            // return an unspecified IPv4Addr
+            SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), 0),
+        ))
     }
 
     pub fn duplicate(&self) -> io::Result<TcpListener> {
@@ -221,8 +246,23 @@ impl TcpListener {
         unsupported()
     }
 
-    pub fn set_nonblocking(&self, _: bool) -> io::Result<()> {
-        unsupported()
+    pub fn set_nonblocking(&self, state: bool) -> io::Result<()> {
+        let fdstat = unsafe {
+            wasi::fd_fdstat_get(self.socket().as_inner().as_raw_fd() as wasi::Fd).map_err(err2io)?
+        };
+
+        let mut flags = fdstat.fs_flags;
+
+        if state {
+            flags |= wasi::FDFLAGS_NONBLOCK;
+        } else {
+            flags &= !wasi::FDFLAGS_NONBLOCK;
+        }
+
+        unsafe {
+            wasi::fd_fdstat_set_flags(self.socket().as_inner().as_raw_fd() as wasi::Fd, flags)
+                .map_err(err2io)
+        }
     }
 
     pub fn socket(&self) -> &Socket {