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
8 use std::cmp::*;
9 use std::collections::HashMap;
10
11 /// Memory checker for wasm guest.
12 pub 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)]
22 pub 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)]
37 pub enum MemState {
38 /// Unallocated memory.
39 Unallocated,
40 /// Initialized but undefined memory.
41 ValidToWrite,
42 /// Initialized and defined memory.
43 ValidToReadWrite,
44 }
45
46 impl Wmemcheck {
47 /// Initializes memory checker instance.
new(mem_size: usize) -> Wmemcheck48 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.
malloc(&mut self, addr: usize, len: usize) -> Result<(), AccessError>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.
read(&mut self, addr: usize, len: usize) -> Result<(), AccessError>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.
write(&mut self, addr: usize, len: usize) -> Result<(), AccessError>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.
free(&mut self, addr: usize) -> Result<(), AccessError>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
is_in_bounds_heap(&self, addr: usize, len: usize) -> bool142 fn is_in_bounds_heap(&self, addr: usize, len: usize) -> bool {
143 self.max_stack_size <= addr && addr + len <= self.metadata.len()
144 }
145
is_in_bounds_stack(&self, addr: usize, len: usize) -> bool146 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.
update_stack_pointer(&mut self, new_sp: usize) -> Result<(), AccessError>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.
memcheck_on(&mut self)171 pub fn memcheck_on(&mut self) {
172 self.flag = true;
173 }
174
175 /// Turns memory checking off.
memcheck_off(&mut self)176 pub fn memcheck_off(&mut self) {
177 self.flag = false;
178 }
179
180 /// Initializes stack and stack pointer in memory checker metadata.
set_stack_size(&mut self, stack_size: usize)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.
update_mem_size(&mut self, num_bytes: usize)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]
basic_wmemcheck()197 fn 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]
read_before_initializing()210 fn 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]
use_after_free()226 fn 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]
double_free()243 fn 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]
out_of_bounds_malloc()256 fn 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]
out_of_bounds_read()277 fn 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]
double_malloc()291 fn 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]
error_type()313 fn 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]
update_sp_no_error()335 fn 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]
bad_stack_malloc()349 fn 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]
stack_full_empty()370 fn 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]
from_test_program()382 fn 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 }
388