From de68cc1726d1c0533cc837a1ce32740b607156a4 Mon Sep 17 00:00:00 2001
From: Ivan Font <ivan@ivanfont.com>
Date: Thu, 9 Feb 2023 16:22:11 -0800
Subject: [PATCH] Add support for WASI sockets to C API (#5624)

* Add support for WASI sockets to C API

Add support for WASI sockets in the C API by adding a new API to handle
preopening sockets for clients. This uses HashMap instead of Vec for
preopened sockets to identify if caller has called in more than once
with the same FD number. If so, then we return false so caller is given
hint that they are attempting to overwrite an already existing socket
FD.

* Apply suggestions from code review

Co-authored-by: Peter Huene <peter@huene.dev>

* s/stdlistener/listener/

---------

Co-authored-by: Peter Huene <peter@huene.dev>
---
 crates/c-api/include/wasi.h | 14 +++++++++++
 crates/c-api/src/wasi.rs    | 48 +++++++++++++++++++++++++++++++++----
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/crates/c-api/include/wasi.h b/crates/c-api/include/wasi.h
index f3d1dfd452..e927e04a44 100644
--- a/crates/c-api/include/wasi.h
+++ b/crates/c-api/include/wasi.h
@@ -7,6 +7,7 @@
 #ifndef WASI_H
 #define WASI_H
 
+#include <stdint.h>
 #include "wasm.h"
 
 #ifndef WASI_API_EXTERN
@@ -156,6 +157,19 @@ WASI_API_EXTERN void wasi_config_inherit_stderr(wasi_config_t* config);
  */
 WASI_API_EXTERN bool wasi_config_preopen_dir(wasi_config_t* config, const char* path, const char* guest_path);
 
+/**
+ * \brief Configures a "preopened" listen socket to be available to WASI APIs.
+ *
+ * By default WASI programs do not have access to open up network sockets on
+ * the host. This API can be used to grant WASI programs access to a network
+ * socket file descriptor on the host.
+ *
+ * The fd_num argument is the number of the file descriptor by which it will be
+ * known in WASM and the host_port is the IP address and port (e.g.
+ * "127.0.0.1:8080") requested to listen on.
+ */
+WASI_API_EXTERN bool wasi_config_preopen_socket(wasi_config_t* config, uint32_t fd_num, const char* host_port);
+
 #undef own
 
 #ifdef __cplusplus
diff --git a/crates/c-api/src/wasi.rs b/crates/c-api/src/wasi.rs
index 3eb6f3bd49..e02063ed21 100644
--- a/crates/c-api/src/wasi.rs
+++ b/crates/c-api/src/wasi.rs
@@ -3,6 +3,7 @@
 use crate::wasm_byte_vec_t;
 use anyhow::Result;
 use cap_std::ambient_authority;
+use std::collections::HashMap;
 use std::ffi::CStr;
 use std::fs::File;
 use std::os::raw::{c_char, c_int};
@@ -10,7 +11,7 @@ use std::path::{Path, PathBuf};
 use std::slice;
 use wasi_common::pipe::ReadPipe;
 use wasmtime_wasi::{
-    sync::{Dir, WasiCtxBuilder},
+    sync::{Dir, TcpListener, WasiCtxBuilder},
     WasiCtx,
 };
 
@@ -18,6 +19,10 @@ unsafe fn cstr_to_path<'a>(path: *const c_char) -> Option<&'a Path> {
     CStr::from_ptr(path).to_str().map(Path::new).ok()
 }
 
+unsafe fn cstr_to_str<'a>(s: *const c_char) -> Option<&'a str> {
+    CStr::from_ptr(s).to_str().ok()
+}
+
 unsafe fn open_file(path: *const c_char) -> Option<File> {
     File::open(cstr_to_path(path)?).ok()
 }
@@ -34,7 +39,8 @@ pub struct wasi_config_t {
     stdin: WasiConfigReadPipe,
     stdout: WasiConfigWritePipe,
     stderr: WasiConfigWritePipe,
-    preopens: Vec<(Dir, PathBuf)>,
+    preopen_dirs: Vec<(Dir, PathBuf)>,
+    preopen_sockets: HashMap<u32, TcpListener>,
     inherit_args: bool,
     inherit_env: bool,
 }
@@ -118,9 +124,12 @@ impl wasi_config_t {
                 builder.stderr(Box::new(file))
             }
         };
-        for (dir, path) in self.preopens {
+        for (dir, path) in self.preopen_dirs {
             builder = builder.preopened_dir(dir, path)?;
         }
+        for (fd_num, listener) in self.preopen_sockets {
+            builder = builder.preopened_socket(fd_num, listener)?;
+        }
         Ok(builder.build())
     }
 }
@@ -266,7 +275,38 @@ pub unsafe extern "C" fn wasi_config_preopen_dir(
         None => return false,
     };
 
-    (*config).preopens.push((dir, guest_path.to_owned()));
+    (*config).preopen_dirs.push((dir, guest_path.to_owned()));
+
+    true
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn wasi_config_preopen_socket(
+    config: &mut wasi_config_t,
+    fd_num: u32,
+    host_port: *const c_char,
+) -> bool {
+    let address = match cstr_to_str(host_port) {
+        Some(s) => s,
+        None => return false,
+    };
+    let listener = match std::net::TcpListener::bind(address) {
+        Ok(listener) => listener,
+        Err(_) => return false,
+    };
+
+    if let Err(_) = listener.set_nonblocking(true) {
+        return false;
+    }
+
+    // Caller cannot call in more than once with the same FD number so return an error.
+    if (*config).preopen_sockets.contains_key(&fd_num) {
+        return false;
+    }
+
+    (*config)
+        .preopen_sockets
+        .insert(fd_num, TcpListener::from_std(listener));
 
     true
 }