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