cargo-fuzztest: support replaying corpus for a certain time

- Add requires = "corpus_db" to replay_corpus_for in FuzzTestOptions.
- Support ExecutionMode::ReplayCorpus in CargoFuzzTestOptions.
- Pass FUZZTEST_REPLAY_CORPUS_FOR and FUZZTEST_TIME_BUDGET_TYPE in FuzztestRunner.
- Add unit, runner, and end-to-end integration tests for corpus replay.

PiperOrigin-RevId: 966059020
diff --git a/rust/cargo_fuzztest/src/lib.rs b/rust/cargo_fuzztest/src/lib.rs
index d3a9c4e..241bf1f 100644
--- a/rust/cargo_fuzztest/src/lib.rs
+++ b/rust/cargo_fuzztest/src/lib.rs
@@ -18,7 +18,8 @@
 use anyhow::{Context, Result};
 use clap::Parser;
 pub use fuzztest_options::{
-    ExecutionMode, FuzzFor, FuzzOptions, FuzzTestOptions, ReplayCrashOptions,
+    ExecutionMode, FuzzFor, FuzzOptions, FuzzTestOptions, ReplayCorpusOptions, ReplayCrashOptions,
+    TimeBudgetType,
 };
 use std::env;
 use std::ffi::OsString;
@@ -82,7 +83,9 @@
                 self.check_centipede_binary_path_is_set()?;
                 mode
             }
-            ExecutionMode::ReplayAllCrashes | ExecutionMode::ReplayCrash(_) => {
+            ExecutionMode::ReplayCorpus(_)
+            | ExecutionMode::ReplayAllCrashes
+            | ExecutionMode::ReplayCrash(_) => {
                 self.check_centipede_binary_path_is_set()?;
                 self.check_corpus_db_is_set()?;
                 mode
@@ -263,12 +266,21 @@
                 cmd.env("FUZZTEST_REPLAY_FINDINGS", "true");
             }
 
+            ExecutionMode::ReplayCorpus(replay_corpus_options) => {
+                cmd.env(
+                    "FUZZTEST_REPLAY_CORPUS_FOR",
+                    replay_corpus_options.replay_corpus_for.to_string(),
+                );
+                let time_budget_str = match replay_corpus_options.time_budget_type {
+                    TimeBudgetType::PerTest => "per-test",
+                    TimeBudgetType::Total => "total",
+                };
+                cmd.env("FUZZTEST_TIME_BUDGET_TYPE", time_budget_str);
+            }
+
             ExecutionMode::SmokeTest => {
                 // nothing to be done
             }
-            _ => {
-                // TODO(the-shank): add support for other modes.
-            }
         }
 
         if let Some(corpus_db) = &self.options.fuzztest_options.corpus_db {
@@ -567,4 +579,112 @@
             Some("/custom/centipede")
         );
     }
