Allow to specify options for certain instances only

.gitignore

@@ -8,3 +8,4 @@ examples/xml*
 examples/seed*
 examples/target/
 
+config.toml

.vscode/launch.json

@@ -0,0 +1,145 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Custom Mutator",
+            "cargo": {
+                "args": [
+                    "build",
+                ],
+                "filter": {
+                    "name": "aflr",
+                    "kind": "bin"
+                }
+            },
+            "args": ["gen", "--config", "/home/user/AFL_Runner/config.toml"],
+            "cwd": "${workspaceFolder}",
+            "env": {"AFL_CUSTOM_MUTATOR_LIBRARY": "/tmp/test"}
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Normal",
+            "cargo": {
+                "args": [
+                    "build",
+                ],
+                "filter": {
+                    "name": "aflr",
+                    "kind": "bin"
+                }
+            },
+            "args": ["gen", "--config", "/home/user/AFL_Runner/config.toml"],
+            "cwd": "${workspaceFolder}",
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Use defaults",
+            "cargo": {
+                "args": [
+                    "build",
+                ],
+                "filter": {
+                    "name": "aflr",
+                    "kind": "bin"
+                }
+            },
+            "args": ["gen", "--use-afl-defaults", "--config", "/home/user/AFL_Runner/config.toml"],
+            "cwd": "${workspaceFolder}",
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Run",
+            "cargo": {
+                "args": [
+                    "build",
+                ],
+                "filter": {
+                    "name": "aflr",
+                    "kind": "bin"
+                }
+            },
+            "args": ["run", "--config", "/home/user/AFL_Runner/config.toml"],
+            "cwd": "${workspaceFolder}",
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Help",
+            "cargo": {
+                "args": [
+                    "build",
+                ],
+                "filter": {
+                    "name": "aflr",
+                    "kind": "bin"
+                }
+            },
+            "args": ["gen", "--help"],
+            "cwd": "${workspaceFolder}",
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Debug unit tests in executable 'dummy'",
+            "cargo": {
+                "args": [
+                    "test",
+                    "--no-run",
+                    "--bin=dummy",
+                    "--package=aflr_demo"
+                ],
+                "filter": {
+                    "name": "dummy",
+                    "kind": "bin"
+                }
+            },
+            "args": [],
+            "cwd": "${workspaceFolder}"
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Debug executable 'aflr'",
+            "cargo": {
+                "args": [
+                    "build",
+                    "--bin=aflr",
+                    "--package=afl_runner"
+                ],
+                "filter": {
+                    "name": "aflr",
+                    "kind": "bin"
+                }
+            },
+            "args": [],
+            "cwd": "${workspaceFolder}"
+        },
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Debug unit tests in executable 'aflr'",
+            "cargo": {
+                "args": [
+                    "test",
+                    "--no-run",
+                    "--bin=aflr",
+                    "--package=afl_runner"
+                ],
+                "filter": {
+                    "name": "aflr",
+                    "kind": "bin"
+                }
+            },
+            "args": [],
+            "cwd": "${workspaceFolder}"
+        }
+    ]
+}

src/afl_cmd_gen.rs

@@ -4,6 +4,7 @@ use std::{fs, process::Command};
 use uuid::Uuid;
 
 use crate::afl_env::{AFLEnv, AFLFlag};
+use crate::cli::FlagGroup;
 use crate::harness::Harness;
 use anyhow::{Context, Result};
 use rand::seq::SliceRandom;
@@ -151,6 +152,44 @@ fn apply_args(cmds: &mut [AflCmd], arg: &str, percentage: f64, rng: &mut impl Rn
     }
 }
 
