From fc8056214beb03ebc2521090ca27097f71f34468 Mon Sep 17 00:00:00 2001
From: s1341 <s1341@users.noreply.github.com>
Date: Fri, 23 Apr 2021 18:25:00 +0300
Subject: [PATCH] Add selectable metadata format, including postcard and JSON
 (#59)

* Add selectable metadata format, including postcard and JSON

* Formatting

* nostd fix
---
 libafl/src/corpus/ondisk.rs | 35 +++++++++++++++++++++++++++--------
 libafl/src/lib.rs           |  8 ++++++++
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/libafl/src/corpus/ondisk.rs b/libafl/src/corpus/ondisk.rs
index e39b51da2b..a71d3d8f7c 100644
--- a/libafl/src/corpus/ondisk.rs
+++ b/libafl/src/corpus/ondisk.rs
@@ -9,6 +9,18 @@ use std::{fs, fs::File, io::Write, path::PathBuf};
 
 use crate::{corpus::Corpus, corpus::Testcase, inputs::Input, state::HasMetadata, Error};
 
+/// Options for the the format of the on-disk metadata
+#[cfg(feature = "std")]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum OnDiskMetadataFormat {
+    /// A binary-encoded postcard
+    Postcard,
+    /// JSON
+    Json,
+    /// JSON formatted for readability
+    JsonPretty,
+}
+
 /// A corpus able to store testcases to disk, and load them from disk, when they are being used.
 #[cfg(feature = "std")]
 #[derive(Default, Serialize, Deserialize, Clone, Debug)]
@@ -20,7 +32,7 @@ where
     entries: Vec<RefCell<Testcase<I>>>,
     current: Option<usize>,
     dir_path: PathBuf,
-    save_meta: bool,
+    meta_format: Option<OnDiskMetadataFormat>,
 }
 
 impl<I> Corpus<I> for OnDiskCorpus<I>
@@ -42,11 +54,15 @@ where
             let filename_str = filename.to_str().expect("Invalid Path");
             testcase.set_filename(filename_str.into());
         };
-        if self.save_meta {
+        if self.meta_format.is_some() {
             let filename = testcase.filename().as_ref().unwrap().to_owned() + ".metadata";
             let mut file = File::create(filename)?;
-            // TODO maybe use JSON
-            let serialized = postcard::to_allocvec(testcase.metadata())?;
+
+            let serialized = match self.meta_format.as_ref().unwrap() {
+                OnDiskMetadataFormat::Postcard => postcard::to_allocvec(testcase.metadata())?,
+                OnDiskMetadataFormat::Json => serde_json::to_vec(testcase.metadata())?,
+                OnDiskMetadataFormat::JsonPretty => serde_json::to_vec_pretty(testcase.metadata())?,
+            };
             file.write_all(&serialized)?;
         }
         testcase
@@ -107,19 +123,22 @@ where
             entries: vec![],
             current: None,
             dir_path,
-            save_meta: false,
+            meta_format: None,
         })
     }
 
-    /// Creates the OnDiskCorpus specifying if metatada must be saved to disk.
+    /// Creates the OnDiskCorpus specifying the type of metatada to be saved to disk.
     /// Will error, if `std::fs::create_dir_all` failed for `dir_path`.
-    pub fn new_save_meta(dir_path: PathBuf, save_meta: bool) -> Result<Self, Error> {
+    pub fn new_save_meta(
+        dir_path: PathBuf,
+        meta_format: Option<OnDiskMetadataFormat>,
+    ) -> Result<Self, Error> {
         fs::create_dir_all(&dir_path)?;
         Ok(Self {
             entries: vec![],
             current: None,
             dir_path,
-            save_meta,
+            meta_format,
         })
     }
 }
diff --git a/libafl/src/lib.rs b/libafl/src/lib.rs
index 0cc2676e27..98ca320fca 100644
--- a/libafl/src/lib.rs
+++ b/libafl/src/lib.rs
@@ -100,6 +100,14 @@ impl From<postcard::Error> for Error {
     }
 }
 
+/// Stringify the json serializer error
+#[cfg(feature = "std")]
+impl From<serde_json::Error> for Error {
+    fn from(err: serde_json::Error) -> Self {
+        Self::Serialize(format!("{:?}", err))
+    }
+}
+
 /// Create an AFL Error from io Error
 #[cfg(feature = "std")]
 impl From<io::Error> for Error {