+
+    #[gtest]
+    fn test_cli_option_parsing_replay_corpus_for_success() {
+        let parsed = CargoFuzzTestOptions::try_parse_from([
+            "cargo-fuzztest",
+            "--replay-corpus-for",
+            "10s",
+            "--corpus-db",
+            "/tmp/corpus_db",
+            "--centipede-binary-path",
+            "/custom/centipede",
+        ])
+        .expect("valid replay-corpus-for options should parse successfully");
+
+        assert_eq!(parsed.fuzztest_options.replay_corpus_for, Some("10s".parse().unwrap()));
+        assert_eq!(parsed.fuzztest_options.time_budget_type, TimeBudgetType::PerTest);
+        assert_eq!(parsed.fuzztest_options.corpus_db.as_deref(), Some("/tmp/corpus_db"));
+
+        let mode = parsed.execution_mode().expect("valid execution mode");
+        assert_eq!(
+            mode,
+            ExecutionMode::ReplayCorpus(ReplayCorpusOptions {
+                replay_corpus_for: "10s".parse().unwrap(),
+                time_budget_type: TimeBudgetType::PerTest,
+                jobs: None,
+            })
+        );
+    }
+
+    #[gtest]
+    fn test_cli_option_parsing_replay_corpus_for_with_time_budget_type() {
+        let parsed = CargoFuzzTestOptions::try_parse_from([
+            "cargo-fuzztest",
+            "--replay-corpus-for",
+            "10s",
+            "--time-budget-type",
+            "total",
+            "--corpus-db",
+            "/tmp/corpus_db",
+            "--centipede-binary-path",
+            "/custom/centipede",
+        ])
+        .expect("valid options with time-budget-type should parse successfully");
+
+        assert_eq!(parsed.fuzztest_options.replay_corpus_for, Some("10s".parse().unwrap()));
+        assert_eq!(parsed.fuzztest_options.time_budget_type, TimeBudgetType::Total);
+
+        let mode = parsed.execution_mode().expect("valid execution mode");
+        assert_eq!(
+            mode,
+            ExecutionMode::ReplayCorpus(ReplayCorpusOptions {
+                replay_corpus_for: "10s".parse().unwrap(),
+                time_budget_type: TimeBudgetType::Total,
+                jobs: None,
+            })
+        );
+    }
+
+    #[gtest]
+    fn test_execution_mode_replay_corpus_missing_centipede_binary_path_errors() {
+        let options = CargoFuzzTestOptions {
+            fuzztest_options: FuzzTestOptions {
+                replay_corpus_for: Some("10s".parse().unwrap()),
+                corpus_db: Some("/tmp/corpus_db".to_string()),
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+        let err =
+            options.execution_mode().expect_err("missing centipede-binary-path should cause error");
+        assert!(err.to_string().contains("`--centipede-binary-path` needs to be specified"));
+    }
+
+    #[gtest]
+    fn test_build_run_command_replay_corpus() {
+        let options = CargoFuzzTestOptions {
+            fuzztest_options: FuzzTestOptions {
+                replay_corpus_for: Some("10s".parse().unwrap()),
+                time_budget_type: TimeBudgetType::Total,
+                corpus_db: Some("/tmp/corpus_db".to_string()),
+                ..Default::default()
+            },
+            centipede_binary_path: Some("/custom/centipede".to_string()),
+            ..Default::default()
+        };
+        let runner = FuzztestRunner::new("x86_64-unknown-linux-gnu".to_string(), options);
+        let cmd =
+            runner.build_run_command(Path::new("/tmp/test_bin")).expect("should build run command");
+
+        let envs: Vec<(String, Option<String>)> = cmd
+            .get_envs()
+            .map(|(k, v)| {
+                (k.to_string_lossy().to_string(), v.map(|s| s.to_string_lossy().to_string()))
+            })
+            .collect();
+
+        assert!(envs.contains(&("FUZZTEST_REPLAY_CORPUS_FOR".to_string(), Some("10s".to_string()))));
+        assert!(
+            envs.contains(&("FUZZTEST_TIME_BUDGET_TYPE".to_string(), Some("total".to_string())))
+        );
+        assert!(
+            envs.contains(&("FUZZTEST_CORPUS_DB".to_string(), Some("/tmp/corpus_db".to_string())))
+        );
+        assert!(envs.contains(&(
+            "FUZZTEST_CENTIPEDE_BINARY_PATH".to_string(),
+            Some("/custom/centipede".to_string())
+        )));
+    }
 }
diff --git a/rust/cargo_fuzztest/tests/e2e_cli_test.rs b/rust/cargo_fuzztest/tests/e2e_cli_test.rs
index 7541d5e..fa73f99 100644
--- a/rust/cargo_fuzztest/tests/e2e_cli_test.rs
+++ b/rust/cargo_fuzztest/tests/e2e_cli_test.rs
@@ -296,6 +296,115 @@
     );
 }
 
