winch_codegen/codegen/
bounds.rs

1//! Exposes heap bounds checks functionality for WebAssembly.
2//! Bounds checks in WebAssembly are critical for safety, so extreme caution is
3//! recommended when working on this area of Winch.
4use super::env::HeapData;
5use crate::{
6    abi::vmctx,
7    codegen::{CodeGenContext, Emission},
8    isa::reg::{Reg, writable},
9    masm::{IntCmpKind, IntScratch, MacroAssembler, OperandSize, RegImm, TrapCode},
10    stack::TypedReg,
11};
12use anyhow::Result;
13use wasmtime_environ::Signed;
14
15/// A newtype to represent an immediate offset argument for a heap access.
16#[derive(Debug, Copy, Clone)]
17pub(crate) struct ImmOffset(u32);
18
19impl ImmOffset {
20    /// Construct an [ImmOffset] from a u32.
21    pub fn from_u32(raw: u32) -> Self {
22        Self(raw)
23    }
24
25    /// Return the underlying u32 value.
26    pub fn as_u32(&self) -> u32 {
27        self.0
28    }
29}
30
31/// An enum to represent the heap bounds.
32#[derive(Debug, Copy, Clone)]
33pub(crate) enum Bounds {
34    /// Static, known ahead-of-time.
35    Static(u64),
36    /// Dynamic. Loaded at runtime.
37    Dynamic(TypedReg),
38}
39
40impl Bounds {
41    /// Construct a [Bounds] from a [TypedReg].
42    pub fn from_typed_reg(tr: TypedReg) -> Self {
43        Self::Dynamic(tr)
44    }
45
46    /// Construct a [Bounds] from a u64.
47    pub fn from_u64(raw: u64) -> Self {
48        Self::Static(raw)
49    }
50
51    /// Return the underlying [TypedReg] value.
52    pub fn as_typed_reg(&self) -> TypedReg {
53        match self {
54            Self::Dynamic(tr) => *tr,
55            _ => panic!(),
56        }
57    }
58
59    /// Return the underlying u64 value.
60    pub fn as_u64(&self) -> u64 {
61        match self {
62            Self::Static(v) => *v,
63            _ => panic!(),
64        }
65    }
66}
67
68/// A newtype to represent a heap access index via a [TypedReg].
69#[derive(Debug, Copy, Clone)]
70pub(crate) struct Index(TypedReg);
71
72impl Index {
73    /// Construct an [Index] from a [TypedReg].
74    pub fn from_typed_reg(tr: TypedReg) -> Self {
75        Self(tr)
76    }
77
78    /// Return the underlying
79    pub fn as_typed_reg(&self) -> TypedReg {
80        self.0
81    }
82}
83
84/// Loads the bounds of the dynamic heap.
85pub(crate) fn load_dynamic_heap_bounds<M>(
86    context: &mut CodeGenContext<Emission>,
87    masm: &mut M,
88    heap: &HeapData,
89    ptr_size: OperandSize,
90) -> Result<Bounds>
91where
92    M: MacroAssembler,
93{
94    let dst = context.any_gpr(masm)?;
95    match heap.memory.static_heap_size() {
96        // Constant size, no need to perform a load.
97        Some(size) => masm.mov(writable!(dst), RegImm::i64(size.signed()), ptr_size)?,
98
99        None => {
100            masm.with_scratch::<IntScratch, _>(|masm, scratch| {
101                let base = if let Some(offset) = heap.import_from {
102                    let addr = masm.address_at_vmctx(offset)?;
103                    masm.load_ptr(addr, scratch.writable())?;
104                    scratch.inner()
105                } else {
106                    vmctx!(M)
107                };
108                let addr = masm.address_at_reg(base, heap.current_length_offset)?;
109                masm.load_ptr(addr, writable!(dst))
110            })?;
111        }
112    }
113
114    Ok(Bounds::from_typed_reg(TypedReg::new(
115        heap.index_type(),
116        dst,
117    )))
118}
119
120/// This function ensures the following:
121/// * The immediate offset and memory access size fit in a single u64. Given:
122///   that the memory access size is a `u8`, we must guarantee that the immediate
123///   offset will fit in a `u32`, making the result of their addition fit in a u64
124///   and overflow safe.
125/// * Adjust the base index to account for the immediate offset via an unsigned
126///   addition and check for overflow in case the previous condition is not met.
127#[inline]
128pub(crate) fn ensure_index_and_offset<M: MacroAssembler>(
129    masm: &mut M,
130    index: Index,
131    offset: u64,
132    heap_ty_size: OperandSize,
133) -> Result<ImmOffset> {
134    match u32::try_from(offset) {
135        // If the immediate offset fits in a u32, then we simply return.
136        Ok(offs) => Ok(ImmOffset::from_u32(offs)),
137        // Else we adjust the index to be index = index + offset, including an
138        // overflow check, and return 0 as the offset.
139        Err(_) => {
140            masm.checked_uadd(
141                writable!(index.as_typed_reg().into()),
142                index.as_typed_reg().into(),
143                RegImm::i64(offset as i64),
144                heap_ty_size,
145                TrapCode::HEAP_OUT_OF_BOUNDS,
146            )?;
147
148            Ok(ImmOffset::from_u32(0))
149        }
150    }
151}
152
153/// Performs the out-of-bounds check and returns the heap address if the access
154/// criteria is in bounds.
155pub(crate) fn load_heap_addr_checked<M, F>(
156    masm: &mut M,
157    context: &mut CodeGenContext<Emission>,
158    ptr_size: OperandSize,
159    heap: &HeapData,
160    enable_spectre_mitigation: bool,
161    bounds: Bounds,
162    index: Index,
163    offset: ImmOffset,
164    mut emit_check_condition: F,
165) -> Result<Reg>
166where
167    M: MacroAssembler,
168    F: FnMut(&mut M, Bounds, Index) -> Result<IntCmpKind>,
169{
170    let cmp_kind = emit_check_condition(masm, bounds, index)?;
171
172    masm.trapif(cmp_kind, TrapCode::HEAP_OUT_OF_BOUNDS)?;
173    let addr = context.any_gpr(masm)?;
174
175    load_heap_addr_unchecked(masm, heap, index, offset, addr, ptr_size)?;
176    if !enable_spectre_mitigation {
177        Ok(addr)
178    } else {
179        // Conditionally assign 0 to the register holding the base address if
180        // the comparison kind is met.
181        let tmp = context.any_gpr(masm)?;
182        masm.mov(writable!(tmp), RegImm::i64(0), ptr_size)?;
183        let cmp_kind = emit_check_condition(masm, bounds, index)?;
184        masm.cmov(writable!(addr), tmp, cmp_kind, ptr_size)?;
185        context.free_reg(tmp);
186        Ok(addr)
187    }
188}
189
190/// Load the requested heap address into the specified destination register.
191/// This function doesn't perform any bounds checks and assumes the caller
192/// performed the right checks.
193pub(crate) fn load_heap_addr_unchecked<M>(
194    masm: &mut M,
195    heap: &HeapData,
196    index: Index,
197    offset: ImmOffset,
198    dst: Reg,
199    ptr_size: OperandSize,
200) -> Result<()>
201where
202    M: MacroAssembler,
203{
204    masm.with_scratch::<IntScratch, _>(|masm, scratch| {
205        let base = if let Some(offset) = heap.import_from {
206            // If the WebAssembly memory is imported, load the address into
207            // the scratch register.
208            masm.load_ptr(masm.address_at_vmctx(offset)?, scratch.writable())?;
209            scratch.inner()
210        } else {
211            // Else if the WebAssembly memory is defined in the current module,
212            // simply use the `VMContext` as the base for subsequent operations.
213            vmctx!(M)
214        };
215
216        // Load the base of the memory into the `addr` register.
217        masm.load_ptr(masm.address_at_reg(base, heap.offset)?, writable!(dst))
218    })?;
219
220    // Start by adding the index to the heap base addr.
221    let index_reg = index.as_typed_reg().reg;
222    masm.add(writable!(dst), dst, index_reg.into(), ptr_size)?;
223
224    if offset.as_u32() > 0 {
225        masm.add(
226            writable!(dst),
227            dst,
228            RegImm::i64(offset.as_u32() as i64),
229            ptr_size,
230        )?;
231    }
232    Ok(())
233}
winch_codegen/codegen/bounds.rs

winch_codegen/codegen/
bounds.rs
