From fd82e3330cc9e63e6c14406e1366a4fa2a1a0d6b Mon Sep 17 00:00:00 2001
From: vanhauser-thc <vh@thc.org>
Date: Tue, 4 Jun 2024 13:45:17 +0200
Subject: [PATCH] add model

---
 GNUmakefile          |  2 +-
 src/afl-fuzz-queue.c | 47 +++++++++++++++++++++++++++++++++-----------
 2 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/GNUmakefile b/GNUmakefile
index dee9bbb3..0b3dbdca 100644
--- a/GNUmakefile
+++ b/GNUmakefile
@@ -472,7 +472,7 @@ src/afl-sharedmem.o : $(COMM_HDR) src/afl-sharedmem.c include/sharedmem.h
 	$(CC) $(CFLAGS) $(CFLAGS_FLTO) $(SPECIAL_PERFORMANCE) -c src/afl-sharedmem.c -o src/afl-sharedmem.o
 
 afl-fuzz: $(COMM_HDR) include/afl-fuzz.h $(AFL_FUZZ_FILES) src/afl-common.o src/afl-sharedmem.o src/afl-forkserver.o src/afl-performance.o src/hashmap.c | test_x86
-	$(CC) $(CFLAGS) $(COMPILE_STATIC) $(CFLAGS_FLTO) $(SPECIAL_PERFORMANCE) -Wno-shift-count-overflow $(AFL_FUZZ_FILES) src/afl-common.o src/afl-sharedmem.o src/afl-forkserver.o src/afl-performance.o src/hashmap.c -o $@ $(PYFLAGS) $(LDFLAGS) -lm
+	$(CC) $(CFLAGS) $(COMPILE_STATIC) $(CFLAGS_FLTO) $(SPECIAL_PERFORMANCE) -Wno-shift-count-overflow $(AFL_FUZZ_FILES) src/afl-common.o src/afl-sharedmem.o src/afl-forkserver.o src/afl-performance.o src/hashmap.c -o $@ $(PYFLAGS) $(LDFLAGS) -lm -lxgboost
 
 afl-showmap: src/afl-showmap.c src/afl-common.o src/afl-sharedmem.o src/afl-forkserver.o src/afl-performance.o $(COMM_HDR) | test_x86
 	$(CC) $(CFLAGS) $(COMPILE_STATIC) $(CFLAGS_FLTO) $(SPECIAL_PERFORMANCE) src/$@.c src/afl-fuzz-mutators.c src/afl-fuzz-python.c src/afl-common.o src/afl-sharedmem.o src/afl-forkserver.o src/afl-performance.o -o $@ $(PYFLAGS) $(LDFLAGS)
diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index 784b377a..dd3ceb6f 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -26,6 +26,7 @@
 #include <limits.h>
 #include <ctype.h>
 #include <math.h>
+#include <xgboost/c_api.h>
 
 #ifdef _STANDALONE_MODULE
 void minimize_bits(afl_state_t *afl, u8 *dst, u8 *src) {
@@ -117,7 +118,7 @@ void create_alias_table(afl_state_t *afl) {
 
     double avg_exec_us = 0.0;
     double avg_bitmap_size = 0.0;
-    double avg_top_size = 0.0;
+    double avg_len = 0.0;
     u32    active = 0;
 
     for (i = 0; i < n; i++) {
@@ -128,8 +129,8 @@ void create_alias_table(afl_state_t *afl) {
       if (likely(!q->disabled)) {
 
         avg_exec_us += q->exec_us;
-        avg_bitmap_size += log(q->bitmap_size);
-        avg_top_size += q->tc_ref;
+        avg_bitmap_size += q->bitmap_size;
+        avg_len += q->len;
         ++active;
 
       }
@@ -138,7 +139,10 @@ void create_alias_table(afl_state_t *afl) {
 
     avg_exec_us /= active;
     avg_bitmap_size /= active;
-    avg_top_size /= active;
+    avg_len /= active;
+
+    float *table = malloc((active + 1) * 3 * sizeof(float));
+    float *pentry = table;
 
     for (i = 0; i < n; i++) {
 
@@ -146,29 +150,48 @@ void create_alias_table(afl_state_t *afl) {
 
       if (likely(!q->disabled)) {
 
-        q->weight =
-            compute_weight(afl, q, avg_exec_us, avg_bitmap_size, avg_top_size);
+        *pentry++ = q->len / avg_len;
+        *pentry++ = q->exec_us / avg_exec_us;
+        *pentry++ = q->bitmap_size / avg_bitmap_size;
         q->perf_score = calculate_score(afl, q);
-        sum += q->weight;
 
       }
 
     }
 
-    if (unlikely(afl->schedule == MMOPT) && afl->queued_discovered) {
+    DMatrixHandle dtest;
+    BoosterHandle booster;
 
-      u32 cnt = afl->queued_discovered >= 5 ? 5 : afl->queued_discovered;
+    // Erstellen einer DMatrix aus dem Array
+    XGDMatrixCreateFromMat((float *)table, 3, active, -1, &dtest);
+    XGBoosterCreate(&dtest, 1, &booster);
+    XGBoosterLoadModel(booster, "./model.bin");
+    
+        bst_ulong out_len;
+    const float *predictions;
+    XGBoosterPredict(booster, dtest, 0, 0, 0, &out_len, &predictions);
 
-      for (i = n - cnt; i < n; i++) {
+    // Ausgabe der Vorhersagen
+    int count = 0;
+    for (i = 0; i < n; i++) {
 
-        struct queue_entry *q = afl->queue_buf[i];
+      struct queue_entry *q = afl->queue_buf[i];
 
-        if (likely(!q->disabled)) { q->weight *= 2.0; }
+      if (likely(!q->disabled)) {
+
+        fprintf(stderr, "Prediction[%u] = %f\n", i, predictions[count]);
+        afl->queue_buf[i]->weight = predictions[count++];
+        sum += predictions[count++];
 
       }
 
     }
 
+    // Freigeben der Ressourcen
+    XGBoosterFree(booster);
+    XGDMatrixFree(dtest);
+    free(table);
+
     for (i = 0; i < n; i++) {
 
       // weight is always 0 for disabled entries