[WASI-Socket] recv_from support

Signed-off-by: Sylveon <[email protected]>
dm4 · Mar 11, 2022 · 6dc0fe6 · 6dc0fe6
1 parent 3674d9d
commit 6dc0fe6
Show file tree

Hide file tree

Showing 9 changed files with 182 additions and 61 deletions.
diff --git a/include/host/wasi/environ.h b/include/host/wasi/environ.h
@@ -925,6 +925,31 @@ class Environ {
     }
   }
 
+  /// Receive a message from a socket.
+  ///
+  /// Note: This is similar to `recv` in POSIX, though it also supports reading
+  /// the data into multiple buffers in the manner of `readv`.
+  ///
+  /// @param[in] RiData List of scatter/gather vectors to which to store data.
+  /// @param[in] RiFlags Message flags.
+  /// @param[in] Address Address of the target.
+  /// @param[in] AddressLength The buffer size of Address.
+  /// @param[out] NRead Return the number of bytes stored in RiData.
+  /// @param[out] RoFlags Return message flags.
+  /// @return Nothing or WASI error.
+  WasiExpect<void> sockRecvFrom(__wasi_fd_t Fd, Span<Span<uint8_t>> RiData,
+                                __wasi_riflags_t RiFlags, uint8_t *Address,
+                                uint8_t AddressLength, __wasi_size_t &NRead,
+                                __wasi_roflags_t &RoFlags) const noexcept {
+    auto Node = getNodeOrNull(Fd);
+    if (unlikely(!Node)) {
+      return WasiUnexpect(__WASI_ERRNO_BADF);
+    } else {
+      return Node->sockRecvFrom(RiData, RiFlags, Address, AddressLength, NRead,
+                                RoFlags);
+    }
+  }
+
   /// Send a message on a socket.
   ///
   /// Note: This is similar to `send` in POSIX, though it also supports writing

diff --git a/include/host/wasi/inode.h b/include/host/wasi/inode.h
@@ -526,6 +526,21 @@ class INode
                             __wasi_riflags_t RiFlags, __wasi_size_t &NRead,
                             __wasi_roflags_t &RoFlags) const noexcept;
 
+  /// Receive a message from a socket.
+  ///
+  /// Note: This is similar to `recv` in POSIX, though it also supports reading
+  /// the data into multiple buffers in the manner of `readv`.
+  ///
+  /// @param[in] RiData List of scatter/gather vectors to which to store data.
+  /// @param[in] RiFlags Message flags.
+  /// @param[out] NRead Return the number of bytes stored in RiData.
+  /// @param[out] RoFlags Return message flags.
+  /// @return Nothing or WASI error.
+  WasiExpect<void> sockRecvFrom(Span<Span<uint8_t>> RiData,
+                                __wasi_riflags_t RiFlags, uint8_t *Address,
+                                uint8_t AddressLength, __wasi_size_t &NRead,
+                                __wasi_roflags_t &RoFlags) const noexcept;
+
   /// Send a message on a socket.
   ///
   /// Note: This is similar to `send` in POSIX, though it also supports writing

diff --git a/include/host/wasi/vinode.h b/include/host/wasi/vinode.h
@@ -576,6 +576,26 @@ class VINode : public std::enable_shared_from_this<VINode> {
     return Node.sockRecv(RiData, RiFlags, NRead, RoFlags);
   }
 
+  /// Receive a message from a socket.
+  ///
+  /// Note: This is similar to `recvfrom` in POSIX, though it also supports
+  /// reading the data into multiple buffers in the manner of `readv`.
+  ///
+  /// @param[in] RiData List of scatter/gather vectors to which to store data.
+  /// @param[in] RiFlags Message flags.
+  /// @param[in] Address Address of the target.
+  /// @param[in] AddressLength The buffer size of Address.
+  /// @param[out] NRead Return the number of bytes stored in RiData.
+  /// @param[out] RoFlags Return message flags.
+  /// @return Nothing or WASI error.
+  WasiExpect<void> sockRecvFrom(Span<Span<uint8_t>> RiData,
+                                __wasi_riflags_t RiFlags, uint8_t *Address,
+                                uint8_t AddressLength, __wasi_size_t &NRead,
+                                __wasi_roflags_t &RoFlags) const noexcept {
+    return Node.sockRecvFrom(RiData, RiFlags, Address, AddressLength, NRead,
+                             RoFlags);
+  }
+
   /// Send a message on a socket.
   ///
   /// Note: This is similar to `send` in POSIX, though it also supports writing

