Make read_memory fail on unitialized memory; add is_memory_accessible

koute · koute · commit fa3478f25cf6 · 2024-12-05T01:40:44.000+09:00
diff --git a/crates/polkavm/src/api.rs b/crates/polkavm/src/api.rs
@@ -1074,6 +1074,56 @@ impl RawInstance {
             .into_result("failed to reset the instance's memory"))
     }
 
+    /// Returns whether a given chunk of memory is accessible through [`read_memory_into`](Self::read_memory_into)/[`write_memory`](Self::write_memory).
+    ///
+    /// If `size` is zero then this will always return `true`.
+    pub fn is_memory_accessible(&self, address: u32, size: u32, is_writable: bool) -> bool {
+        if size == 0 {
+            return true;
+        }
+
+        if address < 0x10000 {
+            return false;
+        }
+
+        if u64::from(address) + cast(size).to_u64() > 0x100000000 {
+            return false;
+        }
+
+        #[inline]
+        fn is_within(range: core::ops::Range<u32>, address: u32, size: u32) -> bool {
+            let address_end = u64::from(address) + cast(size).to_u64();
+            address >= range.start && address_end <= u64::from(range.end)
+        }
+
+        if !self.module.is_dynamic_paging() {
+            let map = self.module.memory_map();
+            if is_within(map.stack_range(), address, size) {
+                return true;
+            }
+
+            let heap_size = self.heap_size();
+            let heap_top = map.heap_base() + heap_size;
+            let heap_top = self.module.round_to_page_size_up(heap_top);
+            if is_within(map.rw_data_address()..heap_top, address, size) {
+                return true;
+            }
+
+            let aux_size = access_backend!(self.backend, |backend| backend.accessible_aux_size());
+            if is_within(map.aux_data_address()..map.aux_data_address() + aux_size, address, size) {
+                return true;
+            }
+
+            if !is_writable && is_within(map.ro_data_range(), address, size) {
+                return true;
+            }
+
+            false
+        } else {
+            access_backend!(self.backend, |backend| backend.is_memory_accessible(address, size, is_writable))
+        }
+    }
+
     /// Reads the VM's memory.
     pub fn read_memory_into<'slice, B>(&self, address: u32, buffer: &'slice mut B) -> Result<&'slice mut [u8], MemoryAccessError>
     where
@@ -1124,6 +1174,21 @@ impl RawInstance {
                 panic!("read_memory: crosscheck mismatch, range = 0x{address:x}..0x{address_end:x}, interpreter = {expected_success}, backend = {success}");
             }
         }
+
+        #[cfg(debug_assertions)]
+        {
+            let is_accessible = self.is_memory_accessible(address, cast(length).assert_always_fits_in_u32(), false);
+            if is_accessible != result.is_ok() {
+                panic!(
+                    "'read_memory_into' doesn't match with 'is_memory_accessible' for 0x{:x}-0x{:x} (read_memory_into = {}, is_memory_accessible = {})",
+                    address,
+                    cast(address).to_usize() + length,
+                    result.is_ok(),
+                    is_accessible,
+                );
+            }
+        }
+
         result
     }
 
@@ -1161,6 +1226,20 @@ impl RawInstance {
             }
         }
 
+        #[cfg(debug_assertions)]
+        {
+            let is_accessible = self.is_memory_accessible(address, cast(data.len()).assert_always_fits_in_u32(), true);
+            if is_accessible != result.is_ok() {
+                panic!(
+                    "'write_memory' doesn't match with 'is_memory_accessible' for 0x{:x}-0x{:x} (write_memory = {}, is_memory_accessible = {})",
+                    address,
+                    cast(address).to_usize() + data.len(),
+                    result.is_ok(),
+                    is_accessible,
+                );
+            }
+        }
+
         result
     }
 
diff --git a/crates/polkavm/src/interpreter.rs b/crates/polkavm/src/interpreter.rs
@@ -131,6 +131,10 @@ impl BasicMemory {
         }
     }
 
+    fn accessible_aux_size(&self) -> u32 {
+        cast(self.accessible_aux_size).assert_always_fits_in_u32()
+    }
+
     fn set_accessible_aux_size(&mut self, size: u32) {
         self.accessible_aux_size = cast(size).to_usize();
     }
@@ -253,38 +257,43 @@ macro_rules! emit_branch {
     };
 }
 
