Skip to content

Commit

Permalink
Merge pull request #5 from erikziyunchi/safer-lib
Browse files Browse the repository at this point in the history 
Making the library code safer
  • Loading branch information
@erikziyunchi
erikziyunchi committed Nov 1, 2023
2 parents bbe6daf + 97d68e7 commit 6611747
Show file tree
Hide file tree
Showing 7 changed files with 105 additions and 35 deletions.
41 changes: 31 additions & 10 deletions crates/water/src/runtime/core.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,10 @@ impl H2O<Host> {
// store.data_mut().preview1_ctx = Some(WasiCtxBuilder::new().inherit_stdio().preopened_dir(path, ".")?.build());
store.data_mut().preview1_ctx = Some(WasiCtxBuilder::new().inherit_stdio().build());

if store.data().preview1_ctx.is_none() {
return Err(anyhow::anyhow!("Failed to retrieve preview1_ctx from Host"));
}

wasmtime_wasi::add_to_linker(&mut linker, |h: &mut Host| h.preview1_ctx.as_mut().unwrap())?;

// initializing stuff for multithreading
Expand All @@ -65,30 +69,37 @@ impl H2O<Host> {
)?));

wasmtime_wasi_threads::add_to_linker(&mut linker, &store, &module, |h: &mut Host| {
h.wasi_threads.as_ref().unwrap()
h.wasi_threads
.as_ref()
.context("Failed to get ref of wasi_threads from Host")?
})?;
}

// export functions -- version dependent -- has to be done before instantiate
match &version {
Some(Version::V0) => {
v0::funcs::export_tcp_connect(&mut linker);
v0::funcs::export_tcplistener_create(&mut linker);
v0::funcs::export_tcp_connect(&mut linker)?;
v0::funcs::export_tcplistener_create(&mut linker)?;
}
Some(Version::V1) => {
v1::funcs::export_tcp_connect(&mut linker);
v1::funcs::export_tcplistener_create(&mut linker);
v1::funcs::export_tcp_connect(&mut linker)?;
v1::funcs::export_tcplistener_create(&mut linker)?;
}
_ => {} // add export funcs for other versions here
}

// export functions -- version independent
version_common::funcs::export_config(&mut linker, conf.config_wasm.clone());
version_common::funcs::export_config(&mut linker, conf.config_wasm.clone())?;

let instance = linker.instantiate(&mut store, &module)?;