diff --git a/include/host/wasi/wasifunc.h b/include/host/wasi/wasifunc.h
@@ -417,9 +417,10 @@ class WasiSockRecvFrom : public Wasi<WasiSockRecvFrom> {
   WasiSockRecvFrom(WASI::Environ &HostEnv) : Wasi(HostEnv) {}
 
   Expect<uint32_t> body(Runtime::Instance::MemoryInstance *MemInst, int32_t Fd,
-                        uint32_t RiDataPtr, int32_t RiDataLen,
+                        uint32_t RiDataPtr, uint32_t RiDataLen,
                         uint32_t AddressPtr, uint32_t RiFlags,
-                        uint32_t /* Out */ RoDataLenPtr);
+                        uint32_t /* Out */ RoDataLenPtr,
+                        uint32_t /* Out */ RoFlagsPtr);
 };
 
 class WasiSockSend : public Wasi<WasiSockSend> {

diff --git a/lib/host/wasi/inode-linux.cpp b/lib/host/wasi/inode-linux.cpp
@@ -991,6 +991,14 @@ WasiExpect<void> INode::sockConnect(uint8_t *Address, uint8_t AddressLength,
 WasiExpect<void> INode::sockRecv(Span<Span<uint8_t>> RiData,
                                  __wasi_riflags_t RiFlags, __wasi_size_t &NRead,
                                  __wasi_roflags_t &RoFlags) const noexcept {
+  return sockRecvFrom(RiData, RiFlags, nullptr, 0, NRead, RoFlags);
+}
+
+WasiExpect<void> INode::sockRecvFrom(Span<Span<uint8_t>> RiData,
+                                     __wasi_riflags_t RiFlags, uint8_t *Address,
+                                     uint8_t AddressLength,
+                                     __wasi_size_t &NRead,
+                                     __wasi_roflags_t &RoFlags) const noexcept {
   int SysRiFlags = 0;
   if (RiFlags & __WASI_RIFLAGS_RECV_PEEK) {
     SysRiFlags |= MSG_PEEK;
@@ -1008,8 +1016,8 @@ WasiExpect<void> INode::sockRecv(Span<Span<uint8_t>> RiData,
   }
 
   msghdr SysMsgHdr;
-  SysMsgHdr.msg_name = nullptr;
-  SysMsgHdr.msg_namelen = 0;
+  SysMsgHdr.msg_name = Address;
+  SysMsgHdr.msg_namelen = AddressLength;
   SysMsgHdr.msg_iov = SysIOVs;
   SysMsgHdr.msg_iovlen = SysIOVsSize;
   SysMsgHdr.msg_control = nullptr;
@@ -1032,41 +1040,16 @@ WasiExpect<void> INode::sockRecv(Span<Span<uint8_t>> RiData,
 }
 
 WasiExpect<void> INode::sockSend(Span<Span<const uint8_t>> SiData,
-                                 __wasi_siflags_t,
+                                 __wasi_siflags_t SiFlags,
                                  __wasi_size_t &NWritten) const noexcept {
-  int SysSiFlags = MSG_NOSIGNAL;
-
-  iovec SysIOVs[kIOVMax];
-  size_t SysIOVsSize = 0;
-  for (auto &IOV : SiData) {
-    SysIOVs[SysIOVsSize].iov_base = const_cast<uint8_t *>(IOV.data());
-    SysIOVs[SysIOVsSize].iov_len = IOV.size();
-    ++SysIOVsSize;
-  }
-
-  msghdr SysMsgHdr;
-  SysMsgHdr.msg_name = nullptr;
-  SysMsgHdr.msg_namelen = 0;
-  SysMsgHdr.msg_iov = SysIOVs;
-  SysMsgHdr.msg_iovlen = SysIOVsSize;
-  SysMsgHdr.msg_control = nullptr;
-  SysMsgHdr.msg_controllen = 0;
-
-  // Store recv bytes length and flags.
-  if (auto Res = ::sendmsg(Fd, &SysMsgHdr, SysSiFlags); unlikely(Res < 0)) {
-    return WasiUnexpect(fromErrNo(errno));
-  } else {
-    NWritten = Res;
-  }
-
-  return {};
+  return sockSendTo(SiData, SiFlags, nullptr, 0, NWritten);
 }
 
 WasiExpect<void> INode::sockSendTo(Span<Span<const uint8_t>> SiData,
                                    __wasi_siflags_t, uint8_t *Address,
                                    uint8_t AddressLength,
                                    __wasi_size_t &NWritten) const noexcept {
-  int SysSiFlags = 0;
+  int SysSiFlags = MSG_NOSIGNAL;
 
   iovec SysIOVs[kIOVMax];
   size_t SysIOVsSize = 0;

diff --git a/lib/host/wasi/inode-macos.cpp b/lib/host/wasi/inode-macos.cpp
@@ -872,6 +872,14 @@ WasiExpect<void> INode::sockConnect(uint8_t *Address, uint8_t AddressLength,
 WasiExpect<void> INode::sockRecv(Span<Span<uint8_t>> RiData,
                                  __wasi_riflags_t RiFlags, __wasi_size_t &NRead,
                                  __wasi_roflags_t &RoFlags) const noexcept {
+  return sockRecvFrom(RiData, RiFlags, nullptr, 0, NRead, RoFlags);
+}
+
+WasiExpect<void> INode::sockRecvFrom(Span<Span<uint8_t>> RiData,
+                                     __wasi_riflags_t RiFlags, uint8_t *Address,
+                                     uint8_t AddressLength,
+                                     __wasi_size_t &NRead,
+                                     __wasi_roflags_t &RoFlags) const noexcept {
   int SysRiFlags = 0;
   if (RiFlags & __WASI_RIFLAGS_RECV_PEEK) {
     SysRiFlags |= MSG_PEEK;
@@ -889,8 +897,8 @@ WasiExpect<void> INode::sockRecv(Span<Span<uint8_t>> RiData,
   }
 
   msghdr SysMsgHdr;
-  SysMsgHdr.msg_name = nullptr;
-  SysMsgHdr.msg_namelen = 0;
+  SysMsgHdr.msg_name = Address;
+  SysMsgHdr.msg_namelen = AddressLength;
   SysMsgHdr.msg_iov = SysIOVs;
   SysMsgHdr.msg_iovlen = SysIOVsSize;
   SysMsgHdr.msg_control = nullptr;
@@ -913,34 +921,9 @@ WasiExpect<void> INode::sockRecv(Span<Span<uint8_t>> RiData,
 }
 
 WasiExpect<void> INode::sockSend(Span<Span<const uint8_t>> SiData,
-                                 __wasi_siflags_t,
+                                 __wasi_siflags_t SiFlags,
                                  __wasi_size_t &NWritten) const noexcept {
-  int SysSiFlags = 0;
-
-  iovec SysIOVs[kIOVMax];
-  size_t SysIOVsSize = 0;
-  for (auto &IOV : SiData) {
-    SysIOVs[SysIOVsSize].iov_base = const_cast<uint8_t *>(IOV.data());
-    SysIOVs[SysIOVsSize].iov_len = IOV.size();
-    ++SysIOVsSize;
-  }
-
-  msghdr SysMsgHdr;
-  SysMsgHdr.msg_name = nullptr;
-  SysMsgHdr.msg_namelen = 0;
-  SysMsgHdr.msg_iov = SysIOVs;
-  SysMsgHdr.msg_iovlen = SysIOVsSize;
-  SysMsgHdr.msg_control = nullptr;
-  SysMsgHdr.msg_controllen = 0;
-
-  // Store recv bytes length and flags.
-  if (auto Res = ::sendmsg(Fd, &SysMsgHdr, SysSiFlags); unlikely(Res < 0)) {
-    return WasiUnexpect(fromErrNo(errno));
-  } else {
-    NWritten = Res;
-  }
-
-  return {};
+  return sockSendTo(SiData, SiFlags, nullptr, 0, NWritten);
 }
 
 WasiExpect<void> INode::sockSendTo(Span<Span<const uint8_t>> SiData,

diff --git a/lib/host/wasi/inode-win.cpp b/lib/host/wasi/inode-win.cpp
@@ -211,6 +211,12 @@ WasiExpect<void> INode::sockRecv(Span<Span<uint8_t>>, __wasi_riflags_t,
   return WasiUnexpect(__WASI_ERRNO_NOSYS);
 }
 
+WasiExpect<void> INode::sockRecvFrom(Span<Span<uint8_t>>, __wasi_riflags_t,
+                                     uint8_t *, uint8_t, __wasi_size_t &,
+                                     __wasi_roflags_t &) const noexcept {
+  return WasiUnexpect(__WASI_ERRNO_NOSYS);
+}
+
 WasiExpect<void> INode::sockSend(Span<Span<const uint8_t>>, __wasi_siflags_t,
                                  __wasi_size_t &) const noexcept {
   return WasiUnexpect(__WASI_ERRNO_NOSYS);

diff --git a/lib/host/wasi/wasifunc.cpp b/lib/host/wasi/wasifunc.cpp
@@ -1804,6 +1804,93 @@ Expect<uint32_t> WasiSockRecv::body(Runtime::Instance::MemoryInstance *MemInst,
   return __WASI_ERRNO_SUCCESS;
 }
 
+Expect<uint32_t> WasiSockRecvFrom::body(
+    Runtime::Instance::MemoryInstance *MemInst, int32_t Fd, uint32_t RiDataPtr,
+    uint32_t RiDataLen, uint32_t AddressPtr, uint32_t RiFlags,
+    uint32_t /* Out */ RoDataLenPtr, uint32_t /* Out */ RoFlagsPtr) {
+  // Check memory instance from module.
+  if (MemInst == nullptr) {
+    return __WASI_ERRNO_FAULT;
+  }
+
+  __wasi_address_t *InnerAddress = MemInst->getPointer<__wasi_address_t *>(
+      AddressPtr, sizeof(__wasi_address_t));
+  if (InnerAddress == nullptr) {
+    return __WASI_ERRNO_FAULT;
+  }
+
+  if (InnerAddress->buf_len != 4 && InnerAddress->buf_len != 16) {
+    return __WASI_ERRNO_INVAL;
+  }
+
+  uint8_t *AddressBuf = MemInst->getPointer<uint8_t *>(
+      InnerAddress->buf, sizeof(uint8_t) * InnerAddress->buf_len);
+  if (AddressBuf == nullptr) {
+    return __WASI_ERRNO_FAULT;
+  }
+
+  __wasi_riflags_t WasiRiFlags;
+  if (auto Res = cast<__wasi_riflags_t>(RiFlags); unlikely(!Res)) {
+    return Res.error();
+  } else {
+    WasiRiFlags = *Res;
+  }
+
+  const __wasi_size_t WasiRiDataLen = RiDataLen;
+  if (unlikely(WasiRiDataLen > WASI::kIOVMax)) {
+    return __WASI_ERRNO_INVAL;
+  }
+
+  // Check for invalid address.
+  auto *const RiDataArray =
+      MemInst->getPointer<__wasi_iovec_t *>(RiDataPtr, WasiRiDataLen);
+  if (unlikely(RiDataArray == nullptr)) {
+    return __WASI_ERRNO_FAULT;
+  }
+
+  auto *const RoDataLen = MemInst->getPointer<__wasi_size_t *>(RoDataLenPtr);
+  if (unlikely(RoDataLen == nullptr)) {
+    return __WASI_ERRNO_FAULT;
+  }
+
+  auto *const RoFlags = MemInst->getPointer<__wasi_roflags_t *>(RoFlagsPtr);
+  if (unlikely(RoFlags == nullptr)) {
+    return __WASI_ERRNO_FAULT;
+  }
+  __wasi_size_t TotalSize = 0;
+  std::array<Span<uint8_t>, WASI::kIOVMax> WasiRiData;
+
+  for (__wasi_size_t I = 0; I < WasiRiDataLen; ++I) {
+    __wasi_iovec_t &RiData = RiDataArray[I];
+
+    // Capping total size.
+    const __wasi_size_t Space =
+        std::numeric_limits<__wasi_size_t>::max() - TotalSize;
+    const __wasi_size_t BufLen =
+        unlikely(RiData.buf_len > Space) ? Space : RiData.buf_len;
+    TotalSize += BufLen;
+
+    // Check for invalid address.
+    auto *const RiDataArr = MemInst->getPointer<uint8_t *>(RiData.buf, BufLen);
+    // Check for invalid address.
+    if (unlikely(RiDataArr == nullptr)) {
+      return __WASI_ERRNO_FAULT;
+    }
+    WasiRiData[I] = {RiDataArr, BufLen};
+  }
+
+  const __wasi_fd_t WasiFd = Fd;
+
+  if (auto Res = Env.sockRecvFrom(
+          WasiFd, {WasiRiData.data(), WasiRiDataLen}, WasiRiFlags, AddressBuf,
+          static_cast<uint8_t>(InnerAddress->buf_len), *RoDataLen, *RoFlags);
+      unlikely(!Res)) {
+    return Res.error();
+  }
+
+  return __WASI_ERRNO_SUCCESS;
+}
+
 Expect<uint32_t> WasiSockSend::body(Runtime::Instance::MemoryInstance *MemInst,
                                     int32_t Fd, uint32_t SiDataPtr,
                                     uint32_t SiDataLen, uint32_t SiFlags,

diff --git a/lib/host/wasi/wasimodule.cpp b/lib/host/wasi/wasimodule.cpp
@@ -65,6 +65,7 @@ WasiModule::WasiModule() : ImportObject("wasi_snapshot_preview1") {
   addHostFunc("sock_listen", std::make_unique<WasiSockListen>(Env));
   addHostFunc("sock_accept", std::make_unique<WasiSockAccept>(Env));
   addHostFunc("sock_recv", std::make_unique<WasiSockRecv>(Env));
+  addHostFunc("sock_recv_from", std::make_unique<WasiSockRecvFrom>(Env));
   addHostFunc("sock_send", std::make_unique<WasiSockSend>(Env));
   addHostFunc("sock_send_to", std::make_unique<WasiSockSendTo>(Env));
   addHostFunc("sock_shutdown", std::make_unique<WasiSockShutdown>(Env));