+/// Applies a list of arguments to a percentage or specific count of AFL commands
+fn apply_args_list(cmds: &mut [AflCmd], args: &[&str], percentage: Option<f64>, count: Option<usize>, rng: &mut impl Rng) {
+    #[allow(clippy::cast_possible_truncation)]
+    #[allow(clippy::cast_sign_loss)]
+    #[allow(clippy::cast_precision_loss)]
+    let count = percentage.map_or_else(|| count.unwrap_or(0), |percentage| (cmds.len() as f64 * percentage) as usize);
+
+    let mut indices = HashSet::new();
+    while indices.len() < count {
+        indices.insert(rng.gen_range(0..cmds.len()));
+    }
+
+    for index in indices {
+        for &arg in args {
+            cmds[index].misc_afl_flags.push(arg.to_string());
+        }
+    }
+}
+
+/// Applies a list of env variables to a percentage or specific count of AFL command environments
+fn apply_env_vars(cmds: &mut [AflCmd], envs: &[&str], percentage: Option<f64>, count: Option<usize>, rng: &mut impl Rng) {
+    #[allow(clippy::cast_possible_truncation)]
+    #[allow(clippy::cast_sign_loss)]
+    #[allow(clippy::cast_precision_loss)]
+    let count = percentage.map_or_else(|| count.unwrap_or(0), |percentage| (cmds.len() as f64 * percentage) as usize);
+
+    let mut indices = HashSet::new();
+    while indices.len() < count {
+        indices.insert(rng.gen_range(0..cmds.len()));
+    }
+
+    for index in indices {
+        for &env in envs {
+            cmds[index].env.push(env.to_string());
+        }
+    }
+}
+
 /// Generates AFL commands based on the provided configuration
 pub struct AFLCmdGenerator {
     /// The harness configuration
@@ -171,7 +210,10 @@ pub struct AFLCmdGenerator {
     /// use the CMPCOV binary
     cmpcov_idxs: Vec<usize>,
     pub ramdisk: Option<String>,
+    // Use AFL defaults (do not variate mutator, power scheduler, trimming, ...)
     pub use_afl_defaults: bool,
+    // AFL additional 
+    pub additional_flags: Option<Vec<FlagGroup>>,
 }
 
 impl AFLCmdGenerator {
@@ -186,6 +228,7 @@ impl AFLCmdGenerator {
         afl_binary: Option<String>,
         is_ramdisk: bool,
         use_afl_defaults: bool,
+        additional_flags: Option<Vec<FlagGroup>>
     ) -> Self {
         let dict = dictionary.and_then(|d| {
             if d.exists() && d.is_file() {
@@ -198,7 +241,7 @@ impl AFLCmdGenerator {
             println!("[*] Attempting to create RAMDisk. Needing elevated privileges.");
             let r = Self::create_ramdisk();
             if let Ok(tmpfs) = r {
-                println!("[+] Using RAMDisk: {}", tmpfs);
+                println!("[+] Using RAMDisk: {tmpfs}");
                 Some(tmpfs)
             } else {
                 println!("[!] Failed to create RAMDisk: {}...", r.err().unwrap());
@@ -219,6 +262,7 @@ impl AFLCmdGenerator {
             cmpcov_idxs: Vec::new(),
             ramdisk: rdisk,
             use_afl_defaults,
+            additional_flags,
         }
     }
 
@@ -300,6 +344,33 @@ impl AFLCmdGenerator {
         self.apply_cmplog(&mut cmds, &mut rng);
         self.apply_target_args(&mut cmds);
         self.apply_cmpcov(&mut cmds, &mut rng);
+
+        // Apply additional flags
+        if let Some(additional_flags) = &self.additional_flags {
+            for flag_group in additional_flags {
+                let (dash_flags, env_flags): (Vec<_>, Vec<_>) = flag_group
+                    .flags
+                    .iter()
+                    .partition(|flag| flag.0.starts_with('-'));
+
+                let dash_flags: Vec<String> = dash_flags
+                    .iter()
+                    .map(|flag| format!("{} {}", flag.0, flag.1))
+                    .collect();
+
+                let env_flags: Vec<String> = env_flags
+                    .iter()
+                    .map(|flag| format!("{}={}", flag.0, flag.1))
+                    .collect();
+
+                apply_args_list(&mut cmds, &dash_flags.iter().map(String::as_str).collect::<Vec<&str>>(), flag_group.probability.map(f64::from), flag_group.count.map(|c| c as usize), &mut rng);
+                apply_env_vars(&mut cmds, &env_flags.iter().map(String::as_str).collect::<Vec<&str>>(), flag_group.probability.map(f64::from), flag_group.count.map(|c| c as usize), &mut rng);
+            }
+        }
+
+        // Verify and merge AFL environment variables
+        self.verify_and_merge_env_vars(&mut cmds);
+
         // NOTE: Needs to called last as it relies on cmpcov/cmplog being already set
         self.apply_fuzzer_roles(&mut cmds);
 
@@ -321,6 +392,33 @@ impl AFLCmdGenerator {
         Ok(cmd_strings)
     }
 
+    /// Verifies and merges AFL environment variables
+    fn verify_and_merge_env_vars(&self, cmds: &mut [AflCmd]) {
+        let allowed_merge_keys = ["AFL_PRELOAD", "LD_PRELOAD"];
+        let mut env_map = std::collections::HashMap::new();
+
+        for cmd in cmds {
+            for env in &cmd.env {
+                let mut split = env.splitn(2, '=');
+                if let (Some(key), Some(value)) = (split.next(), split.next()) {
+                    if allowed_merge_keys.contains(&key) {
+                        env_map.entry(key).or_insert_with(Vec::new).push(value.to_string());
+                    }
+                }
+            }
+        }
+
+        for cmd in &mut *cmds {
+            for key in &allowed_merge_keys {
+                if let Some(values) = env_map.get(key) {
+                    let merged_value = values.join(":");
+                    cmd.env.retain(|env| !env.starts_with(key));
+                    cmd.env.push(format!("{}={}", key, merged_value));
+                }
+            }
+        }
+    }
+
     /// Initializes AFL configurations
     fn initialize_configs(&self, rng: &mut impl Rng) -> Vec<AFLEnv> {
         let mut configs = vec![AFLEnv::new(); self.runners as usize];

src/afl_env.rs

@@ -38,7 +38,7 @@ impl std::fmt::Display for AFLFlag {
             Self::IgnoreSeedProblems => "AFL_IGNORE_SEED_PROBLEMS",
             Self::ImportFirst => "AFL_IMPORT_FIRST",
         };
-        write!(f, "{}", x)
+        write!(f, "{x}")
     }
 }
 
@@ -92,7 +92,7 @@ impl AFLEnv {
         }
 
         for flag in &self.flags {
-            command.push(format!("{}=1", flag.to_string()));
+            command.push(format!("{flag}=1"));
         }
 
         command.push(format!("AFL_TESTCACHE_SIZE={} ", self.testcache_size));

src/cli.rs

@@ -1,7 +1,7 @@
 use serde::Deserialize;
 
 use clap::{ArgAction, Args, Parser, Subcommand, ValueEnum};
-use std::path::PathBuf;
+use std::{collections::HashMap, path::PathBuf};
 
 /// Default corpus directory
 pub const AFL_CORPUS: &str = "/tmp/afl_input";
@@ -367,6 +367,22 @@ pub struct AflConfig {
     pub afl_flags: Option<String>,
     /// Use AFL defaults
     pub use_afl_defaults: Option<bool>,
+    /// Partial AFL flags
+    pub flags_partial: Vec<FlagGroup>,
+}
+
+#[derive(Debug, Deserialize, Clone, Default)]
+pub struct FlagGroup {
+    pub probability: Option<f32>,
+    pub count: Option<u32>,
+    #[serde(flatten)]
+    pub flags: HashMap<String, toml::Value>,
+}
+
+impl FlagGroup {
+    pub const fn is_valid(&self) -> bool {
+        !(self.probability.is_some() && self.count.is_some())
+    }
 }
 
 /// Configuration for tmux