-fn each_page(module: &Module, address: u32, length: u32, callback: impl FnMut(u32, usize, usize, usize)) {
+fn each_page<E>(
+    module: &Module,
+    address: u32,
+    length: u32,
+    callback: impl FnMut(u32, usize, usize, usize) -> Result<(), E>,
+) -> Result<(), E> {
     let page_size = module.memory_map().page_size();
     let page_address_lo = module.round_to_page_size_down(address);
     let page_address_hi = module.round_to_page_size_down(address + (length - 1));
     each_page_impl(page_size, page_address_lo, page_address_hi, address, length, callback)
 }
 
-fn each_page_impl(
+fn each_page_impl<E>(
     page_size: u32,
     page_address_lo: u32,
     page_address_hi: u32,
     address: u32,
     length: u32,
-    mut callback: impl FnMut(u32, usize, usize, usize),
-) {
+    mut callback: impl FnMut(u32, usize, usize, usize) -> Result<(), E>,
+) -> Result<(), E> {
     let page_size = cast(page_size).to_usize();
     let length = cast(length).to_usize();
 
     let initial_page_offset = cast(address).to_usize() - cast(page_address_lo).to_usize();
     let initial_chunk_length = core::cmp::min(length, page_size - initial_page_offset);
-    callback(page_address_lo, initial_page_offset, 0, initial_chunk_length);
+    callback(page_address_lo, initial_page_offset, 0, initial_chunk_length)?;
 
     if page_address_lo == page_address_hi {
-        return;
+        return Ok(());
     }
 
     let mut page_address_lo = cast(page_address_lo).to_u64();
     let page_address_hi = cast(page_address_hi).to_u64();
     page_address_lo += cast(page_size).to_u64();
     let mut buffer_offset = initial_chunk_length;
     while page_address_lo < page_address_hi {
-        callback(cast(page_address_lo).assert_always_fits_in_u32(), 0, buffer_offset, page_size);
+        callback(cast(page_address_lo).assert_always_fits_in_u32(), 0, buffer_offset, page_size)?;
         buffer_offset += page_size;
         page_address_lo += cast(page_size).to_u64();
     }
@@ -304,16 +313,18 @@ fn test_each_page() {
         let page_address_lo = address / page_size * page_size;
         let page_address_hi = (address + (length - 1)) / page_size * page_size;
         let mut output = Vec::new();
-        each_page_impl(
+        each_page_impl::<()>(
             page_size,
             page_address_lo,
             page_address_hi,
             address,
             length,
             |page_address, page_offset, buffer_offset, length| {
                 output.push((page_address, page_offset, buffer_offset, length));
+                Ok(())
             },
-        );
+        )
+        .unwrap();
         output
     }
 
@@ -448,6 +459,11 @@ impl InterpretedInstance {
         self.next_program_counter_changed = true;
     }
 
+    pub fn accessible_aux_size(&self) -> u32 {
+        assert!(!self.module.is_dynamic_paging());
+        self.basic_memory.accessible_aux_size()
+    }
+
     pub fn set_accessible_aux_size(&mut self, size: u32) {
         assert!(!self.module.is_dynamic_paging());
         self.basic_memory.set_accessible_aux_size(size);
@@ -458,6 +474,21 @@ impl InterpretedInstance {
         None
     }
 
+    pub fn is_memory_accessible(&self, address: u32, size: u32, _is_writable: bool) -> bool {
+        assert!(self.module.is_dynamic_paging());
+
+        // TODO: This is very slow.
+        let result = each_page(&self.module, address, size, |page_address, _, _, _| {
+            if !self.dynamic_memory.pages.contains_key(&page_address) {
+                Err(())
+            } else {
+                Ok(())
+            }
+        });
+
+        result.is_ok()
+    }
+
     pub fn read_memory_into<'slice>(
         &self,
         address: u32,
@@ -491,12 +522,16 @@ impl InterpretedInstance {
                         if let Some(page) = page {
                             let src = page.as_ptr().add(page_offset);
                             core::ptr::copy_nonoverlapping(src, dst, length);
+                            Ok(())
                         } else {
-                            core::ptr::write_bytes(dst, 0, length);
+                            Err(MemoryAccessError::OutOfRangeAccess {
+                                address: page_address + cast(page_offset).assert_always_fits_in_u32(),
+                                length: cast(length).to_u64(),
+                            })
                         }
                     }
                 },
-            );
+            )?;
 
             // SAFETY: The buffer was initialized.
             Ok(unsafe { slice_assume_init_mut(buffer) })
@@ -519,15 +554,17 @@ impl InterpretedInstance {
         } else {
             let dynamic_memory = &mut self.dynamic_memory;
             let page_size = self.module.memory_map().page_size();
-            each_page(
+            each_page::<()>(
                 &self.module,
                 address,
                 cast(data.len()).assert_always_fits_in_u32(),
                 move |page_address, page_offset, buffer_offset, length| {
                     let page = dynamic_memory.pages.entry(page_address).or_insert_with(|| empty_page(page_size));
                     page[page_offset..page_offset + length].copy_from_slice(&data[buffer_offset..buffer_offset + length]);
+                    Ok(())
                 },
-            );
+            )
+            .unwrap();
         }
 
         Ok(())
@@ -546,20 +583,23 @@ impl InterpretedInstance {
         } else {
             let dynamic_memory = &mut self.dynamic_memory;
             let page_size = self.module.memory_map().page_size();
-            each_page(
+            each_page::<()>(
                 &self.module,
                 address,
                 length,
                 move |page_address, page_offset, _, length| match dynamic_memory.pages.entry(page_address) {
                     Entry::Occupied(mut entry) => {
                         let page = entry.get_mut();
                         page[page_offset..page_offset + length].fill(0);
+                        Ok(())
                     }
                     Entry::Vacant(entry) => {
                         entry.insert(empty_page(page_size));
+                        Ok(())
                     }
                 },
-            );
+            )
+            .unwrap();
         }
 
         Ok(())
@@ -573,9 +613,11 @@ impl InterpretedInstance {
             todo!()
         } else {
             let dynamic_memory = &mut self.dynamic_memory;
-            each_page(&self.module, address, length, move |page_address, _, _, _| {
+            each_page::<()>(&self.module, address, length, move |page_address, _, _, _| {
                 dynamic_memory.pages.remove(&page_address);
-            });
+                Ok(())
+            })
+            .unwrap();
         }
     }
 
