eliminate race condition for cpu affinity on -M/-S

2 files changed, 98 insertions, 23 deletions

diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index a2e849dc..e51b4729 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -36,14 +36,11 @@ void bind_to_free_cpu(afl_state_t *afl) {
   #if defined(__linux__) || defined(__FreeBSD__) || defined(__DragonFly__)
   cpu_set_t c;
   #elif defined(__NetBSD__)
-  cpuset_t *         c;
+  cpuset_t *c;
   #elif defined(__sun)
-  psetid_t            c;
+  psetid_t c;
   #endif
 
-  u8  cpu_used[4096] = {0};
-  u32 i;
-
   if (afl->cpu_core_count < 2) { return; }
 
   if (afl->afl_env.afl_no_affinity) {
@@ -53,13 +50,46 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
   }
 
+  u8  cpu_used[4096] = {0}, lockfile[PATH_MAX] = "";
+  u32 i;
+
+  if (afl->sync_id) {
+
+    s32 lockfd, first = 1;
+
+    snprintf(lockfile, sizeof(lockfile), "%s/.affinity_lock", afl->sync_dir);
+    setenv(CPU_AFFINITY_ENV_VAR, lockfile, 1);
+
+    do {
+
+      if ((lockfd = open(lockfile, O_RDWR | O_CREAT | O_EXCL, 0600)) < 0) {
+
+        if (first) {
+
+          WARNF("CPU affinity lock file present, waiting ...");
+          first = 0;
+
+        }
+
+        usleep(1000);
+
+      }
+
+    } while (lockfd < 0);
+
+    close(lockfd);
+
+  }
+
   #if defined(__linux__)
+
   DIR *          d;
   struct dirent *de;
   d = opendir("/proc");
 
   if (!d) {
 
+    if (lockfile[0]) unlink(lockfile);
     WARNF("Unable to access /proc - can't scan for free CPU cores.");
     return;
 
@@ -67,11 +97,6 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
   ACTF("Checking CPU core loadout...");
 
-  /* Introduce some jitter, in case multiple AFL tasks are doing the same
-     thing at the same time... */
-
-  usleep(R(1000) * 250);
-
   /* Scan all /proc/<pid>/status entries, checking for Cpus_allowed_list.
      Flag all processes bound to a specific CPU using cpu_used[]. This will
      fail for some exotic binding setups, but is likely good enough in almost
@@ -114,20 +139,29 @@ void bind_to_free_cpu(afl_state_t *afl) {
   }
 
   closedir(d);
+
   #elif defined(__FreeBSD__) || defined(__DragonFly__)
+
   struct kinfo_proc *procs;
   size_t             nprocs;
   size_t             proccount;
   int                s_name[] = {CTL_KERN, KERN_PROC, KERN_PROC_ALL};
   size_t             s_name_l = sizeof(s_name) / sizeof(s_name[0]);
 
-  if (sysctl(s_name, s_name_l, NULL, &nprocs, NULL, 0) != 0) return;
+  if (sysctl(s_name, s_name_l, NULL, &nprocs, NULL, 0) != 0) {
+
+    if (lockfile[0]) unlink(lockfile);
+    return;
+
+  }
+
   proccount = nprocs / sizeof(*procs);
   nprocs = nprocs * 4 / 3;
 
   procs = ck_alloc(nprocs);
   if (sysctl(s_name, s_name_l, procs, &nprocs, NULL, 0) != 0) {
 
+    if (lockfile[0]) unlink(lockfile);
     ck_free(procs);
     return;
 
@@ -136,6 +170,7 @@ void bind_to_free_cpu(afl_state_t *afl) {
   for (i = 0; i < proccount; i++) {
 
     #if defined(__FreeBSD__)
+
     if (!strcmp(procs[i].ki_comm, "idle")) continue;
 
     // fix when ki_oncpu = -1
@@ -145,16 +180,21 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
     if (oncpu != -1 && oncpu < sizeof(cpu_used) && procs[i].ki_pctcpu > 60)
       cpu_used[oncpu] = 1;
+
     #elif defined(__DragonFly__)
+
     if (procs[i].kp_lwp.kl_cpuid < sizeof(cpu_used) &&
         procs[i].kp_lwp.kl_pctcpu > 10)
       cpu_used[procs[i].kp_lwp.kl_cpuid] = 1;
+
     #endif
 
   }
 
   ck_free(procs);
+
   #elif defined(__NetBSD__)
+
   struct kinfo_proc2 *procs;
   size_t              nprocs;
   size_t              proccount;
@@ -163,13 +203,20 @@ void bind_to_free_cpu(afl_state_t *afl) {
       CTL_KERN, KERN_PROC2, KERN_PROC_ALL, 0, sizeof(struct kinfo_proc2), 0};
   size_t s_name_l = sizeof(s_name) / sizeof(s_name[0]);
 
-  if (sysctl(s_name, s_name_l, NULL, &nprocs, NULL, 0) != 0) return;
+  if (sysctl(s_name, s_name_l, NULL, &nprocs, NULL, 0) != 0) {
+
+    if (lockfile[0]) unlink(lockfile);
+    return;
+
+  }
+
   proccount = nprocs / sizeof(struct kinfo_proc2);
   procs = ck_alloc(nprocs * sizeof(struct kinfo_proc2));
   s_name[5] = proccount;
 
   if (sysctl(s_name, s_name_l, procs, &nprocs, NULL, 0) != 0) {
 
+    if (lockfile[0]) unlink(lockfile);
     ck_free(procs);
     return;
 
@@ -183,7 +230,9 @@ void bind_to_free_cpu(afl_state_t *afl) {
   }
 
   ck_free(procs);
+
   #elif defined(__sun)
+
   kstat_named_t *n;
   kstat_ctl_t *  m;
   kstat_t *      k;
@@ -198,6 +247,7 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
   if (!k) {
 
+    if (lockfile[0]) unlink(lockfile);
     kstat_close(m);
     return;
 
@@ -205,6 +255,7 @@ void bind_to_free_cpu(afl_state_t *afl) {
 
   if (kstat_read(m, k, NULL)) {
 
+    if (lockfile[0]) unlink(lockfile);
     kstat_close(m);
     return;
 
@@ -220,6 +271,7 @@ void bind_to_free_cpu(afl_state_t *afl) {
     k = kstat_lookup(m, "cpu_stat", i, NULL);
     if (kstat_read(m, k, &cs)) {
 
+      if (lockfile[0]) unlink(lockfile);
       kstat_close(m);
       return;
 
@@ -233,6 +285,7 @@ void bind_to_free_cpu(afl_state_t *afl) {
   }
 
   kstat_close(m);
+
   #else
     #warning \
         "For this platform we do not have free CPU binding code yet. If possible, please supply a PR to https://github.com/AFLplusplus/AFLplusplus"
@@ -241,7 +294,9 @@ void bind_to_free_cpu(afl_state_t *afl) {
   size_t cpu_start = 0;
 
   try:
+
   #if !defined(__ANDROID__)
+
     for (i = cpu_start; i < afl->cpu_core_count; i++) {
 
       if (!cpu_used[i]) { break; }
@@ -251,6 +306,7 @@ void bind_to_free_cpu(afl_state_t *afl) {
   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) {
@@ -274,18 +330,25 @@ void bind_to_free_cpu(afl_state_t *afl) {
   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");
+
   #endif
 
   #if defined(__linux__)
+
   if (sched_setaffinity(0, sizeof(c), &c)) {
 
     if (cpu_start == afl->cpu_core_count) {
@@ -302,6 +365,7 @@ if (pset_assign(c, i, NULL)) PFATAL("pset_assign failed");
   }
 
   #elif defined(__FreeBSD__) || defined(__DragonFly__)
+
   if (pthread_setaffinity_np(pthread_self(), sizeof(c), &c)) {
 
     if (cpu_start == afl->cpu_core_count)
@@ -314,6 +378,7 @@ if (pset_assign(c, i, NULL)) PFATAL("pset_assign failed");
   }
 
   #elif defined(__NetBSD__)
+
 if (pthread_setaffinity_np(pthread_self(), cpuset_size(c), c)) {
 
   if (cpu_start == afl->cpu_core_count)
@@ -326,7 +391,9 @@ if (pthread_setaffinity_np(pthread_self(), cpuset_size(c), c)) {
 }
 
 cpuset_destroy(c);
+
   #elif defined(__sun)
+
 if (pset_bind(c, P_PID, getpid(), NULL)) {
 
   if (cpu_start == afl->cpu_core_count)
@@ -339,11 +406,17 @@ if (pset_bind(c, P_PID, getpid(), NULL)) {
 }
 
 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 */
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index e4e2669c..f7f247f3 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -42,19 +42,21 @@ static void at_exit() {
 
   int   i;
   char *list[4] = {SHM_ENV_VAR, SHM_FUZZ_ENV_VAR, CMPLOG_SHM_ENV_VAR, NULL};
-  char *ptr = getenv("__AFL_TARGET_PID1");
+  char *ptr;
 
+  ptr = getenv(CPU_AFFINITY_ENV_VAR);
+  if (ptr && *ptr) unlink(ptr);
+
+  ptr = getenv("__AFL_TARGET_PID1");
   if (ptr && *ptr && (i = atoi(ptr)) > 0) kill(i, SIGKILL);
 
   ptr = getenv("__AFL_TARGET_PID2");
-
   if (ptr && *ptr && (i = atoi(ptr)) > 0) kill(i, SIGKILL);
 
   i = 0;
   while (list[i] != NULL) {
 
     ptr = getenv(list[i]);
-
     if (ptr && *ptr) {
 
 #ifdef USEMMAP
@@ -1011,17 +1013,19 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
+  check_crash_handling();
+  check_cpu_governor(afl);
+
   get_core_count(afl);
 
+  atexit(at_exit);
+
+  setup_dirs_fds(afl);
+
   #ifdef HAVE_AFFINITY
   bind_to_free_cpu(afl);
   #endif                                                   /* HAVE_AFFINITY */
 
-  check_crash_handling();
-  check_cpu_governor(afl);
-
-  atexit(at_exit);
-
   afl->fsrv.trace_bits =
       afl_shm_init(&afl->shm, afl->fsrv.map_size, afl->non_instrumented_mode);
 
@@ -1038,12 +1042,10 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
-  setup_dirs_fds(afl);
-
   if (afl->is_secondary_node && check_main_node_exists(afl) == 0) {
 
     WARNF("no -M main node found. You need to run one main instance!");
-    sleep(5);
+    sleep(3);
 
   }