Extend CustomCodeMemory interface allowing for using previously-mapped images

`Module::deserialize_raw` and `CustomCodeMemory` were introduced with aim of
supporting `no_std` platforms. With small tweaks they can also be used for using
Wasmtime on platforms that have full-blown virtual memory capabilities, but
doesn't allow for directly mapping executable pages from user code instead
limiting that capability for system loader.
Adding features necessary for such platforms was previously attempted by #8245.

There are currently two issues with using `Module::deserialize_raw` for
using images either statically linked into embedder executable or dynamically
loaded shared objects:

- `CodeMemory::publish` will initially make entire image read-only,
  destroying executable permissions that cannot be restored by user code.
  This will happen even if `CustomCodeMemory` is provided.
- `CodeMemory::publish` will attempt to unconditionally register
  unwind information. As these are already properly registered by the system
  loader this is superfluous at best, or could fail module loading if system
  decides to return error on attempt for double-registration.

This commit solves these issues by:
- Moving responsibility for making image RO to `CustomCodeMemory` hook.
- Making `CustomCodeMemory` publishing hook return enum that tells
  what steps (only mapping, or mapping with registration) were
  performed, using default implementation only for actions not reported by the hook.
- Additionally `CustomCodeMemory` hooks also receive pointers to the
  entire image, not only executable section. This allows the embedder
  to keep track of the images and unload them in case they were
  loaded as dynamic shared object.

Author: Milek7

crates/wasmtime/src/engine.rs

@@ -3,6 +3,8 @@ use crate::prelude::*;
 #[cfg(feature = "runtime")]
 pub use crate::runtime::code_memory::CustomCodeMemory;
 #[cfg(feature = "runtime")]
+pub use crate::runtime::code_memory::HandledByCustomCodeMemory;
+#[cfg(feature = "runtime")]
 use crate::runtime::type_registry::TypeRegistry;
 #[cfg(feature = "runtime")]
 use crate::runtime::vm::GcRuntime;

crates/wasmtime/src/runtime/code_memory.rs

