diff --git a/crates/environ/src/compile/module_artifacts.rs b/crates/environ/src/compile/module_artifacts.rs index 7d9d132eddd6..9d992e25405c 100644 --- a/crates/environ/src/compile/module_artifacts.rs +++ b/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; + fn finish_object(obj: ObjectBuilder<'_>, state: &Self::State) -> Result; } impl FinishedObject for Vec { - fn finish_object(obj: ObjectBuilder<'_>) -> Result { + type State = (); + fn finish_object(obj: ObjectBuilder<'_>, _state: &Self::State) -> Result { let mut result = ObjectVec::default(); obj.finish(&mut result)?; return Ok(result.0); diff --git a/crates/wasmtime/src/compile.rs b/crates/wasmtime/src/compile.rs index 1222fb9607c4..1598bbf05e51 100644 --- a/crates/wasmtime/src/compile.rs +++ b/crates/wasmtime/src/compile.rs @@ -64,6 +64,7 @@ pub(crate) fn build_artifacts( 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( 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( engine: &Engine, binary: &[u8], _dwarf_package: Option<&[u8]>, + obj_state: &T::State, ) -> Result<(T, Option)> { use wasmtime_environ::component::{ CompiledComponentInfo, ComponentArtifacts, ComponentTypesBuilder, @@ -186,7 +188,7 @@ pub(crate) fn build_component_artifacts( }; object.serialize_info(&artifacts); - let result = T::finish_object(object)?; + let result = T::finish_object(object, obj_state)?; Ok((result, Some(artifacts))) } diff --git a/crates/wasmtime/src/compile/code_builder.rs b/crates/wasmtime/src/compile/code_builder.rs index ec8b543292c7..d6b33fb25f1f 100644 --- a/crates/wasmtime/src/compile/code_builder.rs +++ b/crates/wasmtime/src/compile/code_builder.rs @@ -274,7 +274,7 @@ impl<'a> CodeBuilder<'a> { pub fn compile_module_serialized(&self) -> Result> { 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> { 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) } } diff --git a/crates/wasmtime/src/compile/runtime.rs b/crates/wasmtime/src/compile/runtime.rs index 781be44829e8..23a40eaaf433 100644 --- a/crates/wasmtime/src/compile/runtime.rs +++ b/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( + fn compile_cached( &self, - build_artifacts: fn(&Engine, &[u8], Option<&[u8]>) -> Result<(MmapVecWrapper, Option)>, + build_artifacts: fn( + &Engine, + &[u8], + Option<&[u8]>, + &S, + ) -> Result<(MmapVecWrapper, Option)>, + state: &S, ) -> Result<(Arc, Option)> { 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,7 +93,9 @@ 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 { - 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) } @@ -87,22 +103,42 @@ impl<'a> CodeBuilder<'a> { /// [`Component`] instead of a module. #[cfg(feature = "component-model")] pub fn compile_component(&self) -> Result { - 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> { - let mut code = CodeMemory::new(mmap)?; +fn publish_mmap(engine: &Engine, mmap: MmapVec) -> Result> { + 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 { + type State = CustomAlignment; + fn finish_object(obj: ObjectBuilder<'_>, align: &CustomAlignment) -> Result { 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, len: usize, + alignment: usize, err: Option, } @@ -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); diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index 4ad821ac9591..4778b91cfd01 100644 --- a/crates/wasmtime/src/config.rs +++ b/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>, + #[cfg(feature = "runtime")] + pub(crate) custom_code_memory: Option>, 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>, + ) -> &mut Self { + self.custom_code_memory = custom_code_memory; + self + } + /// Sets the instance allocation strategy to use. /// /// This is notably used in conjunction with diff --git a/crates/wasmtime/src/engine.rs b/crates/wasmtime/src/engine.rs index e2ba2caa7b5c..7f8669c16bca 100644 --- a/crates/wasmtime/src/engine.rs +++ b/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> { + 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> { - 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> { 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)) } diff --git a/crates/wasmtime/src/runtime/code_memory.rs b/crates/wasmtime/src/runtime/code_memory.rs index 8ad04f789b71..390e851cac26 100644 --- a/crates/wasmtime/src/runtime/code_memory.rs +++ b/crates/wasmtime/src/runtime/code_memory.rs @@ -2,6 +2,8 @@ use crate::prelude::*; use crate::runtime::vm::{libcalls, MmapVec, UnwindRegistration}; +use crate::Engine; +use alloc::sync::Arc; use core::ops::Range; use object::endian::Endianness; use object::read::{elf::ElfFile64, Object, ObjectSection}; @@ -22,6 +24,7 @@ pub struct CodeMemory { needs_executable: bool, #[cfg(feature = "debug-builtins")] has_native_debug_info: bool, + custom_code_memory: Option>, relocations: Vec<(usize, obj::LibCall)>, @@ -38,6 +41,14 @@ pub struct CodeMemory { impl Drop for CodeMemory { fn drop(&mut self) { + // 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() { + let text = self.text(); + mem.unpublish_executable(text.as_ptr(), text.len()) + .expect("Executable memory unpublish failed"); + } + // Drop the registrations before `self.mmap` since they (implicitly) refer to it. let _ = self.unwind_registration.take(); #[cfg(feature = "debug-builtins")] @@ -50,13 +61,50 @@ fn _assert() { _assert_send_sync::(); } +/// Interface implemented by an embedder to provide custom +/// implementations of code-memory protection and execute permissions. +pub trait CustomCodeMemory: Send + Sync { + /// The minimal alignment granularity for an address region that + /// can be made executable. + /// + /// Wasmtime does not assume the system page size for this because + /// custom code-memory protection can be used when all other uses + /// of virtual memory are disabled. + fn required_alignment(&self) -> usize; + + /// Publish a region of memory as executable. + /// + /// This should update permissions from the default RW + /// (readable/writable but not executable) to RX + /// (readable/executable but not writable), enforcing W^X + /// discipline. + /// + /// 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. + /// + /// 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 + /// 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<()>; +} + impl CodeMemory { /// Creates a new `CodeMemory` by taking ownership of the provided /// `MmapVec`. /// /// The returned `CodeMemory` manages the internal `MmapVec` and the /// `publish` method is used to actually make the memory executable. - pub fn new(mmap: MmapVec) -> Result { + pub fn new(engine: &Engine, mmap: MmapVec) -> Result { let obj = ElfFile64::::parse(&mmap[..]) .map_err(obj::ObjectCrateErrorWrapper) .with_context(|| "failed to parse internal compilation artifact")?; @@ -140,6 +188,7 @@ impl CodeMemory { _ => log::debug!("ignoring section {name}"), } } + Ok(Self { mmap, unwind_registration: None, @@ -151,6 +200,7 @@ impl CodeMemory { needs_executable, #[cfg(feature = "debug-builtins")] has_native_debug_info, + custom_code_memory: engine.custom_code_memory().cloned(), text, unwind, trap_data, @@ -270,28 +320,30 @@ impl CodeMemory { // Switch the executable portion from readonly to read/execute. if self.needs_executable { - #[cfg(feature = "signals-based-traps")] - { - let text = self.text(); - - use wasmtime_jit_icache_coherence as icache_coherence; - - // Clear the newly allocated code from cache if the processor requires it - // - // Do this before marking the memory as R+X, technically we should be able to do it after - // but there are some CPU's that have had errata about doing this with read only memory. - icache_coherence::clear_cache(text.as_ptr().cast(), text.len()) - .expect("Failed cache clear"); - - self.mmap - .make_executable(self.text.clone(), self.enable_branch_protection) - .context("unable to make memory executable")?; - - // Flush any in-flight instructions from the pipeline - icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush"); + if !self.custom_publish()? { + #[cfg(feature = "signals-based-traps")] + { + let text = self.text(); + + use wasmtime_jit_icache_coherence as icache_coherence; + + // Clear the newly allocated code from cache if the processor requires it + // + // Do this before marking the memory as R+X, technically we should be able to do it after + // but there are some CPU's that have had errata about doing this with read only memory. + icache_coherence::clear_cache(text.as_ptr().cast(), text.len()) + .expect("Failed cache clear"); + + self.mmap + .make_executable(self.text.clone(), self.enable_branch_protection) + .context("unable to make memory executable")?; + + // Flush any in-flight instructions from the pipeline + icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush"); + } + #[cfg(not(feature = "signals-based-traps"))] + bail!("this target requires virtual memory to be enabled"); } - #[cfg(not(feature = "signals-based-traps"))] - bail!("this target requires virtual memory to be enabled"); } // With all our memory set up use the platform-specific @@ -307,6 +359,29 @@ impl CodeMemory { Ok(()) } + fn custom_publish(&mut self) -> Result { + 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) + } else { + Ok(false) + } + } + unsafe fn apply_relocations(&mut self) -> Result<()> { if self.relocations.is_empty() { return Ok(()); diff --git a/crates/wasmtime/src/runtime/vm/mmap_vec.rs b/crates/wasmtime/src/runtime/vm/mmap_vec.rs index f283ef4265bd..33040642384e 100644 --- a/crates/wasmtime/src/runtime/vm/mmap_vec.rs +++ b/crates/wasmtime/src/runtime/vm/mmap_vec.rs @@ -1,13 +1,19 @@ use crate::prelude::*; +#[cfg(not(feature = "signals-based-traps"))] +use crate::runtime::vm::send_sync_ptr::SendSyncPtr; #[cfg(feature = "signals-based-traps")] use crate::runtime::vm::{mmap::UnalignedLength, Mmap}; +#[cfg(not(feature = "signals-based-traps"))] +use alloc::alloc::Layout; use alloc::sync::Arc; use core::ops::{Deref, Range}; +#[cfg(not(feature = "signals-based-traps"))] +use core::ptr::NonNull; #[cfg(feature = "std")] use std::fs::File; /// A type which prefers to store backing memory in an OS-backed memory mapping -/// but can fall back to `Vec` as well. +/// but can fall back to the regular memory allocator as well. /// /// This type is used to store code in Wasmtime and manage read-only and /// executable permissions of compiled images. This is created from either an @@ -20,13 +26,19 @@ use std::fs::File; /// are typically not, then the remaining bytes in the final page for /// mmap-backed instances are unused. /// -/// Note that when `signals-based-traps` is disabled then this type is backed -/// by a normal `Vec`. In such a scenario this type does not support -/// read-only or executable bits and the methods are not available. +/// Note that when `signals-based-traps` is disabled then this type is +/// backed by the regular memory allocator via `alloc` APIs. In such a +/// scenario this type does not support read-only or executable bits +/// and the methods are not available. However, the `CustomCodeMemory` +/// mechanism may be used by the embedder to set up and tear down +/// executable permissions on parts of this storage. pub enum MmapVec { #[doc(hidden)] #[cfg(not(feature = "signals-based-traps"))] - Vec(Vec), + Alloc { + base: SendSyncPtr, + layout: Layout, + }, #[doc(hidden)] #[cfg(feature = "signals-based-traps")] Mmap { @@ -52,20 +64,32 @@ impl MmapVec { } #[cfg(not(feature = "signals-based-traps"))] - fn new_vec(vec: Vec) -> MmapVec { - MmapVec::Vec(vec) + fn new_alloc(len: usize, alignment: usize) -> MmapVec { + let layout = Layout::from_size_align(len, alignment) + .expect("Invalid size or alignment for MmapVec allocation"); + let base = SendSyncPtr::new( + NonNull::new(unsafe { alloc::alloc::alloc_zeroed(layout.clone()) }) + .expect("Allocation of MmapVec storage failed"), + ); + MmapVec::Alloc { base, layout } } - /// Creates a new zero-initialized `MmapVec` with the given `size`. + /// Creates a new zero-initialized `MmapVec` with the given `size` + /// and `alignment`. /// /// This commit will return a new `MmapVec` suitably sized to hold `size` /// bytes. All bytes will be initialized to zero since this is a fresh OS /// page allocation. - pub fn with_capacity(size: usize) -> Result { + pub fn with_capacity_and_alignment(size: usize, alignment: usize) -> Result { #[cfg(feature = "signals-based-traps")] - return Ok(MmapVec::new_mmap(Mmap::with_at_least(size)?, size)); + { + assert!(alignment <= crate::runtime::vm::host_page_size()); + return Ok(MmapVec::new_mmap(Mmap::with_at_least(size)?, size)); + } #[cfg(not(feature = "signals-based-traps"))] - return Ok(MmapVec::new_vec(vec![0; size])); + { + return Ok(MmapVec::new_alloc(size, alignment)); + } } /// Creates a new `MmapVec` from the contents of an existing `slice`. @@ -74,7 +98,21 @@ impl MmapVec { /// `slice` is copied into the new mmap. It's recommended to avoid this /// method if possible to avoid the need to copy data around. pub fn from_slice(slice: &[u8]) -> Result { - let mut result = MmapVec::with_capacity(slice.len())?; + MmapVec::from_slice_with_alignment(slice, 1) + } + + /// Creates a new `MmapVec` from the contents of an existing + /// `slice`, with a minimum alignment. + /// + /// `align` must be a power of two. This is useful when page + /// alignment is required when the system otherwise does not use + /// virtual memory but has a custom code publish handler. + /// + /// A new `MmapVec` is allocated to hold the contents of `slice` and then + /// `slice` is copied into the new mmap. It's recommended to avoid this + /// method if possible to avoid the need to copy data around. pub + pub fn from_slice_with_alignment(slice: &[u8], align: usize) -> Result { + let mut result = MmapVec::with_capacity_and_alignment(slice.len(), align)?; // SAFETY: The mmap hasn't been made readonly yet so this should be // safe to call. unsafe { @@ -132,7 +170,7 @@ impl MmapVec { pub fn original_file(&self) -> Option<&Arc> { match self { #[cfg(not(feature = "signals-based-traps"))] - MmapVec::Vec(_) => None, + MmapVec::Alloc { .. } => None, #[cfg(feature = "signals-based-traps")] MmapVec::Mmap { mmap, .. } => mmap.original_file(), } @@ -155,7 +193,9 @@ impl MmapVec { pub unsafe fn as_mut_slice(&mut self) -> &mut [u8] { match self { #[cfg(not(feature = "signals-based-traps"))] - MmapVec::Vec(v) => v, + MmapVec::Alloc { base, layout } => { + core::slice::from_raw_parts_mut(base.as_mut(), layout.size()) + } #[cfg(feature = "signals-based-traps")] MmapVec::Mmap { mmap, len } => mmap.slice_mut(0..*len), } @@ -169,7 +209,9 @@ impl Deref for MmapVec { fn deref(&self) -> &[u8] { match self { #[cfg(not(feature = "signals-based-traps"))] - MmapVec::Vec(v) => v, + MmapVec::Alloc { base, layout } => unsafe { + core::slice::from_raw_parts(base.as_ptr(), layout.size()) + }, #[cfg(feature = "signals-based-traps")] MmapVec::Mmap { mmap, len } => { // SAFETY: all bytes for this mmap, which is owned by @@ -180,13 +222,28 @@ impl Deref for MmapVec { } } +impl Drop for MmapVec { + fn drop(&mut self) { + match self { + #[cfg(not(feature = "signals-based-traps"))] + MmapVec::Alloc { base, layout, .. } => unsafe { + alloc::alloc::dealloc(base.as_mut(), layout.clone()); + }, + #[cfg(feature = "signals-based-traps")] + MmapVec::Mmap { .. } => { + // Drop impl on the `mmap` takes care of this case. + } + } + } +} + #[cfg(test)] mod tests { use super::MmapVec; #[test] fn smoke() { - let mut mmap = MmapVec::with_capacity(10).unwrap(); + let mut mmap = MmapVec::with_capacity_and_alignment(10, 1).unwrap(); assert_eq!(mmap.len(), 10); assert_eq!(&mmap[..], &[0; 10]); @@ -198,4 +255,11 @@ mod tests { assert_eq!(mmap[0], 1); assert_eq!(mmap[2], 3); } + + #[test] + fn alignment() { + let mmap = MmapVec::with_capacity_and_alignment(10, 4096).unwrap(); + let raw_ptr = &mmap[0] as *const _ as usize; + assert_eq!(raw_ptr & (4096 - 1), 0); + } } diff --git a/tests/all/custom_code_memory.rs b/tests/all/custom_code_memory.rs new file mode 100644 index 000000000000..61afacf61fa6 --- /dev/null +++ b/tests/all/custom_code_memory.rs @@ -0,0 +1,53 @@ +#[cfg(all(not(target_os = "windows"), not(miri)))] +mod not_for_windows { + use rustix::mm::{mprotect, MprotectFlags}; + use rustix::param::page_size; + use std::sync::Arc; + use wasmtime::*; + + struct CustomCodePublish; + impl CustomCodeMemory for CustomCodePublish { + fn required_alignment(&self) -> usize { + page_size() + } + + fn publish_executable(&self, ptr: *const u8, len: usize) -> anyhow::Result<()> { + unsafe { + mprotect( + ptr as *mut _, + len, + MprotectFlags::READ | MprotectFlags::EXEC, + )?; + } + Ok(()) + } + + fn unpublish_executable(&self, ptr: *const u8, len: usize) -> anyhow::Result<()> { + unsafe { + mprotect( + ptr as *mut _, + len, + MprotectFlags::READ | MprotectFlags::WRITE, + )?; + } + Ok(()) + } + } + + #[test] + fn custom_code_publish() { + let mut config = Config::default(); + config.with_custom_code_memory(Some(Arc::new(CustomCodePublish))); + let engine = Engine::new(&config).unwrap(); + let module = Module::new( + &engine, + "(module (func (export \"main\") (result i32) i32.const 42))", + ) + .unwrap(); + let mut store = Store::new(&engine, ()); + let instance = Instance::new(&mut store, &module, &[]).unwrap(); + let func: TypedFunc<(), i32> = instance.get_typed_func(&mut store, "main").unwrap(); + let result = func.call(&mut store, ()).unwrap(); + assert_eq!(result, 42); + } +} diff --git a/tests/all/main.rs b/tests/all/main.rs index 1dd7a2630e89..bee9eb17331f 100644 --- a/tests/all/main.rs +++ b/tests/all/main.rs @@ -8,6 +8,7 @@ mod cli_tests; mod code_too_large; mod component_model; mod coredump; +mod custom_code_memory; mod debug; mod defaults; mod epoch_interruption;