wasmtime_internal_wmemcheck/
lib.rs

1//! > **⚠️ Warning ⚠️**: this crate is an internal-only crate for the Wasmtime
2//! > project and is not intended for general use. APIs are not strictly
3//! > reviewed for safety and usage outside of Wasmtime may have bugs. If
4//! > you're interested in using this feel free to file an issue on the
5//! > Wasmtime repository to start a discussion about doing so, but otherwise
6//! > be aware that your usage of this crate is not supported.
7
8use std::cmp::*;
9use std::collections::HashMap;
10
11/// Memory checker for wasm guest.
12pub struct Wmemcheck {
13    metadata: Vec<MemState>,
14    mallocs: HashMap<usize, usize>,
15    pub stack_pointer: usize,
16    max_stack_size: usize,
17    pub flag: bool,
18}
19
20/// Error types for memory checker.
21#[derive(Debug, PartialEq)]
22pub enum AccessError {
23    /// Malloc over already malloc'd memory.
24    DoubleMalloc { addr: usize, len: usize },
25    /// Read from uninitialized or undefined memory.
26    InvalidRead { addr: usize, len: usize },
27    /// Write to uninitialized memory.
28    InvalidWrite { addr: usize, len: usize },
29    /// Free of non-malloc'd pointer.
30    InvalidFree { addr: usize },
31    /// Access out of bounds of heap or stack.
32    OutOfBounds { addr: usize, len: usize },
33}
34
35/// Memory state for memory checker.
36#[derive(Debug, Clone, PartialEq)]
37pub enum MemState {
38    /// Unallocated memory.
39    Unallocated,
40    /// Initialized but undefined memory.
41    ValidToWrite,
42    /// Initialized and defined memory.
43    ValidToReadWrite,
44}
45
46impl Wmemcheck {
47    /// Initializes memory checker instance.
48    pub fn new(mem_size: usize) -> Wmemcheck {
49        let metadata = vec![MemState::Unallocated; mem_size];
50        let mallocs = HashMap::new();
51        Wmemcheck {
52            metadata,
53            mallocs,
54            stack_pointer: 0,
55            max_stack_size: 0,
56            flag: true,
57        }
58    }
59
60    /// Updates memory checker memory state metadata when malloc is called.
61    pub fn malloc(&mut self, addr: usize, len: usize) -> Result<(), AccessError> {
62        if !self.is_in_bounds_heap(addr, len) {
63            return Err(AccessError::OutOfBounds { addr, len });
64        }
65        for i in addr..addr + len {
66            match self.metadata[i] {
67                MemState::ValidToWrite => {
68                    return Err(AccessError::DoubleMalloc { addr, len });
69                }
70                MemState::ValidToReadWrite => {
71                    return Err(AccessError::DoubleMalloc { addr, len });
72                }
73                _ => {}
74            }
75        }
76        for i in addr..addr + len {
77            self.metadata[i] = MemState::ValidToWrite;
78        }
79        self.mallocs.insert(addr, len);
80        Ok(())
81    }
82
83    /// Updates memory checker memory state metadata when a load occurs.
84    pub fn read(&mut self, addr: usize, len: usize) -> Result<(), AccessError> {
85        if !self.flag {
86            return Ok(());
87        }
88        if !(self.is_in_bounds_stack(addr, len) || self.is_in_bounds_heap(addr, len)) {
89            return Err(AccessError::OutOfBounds { addr, len });
90        }
91        for i in addr..addr + len {
92            match self.metadata[i] {
93                MemState::Unallocated => {
94                    return Err(AccessError::InvalidRead { addr, len });
95                }
96                MemState::ValidToWrite => {
97                    return Err(AccessError::InvalidRead { addr, len });
98                }
99                _ => {}
100            }
101        }
102        Ok(())
103    }
104
105    /// Updates memory checker memory state metadata when a store occurs.
106    pub fn write(&mut self, addr: usize, len: usize) -> Result<(), AccessError> {
107        if !self.flag {
108            return Ok(());
109        }
110        if !(self.is_in_bounds_stack(addr, len) || self.is_in_bounds_heap(addr, len)) {
111            return Err(AccessError::OutOfBounds { addr, len });
112        }
113        for i in addr..addr + len {
114            if let MemState::Unallocated = self.metadata[i] {
115                return Err(AccessError::InvalidWrite { addr, len });
116            }
117        }
118        for i in addr..addr + len {
119            self.metadata[i] = MemState::ValidToReadWrite;
120        }
121        Ok(())
122    }
123
124    /// Updates memory checker memory state metadata when free is called.
125    pub fn free(&mut self, addr: usize) -> Result<(), AccessError> {
126        if !self.mallocs.contains_key(&addr) {
127            return Err(AccessError::InvalidFree { addr });
128        }
129        let len = self.mallocs[&addr];
130        for i in addr..addr + len {
131            if let MemState::Unallocated = self.metadata[i] {
132                return Err(AccessError::InvalidFree { addr });
133            }
134        }
135        self.mallocs.remove(&addr);
136        for i in addr..addr + len {
137            self.metadata[i] = MemState::Unallocated;
138        }
139        Ok(())
140    }
141
142    fn is_in_bounds_heap(&self, addr: usize, len: usize) -> bool {
143        self.max_stack_size <= addr && addr + len <= self.metadata.len()
144    }
145
146    fn is_in_bounds_stack(&self, addr: usize, len: usize) -> bool {
147        self.stack_pointer <= addr && addr + len < self.max_stack_size
148    }
149
150    /// Updates memory checker metadata when stack pointer is updated.
151    pub fn update_stack_pointer(&mut self, new_sp: usize) -> Result<(), AccessError> {
152        if new_sp > self.max_stack_size {
153            return Err(AccessError::OutOfBounds {
154                addr: self.stack_pointer,
155                len: new_sp - self.stack_pointer,
156            });
157        } else if new_sp < self.stack_pointer {
158            for i in new_sp..self.stack_pointer + 1 {
159                self.metadata[i] = MemState::ValidToReadWrite;
160            }
161        } else {
162            for i in self.stack_pointer..new_sp {
163                self.metadata[i] = MemState::Unallocated;
164            }
165        }
166        self.stack_pointer = new_sp;
167        Ok(())
168    }
169
170    /// Turns memory checking on.
171    pub fn memcheck_on(&mut self) {
172        self.flag = true;
173    }
174
175    /// Turns memory checking off.
176    pub fn memcheck_off(&mut self) {
177        self.flag = false;
178    }
179
180    /// Initializes stack and stack pointer in memory checker metadata.
181    pub fn set_stack_size(&mut self, stack_size: usize) {
182        self.max_stack_size = stack_size + 1;
183        // TODO: temporary solution to initialize the entire stack
184        // while keeping stack tracing plumbing in place
185        self.stack_pointer = stack_size;
186        let _ = self.update_stack_pointer(0);
187    }
188
189    /// Updates memory checker metadata size when memory.grow is called.
190    pub fn update_mem_size(&mut self, num_bytes: usize) {
191        let to_append = vec![MemState::Unallocated; num_bytes];
192        self.metadata.extend(to_append);
193    }
194}
195
196#[test]
197fn basic_wmemcheck() {
198    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
199
200    wmemcheck_state.set_stack_size(1024);
201    assert!(wmemcheck_state.malloc(0x1000, 32).is_ok());
202    assert!(wmemcheck_state.write(0x1000, 4).is_ok());
203    assert!(wmemcheck_state.read(0x1000, 4).is_ok());
204    assert_eq!(wmemcheck_state.mallocs, HashMap::from([(0x1000, 32)]));
205    assert!(wmemcheck_state.free(0x1000).is_ok());
206    assert!(wmemcheck_state.mallocs.is_empty());
207}
208
209#[test]
210fn read_before_initializing() {
211    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
212
213    assert!(wmemcheck_state.malloc(0x1000, 32).is_ok());
214    assert_eq!(
215        wmemcheck_state.read(0x1000, 4),
216        Err(AccessError::InvalidRead {
217            addr: 0x1000,
218            len: 4
219        })
220    );
221    assert!(wmemcheck_state.write(0x1000, 4).is_ok());
222    assert!(wmemcheck_state.free(0x1000).is_ok());
223}
224
225#[test]
226fn use_after_free() {
227    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
228
229    assert!(wmemcheck_state.malloc(0x1000, 32).is_ok());
230    assert!(wmemcheck_state.write(0x1000, 4).is_ok());
231    assert!(wmemcheck_state.write(0x1000, 4).is_ok());
232    assert!(wmemcheck_state.free(0x1000).is_ok());
233    assert_eq!(
234        wmemcheck_state.write(0x1000, 4),
235        Err(AccessError::InvalidWrite {
236            addr: 0x1000,
237            len: 4
238        })
239    );
240}
241
242#[test]
243fn double_free() {
244    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
245
246    assert!(wmemcheck_state.malloc(0x1000, 32).is_ok());
247    assert!(wmemcheck_state.write(0x1000, 4).is_ok());
248    assert!(wmemcheck_state.free(0x1000).is_ok());
249    assert_eq!(
250        wmemcheck_state.free(0x1000),
251        Err(AccessError::InvalidFree { addr: 0x1000 })
252    );
253}
254
255#[test]
256fn out_of_bounds_malloc() {
257    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
258
259    assert_eq!(
260        wmemcheck_state.malloc(640 * 1024, 1),
261        Err(AccessError::OutOfBounds {
262            addr: 640 * 1024,
263            len: 1
264        })
265    );
266    assert_eq!(
267        wmemcheck_state.malloc(640 * 1024 - 10, 15),
268        Err(AccessError::OutOfBounds {
269            addr: 640 * 1024 - 10,
270            len: 15
271        })
272    );
273    assert!(wmemcheck_state.mallocs.is_empty());
274}
275
276#[test]
277fn out_of_bounds_read() {
278    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
279
280    assert!(wmemcheck_state.malloc(640 * 1024 - 24, 24).is_ok());
281    assert_eq!(
282        wmemcheck_state.read(640 * 1024 - 24, 25),
283        Err(AccessError::OutOfBounds {
284            addr: 640 * 1024 - 24,
285            len: 25
286        })
287    );
288}
289
290#[test]
291fn double_malloc() {
292    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
293
294    assert!(wmemcheck_state.malloc(0x1000, 32).is_ok());
295    assert_eq!(
296        wmemcheck_state.malloc(0x1000, 32),
297        Err(AccessError::DoubleMalloc {
298            addr: 0x1000,
299            len: 32
300        })
301    );
302    assert_eq!(
303        wmemcheck_state.malloc(0x1002, 32),
304        Err(AccessError::DoubleMalloc {
305            addr: 0x1002,
306            len: 32
307        })
308    );
309    assert!(wmemcheck_state.free(0x1000).is_ok());
310}
311
312#[test]
313fn error_type() {
314    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
315
316    assert!(wmemcheck_state.malloc(0x1000, 32).is_ok());
317    assert_eq!(
318        wmemcheck_state.malloc(0x1000, 32),
319        Err(AccessError::DoubleMalloc {
320            addr: 0x1000,
321            len: 32
322        })
323    );
324    assert_eq!(
325        wmemcheck_state.malloc(640 * 1024, 32),
326        Err(AccessError::OutOfBounds {
327            addr: 640 * 1024,
328            len: 32
329        })
330    );
331    assert!(wmemcheck_state.free(0x1000).is_ok());
332}
333
334#[test]
335fn update_sp_no_error() {
336    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
337
338    wmemcheck_state.set_stack_size(1024);
339    assert!(wmemcheck_state.update_stack_pointer(768).is_ok());
340    assert_eq!(wmemcheck_state.stack_pointer, 768);
341    assert!(wmemcheck_state.malloc(1024 * 2, 32).is_ok());
342    assert!(wmemcheck_state.free(1024 * 2).is_ok());
343    assert!(wmemcheck_state.update_stack_pointer(896).is_ok());
344    assert_eq!(wmemcheck_state.stack_pointer, 896);
345    assert!(wmemcheck_state.update_stack_pointer(1024).is_ok());
346}
347
348#[test]
349fn bad_stack_malloc() {
350    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
351
352    wmemcheck_state.set_stack_size(1024);
353
354    assert!(wmemcheck_state.update_stack_pointer(0).is_ok());
355    assert_eq!(wmemcheck_state.stack_pointer, 0);
356    assert_eq!(
357        wmemcheck_state.malloc(512, 32),
358        Err(AccessError::OutOfBounds { addr: 512, len: 32 })
359    );
360    assert_eq!(
361        wmemcheck_state.malloc(1022, 32),
362        Err(AccessError::OutOfBounds {
363            addr: 1022,
364            len: 32
365        })
366    );
367}
368
369#[test]
370fn stack_full_empty() {
371    let mut wmemcheck_state = Wmemcheck::new(640 * 1024);
372
373    wmemcheck_state.set_stack_size(1024);
374
375    assert!(wmemcheck_state.update_stack_pointer(0).is_ok());
376    assert_eq!(wmemcheck_state.stack_pointer, 0);
377    assert!(wmemcheck_state.update_stack_pointer(1024).is_ok());
378    assert_eq!(wmemcheck_state.stack_pointer, 1024)
379}
380
381#[test]
382fn from_test_program() {
383    let mut wmemcheck_state = Wmemcheck::new(1024 * 1024 * 128);
384    wmemcheck_state.set_stack_size(70864);
385    assert!(wmemcheck_state.write(70832, 1).is_ok());
386    assert!(wmemcheck_state.read(1138, 1).is_ok());
387}
wasmtime_internal_wmemcheck/lib.rs

wasmtime_internal_wmemcheck/
lib.rs