@@ -44,10 +44,17 @@ impl Drop for CodeMemory {
         // If there is a custom code memory handler, restore the
         // original (non-executable) state of the memory.
         if let Some(mem) = self.custom_code_memory.as_ref() {
-            if self.published && self.needs_executable {
-                let text = self.text();
-                mem.unpublish_executable(text.as_ptr(), text.len())
-                    .expect("Executable memory unpublish failed");
+            if self.published {
+                mem.unpublish_image(
+                    self.mmap.as_ptr(),
+                    self.mmap.len(),
+                    if self.needs_executable {
+                        Some(self.text.clone())
+                    } else {
+                        None
+                    },
+                )
+                .expect("Executable memory unpublish failed");
             }
         }
 
@@ -64,6 +71,23 @@ fn _assert() {
     _assert_send_sync::<CodeMemory>();
 }
 
+#[derive(PartialEq, PartialOrd)]
+/// Actions that were performed by `CustomCodeMemory::publish_image` hook.
+pub enum HandledByCustomCodeMemory {
+    /// Nothing was performed, publishing will be done by the
+    /// default Wasmtime implementation.
+    Nothing,
+
+    /// Only memory mapping was performed by the hook,
+    /// registering unwind and debug information will
+    /// be done by default Wasmtime implementation.
+    Mapping,
+
+    /// Everything was performed by the hook and no further
+    /// publishing actions are required.
+    MappingAndRegistration,
+}
+
 /// Interface implemented by an embedder to provide custom
 /// implementations of code-memory protection and execute permissions.
 pub trait CustomCodeMemory: Send + Sync {
@@ -75,30 +99,50 @@ pub trait CustomCodeMemory: Send + Sync {
     /// of virtual memory are disabled.
     fn required_alignment(&self) -> usize;
 
-    /// Publish a region of memory as executable.
+    /// Publish an image.
     ///
-    /// This should update permissions from the default RW
-    /// (readable/writable but not executable) to RX
-    /// (readable/executable but not writable), enforcing W^X
-    /// discipline.
+    /// This should ensure that the provided image region is mapped
+    /// read-only, with the exception of `code` area that (if provided)
+    /// should be mapped as readable and executable.
+    ///
+    /// Return value describes which actions were completed, while other
+    /// steps will be handled by Wasmtime default implementation.
     ///
     /// If the platform requires any data/instruction coherence
     /// action, that should be performed as part of this hook as well.
     ///
-    /// `ptr` and `ptr.offset(len)` are guaranteed to be aligned as
-    /// per `required_alignment()`.
-    fn publish_executable(&self, ptr: *const u8, len: usize) -> anyhow::Result<()>;
-
-    /// Unpublish a region of memory.
+    /// If `image` is the image provided to `Module::deserialize_raw`
+    /// which was already prepared for execution, this hook might
+    /// be a no-op reporting all actions as completed.
     ///
-    /// This should perform the opposite effect of `make_executable`,
-    /// switching a range of memory back from RX (readable/executable)
-    /// to RW (readable/writable). It is guaranteed that no code is
+    /// Values in `code` range are relative to `image` pointer.
+    /// `code` is guaranteed to be aligned as per `required_alignment()`.
+    fn publish_image(
+        &self,
+        image: *const u8,
+        len: usize,
+        code: Option<Range<usize>>,
+    ) -> anyhow::Result<HandledByCustomCodeMemory>;
+
+    /// Unpublish an image.
+    ///
+    /// This should perform the opposite effect of `publish_image`.
+    ///
+    /// If `image` is the image provided to `Module::deserialize_raw`,
+    /// this hook might be used for unloading that image from address space.
+    ///
+    /// Otherwise it should switch a range of memory back to
+    /// readable and writable. It is guaranteed that no code is
     /// running anymore from this region.
     ///
-    /// `ptr` and `ptr.offset(len)` are guaranteed to be aligned as
-    /// per `required_alignment()`.
-    fn unpublish_executable(&self, ptr: *const u8, len: usize) -> anyhow::Result<()>;
+    /// Values in `code` range are relative to `image` pointer.
+    /// `code` is guaranteed to be aligned as per `required_alignment()`.
+    fn unpublish_image(
+        &self,
+        image: *const u8,
+        len: usize,
+        code: Option<Range<usize>>,
+    ) -> anyhow::Result<()>;
 }
 
 impl CodeMemory {
@@ -297,25 +341,29 @@ impl CodeMemory {
         //   both the actual unwinding tables as well as the validity of the
         //   pointers we pass in itself.
         unsafe {
-            // Next freeze the contents of this image by making all of the
-            // memory readonly. Nothing after this point should ever be modified
-            // so commit everything. For a compiled-in-memory image this will
-            // mean IPIs to evict writable mappings from other cores. For
-            // loaded-from-disk images this shouldn't result in IPIs so long as
-            // there weren't any relocations because nothing should have
-            // otherwise written to the image at any point either.
-            //
-            // Note that if virtual memory is disabled this is skipped because
-            // we aren't able to make it readonly, but this is just a
-            // defense-in-depth measure and isn't required for correctness.
-            #[cfg(has_virtual_memory)]
-            if self.mmap.supports_virtual_memory() {
-                self.mmap.make_readonly(0..self.mmap.len())?;
-            }
+            // Potentially call embedder-provided custom publishing implementation,
+            // and only perform publishing steps that weren't done by it.
+            let handled = self.custom_publish()?;
+
+            if handled < HandledByCustomCodeMemory::Mapping {
+                // Next freeze the contents of this image by making all of the
+                // memory readonly. Nothing after this point should ever be modified
+                // so commit everything. For a compiled-in-memory image this will
+                // mean IPIs to evict writable mappings from other cores. For
+                // loaded-from-disk images this shouldn't result in IPIs so long as
+                // there weren't any relocations because nothing should have
+                // otherwise written to the image at any point either.
+                //
+                // Note that if virtual memory is disabled this is skipped because
+                // we aren't able to make it readonly, but this is just a
+                // defense-in-depth measure and isn't required for correctness.
+                #[cfg(has_virtual_memory)]
+                if self.mmap.supports_virtual_memory() {
+                    self.mmap.make_readonly(0..self.mmap.len())?;
+                }
 
-            // Switch the executable portion from readonly to read/execute.
-            if self.needs_executable {
-                if !self.custom_publish()? {
+                // Switch the executable portion from readonly to read/execute.
+                if self.needs_executable {
                     if !self.mmap.supports_virtual_memory() {
                         bail!("this target requires virtual memory to be enabled");
                     }
@@ -343,39 +391,45 @@ impl CodeMemory {
                 }
             }
 
-            // With all our memory set up use the platform-specific
-            // `UnwindRegistration` implementation to inform the general
-            // runtime that there's unwinding information available for all
-            // our just-published JIT functions.
-            self.register_unwind_info()?;
+            if handled < HandledByCustomCodeMemory::MappingAndRegistration {
+                // With all our memory set up use the platform-specific
+                // `UnwindRegistration` implementation to inform the general
+                // runtime that there's unwinding information available for all
+                // our just-published JIT functions.
+                self.register_unwind_info()?;
 
-            #[cfg(feature = "debug-builtins")]
-            self.register_debug_image()?;
+                #[cfg(feature = "debug-builtins")]
+                self.register_debug_image()?;
+            }
         }
 
         Ok(())
     }
 
-    fn custom_publish(&mut self) -> Result<bool> {
+    fn custom_publish(&mut self) -> Result<HandledByCustomCodeMemory> {
         if let Some(mem) = self.custom_code_memory.as_ref() {
-            let text = self.text();
-            // The text section should be aligned to
-            // `custom_code_memory.required_alignment()` due to a
-            // combination of two invariants:
-            //
-            // - MmapVec aligns its start address, even in owned-Vec mode; and
-            // - The text segment inside the ELF image will be aligned according
-            //   to the platform's requirements.
-            let text_addr = text.as_ptr() as usize;
-            assert_eq!(text_addr & (mem.required_alignment() - 1), 0);
-
-            // The custom code memory handler will ensure the
-            // memory is executable and also handle icache
-            // coherence.
-            mem.publish_executable(text.as_ptr(), text.len())?;
-            Ok(true)
+            if self.needs_executable {
+                // The text section should be aligned to
+                // `custom_code_memory.required_alignment()` due to a
+                // combination of two invariants:
+                //
+                // - MmapVec aligns its start address, even in owned-Vec mode; and
+                // - The text segment inside the ELF image will be aligned according
+                //   to the platform's requirements.
+                assert_eq!(
+                    (self.text().as_ptr() as usize) & (mem.required_alignment() - 1),
+                    0
+                );
+
+                // The custom code memory handler will ensure the
+                // memory is executable and also handle icache
+                // coherence.
+                mem.publish_image(self.mmap.as_ptr(), self.mmap.len(), Some(self.text.clone()))
+            } else {
+                mem.publish_image(self.mmap.as_ptr(), self.mmap.len(), None)
+            }
         } else {
-            Ok(false)
+            Ok(HandledByCustomCodeMemory::Nothing)
         }
     }
 

tests/all/custom_code_memory.rs

@@ -1,5 +1,6 @@
 #[cfg(all(not(target_os = "windows"), not(miri)))]
 mod not_for_windows {
+    use core::ops::Range;
     use rustix::mm::{MprotectFlags, mprotect};
     use rustix::param::page_size;
     use std::sync::Arc;
@@ -11,18 +12,31 @@ mod not_for_windows {
             page_size()
         }
 
-        fn publish_executable(&self, ptr: *const u8, len: usize) -> anyhow::Result<()> {
+        fn publish_image(
+            &self,
+            ptr: *const u8,
+            len: usize,
+            code: Option<Range<usize>>,
+        ) -> anyhow::Result<HandledByCustomCodeMemory> {
             unsafe {
-                mprotect(
-                    ptr as *mut _,
-                    len,
-                    MprotectFlags::READ | MprotectFlags::EXEC,
-                )?;
+                mprotect(ptr as *mut _, len, MprotectFlags::READ)?;
+                if let Some(code) = code {
+                    mprotect(
+                        ptr.add(code.start) as *mut _,
+                        code.end - code.start,
+                        MprotectFlags::READ | MprotectFlags::EXEC,
+                    )?;
+                }
             }
-            Ok(())
+            Ok(HandledByCustomCodeMemory::Mapping)
         }
 
-        fn unpublish_executable(&self, ptr: *const u8, len: usize) -> anyhow::Result<()> {
+        fn unpublish_image(
+            &self,
+            ptr: *const u8,
+            len: usize,
+            _code: Option<Range<usize>>,
+        ) -> anyhow::Result<()> {
             unsafe {
                 mprotect(
                     ptr as *mut _,