fix(main): improve the way the kernel calculates how many pages to allocate (#472)

Same as #471 for `main` branch

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: zshipko <zshipko@users.noreply.github.com>

.github/workflows/kernel.yml

@@ -27,18 +27,21 @@ jobs:
 
       - name: Build kernel
         shell: bash
+        continue-on-error: true
         run: |
           cd kernel
           sh build.sh
+          git diff --exit-code
+          export GIT_EXIT_CODE=$?
 
       - uses: peter-evans/create-pull-request@v5
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        if: ${{ env.GIT_EXIT_CODE }} != 0
         with:
           title: "update(kernel): extism-runtime.wasm in ${{ github.event.pull_request.head.ref }}"
           body: "Automated PR to update `runtime/src/extism-runtime.wasm` in PR #${{ github.event.number }}"
           base: "${{ github.event.pull_request.head.ref }}"
           branch: "update-kernel--${{ github.event.pull_request.head.ref }}"
-          add-paths: "runtime/src/extism-runtime.wasm"
           commit-message: "update(kernel): extism-runtime.wasm in ${{ github.event.pull_request.head.ref }}"
           delete-branch: true

kernel/src/lib.rs

@@ -87,19 +87,22 @@ pub enum MemoryStatus {
 ///
 /// The overall layout of the Extism-manged memory is organized like this:
 
-/// |------|-------|---------|-------|--------------|
+/// |------|-------+---------|-------+--------------|
-/// | Root | Block |  Data   | Block |     Data     | ...
+/// | Root | Block +  Data   | Block +     Data     | ...
-/// |------|-------|---------|-------|--------------|
+/// |------|-------+---------|-------+--------------|
 ///
 /// Where `Root` and `Block` are fixed to the size of the `MemoryRoot` and `MemoryBlock` structs. But
 /// the size of `Data` is dependent on the allocation size.
+///
+/// This means that the offset of a `Block` is the size of `Root` plus the size of all existing `Blocks`
+/// including their data.
 #[repr(C)]
 pub struct MemoryRoot {
     /// Position of the bump allocator, relative to `START_PAGE`
     pub position: AtomicU64,
     /// The total size of all data allocated using this allocator
     pub length: AtomicU64,
-    /// A pointer to where the blocks begin
+    /// A pointer to the start of the first block
     pub blocks: [MemoryBlock; 0],
 }
 
@@ -194,6 +197,7 @@ impl MemoryRoot {
         while (block as u64) < self.blocks.as_ptr() as u64 + self_position {
             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);
 
             // An unused block is safe to use
@@ -246,13 +250,13 @@ impl MemoryRoot {
         let curr = self.blocks.as_ptr() as u64 + self_position;
 
         // Get the number of bytes available
-        let mem_left = self_length - self_position;
+        let mem_left = self_length - self_position - core::mem::size_of::<MemoryRoot>() as u64;
 
         // When the allocation is larger than the number of bytes available
         // we will need to try to grow the memory
         if length >= mem_left {
             // Calculate the number of pages needed to cover the remaining bytes
-            let npages = num_pages(length);
+            let npages = num_pages(length - mem_left);
             let x = core::arch::wasm32::memory_grow(0, npages);
             if x == usize::MAX {
                 return None;
@@ -309,7 +313,7 @@ impl MemoryBlock {
     }
 }
 
-// Extism functions - these functions should be
+// Extism functions
 
 /// Allocate a block of memory and return the offset
 #[no_mangle]

runtime/src/tests.rs

@@ -4,6 +4,7 @@ use std::time::Instant;
 const WASM: &[u8] = include_bytes!("../../wasm/code-functions.wasm");
 const WASM_LOOP: &[u8] = include_bytes!("../../wasm/loop.wasm");
 const WASM_GLOBALS: &[u8] = include_bytes!("../../wasm/globals.wasm");
+const WASM_REFLECT: &[u8] = include_bytes!("../../wasm/reflect.wasm");
 
 host_fn!(hello_world (a: String) -> String { a });
 
@@ -289,3 +290,24 @@ fn test_toml_manifest() {
     let count: serde_json::Value = serde_json::from_slice(output).unwrap();
     assert_eq!(count.get("count").unwrap().as_i64().unwrap(), 1);
 }
+
+#[test]
+fn test_fuzz_reflect_plugin() {
+    // assert!(set_log_file("stdout", Some(log::Level::Trace)));
+    let f = Function::new(
+        "host_reflect",
+        [ValType::I64],
+        [ValType::I64],
+        None,
+        hello_world,
+    );
+
+    let mut plugin = Plugin::new(WASM_REFLECT, [f], true).unwrap();
+
+    for i in 1..65540 {
+        let input = "a".repeat(i);
+        let output = plugin.call("reflect", input.clone());
+        let output = std::str::from_utf8(output.unwrap()).unwrap();
+        assert_eq!(output, input);
+    }
+}