From 98ef505a0e3d09022e9fe59d33363f54fa0e7ad3 Mon Sep 17 00:00:00 2001
From: "Dongjia \"toka\" Zhang" <tokazerkje@outlook.com>
Date: Wed, 19 Feb 2025 20:24:15 +0100
Subject: [PATCH] Add `add_metadata_checked` method (#3008)

* new api

* found a bug :>
---
 libafl/src/common/mod.rs                  | 31 +++++++++++++++++++++++
 libafl/src/feedbacks/list.rs              |  2 +-
 libafl/src/feedbacks/map.rs               |  2 +-
 libafl/src/feedbacks/new_hash_feedback.rs |  4 +--
 libafl/src/observers/map/mod.rs           |  4 ---
 5 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/libafl/src/common/mod.rs b/libafl/src/common/mod.rs
index 69345a1fa1..b55701208e 100644
--- a/libafl/src/common/mod.rs
+++ b/libafl/src/common/mod.rs
@@ -26,6 +26,23 @@ pub trait HasMetadata {
         self.metadata_map_mut().insert(meta);
     }
 
+    /// Add a metadata to the metadata map
+    /// Returns error if the metadata is already there
+    #[inline]
+    fn add_metadata_checked<M>(&mut self, meta: M) -> Result<(), Error>
+    where
+        M: SerdeAny,
+    {
+        if self.has_metadata::<M>() {
+            return Err(Error::illegal_argument(format!(
+                "Tried to add a metadata of {}. But this will overwrite the existing metadata",
+                type_name::<M>()
+            )));
+        }
+        self.metadata_map_mut().insert(meta);
+        Ok(())
+    }
+
     /// Gets metadata, or inserts it using the given construction function `default`
     fn metadata_or_insert_with<M>(&mut self, default: impl FnOnce() -> M) -> &mut M
     where
@@ -94,6 +111,20 @@ pub trait HasNamedMetadata {
         self.named_metadata_map_mut().insert(name, meta);
     }
 
+    /// Add a metadata to the metadata map
+    /// Return an error if there already is the metadata with the same name
+    #[inline]
+    fn add_named_metadata_checked<M>(&mut self, name: &str, meta: M) -> Result<(), Error>
+    where
+        M: SerdeAny,
+    {
+        if self.has_named_metadata::<M>(name) {
+            return Err(Error::illegal_argument(format!("Tried to add a metadata of {} named {}. But this will overwrite the existing metadata", type_name::<M>(), name)));
+        }
+        self.named_metadata_map_mut().insert(name, meta);
+        Ok(())
+    }
+
     /// Add a metadata to the metadata map
     #[inline]
     fn remove_named_metadata<M>(&mut self, name: &str) -> Option<Box<M>>
diff --git a/libafl/src/feedbacks/list.rs b/libafl/src/feedbacks/list.rs
index 0a1f76b256..0a24ef1ead 100644
--- a/libafl/src/feedbacks/list.rs
+++ b/libafl/src/feedbacks/list.rs
@@ -118,7 +118,7 @@ where
     T: Debug + Eq + Hash + for<'a> Deserialize<'a> + Serialize + Default + Copy + 'static,
 {
     fn init_state(&mut self, state: &mut S) -> Result<(), Error> {
-        state.add_named_metadata(self.name(), ListFeedbackMetadata::<T>::default());
+        state.add_named_metadata_checked(self.name(), ListFeedbackMetadata::<T>::default())?;
         Ok(())
     }
 }
diff --git a/libafl/src/feedbacks/map.rs b/libafl/src/feedbacks/map.rs
index 61699eff63..59d219cb7a 100644
--- a/libafl/src/feedbacks/map.rs
+++ b/libafl/src/feedbacks/map.rs
@@ -386,7 +386,7 @@ where
     fn init_state(&mut self, state: &mut S) -> Result<(), Error> {
         // Initialize `MapFeedbackMetadata` with an empty vector and add it to the state.
         // The `MapFeedbackMetadata` would be resized on-demand in `is_interesting`
-        state.add_named_metadata(&self.name, MapFeedbackMetadata::<O::Entry>::default());
+        state.add_named_metadata_checked(&self.name, MapFeedbackMetadata::<O::Entry>::default())?;
         Ok(())
     }
 }
diff --git a/libafl/src/feedbacks/new_hash_feedback.rs b/libafl/src/feedbacks/new_hash_feedback.rs
index f79be125fc..a48237c690 100644
--- a/libafl/src/feedbacks/new_hash_feedback.rs
+++ b/libafl/src/feedbacks/new_hash_feedback.rs
@@ -141,10 +141,10 @@ where
     S: HasNamedMetadata,
 {
     fn init_state(&mut self, state: &mut S) -> Result<(), Error> {
-        state.add_named_metadata(
+        state.add_named_metadata_checked(
             &self.name,
             NewHashFeedbackMetadata::with_capacity(self.capacity),
-        );
+        )?;
         Ok(())
     }
 }
diff --git a/libafl/src/observers/map/mod.rs b/libafl/src/observers/map/mod.rs
index 8ee79ac92c..5e76f9bb3b 100644
--- a/libafl/src/observers/map/mod.rs
+++ b/libafl/src/observers/map/mod.rs
@@ -68,7 +68,6 @@ pub use owned_map::*;
 /// // inform the feedback to track indices (required by IndexesLenTimeMinimizerScheduler), but not novelties
 /// // this *MUST* be done before it is passed to MaxMapFeedback!
 /// let edges_observer = edges_observer.track_indices();
-///
 /// // init the feedback
 /// let mut feedback = MaxMapFeedback::new(&edges_observer);
 /// #
@@ -80,9 +79,6 @@ pub use owned_map::*;
 /// #     &mut feedback,
 /// #     &mut (),
 /// # ).unwrap();
-///
-/// # feedback.init_state(&mut state).unwrap();
-///
 /// let scheduler = IndexesLenTimeMinimizerScheduler::new(&edges_observer, QueueScheduler::new());
 /// # scheduler.cull(&state).unwrap();
 /// ```