Refactor

4 files changed, 46 insertions, 39 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index eecfcf75..f341e300 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -635,6 +635,8 @@ typedef struct afl_state {
   u64 plot_prev_qc, plot_prev_uc, plot_prev_uh, plot_prev_ed;
 
   u64 stats_last_stats_ms, stats_last_plot_ms, stats_last_ms, stats_last_execs;
+  // StatsD
+  u64 statsd_last_send_ms;
   double stats_avg_exec;
 
   u8 *clean_trace;
@@ -957,7 +959,7 @@ void show_init_stats(afl_state_t *);
 /* StatsD */
 
 int statsd_socket_init(char *host, int port);
-int send_statsd_metric(afl_state_t *afl);
+int statsd_send_metric(afl_state_t *afl);
 int statsd_format_metric(afl_state_t *afl, char *buff, size_t bufflen);
 
 /* Run */
diff --git a/include/config.h b/include/config.h
index c55dda2b..8cdf0633 100644
--- a/include/config.h
+++ b/include/config.h
@@ -47,6 +47,8 @@ Server config can be adjusted with AFL_STATSD_HOST and AFL_STATSD_PORT env var.
 
 #define USE_STATSD
 #define STATSD_UPDATE_SEC 1
+#define STATSD_DEFAULT_PORT 8125
+#define STATSD_DEFAULT_HOST "127.0.0.1"
 
 /* If you want to have the original afl internal memory corruption checks.
    Disabled by default for speed. it is better to use "make ASAN_BUILD=1". */
diff --git a/src/afl-fuzz-stats.c b/src/afl-fuzz-stats.c
index bbbe6ba6..ffef7647 100644
--- a/src/afl-fuzz-stats.c
+++ b/src/afl-fuzz-stats.c
@@ -423,8 +423,10 @@ void show_stats(afl_state_t *afl) {
   }
 
   #ifdef USE_STATSD
-  if (cur_ms - afl->stats_last_stats_ms > STATSD_UPDATE_SEC * 1000) {
-    if(send_statsd_metric(afl)){
+  if (cur_ms - afl->statsd_last_send_ms > STATSD_UPDATE_SEC * 1000) {
+    /* reset counter, even if send failed. */
+    afl->statsd_last_send_ms = cur_ms;
+    if(statsd_send_metric(afl)){
       WARNF("coundln't send statsd metric.");
     }
   }
diff --git a/src/afl-fuzz-statsd.c b/src/afl-fuzz-statsd.c
index d9cd12d2..5bd1b537 100644
--- a/src/afl-fuzz-statsd.c
+++ b/src/afl-fuzz-statsd.c
@@ -13,11 +13,12 @@
 #define MAX_TAG_LEN 200
 #define METRIC_PREFIX "fuzzing"
 
-int sock = 0;
 struct sockaddr_in server;
 int error = 0;
+int statds_sock = 0;
 
 int statsd_socket_init(char *host, int port){
+    int sock;
     if((sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == -1){
         perror("socket");
         exit(1);
@@ -42,72 +43,72 @@ int statsd_socket_init(char *host, int port){
     memcpy(&(server.sin_addr), &((struct sockaddr_in*)result->ai_addr)->sin_addr, sizeof(struct in_addr));
     freeaddrinfo(result);
 
-    return 0;
+    return sock;
 }
 
-int send_statsd_metric(afl_state_t *afl){
-    /* default port and host.
+int statsd_send_metric(afl_state_t *afl){
+    
+    char buff[MAX_STATSD_PACKET_SIZE] = {0};
+    /* Default port and host.
     Will be overwritten by AFL_STATSD_PORT and AFL_STATSD_HOST environment variable, if they exists.
-    */ 
-    u16 port = 8125;
-    char* host = "127.0.0.1";
+    */
+    u16 port = STATSD_DEFAULT_PORT;
+    char* host = STATSD_DEFAULT_HOST;
 
     char* port_env;
     char* host_env;
     if ((port_env = getenv("AFL_STATSD_PORT")) != NULL) {
-        // sanitization check ?
         port = atoi(port_env);
     }
     if ((host_env = getenv("AFL_STATSD_HOST")) != NULL) {
-        // sanitization check ?
         host = host_env;
     }
-    
-    error = statsd_socket_init(host, port);
-    if (error){
-        perror("Failed to init statsd client. Aborting");
-        return -1;
-    }
-    
-    if(!sock){
-        perror("sock");
-        return -1;
-    }
 
-    char buff[MAX_STATSD_PACKET_SIZE] = {0};
+    /* statds_sock is a global variable. We set it once in the beginning and reuse the socket.
+    If the sendto later fail, we reset it to 0 to be able to recreate it.
+    */
+    if(!statds_sock){
+        statds_sock = statsd_socket_init(host, port);
+        if(!statds_sock){
+            perror("Cannot create socket");
+            return -1;
+        }
+    }
 
     statsd_format_metric(afl, buff, MAX_STATSD_PACKET_SIZE);
-    if (sendto(sock, buff, strlen(buff), 0,
-               (struct sockaddr *)&server, sizeof(server)) == -1) {
-      perror("sendto");
-      return -1;
+    if (sendto(statds_sock, buff, strlen(buff), 0, (struct sockaddr *)&server, sizeof(server)) == -1) {
+        if(!close(statds_sock)){
+            perror("Cannot close socket");
+        }
+        statds_sock = 0;
+        perror("Cannot sendto");
+        return -1;
     }
 
-    close(sock);
-    sock=0;
-
     return 0;
 }
 
 int statsd_format_metric(afl_state_t *afl, char *buff, size_t bufflen){
-    /*
-    metric format:
-    <some.namespaced.name>:<value>|<type>|<tags>
-    tags format:
+    /* Metric format:
+    <some.namespaced.name>:<value>|<type>
     */
     #ifdef USE_STATSD_TAGS
-    /* #key:value,key:value,key
+    /* Tags format: DogStatsD
+    <some.namespaced.name>:<value>|<type>|#key:value,key:value,key
     */
     char tags[MAX_TAG_LEN * 2] = {0};
     snprintf(tags, MAX_TAG_LEN * 2,
         "|#banner:%s,afl_version:%s",
         afl->use_banner,
         VERSION);
-    #else 
+    #else
+    /* No tags.
+    */
     char *tags = "";
     #endif
-    // Sends multiple metrics with one UDP Packet.
-    // bufflen will limit to the max safe size.
+    /* Sends multiple metrics with one UDP Packet.
+    bufflen will limit to the max safe size.
+    */
     snprintf(buff, bufflen,
         METRIC_PREFIX".cycle_done:%llu|g%s\n"
         METRIC_PREFIX".cycles_wo_finds:%llu|g%s\n"