Add RawInstance::set_accessible_aux_size

Author: koute

crates/polkavm-common/src/zygote.rs

@@ -71,6 +71,7 @@ define_address_table! {
     ext_load_program: unsafe extern "C" fn() -> !,
     ext_recycle: unsafe extern "C" fn() -> !,
     ext_fetch_idle_regs: unsafe extern "C" fn() -> !,
+    ext_set_accessible_aux_size: unsafe extern "C" fn() -> !,
 }
 
 pub const FD_DUMMY_STDIN: i32 = 0;
@@ -262,6 +263,7 @@ pub struct VmCtx {
     pub heap_info: VmCtxHeapInfo,
 
     pub arg2: AtomicU32,
+    pub arg3: AtomicU32,
 
     /// Offset in shared memory to this sandbox's memory map.
     pub shm_memory_map_offset: AtomicU64,
@@ -343,6 +345,7 @@ impl VmCtx {
             next_program_counter: AtomicU32::new(0),
             arg: AtomicU32::new(0),
             arg2: AtomicU32::new(0),
+            arg3: AtomicU32::new(0),
             regs: [ATOMIC_U64_ZERO; REG_COUNT],
             jump_into: AtomicU64::new(0),
             next_native_program_counter: AtomicU64::new(0),

crates/polkavm-zygote/src/main.rs

@@ -682,6 +682,7 @@ unsafe fn initialize(mut stack: *mut usize) {
             (if a == linux_raw::SYS_rt_sigreturn => jump @1),
             (if a == linux_raw::SYS_sched_yield => jump @1),
             (if a == linux_raw::SYS_exit => jump @1),
+            (if a == linux_raw::SYS_mprotect => jump @7),
             (seccomp_return_eperm),
 
             // SYS_write
@@ -705,6 +706,13 @@ unsafe fn initialize(mut stack: *mut usize) {
             (if a != linux_raw::PROT_EXEC => jump @0),
             (seccomp_allow),
 
+            // SYS_mprotect
+            ([7]: a = syscall_arg[2]),
+            (if a != linux_raw::PROT_READ => jump @8),
+            (seccomp_allow),
+            ([8]: if a != 0 => jump @0),
+            (seccomp_allow),
+
             ([0]: seccomp_return_eperm),
             ([1]: seccomp_allow),
         };
@@ -913,6 +921,7 @@ pub static EXT_TABLE: ExtTableRaw = ExtTableRaw {
     ext_load_program,
     ext_recycle,
     ext_fetch_idle_regs,
+    ext_set_accessible_aux_size,
 };
 
 #[inline(always)]
@@ -1103,3 +1112,39 @@ pub unsafe extern "C" fn ext_fetch_idle_regs() -> ! {
 
     signal_host_and_longjmp(VMCTX_FUTEX_IDLE);
 }
+
+#[inline(never)]
+pub unsafe extern "C" fn ext_set_accessible_aux_size() -> ! {
+    trace!("Entry point: ext_set_accessible_aux_size");
+    let address = VMCTX.arg.load(Ordering::Relaxed) as usize;
+    let length_accessible = VMCTX.arg2.load(Ordering::Relaxed) as usize;
+    let length_full = VMCTX.arg3.load(Ordering::Relaxed) as usize;
+
+    trace!(
+        "Setting inaccessible: ",
+        Hex(address),
+        "-",
+        Hex(address + length_full),
+        " (",
+        Hex(length_full),
+        ")"
+    );
+
+    linux_raw::sys_mprotect(address as *mut core::ffi::c_void, length_full, 0)
+        .unwrap_or_else(|error| abort_with_error("failed to set accessible aux size: failed to set the region inaccessible", error));
+
+    trace!(
+        "Setting accessible: ",
+        Hex(address),
+        "-",
+        Hex(address + length_accessible),
+        " (",
+        Hex(length_accessible),
+        ")"
+    );
+
+    linux_raw::sys_mprotect(address as *mut core::ffi::c_void, length_accessible, linux_raw::PROT_READ)
+        .unwrap_or_else(|error| abort_with_error("failed to set accessible aux size: failed to set the region read-only", error));
+
+    signal_host_and_longjmp(VMCTX_FUTEX_IDLE);
+}

crates/polkavm/src/api.rs

@@ -336,6 +336,10 @@ impl Module {
         value & !self.state().page_size_mask
     }
 
+    pub(crate) fn round_to_page_size_up(&self, value: u32) -> u32 {
+        self.round_to_page_size_down(value) + (u32::from((value & self.state().page_size_mask) != 0) << self.state().page_shift)
+    }
+
     pub(crate) fn address_to_page(&self, address: u32) -> u32 {
         address >> self.state().page_shift
     }
@@ -1034,6 +1038,31 @@ impl RawInstance {
         }
     }
 
+    /// Sets the accessible region of the aux data, rounded up to the nearest page size.
+    pub fn set_accessible_aux_size(&mut self, size: u32) -> Result<(), Error> {
+        if self.module.is_dynamic_paging() {
+            return Err("setting accessible aux size is only possible on modules without dynamic paging".into());
+        }
+
+        if size > self.module.memory_map().aux_data_size() {
+            return Err(format!(
+                "cannot set accessible aux size: the maximum is {}, while tried to set {}",
+                self.module.memory_map().aux_data_size(),
+                size
+            )
+            .into());
+        }
+
+        let size = self.module.round_to_page_size_up(size);
+        if let Some(ref mut crosscheck) = self.crosscheck_instance {
+            crosscheck.set_accessible_aux_size(size);
+        }
+
+        access_backend!(self.backend, |mut backend| backend
+            .set_accessible_aux_size(size)
+            .into_result("failed to set accessible aux size"))
+    }
+
     /// Resets the VM's memory to its initial state.
     pub fn reset_memory(&mut self) -> Result<(), Error> {
         if let Some(ref mut crosscheck) = self.crosscheck_instance {

crates/polkavm/src/interpreter.rs

@@ -83,6 +83,7 @@ pub(crate) struct BasicMemory {
     aux: Vec<u8>,
     is_memory_dirty: bool,
     heap_size: u32,
+    accessible_aux_size: usize,
 }
 
 impl BasicMemory {
@@ -93,6 +94,7 @@ impl BasicMemory {
             aux: Vec::new(),
             is_memory_dirty: false,
             heap_size: 0,
+            accessible_aux_size: usize::MAX,
         }
     }
 
@@ -116,6 +118,7 @@ impl BasicMemory {
         self.aux.clear();
         self.heap_size = 0;
         self.is_memory_dirty = false;
+        self.accessible_aux_size = 0;
 
         if let Some(interpreted_module) = module.interpreted_module().as_ref() {
             self.rw_data.extend_from_slice(&interpreted_module.rw_data);
@@ -124,21 +127,26 @@ impl BasicMemory {
 
             // TODO: Do this lazily?
             self.aux.resize(cast(module.memory_map().aux_data_size()).to_usize(), 0);
+            self.accessible_aux_size = cast(module.memory_map().aux_data_size()).to_usize();
         }
     }
 
+    fn set_accessible_aux_size(&mut self, size: u32) {
+        self.accessible_aux_size = cast(size).to_usize();
+    }
+
     #[inline]
     fn get_memory_slice<'a>(&'a self, module: &'a Module, address: u32, length: u32) -> Option<&'a [u8]> {
         let memory_map = module.memory_map();
         let (start, memory_slice) = if address >= memory_map.aux_data_address() {
-            (memory_map.aux_data_address(), &self.aux)
+            (memory_map.aux_data_address(), &self.aux[..self.accessible_aux_size])
         } else if address >= memory_map.stack_address_low() {
-            (memory_map.stack_address_low(), &self.stack)
+            (memory_map.stack_address_low(), &self.stack[..])
         } else if address >= memory_map.rw_data_address() {
-            (memory_map.rw_data_address(), &self.rw_data)
+            (memory_map.rw_data_address(), &self.rw_data[..])
         } else if address >= memory_map.ro_data_address() {
             let module = module.interpreted_module().unwrap();
-            (memory_map.ro_data_address(), &module.ro_data)
+            (memory_map.ro_data_address(), &module.ro_data[..])
         } else {
             return None;
         };
@@ -153,11 +161,11 @@ impl BasicMemory {
     fn get_memory_slice_mut<const IS_EXTERNAL: bool>(&mut self, module: &Module, address: u32, length: u32) -> Option<&mut [u8]> {
         let memory_map = module.memory_map();
         let (start, memory_slice) = if IS_EXTERNAL && address >= memory_map.aux_data_address() {
-            (memory_map.aux_data_address(), &mut self.aux)
+            (memory_map.aux_data_address(), &mut self.aux[..self.accessible_aux_size])
         } else if address >= memory_map.stack_address_low() {
-            (memory_map.stack_address_low(), &mut self.stack)
+            (memory_map.stack_address_low(), &mut self.stack[..])
         } else if address >= memory_map.rw_data_address() {
-            (memory_map.rw_data_address(), &mut self.rw_data)
+            (memory_map.rw_data_address(), &mut self.rw_data[..])
         } else {
             return None;
         };
@@ -440,6 +448,11 @@ impl InterpretedInstance {
         self.next_program_counter_changed = true;
     }
 
+    pub fn set_accessible_aux_size(&mut self, size: u32) {
+        assert!(!self.module.is_dynamic_paging());
+        self.basic_memory.set_accessible_aux_size(size);
+    }
+
     #[allow(clippy::unused_self)]
     pub fn next_native_program_counter(&self) -> Option<usize> {
         None

crates/polkavm/src/sandbox.rs

@@ -122,6 +122,7 @@ pub(crate) trait Sandbox: Sized {
     fn next_program_counter(&self) -> Option<ProgramCounter>;
     fn next_native_program_counter(&self) -> Option<usize>;
     fn set_next_program_counter(&mut self, pc: ProgramCounter);
+    fn set_accessible_aux_size(&mut self, size: u32) -> Result<(), Self::Error>;
     fn reset_memory(&mut self) -> Result<(), Self::Error>;
     fn read_memory_into<'slice>(&self, address: u32, slice: &'slice mut [MaybeUninit<u8>]) -> Result<&'slice mut [u8], MemoryAccessError>;
     fn write_memory(&mut self, address: u32, data: &[u8]) -> Result<(), MemoryAccessError>;

crates/polkavm/src/sandbox/linux.rs

@@ -995,6 +995,8 @@ pub struct Sandbox {
     page_set: PageSet,
     dynamic_paging_enabled: bool,
     idle_regs: linux_raw::user_regs_struct,
+    aux_data_address: u32,
+    aux_data_length: u32,
 }
 
 impl Drop for Sandbox {
@@ -1598,6 +1600,8 @@ impl super::Sandbox for Sandbox {
             page_set: PageSet::new(),
             dynamic_paging_enabled: false,
             idle_regs,
+            aux_data_address: 0,
+            aux_data_length: 0,
         })
     }
 
@@ -1705,6 +1709,8 @@ impl super::Sandbox for Sandbox {
             reg.store(0, Ordering::Relaxed);
         }
 
+        self.aux_data_address = module.memory_map().aux_data_address();
+        self.aux_data_length = module.memory_map().aux_data_size();
         self.dynamic_paging_enabled = module.is_dynamic_paging();
         self.is_program_counter_valid = false;
         self.gas_metering = module.gas_metering();
@@ -1904,6 +1910,20 @@ impl super::Sandbox for Sandbox {
         }
     }
 
+    fn set_accessible_aux_size(&mut self, size: u32) -> Result<(), Error> {
+        assert!(!self.dynamic_paging_enabled);
+
+        let module = self.module.as_ref().unwrap();
+        self.aux_data_length = size;
+        self.vmctx().arg.store(self.aux_data_address, Ordering::Relaxed);
+        self.vmctx().arg2.store(size, Ordering::Relaxed);
+        self.vmctx().arg3.store(module.memory_map().aux_data_size(), Ordering::Relaxed);
+        self.vmctx()
+            .jump_into
+            .store(ZYGOTE_TABLES.1.ext_set_accessible_aux_size, Ordering::Relaxed);
+        self.wake_oneshot_and_expect_idle()
+    }
+
     fn reset_memory(&mut self) -> Result<(), Error> {
         if self.module.is_none() {
             return Err(Error::from_str("no module loaded into the sandbox"));

crates/polkavm/src/sandbox/polkavm-zygote

@@ -1823,6 +1823,56 @@ fn aux_data_works(config: Config) {
     assert_eq!(instance.read_u32(module.memory_map().aux_data_address()).unwrap(), 0);
 }
 
+fn aux_data_accessible_area(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::load_indirect_i32(Reg::A1, Reg::A0, 0), asm::ret()], &[]);
+
+    let blob = ProgramBlob::parse(builder.into_vec().into()).unwrap();
+    let mut module_config = ModuleConfig::new();
+    module_config.set_page_size(page_size);
+    module_config.set_aux_data_size(2_u32.pow(24));
+    let module = Module::from_blob(&engine, &module_config, blob).unwrap();
+    let offsets: Vec<_> = module
+        .blob()
+        .instructions(DefaultInstructionSet::default())
+        .map(|inst| inst.offset)
+        .collect();
+
+    let mut instance = module.instantiate().unwrap();
+    instance.set_accessible_aux_size(1).unwrap();
+    instance.write_u32(module.memory_map().aux_data_address(), 0x12345678).unwrap();
+    instance.set_reg(Reg::RA, crate::RETURN_TO_HOST);
+
+    instance.set_reg(Reg::A0, u64::from(module.memory_map().aux_data_address()));
+    instance.set_next_program_counter(offsets[0]);
+    match_interrupt!(instance.run().unwrap(), InterruptKind::Finished);
+    assert_eq!(instance.reg(Reg::A1), 0x12345678);
+
+    instance.set_reg(Reg::A0, u64::from(module.memory_map().aux_data_address() + page_size - 4));
+    instance.set_next_program_counter(offsets[0]);
+    match_interrupt!(instance.run().unwrap(), InterruptKind::Finished);
+
+    assert!(instance.read_u32(module.memory_map().aux_data_address() + page_size - 4).is_ok());
+
+    instance.set_reg(Reg::A0, u64::from(module.memory_map().aux_data_address() + page_size - 3));
+    instance.set_next_program_counter(offsets[0]);
+    match_interrupt!(instance.run().unwrap(), InterruptKind::Trap);
+
+    assert!(instance.read_u32(module.memory_map().aux_data_address() + page_size - 3).is_err());
+
+    instance.set_accessible_aux_size(page_size + 1).unwrap();
+
+    instance.set_reg(Reg::A0, u64::from(module.memory_map().aux_data_address() + page_size - 3));
+    instance.set_next_program_counter(offsets[0]);
+    match_interrupt!(instance.run().unwrap(), InterruptKind::Finished);
+
+    assert!(instance.read_u32(module.memory_map().aux_data_address() + page_size - 3).is_ok());
+}
+
 fn access_memory_from_host(config: Config) {
     let _ = env_logger::try_init();
     let engine = Engine::new(&config).unwrap();
@@ -2733,6 +2783,7 @@ run_tests! {
     invalid_instruction_after_fallthrough
     invalid_branch_target
     aux_data_works
+    aux_data_accessible_area
     access_memory_from_host
     sbrk_knob_works