Skip to content

Commit

Permalink
Fix pipe I/O in forkserver (AFLplusplus#2602)
Browse files Browse the repository at this point in the history 
* Fix read_st_size in forkserver

* fixed read_st, read_st_size, and write_ctl

with read_exact and write_all

* Cleanup redundant return values in forkserver pipe I/O

* Forkserver: avoid using read_exact/write_all for reading/writing integers

* Forkserver: avoid initializing the vec twice when reading a chunk of bytes

* Fix code formatting
  • Loading branch information
@henryhchchc
henryhchchc authored Oct 14, 2024
1 parent da8f17f commit ea4a281
Showing 1 changed file with 61 additions and 51 deletions.
112 changes: 61 additions & 51 deletions libafl/src/executors/forkserver.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -488,27 +488,54 @@ impl Forkserver {
}

/// Read from the st pipe
pub fn read_st(&mut self) -> Result<(usize, i32), Error> {
pub fn read_st(&mut self) -> Result<i32, Error> {
let mut buf: [u8; 4] = [0_u8; 4];

let rlen = self.st_pipe.read(&mut buf)?;
let val: i32 = i32::from_ne_bytes(buf);
Ok((rlen, val))
if rlen == size_of::<i32>() {
Ok(val)
} else {
// NOTE: The underlying API does not guarantee that the read will return
// exactly four bytes, but the chance of this happening is very low.
// This is a sacrifice of correctness for performance.
Err(Error::illegal_state(format!(
"Could not read from st pipe. Expected {} bytes, got {rlen} bytes",
size_of::<i32>()
)))
}
}

/// Read bytes of any length from the st pipe
pub fn read_st_size(&mut self, size: usize) -> Result<(usize, Vec<u8>), Error> {
let mut buf = vec![0; size];

let rlen = self.st_pipe.read(&mut buf)?;
Ok((rlen, buf))
pub fn read_st_size(&mut self, size: usize) -> Result<Vec<u8>, Error> {
let mut buf = Vec::with_capacity(size);
// SAFETY: `buf` will not be returned with `Ok` unless it is filled with `size` bytes.
// So it is ok to set the length to `size` such that the length of `&mut buf` is `size`
// and the `read_exact` call will try to read `size` bytes.
#[allow(
clippy::uninit_vec,
reason = "The vec will be filled right after setting the length."
)]
unsafe {
buf.set_len(size);
}
self.st_pipe.read_exact(&mut buf)?;
Ok(buf)
}

/// Write to the ctl pipe
pub fn write_ctl(&mut self, val: i32) -> Result<usize, Error> {
pub fn write_ctl(&mut self, val: i32) -> Result<(), Error> {
let slen = self.ctl_pipe.write(&val.to_ne_bytes())?;

Ok(slen)
if slen == size_of::<i32>() {
Ok(())
} else {
// NOTE: The underlying API does not guarantee that exactly four bytes
// are written, but the chance of this happening is very low.
// This is a sacrifice of correctness for performance.
Err(Error::illegal_state(format!(
"Could not write to ctl pipe. Expected {} bytes, wrote {slen} bytes",
size_of::<i32>()
)))
}
}

/// Read a message from the child process.
Expand Down Expand Up @@ -846,11 +873,10 @@ where
}
};

let (rlen, version_status) = forkserver.read_st()?; // Initial handshake, read 4-bytes hello message from the forkserver.

if rlen != 4 {
// Initial handshake, read 4-bytes hello message from the forkserver.
let Ok(version_status) = forkserver.read_st() else {
return Err(Error::unknown("Failed to start a forkserver".to_string()));
}
};

