Skip to content

Commit

Permalink
Dedup StdOut/StdErr observer pt.2 (#2876)
Browse files Browse the repository at this point in the history
* dedup stdout stdin

* mm

* Claude was not enough

* but was good enough

* ok

* okok

* okokok

* shorter now

* lol

* fixer
  • Loading branch information
tokatoka authored Jan 21, 2025
1 parent 6eabb79 commit 2842f95
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 18 deletions.
4 changes: 2 additions & 2 deletions libafl/src/feedbacks/stdio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ impl StdOutToMetadataFeedback {
.get(&self.o_ref)
.ok_or(Error::illegal_state("StdOutObserver is missing"))?;
let buffer = observer
.data
.output
.as_ref()
.ok_or(Error::illegal_state("StdOutObserver has no stdout"))?;
let stdout = String::from_utf8_lossy(buffer).into_owned();
Expand Down Expand Up @@ -139,7 +139,7 @@ where
.get(&self.o_ref)
.ok_or(Error::illegal_state("StdErrObserver is missing"))?;
let buffer = observer
.data
.output
.as_ref()
.ok_or(Error::illegal_state("StdErrObserver has no stderr"))?;
let stderr = String::from_utf8_lossy(buffer).into_owned();
Expand Down
33 changes: 17 additions & 16 deletions libafl/src/observers/stdio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,8 +71,8 @@ use crate::{observers::Observer, Error};
/// ) -> Result<bool, Error>
/// {
/// unsafe {
/// STDOUT = observers.get(&self.stdout_observer).unwrap().data.clone();
/// STDERR = observers.get(&self.stderr_observer).unwrap().data.clone();
/// STDOUT = observers.get(&self.stdout_observer).unwrap().output.clone();
/// STDERR = observers.get(&self.stderr_observer).unwrap().output.clone();
/// }
/// Ok(true)
/// }
Expand Down Expand Up @@ -169,58 +169,59 @@ use crate::{observers::Observer, Error};
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct StreamObserver<T> {
pub struct OutputObserver<T> {
/// The name of the observer.
pub name: Cow<'static, str>,
/// The captured stdout/stderr data during last execution.
pub data: Option<Vec<u8>>,
pub output: Option<Vec<u8>>,
/// Phantom data to hold the stream type
phantom: PhantomData<T>,
}

/// Marker traits to distinguish between stdout and stderr
/// Marker traits to mark stdout for the `OutputObserver`
#[derive(Debug, Clone)]
pub struct StdOutMarker;
/// Marker traits to distinguish between stdout and stderr

/// Marker traits to mark stderr for the `OutputObserver`
#[derive(Debug, Clone)]
pub struct StdErrMarker;

impl<T> StreamObserver<T> {
/// Create a new `StreamObserver` with the given name.
impl<T> OutputObserver<T> {
/// Create a new `OutputObserver` with the given name.
#[must_use]
pub fn new(name: &'static str) -> Self {
Self {
name: Cow::from(name),
data: None,
output: None,
phantom: PhantomData,
}
}

/// React to new stream data
pub fn observe(&mut self, data: &[u8]) {
self.data = Some(data.into());
self.output = Some(data.into());
}
}

impl<T> Named for StreamObserver<T> {
impl<T> Named for OutputObserver<T> {
fn name(&self) -> &Cow<'static, str> {
&self.name
}
}

impl<I, S, T> Observer<I, S> for StreamObserver<T> {
impl<I, S, T> Observer<I, S> for OutputObserver<T> {
fn pre_exec_child(&mut self, _state: &mut S, _input: &I) -> Result<(), Error> {
self.data = None;
self.output = None;
Ok(())
}

fn pre_exec(&mut self, _state: &mut S, _input: &I) -> Result<(), Error> {
self.data = None;
self.output = None;
Ok(())
}
}

/// An observer that captures stdout of a target.
pub type StdOutObserver = StreamObserver<StdOutMarker>;
pub type StdOutObserver = OutputObserver<StdOutMarker>;
/// An observer that captures stderr of a target.
pub type StdErrObserver = StreamObserver<StdErrMarker>;
pub type StdErrObserver = OutputObserver<StdErrMarker>;

0 comments on commit 2842f95

Please sign in to comment.