smite-ir: add program minimizers and wire them into AFL++ custom trim

README.md

@@ -66,13 +66,11 @@ printf '\x00' > /tmp/smite-seeds/empty
 AFL_CUSTOM_MUTATOR_LIBRARY=target/release/libsmite_ir_mutator.so \
 AFL_CUSTOM_MUTATOR_ONLY=1 \
 AFL_FRAMESHIFT_DISABLE=1 \
-AFL_DISABLE_TRIM=1 \
 ~/AFLplusplus/afl-fuzz -X -i /tmp/smite-seeds -o /tmp/smite-out -- /tmp/smite-nyx
 ```
 
 `AFL_CUSTOM_MUTATOR_ONLY=1` disables AFL++'s built-in mutators (which would
-corrupt the postcard encoding). `AFL_DISABLE_TRIM=1` prevents AFL++ from
-trimming inputs (which would also corrupt the encoding).
+corrupt the postcard encoding).
 
 ## Running Modes
 

smite-ir-mutator/Cargo.toml

@@ -14,3 +14,5 @@ crate-type = ["cdylib"]
 smite-ir = { path = "../smite-ir" }
 rand.workspace = true
 postcard.workspace = true
+log.workspace = true
+simple_logger.workspace = true

smite-ir-mutator/src/lib.rs

@@ -16,8 +16,15 @@
 //! - `AFL_FRAMESHIFT_DISABLE=1` -- disable AFL++'s `FrameShift` analysis that
 //!   bypasses our custom mutators. This was an AFL++ bug fixed upstream in
 //!   commit eddb2701b022351fb34b696ccf923bb856e9d953.
-//! - `AFL_DISABLE_TRIM=1` -- this library does not implement custom trim and
-//!   AFL++'s default byte-level trim would corrupt our structured programs.
+//!
+//! # Logging
+//!
+//! Logging is opt-in: [`afl_custom_init`] installs a logger only when
+//! `RUST_LOG` is set. With `RUST_LOG` unset, every `log::*!` callsite is a
+//! no-op and AFL's stderr stays clean. Useful filters:
+//! - `RUST_LOG=smite_ir_mutator::trim=debug` -- print the decoded `Program`
+//!   before and after each successful trim.
+//! - `RUST_LOG=debug` -- everything this crate emits.
 //!
 //! # Buffer ownership
 //!
@@ -33,6 +40,7 @@ use rand::rngs::SmallRng;
 use rand::{RngExt, SeedableRng};
 
 use smite_ir::generators::{NodeAnnouncementGenerator, OpenChannelGenerator};
+use smite_ir::minimizers::{CommonSubexpressionEliminator, DeadCodeEliminator, Minimizer};
 use smite_ir::mutators::{InputSwapMutator, OperationParamMutator};
 use smite_ir::{Generator, Mutator, Program, ProgramBuilder};
 
@@ -152,14 +160,26 @@ fn warn_on_unset_afl_env() {
 /// Allocates a new [`MutatorState`] and returns an opaque pointer to it. AFL++
 /// passes this pointer back on every function call as the `data` argument.
 ///
+/// Also installs `simple_logger` if `RUST_LOG` is set, so `log::*!` callsites
+/// in this crate become live. Without `RUST_LOG` no logger is installed and
+/// every callsite is a no-op. Set
+/// `RUST_LOG=smite_ir_mutator::trim=debug` to see trim before/after dumps.
+///
 /// # Safety
 ///
 /// The returned pointer is heap-allocated via `Box::into_raw` and must be freed
 /// by a matching call to [`afl_custom_deinit`].
+///
+/// # Panics
+///
+/// Panics if `RUST_LOG` is set and `simple_logger` fails to initialize.
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn afl_custom_init(_afl: *const c_void, seed: c_uint) -> *mut c_void {
     #[cfg(not(test))]
     warn_on_unset_afl_env();
+    if std::env::var_os("RUST_LOG").is_some() {
+        simple_logger::init_with_env().expect("logger initializes");
+    }
     Box::into_raw(Box::new(MutatorState::new(seed))).cast::<c_void>()
 }
 
@@ -228,6 +248,115 @@ pub unsafe extern "C" fn afl_custom_fuzz(
     len
 }
 
+/// Runs the full minimizer pipeline (`DeadCodeEliminator` then
+/// `CommonSubexpressionEliminator`) on the corpus entry and stages the
+/// resulting candidate for [`afl_custom_trim`] to hand back.
+///
+/// Both minimizers are deterministic in-process transforms safe in IR
+/// semantics, so we don't need iterative AFL feedback. We compose them
+/// once and offer a single candidate. AFL still gets to verify it (its
+/// coverage cksum is the source of truth); on rejection AFL silently
+/// discards the candidate and keeps the original corpus entry.
+///
+/// AFL drives the trim loop with `while (stage_cur < stage_max)`, where
+/// `stage_max` is this function's return value and `stage_cur` is updated
+/// from [`afl_custom_post_trim`]'s return.
+///
+/// # Returns
+///
+/// - `1` if there's a candidate to offer (decode succeeded, validate
+///   passed, and the trim actually shrank the program). AFL enters the
+///   trim loop for one iteration.
+/// - `0` if there's nothing to do (decode/validate failed, or the trim
+///   was a no-op). AFL skips trim entirely.
+/// - Negative would signal a fatal error to AFL; we never produce one.
+///
+/// # Safety
+///
+/// - `data` must be a pointer returned by [`afl_custom_init`].
+/// - `buf` must point to `buf_size` readable bytes.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn afl_custom_init_trim(
+    data: *mut c_void,
+    buf: *mut u8,
+    buf_size: usize,
+) -> i32 {
+    let state = unsafe { &mut *data.cast::<MutatorState>() };
+
+    let input = unsafe { slice::from_raw_parts(buf, buf_size) };
+    let Some(program) = decode_and_validate(input) else {
+        return 0;
+    };
+
+    let before = log::log_enabled!(target: "smite_ir_mutator::trim", log::Level::Debug)
+        .then(|| program.clone());
+
+    let mut trimmed = program;
+    let dce_changed = DeadCodeEliminator.minimize(&mut trimmed);
+    let cse_changed = CommonSubexpressionEliminator.minimize(&mut trimmed);
+    if (!dce_changed && !cse_changed) || !state.serialize(&trimmed, buf_size) {
+        return 0;
+    }
+
+    if let Some(before) = before {
+        log::debug!(
+            target: "smite_ir_mutator::trim",
+            "dce={dce_changed} cse={cse_changed}\n--- before ---\n{before}--- after ---\n{trimmed}---"
+        );
+    }
+
+    1
+}
+
+/// Hands the pre-serialized trimmed candidate back to AFL.
+///
+/// The pointer written into `*out_buf` borrows from `MutatorState::out_buf`
+/// and is valid until the next call into this library; AFL copies the
+/// bytes before re-entering us. We always write a non-null pointer (even
+/// on the zero-length path) to satisfy AFL's `if (unlikely(!retbuf))
+/// FATAL(...)` check.
+///
+/// # Returns
+///
+/// - `> 0` on the first call after [`afl_custom_init_trim`]: the byte
+///   length of the candidate at `*out_buf`.
+/// - `0` afterwards. AFL treats this as "skip this iteration" rather than
+///   a stop signal; the loop terminates via [`afl_custom_post_trim`]'s
+///   return.
+///
+/// # Safety
+///
+/// - `data` must be a pointer returned by [`afl_custom_init`].
+/// - `out_buf` must be a valid, writable pointer to a `*const u8` slot.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn afl_custom_trim(data: *mut c_void, out_buf: *mut *const u8) -> usize {
+    let state = unsafe { &mut *data.cast::<MutatorState>() };
+    unsafe { *out_buf = state.out_buf.as_ptr() };
+    state.out_buf.len()
+}
+
+/// Always returns `1` to terminate AFL's trim loop after a single
+/// iteration.
+///
+/// AFL drives trim with `while (stage_cur < stage_max)` and assigns
+/// `stage_cur` from this function's return value. With `stage_max = 1`
+/// (set by [`afl_custom_init_trim`]), returning `1` makes the condition
+/// `1 < 1` false and breaks the loop.
+///
+/// `success` indicates whether the candidate's coverage cksum matched the
+/// original. We don't need to act on it: AFL itself either persists the
+/// trimmed buffer (on success) or keeps the original corpus entry (on
+/// failure), and we don't track partial state across iterations because
+/// there's only one.
+///
+/// # Safety
+///
+/// - `data` must be a pointer returned by [`afl_custom_init`].
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn afl_custom_post_trim(_data: *mut c_void, _success: u8) -> i32 {
+    1
+}
+
 /// Marker symbol that tells AFL++ not to populate `add_buf` for
 /// [`afl_custom_fuzz`]. AFL++ never actually calls this function -- it only
 /// checks for the symbol's presence via `dlsym` and, if found, skips picking a
