feat: Proof of concept for remote sandbox (#56)

* feat: Proof of concept for remote sandbox

* fix: Make Docker builds work

Author: cdata

Cargo.lock

@@ -38,8 +38,9 @@ url = { version = "2" }
 usuba-bundle = { path = "./rust/usuba-bundle" }
 utoipa = { version = "4" }
 utoipa-swagger-ui = { version = "7" }
+wasmtime = { version = "21" }
 wasm-bindgen = { version = "0.2" }
-wasmtime-environ = { version = "20" }
+wasmtime-environ = { version = "21" }
 web-sys = { version = "0.3" }
 wit-bindgen = { version = "0.25" }
 wit-parser = { version = "0.208" }

Cargo.toml

@@ -10,4 +10,4 @@ crate-type = ["cdylib", "rlib"]
 blake3 = { workspace = true }
 js-component-bindgen = { workspace = true }
 wit-bindgen = { workspace = true }
-wasmtime-environ = { workspace = true,  features = ["component-model", "compile"] }
+wasmtime-environ = { version = "20",  features = ["component-model", "compile"] }

rust/usuba-compat/Cargo.toml

@@ -9,7 +9,7 @@ edition = "2021"
 [dependencies]
 anyhow = { workspace = true }
 async-trait = { workspace = true }
-axum = { workspace = true, features = ["multipart"] }
+axum = { workspace = true, features = ["multipart", "macros"] }
 blake3 = { workspace = true }
 bytes = { workspace = true }
 hyper-util = { workspace = true }
@@ -27,6 +27,8 @@ tracing-subscriber = { workspace = true }
 usuba-bundle = { workspace = true }
 utoipa = { workspace = true, features = ["axum_extras"] }
 utoipa-swagger-ui = { workspace = true, features = ["axum"] }
+wasmtime = { workspace = true }
+wasmtime-wasi = { version = "21" }
 wit-parser = { workspace = true }
 # gcloud-sdk = { version = "0.24.6", features = ["google-cloud-aiplatform-v1"] }
 

rust/usuba/Cargo.toml

@@ -15,9 +15,9 @@ COPY ./rust ./rust
 COPY ./typescript ./typescript
 
 RUN rustup target add wasm32-wasi
-RUN cargo install wasm-tools
+RUN cargo install wasm-tools wit-deps-cli
 
-WORKDIR /build-root/typescript/packages/usuba-ui
+WORKDIR /build-root/typescript
 
 RUN npm install -g @bytecodealliance/jco
 RUN npm ci

rust/usuba/Dockerfile

@@ -4,11 +4,13 @@ extern crate tracing;
 mod bake;
 mod error;
 pub mod openapi;
+mod recipe;
 pub mod routes;
 mod serve;
 mod storage;
 
 pub use bake::*;
 pub use error::*;
+pub use recipe::*;
 pub use serve::*;
 pub use storage::*;

rust/usuba/src/lib.rs

@@ -1,7 +1,10 @@
 use utoipa::OpenApi;
 
 use crate::{
-    routes::{BuildModuleRequest, BuildModuleResponse, BundleRequest},
+    routes::{
+        BuildModuleRequest, BuildModuleResponse, BundleRequest, EvalRecipeRequest,
+        EvalRecipeResponse, JsonValue,
+    },
     ErrorResponse,
 };
 
@@ -10,13 +13,17 @@ use crate::{
     paths(
         crate::routes::build_module,
         crate::routes::retrieve_module,
-        crate::routes::bundle_javascript
+        crate::routes::bundle_javascript,
+        crate::routes::eval_recipe
     ),
     components(
         schemas(BuildModuleResponse),
         schemas(ErrorResponse),
         schemas(BuildModuleRequest),
-        schemas(BundleRequest)
+        schemas(BundleRequest),
+        schemas(EvalRecipeRequest),
+        schemas(EvalRecipeResponse),
+        schemas(JsonValue)
     )
 )]
 pub struct OpenApiDocs;

