96 lines
2.8 KiB
Rust
96 lines
2.8 KiB
Rust
use libafl::prelude::UsesInput;
|
|
use libafl_qemu::CPUArchState;
|
|
use libafl_qemu::Emulator;
|
|
use libafl_qemu::FastSnapshot;
|
|
use libafl_qemu::QemuExecutor;
|
|
use libafl_qemu::QemuHelper;
|
|
use libafl_qemu::QemuHelperTuple;
|
|
use libafl::{executors::ExitKind, inputs::Input, observers::ObserversTuple, state::HasMetadata};
|
|
use libafl_qemu::QemuHooks;
|
|
|
|
use libafl_qemu::{
|
|
emu,
|
|
};
|
|
// TODO be thread-safe maybe with https://amanieu.github.io/thread_local-rs/thread_local/index.html
|
|
#[derive(Debug)]
|
|
pub struct QemuStateRestoreHelper {
|
|
has_snapshot: bool,
|
|
use_snapshot: bool,
|
|
saved_cpu_states: Vec<CPUArchState>,
|
|
fastsnap: Option<FastSnapshot>
|
|
}
|
|
|
|
impl QemuStateRestoreHelper {
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
has_snapshot: false,
|
|
use_snapshot: true,
|
|
saved_cpu_states: vec![],
|
|
fastsnap: None
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for QemuStateRestoreHelper {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl<S> QemuHelper<S> for QemuStateRestoreHelper
|
|
where
|
|
S: UsesInput,
|
|
{
|
|
const HOOKS_DO_SIDE_EFFECTS: bool = true;
|
|
|
|
fn init_hooks<QT>(&self, _hooks: &QemuHooks<'_, QT, S>)
|
|
where
|
|
QT: QemuHelperTuple<S>,
|
|
{
|
|
}
|
|
|
|
fn first_exec<QT>(&self, _hooks: &QemuHooks<'_, QT, S>)
|
|
where
|
|
QT: QemuHelperTuple<S>,
|
|
{
|
|
}
|
|
|
|
fn post_exec<OT>(&mut self, emulator: &Emulator, _input: &S::Input, _observers: &mut OT, _exit_kind: &mut ExitKind) {
|
|
// unsafe { println!("snapshot post {}",emu::icount_get_raw()) };
|
|
}
|
|
|
|
fn pre_exec(&mut self, emulator: &Emulator, _input: &S::Input) {
|
|
// only restore in pre-exec, to preserve the post-execution state for inspection
|
|
#[cfg(feature = "snapshot_restore")]
|
|
{
|
|
#[cfg(feature = "snapshot_fast")]
|
|
match self.fastsnap {
|
|
Some(s) => emulator.restore_fast_snapshot(s),
|
|
None => {self.fastsnap = Some(emulator.create_fast_snapshot(true));},
|
|
}
|
|
#[cfg(not(feature = "snapshot_fast"))]
|
|
if !self.has_snapshot {
|
|
emulator.save_snapshot("Start", true);
|
|
self.has_snapshot = true;
|
|
}
|
|
else
|
|
{
|
|
emulator.load_snapshot("Start", true);
|
|
}
|
|
}
|
|
#[cfg(not(feature = "snapshot_restore"))]
|
|
if !self.has_snapshot {
|
|
self.saved_cpu_states = (0..emulator.num_cpus())
|
|
.map(|i| emulator.cpu_from_index(i).save_state())
|
|
.collect();
|
|
self.has_snapshot = true;
|
|
} else {
|
|
for (i, s) in self.saved_cpu_states.iter().enumerate() {
|
|
emulator.cpu_from_index(i).restore_state(s);
|
|
}
|
|
}
|
|
|
|
// unsafe { println!("snapshot pre {}",emu::icount_get_raw()) };
|
|
}
|
|
} |