You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
282 lines
14 KiB
282 lines
14 KiB
#!/usr/bin/env python3
|
|
# Part of the aflplusplus project, requires Python 3.8+.
|
|
# Author: Chris Ball <chris@printf.net>, ported from Marc "van Hauser" Heuse's "benchmark.sh".
|
|
import argparse, asyncio, json, multiprocessing, os, platform, re, shutil, sys
|
|
from dataclasses import asdict, dataclass
|
|
from decimal import Decimal
|
|
from enum import Enum, auto
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
blue = lambda text: f"\033[1;94m{text}\033[0m"; gray = lambda text: f"\033[1;90m{text}\033[0m"
|
|
green = lambda text: f"\033[0;32m{text}\033[0m"; red = lambda text: f"\033[0;31m{text}\033[0m"
|
|
yellow = lambda text: f"\033[0;33m{text}\033[0m"
|
|
|
|
class Mode(Enum):
|
|
multicore = auto()
|
|
singlecore = auto()
|
|
|
|
@dataclass
|
|
class Target:
|
|
source: Path
|
|
binary: Path
|
|
|
|
@dataclass
|
|
class Run:
|
|
execs_per_sec: float
|
|
execs_total: float
|
|
fuzzers_used: int
|
|
|
|
@dataclass
|
|
class Config:
|
|
afl_persistent_config: bool
|
|
afl_system_config: bool
|
|
afl_version: Optional[str]
|
|
comment: str
|
|
compiler: str
|
|
target_arch: str
|
|
|
|
@dataclass
|
|
class Hardware:
|
|
cpu_fastest_core_mhz: float
|
|
cpu_model: str
|
|
cpu_threads: int
|
|
|
|
@dataclass
|
|
class Results:
|
|
config: Optional[Config]
|
|
hardware: Optional[Hardware]
|
|
targets: Dict[str, Dict[str, Optional[Run]]]
|
|
|
|
all_modes = [Mode.singlecore, Mode.multicore]
|
|
all_targets = [
|
|
Target(source=Path("../utils/persistent_mode/test-instr.c").resolve(), binary=Path("test-instr-persist-shmem")),
|
|
Target(source=Path("../test-instr.c").resolve(), binary=Path("test-instr"))
|
|
]
|
|
modes = [mode.name for mode in all_modes]
|
|
targets = [str(target.binary) for target in all_targets]
|
|
cpu_count = multiprocessing.cpu_count()
|
|
env_vars = {
|
|
"AFL_DISABLE_TRIM": "1", "AFL_I_DONT_CARE_ABOUT_MISSING_CRASHES": "1", "AFL_FAST_CAL": "1",
|
|
"AFL_NO_UI": "1", "AFL_TRY_AFFINITY": "1", "PATH": f'{str(Path("../").resolve())}:{os.environ["PATH"]}',
|
|
}
|
|
|
|
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument("-b", "--basedir", help="directory to use for temp files", type=str, default="/tmp/aflpp-benchmark")
|
|
parser.add_argument("-d", "--debug", help="show verbose debugging output", action="store_true")
|
|
parser.add_argument("-r", "--runs", help="how many runs to average results over", type=int, default=3)
|
|
parser.add_argument("-f", "--fuzzers", help="how many afl-fuzz workers to use", type=int, default=cpu_count)
|
|
parser.add_argument("-m", "--mode", help="pick modes", action="append", default=modes, choices=modes)
|
|
parser.add_argument("-c", "--comment", help="add a comment about your setup", type=str, default="")
|
|
parser.add_argument("--cpu", help="override the detected CPU model name", type=str, default="")
|
|
parser.add_argument("--mhz", help="override the detected CPU MHz", type=str, default="")
|
|
parser.add_argument(
|
|
"-t", "--target", help="pick targets", action="append", default=["test-instr-persist-shmem"], choices=targets
|
|
)
|
|
args = parser.parse_args()
|
|
# Really unsatisfying argparse behavior: we want a default and to allow multiple choices, but if there's a manual choice
|
|
# it should override the default. Seems like we have to remove the default to get that and have correct help text?
|
|
if len(args.target) > 1:
|
|
args.target = args.target[1:]
|
|
if len(args.mode) > 2:
|
|
args.mode = args.mode[2:]
|
|
|
|
chosen_modes = [mode for mode in all_modes if mode.name in args.mode]
|
|
chosen_targets = [target for target in all_targets if str(target.binary) in args.target]
|
|
results = Results(config=None, hardware=None, targets={
|
|
str(t.binary): {m.name: None for m in chosen_modes} for t in chosen_targets}
|
|
)
|
|
debug = lambda text: args.debug and print(blue(text))
|
|
|
|
async def clean_up_tempfiles() -> None:
|
|
shutil.rmtree(f"{args.basedir}/in")
|
|
for target in chosen_targets:
|
|
target.binary.unlink()
|
|
for mode in chosen_modes:
|
|
shutil.rmtree(f"{args.basedir}/out-{mode.name}-{str(target.binary)}")
|
|
|
|
async def check_afl_persistent() -> bool:
|
|
with open("/proc/cmdline", "r") as cmdline:
|
|
return "mitigations=off" in cmdline.read().strip().split(" ")
|
|
|
|
async def check_afl_system() -> bool:
|
|
sysctl = next((s for s in ["sysctl", "/sbin/sysctl"] if shutil.which(s)), None)
|
|
if sysctl:
|
|
(returncode, stdout, _) = await run_command([sysctl, "kernel.randomize_va_space"])
|
|
return returncode == 0 and stdout.decode().rstrip().split(" = ")[1] == "0"
|
|
return False
|
|
|
|
async def prep_env() -> None:
|
|
Path(f"{args.basedir}/in").mkdir(exist_ok=True, parents=True)
|
|
with open(f"{args.basedir}/in/in.txt", "wb") as seed:
|
|
seed.write(b"\x00" * 10240)
|
|
|
|
async def compile_target(source: Path, binary: Path) -> None:
|
|
print(f" [*] Compiling the {binary} fuzzing harness for the benchmark to use.")
|
|
(returncode, stdout, stderr) = await run_command(
|
|
[str(Path("../afl-clang-lto").resolve()), "-o", str(Path(binary.resolve())), str(Path(source).resolve())]
|
|
)
|
|
if returncode == 0:
|
|
return
|
|
print(yellow(f" [*] afl-clang-lto was unable to compile; falling back to afl-cc."))
|
|
(returncode, stdout, stderr) = await run_command(
|
|
[str(Path("../afl-cc").resolve()), "-o", str(Path(binary.resolve())), str(Path(source).resolve())]
|
|
)
|
|
if returncode != 0:
|
|
sys.exit(red(f" [*] Error: afl-cc is unable to compile: {stderr.decode()} {stdout.decode()}"))
|
|
|
|
async def run_command(cmd: List[str]) -> Tuple[Optional[int], bytes, bytes]:
|
|
debug(f"Launching command: {cmd} with env {env_vars}")
|
|
p = await asyncio.create_subprocess_exec(
|
|
*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env_vars
|
|
)
|
|
stdout, stderr = await p.communicate()
|
|
debug(f"Output: {stdout.decode()} {stderr.decode()}")
|
|
return (p.returncode, stdout, stderr)
|
|
|
|
async def check_deps() -> None:
|
|
if not (plat := platform.system()) == "Linux": sys.exit(red(f" [*] {plat} is not supported by this script yet."))
|
|
if not os.access(Path("../afl-fuzz").resolve(), os.X_OK) and os.access(Path("../afl-cc").resolve(), os.X_OK) and (
|
|
os.path.exists(Path("../SanitizerCoveragePCGUARD.so").resolve())):
|
|
sys.exit(red(" [*] Compile AFL++: we need afl-fuzz, afl-clang-fast and SanitizerCoveragePCGUARD.so built."))
|
|
|
|
(returncode, stdout, stderr) = await run_command([str(Path("../afl-cc").resolve()), "-v"])
|
|
if returncode != 0:
|
|
sys.exit(red(f" [*] Error: afl-cc -v returned: {stderr.decode()} {stdout.decode()}"))
|
|
compiler = ""
|
|
target_arch = ""
|
|
for line in stderr.decode().split("\n"):
|
|
if "clang version" in line:
|
|
compiler = line
|
|
elif m := re.match(r"^Target: (.*)", line):
|
|
target_arch = m.group(1)
|
|
|
|
# Pick some sample settings from afl-{persistent,system}-config to try to see whether they were run.
|
|
afl_pc = await check_afl_persistent()
|
|
afl_sc = await check_afl_system()
|
|
if not afl_pc:
|
|
print(yellow(f" [*] afl-persistent-config did not run; run it to improve performance (and decrease security)."))
|
|
if not afl_sc:
|
|
print(yellow(f" [*] afl-system-config did not run; run it to improve performance (and decrease security)."))
|
|
results.config = Config(afl_persistent_config=afl_pc, afl_system_config=afl_sc, afl_version="",
|
|
comment=args.comment, compiler=compiler, target_arch=target_arch)
|
|
|
|
async def colon_values(filename: str, searchKey: str) -> List[str]:
|
|
"""Return a colon-separated value given a key in a file, e.g. 'cpu MHz : 4976.109')"""
|
|
with open(filename, "r") as fh:
|
|
kv_pairs = (line.split(": ", 1) for line in fh if ": " in line)
|
|
v_list = [v.rstrip() for k, v in kv_pairs if k.rstrip() == searchKey]
|
|
return v_list
|
|
|
|
async def describe_afl_config() -> str:
|
|
if results.config is None:
|
|
return "unknown"
|
|
elif results.config.afl_persistent_config and results.config.afl_system_config:
|
|
return "both"
|
|
elif results.config.afl_persistent_config:
|
|
return "persistent"
|
|
elif results.config.afl_system_config:
|
|
return "system"
|
|
else:
|
|
return "none"
|
|
|
|
async def save_benchmark_results() -> None:
|
|
"""Append a single row to the benchmark results in JSON Lines format (which is simple to write and diff)."""
|
|
with open("benchmark-results.jsonl", "a") as jsonfile:
|
|
json.dump(asdict(results), jsonfile, sort_keys=True)
|
|
jsonfile.write("\n")
|
|
print(blue(f" [*] Results have been written to the {jsonfile.name} file."))
|
|
with open("COMPARISON.md", "r+") as comparisonfile:
|
|
described_config = await describe_afl_config()
|
|
aflconfig = described_config.ljust(12)
|
|
if results.hardware is None:
|
|
return
|
|
cpu_model = results.hardware.cpu_model.ljust(51)
|
|
if cpu_model in comparisonfile.read():
|
|
print(blue(f" [*] Results have not been written to the COMPARISON.md file; this CPU is already present."))
|
|
return
|
|
cpu_mhz = str(round(results.hardware.cpu_fastest_core_mhz)).ljust(5)
|
|
if not "test-instr-persist-shmem" in results.targets or \
|
|
not "multicore" in results.targets["test-instr-persist-shmem"] or \
|
|
not "singlecore" in results.targets["test-instr-persist-shmem"] or \
|
|
results.targets["test-instr-persist-shmem"]["singlecore"] is None or \
|
|
results.targets["test-instr-persist-shmem"]["multicore"] is None:
|
|
return
|
|
single = str(round(results.targets["test-instr-persist-shmem"]["singlecore"].execs_per_sec)).ljust(10)
|
|
multi = str(round(results.targets["test-instr-persist-shmem"]["multicore"].execs_per_sec)).ljust(9)
|
|
cores = str(args.fuzzers).ljust(7)
|
|
comparisonfile.write(f"|{cpu_model} | {cpu_mhz} | {cores} | {single} | {multi} | {aflconfig} |\n")
|
|
print(blue(f" [*] Results have been written to the COMPARISON.md file."))
|
|
with open("COMPARISON.md", "r") as comparisonfile:
|
|
print(comparisonfile.read())
|
|
|
|
|
|
async def main() -> None:
|
|
try:
|
|
await clean_up_tempfiles()
|
|
except FileNotFoundError:
|
|
pass
|
|
await check_deps()
|
|
if args.mhz:
|
|
cpu_mhz = float(args.mhz)
|
|
else:
|
|
cpu_mhz_str = await colon_values("/proc/cpuinfo", "cpu MHz")
|
|
if len(cpu_mhz_str) == 0:
|
|
cpu_mhz_str.append("0")
|
|
cpu_mhz = max([float(c) for c in cpu_mhz_str]) # use the fastest CPU MHz for now
|
|
if args.cpu:
|
|
cpu_model = [args.cpu]
|
|
else:
|
|
cpu_model = await colon_values("/proc/cpuinfo", "model name") or [""]
|
|
results.hardware = Hardware(cpu_fastest_core_mhz=cpu_mhz, cpu_model=cpu_model[0], cpu_threads=cpu_count)
|
|
await prep_env()
|
|
print(f" [*] Ready, starting benchmark...")
|
|
for target in chosen_targets:
|
|
await compile_target(target.source, target.binary)
|
|
binary = str(target.binary)
|
|
for mode in chosen_modes:
|
|
if mode == Mode.multicore:
|
|
print(blue(f" [*] Using {args.fuzzers} fuzzers for multicore fuzzing "), end="")
|
|
print(blue("(use --fuzzers to override)." if args.fuzzers == cpu_count else f"(the default is {cpu_count})"))
|
|
execs_per_sec, execs_total = ([] for _ in range(2))
|
|
for run_idx in range(0, args.runs):
|
|
print(gray(f" [*] {mode.name} {binary} run {run_idx+1} of {args.runs}, execs/s: "), end="", flush=True)
|
|
fuzzers = range(0, args.fuzzers if mode == Mode.multicore else 1)
|
|
outdir = f"{args.basedir}/out-{mode.name}-{binary}"
|
|
cmds = []
|
|
for fuzzer_idx, afl in enumerate(fuzzers):
|
|
name = ["-o", outdir, "-M" if fuzzer_idx == 0 else "-S", str(afl)]
|
|
cmds.append(["afl-fuzz", "-i", f"{args.basedir}/in"] + name + ["-s", "123", "-V10", "-D", f"./{binary}"])
|
|
# Prepare the afl-fuzz tasks, and then block while waiting for them to finish.
|
|
fuzztasks = [run_command(cmds[cpu]) for cpu in fuzzers]
|
|
await asyncio.gather(*fuzztasks)
|
|
afl_versions = await colon_values(f"{outdir}/0/fuzzer_stats", "afl_version")
|
|
if results.config:
|
|
results.config.afl_version = afl_versions[0]
|
|
# Our score is the sum of all execs_per_sec entries in fuzzer_stats files for the run.
|
|
sectasks = [colon_values(f"{outdir}/{afl}/fuzzer_stats", "execs_per_sec") for afl in fuzzers]
|
|
all_execs_per_sec = await asyncio.gather(*sectasks)
|
|
execs = sum([Decimal(count[0]) for count in all_execs_per_sec])
|
|
print(green(execs))
|
|
execs_per_sec.append(execs)
|
|
# Also gather execs_total and total_run_time for this run.
|
|
exectasks = [colon_values(f"{outdir}/{afl}/fuzzer_stats", "execs_done") for afl in fuzzers]
|
|
all_execs_total = await asyncio.gather(*exectasks)
|
|
execs_total.append(sum([Decimal(count[0]) for count in all_execs_total]))
|
|
|
|
# (Using float() because Decimal() is not JSON-serializable.)
|
|
avg_afl_execs_per_sec = round(Decimal(sum(execs_per_sec) / len(execs_per_sec)), 2)
|
|
afl_execs_total = int(sum([Decimal(execs) for execs in execs_total]))
|
|
run = Run(execs_per_sec=float(avg_afl_execs_per_sec), execs_total=afl_execs_total, fuzzers_used=len(fuzzers))
|
|
results.targets[binary][mode.name] = run
|
|
print(f" [*] Average execs/sec for this test across all runs was: {green(avg_afl_execs_per_sec)}")
|
|
if (((max(execs_per_sec) - min(execs_per_sec)) / avg_afl_execs_per_sec) * 100) > 15:
|
|
print(yellow(" [*] The difference between your slowest and fastest runs was >15%, maybe try again?"))
|
|
|
|
await clean_up_tempfiles()
|
|
await save_benchmark_results()
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|
|
|