From 26542d57406af58aa2cc7c4b459131aadc9310ff Mon Sep 17 00:00:00 2001 From: zach Date: Wed, 3 Jan 2024 09:29:05 -0800 Subject: [PATCH] feat(kernel): add extism_length_unsafe (#648) - Adds `length_unsafe` function to the extism kernel, a more performant `length` for known-valid memory handles After this is merged I will update go-sdk and js-sdk too. Closes #643 --- kernel/src/lib.rs | 29 +++++++++++++++++++++++++++++ runtime/src/current_plugin.rs | 20 ++++++++++++++++++++ runtime/src/extism-runtime.wasm | Bin 3409 -> 3493 bytes runtime/src/tests/kernel.rs | 18 ++++++++++++++++++ 4 files changed, 67 insertions(+) diff --git a/kernel/src/lib.rs b/kernel/src/lib.rs index ee3be09..7473b6f 100644 --- a/kernel/src/lib.rs +++ b/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::() 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 { diff --git a/runtime/src/current_plugin.rs b/runtime/src/current_plugin.rs index bc70d3c..14728c9 100644 --- a/runtime/src/current_plugin.rs +++ b/runtime/src/current_plugin.rs @@ -246,6 +246,26 @@ impl CurrentPlugin { Ok(len) } + pub fn memory_length_unsafe(&mut self, offs: u64) -> Result { + 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> { &self.vars diff --git a/runtime/src/extism-runtime.wasm b/runtime/src/extism-runtime.wasm index 63e3303d7dbf4feced321481a3e4912c91245ebd..e36ea4282cd120ac3ac2f60a0585580ecf9812f4 100755 GIT binary patch delta 420 zcmYk0J5Rz;7>2*^^g>H}5WE0U6AxWAEa?!-NG#C$fp{;382NvNU2x;71 zjFS`p!QJT3aj>N&aeJTleZR}s=G)e^S*TS3kl@U2O0T3gbL(R|U!sPnDOP=M^R3X>V2MLEs0Am20^kB>uYTa@P+=&h<#}5H{i- zZc8uqYCOSD{Ekn^vrO8g9siJBi{U;%Nm2p9V7N!*hH7G_ER`u|REj-mr*cJ!sTeGA z9xSyp!PA&Y$3!;fQ^`f!CisNoE;Nv|8FX8h6k5d$V Lksg`<$`{%n8!T>l delta 389 zcmYjL%SyvQ6umQ(O!M~HD!Lg11wq^t5#6B{Z_vb~;Od^&J@-EEe(uZ~v{nT`cz)U-GhCa8+fzCeVNdo!sSG2}=i?Jl zwE`U<0kv<#(OB}ORdN3FETCzLjAD_m>3Rt>+0ctzA|#kbG?sZuPAsb|D303JwBS*< zo7;|EX3M8mmb~;=)mY@YHC@DU6#rGOm$;Yd8&)8KD4y`nM2ZA94bFW}dR&A)K_=&Y z;f;95O9a8j*H-lnq4Wv8qmT3kKc?TfUdK&_4gnm3Fkq9ReXO{YsTVe8*|>9Q>jy4V tpPV053U`m1Oks%8LCYGr==ItL(`3LfBM)cIdsk`Oc=194w762++AjuhY%c%+ diff --git a/runtime/src/tests/kernel.rs b/runtime/src/tests/kernel.rs index a9c5a55..213df30 100644 --- a/runtime/src/tests/kernel.rs +++ b/runtime/src/tests/kernel.rs @@ -22,6 +22,20 @@ fn extism_length(mut store: &mut wasmtime::Store, instance: &mut Instance, out[0].unwrap_i64() as u64 } +fn extism_length_unsafe( + mut store: &mut wasmtime::Store, + 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(mut store: &mut wasmtime::Store, 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); }