feat: add rust-protobuf support to extism-convert (#653)

(code from @zshipko, thank you!)

---------

Co-authored-by: Zach Shipko <zach@dylibso.com>

convert/Cargo.toml

@@ -14,6 +14,7 @@ anyhow = "1.0.75"
 base64 = "~0.21"
 bytemuck = {version = "1.14.0", optional = true }
 prost = { version = "0.12.0", optional = true }
+protobuf = { version = "3.2.0", optional = true }
 rmp-serde = { version = "1.1.2", optional = true }
 serde = "1.0.186"
 serde_json = "1.0.105"
@@ -22,7 +23,6 @@ serde_json = "1.0.105"
 serde = { version = "1.0.186", features = ["derive"] }
 
 [features]
-default = ["msgpack", "protobuf", "raw"]
+default = ["msgpack", "prost", "raw"]
 msgpack = ["rmp-serde"]
-protobuf = ["prost"]
 raw = ["bytemuck"]

convert/src/encoding.rs

@@ -112,19 +112,19 @@ impl FromBytesOwned for Base64<String> {
 /// Protobuf encoding
 ///
 /// Allows for `prost` Protobuf messages to be used as arguments to Extism plugin calls
-#[cfg(feature = "protobuf")]
+#[cfg(feature = "prost")]
 #[derive(Debug)]
-pub struct Protobuf<T: prost::Message>(pub T);
+pub struct Prost<T: prost::Message>(pub T);
 
-#[cfg(feature = "protobuf")]
-impl<T: prost::Message> From<T> for Protobuf<T> {
+#[cfg(feature = "prost")]
+impl<T: prost::Message> From<T> for Prost<T> {
     fn from(data: T) -> Self {
         Self(data)
     }
 }
 
-#[cfg(feature = "protobuf")]
-impl<'a, T: prost::Message> ToBytes<'a> for Protobuf<T> {
+#[cfg(feature = "prost")]
+impl<'a, T: prost::Message> ToBytes<'a> for Prost<T> {
     type Bytes = Vec<u8>;
 
     fn to_bytes(&self) -> Result<Self::Bytes, Error> {
@@ -132,10 +132,32 @@ impl<'a, T: prost::Message> ToBytes<'a> for Protobuf<T> {
     }
 }
 
-#[cfg(feature = "protobuf")]
-impl<T: Default + prost::Message> FromBytesOwned for Protobuf<T> {
+#[cfg(feature = "prost")]
+impl<T: Default + prost::Message> FromBytesOwned for Prost<T> {
     fn from_bytes_owned(data: &[u8]) -> Result<Self, Error> {
-        Ok(Protobuf(T::decode(data)?))
+        Ok(Prost(T::decode(data)?))
+    }
+}
+
+/// Protobuf encoding
+///
+/// Allows for `rust-protobuf` Protobuf messages to be used as arguments to Extism plugin calls
+#[cfg(feature = "protobuf")]
+pub struct Protobuf<T: protobuf::Message>(pub T);
+
+#[cfg(feature = "protobuf")]
+impl<'a, T: protobuf::Message> ToBytes<'a> for Protobuf<T> {
+    type Bytes = Vec<u8>;
+
+    fn to_bytes(&self) -> Result<Self::Bytes, Error> {
+        Ok(self.0.write_to_bytes()?)
+    }
+}
+
+#[cfg(feature = "protobuf")]
+impl<T: Default + protobuf::Message> FromBytesOwned for Protobuf<T> {
+    fn from_bytes_owned(data: &[u8]) -> Result<Self, Error> {
+        Ok(Protobuf(T::parse_from_bytes(data)?))
     }
 }
 

convert/src/lib.rs

@@ -18,6 +18,9 @@ pub use encoding::{Base64, Json};
 #[cfg(feature = "msgpack")]
 pub use encoding::Msgpack;
 
+#[cfg(feature = "prost")]
+pub use encoding::Prost;
+
 #[cfg(feature = "protobuf")]
 pub use encoding::Protobuf;
 

kernel/src/lib.rs

@@ -382,11 +382,40 @@ pub unsafe fn free(p: Pointer) {
 }
 
 /// Get the length of an allocated memory block
+///
+/// Note: this should only be called on memory handles returned
+/// by a call to `alloc` - it will return garbage on invalid offsets
+#[no_mangle]
+pub unsafe fn length_unsafe(p: Pointer) -> Length {
+    if p == 0 {
+        return 0;
+    }
+
+    if !MemoryRoot::pointer_in_bounds_fast(p) {
+        return 0;
+    }
+
+    let ptr = p - core::mem::size_of::<MemoryBlock>() as u64;
+    let block = &mut *(ptr as *mut MemoryBlock);
+
+    // Simplest sanity check to verify the pointer is a block
+    if block.status.load(Ordering::Acquire) != MemoryStatus::Active as u8 {
+        return 0;
+    }
+
+    block.used as Length
+}
+
+/// Get the length but returns 0 if the offset is not a valid handle.
+///
+/// Note: this function walks each node in the allocations list, which ensures correctness, but is also
+/// slow
 #[no_mangle]
 pub unsafe fn length(p: Pointer) -> Length {
     if p == 0 {
         return 0;
     }
+
     if let Some(block) = MemoryRoot::new().find_block(p) {
         block.used as Length
     } else {

runtime/Cargo.toml

@@ -9,8 +9,8 @@ repository.workspace = true
 version.workspace = true
 
 [dependencies]
-wasmtime = ">= 14.0.0, < 16.0.0"
-wasmtime-wasi = ">= 14.0.0, < 16.0.0"
+wasmtime = ">= 14.0.0, < 17.0.0"
+wasmtime-wasi = ">= 14.0.0, < 17.0.0"
 anyhow = "1"
 serde = {version = "1", features = ["derive"]}
 serde_json = "1"
@@ -33,7 +33,7 @@ register-filesystem = [] # enables wasm to be loaded from disk
 http = ["ureq"]          # enables extism_http_request
 
 [build-dependencies]
-cbindgen = "0.26"
+cbindgen = { version = "0.26", default-features = false }
 
 [dev-dependencies]
 criterion = "0.5.1"

runtime/src/current_plugin.rs

@@ -246,6 +246,26 @@ impl CurrentPlugin {
         Ok(len)
     }
 
+    pub fn memory_length_unsafe(&mut self, offs: u64) -> Result<u64, Error> {
+        let (linker, mut store) = self.linker_and_store();
+        let output = &mut [Val::I64(0)];
+        if let Some(f) = linker.get(&mut store, EXTISM_ENV_MODULE, "length_unsafe") {
+            f.into_func()
+                .unwrap()
+                .call(&mut store, &[Val::I64(offs as i64)], output)?;
+        } else {
+            anyhow::bail!("unable to locate an extism kernel function: length_unsafe",)
+        }
+        let len = output[0].unwrap_i64() as u64;
+        trace!(
+            plugin = self.id.to_string(),
+            "memory_length_unsafe({}) = {}",
+            offs,
+            len
+        );
+        Ok(len)
+    }
+
     /// Access a plugin's variables
     pub fn vars(&self) -> &std::collections::BTreeMap<String, Vec<u8>> {
         &self.vars

runtime/src/manifest.rs

@@ -117,10 +117,17 @@ pub(crate) fn load(
     match input {
         WasmInput::Data(data) => {
             let has_magic = data.len() >= 4 && data[0..4] == WASM_MAGIC;
-            let is_wat = data.starts_with(b"(module") || data.starts_with(b";;");
+            let s = std::str::from_utf8(&data);
+            let is_wat = s.is_ok_and(|s| {
+                let s = s.trim_start();
+                let starts_with_module = s.len() > 2
+                    && data[0] == b'('   // First character is `(`
+                    && s[1..].trim_start().starts_with("module"); // Then `module` (after any whitespace)
+                starts_with_module || s.starts_with(";;") || s.starts_with("(;")
+            });
             if !has_magic && !is_wat {
                 trace!("Loading manifest");
-                if let Ok(s) = std::str::from_utf8(&data) {
+                if let Ok(s) = s {
                     let t = if let Ok(t) = toml::from_str::<extism_manifest::Manifest>(s) {
                         trace!("Manifest is TOML");
                         modules(engine, &t, &mut mods)?;

runtime/src/plugin.rs

@@ -824,6 +824,24 @@ impl Plugin {
             .and_then(move |_| self.output())
     }
 
+    /// Similar to `Plugin::call`, but returns the Extism error code along with the
+    /// `Error`. It is assumed if `Ok(_)` is returned that the error code was `0`.
+    ///
+    /// All Extism plugin calls return an error code, `Plugin::call` consumes the error code,
+    /// while `Plugin::call_get_error_code` preserves it - this function should only be used
+    /// when you need to inspect the actual return value of a plugin function when it fails.
+    pub fn call_get_error_code<'a, 'b, T: ToBytes<'a>, U: FromBytes<'b>>(
+        &'b mut self,
+        name: impl AsRef<str>,
+        input: T,
+    ) -> Result<U, (Error, i32)> {
+        let lock = self.instance.clone();
+        let mut lock = lock.lock().unwrap();
+        let data = input.to_bytes().map_err(|e| (e, -1))?;
+        self.raw_call(&mut lock, name, data)
+            .and_then(move |_| self.output().map_err(|e| (e, -1)))
+    }
+
     /// Get a `CancelHandle`, which can be used from another thread to cancel a running plugin
     pub fn cancel_handle(&self) -> CancelHandle {
         self.cancel_handle.clone()

runtime/src/tests/kernel.rs

@@ -22,6 +22,20 @@ fn extism_length<T>(mut store: &mut wasmtime::Store<T>, instance: &mut Instance,
     out[0].unwrap_i64() as u64
 }
 
+fn extism_length_unsafe<T>(
+    mut store: &mut wasmtime::Store<T>,
+    instance: &mut Instance,
+    p: u64,
+) -> u64 {
+    let out = &mut [Val::I64(0)];
+    instance
+        .get_func(&mut store, "length_unsafe")
+        .unwrap()
+        .call(&mut store, &[Val::I64(p as i64)], out)
+        .unwrap();
+    out[0].unwrap_i64() as u64
+}
+
 fn extism_load_u8<T>(mut store: &mut wasmtime::Store<T>, instance: &mut Instance, p: u64) -> u8 {
     let out = &mut [Val::I32(0)];
     instance
@@ -173,6 +187,7 @@ fn test_kernel_allocations() {
     let first_alloc = p;
     assert!(p > 0);
     assert_eq!(extism_length(&mut store, instance, p), 1);
+    assert_eq!(extism_length_unsafe(&mut store, instance, p), 1);
     extism_free(&mut store, instance, p);
 
     // 2 bytes
@@ -180,18 +195,21 @@ fn test_kernel_allocations() {
     assert!(x > 0);
     assert!(x != p);
     assert_eq!(extism_length(&mut store, instance, x), 2);
+    assert_eq!(extism_length_unsafe(&mut store, instance, x), 2);
     extism_free(&mut store, instance, x);
 
     for i in 0..64 {
         let p = extism_alloc(&mut store, instance, 64 - i);
         assert!(p > 0);
         assert_eq!(extism_length(&mut store, instance, p), 64 - i);
+        assert_eq!(extism_length_unsafe(&mut store, instance, p), 64 - i);
         extism_free(&mut store, instance, p);
 
         // should re-use the last allocation
         let q = extism_alloc(&mut store, instance, 64 - i);
         assert_eq!(p, q);
         assert_eq!(extism_length(&mut store, instance, q), 64 - i);
+        assert_eq!(extism_length_unsafe(&mut store, instance, q), 64 - i);
         extism_free(&mut store, instance, q);
     }