rust/usuba/src/openapi.rs

@@ -0,0 +1,160 @@
+use std::collections::BTreeMap;
+
+use wasmtime::{
+    component::{bindgen, Resource},
+    Module,
+};
+
+bindgen!({
+  world: "common",
+  path: "../../typescript/common/module/wit"
+});
+
+pub use common::data::types::{Dictionary, Reference, Value};
+use wasmtime_wasi::WasiView;
+
+pub trait InputOutput: Send + Sync + Clone + std::fmt::Debug {
+    fn read(&self, key: &str) -> Option<Value>;
+    fn write(&mut self, key: &str, value: Value);
+}
+
+pub struct ResourceTable<T> {
+    next_index: u32,
+    resources: BTreeMap<u32, T>,
+}
+
+impl<T> ResourceTable<T> {
+    pub fn add(&mut self, resource: T) -> u32 {
+        let index = self.next_index;
+
+        self.next_index = self.next_index + 1;
+        self.resources.insert(index, resource);
+
+        index
+    }
+
+    pub fn lookup(&self, index: u32) -> Option<&T> {
+        self.resources.get(&index)
+    }
+
+    pub fn remove(&mut self, index: u32) {
+        self.resources.remove(&index);
+    }
+}
+
+impl<T> Default for ResourceTable<T> {
+    fn default() -> Self {
+        Self {
+            next_index: Default::default(),
+            resources: BTreeMap::new(),
+        }
+    }
+}
+
+#[repr(transparent)]
+pub struct HostReference(String);
+
+pub struct ModuleEnvironment<Io>
+where
+    Io: InputOutput,
+{
+    io: Io,
+    references: ResourceTable<HostReference>,
+
+    wasi_resources: wasmtime_wasi::ResourceTable,
+    wasi_ctx: wasmtime_wasi::WasiCtx,
+}
+
+impl<Io> ModuleEnvironment<Io>
+where
+    Io: InputOutput,
+{
+    pub fn new(io: Io) -> Self {
+        ModuleEnvironment {
+            io,
+            references: ResourceTable::default(),
+
+            wasi_resources: wasmtime_wasi::ResourceTable::new(),
+            wasi_ctx: wasmtime_wasi::WasiCtx::builder().build(),
+        }
+    }
+
+    pub fn take_io(self) -> Io {
+        self.io
+    }
+}
+
+impl<Io> common::io::state::Host for ModuleEnvironment<Io>
+where
+    Io: InputOutput,
+{
+    fn read(&mut self, name: String) -> Option<Resource<Reference>> {
+        warn!("Reading input: {name}");
+        if self.io.read(&name).is_none() {
+            return None;
+        }
+
+        let reference = HostReference(name);
+        let index = self.references.add(reference);
+
+        Some(Resource::new_own(index))
+    }
+
+    fn write(&mut self, name: String, value: Value) -> () {
+        warn!("Writing output: {name}");
+        self.io.write(&name, value);
+    }
+}
+
+impl<Io> common::data::types::HostDictionary for ModuleEnvironment<Io>
+where
+    Io: InputOutput,
+{
+    fn get(
+        &mut self,
+        _resource: Resource<Dictionary>,
+        _key: String,
+    ) -> Option<wasmtime::component::Resource<Reference>> {
+        unimplemented!("Dictionary not supported yet saaawiii");
+    }
+
+    fn drop(&mut self, _rep: Resource<Dictionary>) -> wasmtime::Result<()> {
+        unimplemented!("Dictionary not supported yet saaawiii");
+    }
+}
+
+impl<Io> common::data::types::HostReference for ModuleEnvironment<Io>
+where
+    Io: InputOutput,
+{
+    /// Dereference a reference to a value
+    /// This call is fallible (for example, if the dereference is not allowed)
+    /// The value may be none (for example, if it is strictly opaque)
+    fn deref(&mut self, resource: Resource<Reference>) -> Result<Option<Value>, String> {
+        let HostReference(key) = self
+            .references
+            .lookup(resource.rep())
+            .ok_or_else(|| String::from("Attempted to deref an untracked Reference"))?;
+
+        Ok(self.io.read(key))
+    }
+
+    fn drop(&mut self, rep: Resource<Reference>) -> wasmtime::Result<()> {
+        Ok(self.references.remove(rep.rep()))
+    }
+}
+
+impl<Io> common::data::types::Host for ModuleEnvironment<Io> where Io: InputOutput {}
+
+impl<Io> WasiView for ModuleEnvironment<Io>
+where
+    Io: InputOutput,
+{
+    fn table(&mut self) -> &mut wasmtime_wasi::ResourceTable {
+        &mut self.wasi_resources
+    }
+
+    fn ctx(&mut self) -> &mut wasmtime_wasi::WasiCtx {
+        &mut self.wasi_ctx
+    }
+}

