doubled the speed of afl_network_proxy

4 files changed, 246 insertions, 56 deletions

diff --git a/examples/afl_network_proxy/Makefile b/examples/afl_network_proxy/Makefile
index eeee1178..0b306dde 100644
--- a/examples/afl_network_proxy/Makefile
+++ b/examples/afl_network_proxy/Makefile
@@ -1,22 +1,2 @@
-PREFIX   ?= /usr/local
-BIN_PATH  = $(PREFIX)/bin
-DOC_PATH  = $(PREFIX)/share/doc/afl
-
-PROGRAMS = afl-network-client afl-network-server
-
-all:	$(PROGRAMS)
-
-afl-network-client:	afl-network-client.c
-	$(CC) -I../../include -o afl-network-client afl-network-client.c
-
-afl-network-server:	afl-network-server.c
-	$(CC) -I../../include -o afl-network-server afl-network-server.c ../../src/afl-forkserver.c ../../src/afl-sharedmem.c ../../src/afl-common.c -DBIN_PATH=\"$(BIN_PATH)\"
-
-clean:
-	rm -f $(PROGRAMS) *~ core
-
-install: all
-	install -d -m 755 $${DESTDIR}$(BIN_PATH) $${DESTDIR}$(DOC_PATH)
-	install -m 755 $(PROGRAMS) $${DESTDIR}$(BIN_PATH)
-	install -m 644 README.md $${DESTDIR}$(DOC_PATH)/README.network_proxy.md
-	
\ No newline at end of file
+all:
+	@echo please use GNU make, thanks!
diff --git a/examples/afl_network_proxy/README.md b/examples/afl_network_proxy/README.md
index c33096be..84ebfa48 100644
--- a/examples/afl_network_proxy/README.md
+++ b/examples/afl_network_proxy/README.md
@@ -12,6 +12,14 @@ Note that the impact on fuzzing speed will be huge, expect a loss of 90%.
 
 ## how to get it running
 
+### Compiling
+
+Just type `make` and let the autodetection do everything for you.
+
+Note that compression is supported but currently disabled. It seems that
+sending 64kb of map data over TCP is faster than compressing it with the
+fastest algorithm and options to 112 byte and sending this. Weird.
+
 ### on the target
 
 Run `afl-network-server` with your target with the -m and -t values you need.
@@ -40,16 +48,11 @@ The TARGET can be an IPv4 or IPv6 address, or a host name that resolves to
 either. Note that also the outgoing interface can be specified with a '%' for
 `afl-network-client`, e.g. `fe80::1234%eth0`.
 
+Also make sure your middle value of `/proc/sys/net/ipv4/tcp_rmem` is larger
+than your MAP_SIZE (130kb is a good value). This is the default TCP window
+size value.
+
 ## how to compile and install
 
 `make && sudo make install`
 
-## Future
-
-It would be much faster and more effective if `afl-network-server` does not
-send the map data back (64kb or more) but the checksum that `afl-fuzz` would
-generate. This change however would make it incompatible with existing
-afl spinoffs.
-
-But in the future this will be implemented and supported as a compile option.
-
diff --git a/examples/afl_network_proxy/afl-network-client.c b/examples/afl_network_proxy/afl-network-client.c
index b9cd88f0..94f6bb42 100644
--- a/examples/afl_network_proxy/afl-network-client.c
+++ b/examples/afl_network_proxy/afl-network-client.c
@@ -41,6 +41,10 @@
 #include <netdb.h>
 #include <fcntl.h>
 
+#ifdef USE_DEFLATE
+#include <libdeflate.h>
+#endif
+
 u8 *__afl_area_ptr;
 
 #ifdef __ANDROID__
