From a659dd821c484c4e0524d0a2bf88023dae5e259b Mon Sep 17 00:00:00 2001
From: Vincent <space_white@yahoo.com>
Date: Thu, 23 Mar 2023 08:21:07 +0900
Subject: [PATCH] Fix testcase set_filename (#1092)

* fix set_filename

* use ? quantifier instead of expect

* fix clippy

* cargo fmt

* add rename old file to new file logic

* add cfg feature std

* add no_std set_filename

* fix create and remove lockfile logic

* fix cargo fmt

* remove unused import

* cargo fmt

* fix clippy

* fix lock filecondition

* remove useless import

---------

Co-authored-by: Dominik Maier <domenukk@gmail.com>
---
 libafl/src/corpus/inmemory_ondisk.rs | 10 +++---
 libafl/src/corpus/testcase.rs        | 50 +++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/libafl/src/corpus/inmemory_ondisk.rs b/libafl/src/corpus/inmemory_ondisk.rs
index d70fe2eb19..818a2d0ca8 100644
--- a/libafl/src/corpus/inmemory_ondisk.rs
+++ b/libafl/src/corpus/inmemory_ondisk.rs
@@ -221,7 +221,6 @@ where
             let mut ctr = 2;
             let filename = loop {
                 let lockfile = format!(".{file}.lafl_lock");
-                // try to create lockfile.
 
                 if OpenOptions::new()
                     .write(true)
@@ -229,15 +228,18 @@ where
                     .open(self.dir_path.join(lockfile))
                     .is_ok()
                 {
-                    break self.dir_path.join(file);
+                    break file;
                 }
 
                 file = format!("{file_orig}-{ctr}");
                 ctr += 1;
             };
 
-            let filename_str = filename.to_str().expect("Invalid Path");
-            testcase.set_filename(filename_str.into());
+            let file_path = self.dir_path.join(filename.clone());
+            let filename_str = file_path.to_str().expect("Invalid Path");
+            testcase.set_filename(filename_str.into())?;
+
+            fs::remove_file(format!(".{filename}.lafl_lock"))?;
         };
         if self.meta_format.is_some() {
             let mut filename = PathBuf::from(testcase.filename().as_ref().unwrap());
diff --git a/libafl/src/corpus/testcase.rs b/libafl/src/corpus/testcase.rs
index 16d62fbbfd..af2240beb2 100644
--- a/libafl/src/corpus/testcase.rs
+++ b/libafl/src/corpus/testcase.rs
@@ -9,6 +9,9 @@ use core::{
     time::Duration,
 };
 
+#[cfg(feature = "std")]
+use std::fs;
+
 use serde::{Deserialize, Serialize};
 
 use crate::{
@@ -144,8 +147,53 @@ where
 
     /// Set the filename
     #[inline]
-    pub fn set_filename(&mut self, filename: String) {
+    #[cfg(feature = "std")]
+    pub fn set_filename(&mut self, filename: String) -> Result<(), Error> {
+        use std::fs::OpenOptions;
+
+        if self.filename.is_some() {
+            let f = self.filename.clone().unwrap();
+            let old_filename = f.as_str();
+
+            let new_filename = filename.as_str();
+
+            // Do operations below when new filename is specified
+            if old_filename.eq(new_filename) {
+                return Ok(());
+            }
+
+            let new_lock_filename = format!(".{new_filename}.lafl_lock");
+
+            // Try to create lock file for new testcases
+            if OpenOptions::new()
+                .create(true)
+                .write(true)
+                .open(&new_lock_filename)
+                .is_err()
+            {
+                return Err(Error::illegal_state(
+                    "unable to create lock file for new testcase",
+                ));
+            }
+
+            fs::rename(old_filename, new_filename)?;
+
+            let old_metadata_filename = format!(".{old_filename}.metadata");
+            let new_metadata_filename = format!(".{new_filename}.metadata");
+            fs::rename(old_metadata_filename, new_metadata_filename)?;
+
+            fs::remove_file(&new_lock_filename)?;
+        }
+
         self.filename = Some(filename);
+        Ok(())
+    }
+
+    #[inline]
+    #[cfg(feature = "no_std")]
+    pub fn set_filename(&mut self, filename: String) -> Result<(), Error> {
+        self.filename = Some(filename);
+        Ok(())
     }
 
     /// Get the execution time of the testcase