add -b option to afl-fuzz

3 files changed, 34 insertions, 8 deletions

diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index 65ad0c9f..ad92dff6 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -53,6 +53,13 @@ void bind_to_free_cpu(afl_state_t *afl) {
   u8  cpu_used[4096] = {0}, lockfile[PATH_MAX] = "";
   u32 i;
 
+  if (afl->cpu_to_bind != -1) {
+
+    i = afl->cpu_to_bind;
+    goto set_cpu;
+
+  }
+
   if (afl->sync_id) {
 
     s32 lockfd, first = 1;
@@ -295,20 +302,23 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
   try:
 
+    if (afl->cpu_to_bind != -1)
+      FATAL("bind to CPU #%d failed!", afl->cpu_to_bind);
+
   #if !defined(__ANDROID__)
 
-    for (i = cpu_start; i < afl->cpu_core_count; i++) {
+  for (i = cpu_start; i < afl->cpu_core_count; i++) {
 
-      if (!cpu_used[i]) { break; }
+    if (!cpu_used[i]) { break; }
 
-    }
+  }
 
   if (i == afl->cpu_core_count) {
 
   #else
 
-    for (i = afl->cpu_core_count - cpu_start - 1; i > -1; i--)
-      if (!cpu_used[i]) break;
+  for (i = afl->cpu_core_count - cpu_start - 1; i > -1; i--)
+    if (!cpu_used[i]) break;
   if (i == -1) {
 
   #endif
@@ -327,6 +337,8 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
   OKF("Found a free CPU core, try binding to #%u.", i);
 
+set_cpu:
+
   afl->cpu_aff = i;
 
   #if defined(__linux__) || defined(__FreeBSD__) || defined(__DragonFly__)
diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index 66280ed1..e2d62bc6 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -94,6 +94,7 @@ void afl_state_init(afl_state_t *afl, uint32_t map_size) {
   afl->havoc_div = 1;                   /* Cycle count divisor for havoc    */
   afl->stage_name = "init";             /* Name of the current fuzz stage   */
   afl->splicing_with = -1;              /* Splicing with which test case?   */
+  afl->cpu_to_bind = -1;
 
 #ifdef HAVE_AFFINITY
   afl->cpu_aff = -1;                    /* Selected CPU core                */
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 5bedf6e1..e33a4bbd 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -143,6 +143,8 @@ static void usage(afl_state_t *afl, u8 *argv0, int more_help) {
       //"  -B bitmap.txt - mutate a specific test case, use the out/fuzz_bitmap
       //" "file\n"
       "  -C            - crash exploration mode (the peruvian rabbit thing)\n"
+      "  -b cpu_id     - bind the fuzzing process to the specified CPU core "
+      "(0-...)\n"
       "  -e ext        - file extension for the fuzz test input file (if "
       "needed)\n\n",
       argv0, EXEC_TIMEOUT, MEM_LIMIT, FOREIGN_SYNCS_MAX);
@@ -271,9 +273,9 @@ int main(int argc, char **argv_orig, char **envp) {
 
   afl->shmem_testcase_mode = 1;  // we always try to perform shmem fuzzing
 
-  while ((opt = getopt(argc, argv,
-                       "+c:i:I:o:f:F:m:t:T:dDnCB:S:M:x:QNUWe:p:s:V:E:L:hRP:")) >
-         0) {
+  while ((opt = getopt(
+              argc, argv,
+              "+b:c:i:I:o:f:F:m:t:T:dDnCB:S:M:x:QNUWe:p:s:V:E:L:hRP:")) > 0) {
 
     switch (opt) {
 
@@ -281,6 +283,17 @@ int main(int argc, char **argv_orig, char **envp) {
         afl->infoexec = optarg;
         break;
 
+      case 'b': {                                          /* bind CPU core */
+
+        if (afl->cpu_to_bind != -1) FATAL("Multiple -b options not supported");
+
+        if (sscanf(optarg, "%u", &afl->cpu_to_bind) < 0 || optarg[0] == '-')
+          FATAL("Bad syntax used for -b");
+
+        break;
+
+      }
+
       case 'c': {
 
         afl->shm.cmplog_mode = 1;