1 use std::{
2 alloc::{Layout, alloc, dealloc},
3 ops::Deref,
4 };
5 use wasmtime::{Result, error::OutOfMemory};
6 use wasmtime_fuzzing::oom::OomTest;
7
8 /// RAII wrapper around a raw allocation to deallocate it on drop.
9 struct Alloc {
10 layout: Layout,
11 ptr: *mut u8,
12 }
13
14 impl Drop for Alloc {
drop(&mut self)15 fn drop(&mut self) {
16 if !self.ptr.is_null() {
17 unsafe {
18 dealloc(self.ptr, self.layout);
19 }
20 }
21 }
22 }
23
24 impl Deref for Alloc {
25 type Target = *mut u8;
26
deref(&self) -> &Self::Target27 fn deref(&self) -> &Self::Target {
28 &self.ptr
29 }
30 }
31
32 impl Alloc {
33 /// Safety: same as `std::alloc::alloc`.
new(layout: Layout) -> Self34 unsafe fn new(layout: Layout) -> Self {
35 let ptr = unsafe { alloc(layout) };
36 Alloc { layout, ptr }
37 }
38 }
39
40 #[test]
smoke_test_ok() -> Result<()>41 pub(crate) fn smoke_test_ok() -> Result<()> {
42 OomTest::new().test(|| Ok(()))
43 }
44
45 #[test]
smoke_test_missed_oom() -> Result<()>46 pub(crate) fn smoke_test_missed_oom() -> Result<()> {
47 let err = OomTest::new()
48 .test(|| unsafe {
49 let _ = Alloc::new(Layout::new::<u64>());
50 Ok(())
51 })
52 .unwrap_err();
53 let err = format!("{err:?}");
54 assert!(
55 err.contains("OOM test function missed an OOM"),
56 "should have missed an OOM, got: {err}"
57 );
58 Ok(())
59 }
60
61 #[test]
smoke_test_disallow_alloc_after_oom() -> Result<()>62 pub(crate) fn smoke_test_disallow_alloc_after_oom() -> Result<()> {
63 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
64 let _ = OomTest::new().test(|| unsafe {
65 let layout = Layout::new::<u64>();
66 let p = Alloc::new(layout);
67 let _q = Alloc::new(layout);
68 if p.is_null() {
69 Err(OutOfMemory::new(layout.size()).into())
70 } else {
71 Ok(())
72 }
73 });
74 }));
75 assert!(result.is_err());
76 Ok(())
77 }
78
79 #[test]
smoke_test_allow_alloc_after_oom() -> Result<()>80 pub(crate) fn smoke_test_allow_alloc_after_oom() -> Result<()> {
81 OomTest::new().allow_alloc_after_oom(true).test(|| unsafe {
82 let layout = Layout::new::<u64>();
83 let p = Alloc::new(layout);
84 let q = Alloc::new(layout);
85 if p.is_null() || q.is_null() {
86 Err(OutOfMemory::new(layout.size()).into())
87 } else {
88 Ok(())
89 }
90 })
91 }
92