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. 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] 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] 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] 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] 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] 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] 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] 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] 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] 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] 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] 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] 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