Bind cpu (#480)

* silence compiletime warning

* refactored cpu binding

* formatted code

1 files changed, 110 insertions, 113 deletions

diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index ad92dff6..1e4f8dee 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -28,10 +28,9 @@
 
 #ifdef HAVE_AFFINITY
 
-/* Build a list of processes bound to specific cores. Returns -1 if nothing
-   can be found. Assumes an upper bound of 4k CPUs. */
+/* bind process to a specific cpu. Returns 0 on failure. */
 
-void bind_to_free_cpu(afl_state_t *afl) {
+static u8 bind_cpu(afl_state_t *afl, s32 cpuid) {
 
   #if defined(__linux__) || defined(__FreeBSD__) || defined(__DragonFly__)
   cpu_set_t c;
@@ -41,25 +40,108 @@ void bind_to_free_cpu(afl_state_t *afl) {
   psetid_t c;
   #endif
 
-  if (afl->cpu_core_count < 2) { return; }
+  afl->cpu_aff = cpuid;
+
+  #if defined(__linux__) || defined(__FreeBSD__) || defined(__DragonFly__)
+
+  CPU_ZERO(&c);
+  CPU_SET(cpuid, &c);
+
+  #elif defined(__NetBSD__)
+
+  c = cpuset_create();
+  if (c == NULL) { PFATAL("cpuset_create failed"); }
+  cpuset_set(cpuid, c);
+
+  #elif defined(__sun)
+
+  pset_create(&c);
+  if (pset_assign(c, cpuid, NULL)) { PFATAL("pset_assign failed"); }
+
+  #endif
+
+  #if defined(__linux__)
+
+  return (sched_setaffinity(0, sizeof(c), &c) == 0);
+
+  #elif defined(__FreeBSD__) || defined(__DragonFly__)
+
+  return (pthread_setaffinity_np(pthread_self(), sizeof(c), &c) == 0);
+
+  #elif defined(__NetBSD__)
+
+  if (pthread_setaffinity_np(pthread_self(), cpuset_size(c), c)) {
+
+    cpuset_destroy(c);
+    return 0;
+
+  }
+
+  cpuset_destroy(c);
+  return 1;
+
+  #elif defined(__sun)
+
+  if (pset_bind(c, P_PID, getpid(), NULL)) {
+
+    pset_destroy(c);
+    return 0;
+
+  }
+
+  pset_destroy(c);
+  return 1;
+
+  #else
+
+  // this will need something for other platforms
+  // TODO: Solaris/Illumos has processor_bind ... might worth a try
+  WARNF("Cannot bind to CPU yet on this platform.");
+  return 1;
+
+  #endif
+
+}
+
+/* Build a list of processes bound to specific cores. Returns -1 if nothing
+   can be found. Assumes an upper bound of 4k CPUs. */
+
+void bind_to_free_cpu(afl_state_t *afl) {
+
+  u8  cpu_used[4096] = {0};
+  u8  lockfile[PATH_MAX] = "";
+  u32 i;
 
   if (afl->afl_env.afl_no_affinity) {
 
+    if (afl->cpu_to_bind != -1) {
+
+      FATAL("-b and AFL_NO_AFFINITY are mututally exclusive.");
+
+    }
+
     WARNF("Not binding to a CPU core (AFL_NO_AFFINITY set).");
     return;
 
   }
 
-  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 (!bind_cpu(afl, afl->cpu_to_bind)) {
+
+      FATAL(
+          "Could not bind to requested CPU %d! Make sure you passed a valid "
+          "-b.",
+          afl->cpu_to_bind);
+
+    }
+
+    return;
 
   }
 
+  if (afl->cpu_core_count < 2) { return; }
+
   if (afl->sync_id) {
 
     s32 lockfd, first = 1;
@@ -300,135 +382,50 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
   size_t cpu_start = 0;
 
-  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++) {
 
-    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;
-  if (i == -1) {
-
-  #endif
-
-    SAYF("\n" cLRD "[-] " cRST
-         "Uh-oh, looks like all %d CPU cores on your system are allocated to\n"
-         "    other instances of afl-fuzz (or similar CPU-locked tasks). "
-         "Starting\n"
-         "    another fuzzer on this machine is probably a bad plan, but if "
-         "you are\n"
-         "    absolutely sure, you can set AFL_NO_AFFINITY and try again.\n",
-         afl->cpu_core_count);
-    FATAL("No more free CPU cores");
-
-  }
-
-  OKF("Found a free CPU core, try binding to #%u.", i);
-
-set_cpu:
+  /* for some reason Android goes backwards */
 
-  afl->cpu_aff = i;
-
-  #if defined(__linux__) || defined(__FreeBSD__) || defined(__DragonFly__)
-
-  CPU_ZERO(&c);
-  CPU_SET(i, &c);
-
-  #elif defined(__NetBSD__)
-
-  c = cpuset_create();
-  if (c == NULL) PFATAL("cpuset_create failed");
-  cpuset_set(i, c);
-
-  #elif defined(__sun)
-
-pset_create(&c);
-if (pset_assign(c, i, NULL)) PFATAL("pset_assign failed");
+  for (i = afl->cpu_core_count - 1; i > -1; i--) {
 
   #endif
 
-  #if defined(__linux__)
+    if (!cpu_used[i]) { continue; }
 
-  if (sched_setaffinity(0, sizeof(c), &c)) {
+    OKF("Found a free CPU core, try binding to #%u.", i);
 
-    if (cpu_start == afl->cpu_core_count) {
+    if (bind_cpu(afl, i)) {
 
-      PFATAL("sched_setaffinity failed for CPU %d, exit", i);
+      /* Success :) */
+      break;
 
     }
 
-    WARNF("sched_setaffinity failed to CPU %d, trying next CPU", i);
+    WARNF("setaffinity failed to CPU %d, trying next CPU", i);
     cpu_start++;
-    goto try
-      ;
 
   }
 
-  #elif defined(__FreeBSD__) || defined(__DragonFly__)
+  if (lockfile[0]) unlink(lockfile);
 
-  if (pthread_setaffinity_np(pthread_self(), sizeof(c), &c)) {
+  if (i == afl->cpu_core_count || i == -1) {
 
-    if (cpu_start == afl->cpu_core_count)
-      PFATAL("pthread_setaffinity failed for cpu %d, exit", i);
-    WARNF("pthread_setaffinity failed to CPU %d, trying next CPU", i);
-    cpu_start++;
-    goto try
-      ;
+    SAYF("\n" cLRD "[-] " cRST
+         "Uh-oh, looks like all %d CPU cores on your system are allocated to\n"
+         "    other instances of afl-fuzz (or similar CPU-locked tasks). "
+         "Starting\n"
+         "    another fuzzer on this machine is probably a bad plan, but if "
+         "you are\n"
+         "    absolutely sure, you can set AFL_NO_AFFINITY and try again.\n",
+         afl->cpu_core_count);
+    FATAL("No more free CPU cores");
 
   }
 
-  #elif defined(__NetBSD__)
-
-if (pthread_setaffinity_np(pthread_self(), cpuset_size(c), c)) {
-
-  if (cpu_start == afl->cpu_core_count)
-    PFATAL("pthread_setaffinity failed for cpu %d, exit", i);
-  WARNF("pthread_setaffinity failed to CPU %d, trying next CPU", i);
-  cpu_start++;
-  goto try
-    ;
-
-}
-
-cpuset_destroy(c);
-
-  #elif defined(__sun)
-
-if (pset_bind(c, P_PID, getpid(), NULL)) {
-
-  if (cpu_start == afl->cpu_core_count)
-    PFATAL("pset_bind failed for cpu %d, exit", i);
-  WARNF("pset_bind failed to CPU %d, trying next CPU", i);
-  cpu_start++;
-  goto try
-    ;
-
-}
-
-pset_destroy(c);
-
-  #else
-
-  // this will need something for other platforms
-  // TODO: Solaris/Illumos has processor_bind ... might worth a try
-
-  #endif
-
-  if (lockfile[0]) unlink(lockfile);
-  // we leave the environment variable to ensure a cleanup for other processes
-
 }
 
 #endif                                                     /* HAVE_AFFINITY */