Additional bounds check in sbrk (#428)

Close #356

Signed-off-by: xermicus <cyrill@parity.io>

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
 

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
 }

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"
+    );
+}

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,