From 45ceab7dc7df0d473570ee6872c65997cff24e9e Mon Sep 17 00:00:00 2001
From: xermicus <cyrill@parity.io>
Date: Wed, 3 Dec 2025 16:00:19 +0100
Subject: [PATCH] Additional bounds check in sbrk (#428)

Close #356

Signed-off-by: xermicus <cyrill@parity.io>
---
 CHANGELOG.md                                  |  1 +
 crates/integration/codesize.json              | 16 +++---
 crates/integration/src/tests.rs               | 49 +++++++++++++++++++
 .../polkavm/context/function/runtime/sbrk.rs  | 24 +++++++--
 4 files changed, 77 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 61067a6..60bfb4d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@ Supported `polkadot-sdk` rev: `2509.0.0`
 
 ### Fixed:
 - The missing `STOP` instruction at the end of `code` blocks.
+- The missing bounds check in the internal sbrk implementation.
 
 ## v0.5.0
 
diff --git a/crates/integration/codesize.json b/crates/integration/codesize.json
index 84c4bfe..8e3ca89 100644
--- a/crates/integration/codesize.json
+++ b/crates/integration/codesize.json
@@ -1,10 +1,10 @@
 {
-  "Baseline": 905,
-  "Computation": 2286,
-  "DivisionArithmetics": 14347,
-  "ERC20": 16929,
-  "Events": 1665,
-  "FibonacciIterative": 1447,
-  "Flipper": 2077,
-  "SHA1": 7721
+  "Baseline": 911,
+  "Computation": 2293,
+  "DivisionArithmetics": 14353,
+  "ERC20": 16936,
+  "Events": 1672,
+  "FibonacciIterative": 1454,
+  "Flipper": 2083,
+  "SHA1": 7727
 }
\ No newline at end of file
diff --git a/crates/integration/src/tests.rs b/crates/integration/src/tests.rs
index 8ec7356..7112ee4 100644
--- a/crates/integration/src/tests.rs
+++ b/crates/integration/src/tests.rs
@@ -619,3 +619,52 @@ fn code_block_with_nested_object_stops() {
     }
     .run();
 }
+
+#[test]
+fn sbrk_bounds_checks() {
+    let code = &build_yul(&[(
+        "poc.yul",
+        r#"object "Test" {
+    code {
+        return(0x4, 0xffffffff)
+        stop()
+    }
+    object "Test_deployed" {
+        code {
+            stop()
+        }
+    }
+}"#,
+    )])
+    .unwrap()["poc.yul:Test"];
+
+    let results = Specs {
+        actions: vec![
+            Instantiate {
+                origin: TestAddress::Alice,
+                value: 0,
+                gas_limit: Some(GAS_LIMIT),
+                storage_deposit_limit: None,
+                code: Code::Bytes(code.to_vec()),
+                data: Default::default(),
+                salt: OptionalHex::default(),
+            },
+            VerifyCall(VerifyCallExpectation {
+                success: false,
+                ..Default::default()
+            }),
+        ],
+        differential: false,
+        ..Default::default()
+    }
+    .run();
+
+    let CallResult::Instantiate { result, .. } = results.last().unwrap() else {
+        unreachable!()
+    };
+
+    assert!(
+        format!("{result:?}").contains("ContractTrapped"),
+        "not seeing a trap means the contract did not catch the OOB"
+    );
+}
diff --git a/crates/llvm-context/src/polkavm/context/function/runtime/sbrk.rs b/crates/llvm-context/src/polkavm/context/function/runtime/sbrk.rs
index 3460088..d8531d7 100644
--- a/crates/llvm-context/src/polkavm/context/function/runtime/sbrk.rs
+++ b/crates/llvm-context/src/polkavm/context/function/runtime/sbrk.rs
@@ -77,6 +77,20 @@ impl RuntimeFunction for Sbrk {
         context.build_unreachable();
 
         context.set_basic_block(offset_in_bounds_block);
+        let size_in_bounds_block = context.append_basic_block("size_in_bounds");
+        let is_size_out_of_bounds = context.builder().build_int_compare(
+            inkwell::IntPredicate::UGT,
+            size,
+            context.heap_size(),
+            "size_in_bounds",
+        )?;
+        context.build_conditional_branch(
+            is_size_out_of_bounds,
+            trap_block,
+            size_in_bounds_block,
+        )?;
+
+        context.set_basic_block(size_in_bounds_block);
         let mask = context
             .xlen_type()
             .const_int(BYTE_LENGTH_WORD as u64 - 1, false);
@@ -88,20 +102,20 @@ impl RuntimeFunction for Sbrk {
             context.builder().build_not(mask, "mask_not")?,
             "memory_size",
         )?;
-        let size_in_bounds_block = context.append_basic_block("size_in_bounds");
-        let is_size_out_of_bounds = context.builder().build_int_compare(
+        let total_size_in_bounds_block = context.append_basic_block("total_size_in_bounds");
+        let is_total_size_out_of_bounds = context.builder().build_int_compare(
             inkwell::IntPredicate::UGT,
             memory_size,
             context.heap_size(),
             "size_out_of_bounds",
         )?;
         context.build_conditional_branch(
-            is_size_out_of_bounds,
+            is_total_size_out_of_bounds,
             trap_block,
-            size_in_bounds_block,
+            total_size_in_bounds_block,
         )?;
 
-        context.set_basic_block(size_in_bounds_block);
+        context.set_basic_block(total_size_in_bounds_block);
         let new_size_block = context.append_basic_block("new_size");
         let is_new_size = context.builder().build_int_compare(
             inkwell::IntPredicate::UGT,