Ok(H2O {
version: version.unwrap(),
version: match version {
Some(v) => v,
None => {
return Err(anyhow::anyhow!("Version is None"));
}
},

engine,
linker,
Expand Down Expand Up @@ -138,8 +149,13 @@ impl H2O<Host> {

// Obtain the directory path and file name from config_wasm
let full_path = Path::new(&config.config_wasm);
let parent_dir = full_path.parent().unwrap(); // Assumes config_wasm has a parent directory
let file_name = full_path.file_name().unwrap().to_str().unwrap(); // Assumes file_name is valid UTF-8
let parent_dir = full_path
.parent()
.ok_or_else(|| anyhow::anyhow!("config_wasm does not have a parent directory"))?; // Assumes config_wasm has a parent directory
let file_name = full_path
.file_name()
.and_then(|os_str| os_str.to_str())
.ok_or_else(|| anyhow::anyhow!("file_name is not valid UTF-8"))?; // Assumes file_name is valid UTF-8

// Open the parent directory
let dir = Dir::open_ambient_dir(parent_dir, ambient_authority())?;
Expand All @@ -148,7 +164,12 @@ impl H2O<Host> {

let wasi_file = wasmtime_wasi::sync::file::File::from_cap_std(wasi_file);

let ctx = self.store.data_mut().preview1_ctx.as_mut().unwrap();
let ctx = self
.store
.data_mut()
.preview1_ctx
.as_mut()
.ok_or(anyhow::anyhow!("preview1_ctx in Store is None"))?;
let config_fd = ctx.push_file(Box::new(wasi_file), FileAccessMode::all())? as i32;

let params = vec![Val::I32(config_fd); config_fn.ty(&self.store).params().len()];
Expand Down
10 changes: 8 additions & 2 deletions crates/water/src/runtime/listener.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,8 @@ impl WATERListener<Host> {
.core
.instance
.get_func(&mut self.core.store, &conf.entry_fn)
.unwrap();
.context(format!("Failed to get function {}", &conf.entry_fn))?;

match fnc.call(&mut self.core.store, &[], &mut []) {
Ok(_) => {}
Err(e) => {
Expand Down Expand Up @@ -153,7 +154,12 @@ impl WATERListener<Host> {
std::mem::forget(water_writer);
std::mem::forget(water_reader);

let ctx = core.store.data_mut().preview1_ctx.as_mut().unwrap();
let ctx = core
.store
.data_mut()
.preview1_ctx
.as_mut()
.ok_or(anyhow::anyhow!("preview1_ctx in Store is None"))?;
let water_reader_fd = ctx.push_file(Box::new(wasi_water_reader), FileAccessMode::all())?;
let water_writer_fd = ctx.push_file(Box::new(wasi_water_writer), FileAccessMode::all())?;

Expand Down
1 change: 0 additions & 1 deletion crates/water/src/runtime/stream.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,6 @@ impl WATERStream<Host> {
info!("[HOST] WATERStream connecting...");

// TODO: add addr:port sharing with WASM, for now WASM is using config.json's remote_addr:port
// let fnc = self.core.instance.get_func(&mut self.core.store, &conf.entry_fn).unwrap();
let fnc = match self.core.instance.get_func(&mut self.core.store, DIAL_FN) {
Some(func) => func,
None => {
Expand Down
41 changes: 31 additions & 10 deletions crates/water/src/runtime/v0/funcs.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
use anyhow::Ok;

use crate::config::wasm_shared_config::StreamConfig;
use crate::runtime::*;
use std::convert::TryInto;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};

// TODO: rename this to dial_v1, since it has the ability to let WASM choose ip:port
pub fn export_tcp_connect(linker: &mut Linker<Host>) {
pub fn export_tcp_connect(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
linker
.func_wrap(
"env",
Expand Down Expand Up @@ -48,22 +50,33 @@ pub fn export_tcp_connect(linker: &mut Linker<Host>) {
addr => std::net::TcpStream::connect(addr),
}
.map(TcpStream::from_std)
.context("failed to connect to endpoint")
.context(format!(
"Failed to connect to {}:{} in Host exported dial",
host, port
))
.unwrap();

// Connecting Tcp
let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();

// Get the WasiCtx of the caller(WASM), then insert_file into it
let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
let ctx: &mut WasiCtx = caller
.data_mut()
.preview1_ctx
.as_mut()
.context("preview1_ctx in Store is None")
.unwrap();
ctx.push_file(socket_file, FileAccessMode::all())
.context("Failed to push file into WASM")
.unwrap() as i32
},
)
.unwrap();
.context("Failed to export Dial function to WASM")?;
Ok(())
}

// TODO: rename this to dial_v1, since it has the ability to let WASM listen on a TcpListener
pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
pub fn export_tcplistener_create(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
linker
.func_wrap(
"env",
Expand All @@ -89,7 +102,7 @@ pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
name: config.name.clone().try_into().unwrap(),
port: config.port as u16,
addr: config.addr.clone(),
}); // TODO: add addr here
});

// Get the pair here addr:port
let (addr, port) = match listener_file {
Expand All @@ -108,11 +121,19 @@ pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();

// Get the WasiCtx of the caller(WASM), then insert_file into it
let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
let ctx: &mut WasiCtx = caller
.data_mut()
.preview1_ctx
.as_mut()
.context("preview1_ctx in Store is None")
.unwrap();
ctx.push_file(socket_file, FileAccessMode::all())
.context("Failed to push file into WASM")
.unwrap() as i32
},
)
.unwrap();
.context("Failed to export TcpListener create function to WASM")?;
Ok(())
}

// Generically link dial functions
Expand Down
40 changes: 31 additions & 9 deletions crates/water/src/runtime/v1/funcs.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
use anyhow::Ok;

use crate::config::wasm_shared_config::StreamConfig;
use crate::runtime::*;
use std::convert::TryInto;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};

// TODO: rename this to dial_v1, since it has the ability to let WASM choose ip:port
pub fn export_tcp_connect(linker: &mut Linker<Host>) {
pub fn export_tcp_connect(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
linker
.func_wrap(
"env",
Expand Down Expand Up @@ -48,22 +50,33 @@ pub fn export_tcp_connect(linker: &mut Linker<Host>) {
addr => std::net::TcpStream::connect(addr),
}
.map(TcpStream::from_std)
.context("failed to connect to endpoint")
.context(format!(
"Failed to connect to {}:{} in Host exported dial",
host, port
))
.unwrap();

// Connecting Tcp
let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();

// Get the WasiCtx of the caller(WASM), then insert_file into it
let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
let ctx: &mut WasiCtx = caller
.data_mut()
.preview1_ctx
.as_mut()
.context("preview1_ctx in Store is None")
.unwrap();
ctx.push_file(socket_file, FileAccessMode::all())
.context("Failed to push file into WASM")
.unwrap() as i32
},
)
.unwrap();
.context("Failed to export Dial function to WASM")?;
Ok(())
}

// TODO: rename this to dial_v1, since it has the ability to let WASM listen on a TcpListener
pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
pub fn export_tcplistener_create(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
linker
.func_wrap(
"env",
Expand Down Expand Up @@ -104,14 +117,23 @@ pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
// Creating Tcp Listener
let tcp = std::net::TcpListener::bind((addr.as_str(), port)).unwrap();
let tcp = TcpListener::from_std(tcp);
// tcp.set_nonblocking(true);
let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();

// Get the WasiCtx of the caller(WASM), then insert_file into it
let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
let ctx: &mut WasiCtx = caller
.data_mut()
.preview1_ctx
.as_mut()
.context("preview1_ctx in Store is None")
.unwrap();
ctx.push_file(socket_file, FileAccessMode::all())
.context("Failed to push file into WASM")
.unwrap() as i32
},
)
.unwrap();
.context("Failed to export TcpListener create function to WASM")?;
Ok(())
}

// Generically link dial functions
Expand Down
5 changes: 3 additions & 2 deletions crates/water/src/runtime/version_common/funcs.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::runtime::*;

// exportint a function for WASM to get CONFIG file
pub fn export_config(linker: &mut Linker<Host>, config_file: String) {
pub fn export_config(linker: &mut Linker<Host>, config_file: String) -> Result<(), anyhow::Error> {
linker
.func_wrap(
"env",
Expand All @@ -22,5 +22,6 @@ pub fn export_config(linker: &mut Linker<Host>, config_file: String) {
.expect("Error with pushing file") as i32
},
)
.unwrap();
.context("Failed to export config function to WASM")?;
Ok(())
}
2 changes: 1 addition & 1 deletion tests/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,4 +20,4 @@ shadowsocks-service = {version = "1.17.0", features = ["server"]}
shadowsocks-rust = "1.17.0"
tokio = { version = "1.24.2", features = ["full", "macros"] }
futures = "0.3.28"
tempfile = "3.8.0"
tempfile = "3.8.0"

0 comments on commit 6611747

Please sign in to comment.