1
1
use core:: { ffi:: c_void, fmt:: Debug } ;
2
- use std:: sync:: atomic:: { AtomicBool , AtomicUsize , Ordering } ;
2
+ use std:: {
3
+ ptr:: { read_volatile, write_volatile} ,
4
+ sync:: atomic:: { AtomicBool , AtomicUsize , Ordering } ,
5
+ } ;
3
6
4
7
use libafl:: {
5
8
bolts:: tuples:: Named ,
@@ -20,29 +23,29 @@ extern "C" {
20
23
}
21
24
22
25
static OOMED : AtomicBool = AtomicBool :: new ( false ) ;
23
- static mut RSS_MAX : usize = 2 << 30 ; // 2GB, which is a typical default
26
+ static RSS_MAX : AtomicUsize = AtomicUsize :: new ( 2 << 30 ) ; // 2GB, which is the default
24
27
25
28
static MALLOC_SIZE : AtomicUsize = AtomicUsize :: new ( 0 ) ;
26
29
27
- extern "C" fn oom_malloc_hook ( ptr : * const c_void , size : usize ) {
30
+ pub extern "C" fn oom_malloc_hook ( ptr : * const c_void , size : usize ) {
28
31
let size = match unsafe { libafl_check_malloc_size ( ptr) } {
29
32
0 => size,
30
33
real => real,
31
34
} ;
32
35
33
36
let total = MALLOC_SIZE . fetch_add ( size, Ordering :: Relaxed ) + size;
34
- if total > unsafe { RSS_MAX } && !OOMED . load ( Ordering :: Relaxed ) {
37
+ if total > RSS_MAX . load ( Ordering :: Relaxed ) && !OOMED . load ( Ordering :: Relaxed ) {
35
38
OOMED . store ( true , Ordering :: Relaxed ) ;
36
39
unsafe {
37
40
// we need to kill the process in a way that immediately triggers the crash handler
38
41
let null = core:: ptr:: null_mut ( ) ;
39
- * null = 0 ;
42
+ write_volatile ( null, 0 ) ;
40
43
panic ! ( "We somehow didn't crash on a null pointer write. Strange..." ) ;
41
44
}
42
45
}
43
46
}
44
47
45
- extern "C" fn oom_free_hook ( ptr : * const c_void ) {
48
+ pub extern "C" fn oom_free_hook ( ptr : * const c_void ) {
46
49
let size = unsafe { libafl_check_malloc_size ( ptr) } ;
47
50
if MALLOC_SIZE
48
51
. fetch_sub ( size, Ordering :: Relaxed )
@@ -63,8 +66,12 @@ pub struct OOMObserver {
63
66
impl OOMObserver {
64
67
pub fn new ( rss_max : usize ) -> Self {
65
68
unsafe {
66
- RSS_MAX = rss_max;
67
- __sanitizer_install_malloc_and_free_hooks ( Some ( oom_malloc_hook) , Some ( oom_free_hook) ) ;
69
+ RSS_MAX . store ( rss_max, Ordering :: Relaxed ) ;
70
+ if __sanitizer_install_malloc_and_free_hooks ( Some ( oom_malloc_hook) , Some ( oom_free_hook) )
71
+ == 0
72
+ {
73
+ panic ! ( "Could not install malloc and free hooks" ) ;
74
+ }
68
75
}
69
76
Self { oomed : false }
70
77
}
0 commit comments