@@ -206,7 +210,12 @@ int main(int argc, char *argv[]) {
   u8 *            interface, *buf, *ptr;
   s32             s = -1;
   struct addrinfo hints, *hres, *aip;
-  u32             len, max_len = 65536;
+  u32 *           lenptr, max_len = 65536;
+#ifdef USE_DEFLATE
+  u8 *   buf2;
+  u32 *  lenptr1, *lenptr2, buf2_len, compress_len;
+  size_t decompress_len;
+#endif
 
   if (argc < 3 || argc > 4) {
 
@@ -235,8 +244,17 @@ int main(int argc, char *argv[]) {
     if ((__afl_map_size = atoi(ptr)) < 8)
       FATAL("illegal map size, may not be < 8 or >= 2^30: %s", ptr);
 
-  if ((buf = malloc(max_len)) == NULL)
-    PFATAL("can not allocate %u memory", max_len);
+  if ((buf = malloc(max_len + 4)) == NULL)
+    PFATAL("can not allocate %u memory", max_len + 4);
+  lenptr = (u32 *)buf;
+
+#ifdef USE_DEFLATE
+  buf2_len = (max_len > __afl_map_size ? max_len : __afl_map_size);
+  if ((buf2 = malloc(buf2_len + 8)) == NULL)
+    PFATAL("can not allocate %u memory", buf2_len + 8);
+  lenptr1 = (u32 *)buf2;
+  lenptr2 = (u32 *)(buf2 + 4);
+#endif
 
   memset(&hints, 0, sizeof(hints));
   hints.ai_socktype = SOCK_STREAM;
@@ -258,28 +276,81 @@ int main(int argc, char *argv[]) {
       fprintf(stderr,
               "Warning: binding to interface is not supported for your OS\n");
 #endif
+
+#ifdef SO_PRIORITY
+      int priority = 7;
+      if (setsockopt(s, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)) <
+          0) {
+
+        priority = 6;
+        if (setsockopt(s, SOL_SOCKET, SO_PRIORITY, &priority,
+                       sizeof(priority)) < 0)
+          WARNF("could not set priority on socket");
+
+      }
+
+#endif
+
       if (connect(s, aip->ai_addr, aip->ai_addrlen) == -1) s = -1;
 
     }
 
   }
 
+#ifdef USE_DEFLATE
+  struct libdeflate_compressor *compressor;
+  compressor = libdeflate_alloc_compressor(1);
+  struct libdeflate_decompressor *decompressor;
+  decompressor = libdeflate_alloc_decompressor();
+  fprintf(stderr, "Compile with compression support\n");
+#endif
+
   if (s == -1)
     FATAL("could not connect to target tcp://%s:%s", argv[1], argv[2]);
+  else
+    fprintf(stderr, "Connected to target tcp://%s:%s\n", argv[1], argv[2]);
 
   /* we initialize the shared memory map and start the forkserver */
   __afl_map_shm();
   __afl_start_forkserver();
 
-  int i = 1, j, status, ret;
+  int i = 1, j, status, ret, received;
+
   // fprintf(stderr, "Waiting for first testcase\n");
-  while ((len = __afl_next_testcase(buf, max_len)) > 0) {
+  while ((*lenptr = __afl_next_testcase(buf + 4, max_len)) > 0) {
+
+    // fprintf(stderr, "Sending testcase with len %u\n", *lenptr);
+#ifdef USE_DEFLATE
+    // we only compress the testcase if it does not fit in the TCP packet
+    if (*lenptr > 1500 - 20 - 32 - 4) {
+
+      // set highest byte to signify compression
+      *lenptr1 = (*lenptr | 0xff000000);
+      *lenptr2 = (u32)libdeflate_deflate_compress(compressor, buf + 4, *lenptr,
+                                                  buf2 + 8, buf2_len);
+      if (send(s, buf2, *lenptr2 + 8, 0) != *lenptr2 + 8)
+        PFATAL("sending test data failed");
+      fprintf(stderr, "COMPRESS (%u->%u):\n", *lenptr, *lenptr2);
+      for (u32 i = 0; i < *lenptr; i++)
+        fprintf(stderr, "%02x", buf[i + 4]);
+      fprintf(stderr, "\n");
+      for (u32 i = 0; i < *lenptr2; i++)
+        fprintf(stderr, "%02x", buf2[i + 8]);
+      fprintf(stderr, "\n");
+
+    } else {
 
-    // fprintf(stderr, "Sending testcase with len %u\n", len);
-    if (send(s, &len, 4, 0) != 4) PFATAL("sending size data %d failed", len);
-    if (send(s, buf, len, 0) != len) PFATAL("sending test data failed");
+#endif
+      if (send(s, buf, *lenptr + 4, 0) != *lenptr + 4)
+        PFATAL("sending test data failed");
+#ifdef USE_DEFLATE
+      // fprintf(stderr, "unCOMPRESS (%u)\n", *lenptr);
 
-    int received = 0;
+    }
+
+#endif
+
+    received = 0;
     while (received < 4 &&
            (ret = recv(s, &status + received, 4 - received, 0)) > 0)
       received += ret;
@@ -288,12 +359,37 @@ int main(int argc, char *argv[]) {
     // fprintf(stderr, "Received status\n");
 
     received = 0;
+#ifdef USE_DEFLATE
+    while (received < 4 &&
+           (ret = recv(s, &compress_len + received, 4 - received, 0)) > 0)
+      received += ret;
+    if (received != 4)
+      FATAL("did not receive compress_len (%d, %d)", received, ret);
+    // fprintf(stderr, "Received status\n");
+
+    received = 0;
+    while (received < compress_len &&
+           (ret = recv(s, buf2 + received, buf2_len - received, 0)) > 0)
+      received += ret;
+    if (received != compress_len)
+      FATAL("did not receive coverage data (%d, %d)", received, ret);
+
+    if (libdeflate_deflate_decompress(decompressor, buf2, compress_len,
+                                      __afl_area_ptr, __afl_map_size,
+                                      &decompress_len) != LIBDEFLATE_SUCCESS ||
+        decompress_len != __afl_map_size)
+      FATAL("decompression failed");
+// fprintf(stderr, "DECOMPRESS (%u->%u): ", compress_len, decompress_len);
+// for (u32 i = 0; i < __afl_map_size; i++) fprintf(stderr, "%02x",
+// __afl_area_ptr[i]); fprintf(stderr, "\n");
+#else
     while (received < __afl_map_size &&
            (ret = recv(s, __afl_area_ptr + received, __afl_map_size - received,
                        0)) > 0)
       received += ret;
     if (received != __afl_map_size)
       FATAL("did not receive coverage data (%d, %d)", received, ret);
+#endif
     // fprintf(stderr, "Received coverage\n");
 
     /* report the test case is done and wait for the next */
@@ -302,6 +398,11 @@ int main(int argc, char *argv[]) {
 
   }
 
+#ifdef USE_DEFLATE
+  libdeflate_free_compressor(compressor);
+  libdeflate_free_decompressor(decompressor);
+#endif
+
   return 0;
 
 }
diff --git a/examples/afl_network_proxy/afl-network-server.c b/examples/afl_network_proxy/afl-network-server.c
index e069af3d..e4c3bc6d 100644
--- a/examples/afl_network_proxy/afl-network-server.c
+++ b/examples/afl_network_proxy/afl-network-server.c
@@ -61,13 +61,21 @@
 #include <sys/socket.h>
 #include <netdb.h>
 
+#ifdef USE_DEFLATE
+#include <libdeflate.h>
+struct libdeflate_compressor *  compressor;
+struct libdeflate_decompressor *decompressor;
+#endif
+
 static u8 *in_file,                    /* Minimizer input test case         */
     *out_file;
 
 static u8 *in_data;                    /* Input data for trimming           */
+static u8 *buf2;
 
-static s32 in_len;
-static u32 map_size = MAP_SIZE;
+static s32    in_len;
+static u32    map_size = MAP_SIZE;
+static size_t buf2_len;
 
 static volatile u8 stop_soon;          /* Ctrl-C pressed?                   */
 
@@ -335,25 +343,64 @@ static void usage(u8 *argv0) {
 
 int recv_testcase(int s, void **buf, size_t *max_len) {
 
-  int size, received = 0, ret;
+  u32    size;
+  s32    ret;
+  size_t received;
 
+  received = 0;
   while (received < 4 && (ret = recv(s, &size + received, 4 - received, 0)) > 0)
     received += ret;
-
   if (received != 4) FATAL("did not receive size information");
-  if (size < 1) FATAL("did not receive valid size information");
+  if (size == 0) FATAL("did not receive valid size information");
   // fprintf(stderr, "received size information of %d\n", size);
 
-  *buf = maybe_grow(buf, max_len, size);
-  // fprintf(stderr, "receiving testcase %p %p max %u\n", buf, *buf, *max_len);
-  received = 0;
-  while (received < size &&
-         (ret = recv(s, ((char *)*buf) + received, size - received, 0)) > 0)
-    received += ret;
+  if ((size && 0xff000000) != 0xff000000) {
+
+    *buf = maybe_grow(buf, max_len, size);
+    received = 0;
+    // fprintf(stderr, "unCOMPRESS (%u)\n", size);
+    while (received < size &&
+           (ret = recv(s, ((char *)*buf) + received, size - received, 0)) > 0)
+      received += ret;
+
+  } else {
+
+#ifdef USE_DEFLATE
+    u32 clen;
+    size = (size & 0x00ffffff);
+    *buf = maybe_grow(buf, max_len, size);
+    received = 0;
+    while (received < 4 &&
+           (ret = recv(s, &clen + received, 4 - received, 0)) > 0)
+      received += ret;
+    if (received != 4) FATAL("did not receive size information");
+    // fprintf(stderr, "received clen information of %d\n", clen);
+    if (clen < 1)
+      FATAL("did not receive valid compressed len information: %u", clen);
+    buf2 = maybe_grow((void **)&buf2, &buf2_len, clen);
+    received = 0;
+    while (received < clen &&
+           (ret = recv(s, buf2 + received, clen - received, 0)) > 0)
+      received += ret;
+    if (received != clen) FATAL("did not receive compressed information");
+    if (libdeflate_deflate_decompress(decompressor, buf2, clen, (char *)*buf,
+                                      *max_len,
+                                      &received) != LIBDEFLATE_SUCCESS)
+      FATAL("decompression failed");
+      // fprintf(stderr, "DECOMPRESS (%u->%u):\n", clen, received);
+      // for (u32 i = 0; i < clen; i++) fprintf(stderr, "%02x", buf2[i]);
+      // fprintf(stderr, "\n");
+      // for (u32 i = 0; i < received; i++) fprintf(stderr, "%02x",
+      // ((u8*)(*buf))[i]); fprintf(stderr, "\n");
+#else
+    FATAL("Received compressed data but not compiled with compression support");
+#endif
+
+  }
 
+  // fprintf(stderr, "receiving testcase %p %p max %u\n", buf, *buf, *max_len);
   if (received != size)
     FATAL("did not receive testcase data %u != %u, %d", received, size, ret);
-
   // fprintf(stderr, "received testcase\n");
   return size;
 
@@ -371,6 +418,10 @@ int main(int argc, char **argv_orig, char **envp) {
   int                 addrlen = sizeof(clientaddr);
   char                str[INET6_ADDRSTRLEN];
   char **             argv = argv_cpy_dup(argc, argv_orig);
+  u8 *                send_buf;
+#ifdef USE_DEFLATE
+  u32 *lenptr;
+#endif
 
   afl_forkserver_t  fsrv_var = {0};
   afl_forkserver_t *fsrv = &fsrv_var;
@@ -378,6 +429,8 @@ int main(int argc, char **argv_orig, char **envp) {
   map_size = get_map_size();
   fsrv->map_size = map_size;
 
+  if ((send_buf = malloc(map_size + 4)) == NULL) PFATAL("malloc");
+
   while ((opt = getopt(argc, argv, "+i:f:m:t:QUWh")) > 0) {
 
     switch (opt) {
@@ -553,6 +606,21 @@ int main(int argc, char **argv_orig, char **envp) {
   }
 
 #endif
+
+#ifdef SO_PRIORITY
+  int priority = 7;
+  if (setsockopt(sock, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)) <
+      0) {
+
+    priority = 6;
+    if (setsockopt(sock, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)) <
+        0)
+      WARNF("could not set priority on socket");
+
+  }
+
+#endif
+
   memset(&serveraddr, 0, sizeof(serveraddr));
   serveraddr.sin6_family = AF_INET6;
   serveraddr.sin6_port = htons(port);
@@ -566,6 +634,14 @@ int main(int argc, char **argv_orig, char **envp) {
   afl_fsrv_start(fsrv, use_argv, &stop_soon,
                  get_afl_env("AFL_DEBUG_CHILD_OUTPUT") ? 1 : 0);
 
+#ifdef USE_DEFLATE
+  compressor = libdeflate_alloc_compressor(1);
+  decompressor = libdeflate_alloc_decompressor();
+  buf2 = maybe_grow((void **)&buf2, &buf2_len, map_size + 16);
+  lenptr = (u32 *)(buf2 + 4);
+  fprintf(stderr, "Compiled with compression support\n");
+#endif
+
   fprintf(stderr,
           "Waiting for incoming connection from afl-network-client on port %d "
           "...\n",
@@ -574,15 +650,40 @@ int main(int argc, char **argv_orig, char **envp) {
   if ((s = accept(sock, NULL, NULL)) < 0) { PFATAL("accept() failed"); }
   fprintf(stderr, "Received connection, starting ...\n");
 
+#ifdef SO_PRIORITY
+  priority = 7;
+  if (setsockopt(s, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)) < 0) {
+
+    priority = 6;
+    if (setsockopt(s, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)) < 0)
+      WARNF("could not set priority on socket");
+
+  }
+
+#endif
+
   while ((in_len = recv_testcase(s, (void **)&in_data, &max_len)) > 0) {
 
     // fprintf(stderr, "received %u\n", in_len);
     run_target(fsrv, use_argv, in_data, in_len, 1);
 
-    if (send(s, &fsrv->child_status, 4, 0) != 4)
-      FATAL("could not send waitpid data");
-    if (send(s, fsrv->trace_bits, fsrv->map_size, 0) != fsrv->map_size)
-      FATAL("could not send coverage data");
+    memcpy(send_buf + 4, fsrv->trace_bits, fsrv->map_size);
+
+#ifdef USE_DEFLATE
+    memcpy(buf2, &fsrv->child_status, 4);
+    *lenptr = (u32)libdeflate_deflate_compress(
+        compressor, send_buf + 4, fsrv->map_size, buf2 + 8, buf2_len - 8);
+    // fprintf(stderr, "COMPRESS (%u->%u): ", fsrv->map_size, *lenptr);
+    // for (u32 i = 0; i < fsrv->map_size; i++) fprintf(stderr, "%02x",
+    // fsrv->trace_bits[i]); fprintf(stderr, "\n");
+    if (send(s, buf2, *lenptr + 8, 0) != 8 + *lenptr)
+      FATAL("could not send data");
+#else
+    memcpy(send_buf, &fsrv->child_status, 4);
+    if (send(s, send_buf, fsrv->map_size + 4, 0) != 4 + fsrv->map_size)
+      FATAL("could not send data");
+#endif
+
     // fprintf(stderr, "sent result\n");
 
   }
@@ -595,6 +696,11 @@ int main(int argc, char **argv_orig, char **envp) {
   afl_fsrv_deinit(fsrv);
   if (fsrv->target_path) { ck_free(fsrv->target_path); }
   if (in_data) { ck_free(in_data); }
+#if USE_DEFLATE
+  if (buf2) { ck_free(buf2); }
+  libdeflate_free_compressor(compressor);
+  libdeflate_free_decompressor(decompressor);
+#endif
 
   argv_cpy_free(argv);