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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
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_config::afl::AflSeed;
use trident_config::TridentConfig;

use super::{Commander, Error};
use rand::RngCore;

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

        // build args without cargo target dir
        let build_args = config.get_afl_build_args();
        // fuzz args without afl workspace in and out
        let fuzz_args = config.get_afl_fuzz_args();

        // cargo target directory
        let cargo_target_dir = config.get_afl_target_dir();

        // afl workspace in and out
        let afl_workspace_in = config.get_afl_workspace_in();
        let afl_workspace_out = config.get_afl_workspace_out();

        let full_target_path = config.get_afl_target_path(&target);

        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(["--target-dir", &cargo_target_dir])
            .args(build_args)
            .args(["--bin", &target])
            .spawn()?;
        Self::handle_child(&mut child).await?;

        let mut child = Command::new("cargo")
            .arg("afl")
            .arg("fuzz")
            .args(["-i", &afl_workspace_in])
            .args(["-o", &afl_workspace_out])
            .args(fuzz_args)
            .arg(&full_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: String) {
        let config = TridentConfig::new();

        let crash_file_path = Path::new(&crash_file);

        let crash_file = if crash_file_path.is_absolute() {
            crash_file_path
        } else {
            let cwd = std::env::current_dir()?;

            &cwd.join(crash_file_path)
        };

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

        // cargo target directory
        let cargo_target_dir = config.get_afl_target_dir();

        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 (bytes, override_file) = obtain_seed(seed);

    let file_path = path.join(&seed.file_name);

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

    Ok(())
}

fn obtain_seed(value: &AflSeed) -> (Vec<u8>, bool) {
    match value.bytes_count {
        Some(number_of_random_bytes) => {
            if number_of_random_bytes > 0 {
                let mut rng = rand::rngs::OsRng;
                let mut seed = vec![0u8; number_of_random_bytes];
                rng.fill_bytes(&mut seed);
                (seed, value.override_file.unwrap_or_default())
            } else {
                let seed_as_bytes = value
                    .seed
                    .clone()
                    .unwrap_or_else(|| panic!("Seed value is empty for seed {}", value.file_name));

                (
                    seed_as_bytes.as_bytes().to_vec(),
                    value.override_file.unwrap_or_default(),
                )
            }
        }
        None => {
            let seed_as_bytes = value
                .seed
                .clone()
                .unwrap_or_else(|| panic!("Seed value is empty for seed {}", value.file_name));
            (
                seed_as_bytes.as_bytes().to_vec(),
                value.override_file.unwrap_or_default(),
            )
        }
    }
}