From 71d367af30cfaafed85e43d28efa241bfad05039 Mon Sep 17 00:00:00 2001
From: Vincent <space_white@yahoo.com>
Date: Wed, 15 Feb 2023 20:38:25 +0900
Subject: [PATCH] TimeObserver: Use Instant::now instead of Duration (#1064)

* Use Instant::now instead of duration

* Use Some

* add custom serde for Instant

* fix linter

* only enable TimeFeedback when std flag is enabled

* fix typo

* fix linter std

* cargo fmt

* allow clippy::trivially_copy_pass_by_ref on custom serde serialize function

* allow TimeObserver and Timefeedback for no_std

* cargo fmt

---------

Co-authored-by: Andrea Fioraldi <andreafioraldi@gmail.com>
---
 libafl/src/observers/mod.rs | 69 +++++++++++++++++++++++++++++++++++--
 1 file changed, 66 insertions(+), 3 deletions(-)

diff --git a/libafl/src/observers/mod.rs b/libafl/src/observers/mod.rs
index b40109127b..31f535db97 100644
--- a/libafl/src/observers/mod.rs
+++ b/libafl/src/observers/mod.rs
@@ -30,12 +30,17 @@ use alloc::{
 };
 use core::{fmt::Debug, time::Duration};
 
+#[cfg(feature = "std")]
+use std::time::Instant;
+
 use serde::{Deserialize, Serialize};
 pub use value::*;
 
+#[cfg(feature = "no_std")]
+use crate::bolts::current_time;
+
 use crate::{
     bolts::{
-        current_time,
         ownedref::OwnedMutPtr,
         tuples::{MatchName, Named},
     },
@@ -411,17 +416,56 @@ where
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct TimeObserver {
     name: String,
+
+    #[cfg(feature = "std")]
+    #[serde(with = "instant_serializer")]
+    start_time: Instant,
+
+    #[cfg(feature = "no_std")]
     start_time: Duration,
+
     last_runtime: Option<Duration>,
 }
 
+#[cfg(feature = "std")]
+mod instant_serializer {
+    use core::time::Duration;
+    use std::time::Instant;
+
+    use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+    #[allow(clippy::trivially_copy_pass_by_ref)]
+    pub fn serialize<S>(instant: &Instant, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let duration = instant.elapsed();
+        duration.serialize(serializer)
+    }
+
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<Instant, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let duration = Duration::deserialize(deserializer)?;
+        let instant = Instant::now().checked_sub(duration).unwrap();
+        Ok(instant)
+    }
+}
+
 impl TimeObserver {
     /// Creates a new [`TimeObserver`] with the given name.
     #[must_use]
     pub fn new(name: &'static str) -> Self {
         Self {
             name: name.to_string(),
+
+            #[cfg(feature = "std")]
+            start_time: Instant::now(),
+
+            #[cfg(feature = "no_std")]
             start_time: Duration::from_secs(0),
+
             last_runtime: None,
         }
     }
@@ -437,12 +481,32 @@ impl<S> Observer<S> for TimeObserver
 where
     S: UsesInput,
 {
+    #[cfg(feature = "std")]
     fn pre_exec(&mut self, _state: &mut S, _input: &S::Input) -> Result<(), Error> {
         self.last_runtime = None;
-        self.start_time = current_time();
+        self.start_time = Instant::now();
         Ok(())
     }
 
+    #[cfg(feature = "no_std")]
+    fn pre_exec(&mut self, _state: &mut S, _input: &S::Input) -> Result<(), Error> {
+        self.last_runtime = None;
+        self.start_time = Duration::from_secs(0);
+        Ok(())
+    }
+
+    #[cfg(feature = "std")]
+    fn post_exec(
+        &mut self,
+        _state: &mut S,
+        _input: &S::Input,
+        _exit_kind: &ExitKind,
+    ) -> Result<(), Error> {
+        self.last_runtime = Some(self.start_time.elapsed());
+        Ok(())
+    }
+
+    #[cfg(feature = "no_std")]
     fn post_exec(
         &mut self,
         _state: &mut S,
@@ -450,7 +514,6 @@ where
         _exit_kind: &ExitKind,
     ) -> Result<(), Error> {
         self.last_runtime = current_time().checked_sub(self.start_time);
-        Ok(())
     }
 }