diff --git a/smite-ir-mutator/src/lib.rs b/smite-ir-mutator/src/lib.rs index 0903ee2..5bdc328 100644 --- a/smite-ir-mutator/src/lib.rs +++ b/smite-ir-mutator/src/lib.rs @@ -30,7 +30,7 @@ use rand::rngs::SmallRng; use rand::{RngExt, SeedableRng}; use smite_ir::generators::OpenChannelGenerator; -use smite_ir::mutators::{InputSwapMutator, OperationParamMutator}; +use smite_ir::mutators::{InputSwapMutator, InstructionDeleteMutator, OperationParamMutator}; use smite_ir::{Generator, Mutator, Program, ProgramBuilder}; /// Mutator state owned by AFL++ across calls. Allocated by [`afl_custom_init`], @@ -77,12 +77,20 @@ impl MutatorState { let stack = 1u32 << self.rng.random_range(0..=4); for _ in 0..stack { // Uniform pick between the available mutators. - let name = if self.rng.random() { - OperationParamMutator.mutate(program, &mut self.rng); - "op-param" - } else { - InputSwapMutator.mutate(program, &mut self.rng); - "input-swap" + let name = match self.rng.random_range(0..3) { + 0 => { + OperationParamMutator.mutate(program, &mut self.rng); + "op-param" + } + 1 => { + InputSwapMutator.mutate(program, &mut self.rng); + "input-swap" + } + 2 => { + InstructionDeleteMutator.mutate(program, &mut self.rng); + "instr-delete" + } + _ => unreachable!(), }; self.last_sequence.push(name); } @@ -369,7 +377,7 @@ mod tests { } for name in suffix.split(',') { assert!( - name == "op-param" || name == "input-swap", + name == "op-param" || name == "input-swap" || name == "instr-delete", "unexpected mutator name in description: {name:?} (full: {s:?})", ); } diff --git a/smite-ir/src/mutators.rs b/smite-ir/src/mutators.rs index 2c076f0..3d3d500 100644 --- a/smite-ir/src/mutators.rs +++ b/smite-ir/src/mutators.rs @@ -4,9 +4,11 @@ //! structural validity. Each mutator makes a small, targeted change. mod input_swap; +mod instruction_delete; mod operation_param; pub use input_swap::InputSwapMutator; +pub use instruction_delete::InstructionDeleteMutator; pub use operation_param::OperationParamMutator; use rand::Rng; diff --git a/smite-ir/src/mutators/instruction_delete.rs b/smite-ir/src/mutators/instruction_delete.rs new file mode 100644 index 0000000..5688246 --- /dev/null +++ b/smite-ir/src/mutators/instruction_delete.rs @@ -0,0 +1,59 @@ +//! Mutator that removes an instruction. + +use rand::{Rng, RngExt, seq::IteratorRandom}; + +use super::Mutator; +use crate::Program; + +/// Deletes a randomly selected instruction by removing it from +/// the instructions list and reindexing the subsequent instructions. +pub struct InstructionDeleteMutator; + +impl Mutator for InstructionDeleteMutator { + fn mutate(&self, program: &mut Program, rng: &mut impl Rng) -> bool { + if program.instructions.is_empty() { + return false; + } + // Pick a random instruction to delete. + let deleted_idx = rng.random_range(0..program.instructions.len()); + let deleted_type = program.instructions[deleted_idx].operation.output_type(); + + // If any instruction downstream depends on the deleted one, pick a prior + // type-matching variable to redirect those inputs to. + let replacement_idx = if program.instructions[(deleted_idx + 1)..] + .iter() + .any(|instr| instr.inputs.contains(&deleted_idx)) + { + match program.instructions[..deleted_idx] + .iter() + .enumerate() + .filter_map(|(i, instr)| { + (instr.operation.output_type() == deleted_type).then_some(i) + }) + .choose(rng) + { + Some(idx) => Some(idx), + // Abort if no valid replacement variable exists in the preceding scope. + None => return false, + } + } else { + None + }; + + // Delete from the program. + program.instructions.remove(deleted_idx); + + // Heal downstream inputs: redirect references to the deleted index, and + // decrement references past it. + for instr in &mut program.instructions[deleted_idx..] { + for input in &mut instr.inputs { + if *input == deleted_idx { + *input = replacement_idx.expect("dependent input implies replacement"); + } else if *input > deleted_idx { + *input -= 1; + } + } + } + true + } +} diff --git a/smite-ir/src/tests.rs b/smite-ir/src/tests.rs index b72bb82..f33c533 100644 --- a/smite-ir/src/tests.rs +++ b/smite-ir/src/tests.rs @@ -6,7 +6,7 @@ use smite::bolt::MAX_MESSAGE_SIZE; use super::*; use generators::OpenChannelGenerator; -use mutators::{InputSwapMutator, OperationParamMutator}; +use mutators::{InputSwapMutator, InstructionDeleteMutator, OperationParamMutator}; use operation::{AcceptChannelField, ChannelTypeVariant, ShutdownScriptVariant}; use program::ValidateError; @@ -1058,3 +1058,134 @@ fn input_swap_preserves_types() { } } } + +// -- InstructionDeleteMutator tests -- + +#[test] +fn instr_delete_changes_values() { + let original = generate_program(0); + let mut program = original.clone(); + let mutator = InstructionDeleteMutator; + let mut rng = SmallRng::seed_from_u64(0); + + for _ in 0..100 { + mutator.mutate(&mut program, &mut rng); + } + assert_ne!( + program, original, + "InstructionDeleteMutator never changed the program" + ); +} + +#[test] +fn instr_delete_false_is_noop() { + let original = generate_program(0); + let mutator = InstructionDeleteMutator; + let mut rng = SmallRng::seed_from_u64(0); + + for _ in 0..100 { + let mut program = original.clone(); + if !mutator.mutate(&mut program, &mut rng) { + assert_eq!(program, original, "program modified on false return"); + } + } +} + +#[test] +fn instr_delete_returns_false_on_empty_program() { + let mut program = Program { + instructions: vec![], + }; + let mutator = InstructionDeleteMutator; + let mut rng = SmallRng::seed_from_u64(0); + assert!(!mutator.mutate(&mut program, &mut rng)); +} + +#[test] +fn instr_delete_returns_false_if_unhealable() { + let original = Program { + instructions: vec![ + Instruction { + operation: Operation::LoadPrivateKey(key(1)), + inputs: vec![], + }, + Instruction { + operation: Operation::DerivePoint, + inputs: vec![0], + }, + ], + }; + let mutator = InstructionDeleteMutator; + let mut rng = SmallRng::seed_from_u64(0); + let mut found_unhealable = false; + + for _ in 0..100 { + let mut program = original.clone(); + if !mutator.mutate(&mut program, &mut rng) { + found_unhealable = true; + break; + } + } + assert!( + found_unhealable, + "mutator never returned false for unhealable instruction" + ); +} + +#[test] +fn instr_delete_shifts_indices_correctly() { + let original = Program { + instructions: vec![ + Instruction { + operation: Operation::LoadPrivateKey(key(1)), + inputs: vec![], + }, + Instruction { + operation: Operation::LoadPrivateKey(key(2)), + inputs: vec![], + }, + Instruction { + operation: Operation::DerivePoint, + inputs: vec![1, 1], + }, + ], + }; + + let mutator = InstructionDeleteMutator; + let mut rng = SmallRng::seed_from_u64(0); + let mut verified_shift = false; + + for _ in 0..100 { + let mut program = original.clone(); + if mutator.mutate(&mut program, &mut rng) && + program.instructions.len() == 2 && + // Check if it deleted a Load* operation, meaning DerivePoint is now at index 1 + program.instructions[1].operation == Operation::DerivePoint + { + // input references must have shifted from [1, 1] down to [0, 0] + assert_eq!(program.instructions[1].inputs, vec![0, 0]); + verified_shift = true; + break; + } + } + assert!( + verified_shift, + "mutator never took the expected target shift path" + ); +} + +#[test] +fn instr_delete_maintains_validity() { + let original = generate_program(0); + let mutator = InstructionDeleteMutator; + let mut rng = SmallRng::seed_from_u64(0); + + for _ in 0..100 { + let mut program = original.clone(); + if mutator.mutate(&mut program, &mut rng) { + program + .validate() + .expect("InstructionDeleteMutator produces valid programs"); + } + } +}