polkavm/
sandbox.rs

1use std::borrow::Cow;
2use std::sync::{Arc, Mutex};
3use core::sync::atomic::{AtomicUsize, Ordering};
4
5use polkavm_common::{
6    error::ExecutionError,
7    zygote::{
8        AddressTable,
9        SandboxMemoryConfig,
10    },
11    utils::{Access, Gas, align_to_next_page_usize}
12};
13
14use crate::api::{BackendAccess, EngineState, ExecuteArgs, Module};
15use crate::compiler::CompiledModule;
16use crate::config::{GasMeteringKind, SandboxKind};
17use crate::utils::GuestInit;
18use crate::error::Error;
19
20macro_rules! get_field_offset {
21    ($struct:expr, |$struct_ident:ident| $get_field:expr) => {{
22        let $struct_ident = $struct;
23        let struct_ref = &$struct_ident;
24        let field_ref = $get_field;
25        let struct_addr = struct_ref as *const _ as usize;
26        let field_addr = field_ref as *const _ as usize;
27        field_addr - struct_addr
28    }}
29}
30
31pub mod generic;
32
33#[cfg(target_os = "linux")]
34pub mod linux;
35
36// This is literally the only thing we need from `libc` on Linux, so instead of including
37// the whole crate let's just define these ourselves.
38#[cfg(target_os = "linux")]
39const _SC_PAGESIZE: core::ffi::c_int = 30;
40
41#[cfg(target_os = "linux")]
42extern "C" {
43    fn sysconf(name: core::ffi::c_int) -> core::ffi::c_long;
44}
45
46#[cfg(not(target_os = "linux"))]
47use libc::{sysconf, _SC_PAGESIZE};
48
49pub(crate) fn get_native_page_size() -> usize {
50    // TODO: Cache this?
51
52    // SAFETY: This function has no safety invariants and should be always safe to call.
53    unsafe { sysconf(_SC_PAGESIZE) as usize }
54}
55
56#[derive(Copy, Clone, PartialEq, Eq, Debug)]
57pub(crate) struct OutOfGas;
58
59pub trait SandboxConfig: Default {
60    fn enable_logger(&mut self, value: bool);
61}
62
63pub trait SandboxAddressSpace {
64    fn native_code_address(&self) -> u64;
65}
66
67pub trait SandboxProgram: Clone {
68    fn machine_code(&self) -> Cow<[u8]>;
69}
70
71pub(crate) trait Sandbox: Sized {
72    const KIND: SandboxKind;
73
74    type Access<'r>: Access<'r> + Into<BackendAccess<'r>> where Self: 'r;
75    type Config: SandboxConfig;
76    type Error: core::fmt::Debug + core::fmt::Display;
77    type Program: SandboxProgram;
78    type AddressSpace: SandboxAddressSpace;
79
80    fn as_sandbox_vec(sandbox_vec: &SandboxVec) -> &Mutex<Vec<Self>>;
81    fn as_compiled_module(module: &Module) -> &CompiledModule<Self>;
82
83    fn reserve_address_space() -> Result<Self::AddressSpace, Self::Error>;
84    fn prepare_program(init: SandboxInit, address_space: Self::AddressSpace) -> Result<Self::Program, Self::Error>;
85    fn spawn(config: &Self::Config) -> Result<Self, Self::Error>;
86    fn execute(&mut self, args: ExecuteArgs) -> Result<(), ExecutionError<Self::Error>>;
87    fn access(&'_ mut self) -> Self::Access<'_>;
88    fn pid(&self) -> Option<u32>;
89    fn address_table() -> AddressTable;
90    fn vmctx_regs_offset() -> usize;
91    fn vmctx_gas_offset() -> usize;
92    fn vmctx_heap_info_offset() -> usize;
93    fn gas_remaining_impl(&self) -> Result<Option<Gas>, OutOfGas>;
94    fn sync(&mut self) -> Result<(), Self::Error>;
95}
96
97#[derive(Copy, Clone, Default)]
98pub struct SandboxInit<'a> {
99    pub guest_init: GuestInit<'a>,
100    pub code: &'a [u8],
101    pub jump_table: &'a [u8],
102    pub sysreturn_address: u64,
103}
104
105impl<'a> SandboxInit<'a> {
106    fn memory_config(&self, native_page_size: usize) -> Result<SandboxMemoryConfig, &'static str> {
107        let memory_map = self.guest_init.memory_map()?;
108        let mut ro_data_fd_size = align_to_next_page_usize(native_page_size, self.guest_init.ro_data.len()).unwrap() as u32;
109        if memory_map.ro_data_size() - ro_data_fd_size < memory_map.page_size() {
110            ro_data_fd_size = memory_map.ro_data_size();
111        }
112
113        let rw_data_fd_size = align_to_next_page_usize(native_page_size, self.guest_init.rw_data.len()).unwrap() as u32;
114        let code_size = align_to_next_page_usize(native_page_size, self.code.len()).unwrap() as u32;
115        let jump_table_size = align_to_next_page_usize(native_page_size, self.jump_table.len()).unwrap() as u32;
116
117        Ok(SandboxMemoryConfig {
118            memory_map,
119            ro_data_fd_size,
120            rw_data_fd_size,
121            code_size,
122            jump_table_size,
123            sysreturn_address: self.sysreturn_address,
124        })
125    }
126}
127
128pub(crate) fn get_gas(args: &ExecuteArgs, gas_metering: Option<GasMeteringKind>) -> Option<i64> {
129    if args.module.is_none() && args.gas.is_none() && gas_metering.is_some() {
130        // Keep whatever value was set there previously.
131        return None;
132    }
133
134    let gas = args.gas.unwrap_or(Gas::MIN);
135    if gas_metering.is_some() {
136        Some(gas.get() as i64)
137    } else {
138        Some(0)
139    }
140}
141
142pub(crate) struct SandboxInstance<S> where S: Sandbox {
143    engine_state: Arc<EngineState>,
144    sandbox: Option<S>
145}
146
147impl<S> SandboxInstance<S> where S: Sandbox {
148    pub fn spawn_and_load_module(engine_state: Arc<EngineState>, module: &Module) -> Result<Self, Error> {
149        let mut sandbox = SandboxInstance {
150            sandbox: Some(reuse_or_spawn_sandbox::<S>(&engine_state, module)?),
151            engine_state,
152        };
153
154        let mut args = ExecuteArgs::new();
155        args.module = Some(module);
156
157        sandbox
158            .execute(args)
159            .map_err(Error::from_display)
160            .map_err(|error| error.context("instantiation failed: failed to upload the program into the sandbox"))?;
161
162        Ok(sandbox)
163    }
164
165    pub fn execute(&mut self, args: ExecuteArgs) -> Result<(), ExecutionError<Error>> {
166        let sandbox = self.sandbox.as_mut().unwrap();
167        let result = match sandbox.execute(args) {
168            Ok(()) => Ok(()),
169            Err(ExecutionError::Trap(trap)) => Err(ExecutionError::Trap(trap)),
170            Err(ExecutionError::Error(error)) => return Err(ExecutionError::Error(Error::from_display(error))),
171            Err(ExecutionError::OutOfGas) => return Err(ExecutionError::OutOfGas),
172        };
173
174        if sandbox.gas_remaining_impl().is_err() {
175            return Err(ExecutionError::OutOfGas);
176        }
177
178        result
179    }
180
181    pub fn access(&'_ mut self) -> S::Access<'_> {
182        self.sandbox.as_mut().unwrap().access()
183    }
184
185    pub fn sandbox(&self) -> &S {
186        self.sandbox.as_ref().unwrap()
187    }
188}
189
190impl<S> Drop for SandboxInstance<S> where S: Sandbox {
191    fn drop(&mut self) {
192        recycle_sandbox::<S>(&self.engine_state, || {
193            let mut sandbox = self.sandbox.take()?;
194            let mut args = ExecuteArgs::new();
195            args.flags |= polkavm_common::VM_RPC_FLAG_CLEAR_PROGRAM_AFTER_EXECUTION;
196            args.gas = Some(polkavm_common::utils::Gas::MIN);
197            args.is_async = true;
198
199            if let Err(error) = sandbox.execute(args) {
200                log::warn!("Failed to cache a sandbox worker process due to an error: {error}");
201                None
202            } else {
203                Some(sandbox)
204            }
205        })
206    }
207}
208
209pub(crate) enum SandboxVec {
210    #[cfg(target_os = "linux")]
211    Linux(Mutex<Vec<crate::sandbox::linux::Sandbox>>),
212    Generic(Mutex<Vec<crate::sandbox::generic::Sandbox>>),
213}
214
215pub(crate) struct SandboxCache {
216    sandboxes: SandboxVec,
217    available_workers: AtomicUsize,
218    worker_limit: usize,
219}
220
221impl SandboxCache {
222    pub(crate) fn new(kind: SandboxKind, worker_count: usize, debug_trace_execution: bool) -> Result<Self, Error> {
223        let sandboxes = match kind {
224            SandboxKind::Linux => {
225                #[cfg(target_os = "linux")]
226                {
227                    SandboxVec::Linux(Mutex::new(spawn_sandboxes(worker_count, debug_trace_execution)?))
228                }
229
230                #[cfg(not(target_os = "linux"))]
231                {
232                    unreachable!()
233                }
234            },
235            SandboxKind::Generic => SandboxVec::Generic(Mutex::new(spawn_sandboxes(worker_count, debug_trace_execution)?)),
236        };
237
238        Ok(SandboxCache {
239            sandboxes,
240            available_workers: AtomicUsize::new(worker_count),
241            worker_limit: worker_count,
242        })
243    }
244
245    fn reuse_sandbox<S>(&self) -> Option<S> where S: Sandbox {
246        if self.available_workers.load(Ordering::Relaxed) == 0 {
247            return None;
248        }
249
250        let sandboxes = S::as_sandbox_vec(&self.sandboxes);
251        let mut sandboxes = match sandboxes.lock() {
252            Ok(sandboxes) => sandboxes,
253            Err(poison) => poison.into_inner(),
254        };
255
256        let mut sandbox = sandboxes.pop()?;
257        self.available_workers.fetch_sub(1, Ordering::Relaxed);
258
259        if let Err(error) = sandbox.sync() {
260            log::warn!("Failed to reuse a sandbox: {error}");
261            None
262        } else {
263            Some(sandbox)
264        }
265    }
266}
267
268fn is_sandbox_logging_enabled() -> bool {
269    cfg!(test) || log::log_enabled!(target: "polkavm", log::Level::Trace) || log::log_enabled!(target: "polkavm::zygote", log::Level::Trace)
270}
271
272fn spawn_sandboxes<S>(count: usize, debug_trace_execution: bool) -> Result<Vec<S>, Error> where S: Sandbox {
273    use crate::sandbox::SandboxConfig;
274
275    let mut sandbox_config = S::Config::default();
276    sandbox_config.enable_logger(is_sandbox_logging_enabled() || debug_trace_execution);
277
278    let mut sandboxes = Vec::with_capacity(count);
279    for nth in 0..count {
280        let sandbox = S::spawn(&sandbox_config)
281            .map_err(crate::Error::from_display)
282            .map_err(|error| error.context(format!("failed to create a worker process ({} out of {})", nth + 1, count)))?;
283
284        sandboxes.push(sandbox);
285    }
286
287    Ok(sandboxes)
288}
289
290fn reuse_or_spawn_sandbox<S>(engine_state: &EngineState, module: &Module) -> Result<S, Error> where S: Sandbox {
291    use crate::sandbox::SandboxConfig;
292
293    let mut sandbox_config = S::Config::default();
294    sandbox_config.enable_logger(is_sandbox_logging_enabled() || module.is_debug_trace_execution_enabled());
295
296    if let Some(sandbox) = engine_state.sandbox_cache().and_then(|cache| cache.reuse_sandbox::<S>()) {
297        Ok(sandbox)
298    } else {
299        S::spawn(&sandbox_config)
300            .map_err(Error::from_display)
301            .map_err(|error| error.context("instantiation failed: failed to create a sandbox"))
302    }
303}
304
305fn recycle_sandbox<S>(engine_state: &EngineState, get_sandbox: impl FnOnce() -> Option<S>) where S: Sandbox {
306    let Some(sandbox_cache) = engine_state.sandbox_cache() else { return };
307    let sandboxes = S::as_sandbox_vec(&sandbox_cache.sandboxes);
308
309    let mut count = sandbox_cache.available_workers.load(Ordering::Relaxed);
310    if count >= sandbox_cache.worker_limit {
311        return;
312    }
313
314    loop {
315        if let Err(new_count) = sandbox_cache.available_workers.compare_exchange(count, count + 1, Ordering::Relaxed, Ordering::Relaxed) {
316            if new_count >= sandbox_cache.worker_limit {
317                return;
318            }
319
320            count = new_count;
321            continue;
322        }
323
324        break;
325    }
326
327    if let Some(sandbox) = get_sandbox() {
328        let mut sandboxes = match sandboxes.lock() {
329            Ok(sandboxes) => sandboxes,
330            Err(poison) => poison.into_inner(),
331        };
332
333        sandboxes.push(sandbox);
334    } else {
335        sandbox_cache.available_workers.fetch_sub(1, Ordering::Relaxed);
336    }
337}
polkavm/sandbox.rs

polkavm/
sandbox.rs
