From ae0d8b8d9064f985765e878f7e99dcb5c174dd62 Mon Sep 17 00:00:00 2001
From: vanhauser-thc <vh@thc.org>
Date: Fri, 19 Jul 2024 13:52:22 +0200
Subject: [PATCH] weight power

---
 include/afl-fuzz.h   | 19 ++++++++++---------
 src/afl-fuzz-queue.c | 11 +++++++++++
 src/afl-fuzz.c       |  4 ++++
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 0f0e45d3..2a133b85 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -338,15 +338,16 @@ enum {
 
 enum {
 
-  /* 00 */ EXPLORE, /* AFL default, Exploration-based constant schedule */
-  /* 01 */ MMOPT,   /* Modified MOPT schedule           */
-  /* 02 */ EXPLOIT, /* AFL's exploitation-based const.  */
-  /* 03 */ FAST,    /* Exponential schedule             */
-  /* 04 */ COE,     /* Cut-Off Exponential schedule     */
-  /* 05 */ LIN,     /* Linear schedule                  */
-  /* 06 */ QUAD,    /* Quadratic schedule               */
-  /* 07 */ RARE,    /* Rare edges                       */
-  /* 08 */ SEEK,    /* EXPLORE that ignores timings     */
+  /* 00 */ EXPLORE, /* AFL default, exploration-based constant schedule */
+  /* 01 */ EXPLOIT, /* AFL's exploitation-based const.   */
+  /* 02 */ WEIGHT,  /* Based on seed weighting algorithm */
+  /* 03 */ MMOPT,   /* Modified MOPT schedule            */
+  /* 04 */ FAST,    /* Exponential schedule              */
+  /* 05 */ COE,     /* Cut-Off Exponential schedule      */
+  /* 06 */ LIN,     /* Linear schedule                   */
+  /* 07 */ QUAD,    /* Quadratic schedule                */
+  /* 08 */ RARE,    /* Rare edges                        */
+  /* 09 */ SEEK,    /* EXPLORE that ignores timings      */
 
   POWER_SCHEDULES_NUM
 
diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index 999929a1..639fb711 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -1020,6 +1020,17 @@ void cull_queue(afl_state_t *afl) {
 
 u32 calculate_score(afl_state_t *afl, struct queue_entry *q) {
 
+  if (likely(afl->schedule == WEIGHT)) {
+
+    u32 val = 100;
+
+    if (unlikely(q->favored)) { val = val << 1; }
+    if (unlikely(!q->was_fuzzed)) { val = val << 1; }
+
+    return val * q->weight;
+
+  }
+
   u32 cal_cycles = afl->total_cal_cycles;
   u32 bitmap_entries = afl->total_bitmap_entries;
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 9867eba3..eecf8f69 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -745,6 +745,10 @@ int main(int argc, char **argv_orig, char **envp) {
 
           afl->schedule = EXPLOIT;
 
+        } else if (!stricmp(optarg, "weight")) {
+
+          afl->schedule = WEIGHT;
+
         } else if (!stricmp(optarg, "lin")) {
 
           afl->schedule = LIN;