polkavm: Add a knob to control sbrk instruction enablement

Currently, we enable sbrk instruction by default, however we do want to
a mechanism to control it for the guest. This patch aims to just do that.

if sbrk is disabled, host should trap the guest program.

A basic test is also added that confirms the same.

Signed-off-by: Aman <[email protected]>

crates/polkavm/src/api.rs

@@ -216,6 +216,7 @@ pub(crate) struct ModulePrivate {
     dynamic_paging: bool,
     page_size_mask: u32,
     page_shift: u32,
+    allow_sbrk: bool,
     #[cfg(feature = "module-cache")]
     pub(crate) module_key: Option<ModuleKey>,
 }
@@ -290,6 +291,10 @@ impl Module {
         self.state().gas_metering
     }
 
+    pub(crate) fn allow_sbrk(&self) -> bool {
+        self.state().allow_sbrk
+    }
+
     pub(crate) fn is_multiple_of_page_size(&self, value: u32) -> bool {
         (value & self.state().page_size_mask) == 0
     }
@@ -520,6 +525,7 @@ impl Module {
             is_strict: config.is_strict,
             step_tracing: config.step_tracing,
             dynamic_paging: config.dynamic_paging,
+            allow_sbrk: config.allow_sbrk,
             crosscheck: engine.crosscheck,
             page_size_mask,
             page_shift,

crates/polkavm/src/compiler.rs

@@ -84,6 +84,7 @@ where
     step_label: Label,
     trap_label: Label,
     invalid_jump_label: Label,
+    allow_sbrk: bool,
 
     _phantom: PhantomData<S>,
 }
@@ -217,6 +218,7 @@ where
             gas_metering_stub_offsets,
             gas_cost_for_basic_block,
             code_length,
+            allow_sbrk: config.allow_sbrk,
             _phantom: PhantomData,
         };
 
@@ -665,10 +667,14 @@ where
 
     #[inline(always)]
     fn sbrk(&mut self, code_offset: u32, args_length: u32, d: RawReg, s: RawReg) -> Self::ReturnTy {
-        self.before_instruction(code_offset);
-        self.gas_visitor.sbrk(d, s);
-        ArchVisitor(self).sbrk(d, s);
-        self.after_instruction::<false>(code_offset, args_length);
+        if self.allow_sbrk {
+            self.before_instruction(code_offset);
+            self.gas_visitor.sbrk(d, s);
+            ArchVisitor(self).sbrk(d, s);
+            self.after_instruction::<false>(code_offset, args_length);
+        } else {
+            self.trap(code_offset, args_length)
+        }
     }
 
     #[inline(always)]

crates/polkavm/src/config.rs

@@ -367,6 +367,7 @@ pub struct ModuleConfig {
     pub(crate) step_tracing: bool,
     pub(crate) dynamic_paging: bool,
     pub(crate) aux_data_size: u32,
+    pub(crate) allow_sbrk: bool,
     cache_by_hash: bool,
 }
 
@@ -386,6 +387,7 @@ impl ModuleConfig {
             step_tracing: false,
             dynamic_paging: false,
             aux_data_size: 0,
+            allow_sbrk: true,
             cache_by_hash: false,
         }
     }
@@ -458,6 +460,18 @@ impl ModuleConfig {
         self
     }
 
+    ///
+    /// Sets whether sbrk instruction is allowed.
+    ///
+    /// When enabled sbrk instruction is not allowed it will lead to a trap, otherwise
+    /// sbrk instruction is emulated.
+    ///
+    /// Default: `true`
+    pub fn set_allow_sbrk(&mut self, enabled: bool) -> &mut Self {
+        self.allow_sbrk = enabled;
+        self
+    }
+
     /// Returns whether the module will be cached by hash.
     pub fn cache_by_hash(&self) -> bool {
         self.cache_by_hash
@@ -484,6 +498,7 @@ impl ModuleConfig {
             is_strict,
             step_tracing,
             dynamic_paging,
+            allow_sbrk,
             // Deliberately ignored.
             cache_by_hash: _,
         } = self;

crates/polkavm/src/interpreter.rs

@@ -13,7 +13,7 @@ use core::mem::MaybeUninit;
 use core::num::NonZeroU32;
 use polkavm_common::abi::VM_ADDR_RETURN_TO_HOST;
 use polkavm_common::operation::*;
-use polkavm_common::program::{asm, InstructionVisitor, RawReg, Reg};
+use polkavm_common::program::{asm, Instruction, InstructionVisitor, RawReg, Reg};
 use polkavm_common::utils::{align_to_next_page_usize, byte_slice_init, slice_assume_init_mut};
 
 type Target = usize;
@@ -689,7 +689,11 @@ impl InterpretedInstance {
         let mut is_first = true;
         let mut is_properly_terminated = false;
         let mut last_program_counter = program_counter;
-        while let Some(instruction) = instructions.next() {
+        while let Some(mut instruction) = instructions.next() {
+            if !self.module.allow_sbrk() && matches!(instruction.kind, Instruction::sbrk(_, _)) {
+                instruction.kind = Instruction::trap;
+            }
+
             self.compiled_offset_for_block
                 .insert(instruction.offset.0, Self::pack_target(self.compiled_handlers.len(), is_first));
 

crates/polkavm/src/tests.rs

@@ -1469,6 +1469,41 @@ fn access_memory_from_host(config: Config) {
     assert!(instance.zero_memory(0xffffffff, 0).is_ok());
 }
 
+fn sbrk_knob_works(config: Config) {
+    let _ = env_logger::try_init();
+    let engine = Engine::new(&config).unwrap();
+    let page_size = get_native_page_size() as u32;
+    let mut builder = ProgramBlobBuilder::new();
+    builder.add_export_by_basic_block(0, b"main");
+    builder.set_code(&[asm::sbrk(Reg::A0, Reg::A0), asm::ret()], &[]);
+
+    let blob = ProgramBlob::parse(builder.into_vec().into()).unwrap();
+
+    for sbrk_allowed in [true, false] {
+        let mut module_config: ModuleConfig = ModuleConfig::new();
+        module_config.set_page_size(page_size);
+        module_config.set_allow_sbrk(sbrk_allowed);
+        module_config.set_gas_metering(Some(GasMeteringKind::Sync));
+        let module = Module::from_blob(&engine, &module_config, blob.clone()).unwrap();
+
+        let mut instance = module.instantiate().unwrap();
+        instance.set_reg(Reg::A0, 0);
+        instance.set_reg(Reg::RA, crate::RETURN_TO_HOST);
+
+        instance.set_gas(5);
+        instance.set_next_program_counter(ProgramCounter(0));
+
+        if sbrk_allowed {
+            match_interrupt!(instance.run().unwrap(), InterruptKind::Finished);
+            assert_eq!(instance.gas(), 3);
+        } else {
+            match_interrupt!(instance.run().unwrap(), InterruptKind::Trap);
+            assert_eq!(instance.program_counter(), Some(ProgramCounter(0)));
+            assert_eq!(instance.gas(), 4);
+        }
+    }
+}
+
 struct TestInstance {
     module: crate::Module,
     instance: crate::Instance,
@@ -2121,6 +2156,7 @@ run_tests! {
     implicit_trap_after_fallthrough
     aux_data_works
     access_memory_from_host
+    sbrk_knob_works
 
     basic_gas_metering_sync
     basic_gas_metering_async