mirror of
https://github.com/AFLplusplus/AFLplusplus.git
synced 2025-06-11 17:51:32 +00:00
Add support for multi-core benchmarking
This commit is contained in:
Chris Ball
@ -2,13 +2,89 @@
|
||||
# Requires Python 3.6+.
|
||||
# Author: Chris Ball <chris@printf.net>
|
||||
# Ported from Marc "van Hauser" Heuse's "benchmark.sh".
|
||||
import asyncio
|
||||
import glob
|
||||
import json
|
||||
import multiprocessing
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from decimal import Decimal
|
||||
|
||||
def colon_value_or_none(filename: str, searchKey: str) -> str | None:
|
||||
debug = False
|
||||
|
||||
targets = [
|
||||
{"source": "../test-instr.c", "binary": "test-instr"},
|
||||
{"source": "../utils/persistent_mode/test-instr.c", "binary": "test-instr-persistent-shmem"},
|
||||
]
|
||||
modes = ["single-core", "multi-core"]
|
||||
results = {}
|
||||
|
||||
colors = {
|
||||
"blue": "\033[1;94m",
|
||||
"gray": "\033[1;90m",
|
||||
"green": "\033[0;32m",
|
||||
"red": "\033[0;31m",
|
||||
"reset": "\033[0m",
|
||||
}
|
||||
|
||||
async def clean_up() -> None:
|
||||
"""Remove temporary files."""
|
||||
shutil.rmtree("in")
|
||||
for target in targets:
|
||||
# os.remove(target["binary"])
|
||||
for mode in modes:
|
||||
for outdir in glob.glob(f"/tmp/out-{mode}-{target['binary']}*"):
|
||||
shutil.rmtree(outdir)
|
||||
|
||||
async def check_deps() -> None:
|
||||
"""Check if the necessary files exist and are executable."""
|
||||
if not (os.access("../afl-fuzz", os.X_OK) and os.access("../afl-cc", os.X_OK) and os.path.exists("../SanitizerCoveragePCGUARD.so")):
|
||||
sys.exit(f"{colors['red']}Error: you need to compile AFL++ first, we need afl-fuzz, afl-clang-fast and SanitizerCoveragePCGUARD.so built.{colors['reset']}")
|
||||
|
||||
async def prep_env() -> dict:
|
||||
# Unset AFL_* environment variables
|
||||
for e in list(os.environ.keys()):
|
||||
if e.startswith("AFL_"):
|
||||
os.environ.pop(e)
|
||||
# Create input directory and file
|
||||
os.makedirs("in", exist_ok=True)
|
||||
with open("in/in.txt", "wb") as f:
|
||||
f.write(b"\x00" * 10240)
|
||||
# Rest of env
|
||||
AFL_PATH = os.path.abspath("../")
|
||||
os.environ["PATH"] = AFL_PATH + ":" + os.environ["PATH"]
|
||||
return {
|
||||
"AFL_BENCH_JUST_ONE": "1",
|
||||
"AFL_DISABLE_TRIM": "1",
|
||||
"AFL_I_DONT_CARE_ABOUT_MISSING_CRASHES": "1",
|
||||
"AFL_NO_UI": "1",
|
||||
"AFL_TRY_AFFINITY": "1",
|
||||
"PATH": f"{AFL_PATH}:{os.environ['PATH']}",
|
||||
}
|
||||
|
||||
async def compile_target(source: str, binary: str) -> None:
|
||||
(returncode, stdout, stderr) = await run_command(
|
||||
["afl-cc", "-o", binary, source],
|
||||
env={"AFL_INSTRUMENT": "PCGUARD", "PATH": os.environ["PATH"]},
|
||||
)
|
||||
if returncode != 0:
|
||||
sys.exit(f"{colors['red']} [*] Error: afl-cc is unable to compile: {stderr} {stdout}{colors['reset']}")
|
||||
|
||||
async def cool_down() -> None:
|
||||
"""Avoid the next test run's results being contaminated by e.g. thermal limits hit on this one."""
|
||||
print(f"{colors['blue']}Taking a five second break to stay cool.{colors['reset']}")
|
||||
await asyncio.sleep(10)
|
||||
|
||||
async def run_command(args, env) -> (int | None, bytes, bytes):
|
||||
if debug:
|
||||
print(f"\n{colors['blue']}Launching command: {args} with env {env}{colors['reset']}")
|
||||
p = await asyncio.create_subprocess_exec(*args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=env)
|
||||
stdout, stderr = await p.communicate()
|
||||
return (p.returncode, stdout, stderr)
|
||||
|
||||
async def colon_value_or_none(filename: str, searchKey: str) -> str | None:
|
||||
"""Read a value (e.g. 'cpu MHz : 4976.109') given its filename and key."""
|
||||
with open(filename, "r") as fh:
|
||||
for line in fh:
|
||||
kv = line.split(": ", 1)
|
||||
@ -20,123 +96,41 @@ def colon_value_or_none(filename: str, searchKey: str) -> str | None:
|
||||
return value
|
||||
return None
|
||||
|
||||
def compile_target(source: str, binary: str) -> None:
|
||||
with open("afl.log", "w") as f:
|
||||
process = subprocess.run(
|
||||
["afl-cc", "-o", binary, source],
|
||||
stdout=f,
|
||||
stderr=subprocess.STDOUT,
|
||||
env={"AFL_INSTRUMENT": "PCGUARD", "PATH": os.environ["PATH"]}
|
||||
)
|
||||
if process.returncode != 0:
|
||||
sys.exit("Error: afl-cc is unable to compile")
|
||||
async def main() -> None:
|
||||
# Remove stale files, if necessary.
|
||||
try:
|
||||
await clean_up()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
# Check if the necessary files exist and are executable
|
||||
if not (
|
||||
os.access("../afl-fuzz", os.X_OK)
|
||||
and os.access("../afl-cc", os.X_OK)
|
||||
and os.path.exists("../SanitizerCoveragePCGUARD.so")
|
||||
):
|
||||
sys.exit("Error: you need to compile AFL++ first, we need afl-fuzz, afl-clang-fast and SanitizerCoveragePCGUARD.so built.")
|
||||
await check_deps()
|
||||
env_vars = await prep_env()
|
||||
cpu_count = multiprocessing.cpu_count()
|
||||
print(f"{colors['gray']} [*] Preparing environment{colors['reset']}")
|
||||
print(f"{colors['gray']} [*] Ready, starting benchmark - this will take approx 1-2 minutes...{colors['reset']}")
|
||||
for target in targets:
|
||||
await compile_target(target["source"], target["binary"])
|
||||
for mode in modes:
|
||||
await cool_down()
|
||||
print(f" [*] {mode} {target['binary']} benchmark starting, execs/s: ", end="", flush=True)
|
||||
if mode == "single-core":
|
||||
cpus = [0]
|
||||
elif mode == "multi-core":
|
||||
cpus = range(0, cpu_count)
|
||||
basedir = f"/tmp/out-{mode}-{target['binary']}-"
|
||||
args = [["afl-fuzz", "-i", "in", "-o", f"{basedir}{cpu}", "-M", f"{cpu}", "-s", "123", "-D", f"./{target['binary']}"] for cpu in cpus]
|
||||
tasks = [run_command(args[cpu], env_vars) for cpu in cpus]
|
||||
output = await asyncio.gather(*tasks)
|
||||
if debug:
|
||||
for _, (_, stdout, stderr) in enumerate(output):
|
||||
print(f"{colors['blue']}Output: {stdout} {stderr}{colors['reset']}")
|
||||
execs = sum([Decimal(await colon_value_or_none(f"{basedir}{cpu}/{cpu}/fuzzer_stats", "execs_per_sec")) for cpu in cpus])
|
||||
print(f"{colors['green']}{execs}{colors['reset']}")
|
||||
|
||||
print("Preparing environment")
|
||||
|
||||
targets = [
|
||||
{"source": "../test-instr.c", "binary": "test-instr"},
|
||||
{"source": "../utils/persistent_mode/test-instr.c", "binary": "test-instr-persistent"}
|
||||
]
|
||||
|
||||
# Unset AFL_* environment variables
|
||||
for e in list(os.environ.keys()):
|
||||
if e.startswith("AFL_"):
|
||||
os.environ.pop(e)
|
||||
|
||||
AFL_PATH = os.path.abspath("../")
|
||||
os.environ["PATH"] = AFL_PATH + ":" + os.environ["PATH"]
|
||||
|
||||
for target in targets:
|
||||
compile_target(target["source"], target["binary"])
|
||||
|
||||
# Create input directory and file
|
||||
os.makedirs("in", exist_ok=True)
|
||||
with open("in/in.txt", "wb") as f:
|
||||
f.write(b"\x00" * 10240)
|
||||
|
||||
print("Ready, starting benchmark - this will take approx 20-30 seconds ...")
|
||||
|
||||
# Run afl-fuzz
|
||||
env_vars = {
|
||||
"AFL_DISABLE_TRIM": "1",
|
||||
"AFL_NO_UI": "1",
|
||||
"AFL_TRY_AFFINITY": "1",
|
||||
"AFL_I_DONT_CARE_ABOUT_MISSING_CRASHES": "1",
|
||||
"AFL_BENCH_JUST_ONE": "1",
|
||||
}
|
||||
|
||||
for target in targets:
|
||||
with open(f"afl-{target['binary']}.log", "a") as f:
|
||||
process = subprocess.run(
|
||||
[
|
||||
"afl-fuzz",
|
||||
"-i",
|
||||
"in",
|
||||
"-o",
|
||||
f"out-{target['binary']}",
|
||||
"-s",
|
||||
"123",
|
||||
"-D",
|
||||
f"./{target['binary']}",
|
||||
],
|
||||
stdout=f,
|
||||
stderr=subprocess.STDOUT,
|
||||
env={**os.environ, **env_vars},
|
||||
)
|
||||
|
||||
print("Analysis:")
|
||||
|
||||
# Extract CPUID from afl.log
|
||||
with open(f"afl-test-instr.log", "r") as f:
|
||||
match = re.search(r".*try binding to.*#(\d+)", f.read())
|
||||
if not match:
|
||||
sys.exit("Couldn't see which CPU# was used in afl.log", 1)
|
||||
cpuid = match.group(1)
|
||||
|
||||
# Print CPU model
|
||||
model = colon_value_or_none("/proc/cpuinfo", "model name")
|
||||
if model:
|
||||
print(" CPU:", model)
|
||||
|
||||
# Print CPU frequency
|
||||
cpu_speed = None
|
||||
with open("/proc/cpuinfo", "r") as fh:
|
||||
current_cpu = None
|
||||
for line in fh:
|
||||
kv = line.split(": ", 1)
|
||||
if kv and len(kv) == 2:
|
||||
(key, value) = kv
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
if key == "processor":
|
||||
current_cpu = value
|
||||
elif key == "cpu MHz" and current_cpu == cpuid:
|
||||
cpu_speed = value
|
||||
if cpu_speed:
|
||||
print(" Mhz:", cpu_speed)
|
||||
|
||||
# Print execs_per_sec from fuzzer_stats
|
||||
for target in targets:
|
||||
execs = colon_value_or_none(f"out-{target['binary']}/default/fuzzer_stats", "execs_per_sec")
|
||||
if execs:
|
||||
print(f" {target['binary']} single-core execs/s:", execs)
|
||||
|
||||
print("\nComparison: (note that values can change by 10-15% per run)")
|
||||
with open("COMPARISON", "r") as f:
|
||||
print(f.read())
|
||||
|
||||
# Clean up
|
||||
os.remove("afl.log")
|
||||
shutil.rmtree("in")
|
||||
for target in targets:
|
||||
shutil.rmtree(f"out-{target['binary']}")
|
||||
os.remove(target["binary"])
|
||||
print("\nComparison: (note that values can change by 10-20% per run)")
|
||||
with open("COMPARISON", "r") as f:
|
||||
print(f.read())
|
||||
await clean_up()
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user