From eb7c8a1174d0bef1275370a1f85093c6e292a6cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B6nke?= <eknoes@users.noreply.github.com>
Date: Sat, 27 Aug 2022 16:05:38 +0200
Subject: [PATCH] Add continous JSON Logging monitor (#738)

* Add simple JSON Monitor

* Add documentation

* Log global state

* Fix formatting

* Save state depending on closure outcome, have file opened all the time

* Make OnDiskJSONMonitor cloneable

* Switch to FnMut to allow stateful closures

* Use &mut M: Monitor for the closure
---
 libafl/src/monitors/disk.rs | 81 ++++++++++++++++++++++++++++++++++++-
 libafl/src/monitors/mod.rs  |  4 +-
 2 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/libafl/src/monitors/disk.rs b/libafl/src/monitors/disk.rs
index c7f795bc98..ad209abf4a 100644
--- a/libafl/src/monitors/disk.rs
+++ b/libafl/src/monitors/disk.rs
@@ -2,7 +2,13 @@
 
 use alloc::{string::String, vec::Vec};
 use core::time::Duration;
-use std::{fs::File, io::Write, path::PathBuf};
+use std::{
+    fs::{File, OpenOptions},
+    io::Write,
+    path::PathBuf,
+};
+
+use serde_json::json;
 
 use crate::{
     bolts::{current_time, format_duration_hms},
@@ -133,3 +139,76 @@ impl OnDiskTOMLMonitor<NopMonitor> {
         Self::new(filename, NopMonitor::new())
     }
 }
+
+#[derive(Debug, Clone)]
+/// Wraps a base monitor and continuously appends the current statistics to a JSON lines file.
+pub struct OnDiskJSONMonitor<F, M>
+where
+    F: FnMut(&mut M) -> bool,
+    M: Monitor,
+{
+    base: M,
+    path: PathBuf,
+    /// A function that has the current runtime as argument and decides, whether a record should be logged
+    log_record: F,
+}
+
+impl<F, M> OnDiskJSONMonitor<F, M>
+where
+    F: FnMut(&mut M) -> bool,
+    M: Monitor,
+{
+    /// Create a new [`OnDiskJSONMonitor`]
+    pub fn new<P>(filename: P, base: M, log_record: F) -> Self
+    where
+        P: Into<PathBuf>,
+    {
+        let path = filename.into();
+
+        Self {
+            base,
+            path,
+            log_record,
+        }
+    }
+}
+
+impl<F, M> Monitor for OnDiskJSONMonitor<F, M>
+where
+    F: FnMut(&mut M) -> bool,
+    M: Monitor,
+{
+    fn client_stats_mut(&mut self) -> &mut Vec<ClientStats> {
+        self.base.client_stats_mut()
+    }
+
+    fn client_stats(&self) -> &[ClientStats] {
+        self.base.client_stats()
+    }
+
+    fn start_time(&mut self) -> Duration {
+        self.base.start_time()
+    }
+
+    fn display(&mut self, event_msg: String, sender_id: u32) {
+        if (self.log_record)(&mut self.base) {
+            let file = OpenOptions::new()
+                .append(true)
+                .create(true)
+                .open(&self.path)
+                .expect("Failed to open logging file");
+
+            let line = json!({
+                "run_time": current_time() - self.base.start_time(),
+                "clients": self.base.client_stats().len(),
+                "corpus": self.base.corpus_size(),
+                "objectives": self.base.objective_size(),
+                "executions": self.base.total_execs(),
+                "exec_sec": self.base.execs_per_sec(),
+                "clients": &self.client_stats()[1..]
+            });
+            writeln!(&file, "{}", line).expect("Unable to write JSON to file");
+        }
+        self.base.display(event_msg, sender_id);
+    }
+}
diff --git a/libafl/src/monitors/mod.rs b/libafl/src/monitors/mod.rs
index 019e902aec..ec2a658cc0 100644
--- a/libafl/src/monitors/mod.rs
+++ b/libafl/src/monitors/mod.rs
@@ -15,7 +15,7 @@ use alloc::{fmt::Debug, string::String, vec::Vec};
 use core::{fmt, time::Duration};
 
 #[cfg(feature = "std")]
-pub use disk::OnDiskTOMLMonitor;
+pub use disk::{OnDiskJSONMonitor, OnDiskTOMLMonitor};
 use hashbrown::HashMap;
 use serde::{Deserialize, Serialize};
 
@@ -56,7 +56,7 @@ impl fmt::Display for UserStats {
 }
 
 /// A simple struct to keep track of client monitor
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone, Default, Serialize)]
 pub struct ClientStats {
     // monitor (maybe we need a separated struct?)
     /// The corpus size for this client