if (version_status & FS_NEW_ERROR) == FS_NEW_ERROR {
report_error_and_exit(version_status & 0x0000ffff)?;
Expand Down Expand Up @@ -896,8 +922,7 @@ where

let xored_status = (status as u32 ^ 0xffffffff) as i32;

let send_len = forkserver.write_ctl(xored_status)?;
if send_len != 4 {
if forkserver.write_ctl(xored_status).is_err() {
return Err(Error::unknown("Writing to forkserver failed.".to_string()));
}

Expand All @@ -906,20 +931,18 @@ where
version
);

let (read_len, status) = forkserver.read_st()?;
if read_len != 4 {
let Ok(status) = forkserver.read_st() else {
return Err(Error::unknown(
"Reading from forkserver failed.".to_string(),
));
}
};

if status & FS_NEW_OPT_MAPSIZE == FS_NEW_OPT_MAPSIZE {
let (read_len, fsrv_map_size) = forkserver.read_st()?;
if read_len != 4 {
let Ok(fsrv_map_size) = forkserver.read_st() else {
return Err(Error::unknown(
"Failed to read map size from forkserver".to_string(),
));
}
};
self.set_map_size(fsrv_map_size)?;
}

Expand All @@ -937,12 +960,11 @@ where
if status & FS_NEW_OPT_AUTODICT != 0 {
// Here unlike shmem input fuzzing, we are forced to read things
// hence no self.autotokens.is_some() to check if we proceed
let (read_len, autotokens_size) = forkserver.read_st()?;
if read_len != 4 {
let Ok(autotokens_size) = forkserver.read_st() else {
return Err(Error::unknown(
"Failed to read autotokens size from forkserver".to_string(),
));
}
};

let tokens_size_max = 0xffffff;

Expand All @@ -952,20 +974,17 @@ where
));
}
log::info!("Autotokens size {autotokens_size:x}");
let (rlen, buf) = forkserver.read_st_size(autotokens_size as usize)?;

if rlen != autotokens_size as usize {
let Ok(buf) = forkserver.read_st_size(autotokens_size as usize) else {
return Err(Error::unknown("Failed to load autotokens".to_string()));
}
};
if let Some(t) = &mut self.autotokens {
t.parse_autodict(&buf, autotokens_size as usize);
}
}

let (read_len, aflx) = forkserver.read_st()?;
if read_len != 4 {
let Ok(aflx) = forkserver.read_st() else {
return Err(Error::unknown("Reading from forkserver failed".to_string()));
}
};

if aflx != keep {
return Err(Error::unknown(format!(
Expand Down Expand Up @@ -1015,18 +1034,16 @@ where
// if send_status is not changed (Options are available but we didn't use any), then don't send the next write_ctl message.
// This is important

let send_len = forkserver.write_ctl(send_status)?;
if send_len != 4 {
if forkserver.write_ctl(send_status).is_err() {
return Err(Error::unknown("Writing to forkserver failed.".to_string()));
}

if (send_status & FS_OPT_AUTODICT) == FS_OPT_AUTODICT {
let (read_len, dict_size) = forkserver.read_st()?;
if read_len != 4 {
let Ok(dict_size) = forkserver.read_st() else {
return Err(Error::unknown(
"Reading from forkserver failed.".to_string(),
));
}
};

if !(2..=0xffffff).contains(&dict_size) {
return Err(Error::illegal_state(
Expand All @@ -1036,11 +1053,9 @@ where

log::info!("Autodict size {dict_size:x}");

let (rlen, buf) = forkserver.read_st_size(dict_size as usize)?;

if rlen != dict_size as usize {
let Ok(buf) = forkserver.read_st_size(dict_size as usize) else {
return Err(Error::unknown("Failed to load autodictionary".to_string()));
}
};
if let Some(t) = &mut self.autotokens {
t.parse_autodict(&buf, dict_size as usize);
}
Expand Down Expand Up @@ -1422,22 +1437,18 @@ where
.write_buf(&input_bytes.as_slice()[..input_size])?;
}

let send_len = self.forkserver.write_ctl(last_run_timed_out)?;

self.forkserver.set_last_run_timed_out(false);

if send_len != 4 {
if self.forkserver.write_ctl(last_run_timed_out).is_err() {
return Err(Error::unknown(
"Unable to request new process from fork server (OOM?)".to_string(),
));
}

let (recv_pid_len, pid) = self.forkserver.read_st()?;
if recv_pid_len != 4 {
let Ok(pid) = self.forkserver.read_st() else {
return Err(Error::unknown(
"Unable to request new process from fork server (OOM?)".to_string(),
));
}
};

if pid <= 0 {
return Err(Error::unknown(
Expand Down Expand Up @@ -1466,8 +1477,7 @@ where

// We need to kill the child in case he has timed out, or we can't get the correct pid in the next call to self.executor.forkserver_mut().read_st()?
let _ = kill(self.forkserver().child_pid(), self.forkserver.kill_signal);
let (recv_status_len, _) = self.forkserver.read_st()?;
if recv_status_len != 4 {
if self.forkserver.read_st().is_err() {
return Err(Error::unknown("Could not kill timed-out child".to_string()));
}
exit_kind = ExitKind::Timeout;
Expand Down

0 comments on commit ea4a281

Please sign in to comment.