+#[gtest]
+fn test_cargo_fuzztest_e2e_replay_corpus() {
+    let centipede_bin = env::var("FUZZTEST_CENTIPEDE_BINARY_PATH")
+        .expect("FUZZTEST_CENTIPEDE_BINARY_PATH needs to be provided");
+
+    let sample_crate_path = get_sample_crate_path("sample_fuzz_crate");
+
+    let temp_target_dir = TempDir::new().expect("Failed to create temporary target directory");
+    fs::create_dir_all(&temp_target_dir).expect("Failed to create target directory");
+
+    let temp_db_dir = TempDir::new().expect("Failed to create temporary corpus db directory");
+    fs::create_dir_all(&temp_db_dir).expect("Failed to create corpus db directory");
+
+    let workdir_root_dir =
+        TempDir::new().expect("Failed to create temporary workdir_root directory");
+    fs::create_dir_all(&workdir_root_dir).expect("Failed to create workdir_root directory");
+
+    // 1. Run Centipede via cargo-fuzztest to fuzz and populate the corpus database.
+    let mut cmd = setup_cargo_fuzztest_command(&sample_crate_path, temp_target_dir.path());
+    cmd.arg("__fuzztest_mod__sample_fuzztest_target::sample_fuzztest_target")
+        .arg("--fuzz-for=3s")
+        .env_remove("FUZZTEST_CENTIPEDE_BINARY_PATH")
+        .arg("--centipede-binary-path")
+        .arg(&centipede_bin)
+        .arg("--corpus-db")
+        .arg(temp_db_dir.path())
+        .arg("--workdir-root")
+        .arg(workdir_root_dir.path());
+    let output = cmd.output().expect("Failed to run cargo-fuzztest to fuzz target");
+    assert!(output.status.success());
+
+    // 2. Run cargo-fuzztest CLI with --replay-corpus-for to verify it replays the corpus.
+    let mut cmd = setup_cargo_fuzztest_command(&sample_crate_path, temp_target_dir.path());
+    cmd.arg("__fuzztest_mod__sample_fuzztest_target::sample_fuzztest_target")
+        .arg("--replay-corpus-for=2s")
+        .arg("--corpus-db")
+        .arg(temp_db_dir.path())
+        .arg("--centipede-binary-path")
+        .arg(&centipede_bin);
+
+    let output = cmd.output().expect("Failed to run cargo-fuzztest in replay-corpus mode");
+    let stderr_str = String::from_utf8_lossy(&output.stderr);
+    let stdout_str = String::from_utf8_lossy(&output.stdout);
+
+    expect_true!(output.status.success());
+    expect_true!(
+        stderr_str.contains(
+            "Replaying __fuzztest_mod__sample_fuzztest_target.sample_fuzztest_target for 2s"
+        ) || stdout_str.contains(
+            "Replaying __fuzztest_mod__sample_fuzztest_target.sample_fuzztest_target for 2s"
+        )
+    );
+}
+
+#[gtest]
+fn test_cargo_fuzztest_e2e_replay_corpus_total_budget() {
+    let centipede_bin = env::var("FUZZTEST_CENTIPEDE_BINARY_PATH")
+        .expect("FUZZTEST_CENTIPEDE_BINARY_PATH needs to be provided");
+
+    let sample_crate_path = get_sample_crate_path("sample_fuzz_crate");
+
+    let temp_target_dir = TempDir::new().expect("Failed to create temporary target directory");
+    fs::create_dir_all(&temp_target_dir).expect("Failed to create target directory");
+
+    let temp_db_dir = TempDir::new().expect("Failed to create temporary corpus db directory");
+    fs::create_dir_all(&temp_db_dir).expect("Failed to create corpus db directory");
+
+    let workdir_root_dir =
+        TempDir::new().expect("Failed to create temporary workdir_root directory");
+    fs::create_dir_all(&workdir_root_dir).expect("Failed to create workdir_root directory");
+
+    // 1. Run Centipede via cargo-fuzztest to fuzz and populate the corpus database.
+    let mut cmd = setup_cargo_fuzztest_command(&sample_crate_path, temp_target_dir.path());
+    cmd.arg("__fuzztest_mod__sample_fuzztest_target::sample_fuzztest_target")
+        .arg("--fuzz-for=3s")
+        .env_remove("FUZZTEST_CENTIPEDE_BINARY_PATH")
+        .arg("--centipede-binary-path")
+        .arg(&centipede_bin)
+        .arg("--corpus-db")
+        .arg(temp_db_dir.path())
+        .arg("--workdir-root")
+        .arg(workdir_root_dir.path());
+    let output = cmd.output().expect("Failed to run cargo-fuzztest to fuzz target");
+    assert!(output.status.success());
+
+    // 2. Run cargo-fuzztest CLI with --replay-corpus-for and --time-budget-type total.
+    let mut cmd = setup_cargo_fuzztest_command(&sample_crate_path, temp_target_dir.path());
+    cmd.arg("__fuzztest_mod__sample_fuzztest_target::sample_fuzztest_target")
+        .arg("--replay-corpus-for=3s")
+        .arg("--time-budget-type=total")
+        .arg("--corpus-db")
+        .arg(temp_db_dir.path())
+        .arg("--centipede-binary-path")
+        .arg(&centipede_bin);
+
+    let output = cmd.output().expect("Failed to run cargo-fuzztest in replay-corpus mode");
+    let stderr_str = String::from_utf8_lossy(&output.stderr);
+    let stdout_str = String::from_utf8_lossy(&output.stdout);
+
+    expect_true!(output.status.success());
+    expect_true!(
+        stderr_str.contains(
+            "Replaying __fuzztest_mod__sample_fuzztest_target.sample_fuzztest_target for 1s"
+        ) || stdout_str.contains(
+            "Replaying __fuzztest_mod__sample_fuzztest_target.sample_fuzztest_target for 1s"
+        )
+    );
+}
+
 fn run_centipede_with_args_expect_termination(centipede_bin: &str, args: &[&str]) -> String {
     // Disable interference from Bazel environment variables.
     let env_diff = [
diff --git a/rust/cargo_fuzztest/tests/runner_test.rs b/rust/cargo_fuzztest/tests/runner_test.rs
index 8498c43..57aaba8 100644
--- a/rust/cargo_fuzztest/tests/runner_test.rs
+++ b/rust/cargo_fuzztest/tests/runner_test.rs
@@ -2,7 +2,7 @@
 
 use cargo_fuzztest::{CargoFuzzTestOptions, FuzztestRunner};
 use common::get_sample_test_bin_path;
-use fuzztest_options::{FuzzFor, FuzzTestOptions};
+use fuzztest_options::{FuzzFor, FuzzTestOptions, TimeBudgetType};
 use googletest::prelude::*;
 
 #[gtest]
@@ -252,6 +252,57 @@
 }
 
 #[gtest]
+fn test_runner_build_run_command_with_replay_corpus() {
+    let binary_path = get_sample_test_bin_path("sample_fuzz_crate");
+    let fuzztest_options = FuzzTestOptions {
+        replay_corpus_for: Some("10s".parse().expect("valid duration")),
+        time_budget_type: TimeBudgetType::Total,
+        corpus_db: Some("/custom/path/to/corpus_db".to_string()),
+        ..Default::default()
+    };
+    let options = CargoFuzzTestOptions {
+        fuzztest_options,
+        centipede_binary_path: Some("/custom/path/to/centipede".to_string()),
+        ..Default::default()
+    };
+    let runner = FuzztestRunner::new("x86_64-unknown-linux-gnu".to_string(), options);
+    let cmd = runner.build_run_command(&binary_path).expect("valid run command");
+
+    let envs: Vec<(String, Option<String>)> = cmd
+        .get_envs()
+        .map(|(k, v)| (k.to_string_lossy().to_string(), v.map(|s| s.to_string_lossy().to_string())))
+        .collect();
+    expect_true!(
+        envs.contains(&("FUZZTEST_REPLAY_CORPUS_FOR".to_string(), Some("10s".to_string())))
+    );
+    expect_true!(
+        envs.contains(&("FUZZTEST_TIME_BUDGET_TYPE".to_string(), Some("total".to_string())))
+    );
+    expect_true!(envs.contains(&(
+        "FUZZTEST_CORPUS_DB".to_string(),
+        Some("/custom/path/to/corpus_db".to_string())
+    )));
+    expect_true!(envs.contains(&(
+        "FUZZTEST_CENTIPEDE_BINARY_PATH".to_string(),
+        Some("/custom/path/to/centipede".to_string())
+    )));
+}
+
+#[gtest]
+fn test_execution_mode_replay_corpus_missing_centipede_binary_path_errors() {
+    let fuzztest_options = FuzzTestOptions {
+        replay_corpus_for: Some("10s".parse().expect("valid duration")),
+        corpus_db: Some("/custom/path/to/corpus_db".to_string()),
+        ..Default::default()
+    };
+    let options = CargoFuzzTestOptions { fuzztest_options, ..Default::default() };
+    let result = options.execution_mode();
+    expect_true!(result.is_err());
+    let err_msg = result.unwrap_err().to_string();
+    expect_true!(err_msg.contains("`--centipede-binary-path` needs to be specified"));
+}
+
+#[gtest]
 fn test_runner_list_command() {
     let binary_path = get_sample_test_bin_path("sample_fuzz_crate");
 
diff --git a/rust/options/src/lib.rs b/rust/options/src/lib.rs
index 94c1b58..d072cef 100644
--- a/rust/options/src/lib.rs
+++ b/rust/options/src/lib.rs
@@ -75,7 +75,7 @@
     pub replay_findings: bool,
 
     /// Replay the corpus for a specified duration.
-    #[arg(env = "FUZZTEST_REPLAY_CORPUS_FOR", long)]
+    #[arg(env = "FUZZTEST_REPLAY_CORPUS_FOR", long, requires = "corpus_db")]
     pub replay_corpus_for: Option<Duration>,
 
     /// Time budget calculation type for replay corpus mode.
@@ -283,12 +283,13 @@
         unsafe {
             std::env::set_var("FUZZTEST_JOBS", "4");
             std::env::set_var("FUZZTEST_REPLAY_CORPUS_FOR", "10s");
+            std::env::set_var("FUZZTEST_CORPUS_DB", "/tmp/corpus_db");
         }
 
         let options = FuzzTestOptions::parse_from(std::iter::empty::<OsString>());
 
         expect_that!(options.jobs, eq(Some(4)));
-        let expected_duration = "10s".parse().expect("valid duration");
+        let expected_duration = "10s".parse().expect("valid duration string");
         expect_that!(
             ExecutionMode::from_fuzztest_options(&options),
             eq(&ExecutionMode::ReplayCorpus(ReplayCorpusOptions {
@@ -302,6 +303,7 @@
         unsafe {
             std::env::remove_var("FUZZTEST_JOBS");
             std::env::remove_var("FUZZTEST_REPLAY_CORPUS_FOR");
+            std::env::remove_var("FUZZTEST_CORPUS_DB");
         }
     }
 
@@ -345,4 +347,70 @@
         expect_that!(options.replay_findings, eq(true));
         expect_that!(options.corpus_db.as_deref(), eq(Some("/tmp/corpus_db")));
     }
+
+    #[gtest]
+    fn test_replay_corpus_for_requires_corpus_db() {
+        // SAFETY: Testing environment parsing in single-threaded context.
+        unsafe {
+            std::env::set_var("FUZZTEST_REPLAY_CORPUS_FOR", "10s");
+            std::env::remove_var("FUZZTEST_CORPUS_DB");
+        }
+
+        let result = FuzzTestOptions::try_parse_from(std::iter::empty::<OsString>());
+
+        // SAFETY: Cleaning up environment variables.
+        unsafe {
+            std::env::remove_var("FUZZTEST_REPLAY_CORPUS_FOR");
+        }
+
+        let err = result.expect_err("parsing should fail when corpus_db is missing");
+        expect_that!(err.kind(), eq(clap::error::ErrorKind::MissingRequiredArgument));
+    }
+
+    #[gtest]
+    fn test_replay_corpus_for_with_corpus_db_succeeds() {
+        // SAFETY: Testing environment parsing in single-threaded context.
+        unsafe {
+            std::env::set_var("FUZZTEST_REPLAY_CORPUS_FOR", "10s");
+            std::env::set_var("FUZZTEST_CORPUS_DB", "/tmp/corpus_db");
+        }
+
+        let result = FuzzTestOptions::try_parse_from(std::iter::empty::<OsString>());
+
+        // SAFETY: Cleaning up environment variables.
+        unsafe {
+            std::env::remove_var("FUZZTEST_REPLAY_CORPUS_FOR");
+            std::env::remove_var("FUZZTEST_CORPUS_DB");
+        }
+
+        let options = result
+            .expect("parsing should succeed when both replay_corpus_for and corpus_db are present");
+        expect_that!(options.replay_corpus_for, eq(Some("10s".parse().unwrap())));
+        expect_that!(options.corpus_db.as_deref(), eq(Some("/tmp/corpus_db")));
+        expect_that!(options.time_budget_type, eq(TimeBudgetType::PerTest));
+    }
+
+    #[gtest]
+    fn test_replay_corpus_for_with_total_time_budget() {
+        // SAFETY: Testing environment parsing in single-threaded context.
+        unsafe {
+            std::env::set_var("FUZZTEST_REPLAY_CORPUS_FOR", "10s");
+            std::env::set_var("FUZZTEST_TIME_BUDGET_TYPE", "total");
+            std::env::set_var("FUZZTEST_CORPUS_DB", "/tmp/corpus_db");
+        }
+
+        let result = FuzzTestOptions::try_parse_from(std::iter::empty::<OsString>());
+
+        // SAFETY: Cleaning up environment variables.
+        unsafe {
+            std::env::remove_var("FUZZTEST_REPLAY_CORPUS_FOR");
+            std::env::remove_var("FUZZTEST_TIME_BUDGET_TYPE");
+            std::env::remove_var("FUZZTEST_CORPUS_DB");
+        }
+
+        let options = result.expect("parsing should succeed with total time budget type");
+        expect_that!(options.replay_corpus_for, eq(Some("10s".parse().unwrap())));
+        expect_that!(options.corpus_db.as_deref(), eq(Some("/tmp/corpus_db")));
+        expect_that!(options.time_budget_type, eq(TimeBudgetType::Total));
+    }
 }