@@ -6,24 +6,31 @@ use std::{
66    sync:: OnceLock , 
77} ; 
88
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 ; 
913use  libafl:: { 
1014    Error , 
1115    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 , 
1519    } , 
1620} ; 
17- use  libafl_bolts:: os:: { ChildHandle ,  ForkResult } ; 
21+ use  libafl_bolts:: { 
22+     os:: { ChildHandle ,  ForkResult } , 
23+     shmem:: { ShMem ,  ShMemId ,  ShMemProvider } , 
24+ } ; 
1825use  nix:: { 
1926    sys:: signal:: { SigHandler ,  Signal } , 
2027    unistd:: Pid , 
2128} ; 
2229
23- #[ cfg( feature = "cmplog" ) ]  
24- use  crate :: cmps:: CMPLOG_MAP_PTR ; 
2530#[ cfg( feature = "cmplog_extended_instrumentation" ) ]  
2631use  crate :: cmps:: EXTENDED_CMPLOG_MAP_PTR ; 
32+ #[ cfg( feature = "cmplog" ) ]  
33+ use  crate :: cmps:: { AflppCmpLogMap ,  CMPLOG_MAP_PTR } ; 
2734
2835use  crate :: coverage:: { __afl_map_size,  EDGES_MAP_PTR ,  INPUT_LENGTH_PTR ,  INPUT_PTR ,  SHM_FUZZING } ; 
2936#[ cfg( any( target_os = "linux" ,  target_vendor = "apple" ) ) ]  
@@ -54,6 +61,7 @@ fn write_to_forkserver(message: &[u8]) -> Result<(), Error> {
5461    } 
5562    Ok ( ( ) ) 
5663} 
64+ #[ cfg( any( target_os = "linux" ,  target_vendor = "apple" ) ) ]  
5765fn  write_all_to_forkserver ( message :  & [ u8 ] )  -> Result < ( ) ,  Error >  { 
5866    let  mut  remain_len = message. len ( ) ; 
5967    while  remain_len > 0  { 
@@ -89,6 +97,39 @@ fn read_u32_from_forkserver() -> Result<u32, Error> {
8997    Ok ( u32:: from_ne_bytes ( buf) ) 
9098} 
9199
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+ 
92133/// Guard [`map_shared_memory`] is invoked only once 
93134static  SHM_MAP_GUARD :  OnceLock < ( ) >  = OnceLock :: new ( ) ; 
94135
@@ -97,31 +138,18 @@ static SHM_MAP_GUARD: OnceLock<()> = OnceLock::new();
97138/// 
98139/// If anything failed, the forkserver will be notified with 
99140/// [`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 >  { 
101142    if  SHM_MAP_GUARD . set ( ( ) ) . is_err ( )  { 
102143        return  Err ( Error :: illegal_state ( "shared memory has been mapped before" ) ) ; 
103144    } 
104-     map_shared_memory_internal ( ) 
145+     map_shared_memory_internal ( shmem_provider ) 
105146} 
106147
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 ) ?; 
123151    unsafe  { 
124-         EDGES_MAP_PTR  = map . cast ( ) ; 
152+         EDGES_MAP_PTR  = target_ptr ; 
125153    } 
126154    Ok ( ( ) ) 
127155} 
@@ -134,32 +162,23 @@ static INPUT_SHM_MAP_GUARD: OnceLock<()> = OnceLock::new();
134162/// 
135163/// If anything failed, the forkserver will be notified with 
136164/// [`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 >  { 
138166    if  INPUT_SHM_MAP_GUARD . set ( ( ) ) . is_err ( )  { 
139167        return  Err ( Error :: illegal_state ( "shared memory has been mapped before" ) ) ; 
140168    } 
141-     map_input_shared_memory_internal ( ) 
169+     map_input_shared_memory_internal ( shmem_provider ) 
142170} 
143171
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 ( ) ; 
163182    unsafe  { 
164183        INPUT_LENGTH_PTR  = map; 
165184        INPUT_PTR  = map. add ( 1 ) . cast ( ) ; 
@@ -177,36 +196,33 @@ static CMPLOG_SHM_MAP_GUARD: OnceLock<()> = OnceLock::new();
177196/// If anything failed, the forkserver will be notified with 
178197/// [`FS_ERROR_SHM_OPEN`]. 
179198#[ 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 >  { 
181200    if  CMPLOG_SHM_MAP_GUARD . set ( ( ) ) . is_err ( )  { 
182201        return  Err ( Error :: illegal_state ( "shared memory has been mapped before" ) ) ; 
183202    } 
184-     map_cmplog_shared_memory_internal ( ) 
203+     map_cmplog_shared_memory_internal ( shmem_provider ) 
185204} 
186205
187206#[ 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 >  { 
189210    let  Ok ( id_str)  = std:: env:: var ( SHM_CMPLOG_ENV_VAR )  else  { 
190211        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+         ) ) ) ; 
198215    } ; 
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) ; 
204220    unsafe  { 
205-         CMPLOG_MAP_PTR  = map . cast ( ) ; 
221+         CMPLOG_MAP_PTR  = target_ptr ; 
206222    } 
207223    #[ cfg( feature = "cmplog_extended_instrumentation" ) ]  
208224    unsafe  { 
209-         EXTENDED_CMPLOG_MAP_PTR  = map . cast ( ) ; 
225+         EXTENDED_CMPLOG_MAP_PTR  = target_ptr ; 
210226    } 
211227    Ok ( ( ) ) 
212228} 
0 commit comments