Make OutFile auto-remove refcounted on drop (#654)

* Make OutFile auto-remove refcounted on drop

* clippy, windows

* remove debug print

* streamlined tmp files names

* outfile -> inputfile
This commit is contained in:
Dominik Maier 2022-05-27 18:01:44 +02:00 committed by GitHub
parent a544bc042d
commit c16738fd10
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 84 additions and 60 deletions

1
.gitignore vendored
View File

@ -6,6 +6,7 @@ vendor
.DS_Store .DS_Store
.env .env
*.tmp
*.o *.o
*.a *.a
*.so *.so

View File

@ -1,21 +1,21 @@
//! `LibAFL` functionality for filesystem interaction //! `LibAFL` functionality for filesystem interaction
#[cfg(feature = "std")]
use alloc::borrow::ToOwned;
use alloc::rc::Rc;
use core::cell::RefCell;
#[cfg(unix)]
use std::os::unix::prelude::{AsRawFd, RawFd};
use std::{ use std::{
fs::{self, remove_file, File, OpenOptions}, fs::{self, remove_file, File, OpenOptions},
io::{Seek, SeekFrom, Write}, io::{Seek, SeekFrom, Write},
path::{Path, PathBuf}, path::{Path, PathBuf},
}; };
#[cfg(unix)]
use std::os::unix::prelude::{AsRawFd, RawFd};
#[cfg(feature = "std")]
use alloc::borrow::ToOwned;
use crate::Error; use crate::Error;
/// The default filename to use to deliver testcases to the target /// The default filename to use to deliver testcases to the target
pub const OUTFILE_STD: &str = ".cur_input"; pub const INPUTFILE_STD: &str = ".cur_input";
/// Creates a `.{file_name}.tmp` file, and writes all bytes to it. /// Creates a `.{file_name}.tmp` file, and writes all bytes to it.
/// After all bytes have been written, the tmp-file is moved to it's original `path`. /// After all bytes have been written, the tmp-file is moved to it's original `path`.
@ -47,37 +47,46 @@ where
inner(path.as_ref(), bytes) inner(path.as_ref(), bytes)
} }
/// An [`OutFile`] to write fuzzer input to. /// An [`InputFile`] to write fuzzer input to.
/// The target/forkserver will read from this file. /// The target/forkserver will read from this file.
#[cfg(feature = "std")] #[cfg(feature = "std")]
#[derive(Debug)] #[derive(Debug)]
pub struct OutFile { pub struct InputFile {
/// The filename/path too this [`OutFile`] /// The filename/path too this [`InputFile`]
pub path: PathBuf, pub path: PathBuf,
/// The underlying file that got created /// The underlying file that got created
pub file: File, pub file: File,
/// The ref count for this [`InputFile`].
/// Once it reaches 0, the underlying [`File`] will be removed.
pub rc: Rc<RefCell<usize>>,
} }
impl Eq for OutFile {} impl Eq for InputFile {}
impl PartialEq for OutFile { impl PartialEq for InputFile {
fn eq(&self, other: &Self) -> bool { fn eq(&self, other: &Self) -> bool {
self.path == other.path self.path == other.path
} }
} }
impl Clone for OutFile { impl Clone for InputFile {
fn clone(&self) -> Self { fn clone(&self) -> Self {
{
let mut rc = self.rc.borrow_mut();
assert_ne!(*rc, usize::MAX, "InputFile rc overflow");
*rc += 1;
}
Self { Self {
path: self.path.clone(), path: self.path.clone(),
file: self.file.try_clone().unwrap(), file: self.file.try_clone().unwrap(),
rc: self.rc.clone(),
} }
} }
} }
#[cfg(feature = "std")] #[cfg(feature = "std")]
impl OutFile { impl InputFile {
/// Creates a new [`OutFile`] /// Creates a new [`InputFile`]
pub fn create<P>(filename: P) -> Result<Self, Error> pub fn create<P>(filename: P) -> Result<Self, Error>
where where
P: AsRef<Path>, P: AsRef<Path>,
@ -91,6 +100,7 @@ impl OutFile {
Ok(Self { Ok(Self {
path: filename.as_ref().to_owned(), path: filename.as_ref().to_owned(),
file: f, file: f,
rc: Rc::new(RefCell::new(1)),
}) })
} }
@ -123,25 +133,38 @@ impl OutFile {
} }
#[cfg(feature = "std")] #[cfg(feature = "std")]
impl Drop for OutFile { impl Drop for InputFile {
fn drop(&mut self) { fn drop(&mut self) {
let mut rc = self.rc.borrow_mut();
assert_ne!(*rc, 0, "InputFile rc should never be 0");
*rc -= 1;
if *rc == 0 {
// try to remove the file, but ignore errors // try to remove the file, but ignore errors
drop(remove_file(&self.path)); drop(remove_file(&self.path));
} }
} }
}
#[cfg(test)] #[cfg(test)]
mod test { mod test {
use crate::bolts::fs::write_file_atomic; use crate::bolts::fs::{write_file_atomic, InputFile};
use std::fs; use std::fs;
#[test] #[test]
fn test_atomic_file_write() { fn test_atomic_file_write() {
let path = "atomic_file_testfile"; let path = "test_atomic_file_write.tmp";
write_file_atomic(&path, b"test").unwrap(); write_file_atomic(&path, b"test").unwrap();
let content = fs::read_to_string(&path).unwrap(); let content = fs::read_to_string(&path).unwrap();
fs::remove_file(&path).unwrap(); fs::remove_file(&path).unwrap();
assert_eq!(content, "test"); assert_eq!(content, "test");
} }
#[test]
fn test_cloned_ref() {
let mut one = InputFile::create("test_cloned_ref.tmp").unwrap();
let two = one.clone();
one.write_buf("Welp".as_bytes()).unwrap();
drop(one);
assert_eq!("Welp", fs::read_to_string(&two.path).unwrap());
}
} }

