From 711d76dbcc37f0c0e2f0529b7ad438ff67cb0f12 Mon Sep 17 00:00:00 2001
From: Matias Gobbi <matiasgobbi24@gmail.com>
Date: Fri, 18 Nov 2022 09:52:41 +0100
Subject: [PATCH 1/2] afl++: fix to CRASH_DIR path

---
 fuzzers/aflplusplus/findings.sh          | 2 +-
 fuzzers/aflplusplus_lto/findings.sh      | 2 +-
 fuzzers/aflplusplus_lto_asan/findings.sh | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/fuzzers/aflplusplus/findings.sh b/fuzzers/aflplusplus/findings.sh
index afc482c57..ef93d9b4b 100755
--- a/fuzzers/aflplusplus/findings.sh
+++ b/fuzzers/aflplusplus/findings.sh
@@ -5,7 +5,7 @@
 # - env SHARED: path to directory shared with host (to store results)
 ##
 
-CRASH_DIR="$SHARED/findings/crashes"
+CRASH_DIR="$SHARED/findings/default/crashes"
 
 if [ ! -d "$CRASH_DIR" ]; then
     exit 1
diff --git a/fuzzers/aflplusplus_lto/findings.sh b/fuzzers/aflplusplus_lto/findings.sh
index afc482c57..ef93d9b4b 100755
--- a/fuzzers/aflplusplus_lto/findings.sh
+++ b/fuzzers/aflplusplus_lto/findings.sh
@@ -5,7 +5,7 @@
 # - env SHARED: path to directory shared with host (to store results)
 ##
 
-CRASH_DIR="$SHARED/findings/crashes"
+CRASH_DIR="$SHARED/findings/default/crashes"
 
 if [ ! -d "$CRASH_DIR" ]; then
     exit 1
diff --git a/fuzzers/aflplusplus_lto_asan/findings.sh b/fuzzers/aflplusplus_lto_asan/findings.sh
index afc482c57..ef93d9b4b 100755
--- a/fuzzers/aflplusplus_lto_asan/findings.sh
+++ b/fuzzers/aflplusplus_lto_asan/findings.sh
@@ -5,7 +5,7 @@
 # - env SHARED: path to directory shared with host (to store results)
 ##
 
-CRASH_DIR="$SHARED/findings/crashes"
+CRASH_DIR="$SHARED/findings/default/crashes"
 
 if [ ! -d "$CRASH_DIR" ]; then
     exit 1

From 8083d0a87261c9cd2eff2886a41e2184d83aadcf Mon Sep 17 00:00:00 2001
From: Adrian Herrera <adrian.herrera02@gmail.com>
Date: Sun, 27 Nov 2022 09:09:09 +1100
Subject: [PATCH 2/2] kscheduler: resize coverage map depending on katz
 centrality analysis

---
 fuzzers/k_scheduler/build.sh      |  7 ++-
 fuzzers/k_scheduler/fetch.sh      | 31 ++++++++++++
 fuzzers/k_scheduler/instrument.sh | 83 ++++++++++++++++++++++---------
 fuzzers/k_scheduler/run.sh        |  2 +-
 4 files changed, 94 insertions(+), 29 deletions(-)

diff --git a/fuzzers/k_scheduler/build.sh b/fuzzers/k_scheduler/build.sh
index 10d882052..5ad4b7417 100755
--- a/fuzzers/k_scheduler/build.sh
+++ b/fuzzers/k_scheduler/build.sh
@@ -14,6 +14,8 @@ fi
 cd "$FUZZER/repo/kscheduler/afl_integration/afl-2.52b_kscheduler"
 make clean
 make -j $(nproc)
+make -j $(nproc) -C llvm_mode ../afl-llvm-rt.o
+cp afl-llvm-rt.o $OUT
 
 export PATH="/usr/local/go/bin:$PATH"
 export GOPATH="$FUZZER/repo/go"
@@ -27,10 +29,7 @@ export LLVM_CXX_NAME="clang++"
 export CC="gclang"
 export CXX="gclang++"
 
-# Build AFL driver and runtime
+# Build AFL driver
 $CXX -std=c++11 -c \
     "$FUZZER/repo/kscheduler/libfuzzer_integration/llvm_11.0.1/compiler-rt/lib/fuzzer/afl/afl_driver.cpp" \
     -o "$OUT/afl_driver.o"
