trident_client/commander/
afl.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
use crate::constants::*;
use fehler::{throw, throws};
use std::io::{Read, Write};
use std::process::Stdio;
use std::{fs::File, path::Path};
use tokio::{io::AsyncWriteExt, process::Command};
use trident_fuzz::config::afl::AflSeed;
use trident_fuzz::config::Config;

use super::{Commander, Error};

impl Commander {
    /// Runs fuzzer on the given target.
    #[throws]
    pub async fn run_afl(&self, target: String) {
        let config = Config::new();

        let build_args = config.get_afl_build_args();
        let fuzz_args = config.get_afl_fuzz_args();

        let mut target_path = config.get_afl_target_path();
        target_path.push_str(&target);

        let afl_workspace_in = config.get_afl_workspace_in();
        let afl_workspace_in_path = Path::new(&afl_workspace_in);
        let initial_seeds = config.get_initial_seed();

        if !afl_workspace_in_path.exists() {
            std::fs::create_dir_all(afl_workspace_in_path)?;

            for x in initial_seeds {
                create_seed_file(afl_workspace_in_path, x)?;
            }
        } else if afl_workspace_in_path.is_dir() {
            for x in initial_seeds {
                create_seed_file(afl_workspace_in_path, x)?;
            }
        } else {
            throw!(Error::BadAFLWorkspace)
        }

        let mut rustflags = std::env::var("RUSTFLAGS").unwrap_or_default();

        rustflags.push_str("--cfg afl");

        let mut child = Command::new("cargo")
            .env("RUSTFLAGS", rustflags)
            .arg("afl")
            .arg("build")
            .args(build_args)
            .args(["--bin", &target])
            .spawn()?;
        Self::handle_child(&mut child).await?;

        let mut child = Command::new("cargo")
            .arg("afl")
            .arg("fuzz")
            .args(fuzz_args)
            .arg(&target_path)
            .spawn()?;

        Self::handle_child(&mut child).await?;
    }

    /// Runs fuzzer on the given target.
    #[throws]
    pub async fn run_afl_debug(&self, target: String, crash_file_path: String) {
        let config = Config::new();

        let crash_file = self.root.join(crash_file_path);

        let cargo_target_dir = config.get_afl_cargo_build_dir();

        if !crash_file.try_exists()? {
            println!("{ERROR} The crash file [{:?}] not found", crash_file);
            throw!(Error::CrashFileNotFound);
        }

        let mut rustflags = std::env::var("RUSTFLAGS").unwrap_or_default();

        rustflags.push_str("--cfg afl");

        let mut file = File::open(crash_file)?;
        let mut file_contents = Vec::new();
        file.read_to_end(&mut file_contents)?;

        // using exec rather than spawn and replacing current process to avoid unflushed terminal output after ctrl+c signal
        let mut child = Command::new("cargo")
            .env("RUSTFLAGS", rustflags)
            .arg("afl")
            .arg("run")
            .args(["--target-dir", &cargo_target_dir])
            .args(["--bin", &target])
            .stdin(Stdio::piped())
            .spawn()?;

        if let Some(mut stdin) = child.stdin.take() {
            stdin.write_all(&file_contents).await?;
        }
        child.wait().await?;
    }
}

fn create_seed_file(path: &Path, seed: &AflSeed) -> std::io::Result<()> {
    let file_path = path.join(&seed.file_name);

    if file_path.exists() {
        if seed.override_file {
            let mut file = File::create(file_path)?;
            file.write_all(&seed.seed)?;
        }
    } else {
        let mut file = File::create(file_path)?;
        file.write_all(&seed.seed)?;
    }

    Ok(())
}