corrected read_timed for values > 4

2 files changed, 62 insertions, 26 deletions

diff --git a/src/afl-common.c b/src/afl-common.c
index d428c9c5..793041b2 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -39,6 +39,7 @@
 #include <limits.h>
 #include <sys/types.h>
 #include <sys/stat.h>
+#include <sys/ioctl.h>
 #include <unistd.h>
 #include <fcntl.h>
 
@@ -869,51 +870,82 @@ u8 *u_stringify_time_diff(u8 *buf, u64 cur_ms, u64 event_ms) {
 
 }
 
-/* Wrapper for select() and read(), reading len bytes.
-  Assumes that all bytes are available on read!
+/* sets a FD to non-blocking mode (used for read_timed) */
+void set_nonblocking(int fd) {
+
+  int ret = 0;
+  int opt = 1;
+  ret = ioctl(fd, FIONBIO, &opt);
+  if (ret == -1) { PFATAL("Could not enable non-blocking mode on fd %d", fd); }
+
+}
+
+
+/* Wrapper for select() and read(), reading exactly len bytes.
+  Should be called on non-blocking fds.
   Returns the time passed to read.
   If the wait times out, returns timeout_ms + 1;
-  Returns 0 if an error occurred (fd closed, signal, ...); */
+  Returns 0 if an error occurred (fd closed, signal, ...);
+  */
 u32 read_timed(s32 fd, void *buf, size_t len, u32 timeout_ms,
                volatile u8 *stop_soon_p) {
 
+  struct timeval timeout;
   fd_set         readfds;
   FD_ZERO(&readfds);
   FD_SET(fd, &readfds);
-  struct timeval timeout;
 
-  timeout.tv_sec = (timeout_ms / 1000);
-  timeout.tv_usec = (timeout_ms % 1000) * 1000;
-#if !defined(__linux__)
-  u64 read_start = get_cur_time_us();
-#endif
+  size_t read_total = 0;
+  ssize_t len_read = 0;
 
-  /* set exceptfds as well to return when a child exited/closed the pipe. */
-  int sret = select(fd + 1, &readfds, NULL, NULL, &timeout);
+  #if defined(__linux__)
+    timeout.tv_sec = (timeout_ms / 1000);
+    timeout.tv_usec = (timeout_ms % 1000) * 1000;
+  #else
+    u64 time_start = get_cur_time_us();
+  #endif
 
-  if (!sret) {
+  while (read_total < len) {
 
-    return timeout_ms + 1;
+    #if !defined(__linux__)
+      u64 time_current = get_cur_time_us();
+      u64 timeout_current = timeout_ms - (time_current - time_start);
+      timeout.tv_sec = (timeout_current / 1000);
+      timeout.tv_usec = (timeout_current % 1000) * 1000;
+    #endif
 
-  } else if (sret < 0) {
+    /* set exceptfds as well to return when a child exited/closed the pipe. */
+    int sret = select(fd + 1, &readfds, NULL, NULL, &timeout);
 
-    return 0;
+    if (!sret) {
 
-  }
+      // printf("Timeout in sret.");
+      return timeout_ms + 1;
+
+    } else if (sret < 0) {
+
+      /* Retry select for all signals other than than ctrl+c */
+      if (errno == EINTR && !*stop_soon_p) { continue; }
+      return 0;
+
+    }
+
+    len_read = read(fd, ((u8 *)buf) + read_total, len - read_total);
+    if (len_read <= 0) { return 0; }
+    read_total += len_read;
 
-  ssize_t len_read = read(fd, ((u8 *)buf), len);
-  if (len_read < len) { return 0; }
+  }
 
-#if defined(__linux__)
-  u32 exec_ms =
+  #if defined(__linux__)
+    s32 exec_ms =
       MIN(timeout_ms,
           ((u64)timeout_ms - (timeout.tv_sec * 1000 + timeout.tv_usec / 1000)));
-#else
-  u32 exec_ms = get_cur_time_us() - read_start;
-#endif
+  #else
+    u32 exec_ms = get_cur_time_us() - time_start;
+  #endif
 
-  // ensure to report 1 ms has passed (0 is an error)
-  return exec_ms > 0 ? exec_ms : 1;
+  return exec_ms > 0 ? exec_ms
+                     : 1;  // at least 1 milli must have passed (0 is an error)
 
 }
 
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index 137a4f99..01774cd0 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -401,6 +401,8 @@ void afl_fsrv_start(afl_forkserver_t *fsrv, char **argv,
   fsrv->fsrv_ctl_fd = ctl_pipe[1];
   fsrv->fsrv_st_fd = st_pipe[0];
 
+  set_nonblocking(fsrv->fsrv_st_fd);
+
   /* Wait for the fork server to come up, but don't wait too long. */
 
   rlen = 0;
@@ -853,7 +855,9 @@ fsrv_run_result_t afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
 
   fsrv->last_run_timed_out = 0;
 
-  if ((res = read(fsrv->fsrv_st_fd, &fsrv->child_pid, 4)) != 4) {
+  res = read_timed(fsrv->fsrv_st_fd, &fsrv->child_pid, 4, timeout, stop_soon_p);
+
+  if (res < 0 || res > timeout) {
 
     if (*stop_soon_p) { return 0; }
     RPFATAL(res, "Unable to request new process from fork server (OOM?)");