-$CC -c -w \
-    "$FUZZER/repo/kscheduler/afl_integration/afl-2.52b_kscheduler/llvm_mode/afl-llvm-rt.o.c" \
-    -o "$OUT/afl-llvm-rt.o"
diff --git a/fuzzers/k_scheduler/fetch.sh b/fuzzers/k_scheduler/fetch.sh
index df4c8f7c1..efa724f2d 100755
--- a/fuzzers/k_scheduler/fetch.sh
+++ b/fuzzers/k_scheduler/fetch.sh
@@ -19,3 +19,34 @@ sed -i '{s/^int main/__attribute__((weak)) &/}' \
     "$FUZZER/repo/kscheduler/libfuzzer_integration/llvm_11.0.1/compiler-rt/lib/fuzzer/afl/afl_driver.cpp"
 sed -i '{s/^int LLVMFuzzerTestOneInput/__attribute__((weak)) &/}' \
     "$FUZZER/repo/kscheduler/libfuzzer_integration/llvm_11.0.1/compiler-rt/lib/fuzzer/afl/afl_driver.cpp"
+sed -i '{s/##SIG_AFL_PERSISTENT##/##SIG_AFL_NOT_PERSISTENT##/}' \
+    "$FUZZER/repo/kscheduler/libfuzzer_integration/llvm_11.0.1/compiler-rt/lib/fuzzer/afl/afl_driver.cpp"
+
+patch -p1 -d "$FUZZER/repo/kscheduler" << EOF
+index 8a09b93b0..794682b86 100644
+--- a/afl_integration/afl-2.52b_kscheduler/config.h
++++ b/afl_integration/afl-2.52b_kscheduler/config.h
+@@ -315,7 +315,9 @@
+    problems with complex programs). You need to recompile the target binary
+    after changing this - otherwise, SEGVs may ensue. */
+
++#if !defined(MAP_SIZE_POW2)
+ #define MAP_SIZE_POW2       16
++#endif
+ #define MAP_SIZE            (1 << MAP_SIZE_POW2)
+
+ /* Maximum allocator request size (keep well under INT_MAX): */
+index e3675d9f8..b287dccb6 100644
+--- a/afl_integration/afl-2.52b_kscheduler_large_bitmap/config.h
++++ b/afl_integration/afl-2.52b_kscheduler_large_bitmap/config.h
+@@ -315,7 +315,9 @@
+    problems with complex programs). You need to recompile the target binary
+    after changing this - otherwise, SEGVs may ensue. */
+
++#if !defined(MAP_SIZE_POW2)
+ #define MAP_SIZE_POW2       17
++#endif
+ #define MAP_SIZE            (1 << MAP_SIZE_POW2)
+
+ /* Maximum allocator request size (keep well under INT_MAX): */
+EOF
diff --git a/fuzzers/k_scheduler/instrument.sh b/fuzzers/k_scheduler/instrument.sh
index 640aae778..7a82b749c 100755
--- a/fuzzers/k_scheduler/instrument.sh
+++ b/fuzzers/k_scheduler/instrument.sh
@@ -10,41 +10,76 @@ set -e
 # - env CFLAGS and CXXFLAGS must be set to link against Magma instrumentation
 ##
 
+ORIG_CFLAGS=$CFLAGS
+ORIG_CXXFLAGS=$CXXFLAGS
+ORIG_LDFLAGS=$LDFLAGS
+ORIG_LIBS=$LIBS
+
 export PATH="/usr/local/go/bin:$PATH"
 export GOPATH="$FUZZER/repo/go"
 export PATH="$GOPATH/bin:$PATH"
 
