@@ -6,24 +6,31 @@ use std::{
6
6
sync:: OnceLock ,
7
7
} ;
8
8
9
+ #[ cfg( any( target_os = "linux" , target_vendor = "apple" ) ) ]
10
+ use libafl:: executors:: forkserver:: FS_NEW_OPT_AUTODTCT ;
11
+ #[ cfg( feature = "cmplog" ) ]
12
+ use libafl:: executors:: forkserver:: SHM_CMPLOG_ENV_VAR ;
9
13
use libafl:: {
10
14
Error ,
11
15
executors:: forkserver:: {
12
- FORKSRV_FD , FS_ERROR_SHM_OPEN , FS_NEW_OPT_AUTODTCT , FS_NEW_OPT_MAPSIZE ,
13
- FS_NEW_OPT_SHDMEM_FUZZ , FS_NEW_VERSION_MAX , FS_OPT_ERROR , SHM_CMPLOG_ENV_VAR , SHM_ENV_VAR ,
14
- SHM_FUZZ_ENV_VAR ,
16
+ AFL_MAP_SIZE_ENV_VAR , FORKSRV_FD , FS_ERROR_SHM_OPEN , FS_NEW_OPT_MAPSIZE ,
17
+ FS_NEW_OPT_SHDMEM_FUZZ , FS_NEW_VERSION_MAX , FS_OPT_ERROR , MAX_INPUT_SIZE_DEFAULT ,
18
+ SHM_ENV_VAR , SHM_FUZZ_ENV_VAR , SHM_FUZZ_MAP_SIZE_ENV_VAR , SHMEM_FUZZ_HDR_SIZE ,
15
19
} ,
16
20
} ;
17
- use libafl_bolts:: os:: { ChildHandle , ForkResult } ;
21
+ use libafl_bolts:: {
22
+ os:: { ChildHandle , ForkResult } ,
23
+ shmem:: { ShMem , ShMemId , ShMemProvider } ,
24
+ } ;
18
25
use nix:: {
19
26
sys:: signal:: { SigHandler , Signal } ,
20
27
unistd:: Pid ,
21
28
} ;
22
29
23
- #[ cfg( feature = "cmplog" ) ]
24
- use crate :: cmps:: CMPLOG_MAP_PTR ;
25
30
#[ cfg( feature = "cmplog_extended_instrumentation" ) ]
26
31
use crate :: cmps:: EXTENDED_CMPLOG_MAP_PTR ;
32
+ #[ cfg( feature = "cmplog" ) ]
33
+ use crate :: cmps:: { AflppCmpLogMap , CMPLOG_MAP_PTR } ;
27
34
28
35
use crate :: coverage:: { __afl_map_size, EDGES_MAP_PTR , INPUT_LENGTH_PTR , INPUT_PTR , SHM_FUZZING } ;
29
36
#[ cfg( any( target_os = "linux" , target_vendor = "apple" ) ) ]
@@ -54,6 +61,7 @@ fn write_to_forkserver(message: &[u8]) -> Result<(), Error> {
54
61
}
55
62
Ok ( ( ) )
56
63
}
64
+ #[ cfg( any( target_os = "linux" , target_vendor = "apple" ) ) ]
57
65
fn write_all_to_forkserver ( message : & [ u8 ] ) -> Result < ( ) , Error > {
58
66
let mut remain_len = message. len ( ) ;
59
67
while remain_len > 0 {
@@ -89,6 +97,39 @@ fn read_u32_from_forkserver() -> Result<u32, Error> {
89
97
Ok ( u32:: from_ne_bytes ( buf) )
90
98
}
91
99
100
+ /// Consume current shared memory structure, and get the raw pointer to
101
+ /// this shared memory.
102
+ ///
103
+ /// Note that calling this method will result in a memory leak.
104
+ fn shmem_into_raw < T : Sized > ( shmem : impl ShMem ) -> * mut T {
105
+ let mut manually_dropped = core:: mem:: ManuallyDrop :: new ( shmem) ;
106
+ manually_dropped. as_mut_ptr ( ) . cast ( )
107
+ }
108
+
109
+ fn map_shared_memory_common < SHM : ShMemProvider > (
110
+ shmem_provider : & mut SHM ,
111
+ map_env_var : & str ,
112
+ map_size_env_var : & str ,
113
+ map_size_default_fallback : usize ,
114
+ ) -> Result < * mut u8 , Error > {
115
+ let Ok ( id_str) = std:: env:: var ( map_env_var) else {
116
+ write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
117
+ return Err ( Error :: illegal_argument ( format ! (
118
+ "Error: shared memory variable {map_env_var} is not set"
119
+ ) ) ) ;
120
+ } ;
121
+ let map_size = if let Ok ( map_size_str) = std:: env:: var ( map_size_env_var) {
122
+ map_size_str
123
+ . parse ( )
124
+ . map_err ( |_| Error :: illegal_argument ( format ! ( "Invalid {map_size_env_var} value" ) ) ) ?
125
+ } else {
126
+ map_size_default_fallback
127
+ } ;
128
+ let shmem = shmem_provider. shmem_from_id_and_size ( ShMemId :: from_string ( & id_str) , map_size) ?;
129
+
130
+ Ok ( shmem_into_raw ( shmem) )
131
+ }
132
+
92
133
/// Guard [`map_shared_memory`] is invoked only once
93
134
static SHM_MAP_GUARD : OnceLock < ( ) > = OnceLock :: new ( ) ;
94
135
@@ -97,31 +138,18 @@ static SHM_MAP_GUARD: OnceLock<()> = OnceLock::new();
97
138
///
98
139
/// If anything failed, the forkserver will be notified with
99
140
/// [`FS_ERROR_SHM_OPEN`].
100
- pub fn map_shared_memory ( ) -> Result < ( ) , Error > {
141
+ pub fn map_shared_memory < SHM : ShMemProvider > ( shmem_provider : & mut SHM ) -> Result < ( ) , Error > {
101
142
if SHM_MAP_GUARD . set ( ( ) ) . is_err ( ) {
102
143
return Err ( Error :: illegal_state ( "shared memory has been mapped before" ) ) ;
103
144
}
104
- map_shared_memory_internal ( )
145
+ map_shared_memory_internal ( shmem_provider )
105
146
}
106
147
107
- fn map_shared_memory_internal ( ) -> Result < ( ) , Error > {
108
- let Ok ( id_str) = std:: env:: var ( SHM_ENV_VAR ) else {
109
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
110
- return Err ( Error :: illegal_argument (
111
- "Error: variable for edge coverage shared memory is not set" ,
112
- ) ) ;
113
- } ;
114
- let Ok ( shm_id) = id_str. parse ( ) else {
115
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
116
- return Err ( Error :: illegal_argument ( "Invalid __AFL_SHM_ID value" ) ) ;
117
- } ;
118
- let map = unsafe { libc:: shmat ( shm_id, core:: ptr:: null ( ) , 0 ) } ;
119
- if map. is_null ( ) || core:: ptr:: eq ( map, libc:: MAP_FAILED ) {
120
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
121
- return Err ( Error :: illegal_state ( "shmat for map" ) ) ;
122
- }
148
+ fn map_shared_memory_internal < SHM : ShMemProvider > ( shmem_provider : & mut SHM ) -> Result < ( ) , Error > {
149
+ let target_ptr =
150
+ map_shared_memory_common ( shmem_provider, SHM_ENV_VAR , AFL_MAP_SIZE_ENV_VAR , 65536 ) ?;
123
151
unsafe {
124
- EDGES_MAP_PTR = map . cast ( ) ;
152
+ EDGES_MAP_PTR = target_ptr ;
125
153
}
126
154
Ok ( ( ) )
127
155
}
@@ -134,32 +162,23 @@ static INPUT_SHM_MAP_GUARD: OnceLock<()> = OnceLock::new();
134
162
///
135
163
/// If anything failed, the forkserver will be notified with
136
164
/// [`FS_ERROR_SHM_OPEN`].
137
- pub fn map_input_shared_memory ( ) -> Result < ( ) , Error > {
165
+ pub fn map_input_shared_memory < SHM : ShMemProvider > ( shmem_provider : & mut SHM ) -> Result < ( ) , Error > {
138
166
if INPUT_SHM_MAP_GUARD . set ( ( ) ) . is_err ( ) {
139
167
return Err ( Error :: illegal_state ( "shared memory has been mapped before" ) ) ;
140
168
}
141
- map_input_shared_memory_internal ( )
169
+ map_input_shared_memory_internal ( shmem_provider )
142
170
}
143
171
144
- fn map_input_shared_memory_internal ( ) -> Result < ( ) , Error > {
145
- let Ok ( id_str) = std:: env:: var ( SHM_FUZZ_ENV_VAR ) else {
146
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
147
- return Err ( Error :: illegal_argument (
148
- "Error: variable for fuzzing shared memory is not set" ,
149
- ) ) ;
150
- } ;
151
- let Ok ( shm_id) = id_str. parse ( ) else {
152
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
153
- return Err ( Error :: illegal_argument ( "Invalid __AFL_SHM_FUZZ_ID value" ) ) ;
154
- } ;
155
- let map = unsafe { libc:: shmat ( shm_id, core:: ptr:: null ( ) , 0 ) } ;
156
- if map. is_null ( ) || core:: ptr:: eq ( map, libc:: MAP_FAILED ) {
157
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
158
- return Err ( Error :: illegal_state (
159
- "Could not access fuzzing shared memory" ,
160
- ) ) ;
161
- }
162
- let map: * mut u32 = map. cast ( ) ;
172
+ fn map_input_shared_memory_internal < SHM : ShMemProvider > (
173
+ shmem_provider : & mut SHM ,
174
+ ) -> Result < ( ) , Error > {
175
+ let target_ptr = map_shared_memory_common (
176
+ shmem_provider,
177
+ SHM_FUZZ_ENV_VAR ,
178
+ SHM_FUZZ_MAP_SIZE_ENV_VAR ,
179
+ MAX_INPUT_SIZE_DEFAULT + SHMEM_FUZZ_HDR_SIZE ,
180
+ ) ?;
181
+ let map: * mut u32 = target_ptr. cast ( ) ;
163
182
unsafe {
164
183
INPUT_LENGTH_PTR = map;
165
184
INPUT_PTR = map. add ( 1 ) . cast ( ) ;
@@ -177,36 +196,33 @@ static CMPLOG_SHM_MAP_GUARD: OnceLock<()> = OnceLock::new();
177
196
/// If anything failed, the forkserver will be notified with
178
197
/// [`FS_ERROR_SHM_OPEN`].
179
198
#[ cfg( feature = "cmplog" ) ]
180
- pub fn map_cmplog_shared_memory ( ) -> Result < ( ) , Error > {
199
+ pub fn map_cmplog_shared_memory < SHM : ShMemProvider > ( shmem_provider : & mut SHM ) -> Result < ( ) , Error > {
181
200
if CMPLOG_SHM_MAP_GUARD . set ( ( ) ) . is_err ( ) {
182
201
return Err ( Error :: illegal_state ( "shared memory has been mapped before" ) ) ;
183
202
}
184
- map_cmplog_shared_memory_internal ( )
203
+ map_cmplog_shared_memory_internal ( shmem_provider )
185
204
}
186
205
187
206
#[ cfg( feature = "cmplog" ) ]
188
- fn map_cmplog_shared_memory_internal ( ) -> Result < ( ) , Error > {
207
+ fn map_cmplog_shared_memory_internal < SHM : ShMemProvider > (
208
+ shmem_provider : & mut SHM ,
209
+ ) -> Result < ( ) , Error > {
189
210
let Ok ( id_str) = std:: env:: var ( SHM_CMPLOG_ENV_VAR ) else {
190
211
write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
191
- return Err ( Error :: illegal_argument (
192
- "Error: variable for cmplog shared memory is not set" ,
193
- ) ) ;
194
- } ;
195
- let Ok ( shm_id) = id_str. parse ( ) else {
196
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
197
- return Err ( Error :: illegal_argument ( "Invalid __AFL_CMPLOG_SHM_ID value" ) ) ;
212
+ return Err ( Error :: illegal_argument ( format ! (
213
+ "Error: shared memory variable {SHM_CMPLOG_ENV_VAR} is not set"
214
+ ) ) ) ;
198
215
} ;
199
- let map = unsafe { libc:: shmat ( shm_id, core:: ptr:: null ( ) , 0 ) } ;
200
- if map. is_null ( ) || core:: ptr:: eq ( map, libc:: MAP_FAILED ) {
201
- write_error_to_forkserver ( FS_ERROR_SHM_OPEN ) ?;
202
- return Err ( Error :: illegal_state ( "shmat for map" ) ) ;
203
- }
216
+ let map_size = size_of :: < AflppCmpLogMap > ( ) ;
217
+ let shmem = shmem_provider. shmem_from_id_and_size ( ShMemId :: from_string ( & id_str) , map_size) ?;
218
+
219
+ let target_ptr = shmem_into_raw ( shmem) ;
204
220
unsafe {
205
- CMPLOG_MAP_PTR = map . cast ( ) ;
221
+ CMPLOG_MAP_PTR = target_ptr ;
206
222
}
207
223
#[ cfg( feature = "cmplog_extended_instrumentation" ) ]
208
224
unsafe {
209
- EXTENDED_CMPLOG_MAP_PTR = map . cast ( ) ;
225
+ EXTENDED_CMPLOG_MAP_PTR = target_ptr ;
210
226
}
211
227
Ok ( ( ) )
212
228
}
0 commit comments