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 63e3303..e36ea42 100755 Binary files a/runtime/src/extism-runtime.wasm and b/runtime/src/extism-runtime.wasm differ 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); }