From fd95729d8d668a8777a45bc779f7d8cd1c3559e7 Mon Sep 17 00:00:00 2001
From: zach <zach@dylibso.com>
Date: Thu, 14 Dec 2023 16:55:13 -0800
Subject: [PATCH] fix(kernel): length function should return 0 for invalid
 offsets (#635)

Fixes #634

- Updates `extism_length` to walks the allocation list to determine
valid offsets instead of assuming the provided offset is valid

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: zshipko <zshipko@users.noreply.github.com>
---
 kernel/src/lib.rs               |  23 ++++++++++++++++++++---
 runtime/src/extism-runtime.wasm | Bin 3296 -> 3409 bytes
 runtime/src/tests/kernel.rs     |   6 +++++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/kernel/src/lib.rs b/kernel/src/lib.rs
index e5b69bc..ee3be09 100644
--- a/kernel/src/lib.rs
+++ b/kernel/src/lib.rs
@@ -306,9 +306,26 @@ impl MemoryRoot {
         if !Self::pointer_in_bounds_fast(offs) {
             return None;
         }
-        let ptr = offs - core::mem::size_of::<MemoryBlock>() as u64;
-        let ptr = ptr as *mut MemoryBlock;
-        Some(&mut *ptr)
+
+        // Get the first block
+        let mut block = self.blocks.as_mut_ptr();
+
+        // Only loop while the block pointer is less then the current position
+        while (block as u64) < self.blocks.as_ptr() as u64 + offs {
+            let b = &mut *block;
+
+            // Get the block status, this lets us know if we are able to re-use it
+            let status = b.status.load(Ordering::Acquire);
+
+            if status == MemoryStatus::Active as u8 && b.data.as_ptr() as Pointer == offs {
+                return Some(b);
+            }
+
+            // Get the next block
+            block = b.next_ptr();
+        }
+
+        None
     }
 }
 
diff --git a/runtime/src/extism-runtime.wasm b/runtime/src/extism-runtime.wasm
index 4ca198ce06e17cfcb58cdc1a72422f96ada367b0..63e3303d7dbf4feced321481a3e4912c91245ebd 100755
GIT binary patch
delta 224
zcmaDLc~NRZEhEz<@y+#&lbBduF)}hw7G>$obv($R$f&@u97ty=GBG(YJ18(Z@?<G8
zGXWX83=EEp?z{{N%xe`G0(qIZ6_^y56&N*`ShAEDy?GgUn7Fxt%9*Ty>KHVc6%`nQ
zc!4Y@26JXjZhmeC2JUM>d+M1SKv01pV6r@mIwQ;EMmCqpcC5l~sCMJCnHk+?Zf++A
PexTi^OdOlzSsU2_VC*GW

delta 101
zcmca8^+0k%EhE!IvCZ|2lbBf6GBPqw7G>#VQ(#!`SdumQA&ZRWB1R_0I>ve@21Q0D
z2POvvh5%j$M<CQ?U{qpqWK>|1wqgLvC@@VnW69tHX{^aIWdf?^R$$zGfwhqx00pfU
ACIA2c

diff --git a/runtime/src/tests/kernel.rs b/runtime/src/tests/kernel.rs
index 564c5f5..a9c5a55 100644
--- a/runtime/src/tests/kernel.rs
+++ b/runtime/src/tests/kernel.rs
@@ -198,6 +198,10 @@ fn test_kernel_allocations() {
     // 512 bytes, test block re-use + splitting
     let p = extism_alloc(&mut store, instance, 512);
     assert_eq!(extism_length(&mut store, instance, p), 512);
+    assert_eq!(extism_length(&mut store, instance, p + 1), 0);
+    assert_eq!(extism_length(&mut store, instance, p + 2), 0);
+    assert_eq!(extism_length(&mut store, instance, p + 3), 0);
+    assert_eq!(extism_length(&mut store, instance, p + 4), 0);
     extism_free(&mut store, instance, p);
 
     // 128 bytes, should be split off the 512 byte block
@@ -210,7 +214,7 @@ fn test_kernel_allocations() {
     let r = extism_alloc(&mut store, instance, 128);
     assert!(p <= r && r < p + 512);
     assert!(r > p);
-    assert_eq!(extism_length(&mut store, instance, q), 128);
+    assert_eq!(extism_length(&mut store, instance, r), 128);
     extism_free(&mut store, instance, q);
 
     // 100 pages