@@ -338,6 +467,19 @@ mod tests {
         postcard::to_allocvec(&builder.build()).expect("postcard serialization")
     }
 
+    /// `seed_program_bytes()` plus an unreferenced `LoadAmount`, so the
+    /// pipeline has something for DCE to drop (and thus `init_trim`
+    /// returns `1`).
+    fn reducible_seed_bytes() -> Vec<u8> {
+        let bytes = seed_program_bytes();
+        let mut program: Program = postcard::from_bytes(&bytes).expect("decode");
+        program.instructions.push(smite_ir::Instruction {
+            operation: smite_ir::Operation::LoadAmount(0xdead_beef),
+            inputs: vec![],
+        });
+        postcard::to_allocvec(&program).expect("encode")
+    }
+
     #[test]
     fn init_returns_nonnull() {
         let state = State::new(0);
@@ -435,4 +577,106 @@ mod tests {
         // crash either.
         unsafe { afl_custom_splice_optout(ptr::null_mut()) };
     }
+
+    // -- Trim tests --
+
+    fn init_trim_via_ffi(state: &State, mut input: Vec<u8>) -> i32 {
+        unsafe { afl_custom_init_trim(state.0, input.as_mut_ptr(), input.len()) }
+    }
+
+    fn trim_via_ffi(state: &State) -> (*const u8, usize) {
+        let mut out: *const u8 = ptr::null();
+        let len = unsafe { afl_custom_trim(state.0, &raw mut out) };
+        (out, len)
+    }
+
+    fn post_trim_via_ffi(state: &State, success: bool) -> i32 {
+        unsafe { afl_custom_post_trim(state.0, u8::from(success)) }
+    }
+
+    #[test]
+    fn trim_init_returns_1_when_reduction_possible() {
+        let state = State::new(0);
+        let rv = init_trim_via_ffi(&state, reducible_seed_bytes());
+        assert_eq!(rv, 1);
+    }
+
+    #[test]
+    fn trim_init_returns_0_when_no_reduction_possible() {
+        // Generator output has no dead code or duplicate loads; the
+        // pipeline is a no-op, so we tell AFL to skip trim entirely.
+        let state = State::new(0);
+        let rv = init_trim_via_ffi(&state, seed_program_bytes());
+        assert_eq!(rv, 0);
+    }
+
+    #[test]
+    fn trim_init_returns_0_for_garbage() {
+        let state = State::new(0);
+        let rv = init_trim_via_ffi(&state, vec![0xFF; 16]);
+        assert_eq!(rv, 0);
+    }
+
+    #[test]
+    fn trim_yields_candidate_after_init() {
+        let state = State::new(0);
+        init_trim_via_ffi(&state, reducible_seed_bytes());
+        let (out, len) = trim_via_ffi(&state);
+        assert!(len > 0);
+        decode_and_validate(out, len);
+    }
+
+    #[test]
+    fn trim_post_trim_returns_1_to_terminate_loop() {
+        let state = State::new(0);
+        init_trim_via_ffi(&state, reducible_seed_bytes());
+        let _ = trim_via_ffi(&state);
+        // post_trim returns 1 unconditionally — it's the load-bearing
+        // termination signal that pushes AFL's `stage_cur` to `stage_max`.
+        assert_eq!(post_trim_via_ffi(&state, true), 1);
+        assert_eq!(post_trim_via_ffi(&state, false), 1);
+    }
+
+    #[test]
+    fn trim_init_does_not_overwrite_sequence() {
+        // Trim is not a mutation; `last_sequence` (used by `describe` to
+        // name queue entries from fuzz) must survive both the no-op and
+        // successful trim paths.
+        for (label, input, expected_rv) in [
+            ("no-op", seed_program_bytes(), 0),
+            ("success", reducible_seed_bytes(), 1),
+        ] {
+            let state = State::new(0);
+            // Run a fuzz call so last_sequence has known contents.
+            let _ = fuzz_via_ffi(&state, Vec::new(), 1 << 16);
+            let before = unsafe { CStr::from_ptr(afl_custom_describe(state.0, 256)) }
+                .to_str()
+                .expect("valid utf-8")
+                .to_string();
+            let rv = init_trim_via_ffi(&state, input);
+            assert_eq!(rv, expected_rv, "{label}");
+            let after = unsafe { CStr::from_ptr(afl_custom_describe(state.0, 256)) }
+                .to_str()
+                .expect("valid utf-8")
+                .to_string();
+            assert_eq!(before, after, "{label}");
+        }
+    }
+
+    #[test]
+    fn trim_candidate_is_smaller_than_input() {
+        let original_bytes = reducible_seed_bytes();
+        let original_program: Program = postcard::from_bytes(&original_bytes).expect("decode");
+
+        let state = State::new(0);
+        init_trim_via_ffi(&state, original_bytes);
+
+        let (out, len) = trim_via_ffi(&state);
+        assert!(len > 0, "trim should yield a candidate");
+        let trimmed = decode_and_validate(out, len);
+        assert!(
+            trimmed.instructions.len() < original_program.instructions.len(),
+            "trim should shrink instruction count"
+        );
+    }
 }

