diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index da5a7cd60d..79c7145dda 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -96,6 +96,23 @@ jobs: run: sudo ./libafl_concolic/test/smoke_test_ubuntu_deps.sh - name: Run smoke test run: ./libafl_concolic/test/smoke_test.sh + bindings: + runs-on: ubuntu-latest + steps: + - uses: actions-rs/toolchain@v1 + with: + profile: minimal + toolchain: stable + - uses: Swatinem/rust-cache@v1 + - name: set mold linker as default linker + uses: rui314/setup-mold@v1 + - name: Install deps + run: sudo apt-get install -y llvm llvm-dev clang ninja-build python3-dev python3-pip python3-venv + - name: Install maturin + run: python3 -m pip install maturin + - uses: actions/checkout@v2 + - name: Run a maturin build + run: cd ./bindings/pylibafl && maturin build fuzzers: env: CC: ccache clang # use ccache in default diff --git a/bindings/pylibafl/Cargo.toml b/bindings/pylibafl/Cargo.toml index 394ecd473b..e41787c2bd 100644 --- a/bindings/pylibafl/Cargo.toml +++ b/bindings/pylibafl/Cargo.toml @@ -1,14 +1,13 @@ [package] name = "pylibafl" -version = "0.7.0" +version = "0.7.1" edition = "2021" [dependencies] pyo3 = { version = "0.15", features = ["extension-module"] } libafl_qemu = { path = "../../libafl_qemu", version = "0.7", features = ["python"] } libafl_sugar = { path = "../../libafl_sugar", version = "0.7", features = ["python"] } -#libafl = { path = "../../libafl", version = "0.7", features = ["python"] } -libafl = { path = "../../libafl", version = "0.7" } +libafl = { path = "../../libafl", version = "0.7", features = ["python"] } [build-dependencies] pyo3-build-config = { version = "0.15" } diff --git a/bindings/pylibafl/src/lib.rs b/bindings/pylibafl/src/lib.rs index af5b387294..4bc7c82116 100644 --- a/bindings/pylibafl/src/lib.rs +++ b/bindings/pylibafl/src/lib.rs @@ -1,22 +1,117 @@ -//use libafl; +use libafl; use libafl_qemu; use libafl_sugar; use pyo3::prelude::*; +use pyo3::types::PyDict; + +const LIBAFL_CODE: &str = r#" +class BaseObserver: + def flush(self): + pass + def pre_exec(self, state, input): + pass + def post_exec(self, state, input, exit_kind): + pass + def pre_exec_child(self, state, input): + pass + def post_exec_child(self, state, input, exit_kind): + pass + def name(self): + return type(self).__name__ + def as_observer(self): + return Observer.new_py(self) + +class BaseFeedback: + def init_state(self, state): + pass + def is_interesting(self, state, mgr, input, observers, exit_kind) -> bool: + return False + def append_metadata(self, state, testcase): + pass + def discard_metadata(self, state, input): + pass + def name(self): + return type(self).__name__ + def as_feedback(self): + return Feedback.new_py(self) + +class BaseExecutor: + def observers(self) -> ObserversTuple: + raise NotImplementedError('Implement this yourself') + def run_target(self, fuzzer, state, mgr, input) -> ExitKind: + raise NotImplementedError('Implement this yourself') + def as_executor(self): + return Executor.new_py(self) + +class BaseStage: + def perform(self, fuzzer, executor, state, manager, corpus_idx): + pass + def as_stage(self): + return Stage.new_py(self) + +class BaseMutator: + def mutate(self, state, input, stage_idx): + pass + def post_exec(self, state, stage_idx, corpus_idx): + pass + def as_mutator(self): + return Mutator.new_py(self) + +class FnStage(BaseStage): + def __init__(self, fn): + self.fn = fn + def __call__(self, fuzzer, executor, state, manager, corpus_idx): + self.fn(fuzzer, executor, state, manager, corpus_idx) + def perform(self, fuzzer, executor, state, manager, corpus_idx): + self.fn(fuzzer, executor, state, manager, corpus_idx) + +def feedback_not(a): + return NotFeedback(a).as_feedback() + +def feedback_and(a, b): + return EagerAndFeedback(a, b).as_feedback() + +def feedback_and_fast(a, b): + return FastAndFeedback(a, b).as_feedback() + +def feedback_or(a, b): + return EagerOrFeedback(a, b).as_feedback() + +def feedback_or_fast(a, b): + return FastOrFeedback(a, b).as_feedback() +"#; #[pymodule] #[pyo3(name = "pylibafl")] pub fn python_module(py: Python, m: &PyModule) -> PyResult<()> { + let modules = py.import("sys")?.getattr("modules")?; + let sugar_module = PyModule::new(py, "sugar")?; libafl_sugar::python_module(py, sugar_module)?; m.add_submodule(sugar_module)?; + modules.set_item("pylibafl.sugar", sugar_module)?; + let qemu_module = PyModule::new(py, "qemu")?; libafl_qemu::python_module(py, qemu_module)?; m.add_submodule(qemu_module)?; - //let libafl_module = PyModule::new(py, "libafl")?; - //libafl::python_module(py, libafl_module)?; - //m.add_submodule(libafl_module)?; + modules.set_item("pylibafl.qemu", qemu_module)?; + + let libafl_module = PyModule::new(py, "libafl")?; + libafl::pybind::python_module(py, libafl_module)?; + + libafl_module.add("__builtins__", py.import("builtins")?)?; + + let locals = PyDict::new(py); + py.run(LIBAFL_CODE, Some(libafl_module.dict()), Some(locals))?; + for (key, val) in locals.iter() { + libafl_module.add(key.extract::<&str>()?, val)?; + } + + m.add_submodule(libafl_module)?; + + modules.set_item("pylibafl.libafl", libafl_module)?; Ok(()) } diff --git a/bindings/pylibafl/test.py b/bindings/pylibafl/test.py new file mode 100644 index 0000000000..de84135200 --- /dev/null +++ b/bindings/pylibafl/test.py @@ -0,0 +1,74 @@ +from pylibafl.libafl import * +import ctypes + +class FooObserver(BaseObserver): + def __init__(self): + self.n = 0 + def name(self): + return "Foo" + def pre_exec(self, state, input): + if self.n % 10000 == 0: + print("FOO!", self.n, input) + self.n += 1 + +class FooFeedback(BaseFeedback): + def is_interesting(self, state, mgr, input, observers, exit_kind): + ob = observers.match_name("Foo").unwrap_py() + return ob.n % 10000 == 0 + +class FooExecutor(BaseExecutor): + def __init__(self, harness, observers: ObserversTuple): + self.h = harness + self.o = observers + def observers(self): + return self.o + def run_target(self, fuzzer, state, mgr, input) -> ExitKind: + return (self.h)(input) + +libc = ctypes.cdll.LoadLibrary("libc.so.6") + +area_ptr = libc.calloc(1, 4096) + +observer = StdMapObserverI8("mymap", area_ptr, 4096) + +m = observer.as_map_observer() + +observers = ObserversTuple([observer.as_map_observer().as_observer(), FooObserver().as_observer()]) + +feedback = feedback_or(MaxMapFeedbackI8(m).as_feedback(), FooFeedback().as_feedback()) + +objective = feedback_and_fast(CrashFeedback().as_feedback(), MaxMapFeedbackI8(m).as_feedback()) + +fuzzer = StdFuzzer(feedback, objective) + +rand = StdRand.with_current_nanos() + +state = StdState(rand.as_rand(), InMemoryCorpus().as_corpus(), InMemoryCorpus().as_corpus(), feedback, objective) + +monitor = SimpleMonitor(lambda s: print(s)) + +mgr = SimpleEventManager(monitor.as_monitor()) + +def harness(buf) -> ExitKind: + #print(buf) + m[0] = 1 + if len(buf) > 0 and buf[0] == ord('a'): + m[1] = 1 + if len(buf) > 1 and buf[1] == ord('b'): + m[2] = 1 + if len(buf) > 2 and buf[2] == ord('c'): + m[3] = 1 + return ExitKind.crash() + return ExitKind.ok() + +# executor = InProcessExecutor(harness, observers, fuzzer, state, mgr.as_manager()) + +executor = FooExecutor(harness, observers) + +stage = StdMutationalStage(StdHavocMutator().as_mutator()) + +stage_tuple_list = StagesTuple([stage.as_stage()]) + +fuzzer.add_input(state, executor.as_executor(), mgr.as_manager(), b'\0\0') + +fuzzer.fuzz_loop(executor.as_executor(), state, mgr.as_manager(), stage_tuple_list) diff --git a/fuzzers/libfuzzer_libpng/Makefile.toml b/fuzzers/libfuzzer_libpng/Makefile.toml index 4d9cfaaa3f..f95c932bf1 100644 --- a/fuzzers/libfuzzer_libpng/Makefile.toml +++ b/fuzzers/libfuzzer_libpng/Makefile.toml @@ -22,7 +22,7 @@ windows_alias = "unsupported" condition = { files_not_exist = ["./libpng-1.6.37"]} script_runner="@shell" script=''' -wget https://deac-fra.dl.sourceforge.net/project/libpng/libpng16/1.6.37/libpng-1.6.37.tar.xz +curl https://deac-fra.dl.sourceforge.net/project/libpng/libpng16/1.6.37/libpng-1.6.37.tar.xz --output libpng-1.6.37.tar.xz tar -xvf libpng-1.6.37.tar.xz ''' diff --git a/fuzzers/libfuzzer_libpng/src/lib.rs b/fuzzers/libfuzzer_libpng/src/lib.rs index 4f804a5c8a..ab9bceaf10 100644 --- a/fuzzers/libfuzzer_libpng/src/lib.rs +++ b/fuzzers/libfuzzer_libpng/src/lib.rs @@ -56,6 +56,7 @@ pub fn libafl_main() { } /// The actual fuzzer +#[cfg(not(test))] fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Result<(), Error> { // 'While the stats are state, they are usually used in the broker - which is likely never restarted let monitor = MultiMonitor::new(|s| println!("{}", s)); diff --git a/libafl/Cargo.toml b/libafl/Cargo.toml index 6c5b1e2c23..42502d2949 100644 --- a/libafl/Cargo.toml +++ b/libafl/Cargo.toml @@ -19,7 +19,7 @@ fork = [] # uses the fork() syscall to spawn children, instead of launching a ne rand_trait = ["rand_core"] # If set, libafl's rand implementations will implement `rand::Rng` introspection = [] # Include performance statistics of the fuzzing pipeline concolic_mutation = ["z3"] # include a simple concolic mutator based on z3 -#python = ["pyo3"] +python = ["pyo3"] tui_monitor = ["tui", "crossterm"] # enable TuiMonitor with crossterm cli = ["clap"] # expose bolts::cli qemu_cli = ["cli"] @@ -84,7 +84,8 @@ wait-timeout = { version = "0.2", optional = true } # used by CommandExecutor to z3 = { version = "0.11", features = ["static-link-z3"], optional = true } # for concolic mutation -pyo3 = { version = "0.15", optional = true } +pyo3 = { version = "0.15", optional = true, features = ["serde", "macros"] } +concat-idents = "1.1.3" # AGPL # !!! this create requires nightly diff --git a/libafl/src/bolts/os/unix_signals.rs b/libafl/src/bolts/os/unix_signals.rs index 7937920dcb..815fa25152 100644 --- a/libafl/src/bolts/os/unix_signals.rs +++ b/libafl/src/bolts/os/unix_signals.rs @@ -100,6 +100,14 @@ use crate::Error; pub use libc::{c_void, siginfo_t}; extern "C" { + /// The `libc` `getcontext` + /// + /// # Notes + /// + /// Somehow, `MacOS` on `aarch64` claims this type uses a 128 bit `int` + /// (@`rustc` 1.59.0) + /// Not much we can about it, for now. + #[allow(improper_ctypes)] fn getcontext(ucp: *mut ucontext_t) -> c_int; } diff --git a/libafl/src/bolts/rands.rs b/libafl/src/bolts/rands.rs index c8c726b651..c0474e6815 100644 --- a/libafl/src/bolts/rands.rs +++ b/libafl/src/bolts/rands.rs @@ -432,6 +432,7 @@ mod tests { } #[cfg(feature = "python")] +#[allow(missing_docs)] /// `Rand` Python bindings pub mod pybind { use super::Rand; @@ -444,7 +445,7 @@ pub mod pybind { /// Python class for StdRand pub struct PythonStdRand { /// Rust wrapped StdRand object - pub std_rand: StdRand, + pub inner: StdRand, } #[pymethods] @@ -452,52 +453,58 @@ pub mod pybind { #[staticmethod] fn with_current_nanos() -> Self { Self { - std_rand: StdRand::with_seed(current_nanos()), + inner: StdRand::with_seed(current_nanos()), } } #[staticmethod] fn with_seed(seed: u64) -> Self { Self { - std_rand: StdRand::with_seed(seed), + inner: StdRand::with_seed(seed), } } + + fn as_rand(slf: Py) -> PythonRand { + PythonRand::new_std(slf) + } } #[derive(Serialize, Deserialize, Debug, Clone)] enum PythonRandWrapper { - StdRand(PythonStdRand), + Std(Py), } /// Rand Trait binding #[pyclass(unsendable, name = "Rand")] #[derive(Serialize, Deserialize, Debug, Clone)] pub struct PythonRand { - rand: PythonRandWrapper, + wrapper: PythonRandWrapper, + } + + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonRandWrapper, { Std }) + }; } #[pymethods] impl PythonRand { #[staticmethod] - fn new_from_std(py_std_rand: PythonStdRand) -> Self { + fn new_std(py_std_rand: Py) -> Self { Self { - rand: PythonRandWrapper::StdRand(py_std_rand), + wrapper: PythonRandWrapper::Std(py_std_rand), } } } impl Rand for PythonRand { fn set_seed(&mut self, seed: u64) { - match &mut self.rand { - PythonRandWrapper::StdRand(py_std_rand) => py_std_rand.std_rand.set_seed(seed), - } + unwrap_me_mut!(self.wrapper, r, { r.set_seed(seed) }); } #[inline] fn next(&mut self) -> u64 { - match &mut self.rand { - PythonRandWrapper::StdRand(py_std_rand) => py_std_rand.std_rand.next(), - } + unwrap_me_mut!(self.wrapper, r, { r.next() }) } } diff --git a/libafl/src/bolts/tuples.rs b/libafl/src/bolts/tuples.rs index d232119a09..2dc59ed63a 100644 --- a/libafl/src/bolts/tuples.rs +++ b/libafl/src/bolts/tuples.rs @@ -9,9 +9,11 @@ use core::{ use xxhash_rust::xxh3::xxh3_64; -#[rustversion::nightly] +/// Returns if the type `T` is equal to `U` /// From -const fn type_eq() -> bool { +#[rustversion::nightly] +#[must_use] +pub const fn type_eq() -> bool { // Helper trait. `VALUE` is false, except for the specialization of the // case where `T == U`. trait TypeEq { @@ -31,8 +33,10 @@ const fn type_eq() -> bool { >::VALUE } +/// Returns if the type `T` is equal to `U` #[rustversion::not(nightly)] -const fn type_eq() -> bool { +#[must_use] +pub const fn type_eq() -> bool { // BEWARE! This is not unsafe, it is SUPER UNSAFE true } diff --git a/libafl/src/corpus/cached.rs b/libafl/src/corpus/cached.rs index df9fd9adc2..223b2205ba 100644 --- a/libafl/src/corpus/cached.rs +++ b/libafl/src/corpus/cached.rs @@ -137,19 +137,19 @@ where /// ``CachedOnDiskCorpus`` Python bindings #[cfg(feature = "python")] pub mod pybind { - use std::path::PathBuf; - + use crate::corpus::pybind::PythonCorpus; use crate::corpus::CachedOnDiskCorpus; use crate::inputs::BytesInput; use pyo3::prelude::*; use serde::{Deserialize, Serialize}; + use std::path::PathBuf; #[pyclass(unsendable, name = "CachedOnDiskCorpus")] #[derive(Serialize, Deserialize, Debug, Clone)] /// Python class for CachedOnDiskCorpus pub struct PythonCachedOnDiskCorpus { /// Rust wrapped CachedOnDiskCorpus object - pub cached_on_disk_corpus: CachedOnDiskCorpus, + pub inner: CachedOnDiskCorpus, } #[pymethods] @@ -157,10 +157,13 @@ pub mod pybind { #[new] fn new(path: String, cache_max_len: usize) -> Self { Self { - cached_on_disk_corpus: CachedOnDiskCorpus::new(PathBuf::from(path), cache_max_len) - .unwrap(), + inner: CachedOnDiskCorpus::new(PathBuf::from(path), cache_max_len).unwrap(), } } + + fn as_corpus(slf: Py) -> PythonCorpus { + PythonCorpus::new_cached_on_disk(slf) + } } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { diff --git a/libafl/src/corpus/inmemory.rs b/libafl/src/corpus/inmemory.rs index 937d415f0a..e3bbcf084a 100644 --- a/libafl/src/corpus/inmemory.rs +++ b/libafl/src/corpus/inmemory.rs @@ -91,6 +91,7 @@ where /// `InMemoryCorpus` Python bindings #[cfg(feature = "python")] pub mod pybind { + use crate::corpus::pybind::PythonCorpus; use crate::corpus::InMemoryCorpus; use crate::inputs::BytesInput; use pyo3::prelude::*; @@ -101,7 +102,7 @@ pub mod pybind { /// Python class for InMemoryCorpus pub struct PythonInMemoryCorpus { /// Rust wrapped InMemoryCorpus object - pub in_memory_corpus: InMemoryCorpus, + pub inner: InMemoryCorpus, } #[pymethods] @@ -109,9 +110,13 @@ pub mod pybind { #[new] fn new() -> Self { Self { - in_memory_corpus: InMemoryCorpus::new(), + inner: InMemoryCorpus::new(), } } + + fn as_corpus(slf: Py) -> PythonCorpus { + PythonCorpus::new_in_memory(slf) + } } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index bca5ed8a71..36a0a676a6 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -54,8 +54,10 @@ where /// `Corpus` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { use crate::corpus::inmemory::pybind::PythonInMemoryCorpus; + use crate::corpus::testcase::pybind::PythonTestcaseWrapper; use crate::corpus::{Corpus, Testcase}; use crate::inputs::BytesInput; use crate::Error; @@ -68,148 +70,138 @@ pub mod pybind { #[derive(Serialize, Deserialize, Debug, Clone)] enum PythonCorpusWrapper { - InMemory(PythonInMemoryCorpus), - CachedOnDisk(PythonCachedOnDiskCorpus), - OnDisk(PythonOnDiskCorpus), + InMemory(Py), + CachedOnDisk(Py), + OnDisk(Py), } /// Corpus Trait binding #[pyclass(unsendable, name = "Corpus")] + #[allow(clippy::unsafe_derive_deserialize)] #[derive(Serialize, Deserialize, Debug, Clone)] pub struct PythonCorpus { - corpus: PythonCorpusWrapper, + wrapper: PythonCorpusWrapper, + } + + macro_rules! unwrap_me { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_body!( + $wrapper, + $name, + $body, + PythonCorpusWrapper, + { + InMemory, + CachedOnDisk, + OnDisk + } + ) + }; + } + + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!( + $wrapper, + $name, + $body, + PythonCorpusWrapper, + { + InMemory, + CachedOnDisk, + OnDisk + } + ) + }; } #[pymethods] impl PythonCorpus { #[staticmethod] - fn new_from_in_memory(py_in_memory_corpus: PythonInMemoryCorpus) -> Self { + #[must_use] + pub fn new_in_memory(py_in_memory_corpus: Py) -> Self { Self { - corpus: PythonCorpusWrapper::InMemory(py_in_memory_corpus), + wrapper: PythonCorpusWrapper::InMemory(py_in_memory_corpus), } } #[staticmethod] - fn new_from_cached_on_disk(py_cached_on_disk_corpus: PythonCachedOnDiskCorpus) -> Self { + #[must_use] + pub fn new_cached_on_disk(py_cached_on_disk_corpus: Py) -> Self { Self { - corpus: PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus), + wrapper: PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus), } } #[staticmethod] - fn new_from_on_disk(py_on_disk_corpus: PythonOnDiskCorpus) -> Self { + #[must_use] + pub fn new_on_disk(py_on_disk_corpus: Py) -> Self { Self { - corpus: PythonCorpusWrapper::OnDisk(py_on_disk_corpus), + wrapper: PythonCorpusWrapper::OnDisk(py_on_disk_corpus), } } + + #[pyo3(name = "count")] + fn pycount(&self) -> usize { + self.count() + } + + #[pyo3(name = "current")] + fn pycurrent(&self) -> Option { + *self.current() + } + + #[pyo3(name = "get")] + fn pyget(&self, idx: usize) -> PythonTestcaseWrapper { + let t: &mut Testcase = unwrap_me!(self.wrapper, c, { + c.get(idx) + .map(|v| unsafe { v.as_ptr().as_mut().unwrap() }) + .expect("PythonCorpus::get failed") + }); + PythonTestcaseWrapper::wrap(t) + } } impl Corpus for PythonCorpus { #[inline] fn count(&self) -> usize { - match &self.corpus { - PythonCorpusWrapper::InMemory(py_in_memory_corpus) => { - py_in_memory_corpus.in_memory_corpus.count() - } - PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus) => { - py_cached_on_disk_corpus.cached_on_disk_corpus.count() - } - PythonCorpusWrapper::OnDisk(py_on_disk_corpus) => { - py_on_disk_corpus.on_disk_corpus.count() - } - } + unwrap_me!(self.wrapper, c, { c.count() }) } #[inline] fn add(&mut self, testcase: Testcase) -> Result { - match &mut self.corpus { - PythonCorpusWrapper::InMemory(py_in_memory_corpus) => { - py_in_memory_corpus.in_memory_corpus.add(testcase) - } - PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus) => { - py_cached_on_disk_corpus.cached_on_disk_corpus.add(testcase) - } - PythonCorpusWrapper::OnDisk(py_on_disk_corpus) => { - py_on_disk_corpus.on_disk_corpus.add(testcase) - } - } + unwrap_me_mut!(self.wrapper, c, { c.add(testcase) }) } #[inline] fn replace(&mut self, idx: usize, testcase: Testcase) -> Result<(), Error> { - match &mut self.corpus { - PythonCorpusWrapper::InMemory(py_in_memory_corpus) => { - py_in_memory_corpus.in_memory_corpus.replace(idx, testcase) - } - PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus) => { - py_cached_on_disk_corpus - .cached_on_disk_corpus - .replace(idx, testcase) - } - PythonCorpusWrapper::OnDisk(py_on_disk_corpus) => { - py_on_disk_corpus.on_disk_corpus.replace(idx, testcase) - } - } + unwrap_me_mut!(self.wrapper, c, { c.replace(idx, testcase) }) } #[inline] fn remove(&mut self, idx: usize) -> Result>, Error> { - match &mut self.corpus { - PythonCorpusWrapper::InMemory(py_in_memory_corpus) => { - py_in_memory_corpus.in_memory_corpus.remove(idx) - } - PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus) => { - py_cached_on_disk_corpus.cached_on_disk_corpus.remove(idx) - } - PythonCorpusWrapper::OnDisk(py_on_disk_corpus) => { - py_on_disk_corpus.on_disk_corpus.remove(idx) - } - } + unwrap_me_mut!(self.wrapper, c, { c.remove(idx) }) } #[inline] fn get(&self, idx: usize) -> Result<&RefCell>, Error> { - match &self.corpus { - PythonCorpusWrapper::InMemory(py_in_memory_corpus) => { - py_in_memory_corpus.in_memory_corpus.get(idx) - } - PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus) => { - py_cached_on_disk_corpus.cached_on_disk_corpus.get(idx) - } - PythonCorpusWrapper::OnDisk(py_on_disk_corpus) => { - py_on_disk_corpus.on_disk_corpus.get(idx) - } - } + let ptr = unwrap_me!(self.wrapper, c, { + c.get(idx) + .map(|v| v as *const RefCell>) + })?; + Ok(unsafe { ptr.as_ref().unwrap() }) } #[inline] fn current(&self) -> &Option { - match &self.corpus { - PythonCorpusWrapper::InMemory(py_in_memory_corpus) => { - py_in_memory_corpus.in_memory_corpus.current() - } - PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus) => { - py_cached_on_disk_corpus.cached_on_disk_corpus.current() - } - PythonCorpusWrapper::OnDisk(py_on_disk_corpus) => { - py_on_disk_corpus.on_disk_corpus.current() - } - } + let ptr = unwrap_me!(self.wrapper, c, { c.current() as *const Option }); + unsafe { ptr.as_ref().unwrap() } } #[inline] fn current_mut(&mut self) -> &mut Option { - match &mut self.corpus { - PythonCorpusWrapper::InMemory(py_in_memory_corpus) => { - py_in_memory_corpus.in_memory_corpus.current_mut() - } - PythonCorpusWrapper::CachedOnDisk(py_cached_on_disk_corpus) => { - py_cached_on_disk_corpus.cached_on_disk_corpus.current_mut() - } - PythonCorpusWrapper::OnDisk(py_on_disk_corpus) => { - py_on_disk_corpus.on_disk_corpus.current_mut() - } - } + let ptr = unwrap_me_mut!(self.wrapper, c, { c.current_mut() as *mut Option }); + unsafe { ptr.as_mut().unwrap() } } } diff --git a/libafl/src/corpus/ondisk.rs b/libafl/src/corpus/ondisk.rs index 90ed986c90..1b7a0921ee 100644 --- a/libafl/src/corpus/ondisk.rs +++ b/libafl/src/corpus/ondisk.rs @@ -208,19 +208,19 @@ where #[cfg(feature = "python")] /// `OnDiskCorpus` Python bindings pub mod pybind { - use std::path::PathBuf; - + use crate::corpus::pybind::PythonCorpus; use crate::corpus::OnDiskCorpus; use crate::inputs::BytesInput; use pyo3::prelude::*; use serde::{Deserialize, Serialize}; + use std::path::PathBuf; #[pyclass(unsendable, name = "OnDiskCorpus")] #[derive(Serialize, Deserialize, Debug, Clone)] /// Python class for OnDiskCorpus pub struct PythonOnDiskCorpus { /// Rust wrapped OnDiskCorpus object - pub on_disk_corpus: OnDiskCorpus, + pub inner: OnDiskCorpus, } #[pymethods] @@ -228,9 +228,13 @@ pub mod pybind { #[new] fn new(path: String) -> Self { Self { - on_disk_corpus: OnDiskCorpus::new(PathBuf::from(path)).unwrap(), + inner: OnDiskCorpus::new(PathBuf::from(path)).unwrap(), } } + + fn as_corpus(slf: Py) -> PythonCorpus { + PythonCorpus::new_on_disk(slf) + } } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { diff --git a/libafl/src/corpus/testcase.rs b/libafl/src/corpus/testcase.rs index 2d7a82feb5..994055668e 100644 --- a/libafl/src/corpus/testcase.rs +++ b/libafl/src/corpus/testcase.rs @@ -345,3 +345,98 @@ impl SchedulerTestcaseMetaData { } crate::impl_serdeany!(SchedulerTestcaseMetaData); + +#[cfg(feature = "python")] +#[allow(missing_docs)] +/// `Testcase` Python bindings +pub mod pybind { + use super::{HasMetadata, Testcase}; + use crate::bolts::ownedref::OwnedPtrMut; + use crate::inputs::{BytesInput, HasBytesVec}; + use crate::pybind::PythonMetadata; + use pyo3::prelude::*; + use pyo3::types::PyDict; + + /// `PythonTestcase` with fixed generics + pub type PythonTestcase = Testcase; + + #[pyclass(unsendable, name = "Testcase")] + #[derive(Debug)] + /// Python class for Testcase + pub struct PythonTestcaseWrapper { + /// Rust wrapped Testcase object + pub inner: OwnedPtrMut, + } + + impl PythonTestcaseWrapper { + pub fn wrap(r: &mut PythonTestcase) -> Self { + Self { + inner: OwnedPtrMut::Ptr(r), + } + } + + #[must_use] + pub fn unwrap(&self) -> &PythonTestcase { + self.inner.as_ref() + } + + pub fn unwrap_mut(&mut self) -> &mut PythonTestcase { + self.inner.as_mut() + } + } + + #[pymethods] + impl PythonTestcaseWrapper { + #[new] + fn new(input: Vec) -> Self { + Self { + inner: OwnedPtrMut::Owned(Box::new(PythonTestcase::new(BytesInput::new(input)))), + } + } + + fn load_input(&mut self) -> &[u8] { + self.inner + .as_mut() + .load_input() + .expect("Failed to load input") + .bytes() + } + + #[getter] + fn exec_time_ms(&self) -> Option { + self.inner.as_ref().exec_time().map(|t| t.as_millis()) + } + + #[getter] + fn executions(&self) -> usize { + *self.inner.as_ref().executions() + } + + #[getter] + fn fuzz_level(&self) -> usize { + self.inner.as_ref().fuzz_level() + } + + #[getter] + fn fuzzed(&self) -> bool { + self.inner.as_ref().fuzzed() + } + + fn metadata(&mut self) -> PyObject { + let meta = self.inner.as_mut().metadata_mut(); + if !meta.contains::() { + Python::with_gil(|py| { + let dict: Py = PyDict::new(py).into(); + meta.insert(PythonMetadata::new(dict.to_object(py))); + }); + } + meta.get::().unwrap().map.clone() + } + } + + /// Register the classes to the python module + pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; + Ok(()) + } +} diff --git a/libafl/src/events/mod.rs b/libafl/src/events/mod.rs index 1850a0560e..593c08f36a 100644 --- a/libafl/src/events/mod.rs +++ b/libafl/src/events/mod.rs @@ -543,208 +543,99 @@ mod tests { }; } } + /// `EventManager` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { use crate::events::simple::pybind::PythonSimpleEventManager; use crate::events::{ Event, EventFirer, EventManager, EventManagerId, EventProcessor, EventRestarter, HasEventManagerId, ProgressReporter, }; + use crate::executors::pybind::PythonExecutor; + use crate::fuzzer::pybind::PythonStdFuzzer; use crate::inputs::BytesInput; + use crate::state::pybind::PythonStdState; use crate::Error; use pyo3::prelude::*; - macro_rules! define_python_event_manager { - ($struct_name_trait:ident, $py_name_trait:tt, $wrapper_name: ident, $std_state_name: ident, $executor_name: ident, $my_std_fuzzer_type_name: ident) => { - use crate::executors::pybind::$executor_name; - use crate::pybind::$my_std_fuzzer_type_name; - use crate::state::pybind::$std_state_name; + #[derive(Debug, Clone)] + pub enum PythonEventManagerWrapper { + Simple(Py), + } - #[derive(Debug, Clone)] - enum $wrapper_name { - Simple(*mut PythonSimpleEventManager), - } + /// EventManager Trait binding + #[pyclass(unsendable, name = "EventManager")] + #[derive(Debug, Clone)] + pub struct PythonEventManager { + pub wrapper: PythonEventManagerWrapper, + } - /// EventManager Trait binding - #[pyclass(unsendable, name = $py_name_trait)] - #[derive(Debug, Clone)] - pub struct $struct_name_trait { - event_manager: $wrapper_name, - } - - impl $struct_name_trait { - fn get_event_manager( - &self, - ) -> &impl EventManager< - $executor_name, - BytesInput, - $std_state_name, - $my_std_fuzzer_type_name, - > { - unsafe { - match self.event_manager { - $wrapper_name::Simple(py_simple_event_manager) => { - &(*py_simple_event_manager).simple_event_manager - } - } - } - } - - fn get_mut_event_manager( - &mut self, - ) -> &mut impl EventManager< - $executor_name, - BytesInput, - $std_state_name, - $my_std_fuzzer_type_name, - > { - unsafe { - match self.event_manager { - $wrapper_name::Simple(py_simple_event_manager) => { - &mut (*py_simple_event_manager).simple_event_manager - } - } - } - } - } - - #[pymethods] - impl $struct_name_trait { - #[staticmethod] - fn new_from_simple(py_simple_event_manager: &mut PythonSimpleEventManager) -> Self { - Self { - event_manager: $wrapper_name::Simple(py_simple_event_manager), - } - } - } - - impl EventFirer for $struct_name_trait { - fn fire( - &mut self, - _state: &mut S, - event: Event, - ) -> Result<(), Error> { - self.get_mut_event_manager().fire(_state, event) - } - } - - impl EventRestarter for $struct_name_trait {} - - impl - EventProcessor< - $executor_name, - BytesInput, - $std_state_name, - $my_std_fuzzer_type_name, - > for $struct_name_trait - { - fn process( - &mut self, - _fuzzer: &mut $my_std_fuzzer_type_name, - state: &mut $std_state_name, - _executor: &mut $executor_name, - ) -> Result { - self.get_mut_event_manager() - .process(_fuzzer, state, _executor) - } - } - - impl ProgressReporter for $struct_name_trait {} - - impl HasEventManagerId for $struct_name_trait { - fn mgr_id(&self) -> EventManagerId { - self.get_event_manager().mgr_id() - } - } - - impl EventManager<$executor_name, BytesInput, $std_state_name, $my_std_fuzzer_type_name> - for $struct_name_trait - { - } + macro_rules! unwrap_me { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_body!($wrapper, $name, $body, PythonEventManagerWrapper, { + Simple + }) }; } - define_python_event_manager!( - PythonEventManagerI8, - "EventManagerI8", - PythonEventManagerWrapperI8, - MyStdStateI8, - PythonExecutorI8, - MyStdFuzzerI8 - ); + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonEventManagerWrapper, { + Simple + }) + }; + } - define_python_event_manager!( - PythonEventManagerI16, - "EventManagerI16", - PythonEventManagerWrapperI16, - MyStdStateI16, - PythonExecutorI16, - MyStdFuzzerI16 - ); + #[pymethods] + impl PythonEventManager { + #[staticmethod] + #[must_use] + pub fn new_simple(mgr: Py) -> Self { + Self { + wrapper: PythonEventManagerWrapper::Simple(mgr), + } + } + } - define_python_event_manager!( - PythonEventManagerI32, - "EventManagerI32", - PythonEventManagerWrapperI32, - MyStdStateI32, - PythonExecutorI32, - MyStdFuzzerI32 - ); + impl EventFirer for PythonEventManager { + fn fire(&mut self, state: &mut S, event: Event) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, e, { e.fire(state, event) }) + } + } - define_python_event_manager!( - PythonEventManagerI64, - "EventManagerI64", - PythonEventManagerWrapperI64, - MyStdStateI64, - PythonExecutorI64, - MyStdFuzzerI64 - ); + impl EventRestarter for PythonEventManager {} - define_python_event_manager!( - PythonEventManagerU8, - "EventManagerU8", - PythonEventManagerWrapperU8, - MyStdStateU8, - PythonExecutorU8, - MyStdFuzzerU8 - ); - define_python_event_manager!( - PythonEventManagerU16, - "EventManagerU16", - PythonEventManagerWrapperU16, - MyStdStateU16, - PythonExecutorU16, - MyStdFuzzerU16 - ); - define_python_event_manager!( - PythonEventManagerU32, - "EventManagerU32", - PythonEventManagerWrapperU32, - MyStdStateU32, - PythonExecutorU32, - MyStdFuzzerU32 - ); - define_python_event_manager!( - PythonEventManagerU64, - "EventManagerU64", - PythonEventManagerWrapperU64, - MyStdStateU64, - PythonExecutorU64, - MyStdFuzzerU64 - ); + impl EventProcessor + for PythonEventManager + { + fn process( + &mut self, + fuzzer: &mut PythonStdFuzzer, + state: &mut PythonStdState, + executor: &mut PythonExecutor, + ) -> Result { + unwrap_me_mut!(self.wrapper, e, { e.process(fuzzer, state, executor) }) + } + } + + impl ProgressReporter for PythonEventManager {} + + impl HasEventManagerId for PythonEventManager { + fn mgr_id(&self) -> EventManagerId { + unwrap_me!(self.wrapper, e, { e.mgr_id() }) + } + } + + impl EventManager + for PythonEventManager + { + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; Ok(()) } } diff --git a/libafl/src/events/simple.rs b/libafl/src/events/simple.rs index 0d1a9b895d..605cd19f3c 100644 --- a/libafl/src/events/simple.rs +++ b/libafl/src/events/simple.rs @@ -425,9 +425,12 @@ where Ok((state, mgr)) } } + /// `SimpleEventManager` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { + use crate::events::pybind::PythonEventManager; use crate::events::SimpleEventManager; use crate::inputs::BytesInput; use crate::monitors::pybind::PythonMonitor; @@ -438,7 +441,7 @@ pub mod pybind { /// Python class for SimpleEventManager pub struct PythonSimpleEventManager { /// Rust wrapped SimpleEventManager object - pub simple_event_manager: SimpleEventManager, + pub inner: SimpleEventManager, } #[pymethods] @@ -446,9 +449,13 @@ pub mod pybind { #[new] fn new(py_monitor: PythonMonitor) -> Self { Self { - simple_event_manager: SimpleEventManager::new(py_monitor), + inner: SimpleEventManager::new(py_monitor), } } + + fn as_manager(slf: Py) -> PythonEventManager { + PythonEventManager::new_simple(slf) + } } /// Register the classes to the python module diff --git a/libafl/src/executors/inprocess.rs b/libafl/src/executors/inprocess.rs index 3632eb98e4..1643110e00 100644 --- a/libafl/src/executors/inprocess.rs +++ b/libafl/src/executors/inprocess.rs @@ -1673,153 +1673,66 @@ mod tests { } #[cfg(feature = "python")] +#[allow(missing_docs)] /// `InProcess` Python bindings pub mod pybind { - use crate::bolts::tuples::tuple_list; + use crate::events::pybind::PythonEventManager; + use crate::executors::pybind::PythonExecutor; use crate::executors::{inprocess::OwnedInProcessExecutor, ExitKind}; + use crate::fuzzer::pybind::PythonStdFuzzerWrapper; use crate::inputs::{BytesInput, HasBytesVec}; + use crate::observers::pybind::PythonObserversTuple; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; use pyo3::prelude::*; use pyo3::types::PyBytes; - macro_rules! define_python_in_process_executor { - ($struct_name:ident, $py_name:tt, $my_std_state_type_name: ident, $std_state_name: ident, $event_manager_name: ident, $observer_name: ident, $std_fuzzer_name: ident) => { - use crate::events::pybind::$event_manager_name; - use crate::fuzzer::pybind::$std_fuzzer_name; - use crate::observers::map::pybind::$observer_name; - use crate::state::pybind::{$my_std_state_type_name, $std_state_name}; - - #[pyclass(unsendable, name = $py_name)] - #[derive(Debug)] - /// Python class for OwnedInProcessExecutor (i.e. InProcessExecutor with owned harness) - pub struct $struct_name { - /// Rust wrapped OwnedInProcessExecutor object - pub owned_in_process_executor: OwnedInProcessExecutor< - BytesInput, - ($observer_name, ()), - $my_std_state_type_name, - >, - } - - #[pymethods] - impl $struct_name { - #[new] - fn new( - harness: PyObject, - py_observer: $observer_name, - py_fuzzer: &mut $std_fuzzer_name, - py_state: &mut $std_state_name, - py_event_manager: &mut $event_manager_name, - ) -> Self { - Self { - owned_in_process_executor: OwnedInProcessExecutor::new( - Box::new(move |input: &BytesInput| { - Python::with_gil(|py| -> PyResult<()> { - let args = (PyBytes::new(py, input.bytes()),); - harness.call1(py, args)?; - Ok(()) - }) - .unwrap(); - ExitKind::Ok - }), - tuple_list!(py_observer), - &mut py_fuzzer.std_fuzzer, - &mut py_state.std_state, - py_event_manager, - ) - .expect("Failed to create the Executor".into()), - } - } - } - }; + #[pyclass(unsendable, name = "InProcessExecutor")] + #[derive(Debug)] + /// Python class for OwnedInProcessExecutor (i.e. InProcessExecutor with owned harness) + pub struct PythonOwnedInProcessExecutor { + /// Rust wrapped OwnedInProcessExecutor object + pub inner: OwnedInProcessExecutor, } - define_python_in_process_executor!( - PythonOwnedInProcessExecutor, - "OwnedInProcessExecutor", - PythonStdState, - PythonStdStateI8, - PythonEventManagerI8, - PythonMapObserverI8, - PythonStdFuzzerI8 - ); + #[pymethods] + impl PythonOwnedInProcessExecutor { + #[new] + fn new( + harness: PyObject, + py_observers: PythonObserversTuple, + py_fuzzer: &mut PythonStdFuzzerWrapper, + py_state: &mut PythonStdStateWrapper, + py_event_manager: &mut PythonEventManager, + ) -> Self { + Self { + inner: OwnedInProcessExecutor::new( + Box::new(move |input: &BytesInput| { + Python::with_gil(|py| -> PyResult<()> { + let args = (PyBytes::new(py, input.bytes()),); + harness.call1(py, args)?; + Ok(()) + }) + .unwrap(); + ExitKind::Ok + }), + py_observers, + py_fuzzer.unwrap_mut(), + py_state.unwrap_mut(), + py_event_manager, + ) + .expect("Failed to create the Executor"), + } + } - define_python_in_process_executor!( - PythonOwnedInProcessExecutorI16, - "OwnedInProcessExecutorI16", - MyStdStateI16, - PythonStdStateI16, - PythonEventManagerI16, - PythonMapObserverI16, - PythonStdFuzzerI16 - ); - define_python_in_process_executor!( - PythonOwnedInProcessExecutorI32, - "OwnedInProcessExecutorI32", - MyStdStateI32, - PythonStdStateI32, - PythonEventManagerI32, - PythonMapObserverI32, - PythonStdFuzzerI32 - ); - define_python_in_process_executor!( - PythonOwnedInProcessExecutorI64, - "OwnedInProcessExecutorI64", - MyStdStateI64, - PythonStdStateI64, - PythonEventManagerI64, - PythonMapObserverI64, - PythonStdFuzzerI64 - ); - - define_python_in_process_executor!( - PythonOwnedInProcessExecutorU8, - "OwnedInProcessExecutorU8", - MyStdStateU8, - PythonStdStateU8, - PythonEventManagerU8, - PythonMapObserverU8, - PythonStdFuzzerU8 - ); - - define_python_in_process_executor!( - PythonOwnedInProcessExecutorU16, - "OwnedInProcessExecutorU16", - MyStdStateU16, - PythonStdStateU16, - PythonEventManagerU16, - PythonMapObserverU16, - PythonStdFuzzerU16 - ); - define_python_in_process_executor!( - PythonOwnedInProcessExecutorU32, - "OwnedInProcessExecutorU32", - MyStdStateU32, - PythonStdStateU32, - PythonEventManagerU32, - PythonMapObserverU32, - PythonStdFuzzerU32 - ); - define_python_in_process_executor!( - PythonOwnedInProcessExecutorU64, - "OwnedInProcessExecutorU64", - MyStdStateU64, - PythonStdStateU64, - PythonEventManagerU64, - PythonMapObserverU64, - PythonStdFuzzerU64 - ); + #[must_use] + pub fn as_executor(slf: Py) -> PythonExecutor { + PythonExecutor::new_inprocess(slf) + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; Ok(()) } } diff --git a/libafl/src/executors/mod.rs b/libafl/src/executors/mod.rs index 3ff9d6b9d2..b04d80e78a 100644 --- a/libafl/src/executors/mod.rs +++ b/libafl/src/executors/mod.rs @@ -186,127 +186,253 @@ mod test { } #[cfg(feature = "python")] +#[allow(missing_docs)] /// `Executor` Python bindings pub mod pybind { + use crate::events::pybind::PythonEventManager; + use crate::executors::inprocess::pybind::PythonOwnedInProcessExecutor; use crate::executors::{Executor, ExitKind, HasObservers}; - use crate::inputs::BytesInput; + use crate::fuzzer::pybind::{PythonStdFuzzer, PythonStdFuzzerWrapper}; + use crate::inputs::{BytesInput, HasBytesVec}; + use crate::observers::pybind::PythonObserversTuple; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; use crate::Error; use pyo3::prelude::*; + use serde::{Deserialize, Serialize}; - macro_rules! define_python_executor { - ($struct_name_trait:ident, $py_name_trait:tt, $wrapper_name: ident, $my_std_state_type_name: ident, $my_std_fuzzer_type_name: ident, - $event_manager_name: ident, $in_process_executor_name: ident, $observer_name: ident) => { - use crate::events::pybind::$event_manager_name; - use crate::executors::inprocess::pybind::$in_process_executor_name; - use crate::fuzzer::pybind::$my_std_fuzzer_type_name; - use crate::observers::pybind::$observer_name; - use crate::state::pybind::$my_std_state_type_name; + #[pyclass(unsendable, name = "ExitKind")] + #[derive(Clone, Debug, Serialize, Deserialize)] + pub struct PythonExitKind { + pub inner: ExitKind, + } - #[derive(Debug)] - enum $wrapper_name { - OwnedInProcess(*mut $in_process_executor_name), + impl From for PythonExitKind { + fn from(inner: ExitKind) -> Self { + Self { inner } + } + } + + #[pymethods] + impl PythonExitKind { + fn __eq__(&self, other: &PythonExitKind) -> bool { + self.inner == other.inner + } + + #[must_use] + fn is_ok(&self) -> bool { + self.inner == ExitKind::Ok + } + + #[must_use] + fn is_crash(&self) -> bool { + self.inner == ExitKind::Crash + } + + #[must_use] + fn is_oom(&self) -> bool { + self.inner == ExitKind::Oom + } + + #[must_use] + fn is_timeout(&self) -> bool { + self.inner == ExitKind::Timeout + } + + #[staticmethod] + #[must_use] + fn ok() -> Self { + Self { + inner: ExitKind::Ok, } + } - #[pyclass(unsendable, name = $py_name_trait)] - #[derive(Debug)] - /// Executor + HasObservers Trait binding - pub struct $struct_name_trait { - wrapper: $wrapper_name, + #[staticmethod] + #[must_use] + fn crash() -> Self { + Self { + inner: ExitKind::Crash, } + } - impl $struct_name_trait { - fn unwrap( - &self, - ) -> &(impl Executor< - $event_manager_name, - BytesInput, - $my_std_state_type_name, - $my_std_fuzzer_type_name, - > + HasObservers) { - unsafe { - match self.wrapper { - $wrapper_name::OwnedInProcess(py_wrapper) => &(*py_wrapper).upcast(), - } + #[staticmethod] + #[must_use] + fn oom() -> Self { + Self { + inner: ExitKind::Oom, + } + } + + #[staticmethod] + #[must_use] + fn timeout() -> Self { + Self { + inner: ExitKind::Timeout, + } + } + } + + #[derive(Clone, Debug)] + pub struct PyObjectExecutor { + inner: PyObject, + tuple: PythonObserversTuple, + } + + impl PyObjectExecutor { + #[must_use] + pub fn new(obj: PyObject) -> Self { + let tuple = Python::with_gil(|py| -> PyResult { + obj.call_method1(py, "observers", ())?.extract(py) + }) + .unwrap(); + PyObjectExecutor { inner: obj, tuple } + } + } + + impl HasObservers for PyObjectExecutor { + #[inline] + fn observers(&self) -> &PythonObserversTuple { + &self.tuple + } + + #[inline] + fn observers_mut(&mut self) -> &mut PythonObserversTuple { + &mut self.tuple + } + } + + impl Executor + for PyObjectExecutor + { + #[inline] + fn run_target( + &mut self, + fuzzer: &mut PythonStdFuzzer, + state: &mut PythonStdState, + mgr: &mut PythonEventManager, + input: &BytesInput, + ) -> Result { + let ek = Python::with_gil(|py| -> PyResult<_> { + let ek: PythonExitKind = self + .inner + .call_method1( + py, + "run_target", + ( + PythonStdFuzzerWrapper::wrap(fuzzer), + PythonStdStateWrapper::wrap(state), + mgr.clone(), + input.bytes(), + ), + )? + .extract(py)?; + Ok(ek) + })?; + Ok(ek.inner) + } + } + + #[derive(Clone, Debug)] + enum PythonExecutorWrapper { + InProcess(Py), + Python(PyObjectExecutor), + } + + #[pyclass(unsendable, name = "Executor")] + #[derive(Clone, Debug)] + /// Executor + HasObservers Trait binding + pub struct PythonExecutor { + wrapper: PythonExecutorWrapper, + } + + macro_rules! unwrap_me { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_body!($wrapper, $name, $body, PythonExecutorWrapper, + { InProcess }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body } } - - fn unwrap_mut( - &mut self, - ) -> &mut (impl Executor< - $event_manager_name, - BytesInput, - $my_std_state_type_name, - $my_std_fuzzer_type_name, - > + HasObservers) { - unsafe { - match self.wrapper { - $wrapper_name::OwnedInProcess(py_wrapper) => { - &mut (*py_wrapper).upcast_mut() - } - } - } - } - } - - #[pymethods] - impl $struct_name_trait { - #[staticmethod] - fn new_from_inprocess( - owned_inprocess_executor: &mut $in_process_executor_name, - ) -> Self { - Self { - wrapper: $wrapper_name::OwnedInProcess(owned_inprocess_executor), - } - } - } - - impl HasObservers for $struct_name_trait { - // #[inline] - fn observers(&self) -> &($observer_name, ()) { - self.unwrap().observers() - } - - #[inline] - fn observers_mut(&mut self) -> &mut ($observer_name, ()) { - self.unwrap_mut().observers_mut() - } - } - - impl - Executor< - $event_manager_name, - BytesInput, - $my_std_state_type_name, - $my_std_fuzzer_type_name, - > for $struct_name_trait - { - #[inline] - fn run_target( - &mut self, - fuzzer: &mut $my_std_fuzzer_type_name, - state: &mut $my_std_state_type_name, - mgr: &mut $event_manager_name, - input: &BytesInput, - ) -> Result { - self.unwrap_mut().run_target(fuzzer, state, mgr, input) - } - } + ) }; } - define_python_executor!( - PythonExecutor, - "Executor", - PythonExecutorWrapper, - PythonStdState, - PythonStdFuzzer, - PythonEventManager, - PythonOwnedInProcessExecutor, - PythonObserver - ); + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonExecutorWrapper, + { InProcess }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body + } + } + ) + }; + } + + #[pymethods] + impl PythonExecutor { + #[staticmethod] + #[must_use] + pub fn new_inprocess(owned_inprocess_executor: Py) -> Self { + Self { + wrapper: PythonExecutorWrapper::InProcess(owned_inprocess_executor), + } + } + + #[staticmethod] + #[must_use] + pub fn new_py(obj: PyObject) -> Self { + Self { + wrapper: PythonExecutorWrapper::Python(PyObjectExecutor::new(obj)), + } + } + + #[must_use] + pub fn unwrap_py(&self) -> Option { + match &self.wrapper { + PythonExecutorWrapper::Python(pyo) => Some(pyo.inner.clone()), + PythonExecutorWrapper::InProcess(_) => None, + } + } + } + + impl HasObservers for PythonExecutor { + #[inline] + fn observers(&self) -> &PythonObserversTuple { + let ptr = unwrap_me!(self.wrapper, e, { + e.observers() as *const PythonObserversTuple + }); + unsafe { ptr.as_ref().unwrap() } + } + + #[inline] + fn observers_mut(&mut self) -> &mut PythonObserversTuple { + let ptr = unwrap_me_mut!(self.wrapper, e, { + e.observers_mut() as *mut PythonObserversTuple + }); + unsafe { ptr.as_mut().unwrap() } + } + } + + impl Executor for PythonExecutor { + #[inline] + fn run_target( + &mut self, + fuzzer: &mut PythonStdFuzzer, + state: &mut PythonStdState, + mgr: &mut PythonEventManager, + input: &BytesInput, + ) -> Result { + unwrap_me_mut!(self.wrapper, e, { e.run_target(fuzzer, state, mgr, input) }) + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/libafl/src/feedbacks/map.rs b/libafl/src/feedbacks/map.rs index 58ef5585f3..47328db9e7 100644 --- a/libafl/src/feedbacks/map.rs +++ b/libafl/src/feedbacks/map.rs @@ -666,451 +666,117 @@ mod tests { /// `MapFeedback` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { - use crate::bolts::{tuples::Named, AsMutIterator, AsRefIterator, HasLen}; - use crate::observers::{map::OwnedMapObserver, MapObserver, Observer}; - use crate::Error; + use super::{Debug, HasObserverName, MaxMapFeedback}; + use crate::feedbacks::pybind::PythonFeedback; + use crate::inputs::BytesInput; + use crate::state::pybind::PythonStdState; + use concat_idents::concat_idents; use pyo3::prelude::*; - use serde::{Deserialize, Serialize}; - use std::slice::{Iter, IterMut}; - macro_rules! define_python_map_observer { - ($struct_name:ident, $py_name:tt, $struct_name_trait:ident, $py_name_trait:tt, $datatype:ty, $wrapper_name: ident) => { + macro_rules! define_python_map_feedback { + ($struct_name:ident, $py_name:tt, $datatype:ty, $map_observer_type_name: ident, $my_std_state_type_name: ident) => { + use crate::observers::map::pybind::$map_observer_type_name; + #[pyclass(unsendable, name = $py_name)] - #[derive(Serialize, Deserialize, Debug, Clone)] - /// Python class for OwnedMapObserver (i.e. StdMapObserver with owned map) + #[derive(Debug, Clone)] + /// Python class for MaxMapFeedback pub struct $struct_name { - /// Rust wrapped OwnedMapObserver object - pub inner: OwnedMapObserver<$datatype>, + /// Rust wrapped MaxMapFeedback object + pub inner: MaxMapFeedback< + BytesInput, + $map_observer_type_name, /* PythonMapObserverI8 */ + $my_std_state_type_name, + $datatype, + >, } #[pymethods] impl $struct_name { #[new] - fn new(name: String, map: Vec<$datatype>) -> Self { + fn new(observer: &$map_observer_type_name) -> Self { Self { - //TODO: Not leak memory - inner: OwnedMapObserver::new(Box::leak(name.into_boxed_str()), map), + inner: MaxMapFeedback::new(observer), } } + + #[must_use] + pub fn as_feedback(slf: Py) -> PythonFeedback { + concat_idents!(func = new_max_map_,$datatype { + PythonFeedback::func(slf) + }) + } } - #[derive(Serialize, Deserialize, Debug, Clone)] - enum $wrapper_name { - Owned($struct_name), - } - - // Should not be exposed to user - #[pyclass(unsendable, name = $py_name_trait)] - #[derive(Serialize, Deserialize, Debug, Clone)] - /// MapObserver + Observer Trait binding - pub struct $struct_name_trait { - pub wrapper: $wrapper_name, - } - - impl $struct_name_trait { - fn unwrap(&self) -> &impl MapObserver { - unsafe { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => &py_wrapper.inner, - } - } - } - - fn unwrap_mut(&mut self) -> &mut impl MapObserver { - unsafe { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => &mut py_wrapper.inner, - } - } - } - - fn upcast(&self) -> &impl Observer { - unsafe { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => &py_wrapper.inner, - } - } - } - - fn upcast_mut(&mut self) -> &mut impl Observer { - unsafe { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => &mut py_wrapper.inner, - } - } - } - } - - #[pymethods] - impl $struct_name_trait { - #[staticmethod] - fn new_from_owned(owned_map: $struct_name) -> Self { - Self { - wrapper: $wrapper_name::Owned(owned_map), - } - } - } - - impl<'it> AsRefIterator<'it> for $struct_name_trait { - type Item = $datatype; - type IntoIter = Iter<'it, $datatype>; - - fn as_ref_iter(&'it self) -> Self::IntoIter { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.as_ref_iter(), - } - } - } - - impl<'it> AsMutIterator<'it> for $struct_name_trait { - type Item = $datatype; - type IntoIter = IterMut<'it, $datatype>; - - fn as_mut_iter(&'it mut self) -> Self::IntoIter { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.as_mut_iter(), - } - } - } - - impl MapObserver for $struct_name_trait { - type Entry = $datatype; - - #[inline] - fn get(&self, idx: usize) -> &$datatype { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => &py_wrapper.inner.get(idx), - } - } - - #[inline] - fn get_mut(&mut self, idx: usize) -> &mut $datatype { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.get_mut(idx), - } - } - - #[inline] - fn usable_count(&self) -> usize { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.wrapper.usable_count(), - } - } - - fn hash(&self) -> u64 { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.hash(), - } - } - - #[inline] - fn initial(&self) -> $datatype { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.initial(), - } - } - - #[inline] - fn initial_mut(&mut self) -> &mut $datatype { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.initial_mut(), - } - } - - #[inline] - fn set_initial(&mut self, initial: $datatype) { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => { - py_wrapper.inner.set_initial(initial); - } - } - } - - fn to_vec(&self) -> Vec<$datatype> { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.to_vec(), - } - } - } - - impl Named for $struct_name_trait { - #[inline] - fn name(&self) -> &str { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.name(), - } - } - } - - impl HasLen for $struct_name_trait { - #[inline] - fn len(&self) -> usize { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.len(), - } - } - } - - impl Observer for $struct_name_trait - where - Self: MapObserver, - { - #[inline] - fn pre_exec(&mut self, _state: &mut S, _input: &I) -> Result<(), Error> { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => { - py_wrapper.inner.pre_exec(_state, _input) - } - } - } - } - }; - } - - define_python_map_observer!( - PythonOwnedMapObserverI8, - "OwnedMapObserverI8", - PythonMapObserverI8, - "MapObserverI8", - i8, - PythonMapObserverWrapperI8 - ); - define_python_map_observer!( - PythonOwnedMapObserverI16, - "OwnedMapObserverI16", - PythonMapObserverI16, - "MapObserverI16", - i16, - PythonMapObserverWrapperI16 - ); - define_python_map_observer!( - PythonOwnedMapObserverI32, - "OwnedMapObserverI32", - PythonMapObserverI32, - "MapObserverI32", - i32, - PythonMapObserverWrapperI32 - ); - define_python_map_observer!( - PythonOwnedMapObserverI64, - "OwnedMapObserverI64", - PythonMapObserverI64, - "MapObserverI64", - i64, - PythonMapObserverWrapperI64 - ); - - define_python_map_observer!( - PythonOwnedMapObserverU8, - "OwnedMapObserverU8", - PythonMapObserverU8, - "MapObserverU8", - u8, - PythonMapObserverWrapperU8 - ); - define_python_map_observer!( - PythonOwnedMapObserverU16, - "OwnedMapObserverU16", - PythonMapObserverU16, - "MapObserverU16", - u16, - PythonMapObserverWrapperU16 - ); - define_python_map_observer!( - PythonOwnedMapObserverU32, - "OwnedMapObserverU32", - PythonMapObserverU32, - "MapObserverU32", - u32, - PythonMapObserverWrapperU32 - ); - define_python_map_observer!( - PythonOwnedMapObserverU64, - "OwnedMapObserverU64", - PythonMapObserverU64, - "MapObserverU64", - u64, - PythonMapObserverWrapperU64 - ); - - /// Register the classes to the python module - pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - Ok(()) - } -} - -#[cfg(feature = "python")] -/// Map Feedback Python bindings -pub mod pybind { - use crate::feedbacks::map::{MapFeedbackMetadata, MaxMapFeedback}; - use crate::inputs::BytesInput; - use pyo3::prelude::*; - - macro_rules! define_python_map_feedback { - ($map_feedback_state_struct_name:ident, $map_feedback_state_py_name:tt, $max_map_feedback_struct_name:ident, - $max_map_feedback_py_name:tt, $datatype:ty, $map_observer_name: ident, $std_state_name: ident) => { - use crate::observers::map::pybind::$map_observer_name; - use crate::state::pybind::$std_state_name; - - #[pyclass(unsendable, name = $map_feedback_state_py_name)] - #[derive(Clone, Debug)] - /// Python class for MapFeedbackMetadata - pub struct $map_feedback_state_struct_name { - /// Rust wrapped MapFeedbackMetadata object - pub map_feedback_state: MapFeedbackMetadata<$datatype>, - } - - #[pymethods] - impl $map_feedback_state_struct_name { - #[staticmethod] - fn with_observer(py_observer: &$map_observer_name) -> Self { - Self { - map_feedback_state: MapFeedbackMetadata::with_observer(py_observer), - } - } - } - - #[pyclass(unsendable, name = $max_map_feedback_py_name)] - #[derive(Debug)] - /// Python class for MaxMapFeedback - pub struct $max_map_feedback_struct_name { - /// Rust wrapped MaxMapFeedback object - pub max_map_feedback: - MaxMapFeedback, - } - - impl Clone for $max_map_feedback_struct_name { - fn clone(&self) -> Self { - Self { - max_map_feedback: self.max_map_feedback.clone(), - } - } - } - - #[pymethods] - impl $max_map_feedback_struct_name { - #[new] - fn new( - py_feedback_state: &$map_feedback_state_struct_name, - py_observer: &$map_observer_name, - ) -> Self { - Self { - max_map_feedback: MaxMapFeedback::new( - &py_feedback_state.map_feedback_state, - py_observer, - ), - } + impl HasObserverName for $struct_name { + fn observer_name(&self) -> &str { + self.inner.observer_name() } } }; } define_python_map_feedback!( - PythonMapFeedbackMetadataI8, - "MapFeedbackMetadataI8", PythonMaxMapFeedbackI8, "MaxMapFeedbackI8", i8, PythonMapObserverI8, - MyStdStateI8 + PythonStdState ); - define_python_map_feedback!( - PythonMapFeedbackMetadataI16, - "MapFeedbackMetadataI16", PythonMaxMapFeedbackI16, "MaxMapFeedbackI16", i16, PythonMapObserverI16, - MyStdStateI16 + PythonStdState ); define_python_map_feedback!( - PythonMapFeedbackMetadataI32, - "MapFeedbackMetadataI32", PythonMaxMapFeedbackI32, "MaxMapFeedbackI32", i32, PythonMapObserverI32, - MyStdStateI32 + PythonStdState ); define_python_map_feedback!( - PythonMapFeedbackMetadataI64, - "MapFeedbackMetadataI64", PythonMaxMapFeedbackI64, "MaxMapFeedbackI64", i64, PythonMapObserverI64, - MyStdStateI64 + PythonStdState ); define_python_map_feedback!( - PythonMapFeedbackMetadataU8, - "MapFeedbackMetadataU8", PythonMaxMapFeedbackU8, "MaxMapFeedbackU8", u8, PythonMapObserverU8, - MyStdStateU8 + PythonStdState ); - define_python_map_feedback!( - PythonMapFeedbackMetadataU16, - "MapFeedbackMetadataU16", PythonMaxMapFeedbackU16, "MaxMapFeedbackU16", u16, PythonMapObserverU16, - MyStdStateU16 + PythonStdState ); define_python_map_feedback!( - PythonMapFeedbackMetadataU32, - "MapFeedbackMetadataU32", PythonMaxMapFeedbackU32, "MaxMapFeedbackU32", u32, PythonMapObserverU32, - MyStdStateU32 + PythonStdState ); define_python_map_feedback!( - PythonMapFeedbackMetadataU64, - "MapFeedbackMetadataU64", PythonMaxMapFeedbackU64, "MaxMapFeedbackU64", u64, PythonMapObserverU64, - MyStdStateU64 + PythonStdState ); /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/libafl/src/feedbacks/mod.rs b/libafl/src/feedbacks/mod.rs index 741cb66a6c..a17dddac1e 100644 --- a/libafl/src/feedbacks/mod.rs +++ b/libafl/src/feedbacks/mod.rs @@ -1038,112 +1038,577 @@ impl From for ConstFeedback { /// `Feedback` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { - use crate::inputs::BytesInput; + use super::{ + ConstFeedback, CrashFeedback, Debug, EagerAndFeedback, EagerOrFeedback, FastAndFeedback, + FastOrFeedback, Feedback, NotFeedback, String, ToString, + }; + use crate::corpus::testcase::pybind::{PythonTestcase, PythonTestcaseWrapper}; + use crate::events::pybind::PythonEventManager; + use crate::executors::pybind::PythonExitKind; + use crate::feedbacks::map::pybind::{ + PythonMaxMapFeedbackI16, PythonMaxMapFeedbackI32, PythonMaxMapFeedbackI64, + PythonMaxMapFeedbackI8, PythonMaxMapFeedbackU16, PythonMaxMapFeedbackU32, + PythonMaxMapFeedbackU64, PythonMaxMapFeedbackU8, + }; + use crate::inputs::HasBytesVec; + use crate::observers::pybind::PythonObserversTuple; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; use crate::{ bolts::tuples::Named, corpus::Testcase, events::EventFirer, executors::ExitKind, - feedbacks::Feedback, observers::ObserversTuple, Error, + inputs::BytesInput, observers::ObserversTuple, Error, }; use pyo3::prelude::*; + use std::cell::UnsafeCell; - macro_rules! define_python_feedback { - ($struct_name_trait:ident, $py_name_trait:tt, $wrapper_name: ident, $my_std_state_type_name: ident) => { - use crate::observers::map::pybind::PythonMaxMapFeedbackI8; - use crate::state::pybind::$my_std_state_type_name; + #[derive(Debug)] + pub struct PyObjectFeedback { + inner: PyObject, + name: UnsafeCell, + } - #[derive(Debug)] - enum $wrapper_name { - MaxMapI8(*mut PythonMaxMapFeedbackI8), + impl Clone for PyObjectFeedback { + fn clone(&self) -> PyObjectFeedback { + PyObjectFeedback { + inner: self.inner.clone(), + name: UnsafeCell::new(String::new()), } + } + } - #[pyclass(unsendable, name = $py_name_trait)] - #[derive(Debug)] - /// Observer Trait binding - pub struct $struct_name_trait { - pub wrapper: $wrapper_name, + impl PyObjectFeedback { + #[must_use] + pub fn new(obj: PyObject) -> Self { + PyObjectFeedback { + inner: obj, + name: UnsafeCell::new(String::new()), } + } + } - impl $struct_name_trait { - fn unwrap(&self) -> &impl Feedback { - unsafe { - match self.wrapper { - $wrapper_name::MaxMapI8(py_wrapper) => &(*py_wrapper).upcast(), - } - } - } + // crate::impl_serde_pyobjectwrapper!(PyObjectObserver, inner); - fn unwrap_mut( - &mut self, - ) -> &mut impl Feedback { - unsafe { - match self.wrapper { - $wrapper_name::MaxMapI8(py_wrapper) => &mut (*py_wrapper).upcast_mut(), - } - } - } + impl Named for PyObjectFeedback { + fn name(&self) -> &str { + let s = Python::with_gil(|py| -> PyResult { + let s: String = self.inner.call_method0(py, "name")?.extract(py)?; + Ok(s) + }) + .unwrap(); + unsafe { + *self.name.get() = s; + &*self.name.get() + } + } + } + + impl Feedback for PyObjectFeedback { + fn init_state(&mut self, state: &mut PythonStdState) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner + .call_method1(py, "init_state", (PythonStdStateWrapper::wrap(state),))?; + Ok(()) + })?; + Ok(()) + } + + fn is_interesting( + &mut self, + state: &mut PythonStdState, + manager: &mut EM, + input: &BytesInput, + observers: &OT, + exit_kind: &ExitKind, + ) -> Result + where + EM: EventFirer, + OT: ObserversTuple, + { + // SAFETY: We use this observer in Python ony when the ObserverTuple is PythonObserversTuple + let dont_look_at_this: &PythonObserversTuple = + unsafe { &*(observers as *const OT as *const PythonObserversTuple) }; + let dont_look_at_this2: &PythonEventManager = + unsafe { &*(manager as *mut EM as *const PythonEventManager) }; + Ok(Python::with_gil(|py| -> PyResult { + let r: bool = self + .inner + .call_method1( + py, + "is_interesting", + ( + PythonStdStateWrapper::wrap(state), + dont_look_at_this2.clone(), + input.bytes(), + dont_look_at_this.clone(), + PythonExitKind::from(*exit_kind), + ), + )? + .extract(py)?; + Ok(r) + })?) + } + + fn append_metadata( + &mut self, + state: &mut PythonStdState, + testcase: &mut PythonTestcase, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "append_metadata", + ( + PythonStdStateWrapper::wrap(state), + PythonTestcaseWrapper::wrap(testcase), + ), + )?; + Ok(()) + })?; + Ok(()) + } + + fn discard_metadata( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "discard_metadata", + (PythonStdStateWrapper::wrap(state), input.bytes()), + )?; + Ok(()) + })?; + Ok(()) + } + } + + #[derive(Clone, Debug)] + #[pyclass(unsendable, name = "CrashFeedback")] + pub struct PythonCrashFeedback { + pub inner: CrashFeedback, + } + + #[pymethods] + impl PythonCrashFeedback { + #[new] + fn new() -> Self { + Self { + inner: CrashFeedback::new(), + } + } + + #[must_use] + pub fn as_feedback(slf: Py) -> PythonFeedback { + PythonFeedback::new_crash(slf) + } + } + + #[derive(Clone, Debug)] + #[pyclass(unsendable, name = "ConstFeedback")] + pub struct PythonConstFeedback { + pub inner: ConstFeedback, + } + + #[pymethods] + impl PythonConstFeedback { + #[new] + fn new(v: bool) -> Self { + Self { + inner: ConstFeedback::new(v), + } + } + + #[must_use] + pub fn as_feedback(slf: Py) -> PythonFeedback { + PythonFeedback::new_const(slf) + } + } + + #[derive(Debug)] + #[pyclass(unsendable, name = "NotFeedback")] + pub struct PythonNotFeedback { + pub inner: NotFeedback, + } + + #[pymethods] + impl PythonNotFeedback { + #[new] + fn new(feedback: PythonFeedback) -> Self { + Self { + inner: NotFeedback::new(feedback), + } + } + + #[must_use] + pub fn as_feedback(slf: Py) -> PythonFeedback { + PythonFeedback::new_not(slf) + } + } + + macro_rules! define_combined { + ($feed:ident, $pyname:ident, $pystring:tt, $method:ident) => { + #[derive(Debug)] + #[pyclass(unsendable, name = $pystring)] + pub struct $pyname { + pub inner: $feed, } #[pymethods] - impl $struct_name_trait { - #[staticmethod] - fn new_map(map_feedback: &mut PythonMaxMapFeedbackI8) -> Self { + impl $pyname { + #[new] + fn new(a: PythonFeedback, b: PythonFeedback) -> Self { Self { - observer: $wrapper_name::MaxMapI8(map_feedback), + inner: $feed::new(a, b), } } - } - impl Named for $struct_name_trait { - fn name(&self) -> &str { - self.unwrap().name() - } - } - - impl Feedback for $struct_name_trait { - fn init_state(&mut self, state: &mut S) -> Result<(), Error> { - self.unwrap_mut().init_state(state) - } - - fn is_interesting( - &mut self, - state: &mut S, - manager: &mut EM, - input: &I, - observers: &OT, - exit_kind: &ExitKind, - ) -> Result - where - EM: EventFirer, - OT: ObserversTuple, - { - self.unwrap_mut() - .is_interesting(state, manager, input, observers, exit_kind) - } - - fn append_metadata( - &mut self, - state: &mut S, - testcase: &mut Testcase, - ) -> Result<(), Error> { - self.unwrap_mut().append_metadata(state, testcase) - } - - fn discard_metadata(&mut self, state: &mut S, input: &I) -> Result<(), Error> { - self.unwrap_mut().discard_metadata(state, input) + #[must_use] + pub fn as_feedback(slf: Py) -> PythonFeedback { + PythonFeedback::$method(slf) } } }; } - define_python_feedback!( - PythonFeedback, - "Feedback", - PythonFeedbackWrapper, - PythonStdState, + define_combined!( + EagerAndFeedback, + PythonEagerAndFeedback, + "EagerAndFeedback", + new_and ); + define_combined!( + FastAndFeedback, + PythonFastAndFeedback, + "FastAndFeedback", + new_fast_and + ); + define_combined!( + EagerOrFeedback, + PythonEagerOrFeedback, + "EagerOrFeedback", + new_or + ); + define_combined!( + FastOrFeedback, + PythonFastOrFeedback, + "FastOrFeedback", + new_fast_or + ); + + #[derive(Clone, Debug)] + pub enum PythonFeedbackWrapper { + MaxMapI8(Py), + MaxMapI16(Py), + MaxMapI32(Py), + MaxMapI64(Py), + MaxMapU8(Py), + MaxMapU16(Py), + MaxMapU32(Py), + MaxMapU64(Py), + Crash(Py), + Const(Py), + Not(Py), + And(Py), + FastAnd(Py), + Or(Py), + FastOr(Py), + Python(PyObjectFeedback), + } + + #[pyclass(unsendable, name = "Feedback")] + #[derive(Debug)] + /// Observer Trait binding + pub struct PythonFeedback { + pub wrapper: PythonFeedbackWrapper, + name: UnsafeCell, + } + + macro_rules! unwrap_me { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_body!($wrapper, $name, $body, PythonFeedbackWrapper, + { + MaxMapI8, + MaxMapI16, + MaxMapI32, + MaxMapI64, + MaxMapU8, + MaxMapU16, + MaxMapU32, + MaxMapU64, + Crash, + Const, + Not, + And, + FastAnd, + Or, + FastOr + }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body + } + } + ) + }; + } + + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonFeedbackWrapper, + { + MaxMapI8, + MaxMapI16, + MaxMapI32, + MaxMapI64, + MaxMapU8, + MaxMapU16, + MaxMapU32, + MaxMapU64, + Crash, + Const, + Not, + And, + FastAnd, + Or, + FastOr + }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body + } + } + ) + }; + } + + impl Clone for PythonFeedback { + fn clone(&self) -> PythonFeedback { + PythonFeedback { + wrapper: self.wrapper.clone(), + name: UnsafeCell::new(String::new()), + } + } + } + + #[pymethods] + impl PythonFeedback { + #[staticmethod] + #[must_use] + pub fn new_max_map_i8(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapI8(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_max_map_i16(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapI16(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_max_map_i32(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapI32(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_max_map_i64(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapI64(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_max_map_u8(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapU8(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_max_map_u16(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapU16(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_max_map_u32(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapU32(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_max_map_u64(map_feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::MaxMapU64(map_feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_crash(feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::Crash(feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_const(feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::Const(feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_not(feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::Not(feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_and(feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::And(feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_fast_and(feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::FastAnd(feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_or(feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::Or(feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_fast_or(feedback: Py) -> Self { + Self { + wrapper: PythonFeedbackWrapper::FastOr(feedback), + name: UnsafeCell::new(String::new()), + } + } + + #[staticmethod] + #[must_use] + pub fn new_py(obj: PyObject) -> Self { + Self { + wrapper: PythonFeedbackWrapper::Python(PyObjectFeedback::new(obj)), + name: UnsafeCell::new(String::new()), + } + } + + pub fn unwrap_py(&self) -> Option { + match &self.wrapper { + PythonFeedbackWrapper::Python(pyo) => Some(pyo.inner.clone()), + _ => None, + } + } + } + + impl Named for PythonFeedback { + fn name(&self) -> &str { + let s = unwrap_me!(self.wrapper, f, { f.name().to_string() }); + unsafe { + *self.name.get() = s; + &*self.name.get() + } + } + } + + impl Feedback for PythonFeedback { + fn init_state(&mut self, state: &mut PythonStdState) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, f, { + Feedback::::init_state(f, state) + }) + } + + fn is_interesting( + &mut self, + state: &mut PythonStdState, + manager: &mut EM, + input: &BytesInput, + observers: &OT, + exit_kind: &ExitKind, + ) -> Result + where + EM: EventFirer, + OT: ObserversTuple, + { + unwrap_me_mut!(self.wrapper, f, { + f.is_interesting(state, manager, input, observers, exit_kind) + }) + } + + fn append_metadata( + &mut self, + state: &mut PythonStdState, + testcase: &mut Testcase, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, f, { f.append_metadata(state, testcase) }) + } + + fn discard_metadata( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, f, { f.discard_metadata(state, input) }) + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/libafl/src/fuzzer/mod.rs b/libafl/src/fuzzer/mod.rs index be9b538264..5384636632 100644 --- a/libafl/src/fuzzer/mod.rs +++ b/libafl/src/fuzzer/mod.rs @@ -669,197 +669,104 @@ where } #[cfg(feature = "python")] +#[allow(missing_docs)] /// `Fuzzer` Python bindings pub mod pybind { - use crate::feedbacks::{CrashFeedback, MaxMapFeedback}; - use crate::fuzzer::{Fuzzer, StdFuzzer}; + use crate::bolts::ownedref::OwnedPtrMut; + use crate::events::pybind::PythonEventManager; + use crate::executors::pybind::PythonExecutor; + use crate::feedbacks::pybind::PythonFeedback; + use crate::fuzzer::{Evaluator, Fuzzer, StdFuzzer}; use crate::inputs::BytesInput; + use crate::observers::pybind::PythonObserversTuple; use crate::schedulers::QueueScheduler; + use crate::stages::pybind::PythonStagesTuple; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; use pyo3::prelude::*; - macro_rules! define_python_fuzzer { - ($type_name:ident, $struct_name:ident, $py_name:tt, $datatype:ty, $my_std_state_type_name: ident, $std_state_name: ident, - $event_manager_name: ident, $map_observer_name: ident, $max_map_feedback_py_name: ident, $executor_name: ident, $stage_tuple_name: ident) => { - use crate::events::pybind::$event_manager_name; - use crate::executors::pybind::$executor_name; - use crate::feedbacks::map::pybind::$max_map_feedback_py_name; - use crate::observers::map::pybind::$map_observer_name; - use crate::stages::owned::pybind::$stage_tuple_name; - use crate::state::pybind::{$my_std_state_type_name, $std_state_name}; + /// `StdFuzzer` with fixed generics + pub type PythonStdFuzzer = StdFuzzer< + QueueScheduler, + PythonFeedback, + BytesInput, + PythonFeedback, + PythonObserversTuple, + PythonStdState, + >; - /// `StdFuzzer` with fixed generics - pub type $type_name = StdFuzzer< - QueueScheduler, - MaxMapFeedback, - BytesInput, - CrashFeedback, - ($map_observer_name, ()), - $my_std_state_type_name, - >; - /// Python class for StdFuzzer - #[pyclass(unsendable, name = $py_name)] - #[derive(Debug)] - pub struct $struct_name { - /// Rust wrapped StdFuzzer object - pub std_fuzzer: $type_name, - } - - #[pymethods] - impl $struct_name { - #[new] - fn new(py_max_map_feedback: $max_map_feedback_py_name) -> Self { - Self { - std_fuzzer: StdFuzzer::new( - QueueScheduler::new(), - py_max_map_feedback.max_map_feedback, - CrashFeedback::new(), - ), - } - } - - fn fuzz_loop( - &mut self, - py_executor: &mut $executor_name, - py_state: &mut $std_state_name, - py_mgr: &mut $event_manager_name, - stage_tuple: &mut $stage_tuple_name, - ) { - self.std_fuzzer - .fuzz_loop( - &mut stage_tuple.stages_owned_list, - py_executor, - &mut py_state.std_state, - py_mgr, - ) - .expect("Failed to generate the initial corpus".into()); - } - } - }; + /// Python class for StdFuzzer + #[pyclass(unsendable, name = "StdFuzzer")] + #[derive(Debug)] + pub struct PythonStdFuzzerWrapper { + /// Rust wrapped StdFuzzer object + pub inner: OwnedPtrMut, } - define_python_fuzzer!( - MyStdFuzzerI8, - PythonStdFuzzerI8, - "StdFuzzerI8", - i8, - MyStdStateI8, - PythonStdStateI8, - PythonEventManagerI8, - PythonMapObserverI8, - PythonMaxMapFeedbackI8, - PythonExecutorI8, - PythonStagesOwnedListI8 - ); + impl PythonStdFuzzerWrapper { + pub fn wrap(r: &mut PythonStdFuzzer) -> Self { + Self { + inner: OwnedPtrMut::Ptr(r), + } + } - define_python_fuzzer!( - MyStdFuzzerI16, - PythonStdFuzzerI16, - "StdFuzzerI16", - i16, - MyStdStateI16, - PythonStdStateI16, - PythonEventManagerI16, - PythonMapObserverI16, - PythonMaxMapFeedbackI16, - PythonExecutorI16, - PythonStagesOwnedListI16 - ); + #[must_use] + pub fn unwrap(&self) -> &PythonStdFuzzer { + self.inner.as_ref() + } - define_python_fuzzer!( - MyStdFuzzerI32, - PythonStdFuzzerI32, - "StdFuzzerI32", - i32, - MyStdStateI32, - PythonStdStateI32, - PythonEventManagerI32, - PythonMapObserverI32, - PythonMaxMapFeedbackI32, - PythonExecutorI32, - PythonStagesOwnedListI32 - ); + pub fn unwrap_mut(&mut self) -> &mut PythonStdFuzzer { + self.inner.as_mut() + } + } - define_python_fuzzer!( - MyStdFuzzerI64, - PythonStdFuzzerI64, - "StdFuzzerI64", - i64, - MyStdStateI64, - PythonStdStateI64, - PythonEventManagerI64, - PythonMapObserverI64, - PythonMaxMapFeedbackI64, - PythonExecutorI64, - PythonStagesOwnedListI64 - ); + #[pymethods] + impl PythonStdFuzzerWrapper { + #[new] + fn new(py_feedback: PythonFeedback, py_objective: PythonFeedback) -> Self { + Self { + inner: OwnedPtrMut::Owned(Box::new(StdFuzzer::new( + QueueScheduler::new(), + py_feedback, + py_objective, + ))), + } + } - define_python_fuzzer!( - MyStdFuzzerU8, - PythonStdFuzzerU8, - "StdFuzzerU8", - u8, - MyStdStateU8, - PythonStdStateU8, - PythonEventManagerU8, - PythonMapObserverU8, - PythonMaxMapFeedbackU8, - PythonExecutorU8, - PythonStagesOwnedListU8 - ); + fn add_input( + &mut self, + py_state: &mut PythonStdStateWrapper, + py_executor: &mut PythonExecutor, + py_mgr: &mut PythonEventManager, + input: Vec, + ) -> usize { + self.inner + .as_mut() + .add_input( + py_state.unwrap_mut(), + py_executor, + py_mgr, + BytesInput::new(input), + ) + .expect("Failed to add input") + } - define_python_fuzzer!( - MyStdFuzzerU16, - PythonStdFuzzerU16, - "StdFuzzerU16", - u16, - MyStdStateU16, - PythonStdStateU16, - PythonEventManagerU16, - PythonMapObserverU16, - PythonMaxMapFeedbackU16, - PythonExecutorU16, - PythonStagesOwnedListU16 - ); - - define_python_fuzzer!( - MyStdFuzzerU32, - PythonStdFuzzerU32, - "StdFuzzerU32", - u32, - MyStdStateU32, - PythonStdStateU32, - PythonEventManagerU32, - PythonMapObserverU32, - PythonMaxMapFeedbackU32, - PythonExecutorU32, - PythonStagesOwnedListU32 - ); - - define_python_fuzzer!( - MyStdFuzzerU64, - PythonStdFuzzerU64, - "StdFuzzerU64", - u64, - MyStdStateU64, - PythonStdStateU64, - PythonEventManagerU64, - PythonMapObserverU64, - PythonMaxMapFeedbackU64, - PythonExecutorU64, - PythonStagesOwnedListU64 - ); + fn fuzz_loop( + &mut self, + py_executor: &mut PythonExecutor, + py_state: &mut PythonStdStateWrapper, + py_mgr: &mut PythonEventManager, + stages_tuple: &mut PythonStagesTuple, + ) { + self.inner + .as_mut() + .fuzz_loop(stages_tuple, py_executor, py_state.unwrap_mut(), py_mgr) + .expect("Failed to generate the initial corpus"); + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; Ok(()) } } diff --git a/libafl/src/generators/mod.rs b/libafl/src/generators/mod.rs index 3cd9c12229..d8f18e7bd8 100644 --- a/libafl/src/generators/mod.rs +++ b/libafl/src/generators/mod.rs @@ -127,88 +127,199 @@ where } /// `Generator` Python bindings +#[allow(missing_docs)] #[cfg(feature = "python")] pub mod pybind { - use crate::generators::RandPrintablesGenerator; + use crate::generators::{Generator, RandBytesGenerator, RandPrintablesGenerator}; + use crate::inputs::{BytesInput, HasBytesVec}; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; + use crate::Error; use pyo3::prelude::*; - macro_rules! define_python_event_manager { - ($struct_name:ident, $py_name:tt, $my_std_state_type_name: ident) => { - use crate::state::pybind::$my_std_state_type_name; + #[derive(Clone, Debug)] + pub struct PyObjectGenerator { + inner: PyObject, + } - #[pyclass(unsendable, name = $py_name)] - /// Python class for RandPrintablesGenerator - #[derive(Debug)] - pub struct $struct_name { - /// Rust wrapped SimpleEventManager object - pub rand_printable_generator: RandPrintablesGenerator<$my_std_state_type_name>, + impl PyObjectGenerator { + #[must_use] + pub fn new(obj: PyObject) -> Self { + PyObjectGenerator { inner: obj } + } + } + + impl Generator for PyObjectGenerator { + fn generate(&mut self, state: &mut PythonStdState) -> Result { + let bytes = Python::with_gil(|py| -> PyResult> { + self.inner + .call_method1(py, "generate", (PythonStdStateWrapper::wrap(state),))? + .extract(py) + }) + .unwrap(); + Ok(BytesInput::new(bytes)) + } + + fn generate_dummy(&self, state: &mut PythonStdState) -> BytesInput { + let bytes = Python::with_gil(|py| -> PyResult> { + self.inner + .call_method1(py, "generate_dummy", (PythonStdStateWrapper::wrap(state),))? + .extract(py) + }) + .unwrap(); + BytesInput::new(bytes) + } + } + + #[pyclass(unsendable, name = "RandBytesGenerator")] + #[derive(Debug, Clone)] + /// Python class for RandBytesGenerator + pub struct PythonRandBytesGenerator { + /// Rust wrapped RandBytesGenerator object + pub inner: RandBytesGenerator, + } + + #[pymethods] + impl PythonRandBytesGenerator { + #[new] + fn new(max_size: usize) -> Self { + Self { + inner: RandBytesGenerator::new(max_size), } + } - #[pymethods] - impl $struct_name { - #[new] - fn new(max_size: usize) -> Self { - Self { - rand_printable_generator: RandPrintablesGenerator::new(max_size), + fn generate(&mut self, state: &mut PythonStdStateWrapper) -> Vec { + self.inner + .generate(state.unwrap_mut()) + .expect("PythonRandBytesGenerator::generate failed") + .bytes() + .to_vec() + } + + fn as_generator(slf: Py) -> PythonGenerator { + PythonGenerator::new_rand_bytes(slf) + } + } + + #[pyclass(unsendable, name = "RandPrintablesGenerator")] + #[derive(Debug, Clone)] + /// Python class for RandPrintablesGenerator + pub struct PythonRandPrintablesGenerator { + /// Rust wrapped RandPrintablesGenerator object + pub inner: RandPrintablesGenerator, + } + + #[pymethods] + impl PythonRandPrintablesGenerator { + #[new] + fn new(max_size: usize) -> Self { + Self { + inner: RandPrintablesGenerator::new(max_size), + } + } + + fn generate(&mut self, state: &mut PythonStdStateWrapper) -> Vec { + self.inner + .generate(state.unwrap_mut()) + .expect("PythonRandPrintablesGenerator::generate failed") + .bytes() + .to_vec() + } + + fn as_generator(slf: Py) -> PythonGenerator { + PythonGenerator::new_rand_printables(slf) + } + } + + #[derive(Debug, Clone)] + enum PythonGeneratorWrapper { + RandBytes(Py), + RandPrintables(Py), + Python(PyObjectGenerator), + } + + /// Rand Trait binding + #[pyclass(unsendable, name = "Generator")] + #[derive(Debug, Clone)] + pub struct PythonGenerator { + wrapper: PythonGeneratorWrapper, + } + + macro_rules! unwrap_me { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_body!($wrapper, $name, $body, PythonGeneratorWrapper, + { RandBytes, RandPrintables }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body } } - } + ) }; } - define_python_event_manager!( - PythonRandPrintablesGeneratorI8, - "RandPrintablesGeneratorI8", - MyStdStateI8 - ); - define_python_event_manager!( - PythonRandPrintablesGeneratorI16, - "RandPrintablesGeneratorI16", - MyStdStateI16 - ); - define_python_event_manager!( - PythonRandPrintablesGeneratorI32, - "RandPrintablesGeneratorI32", - MyStdStateI32 - ); - define_python_event_manager!( - PythonRandPrintablesGeneratorI64, - "RandPrintablesGeneratorI64", - MyStdStateI64 - ); + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonGeneratorWrapper, + { RandBytes, RandPrintables }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body + } + } + ) + }; + } - define_python_event_manager!( - PythonRandPrintablesGeneratorU8, - "RandPrintablesGeneratorU8", - MyStdStateU8 - ); - define_python_event_manager!( - PythonRandPrintablesGeneratorU16, - "RandPrintablesGeneratorU16", - MyStdStateU16 - ); - define_python_event_manager!( - PythonRandPrintablesGeneratorU32, - "RandPrintablesGeneratorU32", - MyStdStateU32 - ); - define_python_event_manager!( - PythonRandPrintablesGeneratorU64, - "RandPrintablesGeneratorU64", - MyStdStateU64 - ); + #[pymethods] + impl PythonGenerator { + #[staticmethod] + fn new_rand_bytes(py_gen: Py) -> Self { + Self { + wrapper: PythonGeneratorWrapper::RandBytes(py_gen), + } + } + + #[staticmethod] + fn new_rand_printables(py_gen: Py) -> Self { + Self { + wrapper: PythonGeneratorWrapper::RandPrintables(py_gen), + } + } + + #[staticmethod] + #[must_use] + pub fn new_py(obj: PyObject) -> Self { + Self { + wrapper: PythonGeneratorWrapper::Python(PyObjectGenerator::new(obj)), + } + } + + #[must_use] + pub fn unwrap_py(&self) -> Option { + match &self.wrapper { + PythonGeneratorWrapper::Python(pyo) => Some(pyo.inner.clone()), + _ => None, + } + } + } + + impl Generator for PythonGenerator { + fn generate(&mut self, state: &mut PythonStdState) -> Result { + unwrap_me_mut!(self.wrapper, g, { g.generate(state) }) + } + + fn generate_dummy(&self, state: &mut PythonStdState) -> BytesInput { + unwrap_me!(self.wrapper, g, { g.generate_dummy(state) }) + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } } diff --git a/libafl/src/lib.rs b/libafl/src/lib.rs index 8327ad0602..cdb51ef197 100644 --- a/libafl/src/lib.rs +++ b/libafl/src/lib.rs @@ -371,6 +371,22 @@ impl From for Error { } } +#[cfg(feature = "python")] +impl From for Error { + fn from(err: pyo3::PyErr) -> Self { + pyo3::Python::with_gil(|py| { + if err.matches( + py, + pyo3::types::PyType::new::(py), + ) { + Self::shutting_down() + } else { + Self::illegal_state(format!("Python exception: {:?}", err)) + } + }) + } +} + #[cfg(feature = "std")] impl std::error::Error for Error {} @@ -464,30 +480,180 @@ pub extern "C" fn external_current_millis() -> u64 { } #[cfg(feature = "python")] -use pyo3::prelude::*; +#[allow(missing_docs)] +pub mod pybind { + use super::{ + bolts, corpus, events, executors, feedbacks, fuzzer, generators, monitors, mutators, + observers, stages, state, + }; + use pyo3::prelude::*; -#[cfg(feature = "python")] -#[pymodule] -#[pyo3(name = "libafl")] -/// Register the classes to the python module -pub fn python_module(py: Python, m: &PyModule) -> PyResult<()> { - observers::map::pybind::register(py, m)?; - feedbacks::map::pybind::register(py, m)?; - state::pybind::register(py, m)?; - monitors::pybind::register(py, m)?; - events::pybind::register(py, m)?; - events::simple::pybind::register(py, m)?; - fuzzer::pybind::register(py, m)?; - executors::pybind::register(py, m)?; - executors::inprocess::pybind::register(py, m)?; - generators::pybind::register(py, m)?; - corpus::pybind::register(py, m)?; - corpus::ondisk::pybind::register(py, m)?; - corpus::inmemory::pybind::register(py, m)?; - corpus::cached::pybind::register(py, m)?; - bolts::rands::pybind::register(py, m)?; - stages::pybind::register(py, m)?; - stages::owned::pybind::register(py, m)?; - stages::mutational::pybind::register(py, m)?; - Ok(()) + #[derive(Debug, Clone)] + pub struct PythonMetadata { + pub map: PyObject, + } + + crate::impl_serde_pyobjectwrapper!(PythonMetadata, map); + crate::impl_serdeany!(PythonMetadata); + + impl PythonMetadata { + #[must_use] + pub fn new(map: PyObject) -> Self { + Self { map } + } + } + + #[macro_export] + macro_rules! unwrap_me_body { + ($wrapper:expr, $name:ident, $body:block, $wrapper_type:ident, { $($wrapper_option:tt),* }) => { + match &$wrapper { + $( + $wrapper_type::$wrapper_option(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + let $name = &borrowed.inner; + Ok($body) + }) + .unwrap() + } + )* + } + }; + ($wrapper:expr, $name:ident, $body:block, $wrapper_type:ident, { $($wrapper_option:tt),* }, { $($wrapper_optional:tt($pw:ident) => $code_block:block)* }) => { + match &$wrapper { + $( + $wrapper_type::$wrapper_option(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + let $name = &borrowed.inner; + Ok($body) + }) + .unwrap() + } + )* + $($wrapper_type::$wrapper_optional($pw) => { $code_block })* + } + }; + } + + #[macro_export] + macro_rules! unwrap_me_mut_body { + ($wrapper:expr, $name:ident, $body:block, $wrapper_type:ident, { $($wrapper_option:tt),*}) => { + match &mut $wrapper { + $( + $wrapper_type::$wrapper_option(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + let $name = &mut borrowed.inner; + Ok($body) + }) + .unwrap() + } + )* + } + }; + ($wrapper:expr, $name:ident, $body:block, $wrapper_type:ident, { $($wrapper_option:tt),*}, { $($wrapper_optional:tt($pw:ident) => $code_block:block)* }) => { + match &mut $wrapper { + $( + $wrapper_type::$wrapper_option(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + let $name = &mut borrowed.inner; + Ok($body) + }) + .unwrap() + } + )* + $($wrapper_type::$wrapper_optional($pw) => { $code_block })* + } + }; + } + + #[macro_export] + macro_rules! impl_serde_pyobjectwrapper { + ($struct_name:ident, $inner:tt) => { + const _: () = { + use pyo3::prelude::*; + use serde::{Deserialize, Deserializer, Serialize, Serializer}; + + impl Serialize for $struct_name { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let buf = Python::with_gil(|py| -> PyResult> { + let pickle = PyModule::import(py, "pickle")?; + let buf: Vec = + pickle.getattr("dumps")?.call1((&self.$inner,))?.extract()?; + Ok(buf) + }) + .unwrap(); + serializer.serialize_bytes(&buf) + } + } + + struct PyObjectVisitor; + + impl<'de> serde::de::Visitor<'de> for PyObjectVisitor { + type Value = $struct_name; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter + .write_str("Expecting some bytes to deserialize from the Python side") + } + + fn visit_byte_buf(self, v: Vec) -> Result + where + E: serde::de::Error, + { + let obj = Python::with_gil(|py| -> PyResult { + let pickle = PyModule::import(py, "pickle")?; + let obj = pickle.getattr("loads")?.call1((v,))?.to_object(py); + Ok(obj) + }) + .unwrap(); + Ok($struct_name::new(obj)) + } + } + + impl<'de> Deserialize<'de> for $struct_name { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_byte_buf(PyObjectVisitor) + } + } + }; + }; + } + + #[pymodule] + #[pyo3(name = "libafl")] + /// Register the classes to the python module + pub fn python_module(py: Python, m: &PyModule) -> PyResult<()> { + observers::map::pybind::register(py, m)?; + observers::pybind::register(py, m)?; + feedbacks::map::pybind::register(py, m)?; + feedbacks::pybind::register(py, m)?; + state::pybind::register(py, m)?; + monitors::pybind::register(py, m)?; + events::pybind::register(py, m)?; + events::simple::pybind::register(py, m)?; + fuzzer::pybind::register(py, m)?; + executors::pybind::register(py, m)?; + executors::inprocess::pybind::register(py, m)?; + generators::pybind::register(py, m)?; + mutators::pybind::register(py, m)?; + mutators::scheduled::pybind::register(py, m)?; + corpus::pybind::register(py, m)?; + corpus::testcase::pybind::register(py, m)?; + corpus::ondisk::pybind::register(py, m)?; + corpus::inmemory::pybind::register(py, m)?; + corpus::cached::pybind::register(py, m)?; + bolts::rands::pybind::register(py, m)?; + stages::pybind::register(py, m)?; + stages::mutational::pybind::register(py, m)?; + Ok(()) + } } diff --git a/libafl/src/monitors/mod.rs b/libafl/src/monitors/mod.rs index 845094c02e..5c6e83b201 100644 --- a/libafl/src/monitors/mod.rs +++ b/libafl/src/monitors/mod.rs @@ -815,95 +815,136 @@ impl Default for ClientPerfMonitor { } /// `Monitor` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { use crate::monitors::{Monitor, SimpleMonitor}; use pyo3::prelude::*; + use pyo3::types::PyUnicode; use super::ClientStats; use core::time::Duration; + // TODO create a PyObjectFnMut to pass, track stabilization of https://github.com/rust-lang/rust/issues/29625 + #[pyclass(unsendable, name = "SimpleMonitor")] - #[derive(Clone, Debug)] /// Python class for SimpleMonitor pub struct PythonSimpleMonitor { /// Rust wrapped SimpleMonitor object - pub simple_monitor: SimpleMonitor, + pub inner: SimpleMonitor>, + print_fn: PyObject, + } + + impl std::fmt::Debug for PythonSimpleMonitor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PythonSimpleMonitor") + .field("print_fn", &self.print_fn) + .finish() + } + } + + impl Clone for PythonSimpleMonitor { + fn clone(&self) -> PythonSimpleMonitor { + let py_print_fn = self.print_fn.clone(); + let closure = move |s: String| { + Python::with_gil(|py| -> PyResult<()> { + py_print_fn.call1(py, (PyUnicode::new(py, &s),))?; + Ok(()) + }) + .unwrap(); + }; + + PythonSimpleMonitor { + inner: SimpleMonitor { + print_fn: Box::new(closure), + start_time: self.inner.start_time, + client_stats: self.inner.client_stats.clone(), + }, + print_fn: self.print_fn.clone(), + } + } } #[pymethods] impl PythonSimpleMonitor { #[new] - fn new(/*py_print_fn: PyObject */) -> Self { - // TODO: Find a fix to: closures can only be coerced to `fn` types if they - // do not capture any variables and print_fn expected to be fn pointer. - // fn printf_fn (s: String){ - // Python::with_gil(|py| -> PyResult<()> { - // print_fn.call1(py, (PyUnicode::new(py, &s),)); - // Ok(()) - // }); - // } + fn new(py_print_fn: PyObject) -> Self { + let py_print_fn1 = py_print_fn.clone(); + let closure = move |s: String| { + Python::with_gil(|py| -> PyResult<()> { + py_print_fn1.call1(py, (PyUnicode::new(py, &s),))?; + Ok(()) + }) + .unwrap(); + }; Self { - simple_monitor: SimpleMonitor::new(|s| println!("{}", s)), + inner: SimpleMonitor::new(Box::new(closure)), + print_fn: py_print_fn, } } + + #[must_use] + pub fn as_monitor(slf: Py) -> PythonMonitor { + PythonMonitor::new_simple(slf) + } } #[derive(Clone, Debug)] enum PythonMonitorWrapper { - Simple(PythonSimpleMonitor), + Simple(Py), } #[pyclass(unsendable, name = "Monitor")] #[derive(Clone, Debug)] /// EventManager Trait binding pub struct PythonMonitor { - monitor: PythonMonitorWrapper, + wrapper: PythonMonitorWrapper, } - impl PythonMonitor { - fn get_monitor(&self) -> &impl Monitor { - match &self.monitor { - PythonMonitorWrapper::Simple(py_simple_monitor) => { - &py_simple_monitor.simple_monitor - } - } - } + macro_rules! unwrap_me { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_body!($wrapper, $name, $body, PythonMonitorWrapper, { Simple }) + }; + } - fn get_mut_monitor(&mut self) -> &mut impl Monitor { - match &mut self.monitor { - PythonMonitorWrapper::Simple(py_simple_monitor) => { - &mut py_simple_monitor.simple_monitor - } - } - } + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonMonitorWrapper, { Simple }) + }; } #[pymethods] impl PythonMonitor { #[staticmethod] - fn new_from_simple(simple_monitor: PythonSimpleMonitor) -> Self { + #[must_use] + pub fn new_simple(simple_monitor: Py) -> Self { Self { - monitor: PythonMonitorWrapper::Simple(simple_monitor), + wrapper: PythonMonitorWrapper::Simple(simple_monitor), } } } impl Monitor for PythonMonitor { fn client_stats_mut(&mut self) -> &mut Vec { - self.get_mut_monitor().client_stats_mut() + let ptr = unwrap_me_mut!(self.wrapper, m, { + m.client_stats_mut() as *mut Vec + }); + unsafe { ptr.as_mut().unwrap() } } fn client_stats(&self) -> &[ClientStats] { - self.get_monitor().client_stats() + let ptr = unwrap_me!(self.wrapper, m, { + m.client_stats() as *const [ClientStats] + }); + unsafe { ptr.as_ref().unwrap() } } /// Time this fuzzing run stated fn start_time(&mut self) -> Duration { - self.get_mut_monitor().start_time() + unwrap_me_mut!(self.wrapper, m, { m.start_time() }) } fn display(&mut self, event_msg: String, sender_id: u32) { - self.get_mut_monitor().display(event_msg, sender_id); + unwrap_me_mut!(self.wrapper, m, { m.display(event_msg, sender_id) }); } } /// Register the classes to the python module diff --git a/libafl/src/mutators/mod.rs b/libafl/src/mutators/mod.rs index 4b222314bb..8c9e8079de 100644 --- a/libafl/src/mutators/mod.rs +++ b/libafl/src/mutators/mod.rs @@ -207,3 +207,150 @@ where } } } + +/// `Mutator` Python bindings +#[cfg(feature = "python")] +#[allow(missing_docs)] +pub mod pybind { + use super::{MutationResult, Mutator}; + use crate::inputs::{BytesInput, HasBytesVec}; + use crate::mutators::scheduled::pybind::PythonStdHavocMutator; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; + use crate::Error; + use pyo3::prelude::*; + + #[derive(Clone, Debug)] + pub struct PyObjectMutator { + inner: PyObject, + } + + impl PyObjectMutator { + #[must_use] + pub fn new(obj: PyObject) -> Self { + PyObjectMutator { inner: obj } + } + } + + impl Mutator for PyObjectMutator { + fn mutate( + &mut self, + state: &mut PythonStdState, + input: &mut BytesInput, + stage_idx: i32, + ) -> Result { + let mutated = Python::with_gil(|py| -> PyResult { + self.inner + .call_method1( + py, + "mutate", + (PythonStdStateWrapper::wrap(state), input.bytes(), stage_idx), + )? + .extract(py) + })?; + Ok(if mutated { + MutationResult::Mutated + } else { + MutationResult::Skipped + }) + } + + fn post_exec( + &mut self, + state: &mut PythonStdState, + stage_idx: i32, + corpus_idx: Option, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "post_exec", + (PythonStdStateWrapper::wrap(state), stage_idx, corpus_idx), + )?; + Ok(()) + })?; + Ok(()) + } + } + + #[derive(Debug, Clone)] + pub enum PythonMutatorWrapper { + StdHavoc(Py), + Python(PyObjectMutator), + } + + /// Mutator Trait binding + #[pyclass(unsendable, name = "Mutator")] + #[derive(Debug, Clone)] + pub struct PythonMutator { + pub wrapper: PythonMutatorWrapper, + } + + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonMutatorWrapper, { + StdHavoc + }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body + } + }) + }; + } + + #[pymethods] + impl PythonMutator { + #[staticmethod] + #[must_use] + pub fn new_std_havoc(mgr: Py) -> Self { + Self { + wrapper: PythonMutatorWrapper::StdHavoc(mgr), + } + } + + #[staticmethod] + #[must_use] + pub fn new_py(obj: PyObject) -> Self { + Self { + wrapper: PythonMutatorWrapper::Python(PyObjectMutator::new(obj)), + } + } + + #[must_use] + pub fn unwrap_py(&self) -> Option { + match &self.wrapper { + PythonMutatorWrapper::Python(pyo) => Some(pyo.inner.clone()), + PythonMutatorWrapper::StdHavoc(_) => None, + } + } + } + + impl Mutator for PythonMutator { + fn mutate( + &mut self, + state: &mut PythonStdState, + input: &mut BytesInput, + stage_idx: i32, + ) -> Result { + unwrap_me_mut!(self.wrapper, m, { m.mutate(state, input, stage_idx) }) + } + + fn post_exec( + &mut self, + state: &mut PythonStdState, + stage_idx: i32, + corpus_idx: Option, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, m, { + m.post_exec(state, stage_idx, corpus_idx) + }) + } + } + + /// Register the classes to the python module + pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; + Ok(()) + } +} diff --git a/libafl/src/mutators/scheduled.rs b/libafl/src/mutators/scheduled.rs index b937d9e2cf..92d499fd83 100644 --- a/libafl/src/mutators/scheduled.rs +++ b/libafl/src/mutators/scheduled.rs @@ -208,9 +208,8 @@ where } } -/// Get the mutations that compose the Havoc mutator -#[must_use] -pub fn havoc_mutations() -> tuple_list_type!( +/// Tuple type of the mutations that compose the Havoc mutator +pub type HavocMutationsType = tuple_list_type!( BitFlipMutator, ByteFlipMutator, ByteIncMutator, @@ -238,7 +237,11 @@ pub fn havoc_mutations() -> tuple_list_type!( BytesSwapMutator, CrossoverInsertMutator, CrossoverReplaceMutator, -) { +); + +/// Get the mutations that compose the Havoc mutator +#[must_use] +pub fn havoc_mutations() -> HavocMutationsType { tuple_list!( BitFlipMutator::new(), ByteFlipMutator::new(), @@ -497,3 +500,42 @@ mod tests { } } } + +/// `SchedulerMutator` Python bindings +#[cfg(feature = "python")] +#[allow(missing_docs)] +pub mod pybind { + use super::{havoc_mutations, Debug, HavocMutationsType, StdScheduledMutator}; + use crate::inputs::BytesInput; + use crate::mutators::pybind::PythonMutator; + use crate::state::pybind::PythonStdState; + use pyo3::prelude::*; + + #[pyclass(unsendable, name = "StdHavocMutator")] + #[derive(Debug)] + /// Python class for StdHavocMutator + pub struct PythonStdHavocMutator { + /// Rust wrapped StdHavocMutator object + pub inner: StdScheduledMutator, + } + + #[pymethods] + impl PythonStdHavocMutator { + #[new] + fn new() -> Self { + Self { + inner: StdScheduledMutator::new(havoc_mutations()), + } + } + + fn as_mutator(slf: Py) -> PythonMutator { + PythonMutator::new_std_havoc(slf) + } + } + + /// Register the classes to the python module + pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; + Ok(()) + } +} diff --git a/libafl/src/observers/map.rs b/libafl/src/observers/map.rs index 5cc8de407a..d51b9e9024 100644 --- a/libafl/src/observers/map.rs +++ b/libafl/src/observers/map.rs @@ -181,7 +181,7 @@ where /// The Map Observer retrieves the state of a map, /// that will get updated by the target. /// A well-known example is the AFL-Style coverage map. -#[derive(Serialize, Deserialize, Debug)] +#[derive(Clone, Serialize, Deserialize, Debug)] #[serde(bound = "T: serde::de::DeserializeOwned")] #[allow(clippy::unsafe_derive_deserialize)] pub struct StdMapObserver<'a, T> @@ -384,20 +384,26 @@ where { /// Creates a new [`MapObserver`] #[must_use] - pub fn new(name: &'static str, map: &'a mut [T]) -> Self { + pub fn new(name: S, map: &'a mut [T]) -> Self + where + S: Into, + { Self { map: OwnedSliceMut::from(map), - name: name.to_string(), + name: name.into(), initial: T::default(), } } /// Creates a new [`MapObserver`] with an owned map #[must_use] - pub fn new_owned(name: &'static str, map: Vec) -> Self { + pub fn new_owned(name: S, map: Vec) -> Self + where + S: Into, + { Self { map: OwnedSliceMut::from(map), - name: name.to_string(), + name: name.into(), initial: T::default(), } } @@ -407,10 +413,13 @@ where /// # Safety /// Will dereference the owned slice with up to len elements. #[must_use] - pub fn new_from_ownedref(name: &'static str, map: OwnedSliceMut<'a, T>) -> Self { + pub fn new_from_ownedref(name: S, map: OwnedSliceMut<'a, T>) -> Self + where + S: Into, + { Self { map, - name: name.to_string(), + name: name.into(), initial: T::default(), } } @@ -419,10 +428,13 @@ where /// /// # Safety /// Will dereference the `map_ptr` with up to len elements. - pub unsafe fn new_from_ptr(name: &'static str, map_ptr: *mut T, len: usize) -> Self { + pub unsafe fn new_from_ptr(name: S, map_ptr: *mut T, len: usize) -> Self + where + S: Into, + { StdMapObserver { map: OwnedSliceMut::from_raw_parts_mut(map_ptr, len), - name: name.to_string(), + name: name.into(), initial: T::default(), } } @@ -1490,26 +1502,128 @@ where /// `MapObserver` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { - use crate::bolts::{tuples::Named, AsMutIterator, AsRefIterator, HasLen}; - use crate::observers::{map::OwnedMapObserver, MapObserver, Observer}; - use crate::Error; + use super::{ + AsMutIterator, AsRefIterator, Debug, Error, HasLen, Iter, IterMut, MapObserver, Named, + Observer, OwnedMapObserver, StdMapObserver, String, Vec, + }; + use crate::observers::pybind::PythonObserver; + use concat_idents::concat_idents; use pyo3::prelude::*; use serde::{Deserialize, Serialize}; - use std::slice::{Iter, IterMut}; + + #[macro_export] + macro_rules! mapob_unwrap_me { + ($wrapper_name:ident, $wrapper:expr, $name:ident, $body:block) => { + match &$wrapper { + $wrapper_name::Std(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + let $name = &borrowed.inner; + Ok($body) + }) + .unwrap(), + $wrapper_name::Owned(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + let $name = &borrowed.inner; + Ok($body) + }) + .unwrap(), + $wrapper_name::None => panic!("Serde is not supported ATM"), + } + }; + } + + #[macro_export] + macro_rules! mapob_unwrap_me_mut { + ($wrapper_name:ident, $wrapper:expr, $name:ident, $body:block) => { + match &mut $wrapper { + $wrapper_name::Std(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + let $name = &mut borrowed.inner; + Ok($body) + }) + .unwrap(), + $wrapper_name::Owned(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + let $name = &mut borrowed.inner; + Ok($body) + }) + .unwrap(), + $wrapper_name::None => panic!("Serde is not supported ATM"), + } + }; + } macro_rules! define_python_map_observer { - ($struct_name:ident, $py_name:tt, $struct_name_trait:ident, $py_name_trait:tt, $datatype:ty, $wrapper_name: ident) => { - #[pyclass(unsendable, name = $py_name)] + ($struct_name1:ident, $py_name1:tt, $struct_name2:ident, $py_name2:tt, $struct_name_trait:ident, $py_name_trait:tt, $datatype:ty, $wrapper_name: ident) => { + #[pyclass(unsendable, name = $py_name1)] + #[allow(clippy::unsafe_derive_deserialize)] + #[derive(Serialize, Deserialize, Debug, Clone)] + /// Python class for StdMapObserver + pub struct $struct_name1 { + /// Rust wrapped StdMapObserver object + pub inner: StdMapObserver<'static, $datatype>, + } + + #[pymethods] + impl $struct_name1 { + #[new] + fn new(name: String, ptr: usize, size: usize) -> Self { + Self { + inner: unsafe { StdMapObserver::new_from_ptr(name, ptr as *mut $datatype, size) } + } + } + + #[must_use] + pub fn as_map_observer(slf: Py) -> $struct_name_trait { + $struct_name_trait::new_std(slf) + } + + #[must_use] + pub fn as_observer(slf: Py) -> PythonObserver { + let m = Self::as_map_observer(slf); + Python::with_gil(|py| -> PyResult { + let p: Py<_> = Py::new(py, m)?; + Ok($struct_name_trait::as_observer(p)) + }).unwrap() + } + + fn __getitem__(&self, idx: usize) -> $datatype { + *self.inner.get(idx) + } + + fn __setitem__(&mut self, idx: usize, val: $datatype) { + *self.inner.get_mut(idx) = val; + } + + #[pyo3(name = "usable_count")] + fn pyusable_count(&self) -> usize { + self.inner.usable_count() + } + + #[pyo3(name = "len")] + fn pylen(&self) -> usize { + self.inner.len() + } + + #[pyo3(name = "name")] + fn pyname(&self) -> &str { + self.inner.name() + } + + } + + #[pyclass(unsendable, name = $py_name2)] #[derive(Serialize, Deserialize, Debug, Clone)] /// Python class for OwnedMapObserver (i.e. StdMapObserver with owned map) - pub struct $struct_name { + pub struct $struct_name2 { /// Rust wrapped OwnedMapObserver object pub inner: OwnedMapObserver<$datatype>, } #[pymethods] - impl $struct_name { + impl $struct_name2 { #[new] fn new(name: String, map: Vec<$datatype>) -> Self { Self { @@ -1517,11 +1631,56 @@ pub mod pybind { inner: OwnedMapObserver::new(Box::leak(name.into_boxed_str()), map), } } + + #[must_use] + pub fn as_map_observer(slf: Py) -> $struct_name_trait { + $struct_name_trait::new_owned(slf) + } + + #[must_use] + pub fn as_observer(slf: Py) -> PythonObserver { + let m = Self::as_map_observer(slf); + Python::with_gil(|py| -> PyResult { + let p: Py<_> = Py::new(py, m)?; + Ok($struct_name_trait::as_observer(p)) + }).unwrap() + } + + fn __getitem__(&self, idx: usize) -> $datatype { + *self.inner.get(idx) + } + + fn __setitem__(&mut self, idx: usize, val: $datatype) { + *self.inner.get_mut(idx) = val; + } + + #[pyo3(name = "usable_count")] + fn pyusable_count(&self) -> usize { + self.inner.usable_count() + } + + #[pyo3(name = "len")] + fn pylen(&self) -> usize { + self.inner.len() + } + + #[pyo3(name = "name")] + fn pyname(&self) -> &str { + self.inner.name() + } } - #[derive(Serialize, Deserialize, Debug, Clone)] - enum $wrapper_name { - Owned($struct_name), + #[derive(Debug, Clone)] + pub enum $wrapper_name { + Std(Py<$struct_name1>), + Owned(Py<$struct_name2>), + None + } + + impl Default for $wrapper_name { + fn default() -> Self { + $wrapper_name::None + } } // Should not be exposed to user @@ -1529,51 +1688,55 @@ pub mod pybind { #[derive(Serialize, Deserialize, Debug, Clone)] /// MapObserver + Observer Trait binding pub struct $struct_name_trait { + #[serde(skip)] pub wrapper: $wrapper_name, } - impl $struct_name_trait { - fn unwrap(&self) -> &impl MapObserver { - unsafe { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => &py_wrapper.inner, - } - } - } - - fn unwrap_mut(&mut self) -> &mut impl MapObserver { - unsafe { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => &mut py_wrapper.inner, - } - } - } - - fn upcast(&self) -> &impl Observer { - unsafe { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => &py_wrapper.inner, - } - } - } - - fn upcast_mut(&mut self) -> &mut impl Observer { - unsafe { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => &mut py_wrapper.inner, - } - } - } - } - #[pymethods] impl $struct_name_trait { #[staticmethod] - fn new_from_owned(owned_map: $struct_name) -> Self { + fn new_std(std_map: Py<$struct_name1>) -> Self { + Self { + wrapper: $wrapper_name::Std(std_map), + } + } + + #[staticmethod] + fn new_owned(owned_map: Py<$struct_name2>) -> Self { Self { wrapper: $wrapper_name::Owned(owned_map), } } + + #[must_use] + pub fn as_observer(slf: Py) -> PythonObserver { + concat_idents!(func = new_map_,$datatype { + PythonObserver::func(slf) + }) + } + + fn __getitem__(&self, idx: usize) -> $datatype { + *self.get(idx) + } + + fn __setitem__(&mut self, idx: usize, val: $datatype) { + *self.get_mut(idx) = val; + } + + #[pyo3(name = "usable_count")] + fn pyusable_count(&self) -> usize { + self.usable_count() + } + + #[pyo3(name = "len")] + fn pylen(&self) -> usize { + self.len() + } + + #[pyo3(name = "name")] + fn pyname(&self) -> &str { + self.name() + } } impl<'it> AsRefIterator<'it> for $struct_name_trait { @@ -1581,9 +1744,7 @@ pub mod pybind { type IntoIter = Iter<'it, $datatype>; fn as_ref_iter(&'it self) -> Self::IntoIter { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.as_ref_iter(), - } + mapob_unwrap_me!($wrapper_name, self.wrapper, m, { unsafe { std::mem::transmute::<_, Self::IntoIter>(m.as_ref_iter()) } }) } } @@ -1592,9 +1753,7 @@ pub mod pybind { type IntoIter = IterMut<'it, $datatype>; fn as_mut_iter(&'it mut self) -> Self::IntoIter { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.as_mut_iter(), - } + mapob_unwrap_me_mut!($wrapper_name, self.wrapper, m, { unsafe { std::mem::transmute::<_, Self::IntoIter>(m.as_mut_iter()) } }) } } @@ -1603,76 +1762,58 @@ pub mod pybind { #[inline] fn get(&self, idx: usize) -> &$datatype { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => &py_wrapper.inner.get(idx), - } + let ptr = mapob_unwrap_me!($wrapper_name, self.wrapper, m, { m.get(idx) as *const $datatype }); + unsafe { ptr.as_ref().unwrap() } } #[inline] fn get_mut(&mut self, idx: usize) -> &mut $datatype { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.get_mut(idx), - } + let ptr = mapob_unwrap_me_mut!($wrapper_name, self.wrapper, m, { m.get_mut(idx) as *mut $datatype }); + unsafe { ptr.as_mut().unwrap() } } #[inline] fn usable_count(&self) -> usize { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.wrapper.usable_count(), - } + mapob_unwrap_me!($wrapper_name, self.wrapper, m, { m.usable_count() }) } fn hash(&self) -> u64 { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.hash(), - } + mapob_unwrap_me!($wrapper_name, self.wrapper, m, { m.hash() }) } #[inline] fn initial(&self) -> $datatype { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.initial(), - } + mapob_unwrap_me!($wrapper_name, self.wrapper, m, { m.initial() }) } #[inline] fn initial_mut(&mut self) -> &mut $datatype { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.initial_mut(), - } + let ptr = mapob_unwrap_me_mut!($wrapper_name, self.wrapper, m, { m.initial_mut() as *mut $datatype }); + unsafe { ptr.as_mut().unwrap() } } #[inline] fn set_initial(&mut self, initial: $datatype) { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => { - py_wrapper.inner.set_initial(initial); - } - } + mapob_unwrap_me_mut!($wrapper_name, self.wrapper, m, { m.set_initial(initial) }); } fn to_vec(&self) -> Vec<$datatype> { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.to_vec(), - } + mapob_unwrap_me!($wrapper_name, self.wrapper, m, { m.to_vec() }) } } impl Named for $struct_name_trait { #[inline] fn name(&self) -> &str { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.name(), - } + let ptr = mapob_unwrap_me!($wrapper_name, self.wrapper, m, { m.name() as *const str }); + unsafe { ptr.as_ref().unwrap() } } } impl HasLen for $struct_name_trait { #[inline] fn len(&self) -> usize { - match &self.wrapper { - $wrapper_name::Owned(py_wrapper) => py_wrapper.inner.len(), - } + mapob_unwrap_me!($wrapper_name, self.wrapper, m, { m.len() }) } } @@ -1681,18 +1822,16 @@ pub mod pybind { Self: MapObserver, { #[inline] - fn pre_exec(&mut self, _state: &mut S, _input: &I) -> Result<(), Error> { - match &mut self.wrapper { - $wrapper_name::Owned(py_wrapper) => { - py_wrapper.inner.pre_exec(_state, _input) - } - } + fn pre_exec(&mut self, state: &mut S, input: &I) -> Result<(), Error> { + mapob_unwrap_me_mut!($wrapper_name, self.wrapper, m, { m.pre_exec(state, input) }) } } }; } define_python_map_observer!( + PythonStdMapObserverI8, + "StdMapObserverI8", PythonOwnedMapObserverI8, "OwnedMapObserverI8", PythonMapObserverI8, @@ -1701,6 +1840,8 @@ pub mod pybind { PythonMapObserverWrapperI8 ); define_python_map_observer!( + PythonStdMapObserverI16, + "StdMapObserverI16", PythonOwnedMapObserverI16, "OwnedMapObserverI16", PythonMapObserverI16, @@ -1709,6 +1850,8 @@ pub mod pybind { PythonMapObserverWrapperI16 ); define_python_map_observer!( + PythonStdMapObserverI32, + "StdMapObserverI32", PythonOwnedMapObserverI32, "OwnedMapObserverI32", PythonMapObserverI32, @@ -1717,6 +1860,8 @@ pub mod pybind { PythonMapObserverWrapperI32 ); define_python_map_observer!( + PythonStdMapObserverI64, + "StdMapObserverI64", PythonOwnedMapObserverI64, "OwnedMapObserverI64", PythonMapObserverI64, @@ -1726,6 +1871,8 @@ pub mod pybind { ); define_python_map_observer!( + PythonStdMapObserverU8, + "StdMapObserverU8", PythonOwnedMapObserverU8, "OwnedMapObserverU8", PythonMapObserverU8, @@ -1734,6 +1881,8 @@ pub mod pybind { PythonMapObserverWrapperU8 ); define_python_map_observer!( + PythonStdMapObserverU16, + "StdMapObserverU16", PythonOwnedMapObserverU16, "OwnedMapObserverU16", PythonMapObserverU16, @@ -1742,6 +1891,8 @@ pub mod pybind { PythonMapObserverWrapperU16 ); define_python_map_observer!( + PythonStdMapObserverU32, + "StdMapObserverU32", PythonOwnedMapObserverU32, "OwnedMapObserverU32", PythonMapObserverU32, @@ -1750,6 +1901,8 @@ pub mod pybind { PythonMapObserverWrapperU32 ); define_python_map_observer!( + PythonStdMapObserverU64, + "StdMapObserverU64", PythonOwnedMapObserverU64, "OwnedMapObserverU64", PythonMapObserverU64, @@ -1760,21 +1913,29 @@ pub mod pybind { /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/libafl/src/observers/mod.rs b/libafl/src/observers/mod.rs index f26099156d..65e7baab85 100644 --- a/libafl/src/observers/mod.rs +++ b/libafl/src/observers/mod.rs @@ -292,119 +292,769 @@ where /// `Observer` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { - use crate::bolts::tuples::Named; + use super::{Debug, Observer, ObserversTuple, String, Vec}; + use crate::bolts::tuples::{type_eq, MatchName, Named}; + use crate::executors::pybind::PythonExitKind; use crate::executors::ExitKind; use crate::inputs::BytesInput; - use crate::observers::Observer; + use crate::inputs::HasBytesVec; + use crate::observers::map::pybind::{ + PythonMapObserverI16, PythonMapObserverI32, PythonMapObserverI64, PythonMapObserverI8, + PythonMapObserverU16, PythonMapObserverU32, PythonMapObserverU64, PythonMapObserverU8, + PythonMapObserverWrapperI16, PythonMapObserverWrapperI32, PythonMapObserverWrapperI64, + PythonMapObserverWrapperI8, PythonMapObserverWrapperU16, PythonMapObserverWrapperU32, + PythonMapObserverWrapperU64, PythonMapObserverWrapperU8, + }; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; use crate::Error; use pyo3::prelude::*; + use serde::{Deserialize, Serialize}; + use std::cell::UnsafeCell; - macro_rules! define_python_observer { - ($struct_name_trait:ident, $py_name_trait:tt, $wrapper_name: ident, $my_std_state_type_name: ident) => { - use crate::observers::map::pybind::PythonMapObserverI8; - use crate::state::pybind::$my_std_state_type_name; + #[derive(Debug)] + pub struct PyObjectObserver { + inner: PyObject, + name: UnsafeCell, + } - #[derive(Debug)] - enum $wrapper_name { - MapI8(*mut PythonMapObserverI8), + impl Clone for PyObjectObserver { + fn clone(&self) -> PyObjectObserver { + PyObjectObserver { + inner: self.inner.clone(), + name: UnsafeCell::new(String::new()), } + } + } - #[pyclass(unsendable, name = $py_name_trait)] - #[derive(Debug)] - /// Observer Trait binding - pub struct $struct_name_trait { - pub wrapper: $wrapper_name, + impl PyObjectObserver { + #[must_use] + pub fn new(obj: PyObject) -> Self { + PyObjectObserver { + inner: obj, + name: UnsafeCell::new(String::new()), } + } + } - impl $struct_name_trait { - fn unwrap(&self) -> &impl Observer { - unsafe { - match self.wrapper { - $wrapper_name::MapI8(py_wrapper) => &(*py_wrapper).upcast(), - } - } - } + crate::impl_serde_pyobjectwrapper!(PyObjectObserver, inner); - fn unwrap_mut( - &mut self, - ) -> &mut impl Observer { - unsafe { - match self.wrapper { - $wrapper_name::MapI8(py_wrapper) => &mut (*py_wrapper).upcast_mut(), - } - } - } + impl Named for PyObjectObserver { + fn name(&self) -> &str { + let s = Python::with_gil(|py| -> PyResult { + let s: String = self.inner.call_method0(py, "name")?.extract(py)?; + Ok(s) + }) + .unwrap(); + unsafe { + *self.name.get() = s; + &*self.name.get() } + } + } - #[pymethods] - impl $struct_name_trait { - #[staticmethod] - fn new_map(map_observer: &mut PythonMapObserverI8) -> Self { - Self { - observer: $wrapper_name::MapI8(map_observer), - } + impl Observer for PyObjectObserver { + fn flush(&mut self) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method0(py, "flush")?; + Ok(()) + }) + .unwrap(); + Ok(()) + } + + fn pre_exec( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "pre_exec", + (PythonStdStateWrapper::wrap(state), input.bytes()), + )?; + Ok(()) + })?; + Ok(()) + } + + fn post_exec( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + exit_kind: &ExitKind, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "post_exec", + ( + PythonStdStateWrapper::wrap(state), + input.bytes(), + PythonExitKind::from(*exit_kind), + ), + )?; + Ok(()) + })?; + Ok(()) + } + + fn pre_exec_child( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "pre_exec_child", + (PythonStdStateWrapper::wrap(state), input.bytes()), + )?; + Ok(()) + })?; + Ok(()) + } + + fn post_exec_child( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + exit_kind: &ExitKind, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "post_exec_child", + ( + PythonStdStateWrapper::wrap(state), + input.bytes(), + PythonExitKind::from(*exit_kind), + ), + )?; + Ok(()) + })?; + Ok(()) + } + } + + #[derive(Serialize, Deserialize, Clone, Debug)] + pub enum PythonObserverWrapper { + MapI8(Py), + MapI16(Py), + MapI32(Py), + MapI64(Py), + MapU8(Py), + MapU16(Py), + MapU32(Py), + MapU64(Py), + Python(PyObjectObserver), + } + + #[pyclass(unsendable, name = "Observer")] + #[derive(Serialize, Deserialize, Clone, Debug)] + /// Observer Trait binding + pub struct PythonObserver { + pub wrapper: PythonObserverWrapper, + } + + macro_rules! unwrap_me { + ($wrapper:expr, $name:ident, $body:block) => { + match &$wrapper { + PythonObserverWrapper::MapI8(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperI8, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap(), + PythonObserverWrapper::MapI16(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperI16, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() } - } - - impl Named for $struct_name_trait { - fn name(&self) -> &str { - self.unwrap().name() + PythonObserverWrapper::MapI32(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperI32, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() } - } - - impl Observer for $struct_name_trait { - fn flush(&mut self) -> Result<(), Error> { - self.unwrap_mut().flush() + PythonObserverWrapper::MapI64(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperI64, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() } - - fn pre_exec( - &mut self, - state: &mut $my_std_state_type_name, - input: &BytesInput, - ) -> Result<(), Error> { - self.unwrap_mut().pre_exec(state, input) + PythonObserverWrapper::MapU8(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperU8, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap(), + PythonObserverWrapper::MapU16(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperU16, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() } - - fn post_exec( - &mut self, - state: &mut $my_std_state_type_name, - input: &BytesInput, - exit_kind: &ExitKind, - ) -> Result<(), Error> { - self.unwrap_mut().post_exec(state, input, exit_kind) + PythonObserverWrapper::MapU32(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperU32, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() } - - fn pre_exec_child( - &mut self, - state: &mut $my_std_state_type_name, - input: &BytesInput, - ) -> Result<(), Error> { - self.unwrap_mut().pre_exec(state, input) + PythonObserverWrapper::MapU64(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let borrowed = py_wrapper.borrow(py); + Ok(crate::mapob_unwrap_me!( + PythonMapObserverWrapperU64, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() } - - fn post_exec_child( - &mut self, - state: &mut $my_std_state_type_name, - input: &BytesInput, - exit_kind: &ExitKind, - ) -> Result<(), Error> { - self.unwrap_mut().post_exec_child(state, input, exit_kind) + PythonObserverWrapper::Python(py_wrapper) => { + let $name = py_wrapper; + $body } } }; } - define_python_observer!( - PythonObserver, - "Observer", - PythonObserverWrapper, - PythonStdState, - ); + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + match &mut $wrapper { + PythonObserverWrapper::MapI8(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperI8, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap(), + PythonObserverWrapper::MapI16(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperI16, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() + } + PythonObserverWrapper::MapI32(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperI32, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() + } + PythonObserverWrapper::MapI64(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperI64, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() + } + PythonObserverWrapper::MapU8(py_wrapper) => Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperU8, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap(), + PythonObserverWrapper::MapU16(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperU16, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() + } + PythonObserverWrapper::MapU32(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperU32, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() + } + PythonObserverWrapper::MapU64(py_wrapper) => { + Python::with_gil(|py| -> PyResult<_> { + let mut borrowed = py_wrapper.borrow_mut(py); + Ok(crate::mapob_unwrap_me_mut!( + PythonMapObserverWrapperU64, + borrowed.wrapper, + $name, + $body + )) + }) + .unwrap() + } + PythonObserverWrapper::Python(py_wrapper) => { + let $name = py_wrapper; + $body + } + } + }; + } + + #[pymethods] + impl PythonObserver { + #[staticmethod] + #[must_use] + pub fn new_map_i8(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapI8(map_observer), + } + } + #[staticmethod] + #[must_use] + pub fn new_map_i16(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapI16(map_observer), + } + } + #[staticmethod] + #[must_use] + pub fn new_map_i32(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapI32(map_observer), + } + } + #[staticmethod] + #[must_use] + pub fn new_map_i64(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapI64(map_observer), + } + } + + #[staticmethod] + #[must_use] + pub fn new_map_u8(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapU8(map_observer), + } + } + #[staticmethod] + #[must_use] + pub fn new_map_u16(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapU16(map_observer), + } + } + #[staticmethod] + #[must_use] + pub fn new_map_u32(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapU32(map_observer), + } + } + #[staticmethod] + #[must_use] + pub fn new_map_u64(map_observer: Py) -> Self { + Self { + wrapper: PythonObserverWrapper::MapU64(map_observer), + } + } + #[staticmethod] + #[must_use] + pub fn new_py(py_observer: PyObject) -> Self { + Self { + wrapper: PythonObserverWrapper::Python(PyObjectObserver::new(py_observer)), + } + } + + pub fn unwrap_py(&self) -> Option { + match &self.wrapper { + PythonObserverWrapper::Python(pyo) => Some(pyo.inner.clone()), + _ => None, + } + } + } + + impl Named for PythonObserver { + fn name(&self) -> &str { + let ptr = unwrap_me!(self.wrapper, o, { o.name() as *const str }); + unsafe { ptr.as_ref().unwrap() } + } + } + + impl Observer for PythonObserver { + fn flush(&mut self) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, o, { + Observer::::flush(o) + }) + } + + fn pre_exec( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, o, { o.pre_exec(state, input) }) + } + + fn post_exec( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + exit_kind: &ExitKind, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, o, { o.post_exec(state, input, exit_kind) }) + } + + fn pre_exec_child( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, o, { o.pre_exec_child(state, input) }) + } + + fn post_exec_child( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + exit_kind: &ExitKind, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, o, { + o.post_exec_child(state, input, exit_kind) + }) + } + } + + #[derive(Serialize, Deserialize, Clone, Debug)] + #[pyclass(unsendable, name = "ObserversTuple")] + pub struct PythonObserversTuple { + list: Vec, + } + + #[pymethods] + impl PythonObserversTuple { + #[new] + fn new(list: Vec) -> Self { + Self { list } + } + + fn len(&self) -> usize { + self.list.len() + } + + fn __getitem__(&self, idx: usize) -> PythonObserver { + self.list[idx].clone() + } + + #[pyo3(name = "match_name")] + fn pymatch_name(&self, name: &str) -> Option { + for ob in &self.list { + if ob.name() == name { + return Some(ob.clone()); + } + } + None + } + } + + impl ObserversTuple for PythonObserversTuple { + fn pre_exec_all( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + for ob in &mut self.list { + ob.pre_exec(state, input)?; + } + Ok(()) + } + + fn post_exec_all( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + exit_kind: &ExitKind, + ) -> Result<(), Error> { + for ob in &mut self.list { + ob.post_exec(state, input, exit_kind)?; + } + Ok(()) + } + + fn pre_exec_child_all( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + ) -> Result<(), Error> { + for ob in &mut self.list { + ob.pre_exec_child(state, input)?; + } + Ok(()) + } + + fn post_exec_child_all( + &mut self, + state: &mut PythonStdState, + input: &BytesInput, + exit_kind: &ExitKind, + ) -> Result<(), Error> { + for ob in &mut self.list { + ob.post_exec_child(state, input, exit_kind)?; + } + Ok(()) + } + } + + impl MatchName for PythonObserversTuple { + fn match_name(&self, name: &str) -> Option<&T> { + unsafe { + let mut r = None; + for ob in &self.list { + Python::with_gil(|py| -> PyResult<_> { + match &ob.wrapper { + PythonObserverWrapper::MapI8(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + PythonObserverWrapper::MapI16(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + PythonObserverWrapper::MapI32(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + PythonObserverWrapper::MapI64(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + + PythonObserverWrapper::MapU8(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + PythonObserverWrapper::MapU16(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + PythonObserverWrapper::MapU32(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + PythonObserverWrapper::MapU64(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow(py)) as *const T) + .as_ref(); + } + } + PythonObserverWrapper::Python(py_wrapper) => { + if type_eq::() && py_wrapper.name() == name { + r = (py_wrapper as *const _ as *const T).as_ref(); + } + } + } + Ok(()) + }) + .unwrap(); + } + r + } + } + + fn match_name_mut(&mut self, name: &str) -> Option<&mut T> { + unsafe { + let mut r = None; + for ob in &mut self.list { + Python::with_gil(|py| -> PyResult<_> { + match &mut ob.wrapper { + PythonObserverWrapper::MapI8(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + PythonObserverWrapper::MapI16(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + PythonObserverWrapper::MapI32(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + PythonObserverWrapper::MapI64(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + + PythonObserverWrapper::MapU8(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + PythonObserverWrapper::MapU16(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + PythonObserverWrapper::MapU32(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + PythonObserverWrapper::MapU64(py_wrapper) => { + if type_eq::() + && py_wrapper.borrow(py).name() == name + { + r = (std::ptr::addr_of!(*(*py_wrapper).borrow_mut(py)) + as *mut T) + .as_mut(); + } + } + PythonObserverWrapper::Python(py_wrapper) => { + if type_eq::() && py_wrapper.name() == name { + r = (py_wrapper as *mut _ as *mut T).as_mut(); + } + } + } + Ok(()) + }) + .unwrap(); + } + r + } + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { m.add_class::()?; + m.add_class::()?; Ok(()) } } diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 1373632ab1..9f9584c411 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -259,187 +259,174 @@ where /// `Stage` Python bindings #[cfg(feature = "python")] +#[allow(missing_docs)] pub mod pybind { - use crate::impl_asany; - use crate::stages::Stage; + use crate::events::pybind::PythonEventManager; + use crate::executors::pybind::PythonExecutor; + use crate::fuzzer::pybind::{PythonStdFuzzer, PythonStdFuzzerWrapper}; + use crate::stages::mutational::pybind::PythonStdMutationalStage; + use crate::stages::{Stage, StagesTuple}; + use crate::state::pybind::{PythonStdState, PythonStdStateWrapper}; use crate::Error; use pyo3::prelude::*; - use super::owned::AnyStage; + #[derive(Clone, Debug)] + pub struct PyObjectStage { + inner: PyObject, + } - macro_rules! define_python_stage { - ($struct_name_trait:ident, $py_name_trait:tt, $wrapper_name: ident, $std_havoc_mutations_stage_name: ident, $my_std_state_type_name: ident, - $my_std_fuzzer_type_name: ident, $executor_name: ident, $event_manager_name: ident) => { - use crate::events::pybind::$event_manager_name; - use crate::executors::pybind::$executor_name; - use crate::fuzzer::pybind::$my_std_fuzzer_type_name; - use crate::stages::mutational::pybind::$std_havoc_mutations_stage_name; - use crate::state::pybind::$my_std_state_type_name; + impl PyObjectStage { + #[must_use] + pub fn new(obj: PyObject) -> Self { + PyObjectStage { inner: obj } + } + } - #[derive(Debug)] - enum $wrapper_name { - StdHavocMutations(*mut $std_havoc_mutations_stage_name), - } + impl Stage for PyObjectStage { + #[inline] + fn perform( + &mut self, + fuzzer: &mut PythonStdFuzzer, + executor: &mut PythonExecutor, + state: &mut PythonStdState, + manager: &mut PythonEventManager, + corpus_idx: usize, + ) -> Result<(), Error> { + Python::with_gil(|py| -> PyResult<()> { + self.inner.call_method1( + py, + "perform", + ( + PythonStdFuzzerWrapper::wrap(fuzzer), + executor.clone(), + PythonStdStateWrapper::wrap(state), + manager.clone(), + corpus_idx, + ), + )?; + Ok(()) + })?; + Ok(()) + } + } - /// Stage Trait binding - #[pyclass(unsendable, name = $py_name_trait)] - #[derive(Debug)] - pub struct $struct_name_trait { - stage: $wrapper_name, - } + #[derive(Clone, Debug)] + pub enum PythonStageWrapper { + StdMutational(Py), + Python(PyObjectStage), + } - #[pymethods] - impl $struct_name_trait { - #[staticmethod] - fn new_from_std_scheduled( - py_std_havoc_mutations_stage: &mut $std_havoc_mutations_stage_name, - ) -> Self { - Self { - stage: $wrapper_name::StdHavocMutations(py_std_havoc_mutations_stage), + /// Stage Trait binding + #[pyclass(unsendable, name = "Stage")] + #[derive(Clone, Debug)] + pub struct PythonStage { + wrapper: PythonStageWrapper, + } + + macro_rules! unwrap_me_mut { + ($wrapper:expr, $name:ident, $body:block) => { + crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonStageWrapper, + { StdMutational }, + { + Python(py_wrapper) => { + let $name = py_wrapper; + $body } } - } - - impl - Stage< - $executor_name, - $event_manager_name, - $my_std_state_type_name, - $my_std_fuzzer_type_name, - > for $struct_name_trait - { - #[inline] - #[allow(clippy::let_and_return)] - fn perform( - &mut self, - fuzzer: &mut $my_std_fuzzer_type_name, - executor: &mut $executor_name, - state: &mut $my_std_state_type_name, - manager: &mut $event_manager_name, - corpus_idx: usize, - ) -> Result<(), Error> { - unsafe { - match self.stage { - $wrapper_name::StdHavocMutations(py_std_havoc_mutations_stage) => { - (*py_std_havoc_mutations_stage) - .std_mutational_stage - .perform(fuzzer, executor, state, manager, corpus_idx) - } - } - } - } - } - - impl_asany!($struct_name_trait); - - impl - AnyStage< - $executor_name, - $event_manager_name, - $my_std_state_type_name, - $my_std_fuzzer_type_name, - > for $struct_name_trait - { - } + ) }; } - define_python_stage!( - PythonStageI8, - "StageI8", - PythonStageWrapperI8, - PythonStdScheduledHavocMutationsStageI8, - MyStdStateI8, - MyStdFuzzerI8, - PythonExecutorI8, - PythonEventManagerI8 - ); + #[pymethods] + impl PythonStage { + #[staticmethod] + #[must_use] + pub fn new_std_mutational( + py_std_havoc_mutations_stage: Py, + ) -> Self { + Self { + wrapper: PythonStageWrapper::StdMutational(py_std_havoc_mutations_stage), + } + } - define_python_stage!( - PythonStageI16, - "StageI16", - PythonStageWrapperI16, - PythonStdScheduledHavocMutationsStageI16, - MyStdStateI16, - MyStdFuzzerI16, - PythonExecutorI16, - PythonEventManagerI16 - ); + #[staticmethod] + #[must_use] + pub fn new_py(obj: PyObject) -> Self { + Self { + wrapper: PythonStageWrapper::Python(PyObjectStage::new(obj)), + } + } - define_python_stage!( - PythonStageI32, - "StageI32", - PythonStageWrapperI32, - PythonStdScheduledHavocMutationsStageI32, - MyStdStateI32, - MyStdFuzzerI32, - PythonExecutorI32, - PythonEventManagerI32 - ); + #[must_use] + pub fn unwrap_py(&self) -> Option { + match &self.wrapper { + PythonStageWrapper::Python(pyo) => Some(pyo.inner.clone()), + PythonStageWrapper::StdMutational(_) => None, + } + } + } - define_python_stage!( - PythonStageI64, - "StageI64", - PythonStageWrapperI64, - PythonStdScheduledHavocMutationsStageI64, - MyStdStateI64, - MyStdFuzzerI64, - PythonExecutorI64, - PythonEventManagerI64 - ); + impl Stage for PythonStage { + #[inline] + #[allow(clippy::let_and_return)] + fn perform( + &mut self, + fuzzer: &mut PythonStdFuzzer, + executor: &mut PythonExecutor, + state: &mut PythonStdState, + manager: &mut PythonEventManager, + corpus_idx: usize, + ) -> Result<(), Error> { + unwrap_me_mut!(self.wrapper, s, { + s.perform(fuzzer, executor, state, manager, corpus_idx) + }) + } + } - define_python_stage!( - PythonStageU8, - "StageU8", - PythonStageWrapperU8, - PythonStdScheduledHavocMutationsStageU8, - MyStdStateU8, - MyStdFuzzerU8, - PythonExecutorU8, - PythonEventManagerU8 - ); - define_python_stage!( - PythonStageU16, - "StageU16", - PythonStageWrapperU16, - PythonStdScheduledHavocMutationsStageU16, - MyStdStateU16, - MyStdFuzzerU16, - PythonExecutorU16, - PythonEventManagerU16 - ); - define_python_stage!( - PythonStageU32, - "StageU32", - PythonStageWrapperU32, - PythonStdScheduledHavocMutationsStageU32, - MyStdStateU32, - MyStdFuzzerU32, - PythonExecutorU32, - PythonEventManagerU32 - ); - define_python_stage!( - PythonStageU64, - "StageU64", - PythonStageWrapperU64, - PythonStdScheduledHavocMutationsStageU64, - MyStdStateU64, - MyStdFuzzerU64, - PythonExecutorU64, - PythonEventManagerU64 - ); + #[derive(Clone, Debug)] + #[pyclass(unsendable, name = "StagesTuple")] + pub struct PythonStagesTuple { + list: Vec, + } + + #[pymethods] + impl PythonStagesTuple { + #[new] + fn new(list: Vec) -> Self { + Self { list } + } + + fn len(&self) -> usize { + self.list.len() + } + + fn __getitem__(&self, idx: usize) -> PythonStage { + self.list[idx].clone() + } + } + + impl StagesTuple + for PythonStagesTuple + { + fn perform_all( + &mut self, + fuzzer: &mut PythonStdFuzzer, + executor: &mut PythonExecutor, + state: &mut PythonStdState, + manager: &mut PythonEventManager, + corpus_idx: usize, + ) -> Result<(), Error> { + for s in &mut self.list { + s.perform(fuzzer, executor, state, manager, corpus_idx)?; + } + Ok(()) + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } } diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index 1513756f10..6c364e80b8 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -163,161 +163,51 @@ where } #[cfg(feature = "python")] +#[allow(missing_docs)] /// `StdMutationalStage` Python bindings pub mod pybind { - use crate::bolts::tuples::tuple_list_type; + use crate::events::pybind::PythonEventManager; + use crate::executors::pybind::PythonExecutor; + use crate::fuzzer::pybind::PythonStdFuzzer; use crate::inputs::BytesInput; - pub use crate::mutators::mutations::*; - pub use crate::mutators::mutations::*; - use crate::mutators::{havoc_mutations, StdScheduledMutator}; + use crate::mutators::pybind::PythonMutator; + use crate::stages::pybind::PythonStage; use crate::stages::StdMutationalStage; + use crate::state::pybind::PythonStdState; use pyo3::prelude::*; - type HavocMutationsType = tuple_list_type!( - BitFlipMutator, - ByteFlipMutator, - ByteIncMutator, - ByteDecMutator, - ByteNegMutator, - ByteRandMutator, - ByteAddMutator, - WordAddMutator, - DwordAddMutator, - QwordAddMutator, - ByteInterestingMutator, - WordInterestingMutator, - DwordInterestingMutator, - BytesDeleteMutator, - BytesDeleteMutator, - BytesDeleteMutator, - BytesDeleteMutator, - BytesExpandMutator, - BytesInsertMutator, - BytesRandInsertMutator, - BytesSetMutator, - BytesRandSetMutator, - BytesCopyMutator, - BytesInsertCopyMutator, - BytesSwapMutator, - CrossoverInsertMutator, - CrossoverReplaceMutator, - ); - - macro_rules! define_python_std_mutational_stage { - ($struct_name:ident, $py_name:tt, $my_std_state_type_name: ident, $my_std_fuzzer_type_name: ident, $executor_name: ident, $event_manager_name: ident) => { - use crate::events::pybind::$event_manager_name; - use crate::executors::pybind::$executor_name; - use crate::fuzzer::pybind::$my_std_fuzzer_type_name; - use crate::state::pybind::$my_std_state_type_name; - - #[pyclass(unsendable, name = $py_name)] - #[derive(Debug)] - /// Python class for StdMutationalStage - pub struct $struct_name { - /// Rust wrapped StdMutationalStage object - pub std_mutational_stage: StdMutationalStage< - $executor_name, - $event_manager_name, - BytesInput, - StdScheduledMutator, - $my_std_state_type_name, - $my_std_fuzzer_type_name, - >, - } - - #[pymethods] - impl $struct_name { - #[staticmethod] - fn new_from_scheduled_havoc_mutations() -> Self { - Self { - std_mutational_stage: StdMutationalStage::new(StdScheduledMutator::new( - havoc_mutations(), - )), - } - } - } - }; + #[pyclass(unsendable, name = "StdMutationalStage")] + #[derive(Debug)] + /// Python class for StdMutationalStage + pub struct PythonStdMutationalStage { + /// Rust wrapped StdMutationalStage object + pub inner: StdMutationalStage< + PythonExecutor, + PythonEventManager, + BytesInput, + PythonMutator, + PythonStdState, + PythonStdFuzzer, + >, } - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageI8, - "StdScheduledHavocMutationsStageI8", - MyStdStateI8, - MyStdFuzzerI8, - PythonExecutorI8, - PythonEventManagerI8 - ); + #[pymethods] + impl PythonStdMutationalStage { + #[new] + fn new(mutator: PythonMutator) -> Self { + Self { + inner: StdMutationalStage::new(mutator), + } + } - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageI16, - "StdScheduledHavocMutationsStageI16", - MyStdStateI16, - MyStdFuzzerI16, - PythonExecutorI16, - PythonEventManagerI16 - ); - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageI32, - "StdScheduledHavocMutationsStageI32", - MyStdStateI32, - MyStdFuzzerI32, - PythonExecutorI32, - PythonEventManagerI32 - ); - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageI64, - "StdScheduledHavocMutationsStageI64", - MyStdStateI64, - MyStdFuzzerI64, - PythonExecutorI64, - PythonEventManagerI64 - ); - - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageU8, - "StdScheduledHavocMutationsStageU8", - MyStdStateU8, - MyStdFuzzerU8, - PythonExecutorU8, - PythonEventManagerU8 - ); - - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageU16, - "StdScheduledHavocMutationsStageU16", - MyStdStateU16, - MyStdFuzzerU16, - PythonExecutorU16, - PythonEventManagerU16 - ); - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageU32, - "StdScheduledHavocMutationsStageU32", - MyStdStateU32, - MyStdFuzzerU32, - PythonExecutorU32, - PythonEventManagerU32 - ); - define_python_std_mutational_stage!( - PythonStdScheduledHavocMutationsStageU64, - "StdScheduledHavocMutationsStageU64", - MyStdStateU64, - MyStdFuzzerU64, - PythonExecutorU64, - PythonEventManagerU64 - ); + fn as_stage(slf: Py) -> PythonStage { + PythonStage::new_std_mutational(slf) + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; Ok(()) } } diff --git a/libafl/src/stages/owned.rs b/libafl/src/stages/owned.rs index 5dc2a68b29..efa4f5841b 100644 --- a/libafl/src/stages/owned.rs +++ b/libafl/src/stages/owned.rs @@ -42,147 +42,3 @@ impl StagesOwnedList { Self { list } } } - -#[cfg(feature = "python")] -/// `StagesOwnedList` Python bindings -pub mod pybind { - use crate::stages::owned::StagesOwnedList; - use pyo3::prelude::*; - - macro_rules! define_python_stage_owned_list { - ($struct_name:ident, $py_name:tt, $my_std_state_type_name: ident, $my_std_fuzzer_type_name: ident, $event_manager_name: ident, - $executor_name: ident, $stage_name: ident) => { - use crate::events::pybind::$event_manager_name; - use crate::executors::pybind::$executor_name; - use crate::fuzzer::pybind::$my_std_fuzzer_type_name; - use crate::stages::pybind::$stage_name; - use crate::state::pybind::$my_std_state_type_name; - #[pyclass(unsendable, name = $py_name)] - - /// Python class for StagesOwnedList - #[allow(missing_debug_implementations)] - pub struct $struct_name { - /// Rust wrapped StagesOwnedList object - pub stages_owned_list: StagesOwnedList< - $executor_name, - $event_manager_name, - $my_std_state_type_name, - $my_std_fuzzer_type_name, - >, - } - - #[pymethods] - impl $struct_name { - //TODO: Add new from list - #[new] - fn new(stage: &$stage_name) -> Self { - // TODO: Be safe - unsafe { - Self { - stages_owned_list: StagesOwnedList { - list: vec![Box::new(std::mem::transmute_copy::< - $stage_name, - $stage_name, - >(stage))], - }, - } - } - } - } - }; - } - - define_python_stage_owned_list!( - PythonStagesOwnedListI8, - "StagesOwnedListI8", - MyStdStateI8, - MyStdFuzzerI8, - PythonEventManagerI8, - PythonExecutorI8, - PythonStageI8 - ); - - define_python_stage_owned_list!( - PythonStagesOwnedListI16, - "StagesOwnedListI16", - MyStdStateI16, - MyStdFuzzerI16, - PythonEventManagerI16, - PythonExecutorI16, - PythonStageI16 - ); - - define_python_stage_owned_list!( - PythonStagesOwnedListI32, - "StagesOwnedListI32", - MyStdStateI32, - MyStdFuzzerI32, - PythonEventManagerI32, - PythonExecutorI32, - PythonStageI32 - ); - - define_python_stage_owned_list!( - PythonStagesOwnedListI64, - "StagesOwnedListI64", - MyStdStateI64, - MyStdFuzzerI64, - PythonEventManagerI64, - PythonExecutorI64, - PythonStageI64 - ); - - define_python_stage_owned_list!( - PythonStagesOwnedListU8, - "StagesOwnedListU8", - MyStdStateU8, - MyStdFuzzerU8, - PythonEventManagerU8, - PythonExecutorU8, - PythonStageU8 - ); - - define_python_stage_owned_list!( - PythonStagesOwnedListU16, - "StagesOwnedListU16", - MyStdStateU16, - MyStdFuzzerU16, - PythonEventManagerU16, - PythonExecutorU16, - PythonStageU16 - ); - - define_python_stage_owned_list!( - PythonStagesOwnedListU32, - "StagesOwnedListU32", - MyStdStateU32, - MyStdFuzzerU32, - PythonEventManagerU32, - PythonExecutorU32, - PythonStageU32 - ); - - define_python_stage_owned_list!( - PythonStagesOwnedListU64, - "StagesOwnedListU64", - MyStdStateU64, - MyStdFuzzerU64, - PythonEventManagerU64, - PythonExecutorU64, - PythonStageU64 - ); - - /// Register the classes to the python module - pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - Ok(()) - } -} diff --git a/libafl/src/state/mod.rs b/libafl/src/state/mod.rs index ec3a061414..f084fc094e 100644 --- a/libafl/src/state/mod.rs +++ b/libafl/src/state/mod.rs @@ -652,72 +652,137 @@ where } #[cfg(feature = "python")] +#[allow(missing_docs)] /// `State` Python bindings pub mod pybind { + use crate::bolts::ownedref::OwnedPtrMut; use crate::bolts::rands::pybind::PythonRand; - use crate::bolts::tuples::tuple_list; use crate::corpus::pybind::PythonCorpus; + use crate::events::pybind::PythonEventManager; + use crate::executors::pybind::PythonExecutor; use crate::feedbacks::pybind::PythonFeedback; + use crate::fuzzer::pybind::PythonStdFuzzerWrapper; + use crate::generators::pybind::PythonGenerator; use crate::inputs::BytesInput; - use crate::state::StdState; + use crate::pybind::PythonMetadata; + use crate::state::{ + HasCorpus, HasExecutions, HasMaxSize, HasMetadata, HasRand, HasSolutions, StdState, + }; use pyo3::prelude::*; + use pyo3::types::PyDict; + use std::path::PathBuf; - macro_rules! define_python_state { - ($type_name:ident, $struct_name:ident, $py_name:tt) => { - use crate::events::pybind::$event_manager_name; - use crate::executors::pybind::$executor_name; - use crate::fuzzer::pybind::$fuzzer_name; - use crate::generators::pybind::$rand_printable_generator; + /// `StdState` with fixed generics + pub type PythonStdState = StdState; - /// `StdState` with fixed generics - pub type $type_name = StdState; - - #[pyclass(unsendable, name = $py_name)] - #[derive(Debug)] - /// Python class for StdState - pub struct $struct_name { - /// Rust wrapped StdState object - pub std_state: $type_name, - } - - #[pymethods] - impl $struct_name { - #[new] - fn new( - py_rand: PythonRand, - corpus: PythonCorpus, - solutions: PythonCorpus, - feedback: &mut PythonFeedback, - objective: &mut PythonFeedback, - ) -> Self { - Self { - std_state: StdState::new(py_rand, corpus, solutions, feedback, objective), - } - } - - fn generate_initial_inputs( - &mut self, - py_fuzzer: &mut PythonFuzzer, - py_executor: &mut PythonExecutor, - py_generator: &mut PythonGenerator, - py_mgr: &mut PythonEventManager, - num: usize, - ) { - self.std_state - .generate_initial_inputs( - &mut py_fuzzer.std_fuzzer, - py_executor, - &mut py_generator.rand_printable_generator, - py_mgr, - num, - ) - .expect("Failed to generate the initial corpus".into()); - } - } - }; + #[pyclass(unsendable, name = "StdState")] + #[derive(Debug)] + /// Python class for StdState + pub struct PythonStdStateWrapper { + /// Rust wrapped StdState object + pub inner: OwnedPtrMut, } - define_python_state!(PythonStdState, PythonStdStateWrapper, "StdState",); + impl PythonStdStateWrapper { + pub fn wrap(r: &mut PythonStdState) -> Self { + Self { + inner: OwnedPtrMut::Ptr(r), + } + } + + #[must_use] + pub fn unwrap(&self) -> &PythonStdState { + self.inner.as_ref() + } + + pub fn unwrap_mut(&mut self) -> &mut PythonStdState { + self.inner.as_mut() + } + } + + #[pymethods] + impl PythonStdStateWrapper { + #[new] + fn new( + py_rand: PythonRand, + corpus: PythonCorpus, + solutions: PythonCorpus, + feedback: &mut PythonFeedback, + objective: &mut PythonFeedback, + ) -> Self { + Self { + inner: OwnedPtrMut::Owned(Box::new( + StdState::new(py_rand, corpus, solutions, feedback, objective) + .expect("Failed to create a new StdState"), + )), + } + } + + fn metadata(&mut self) -> PyObject { + let meta = self.inner.as_mut().metadata_mut(); + if !meta.contains::() { + Python::with_gil(|py| { + let dict: Py = PyDict::new(py).into(); + meta.insert(PythonMetadata::new(dict.to_object(py))); + }); + } + meta.get::().unwrap().map.clone() + } + + fn rand(&self) -> PythonRand { + self.inner.as_ref().rand().clone() + } + + fn corpus(&self) -> PythonCorpus { + self.inner.as_ref().corpus().clone() + } + + fn solutions(&self) -> PythonCorpus { + self.inner.as_ref().solutions().clone() + } + + fn executions(&self) -> usize { + *self.inner.as_ref().executions() + } + + fn max_size(&self) -> usize { + self.inner.as_ref().max_size() + } + + fn generate_initial_inputs( + &mut self, + py_fuzzer: &mut PythonStdFuzzerWrapper, + py_executor: &mut PythonExecutor, + py_generator: &mut PythonGenerator, + py_mgr: &mut PythonEventManager, + num: usize, + ) { + self.inner + .as_mut() + .generate_initial_inputs( + py_fuzzer.unwrap_mut(), + py_executor, + py_generator, + py_mgr, + num, + ) + .expect("Failed to generate the initial corpus"); + } + + #[allow(clippy::needless_pass_by_value)] + fn load_initial_inputs( + &mut self, + py_fuzzer: &mut PythonStdFuzzerWrapper, + py_executor: &mut PythonExecutor, + py_mgr: &mut PythonEventManager, + in_dirs: Vec, + ) { + self.inner + .as_mut() + .load_initial_inputs(py_fuzzer.unwrap_mut(), py_executor, py_mgr, &in_dirs) + .expect("Failed to load the initial corpus"); + } + } /// Register the classes to the python module pub fn register(_py: Python, m: &PyModule) -> PyResult<()> { diff --git a/libafl_frida/src/coverage_rt.rs b/libafl_frida/src/coverage_rt.rs index 016c770572..b9d2df51f8 100644 --- a/libafl_frida/src/coverage_rt.rs +++ b/libafl_frida/src/coverage_rt.rs @@ -3,9 +3,6 @@ use core::ptr::addr_of_mut; use dynasmrt::{dynasm, DynasmApi, DynasmLabelApi}; use rangemap::RangeMap; -#[cfg(target_arch = "aarch64")] -use std::ffi::c_void; - #[cfg(target_arch = "x86_64")] use frida_gum::instruction_writer::X86Register; #[cfg(target_arch = "aarch64")]