From 7b908733054d809671d5e18972c422f8e35e1bd3 Mon Sep 17 00:00:00 2001
From: Aarnav <aarnavbos@gmail.com>
Date: Thu, 30 May 2024 19:05:53 +0200
Subject: [PATCH] introduce CustomFilenameToTestcaseFeedback (#2256)

* introduce CustomTestcaseFilenameFeedback

* rename CustomTestcaseFilenameFeedback to CustomFilenameToTestcaseFeedback

* rename custom_testcase_filename to custom_filename

---------

Co-authored-by: Romain Malmain <romain.malmain@pm.me>
---
 libafl/src/feedbacks/custom_filename.rs | 134 ++++++++++++++++++++++++
 libafl/src/feedbacks/mod.rs             |   8 +-
 2 files changed, 139 insertions(+), 3 deletions(-)
 create mode 100644 libafl/src/feedbacks/custom_filename.rs

diff --git a/libafl/src/feedbacks/custom_filename.rs b/libafl/src/feedbacks/custom_filename.rs
new file mode 100644
index 0000000000..f1073506a4
--- /dev/null
+++ b/libafl/src/feedbacks/custom_filename.rs
@@ -0,0 +1,134 @@
+use alloc::{borrow::Cow, string::String};
+use core::{
+    fmt::{self, Debug, Formatter},
+    marker::PhantomData,
+};
+
+use libafl_bolts::Named;
+use serde::{Deserialize, Serialize};
+
+use crate::{
+    corpus::Testcase,
+    events::EventFirer,
+    executors::ExitKind,
+    feedbacks::{Feedback, FeedbackFactory},
+    inputs::Input,
+    observers::ObserversTuple,
+    state::State,
+    Error,
+};
+
+/// A [`CustomFilenameToTestcaseFeedback`] takes a closure which returns a filename for the testcase.
+/// Is never interesting (use with an Eager OR).
+/// Note: Use only in conjunction with a `Corpus` type that writes to disk.
+/// Note: If used as part of the `Objective` chain, then it will only apply to testcases which are
+/// `Objectives`, vice versa for `Feedback`.
+#[derive(Serialize, Deserialize)]
+pub struct CustomFilenameToTestcaseFeedback<F, I, S>
+where
+    I: Input,
+    S: State<Input = I>,
+    F: FnMut(&mut S, &mut Testcase<I>) -> Result<String, Error>,
+{
+    /// Closure that returns the filename.
+    func: F,
+    phantomm: PhantomData<(I, S)>,
+}
+
+impl<F, I, S> CustomFilenameToTestcaseFeedback<F, I, S>
+where
+    I: Input,
+    S: State<Input = I>,
+    F: FnMut(&mut S, &mut Testcase<I>) -> Result<String, Error>,
+{
+    /// Create a new [`CustomFilenameToTestcaseFeedback`].
+    pub fn new(func: F) -> Self {
+        Self {
+            func,
+            phantomm: PhantomData,
+        }
+    }
+}
+
+impl<F, I, S, T> FeedbackFactory<CustomFilenameToTestcaseFeedback<F, I, S>, S, T>
+    for CustomFilenameToTestcaseFeedback<F, I, S>
+where
+    I: Input,
+    S: State<Input = I>,
+    F: FnMut(&mut S, &mut Testcase<I>) -> Result<String, Error> + Clone,
+{
+    fn create_feedback(&self, _ctx: &T) -> CustomFilenameToTestcaseFeedback<F, I, S> {
+        Self {
+            func: self.func.clone(),
+            phantomm: self.phantomm,
+        }
+    }
+}
+
+impl<F, I, S> Named for CustomFilenameToTestcaseFeedback<F, I, S>
+where
+    I: Input,
+    S: State<Input = I>,
+    F: FnMut(&mut S, &mut Testcase<I>) -> Result<String, Error>,
+{
+    fn name(&self) -> &Cow<'static, str> {
+        static NAME: Cow<'static, str> = Cow::Borrowed("CustomFilenameToTestcaseFeedback");
+        &NAME
+    }
+}
+
+impl<F, I, S> Debug for CustomFilenameToTestcaseFeedback<F, I, S>
+where
+    I: Input,
+    S: State<Input = I>,
+    F: FnMut(&mut S, &mut Testcase<I>) -> Result<String, Error>,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("CustomFilenameToTestcaseFeedback")
+            .finish_non_exhaustive()
+    }
+}
+
+impl<F, I, S> Feedback<S> for CustomFilenameToTestcaseFeedback<F, I, S>
+where
+    S: State<Input = I>,
+    F: FnMut(&mut S, &mut Testcase<S::Input>) -> Result<String, Error>,
+    I: Input,
+{
+    #[allow(clippy::wrong_self_convention)]
+    #[inline]
+    fn is_interesting<EM, OT>(
+        &mut self,
+        _state: &mut S,
+        _manager: &mut EM,
+        _input: &I,
+        _observers: &OT,
+        _exit_kind: &ExitKind,
+    ) -> Result<bool, Error>
+    where
+        EM: EventFirer<State = S>,
+    {
+        Ok(false)
+    }
+
+    fn append_metadata<EM, OT>(
+        &mut self,
+        state: &mut S,
+        _manager: &mut EM,
+        _observers: &OT,
+        testcase: &mut Testcase<<S>::Input>,
+    ) -> Result<(), Error>
+    where
+        OT: ObserversTuple<S>,
+        EM: EventFirer<State = S>,
+    {
+        *testcase.filename_mut() = Some((self.func)(state, testcase)?);
+        Ok(())
+    }
+
+    #[cfg(feature = "track_hit_feedbacks")]
+    #[inline]
+    fn last_result(&self) -> Result<bool, Error> {
+        Ok(false)
+    }
+}
diff --git a/libafl/src/feedbacks/mod.rs b/libafl/src/feedbacks/mod.rs
index 4a8db96c7d..5a80358930 100644
--- a/libafl/src/feedbacks/mod.rs
+++ b/libafl/src/feedbacks/mod.rs
@@ -39,7 +39,12 @@ use crate::{
 };
 #[cfg(feature = "std")]
 pub mod concolic;
+#[cfg(feature = "std")]
+/// The module for `CustomFilenameToTestcaseFeedback`
+pub mod custom_filename;
 pub mod differential;
+/// The module for list feedback
+pub mod list;
 pub mod map;
 #[cfg(feature = "nautilus")]
 pub mod nautilus;
@@ -49,9 +54,6 @@ pub mod new_hash_feedback;
 pub mod stdio;
 pub mod transferred;
 
-/// The module for list feedback
-pub mod list;
-
 /// Feedbacks evaluate the observers.
 /// Basically, they reduce the information provided by an observer to a value,
 /// indicating the "interestingness" of the last run.