rust/usuba/src/recipe/common.rs

@@ -0,0 +1,6 @@
+// pub mod common;
+pub mod common;
+
+mod runtime;
+
+pub use runtime::*;

rust/usuba/src/recipe/mod.rs

@@ -0,0 +1,78 @@
+use wasmtime::component::{Component, Linker};
+use wasmtime::{Engine, Store};
+
+use crate::common::exports::common::module::module::GuestBody;
+use crate::{Bake, Baker, UsubaError};
+
+use super::common::Common;
+pub use super::common::{Dictionary, InputOutput, ModuleEnvironment, Value};
+
+const COMMON_MODULE_WIT: &[u8] =
+    include_bytes!("../../../../typescript/common/module/wit/module.wit");
+
+const COMMON_IO_WIT: &[u8] = include_bytes!("../../../../typescript/common/io/wit/io.wit");
+
+const COMMON_DATA_WIT: &[u8] = include_bytes!("../../../../typescript/common/data/wit/data.wit");
+
+pub struct Runtime {}
+
+impl Runtime {
+    pub async fn eval<Io: InputOutput + 'static>(
+        &mut self,
+        content_type: String,
+        source_code: String,
+        io: Io,
+    ) -> Result<Io, UsubaError> {
+        let component_baker = match content_type.as_str() {
+            "text/javascript" => Baker::JavaScript,
+            "text/x-python" => Baker::Python,
+            _ => return Err(UsubaError::BadRequest),
+        };
+
+        let component_wasm = component_baker
+            .bake(
+                "common",
+                vec![COMMON_MODULE_WIT.into()],
+                source_code.into(),
+                vec![COMMON_DATA_WIT.into(), COMMON_IO_WIT.into()],
+            )
+            .await?;
+
+        let mut config = wasmtime::Config::default();
+        config.async_support(false);
+
+        let engine = Engine::new(&config)?;
+
+        let mut store = Store::new(&engine, ModuleEnvironment::new(io));
+
+        let component = Component::new(&engine, component_wasm)?;
+        let mut linker = Linker::new(&engine);
+
+        wasmtime_wasi::add_to_linker_sync(&mut linker)?;
+
+        Common::add_to_linker(&mut linker, |environment| environment)?;
+
+        let (common, _inst) = Common::instantiate(&mut store, &component, &linker)?;
+
+        let store = tokio::task::spawn_blocking(move || {
+            let common_module = common.common_module_module();
+
+            match common_module.call_create(&mut store) {
+                Ok(body_resource) => {
+                    common
+                        .common_module_module()
+                        .body()
+                        .call_run(&mut store, body_resource)?;
+                }
+                Err(error) => {
+                    error!("Create failed: {}", error);
+                }
+            };
+
+            Ok(store) as wasmtime::Result<Store<ModuleEnvironment<Io>>, wasmtime::Error>
+        })
+        .await??;
+
+        Ok(store.into_data().take_io())
+    }
+}