smite-ir/src/instruction.rs

@@ -12,7 +12,7 @@ use super::Operation;
 ///
 /// In SSA form, each instruction produces at most one variable (at the index
 /// equal to the instruction's position in the program).
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct Instruction {
     /// The operation to perform.
     pub operation: Operation,

smite-ir/src/lib.rs

@@ -7,13 +7,15 @@
 //!
 //! # Modules
 //! - [`instruction`] - Single IR instruction (operation + input references).
+//! - [`minimizers`] - Shrink a program while preserving interesting behaviour.
 //! - [`operation`] - Operations that load, compute, build or act.
 //! - [`program`] - Ordered list of instructions.
 //! - [`variable`] - Typed runtime values and lightweight type tags.
 
 pub mod builder;
 pub mod generators;
 pub mod instruction;
+pub mod minimizers;
 pub mod mutators;
 pub mod operation;
 pub mod program;
@@ -22,6 +24,7 @@ pub mod variable;
 pub use builder::ProgramBuilder;
 pub use generators::Generator;
 pub use instruction::Instruction;
+pub use minimizers::Minimizer;
 pub use mutators::Mutator;
 pub use operation::Operation;
 pub use program::Program;

smite-ir/src/minimizers.rs

@@ -0,0 +1,24 @@
+//! IR program minimizers.
+//!
+//! A [`Minimizer`] reduces a [`Program`] to a smaller, behaviourally
+//! equivalent version in a single pass. Both transforms are safe in IR
+//! semantics, so they don't need an oracle to drive the search.
+//!
+//! Run them in pipeline order for best results:
+//! 1. [`DeadCodeEliminator`] — drop dead instructions and reindex
+//! 2. [`CommonSubexpressionEliminator`] — merge equivalent pure expressions
+
+mod cse;
+mod dead_code;
+
+pub use cse::CommonSubexpressionEliminator;
+pub use dead_code::DeadCodeEliminator;
+
+use super::Program;
+
+/// A minimizer that reduces an IR program in one call.
+pub trait Minimizer {
+    /// Reduces `program` in place to a smaller, behaviourally equivalent
+    /// version. Returns `true` if the program was modified.
+    fn minimize(&self, program: &mut Program) -> bool;
+}