420 lines
16 KiB
Rust
420 lines
16 KiB
Rust
#![feature(int_log)]
|
|
|
|
mod apu;
|
|
mod audio;
|
|
mod cartridge;
|
|
mod cpu;
|
|
mod fuzzing_state;
|
|
mod input;
|
|
mod ppu;
|
|
mod screen;
|
|
mod state;
|
|
|
|
use apu::Apu;
|
|
use ppu::Ppu;
|
|
|
|
use crate::cpu::Cpu;
|
|
use crate::fuzzing_state::{FuzzingInputState, FuzzingState};
|
|
use crate::input::ConsoleAction::SoftReset;
|
|
use crate::input::{ConsoleAction, FuzzingInput};
|
|
use crate::screen::Screen;
|
|
use minifb::{ScaleMode, Window, WindowOptions};
|
|
use priority_queue::double_priority_queue::DoublePriorityQueue;
|
|
use priority_queue::priority_queue::PriorityQueue;
|
|
use std::collections::HashSet;
|
|
use std::sync::mpsc::{Receiver, Sender};
|
|
use std::sync::{mpsc, Arc, Mutex};
|
|
use std::{fs, thread};
|
|
|
|
// The number of cpu instances to spawn..
|
|
const NUM_THREADS: usize = 28;
|
|
|
|
// The number of frames to fuzz and process
|
|
// A small number exploits the current point more at the expense of
|
|
// large exploration - and vice versa.
|
|
const FRAMES_TO_CONSIDER: usize = 500;
|
|
|
|
// Same input should generate the same output...
|
|
// (I make no guarantee of that at the moment)
|
|
const RNG_SEED: u32 = 0x1334357;
|
|
|
|
// If set to a low number, this disables start presses after the given frame
|
|
// Useful for some games where pausing does nothing to advance the game...
|
|
const DISABLE_START_PRESSES_AFTER: usize = 100;
|
|
|
|
// The rate at which seed inputs become corrupted..
|
|
const MUTATION_RATE: f64 = 0.10;
|
|
|
|
// The rate at which seed inputs may become soft resets..
|
|
const MUTATION_RATE_SOFT_RESET: f64 = 0.00;
|
|
|
|
// The number of cases to fuzz from a given seed input at each consideration point..
|
|
const SEED_CASES: usize = 10;
|
|
|
|
// Only add the original seed case top the queue (hammer a single state)
|
|
const SINGLE_SHOT: bool = true;
|
|
|
|
// Fast forward the seed state to a given frame
|
|
const PLAY_FROM: usize = 1900;
|
|
|
|
fn main() -> Result<(), String> {
|
|
let argv = std::env::args().collect::<Vec<String>>();
|
|
|
|
if argv.len() < 3 {
|
|
println!("usage: fuzz <rom> <fm2 seed input>")
|
|
}
|
|
|
|
// The rom file...mostly tested with Super Mario Bros
|
|
let rom_filename = argv[1].to_string();
|
|
// The tas file to use as seed input e.g. happylee-supermariobros,warped.fm2
|
|
let fm2_file = argv[2].to_string();
|
|
|
|
// Seed for our Random Number Generator...
|
|
let mut rng = Rng::init(RNG_SEED);
|
|
|
|
// Mutex for launching windows others we get double frees in the
|
|
// windowing library...(yup)
|
|
let mutex = Arc::new(Mutex::new(0));
|
|
|
|
let mut seed_input = FuzzingInput::load(fm2_file.as_str());
|
|
seed_input.disable_start_after(DISABLE_START_PRESSES_AFTER);
|
|
|
|
let mut starting_state = FuzzingState::default(rom_filename.clone());
|
|
|
|
let mut previous_states: HashSet<usize> = HashSet::new();
|
|
|
|
if PLAY_FROM > 0 {
|
|
println!("Playing {} frames prelude...", PLAY_FROM);
|
|
let (mut cpu, _) = starting_state.load_state(rom_filename.clone());
|
|
play_frames(&mut cpu, PLAY_FROM, &seed_input);
|
|
starting_state = FuzzingState::save_state(cpu, PLAY_FROM);
|
|
for addr in starting_state.cpu_state.addresses_fetched.iter() {
|
|
if previous_states.insert(*addr) {}
|
|
}
|
|
}
|
|
|
|
let mut fuzzing_queue: PriorityQueue<FuzzingInputState, u64> = PriorityQueue::new();
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(seed_input.clone(), starting_state.clone()),
|
|
10000,
|
|
);
|
|
for _i in 1..NUM_THREADS {
|
|
let mut mutated_input = seed_input.clone();
|
|
mutated_input.mutate_from(&mut rng, PLAY_FROM, FRAMES_TO_CONSIDER);
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(mutated_input.clone(), starting_state.clone()),
|
|
rng.next() as u64,
|
|
);
|
|
}
|
|
|
|
let mut result_channels = vec![];
|
|
for i in 0..NUM_THREADS {
|
|
let i = i.clone();
|
|
let mutex = mutex.clone();
|
|
let rom_filename = rom_filename.clone();
|
|
let (tx_inputs, rx_input): (
|
|
Sender<(FuzzingInput, FuzzingState)>,
|
|
Receiver<(FuzzingInput, FuzzingState)>,
|
|
) = mpsc::channel();
|
|
let (tx_results, rx_results): (
|
|
Sender<(FuzzingInput, FuzzingState)>,
|
|
Receiver<(FuzzingInput, FuzzingState)>,
|
|
) = mpsc::channel();
|
|
result_channels.push((tx_inputs, rx_results));
|
|
thread::spawn(move || {
|
|
run_game(i as isize, mutex, rom_filename, rx_input, tx_results);
|
|
});
|
|
}
|
|
|
|
let mut loop_count = 0;
|
|
|
|
loop {
|
|
println!("Prospective Cases: {}", fuzzing_queue.len());
|
|
|
|
for i in 0..NUM_THREADS {
|
|
if fuzzing_queue.is_empty() == false {
|
|
let (state, score) = fuzzing_queue.pop().unwrap();
|
|
// Send to thread for process...
|
|
result_channels[i]
|
|
.0
|
|
.send((state.0.clone(), state.1.clone()))
|
|
.expect("error sending new fuzzing input to thread...");
|
|
} else {
|
|
result_channels[i]
|
|
.0
|
|
.send((seed_input.clone(), starting_state.clone()))
|
|
.expect("error sending new fuzzing input to thread");
|
|
}
|
|
}
|
|
|
|
// Gather results from each thread
|
|
for i in 0..NUM_THREADS {
|
|
match result_channels[i].1.recv() {
|
|
Ok((fuzzing_input, fuzzing_state)) => {
|
|
/* let mut lowest_similarity = u64::MAX;
|
|
for (_j, state) in previous_states.iter().enumerate() {
|
|
let similiarty =
|
|
hamming::distance(&state, &fuzzing_state.cpu_state.mem);
|
|
|
|
if similiarty < lowest_similarity {
|
|
lowest_similarity = similiarty
|
|
}
|
|
}
|
|
|
|
previous_states.insert(fuzzing_state.cpu_state.mem.clone());*/
|
|
let mut lowest_similarity = 0u64;
|
|
for addr in fuzzing_state.cpu_state.addresses_fetched.iter() {
|
|
if previous_states.insert(*addr) {
|
|
lowest_similarity += 1;
|
|
}
|
|
}
|
|
|
|
if i == 3 && loop_count == 3 {
|
|
fuzzing_input.export(
|
|
"megaman_interestng.fm2".parse().unwrap(),
|
|
fuzzing_state.frames,
|
|
);
|
|
}
|
|
|
|
if SINGLE_SHOT {
|
|
if lowest_similarity > 0 {
|
|
println!("{} ", lowest_similarity);
|
|
for i in 0..lowest_similarity {
|
|
let mut mutated_input: FuzzingInput = fuzzing_input.clone();
|
|
mutated_input.mutate_from(
|
|
&mut rng,
|
|
starting_state.frames,
|
|
FRAMES_TO_CONSIDER,
|
|
);
|
|
// Add the mutated input and the regular input to the queue...
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(
|
|
mutated_input.clone(),
|
|
starting_state.clone(),
|
|
),
|
|
lowest_similarity,
|
|
);
|
|
}
|
|
}
|
|
continue;
|
|
} else {
|
|
// Seed Thread
|
|
if fuzzing_input.mutated == false {
|
|
// Seed input handling...
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(fuzzing_input.clone(), fuzzing_state.clone()),
|
|
10000,
|
|
);
|
|
|
|
// Add seed cases..
|
|
for _i in 0..SEED_CASES {
|
|
let mut mutated_input: FuzzingInput = fuzzing_input.clone();
|
|
mutated_input.mutate_from(
|
|
&mut rng,
|
|
fuzzing_state.frames,
|
|
FRAMES_TO_CONSIDER,
|
|
);
|
|
// Add the mutated input and the regular input to the queue...
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(mutated_input.clone(), fuzzing_state.clone()),
|
|
lowest_similarity,
|
|
);
|
|
}
|
|
} else if lowest_similarity > 0 {
|
|
// Only add to the queue if strictly better...
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(fuzzing_input.clone(), fuzzing_state.clone()),
|
|
lowest_similarity,
|
|
);
|
|
|
|
let mut mutated_input: FuzzingInput = fuzzing_input.clone();
|
|
mutated_input.mutate_from(
|
|
&mut rng,
|
|
fuzzing_state.frames,
|
|
FRAMES_TO_CONSIDER,
|
|
);
|
|
// Add the mutated input and the regular input to the queue...
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(mutated_input.clone(), fuzzing_state.clone()),
|
|
lowest_similarity,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
if SINGLE_SHOT {
|
|
// Add seed cases..
|
|
for _i in 0..SEED_CASES {
|
|
let mut mutated_input: FuzzingInput = seed_input.clone();
|
|
mutated_input.mutate_from(&mut rng, starting_state.frames, FRAMES_TO_CONSIDER);
|
|
// Add the mutated input and the regular input to the queue...
|
|
fuzzing_queue.push(
|
|
FuzzingInputState(mutated_input.clone(), starting_state.clone()),
|
|
rng.next() as u64,
|
|
);
|
|
}
|
|
}
|
|
loop_count += 1;
|
|
}
|
|
}
|
|
|
|
pub struct Rng {
|
|
state: u32,
|
|
}
|
|
|
|
impl Rng {
|
|
pub fn init(state: u32) -> Rng {
|
|
Rng { state }
|
|
}
|
|
|
|
pub fn next(&mut self) -> u32 {
|
|
self.state ^= self.state << 13;
|
|
self.state ^= self.state >> 17;
|
|
self.state ^= self.state << 5;
|
|
return self.state;
|
|
}
|
|
}
|
|
|
|
fn run_game(
|
|
i: isize,
|
|
mutex: Arc<Mutex<u8>>,
|
|
filename: String,
|
|
refresh_inputs: Receiver<(FuzzingInput, FuzzingState)>,
|
|
send_results: Sender<(FuzzingInput, FuzzingState)>,
|
|
) -> FuzzingState {
|
|
let mut window = match mutex.lock() {
|
|
Ok(_x) => {
|
|
let mut window = Window::new(
|
|
"NesFuzz",
|
|
256,
|
|
240,
|
|
WindowOptions {
|
|
resize: true,
|
|
borderless: true,
|
|
scale_mode: ScaleMode::UpperLeft,
|
|
..WindowOptions::default()
|
|
},
|
|
)
|
|
.unwrap_or_else(|x| panic!("Unable to create window {}", x));
|
|
|
|
window.set_position(10 + ((i % 7) * 260), (240 * (i / 7)) + 100);
|
|
window
|
|
}
|
|
_ => {
|
|
panic!("unable to acquire mutex")
|
|
}
|
|
};
|
|
|
|
let mut screen = Screen::new();
|
|
|
|
loop {
|
|
match refresh_inputs.recv() {
|
|
Err(_x) => {}
|
|
Ok((fuzzing_input, fuzzing_state)) => {
|
|
let mut new_frames = 0;
|
|
|
|
// Initialize hardware components
|
|
let (mut cpu, mut frames) = fuzzing_state.load_state(filename.clone());
|
|
let mut prev_frame = frames;
|
|
while new_frames < FRAMES_TO_CONSIDER {
|
|
// step CPU: perform 1 cpu instruction, getting back number of clock cycles it took
|
|
|
|
let cpu_cycles = cpu.step();
|
|
// clock PPU three times for every CPU cycle
|
|
for _ in 0..cpu_cycles * 3 {
|
|
let (pixel, end_of_frame) = cpu.ppu.clock();
|
|
match pixel {
|
|
Some((x, y, color)) => screen.plot_pixel(
|
|
x,
|
|
y,
|
|
color.0 as u32,
|
|
color.1 as u32,
|
|
color.2 as u32,
|
|
),
|
|
None => (),
|
|
};
|
|
if end_of_frame {
|
|
screen.draw_string(10, 210, &*format!("Frames: {}", new_frames));
|
|
window
|
|
.update_with_buffer(&screen.display, 256, 240)
|
|
.unwrap();
|
|
frames += 1;
|
|
new_frames += 1;
|
|
}
|
|
}
|
|
// Checking for Inputs...
|
|
match fuzzing_input.get_frame_input(frames) {
|
|
Some(input) => {
|
|
match input.console_action {
|
|
ConsoleAction::None => {}
|
|
ConsoleAction::SoftReset => {
|
|
println!("soft reset");
|
|
cpu.soft_reset();
|
|
frames += 1;
|
|
}
|
|
}
|
|
//if cpu.strobe {
|
|
cpu.button_states = input.player_1_input;
|
|
// FIXME PLayer 2 doesn't play nicely with some games (e.g. mario)
|
|
// So to enable player 2 controls you also have to uncomment the
|
|
// bus in cpu/mod.rs
|
|
cpu.button_states2 = input.player_2_input;
|
|
//}
|
|
}
|
|
_ => {}
|
|
};
|
|
}
|
|
|
|
send_results
|
|
.send((fuzzing_input.clone(), FuzzingState::save_state(cpu, frames)))
|
|
.expect("error sending result to thread");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn play_frames(cpu: &mut Cpu, num_frames: usize, fuzzing_input: &FuzzingInput) {
|
|
let mut played_frames = 0;
|
|
let mut prev_frame = 0;
|
|
while played_frames < num_frames {
|
|
// step CPU: perform 1 cpu instruction, getting back number of clock cycles it took
|
|
|
|
// A the start of each new frame set our input for that frame...
|
|
if prev_frame == 0 || prev_frame != played_frames {
|
|
match fuzzing_input.get_frame_input(played_frames) {
|
|
Some(input) => {
|
|
match input.console_action {
|
|
ConsoleAction::None => {}
|
|
ConsoleAction::SoftReset => {
|
|
cpu.soft_reset();
|
|
played_frames += 1;
|
|
}
|
|
}
|
|
|
|
cpu.button_states = input.player_1_input;
|
|
cpu.button_states2 = input.player_2_input;
|
|
}
|
|
_ => {
|
|
cpu.button_states = 0;
|
|
cpu.button_states2 = 0;
|
|
}
|
|
};
|
|
prev_frame = played_frames;
|
|
}
|
|
|
|
let cpu_cycles = cpu.step();
|
|
|
|
// clock PPU three times for every CPU cycle
|
|
for _ in 0..cpu_cycles * 3 {
|
|
let (_, end_of_frame) = cpu.ppu.clock();
|
|
if end_of_frame {
|
|
prev_frame = played_frames;
|
|
played_frames += 1;
|
|
}
|
|
}
|
|
}
|
|
}
|