Wasmtime: support a notion of "custom code publisher".

crates/environ/src/compile/module_artifacts.rs

@@ -274,12 +274,16 @@ impl<'a> ObjectBuilder<'a> {
 
 /// A type which can be the result of serializing an object.
 pub trait FinishedObject: Sized {
+    /// State required for `finish_object`, if any.
+    type State;
+
     /// Emit the object as `Self`.
-    fn finish_object(obj: ObjectBuilder<'_>) -> Result<Self>;
+    fn finish_object(obj: ObjectBuilder<'_>, state: &Self::State) -> Result<Self>;
 }
 
 impl FinishedObject for Vec<u8> {
-    fn finish_object(obj: ObjectBuilder<'_>) -> Result<Self> {
+    type State = ();
+    fn finish_object(obj: ObjectBuilder<'_>, _state: &Self::State) -> Result<Self> {
         let mut result = ObjectVec::default();
         obj.finish(&mut result)?;
         return Ok(result.0);

crates/wasmtime/src/compile.rs

@@ -64,6 +64,7 @@ pub(crate) fn build_artifacts<T: FinishedObject>(
     engine: &Engine,
     wasm: &[u8],
     dwarf_package: Option<&[u8]>,
+    obj_state: &T::State,
 ) -> Result<(T, Option<(CompiledModuleInfo, ModuleTypes)>)> {
     let tunables = engine.tunables();
 
@@ -111,7 +112,7 @@ pub(crate) fn build_artifacts<T: FinishedObject>(
     let info = compilation_artifacts.unwrap_as_module_info();
     let types = types.finish();
     object.serialize_info(&(&info, &types));
-    let result = T::finish_object(object)?;
+    let result = T::finish_object(object, obj_state)?;
 
     Ok((result, Some((info, types))))
 }
@@ -128,6 +129,7 @@ pub(crate) fn build_component_artifacts<T: FinishedObject>(
     engine: &Engine,
     binary: &[u8],
     _dwarf_package: Option<&[u8]>,
+    obj_state: &T::State,
 ) -> Result<(T, Option<wasmtime_environ::component::ComponentArtifacts>)> {
     use wasmtime_environ::component::{
         CompiledComponentInfo, ComponentArtifacts, ComponentTypesBuilder,
@@ -186,7 +188,7 @@ pub(crate) fn build_component_artifacts<T: FinishedObject>(
     };
     object.serialize_info(&artifacts);
 
-    let result = T::finish_object(object)?;
+    let result = T::finish_object(object, obj_state)?;
     Ok((result, Some(artifacts)))
 }
 

crates/wasmtime/src/compile/code_builder.rs

@@ -274,7 +274,7 @@ impl<'a> CodeBuilder<'a> {
     pub fn compile_module_serialized(&self) -> Result<Vec<u8>> {
         let wasm = self.get_wasm()?;
         let dwarf_package = self.get_dwarf_package();
-        let (v, _) = super::build_artifacts(self.engine, &wasm, dwarf_package.as_deref())?;
+        let (v, _) = super::build_artifacts(self.engine, &wasm, dwarf_package.as_deref(), &())?;
         Ok(v)
     }
 
@@ -284,7 +284,7 @@ impl<'a> CodeBuilder<'a> {
     #[cfg(feature = "component-model")]
     pub fn compile_component_serialized(&self) -> Result<Vec<u8>> {
         let bytes = self.get_wasm()?;
-        let (v, _) = super::build_component_artifacts(self.engine, &bytes, None)?;
+        let (v, _) = super::build_component_artifacts(self.engine, &bytes, None, &())?;
         Ok(v)
     }
 }

crates/wasmtime/src/compile/runtime.rs

@@ -9,9 +9,15 @@ use std::sync::Arc;
 use wasmtime_environ::{FinishedObject, ObjectBuilder, ObjectKind};
 
 impl<'a> CodeBuilder<'a> {
-    fn compile_cached<T>(
+    fn compile_cached<T, S>(
         &self,
-        build_artifacts: fn(&Engine, &[u8], Option<&[u8]>) -> Result<(MmapVecWrapper, Option<T>)>,
+        build_artifacts: fn(
+            &Engine,
+            &[u8],
+            Option<&[u8]>,
+            &S,
+        ) -> Result<(MmapVecWrapper, Option<T>)>,
+        state: &S,
     ) -> Result<(Arc<CodeMemory>, Option<T>)> {
         let wasm = self.get_wasm()?;
         let dwarf_package = self.get_dwarf_package();
@@ -28,24 +34,32 @@ impl<'a> CodeBuilder<'a> {
                 &dwarf_package,
                 // Don't hash this as it's just its own "pure" function pointer.
                 NotHashed(build_artifacts),
+                // Don't hash the FinishedObject state: this contains
+                // things like required runtime alignment, and does
+                // not impact the compilation result itself.
+                NotHashed(state),
             );
             let (code, info_and_types) =
                 wasmtime_cache::ModuleCacheEntry::new("wasmtime", self.engine.cache_config())
                     .get_data_raw(
                         &state,
                         // Cache miss, compute the actual artifacts
-                        |(engine, wasm, dwarf_package, build_artifacts)| -> Result<_> {
-                            let (mmap, info) =
-                                (build_artifacts.0)(engine.0, wasm, dwarf_package.as_deref())?;
-                            let code = publish_mmap(mmap.0)?;
+                        |(engine, wasm, dwarf_package, build_artifacts, state)| -> Result<_> {
+                            let (mmap, info) = (build_artifacts.0)(
+                                engine.0,
+                                wasm,
+                                dwarf_package.as_deref(),
+                                state.0,
+                            )?;
+                            let code = publish_mmap(engine.0, mmap.0)?;
                             Ok((code, info))
                         },
                         // Implementation of how to serialize artifacts
-                        |(_engine, _wasm, _, _), (code, _info_and_types)| {
+                        |(_engine, _wasm, _, _, _), (code, _info_and_types)| {
                             Some(code.mmap().to_vec())
                         },
                         // Cache hit, deserialize the provided artifacts
-                        |(engine, wasm, _, _), serialized_bytes| {
+                        |(engine, wasm, _, _, _), serialized_bytes| {
                             let kind = if wasmparser::Parser::is_component(&wasm) {
                                 ObjectKind::Component
                             } else {
@@ -61,8 +75,8 @@ impl<'a> CodeBuilder<'a> {
         #[cfg(not(feature = "cache"))]
         {
             let (mmap, info_and_types) =
-                build_artifacts(self.engine, &wasm, dwarf_package.as_deref())?;
-            let code = publish_mmap(mmap.0)?;
+                build_artifacts(self.engine, &wasm, dwarf_package.as_deref(), state)?;
+            let code = publish_mmap(self.engine, mmap.0)?;
             return Ok((code, info_and_types));
         }
 
@@ -79,30 +93,52 @@ impl<'a> CodeBuilder<'a> {
     /// Note that this method will cache compilations if the `cache` feature is
     /// enabled and turned on in [`Config`](crate::Config).
     pub fn compile_module(&self) -> Result<Module> {
-        let (code, info_and_types) = self.compile_cached(super::build_artifacts)?;
+        let custom_alignment = self.custom_alignment();
+        let (code, info_and_types) =
+            self.compile_cached(super::build_artifacts, &custom_alignment)?;
         Module::from_parts(self.engine, code, info_and_types)
     }
 
     /// Same as [`CodeBuilder::compile_module`] except that it compiles a
     /// [`Component`] instead of a module.
     #[cfg(feature = "component-model")]
     pub fn compile_component(&self) -> Result<Component> {
-        let (code, artifacts) = self.compile_cached(super::build_component_artifacts)?;
+        let custom_alignment = self.custom_alignment();
+        let (code, artifacts) =
+            self.compile_cached(super::build_component_artifacts, &custom_alignment)?;
         Component::from_parts(self.engine, code, artifacts)
     }
+
+    fn custom_alignment(&self) -> CustomAlignment {
+        CustomAlignment {
+            alignment: self
+                .engine
+                .custom_code_memory()
+                .map(|c| c.required_alignment())
+                .unwrap_or(1),
+        }
+    }
 }
 
-fn publish_mmap(mmap: MmapVec) -> Result<Arc<CodeMemory>> {
-    let mut code = CodeMemory::new(mmap)?;
+fn publish_mmap(engine: &Engine, mmap: MmapVec) -> Result<Arc<CodeMemory>> {
+    let mut code = CodeMemory::new(engine, mmap)?;
     code.publish()?;
     Ok(Arc::new(code))
 }
 
 pub(crate) struct MmapVecWrapper(pub MmapVec);
 
+/// Custom alignment requirements from the Engine for
+/// produced-at-runtime-in-memory code artifacts.
+pub(crate) struct CustomAlignment {
+    alignment: usize,
+}
+
 impl FinishedObject for MmapVecWrapper {
-    fn finish_object(obj: ObjectBuilder<'_>) -> Result<Self> {
+    type State = CustomAlignment;
+    fn finish_object(obj: ObjectBuilder<'_>, align: &CustomAlignment) -> Result<Self> {
         let mut result = ObjectMmap::default();
+        result.alignment = align.alignment;
         return match obj.finish(&mut result) {
             Ok(()) => {
                 assert!(result.mmap.is_some(), "no reserve");
@@ -127,6 +163,7 @@ impl FinishedObject for MmapVecWrapper {
         struct ObjectMmap {
             mmap: Option<MmapVec>,
             len: usize,
+            alignment: usize,
             err: Option<Error>,
         }
 
@@ -137,7 +174,7 @@ impl FinishedObject for MmapVecWrapper {
 
             fn reserve(&mut self, additional: usize) -> Result<(), ()> {
                 assert!(self.mmap.is_none(), "cannot reserve twice");
-                self.mmap = match MmapVec::with_capacity(additional) {
+                self.mmap = match MmapVec::with_capacity_and_alignment(additional, self.alignment) {
                     Ok(mmap) => Some(mmap),
                     Err(e) => {
                         self.err = Some(e);

crates/wasmtime/src/config.rs

@@ -29,6 +29,8 @@ use crate::stack::{StackCreator, StackCreatorProxy};
 #[cfg(feature = "async")]
 use wasmtime_fiber::RuntimeFiberStackCreator;
 
+#[cfg(feature = "runtime")]
+pub use crate::runtime::code_memory::CustomCodeMemory;
 #[cfg(feature = "pooling-allocator")]
 pub use crate::runtime::vm::MpkEnabled;
 #[cfg(all(feature = "incremental-cache", feature = "cranelift"))]
@@ -133,6 +135,8 @@ pub struct Config {
     pub(crate) cache_config: CacheConfig,
     #[cfg(feature = "runtime")]
     pub(crate) mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>,
+    #[cfg(feature = "runtime")]
+    pub(crate) custom_code_memory: Option<Arc<dyn CustomCodeMemory>>,
     pub(crate) allocation_strategy: InstanceAllocationStrategy,
     pub(crate) max_wasm_stack: usize,
     /// Explicitly enabled features via `Config::wasm_*` methods. This is a
@@ -233,6 +237,8 @@ impl Config {
             profiling_strategy: ProfilingStrategy::None,
             #[cfg(feature = "runtime")]
             mem_creator: None,
+            #[cfg(feature = "runtime")]
+            custom_code_memory: None,
             allocation_strategy: InstanceAllocationStrategy::OnDemand,
             // 512k of stack -- note that this is chosen currently to not be too
             // big, not be too small, and be a good default for most platforms.
@@ -1336,6 +1342,33 @@ impl Config {
         self
     }
 
+    /// Sets a custom executable-memory publisher.
+    ///
+    /// Custom executable-memory publishers are hooks that allow
+    /// Wasmtime to make certain regions of memory executable when
+    /// loading precompiled modules or compiling new modules
+    /// in-process. In most modern operating systems, memory allocated
+    /// for heap usage is readable and writable by default but not
+    /// executable. To jump to machine code stored in that memory, we
+    /// need to make it executable. For security reasons, we usually
+    /// also make it read-only at the same time, so the executing code
+    /// can't be modified later.
+    ///
+    /// By default, Wasmtime will use the appropriate system calls on
+    /// the host platform for this work. However, it also allows
+    /// plugging in a custom implementation via this configuration
+    /// option. This may be useful on custom or `no_std` platforms,
+    /// for example, especially where virtual memory is not otherwise
+    /// used by Wasmtime (no `signals-and-traps` feature).
+    #[cfg(feature = "runtime")]
+    pub fn with_custom_code_memory(
+        &mut self,
+        custom_code_memory: Option<Arc<dyn CustomCodeMemory>>,
+    ) -> &mut Self {
+        self.custom_code_memory = custom_code_memory;
+        self
+    }
+
     /// Sets the instance allocation strategy to use.
     ///
     /// This is notably used in conjunction with

crates/wasmtime/src/engine.rs

@@ -1,5 +1,7 @@
 use crate::prelude::*;
 #[cfg(feature = "runtime")]
+pub use crate::runtime::code_memory::CustomCodeMemory;
+#[cfg(feature = "runtime")]
 use crate::runtime::type_registry::TypeRegistry;
 #[cfg(feature = "runtime")]
 use crate::runtime::vm::GcRuntime;
@@ -655,6 +657,11 @@ impl Engine {
         &self.inner.signatures
     }
 
+    #[cfg(feature = "runtime")]
+    pub(crate) fn custom_code_memory(&self) -> Option<&Arc<dyn CustomCodeMemory>> {
+        self.config().custom_code_memory.as_ref()
+    }
+
     pub(crate) fn epoch_counter(&self) -> &AtomicU64 {
         &self.inner.epoch
     }
@@ -722,6 +729,15 @@ impl Engine {
         (f1(), f2())
     }
 
+    /// Returns the required alignment for a code image, if we
+    /// allocate in a way that is not a system `mmap()` that naturally
+    /// aligns it.
+    fn required_code_alignment(&self) -> usize {
+        self.custom_code_memory()
+            .map(|c| c.required_alignment())
+            .unwrap_or(1)
+    }
+
     /// Loads a `CodeMemory` from the specified in-memory slice, copying it to a
     /// uniquely owned mmap.
     ///
@@ -732,7 +748,13 @@ impl Engine {
         bytes: &[u8],
         expected: ObjectKind,
     ) -> Result<Arc<crate::CodeMemory>> {
-        self.load_code(crate::runtime::vm::MmapVec::from_slice(bytes)?, expected)
+        self.load_code(
+            crate::runtime::vm::MmapVec::from_slice_with_alignment(
+                bytes,
+                self.required_code_alignment(),
+            )?,
+            expected,
+        )
     }
 
     /// Like `load_code_bytes`, but creates a mmap from a file on disk.
@@ -755,7 +777,7 @@ impl Engine {
         expected: ObjectKind,
     ) -> Result<Arc<crate::CodeMemory>> {
         serialization::check_compatible(self, &mmap, expected)?;
-        let mut code = crate::CodeMemory::new(mmap)?;
+        let mut code = crate::CodeMemory::new(self, mmap)?;
         code.publish()?;
         Ok(Arc::new(code))
     }