-export CFLAGS="$CFLAGS -O2 -fsanitize-coverage=trace-pc-guard,no-prune -fno-omit-frame-pointer -gline-tables-only -fsanitize=fuzzer-no-link"
-export CXXFLAGS="$CXXFLAGS -O2 -fsanitize-coverage=trace-pc-guard,no-prune -fno-omit-frame-pointer -gline-tables-only -fsanitize=fuzzer-no-link"
-export LDFLAGS="$LDFLAGS -fsanitize=fuzzer-no-link"
+(
+  export CFLAGS="$ORIG_CFLAGS -O2 -fsanitize-coverage=trace-pc-guard,no-prune -fno-omit-frame-pointer -gline-tables-only -fsanitize=fuzzer-no-link"
+  export CXXFLAGS="$ORIG_CXXFLAGS -O2 -fsanitize-coverage=trace-pc-guard,no-prune -fno-omit-frame-pointer -gline-tables-only -fsanitize=fuzzer-no-link"
+  export LDFLAGS="$ORIG_LDFLAGS -fsanitize=fuzzer-no-link"
+  export LIBS="$ORIG_LIBS -l:afl_driver.o -l:afl-llvm-rt.o -lstdc++"
+
+  export LLVM_CC_NAME="clang"
+  export LLVM_CXX_NAME="clang++"
+  export CC="gclang"
+  export CXX="gclang++"
 
-export LLVM_CC_NAME="clang"
-export LLVM_CXX_NAME="clang++"
-export CC="gclang"
-export CXX="gclang++"
+  "$MAGMA/build.sh"
+  "$TARGET/build.sh"
 
-export LIBS="$LIBS -l:afl_driver.o -l:afl-llvm-rt.o -lstdc++"
+  cd $OUT
+  source "$TARGET/configrc"
 
-"$MAGMA/build.sh"
-"$TARGET/build.sh"
+  for P in "${PROGRAMS[@]}"; do
+    mkdir -p "$OUT/${P}_out"
+    cd "$OUT/${P}_out"
 
-cd $OUT
-source "$TARGET/configrc"
+    get-bc -o "$P.bc" "$OUT/$P"
+    llvm-dis "$P.bc"
+    python3 "$FUZZER/repo/kscheduler/afl_integration/build_example/fix_long_fun_name.py" "$P.ll"
+    opt-11 -dot-cfg "${P}_fix.ll"
 
-for P in "${PROGRAMS[@]}"; do
-  get-bc "$P"
+    mkdir -p cfgs
+    for f in $(ls -a | grep '^\.*' | grep dot); do mv $f "cfgs/${f:1}"; done
 
-  llvm-dis "$P.bc"
-  python3 "$FUZZER/repo/kscheduler/afl_integration/build_example/fix_long_fun_name.py" "$P.ll"
-  mkdir -p "$OUT/cfg_out_$P"
-  cd "$OUT/cfg_out_$P"
-  opt-11 -dot-cfg "$OUT/${P}_fix.ll"
-  for f in $(ls -a | grep '^\.*' | grep dot); do mv $f ${f:1}; done
+    python3 "$FUZZER/repo/kscheduler/afl_integration/build_example/gen_graph.py" \
+        "${P}_fix.ll" cfgs
 
-  cd $OUT
-  python3 "$FUZZER/repo/kscheduler/afl_integration/build_example/gen_graph.py" \
-      "${P}_fix.ll" "cfg_out_$P"
-done
+    # We need to configure the AFL map so that it fits all of the CFG edges. So
+    # save the size required for this program
+    MAP_SIZE="$(wc -l < katz_cent)"
+    MAP_SIZE_POW2=$(python3 -c "from math import ceil, log2; print('%d' % ceil(log2(${MAP_SIZE})))")
+    echo $MAP_SIZE_POW2 >> "$OUT/map_sizes"
+  done
+)
+
+# Determine the largest map size (amongst all the programs) and recompile AFL
+# and the target with that map size
+MAP_SIZE_POW2=$(sort -nr "$OUT/map_sizes" | head -n1)
+if [[ "${MAP_SIZE_POW2}" -gt "16" ]]; then
+  (
+    export CFLAGS="-DMAP_SIZE_POW2=${MAP_SIZE_POW2}"
+    "$FUZZER/build.sh"
+
+    export CFLAGS="$ORIG_CFLAGS -O2 -fsanitize-coverage=trace-pc-guard,no-prune -fno-omit-frame-pointer -gline-tables-only -fsanitize=fuzzer-no-link"
+    export CXXFLAGS="$ORIG_CXXFLAGS -O2 -fsanitize-coverage=trace-pc-guard,no-prune -fno-omit-frame-pointer -gline-tables-only -fsanitize=fuzzer-no-link"
+    export LDFLAGS="$ORIG_LDFLAGS -fsanitize=fuzzer-no-link"
+    export LIBS="$ORIG_LIBS -l:afl_driver.o -l:afl-llvm-rt.o -lstdc++"
+
+    export LLVM_CC_NAME="clang"
+    export LLVM_CXX_NAME="clang++"
+    export CC="gclang"
+    export CXX="gclang++"
+
+    "$TARGET/build.sh"
+  )
+fi
 
 # NOTE: We pass $OUT directly to the target build.sh script, since the artifact
 #       itself is the fuzz target. In the case of Angora, we might need to
diff --git a/fuzzers/k_scheduler/run.sh b/fuzzers/k_scheduler/run.sh
index 53db18e28..e8ce27e2f 100755
--- a/fuzzers/k_scheduler/run.sh
+++ b/fuzzers/k_scheduler/run.sh
@@ -13,7 +13,7 @@
 
 mkdir -p "$SHARED/findings"
 
-cd $OUT
+cd "$OUT/${PROGRAM}_out"
 python3 "$FUZZER/repo/kscheduler/afl_integration/build_example/gen_dyn_weight.py" 2>&1 &
 sleep 5s