diff --git a/crates/polkavm/src/sandbox.rs b/crates/polkavm/src/sandbox.rs
@@ -122,7 +122,9 @@ 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 accessible_aux_size(&self) -> u32;
     fn set_accessible_aux_size(&mut self, size: u32) -> Result<(), Self::Error>;
+    fn is_memory_accessible(&self, address: u32, size: u32, is_writable: bool) -> bool;
     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>;
diff --git a/crates/polkavm/src/sandbox/linux.rs b/crates/polkavm/src/sandbox/linux.rs
@@ -1910,6 +1910,11 @@ impl super::Sandbox for Sandbox {
         }
     }
 
+    fn accessible_aux_size(&self) -> u32 {
+        assert!(!self.dynamic_paging_enabled);
+        self.aux_data_length
+    }
+
     fn set_accessible_aux_size(&mut self, size: u32) -> Result<(), Error> {
         assert!(!self.dynamic_paging_enabled);
 
@@ -1924,6 +1929,15 @@ impl super::Sandbox for Sandbox {
         self.wake_oneshot_and_expect_idle()
     }
 
+    fn is_memory_accessible(&self, address: u32, size: u32, _is_writable: bool) -> bool {
+        assert!(self.dynamic_paging_enabled);
+
+        let module = self.module.as_ref().unwrap();
+        let page_start = module.address_to_page(module.round_to_page_size_down(address));
+        let page_end = module.address_to_page(module.round_to_page_size_down(address + size));
+        self.page_set.contains((page_start, page_end))
+    }
+
     fn reset_memory(&mut self) -> Result<(), Error> {
         if self.module.is_none() {
             return Err(Error::from_str("no module loaded into the sandbox"));
@@ -1957,9 +1971,7 @@ impl super::Sandbox for Sandbox {
             let page_start = module.address_to_page(module.round_to_page_size_down(address));
             let page_end = module.address_to_page(module.round_to_page_size_down(address + slice.len() as u32));
             if !self.page_set.contains((page_start, page_end)) {
-                unsafe {
-                    core::ptr::write_bytes(slice.as_mut_ptr().cast::<u8>(), 0, slice.len());
-                }
+                return Err(MemoryAccessError::Error("incomplete read".into()));
             } else {
                 let memory: &[core::mem::MaybeUninit<u8>] =
                     unsafe { core::slice::from_raw_parts(self.memory_mmap.as_ptr().cast(), self.memory_mmap.len()) };
diff --git a/crates/polkavm/src/tests.rs b/crates/polkavm/src/tests.rs
@@ -927,7 +927,7 @@ fn dynamic_paging_reading_does_not_resolve_segfaults(mut engine_config: Config)
     instance.set_next_program_counter(offsets[0]);
     let segfault = expect_segfault(instance.run().unwrap());
     assert_eq!(segfault.page_address, 0x10000);
-    assert_eq!(instance.read_u32(0x10000).unwrap(), 0x00000000);
+    assert!(instance.read_u32(0x10000).is_err());
 
     let segfault = expect_segfault(instance.run().unwrap());
     assert_eq!(segfault.page_address, 0x10000);
@@ -1175,7 +1175,11 @@ fn dynamic_paging_worker_recycle_turn_dynamic_paging_on_and_off(mut engine_confi
             module_static.instantiate().unwrap()
         };
 
-        assert_eq!(instance.read_u32(0x20000).unwrap(), 0);
+        if !is_dynamic {
+            assert_eq!(instance.read_u32(0x20000).unwrap(), 0);
+        } else {
+            assert!(instance.read_u32(0x20000).is_err());
+        }
 
         instance.set_reg(Reg::RA, crate::RETURN_TO_HOST);
         instance.set_next_program_counter(ProgramCounter(0));
