Add Caller::read_u32

paritytech · Sep 5, 2023 · a8bb328 · a8bb328
1 parent ef48544
commit a8bb328
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/crates/polkavm/src/caller.rs b/crates/polkavm/src/caller.rs
@@ -1,5 +1,6 @@
 use crate::api::BackendAccess;
 use crate::tracer::Tracer;
+use core::mem::MaybeUninit;
 use polkavm_common::error::Trap;
 use polkavm_common::program::Reg;
 use polkavm_common::utils::{Access, AsUninitSliceMut};
@@ -90,6 +91,15 @@ impl CallerRaw {
         unsafe { self.access() }.read_memory_into_new_vec(address, length)
     }
 
+    unsafe fn read_u32(&self, address: u32) -> Result<u32, Trap> {
+        let mut buffer: MaybeUninit<[u8; 4]> = MaybeUninit::uninit();
+
+        // SAFETY: The caller will make sure that the invariants hold.
+        let slice = unsafe { self.read_memory_into_slice(address, &mut buffer) }?;
+        let value = u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]);
+        Ok(value)
+    }
+
     unsafe fn write_memory(&mut self, address: u32, data: &[u8]) -> Result<(), Trap> {
         log::trace!(
             "Writing memory (during hostcall): 0x{:x}-0x{:x} ({} bytes)",
@@ -192,6 +202,11 @@ impl<'a, T> Caller<'a, T> {
         unsafe { self.raw.read_memory_into_new_vec(address, length) }
     }
 
+    pub fn read_u32(&self, address: u32) -> Result<u32, Trap> {
+        // SAFETY: This can only be called from inside of `Caller::wrap` so this is always valid.
+        unsafe { self.raw.read_u32(address) }
+    }
+
     pub fn write_memory(&mut self, address: u32, data: &[u8]) -> Result<(), Trap> {
         // SAFETY: This can only be called from inside of `Caller::wrap` so this is always valid.
         unsafe { self.raw.write_memory(address, data) }
@@ -258,6 +273,13 @@ impl<T> CallerRef<T> {
         unsafe { (*self.raw).read_memory_into_new_vec(address, length) }
     }
 
+    pub fn read_u32(&self, address: u32) -> Result<u32, Trap> {
+        self.check_lifetime_or_panic();
+
+        // SAFETY: We've made sure the lifetime is valid.
+        unsafe { (*self.raw).read_u32(address) }
+    }
+
     pub fn write_memory(&mut self, address: u32, data: &[u8]) -> Result<(), Trap> {
         self.check_lifetime_or_panic();
 

diff --git a/crates/polkavm/src/lib.rs b/crates/polkavm/src/lib.rs
@@ -56,17 +56,12 @@ mod tests {
         linker
             .func_wrap("get_third_number", move |caller: Caller<State>| -> Result<u32, Trap> {
                 {
-                    let mut buffer = [0_u8; 4];
-                    caller.read_memory_into_slice(polkavm_common::abi::VM_ADDR_USER_STACK_HIGH - 4, &mut buffer)?;
-                    let value = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
+                    let value = caller.read_u32(polkavm_common::abi::VM_ADDR_USER_STACK_HIGH - 4)?;
                     assert_eq!(value, polkavm_common::abi::VM_ADDR_RETURN_TO_HOST);
                 }
                 {
                     let caller = caller.into_ref();
-
-                    let mut buffer = [0_u8; 4];
-                    caller.read_memory_into_slice(polkavm_common::abi::VM_ADDR_USER_STACK_HIGH - 4, &mut buffer)?;
-                    let value = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
+                    let value = caller.read_u32(polkavm_common::abi::VM_ADDR_USER_STACK_HIGH - 4)?;
                     assert_eq!(value, polkavm_common::abi::VM_ADDR_RETURN_TO_HOST);
 
                     let illegal_contraband = caller.data().illegal_contraband.clone();