From 46b75747ef22e164f61cd334e1fb6da567985a22 Mon Sep 17 00:00:00 2001
From: Michael Rodler <contact@f0rki.at>
Date: Thu, 16 Feb 2023 17:29:33 +0100
Subject: [PATCH] Make sure input was loaded to avoid panic on unwrap in
 MutatedTransform (#1077)

* make sure input was loaded to avoid panic on unwrap
fixes issue #1059

* avoid unnecessary clone, avoid unnecessary branching

---------

Co-authored-by: Michael Rodler <m@mrodler.eu>
Co-authored-by: Addison Crump <addison.crump@cispa.de>
---
 libafl/src/inputs/generalized.rs |  2 +-
 libafl/src/stages/mutational.rs  | 10 +++++-----
 libafl/src/stages/power.rs       |  4 ++--
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/libafl/src/inputs/generalized.rs b/libafl/src/inputs/generalized.rs
index 03b30e40b4..5920e5ce49 100644
--- a/libafl/src/inputs/generalized.rs
+++ b/libafl/src/inputs/generalized.rs
@@ -108,7 +108,7 @@ where
     type Post = Self;
 
     fn try_transform_from(
-        base: &Testcase<BytesInput>,
+        base: &mut Testcase<BytesInput>,
         _state: &S,
         corpus_idx: CorpusId,
     ) -> Result<Self, Error> {
diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs
index d5120313e5..79628b66b8 100644
--- a/libafl/src/stages/mutational.rs
+++ b/libafl/src/stages/mutational.rs
@@ -51,7 +51,7 @@ where
 
     /// Transform the provided testcase into this type
     fn try_transform_from(
-        base: &Testcase<I>,
+        base: &mut Testcase<I>,
         state: &S,
         corpus_idx: CorpusId,
     ) -> Result<Self, Error>;
@@ -69,11 +69,11 @@ where
 
     #[inline]
     fn try_transform_from(
-        base: &Testcase<I>,
+        base: &mut Testcase<I>,
         _state: &S,
         _corpus_idx: CorpusId,
     ) -> Result<Self, Error> {
-        Ok(base.input().as_ref().unwrap().clone())
+        Ok(base.load_input()?.clone())
     }
 
     #[inline]
@@ -116,8 +116,8 @@ where
         let num = self.iterations(state, corpus_idx)?;
 
         start_timer!(state);
-        let testcase = state.corpus().get(corpus_idx)?.borrow();
-        let Ok(input) = I::try_transform_from(&testcase, state, corpus_idx) else { return Ok(()); };
+        let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut();
+        let Ok(input) = I::try_transform_from(&mut testcase, state, corpus_idx) else { return Ok(()); };
         drop(testcase);
         mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
 
diff --git a/libafl/src/stages/power.rs b/libafl/src/stages/power.rs
index c95382cd0a..78fba0efe5 100644
--- a/libafl/src/stages/power.rs
+++ b/libafl/src/stages/power.rs
@@ -82,8 +82,8 @@ where
     ) -> Result<(), Error> {
         let num = self.iterations(state, corpus_idx)?;
 
-        let testcase = state.corpus().get(corpus_idx)?.borrow();
-        let Ok(input) = I::try_transform_from(&testcase, state, corpus_idx) else { return Ok(()); };
+        let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut();
+        let Ok(input) = I::try_transform_from(&mut testcase, state, corpus_idx) else { return Ok(()); };
         drop(testcase);
 
         for i in 0..num {