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::{scratch, vmctx},
7    codegen::{CodeGenContext, Emission},
8    isa::reg::{writable, Reg},
9    masm::{IntCmpKind, 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            let scratch = scratch!(M);
101            let base = if let Some(offset) = heap.import_from {
102                let addr = masm.address_at_vmctx(offset)?;
103                masm.load_ptr(addr, writable!(scratch))?;
104                scratch
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    Ok(Bounds::from_typed_reg(TypedReg::new(
114        heap.index_type(),
115        dst,
116    )))
117}
118
119/// This function ensures the following:
120/// * The immediate offset and memory access size fit in a single u64. Given:
121///   that the memory access size is a `u8`, we must guarantee that the immediate
122///   offset will fit in a `u32`, making the result of their addition fit in a u64
123///   and overflow safe.
124/// * Adjust the base index to account for the immediate offset via an unsigned
125///   addition and check for overflow in case the previous condition is not met.
126#[inline]
127pub(crate) fn ensure_index_and_offset<M: MacroAssembler>(
128    masm: &mut M,
129    index: Index,
130    offset: u64,
131    heap_ty_size: OperandSize,
132) -> Result<ImmOffset> {
133    match u32::try_from(offset) {
134        // If the immediate offset fits in a u32, then we simply return.
135        Ok(offs) => Ok(ImmOffset::from_u32(offs)),
136        // Else we adjust the index to be index = index + offset, including an
137        // overflow check, and return 0 as the offset.
138        Err(_) => {
139            masm.checked_uadd(
140                writable!(index.as_typed_reg().into()),
141                index.as_typed_reg().into(),
142                RegImm::i64(offset as i64),
143                heap_ty_size,
144                TrapCode::HEAP_OUT_OF_BOUNDS,
145            )?;
146
147            Ok(ImmOffset::from_u32(0))
148        }
149    }
150}
151
152/// Performs the out-of-bounds check and returns the heap address if the access
153/// criteria is in bounds.
154pub(crate) fn load_heap_addr_checked<M, F>(
155    masm: &mut M,
156    context: &mut CodeGenContext<Emission>,
157    ptr_size: OperandSize,
158    heap: &HeapData,
159    enable_spectre_mitigation: bool,
160    bounds: Bounds,
161    index: Index,
162    offset: ImmOffset,
163    mut emit_check_condition: F,
164) -> Result<Reg>
165where
166    M: MacroAssembler,
167    F: FnMut(&mut M, Bounds, Index) -> Result<IntCmpKind>,
168{
169    let cmp_kind = emit_check_condition(masm, bounds, index)?;
170
171    masm.trapif(cmp_kind, TrapCode::HEAP_OUT_OF_BOUNDS)?;
172    let addr = context.any_gpr(masm)?;
173
174    load_heap_addr_unchecked(masm, heap, index, offset, addr, ptr_size)?;
175    if !enable_spectre_mitigation {
176        Ok(addr)
177    } else {
178        // Conditionally assign 0 to the register holding the base address if
179        // the comparison kind is met.
180        let tmp = context.any_gpr(masm)?;
181        masm.mov(writable!(tmp), RegImm::i64(0), ptr_size)?;
182        let cmp_kind = emit_check_condition(masm, bounds, index)?;
183        masm.cmov(writable!(addr), tmp, cmp_kind, ptr_size)?;
184        context.free_reg(tmp);
185        Ok(addr)
186    }
187}
188
189/// Load the requested heap address into the specified destination register.
190/// This function doesn't perform any bounds checks and assumes the caller
191/// performed the right checks.
192pub(crate) fn load_heap_addr_unchecked<M>(
193    masm: &mut M,
194    heap: &HeapData,
195    index: Index,
196    offset: ImmOffset,
197    dst: Reg,
198    ptr_size: OperandSize,
199) -> Result<()>
200where
201    M: MacroAssembler,
202{
203    let base = if let Some(offset) = heap.import_from {
204        // If the WebAssembly memory is imported, load the address into
205        // the scratch register.
206        let scratch = scratch!(M);
207        masm.load_ptr(masm.address_at_vmctx(offset)?, writable!(scratch))?;
208        scratch
209    } else {
210        // Else if the WebAssembly memory is defined in the current module,
211        // simply use the `VMContext` as the base for subsequent operations.
212        vmctx!(M)
213    };
214
215    // Load the base of the memory into the `addr` register.
216    masm.load_ptr(masm.address_at_reg(base, heap.offset)?, writable!(dst))?;
217    // Start by adding the index to the heap base addr.
218    let index_reg = index.as_typed_reg().reg;
219    masm.add(writable!(dst), dst, index_reg.into(), ptr_size)?;
220
221    if offset.as_u32() > 0 {
222        masm.add(
223            writable!(dst),
224            dst,
225            RegImm::i64(offset.as_u32() as i64),
226            ptr_size,
227        )?;
228    }
229    Ok(())
230}
winch_codegen/codegen/bounds.rs

winch_codegen/codegen/
bounds.rs
