Add network call.

Cargo.toml

@@ -27,6 +27,8 @@ members = [
     "crates/harness/runner",
     "crates/harness/plot",
     "crates/tlsn",
+    "crates/wasmtime-plugin",
+    "crates/wasmtime-host",
 ]
 resolver = "2"
 

crates/wasmtime-host/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "wasmtime-host"
+version = "0.1.0"
+edition = "2024"
+
+[lints]
+workspace = true
+
+[dependencies]
+anyhow = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread", "net", "io-util"] }
+wasmtime = { git = "https://github.com/bytecodealliance/wasmtime", branch = "main", features = ["component-model-async"] }

crates/wasmtime-host/src/main.rs

@@ -0,0 +1,149 @@
+use std::{pin::Pin, task::{Context, Waker}};
+
+use anyhow::{Error};
+use serde::{Deserialize, Serialize};
+use tokio::{io::{AsyncWrite}, net::TcpStream};
+use wasmtime::{component::{bindgen, Component, HasSelf, Linker, Resource, ResourceTable}, Config, Engine, Store, Result};
+
+use component::wasmtime_plugin::io::{Host, HostNetworkIo};
+
+const WASM_PLUGIN_PATH: &str = "../../target/wasm32-unknown-unknown/release/wasmtime_plugin_component.wasm";
+
+bindgen!({
+    path: "../wasmtime-plugin/wit/plugin.wit",
+    imports: { 
+        default: async | trappable,
+    },
+    exports: {
+        default: async,
+    },
+    with: {
+        "component:wasmtime-plugin/io/network-io": NetworkIo,
+    },
+});
+
+#[derive(Serialize, Debug)]
+struct Input {
+    host: String,
+    port: u32,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+struct Output {
+    result: bool
+}
+
+pub struct NetworkIo {
+    inner: TcpStream
+}
+
+#[derive(Default)]
+struct HostState {
+    table: ResourceTable,
+}
+
+impl Host for HostState {}
+
+impl HostNetworkIo for HostState {
+    async fn new(&mut self, host: String, port: u32) -> Result<Resource<NetworkIo>> {
+        let connection = TcpStream::connect(format!("{host}:{port}"))
+            .await
+            .map_err(|err| Error::msg(err))?;
+
+        println!("Connection established");
+
+        let id = self.table.push(NetworkIo { inner: connection })?;
+        Ok(id)
+    }
+
+    async fn read(&mut self, network_io: Resource<NetworkIo>, max: u32) -> Result<Vec<u8>> {
+        debug_assert!(!network_io.owned());
+
+        println!("Trying to reads");
+
+        let conn = self.table.get(&network_io)?;
+        let mut buf = vec![0u8; max as usize];
+        loop {
+            conn.inner.readable().await?;
+            match conn.inner.try_read(&mut buf) {
+                Ok(0) => return Ok(Vec::new()),       // EOF
+                Ok(n) => return Ok(buf[..n].to_vec()),// got bytes
+                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => continue,
+                Err(e) => return Err(e.into()),
+            }
+        }
+    }
+
+    async fn write(&mut self, network_io: Resource<NetworkIo>, bytes: Vec<u8>) -> Result<u32> {
+        debug_assert!(!network_io.owned());
+
+        println!("Trying to write");
+
+        let conn = self.table.get(&network_io)?;
+        let mut offset = 0;
+        while offset < bytes.len() {
+            conn.inner.writable().await?;
+            match conn.inner.try_write(&bytes[offset..]) {
+                Ok(0) => break,
+                Ok(n) => offset += n,
+                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => continue,
+                Err(e) => return Err(e.into()),
+            }
+        }
+        Ok(offset as u32)
+    }
+
+    async  fn shutdown(&mut self, network_io: Resource<NetworkIo>,) -> Result<()> {
+        debug_assert!(!network_io.owned());
+
+        println!("Trying to shut down");
+
+        let conn = self.table.get_mut(&network_io)?;
+        let mut context = Context::from_waker(Waker::noop());
+        
+        let _ = Pin::new(&mut conn.inner).poll_shutdown(&mut context);
+
+        Ok(())
+    }
+
+    async fn drop(&mut self, network_io: Resource<NetworkIo>) -> Result<()> {
+        debug_assert!(network_io.owned());
+
+        println!("Trying to drop");
+
+        self.table.delete(network_io)?;
+        Ok(())
+    }
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    println!("Starting wasmtime");
+
+    let mut config = Config::new();
+    config
+        .async_support(true)
+        .wasm_component_model_async(true);
+
+    let engine = Engine::new(&config)?;
+
+    let component = Component::from_file(&engine, WASM_PLUGIN_PATH)?;
+    let mut linker = Linker::new(&engine);
+    
+    Plugin::add_to_linker::<_, HasSelf<_>>(&mut linker, |state| state)?;
+
+    let mut store = Store::new(&engine, HostState::default());
+
+    let plugin = Plugin::instantiate_async(&mut store, &component, &linker).await?;
+
+    let input = Input { host: "0.0.0.0".into(), port: 7044 };
+    println!("Inputs: {:?}", input);
+    let input_bytes = serde_json::to_vec(&input)?;
+
+    let res_byte = plugin.call_main(&mut store, &input_bytes).await?;
+    let res: Output = serde_json::from_slice(&res_byte)?;
+
+    println!("Output: {:?}", res);
+
+    Ok(())
+}

crates/wasmtime-plugin/.cargo/config.toml

@@ -0,0 +1,5 @@
+[build]
+target = "wasm32-unknown-unknown"
+
+[unstable]
+build-std = ["panic_abort", "std"]

crates/wasmtime-plugin/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "wasmtime-plugin"
+version = "0.1.0"
+edition = "2024"
+
+[lints]
+workspace = true
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+http-body-util = { workspace = true }
+hyper = { workspace = true, features = ["client", "http1"] }
+futures = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+wit-bindgen = { version = "0.43.0" }

crates/wasmtime-plugin/build.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+PACKAGE_NAME="wasmtime_plugin"
+
+set -e
+
+# Build the project
+cargo build --release
+
+# Convert the wasm binary to a component
+wasm-tools component new ../../target/wasm32-unknown-unknown/release/${PACKAGE_NAME}.wasm -o ../../target/wasm32-unknown-unknown/release/${PACKAGE_NAME}_component.wasm
+
+echo "Component created: ${PACKAGE_NAME}_component.wasm"

crates/wasmtime-plugin/src/lib.rs

@@ -0,0 +1,116 @@
+use core::slice;
+use http_body_util::Empty;
+use hyper::{body::Bytes, client::conn::http1, Request, StatusCode};
+use serde::{Deserialize, Serialize};
+use wit_bindgen::spawn;
+use std::{pin::Pin, task::{Context, Poll}};
+
+use crate::component::wasmtime_plugin::io::NetworkIo;
+
+wit_bindgen::generate!({
+    path: "wit/plugin.wit",
+    async: true
+});
+
+struct Component;
+
+#[derive(Deserialize)]
+struct Input {
+    host: String,
+    port: u32,
+}
+
+#[derive(Serialize)]
+struct Output {
+    result: bool
+}
+
+struct HyperIo {
+    inner: NetworkIo
+}
+
+impl hyper::rt::Write for HyperIo
+{
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        let fut = self.inner.write(buf.to_vec());
+        match std::pin::pin!(fut).poll(cx) {
+            std::task::Poll::Pending => std::task::Poll::Pending,
+            std::task::Poll::Ready(n) => {
+                std::task::Poll::Ready(Ok(n as usize))
+            }
+        }
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>
+    ) -> Poll<Result<(), std::io::Error>> {
+        std::task::Poll::Ready(Ok(()))
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let fut = self.inner.shutdown();
+        match std::pin::pin!(fut).poll(cx) {
+            std::task::Poll::Pending => std::task::Poll::Pending,
+            std::task::Poll::Ready(_) => std::task::Poll::Ready(Ok(()))
+        }
+    }
+}
+
+impl hyper::rt::Read for HyperIo
+{
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        mut buf: hyper::rt::ReadBufCursor<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        let buf_len = unsafe { buf.as_mut().len() };
+        let buf_slice = unsafe {
+            slice::from_raw_parts_mut(buf.as_mut().as_mut_ptr() as *mut u8, buf_len)
+        };
+        let fut = self.inner.read(buf_len as u32);
+        match std::pin::pin!(fut).poll(cx) {
+            std::task::Poll::Pending => std::task::Poll::Pending,
+            std::task::Poll::Ready(bytes) => {
+                let n = bytes.len().min(buf_len);
+                buf_slice[..n].copy_from_slice(&bytes[..n]);
+                unsafe {
+                    buf.advance(n);
+                }
+                std::task::Poll::Ready(Ok(()))
+            }
+        }
+    }
+}
+
+impl Guest for Component {
+    async fn main(input: Vec<u8>) -> Vec<u8> {
+        let input: Input = serde_json::from_slice(&input).unwrap();
+
+        let io = NetworkIo::new(input.host, input.port).await;
+        let conn = HyperIo { inner: io };
+
+        let (mut request_sender, conn) = http1::handshake(conn).await.unwrap();
+
+        spawn(async move { conn.await.expect("connection should finish") });
+
+        let request_builder = Request::builder()
+            .uri("/");
+        let request = request_builder.body(Empty::<Bytes>::new()).unwrap();
+
+        let response = request_sender.send_request(request).await.unwrap();
+        assert!(response.status() == StatusCode::OK);
+
+        let output = Output { result: true };
+        serde_json::to_vec(&output).unwrap()
+    }
+}
+
+export!(Component);

crates/wasmtime-plugin/wit/plugin.wit

@@ -0,0 +1,19 @@
+package component:wasmtime-plugin;
+
+interface io {
+    resource network-io {
+        constructor(host: string, port: u32);
+
+        read: func(max: u32) -> list<u8>;
+
+        write: func(bytes: list<u8>) -> u32;
+
+        shutdown: func();
+    }
+}
+
+world plugin {
+    import io;
+
+    export main: func(input: list<u8>) -> list<u8>;
+}