linera_wasmer_compiler/engine/
code_memory.rs

1// This file contains code from external sources.
2// Attributions: https://github.com/wasmerio/wasmer/blob/main/docs/ATTRIBUTIONS.md
3
4//! Memory management for executable code.
5use super::unwind::UnwindRegistry;
6use crate::GlobalFrameInfoRegistration;
7use wasmer_types::{
8    compilation::unwind::CompiledFunctionUnwindInfoLike, CompiledFunctionUnwindInfoReference,
9    CustomSectionLike, FunctionBodyLike,
10};
11use wasmer_vm::{Mmap, VMFunctionBody};
12
13/// The optimal alignment for functions.
14///
15/// On x86-64, this is 16 since it's what the optimizations assume.
16/// When we add support for other architectures, we should also figure out their
17/// optimal alignment values.
18const ARCH_FUNCTION_ALIGNMENT: usize = 16;
19
20/// The optimal alignment for data.
21///
22const DATA_SECTION_ALIGNMENT: usize = 64;
23
24/// Memory manager for executable code.
25pub struct CodeMemory {
26    // frame info is placed first, to ensure it's dropped before the mmap
27    frame_info_registration: Option<GlobalFrameInfoRegistration>,
28    unwind_registry: UnwindRegistry,
29    mmap: Mmap,
30    start_of_nonexecutable_pages: usize,
31}
32
33impl CodeMemory {
34    /// Create a new `CodeMemory` instance.
35    pub fn new() -> Self {
36        Self {
37            unwind_registry: UnwindRegistry::new(),
38            mmap: Mmap::new(),
39            start_of_nonexecutable_pages: 0,
40            frame_info_registration: None,
41        }
42    }
43
44    /// Mutably get the UnwindRegistry.
45    pub fn unwind_registry_mut(&mut self) -> &mut UnwindRegistry {
46        &mut self.unwind_registry
47    }
48
49    /// Allocate a single contiguous block of memory at a fixed virtual address for the functions and custom sections, and copy the data in place.
50    #[allow(clippy::type_complexity)]
51    pub fn allocate<'module, 'memory, FunctionBody, CustomSection>(
52        &'memory mut self,
53        functions: &'memory [&'module FunctionBody],
54        executable_sections: &'memory [&'module CustomSection],
55        data_sections: &'memory [&'module CustomSection],
56    ) -> Result<
57        (
58            Vec<&'memory mut [VMFunctionBody]>,
59            Vec<&'memory mut [u8]>,
60            Vec<&'memory mut [u8]>,
61        ),
62        String,
63    >
64    where
65        FunctionBody: FunctionBodyLike<'module> + 'module,
66        CustomSection: CustomSectionLike<'module> + 'module,
67    {
68        let mut function_result = vec![];
69        let mut data_section_result = vec![];
70        let mut executable_section_result = vec![];
71
72        let page_size = region::page::size();
73
74        // 1. Calculate the total size, that is:
75        // - function body size, including all trampolines
76        // -- windows unwind info
77        // -- padding between functions
78        // - executable section body
79        // -- padding between executable sections
80        // - padding until a new page to change page permissions
81        // - data section body size
82        // -- padding between data sections
83
84        let total_len = round_up(
85            functions.iter().fold(0, |acc, func| {
86                round_up(
87                    acc + Self::function_allocation_size(*func),
88                    ARCH_FUNCTION_ALIGNMENT,
89                )
90            }) + executable_sections.iter().fold(0, |acc, exec| {
91                round_up(acc + exec.bytes().len(), ARCH_FUNCTION_ALIGNMENT)
92            }),
93            page_size,
94        ) + data_sections.iter().fold(0, |acc, data| {
95            round_up(acc + data.bytes().len(), DATA_SECTION_ALIGNMENT)
96        });
97
98        // 2. Allocate the pages. Mark them all read-write.
99
100        self.mmap = Mmap::with_at_least(total_len)?;
101
102        // 3. Determine where the pointers to each function, executable section
103        // or data section are. Copy the functions. Collect the addresses of each and return them.
104
105        let mut bytes = 0;
106        let mut buf = self.mmap.as_mut_slice();
107        for func in functions {
108            let len = round_up(
109                Self::function_allocation_size(*func),
110                ARCH_FUNCTION_ALIGNMENT,
111            );
112            let (func_buf, next_buf) = buf.split_at_mut(len);
113            buf = next_buf;
114            bytes += len;
115
116            let vmfunc = Self::copy_function(&mut self.unwind_registry, *func, func_buf);
117            assert_eq!(vmfunc.as_ptr() as usize % ARCH_FUNCTION_ALIGNMENT, 0);
118            function_result.push(vmfunc);
119        }
120        for section in executable_sections {
121            let section = section.bytes();
122            assert_eq!(buf.as_mut_ptr() as usize % ARCH_FUNCTION_ALIGNMENT, 0);
123            let len = round_up(section.len(), ARCH_FUNCTION_ALIGNMENT);
124            let (s, next_buf) = buf.split_at_mut(len);
125            buf = next_buf;
126            bytes += len;
127            s[..section.len()].copy_from_slice(section);
128            executable_section_result.push(s);
129        }
130
131        self.start_of_nonexecutable_pages = bytes;
132
133        if !data_sections.is_empty() {
134            // Data sections have different page permissions from the executable
135            // code that came before it, so they need to be on different pages.
136            let padding = round_up(bytes, page_size) - bytes;
137            buf = buf.split_at_mut(padding).1;
138
139            for section in data_sections {
140                let section = section.bytes();
141                assert_eq!(buf.as_mut_ptr() as usize % DATA_SECTION_ALIGNMENT, 0);
142                let len = round_up(section.len(), DATA_SECTION_ALIGNMENT);
143                let (s, next_buf) = buf.split_at_mut(len);
144                buf = next_buf;
145                s[..section.len()].copy_from_slice(section);
146                data_section_result.push(s);
147            }
148        }
149
150        Ok((
151            function_result,
152            executable_section_result,
153            data_section_result,
154        ))
155    }
156
157    /// Apply the page permissions.
158    pub fn publish(&mut self) {
159        if self.mmap.is_empty() || self.start_of_nonexecutable_pages == 0 {
160            return;
161        }
162        assert!(self.mmap.len() >= self.start_of_nonexecutable_pages);
163        unsafe {
164            region::protect(
165                self.mmap.as_mut_ptr(),
166                self.start_of_nonexecutable_pages,
167                region::Protection::READ_EXECUTE,
168            )
169        }
170        .expect("unable to make memory readonly and executable");
171    }
172
173    /// Calculates the allocation size of the given compiled function.
174    fn function_allocation_size<'a>(func: &'a impl FunctionBodyLike<'a>) -> usize {
175        match &func.unwind_info().map(|o| o.get()) {
176            Some(CompiledFunctionUnwindInfoReference::WindowsX64(info)) => {
177                // Windows unwind information is required to be emitted into code memory
178                // This is because it must be a positive relative offset from the start of the memory
179                // Account for necessary unwind information alignment padding (32-bit alignment)
180                ((func.body().len() + 3) & !3) + info.len()
181            }
182            _ => func.body().len(),
183        }
184    }
185
186    /// Copies the data of the compiled function to the given buffer.
187    ///
188    /// This will also add the function to the current function table.
189    fn copy_function<'module, 'memory>(
190        registry: &mut UnwindRegistry,
191        func: &'module impl FunctionBodyLike<'module>,
192        buf: &'memory mut [u8],
193    ) -> &'memory mut [VMFunctionBody] {
194        assert_eq!(buf.as_ptr() as usize % ARCH_FUNCTION_ALIGNMENT, 0);
195
196        let func_len = func.body().len();
197
198        let (body, remainder) = buf.split_at_mut(func_len);
199        body.copy_from_slice(func.body());
200        let vmfunc = Self::view_as_mut_vmfunc_slice(body);
201
202        let unwind_info = func.unwind_info().map(|o| o.get());
203        if let Some(CompiledFunctionUnwindInfoReference::WindowsX64(info)) = unwind_info {
204            // Windows unwind information is written following the function body
205            // Keep unwind information 32-bit aligned (round up to the nearest 4 byte boundary)
206            let unwind_start = (func_len + 3) & !3;
207            let unwind_size = info.len();
208            let padding = unwind_start - func_len;
209            assert_eq!((func_len + padding) % 4, 0);
210            let slice = remainder.split_at_mut(padding + unwind_size).0;
211            slice[padding..].copy_from_slice(info);
212        }
213
214        if let Some(ref info) = unwind_info {
215            registry
216                .register(vmfunc.as_ptr() as usize, 0, func_len as u32, info)
217                .expect("failed to register unwind information");
218        }
219
220        vmfunc
221    }
222
223    /// Convert mut a slice from u8 to VMFunctionBody.
224    fn view_as_mut_vmfunc_slice(slice: &mut [u8]) -> &mut [VMFunctionBody] {
225        let byte_ptr: *mut [u8] = slice;
226        let body_ptr = byte_ptr as *mut [VMFunctionBody];
227        unsafe { &mut *body_ptr }
228    }
229
230    /// Register the frame info, so it's free when the mememory gets freed
231    pub fn register_frame_info(&mut self, frame_info: GlobalFrameInfoRegistration) {
232        self.frame_info_registration = Some(frame_info);
233    }
234}
235
236fn round_up(size: usize, multiple: usize) -> usize {
237    debug_assert!(multiple.is_power_of_two());
238    (size + (multiple - 1)) & !(multiple - 1)
239}
240
241#[cfg(test)]
242mod tests {
243    use super::CodeMemory;
244    fn _assert() {
245        fn _assert_send_sync<T: Send + Sync>() {}
246        _assert_send_sync::<CodeMemory>();
247    }
248}