View File

@ -19,7 +19,7 @@ use std::{
use crate::{ use crate::{
bolts::{ bolts::{
fs::{OutFile, OUTFILE_STD}, fs::{InputFile, INPUTFILE_STD},
tuples::MatchName, tuples::MatchName,
AsSlice, AsSlice,
}, },
@ -39,7 +39,7 @@ use super::HasObservers;
/// How to deliver input to an external program /// How to deliver input to an external program
/// `StdIn`: The traget reads from stdin /// `StdIn`: The traget reads from stdin
/// `File`: The target reads from the specified [`OutFile`] /// `File`: The target reads from the specified [`InputFile`]
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
enum InputLocation { enum InputLocation {
/// Mutate a commandline argument to deliver an input /// Mutate a commandline argument to deliver an input
@ -49,11 +49,11 @@ enum InputLocation {
}, },
/// Deliver input via `StdIn` /// Deliver input via `StdIn`
StdIn, StdIn,
/// Deliver the iniput via the specified [`OutFile`] /// Deliver the input via the specified [`InputFile`]
/// You can use specify [`OutFile::create(OUTFILE_STD)`] to use a default filename. /// You can use specify [`InputFile::create(INPUTFILE_STD)`] to use a default filename.
File { File {
/// The fiel to write input to. The target should read input from this location. /// The file to write input to. The target should read input from this location.
out_file: OutFile, input_file: InputFile,
}, },
} }
@ -133,8 +133,8 @@ impl CommandConfigurator for StdCommandConfigurator {
drop(stdin); drop(stdin);
Ok(handle) Ok(handle)
} }
InputLocation::File { out_file } => { InputLocation::File { input_file } => {
out_file.write_buf(input.target_bytes().as_slice())?; input_file.write_buf(input.target_bytes().as_slice())?;
Ok(self.command.spawn()?) Ok(self.command.spawn()?)
} }
} }
@ -250,7 +250,7 @@ where
has_asan_observer, has_asan_observer,
configurer: StdCommandConfigurator { configurer: StdCommandConfigurator {
input_location: InputLocation::File { input_location: InputLocation::File {
out_file: OutFile::create(path)?, input_file: InputFile::create(path)?,
}, },
command, command,
debug_child, debug_child,
@ -461,7 +461,7 @@ impl CommandExecutorBuilder {
/// Uses a default filename. /// Uses a default filename.
/// Use [`Self::arg_input_file`] to specify a custom filename. /// Use [`Self::arg_input_file`] to specify a custom filename.
pub fn arg_input_file_std(&mut self) -> &mut Self { pub fn arg_input_file_std(&mut self) -> &mut Self {
self.arg_input_file(OUTFILE_STD); self.arg_input_file(INPUTFILE_STD);
self self
} }
@ -469,9 +469,9 @@ impl CommandExecutorBuilder {
/// and adds the filename as arg to at the current position. /// and adds the filename as arg to at the current position.
pub fn arg_input_file<P: AsRef<Path>>(&mut self, path: P) -> &mut Self { pub fn arg_input_file<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
self.arg(path.as_ref()); self.arg(path.as_ref());
let out_file_std = OutFile::create(path.as_ref()).unwrap(); let input_file_std = InputFile::create(path.as_ref()).unwrap();
self.input(InputLocation::File { self.input(InputLocation::File {
out_file: out_file_std, input_file: input_file_std,
}); });
self self
} }

View File

@ -15,7 +15,7 @@ use std::{
use crate::{ use crate::{
bolts::{ bolts::{
fs::{OutFile, OUTFILE_STD}, fs::{InputFile, INPUTFILE_STD},
os::{dup2, pipes::Pipe}, os::{dup2, pipes::Pipe},
shmem::{ShMem, ShMemProvider, StdShMemProvider}, shmem::{ShMem, ShMemProvider, StdShMemProvider},
AsMutSlice, AsSlice, AsMutSlice, AsSlice,
@ -175,7 +175,7 @@ impl Forkserver {
target: OsString, target: OsString,
args: Vec<OsString>, args: Vec<OsString>,
envs: Vec<(OsString, OsString)>, envs: Vec<(OsString, OsString)>,
out_filefd: RawFd, input_filefd: RawFd,
use_stdin: bool, use_stdin: bool,
memlimit: u64, memlimit: u64,
debug_output: bool, debug_output: bool,
@ -199,7 +199,7 @@ impl Forkserver {
.envs(envs) .envs(envs)
.setlimit(memlimit) .setlimit(memlimit)
.setsid() .setsid()
.setstdin(out_filefd, use_stdin) .setstdin(input_filefd, use_stdin)
.setpipe( .setpipe(
st_pipe.read_end().unwrap(), st_pipe.read_end().unwrap(),
st_pipe.write_end().unwrap(), st_pipe.write_end().unwrap(),
@ -337,10 +337,10 @@ pub trait HasForkserver {
fn forkserver_mut(&mut self) -> &mut Forkserver; fn forkserver_mut(&mut self) -> &mut Forkserver;
/// The file the forkserver is reading from /// The file the forkserver is reading from
fn out_file(&self) -> &OutFile; fn input_file(&self) -> &InputFile;
/// The file the forkserver is reading from, mutable /// The file the forkserver is reading from, mutable
fn out_file_mut(&mut self) -> &mut OutFile; fn input_file_mut(&mut self) -> &mut InputFile;
/// The map of the fuzzer /// The map of the fuzzer
fn shmem(&self) -> &Option<<<Self as HasForkserver>::SP as ShMemProvider>::ShMem>; fn shmem(&self) -> &Option<<<Self as HasForkserver>::SP as ShMemProvider>::ShMem>;
@ -405,7 +405,7 @@ where
} }
None => { None => {
self.executor self.executor
.out_file_mut() .input_file_mut()
.write_buf(input.target_bytes().as_slice())?; .write_buf(input.target_bytes().as_slice())?;
} }
} }
@ -479,12 +479,12 @@ where
{ {
target: OsString, target: OsString,
args: Vec<OsString>, args: Vec<OsString>,
out_file: OutFile, input_file: InputFile,
forkserver: Forkserver, forkserver: Forkserver,
observers: OT, observers: OT,
map: Option<SP::ShMem>, map: Option<SP::ShMem>,
phantom: PhantomData<(I, S)>, phantom: PhantomData<(I, S)>,
/// Cache that indicates if we have a asan observer registered. /// Cache that indicates if we have a `ASan` observer registered.
has_asan_observer: Option<bool>, has_asan_observer: Option<bool>,
} }
@ -497,7 +497,7 @@ where
f.debug_struct("ForkserverExecutor") f.debug_struct("ForkserverExecutor")
.field("target", &self.target) .field("target", &self.target)
.field("args", &self.args) .field("args", &self.args)
.field("out_file", &self.out_file) .field("input_file", &self.input_file)
.field("forkserver", &self.forkserver) .field("forkserver", &self.forkserver)
.field("observers", &self.observers) .field("observers", &self.observers)
.field("map", &self.map) .field("map", &self.map)
@ -534,9 +534,9 @@ where
&self.forkserver &self.forkserver
} }
/// The [`OutFile`] used by this [`Executor`]. /// The [`InputFile`] used by this [`Executor`].
pub fn out_file(&self) -> &OutFile { pub fn input_file(&self) -> &InputFile {
&self.out_file &self.input_file
} }
} }
@ -549,7 +549,7 @@ pub struct ForkserverExecutorBuilder<'a, SP> {
debug_child: bool, debug_child: bool,
use_stdin: bool, use_stdin: bool,
autotokens: Option<&'a mut Tokens>, autotokens: Option<&'a mut Tokens>,
out_filename: Option<OsString>, input_filename: Option<OsString>,
shmem_provider: Option<&'a mut SP>, shmem_provider: Option<&'a mut SP>,
} }
@ -565,12 +565,12 @@ impl<'a, SP> ForkserverExecutorBuilder<'a, SP> {
OT: ObserversTuple<I, S>, OT: ObserversTuple<I, S>,
SP: ShMemProvider, SP: ShMemProvider,
{ {
let out_filename = match &self.out_filename { let input_filename = match &self.input_filename {
Some(name) => name.clone(), Some(name) => name.clone(),
None => OsString::from(".cur_input"), None => OsString::from(".cur_input"),
}; };
let out_file = OutFile::create(&out_filename)?; let input_file = InputFile::create(&input_filename)?;
let map = match &mut self.shmem_provider { let map = match &mut self.shmem_provider {
None => None, None => None,
@ -591,7 +591,7 @@ impl<'a, SP> ForkserverExecutorBuilder<'a, SP> {
t.clone(), t.clone(),
self.arguments.clone(), self.arguments.clone(),
self.envs.clone(), self.envs.clone(),
out_file.as_raw_fd(), input_file.as_raw_fd(),
self.use_stdin, self.use_stdin,
0, 0,
self.debug_child, self.debug_child,
@ -671,7 +671,7 @@ impl<'a, SP> ForkserverExecutorBuilder<'a, SP> {
Ok(ForkserverExecutor { Ok(ForkserverExecutor {
target, target,
args: self.arguments.clone(), args: self.arguments.clone(),
out_file, input_file,
forkserver, forkserver,
observers, observers,
map, map,
@ -701,7 +701,7 @@ impl<'a, SP> ForkserverExecutorBuilder<'a, SP> {
if item.as_ref() == "@@" && use_stdin { if item.as_ref() == "@@" && use_stdin {
use_stdin = false; use_stdin = false;
res.push(OsString::from(".cur_input")); res.push(OsString::from(".cur_input"));
} else if let Some(name) = &self.out_filename { } else if let Some(name) = &self.input_filename {
if name == item.as_ref() && use_stdin { if name == item.as_ref() && use_stdin {
use_stdin = false; use_stdin = false;
res.push(name.clone()); res.push(name.clone());
@ -734,7 +734,7 @@ impl<'a> ForkserverExecutorBuilder<'a, StdShMemProvider> {
debug_child: false, debug_child: false,
use_stdin: true, use_stdin: true,
autotokens: None, autotokens: None,
out_filename: None, input_filename: None,
shmem_provider: None, shmem_provider: None,
} }
} }
@ -806,14 +806,14 @@ impl<'a> ForkserverExecutorBuilder<'a, StdShMemProvider> {
/// Place the input at this position and set the filename for the input. /// Place the input at this position and set the filename for the input.
pub fn arg_input_file<P: AsRef<Path>>(self, path: P) -> Self { pub fn arg_input_file<P: AsRef<Path>>(self, path: P) -> Self {
let mut moved = self.arg(path.as_ref()); let mut moved = self.arg(path.as_ref());
moved.out_filename = Some(path.as_ref().as_os_str().to_os_string()); moved.input_filename = Some(path.as_ref().as_os_str().to_os_string());
moved moved
} }
#[must_use] #[must_use]
/// Place the input at this position and set the default filename for the input. /// Place the input at this position and set the default filename for the input.
pub fn arg_input_file_std(self) -> Self { pub fn arg_input_file_std(self) -> Self {
self.arg_input_file(OUTFILE_STD) self.arg_input_file(INPUTFILE_STD)
} }
#[must_use] #[must_use]
@ -835,7 +835,7 @@ impl<'a> ForkserverExecutorBuilder<'a, StdShMemProvider> {
debug_child: self.debug_child, debug_child: self.debug_child,
use_stdin: self.use_stdin, use_stdin: self.use_stdin,
autotokens: self.autotokens, autotokens: self.autotokens,
out_filename: self.out_filename, input_filename: self.input_filename,
shmem_provider: Some(shmem_provider), shmem_provider: Some(shmem_provider),
} }
} }
@ -875,7 +875,7 @@ where
.copy_from_slice(target_bytes.as_slice()); .copy_from_slice(target_bytes.as_slice());
} }
None => { None => {
self.out_file.write_buf(input.target_bytes().as_slice())?; self.input_file.write_buf(input.target_bytes().as_slice())?;
} }
} }
@ -971,13 +971,13 @@ where
} }
#[inline] #[inline]
fn out_file(&self) -> &OutFile { fn input_file(&self) -> &InputFile {
&self.out_file &self.input_file
} }
#[inline] #[inline]
fn out_file_mut(&mut self) -> &mut OutFile { fn input_file_mut(&mut self) -> &mut InputFile {
&mut self.out_file &mut self.input_file
} }
#[inline] #[inline]