From d4c0af56a46a2998ec0bdc29d2ff45010051028f Mon Sep 17 00:00:00 2001
From: Quentin Dufour <quentin@deuxfleurs.fr>
Date: Thu, 14 Jan 2021 17:43:38 +0100
Subject: [PATCH] Fix some bugs

---
 src/meas_lat.c | 49 +++++++++++++++++++++++++------------------------
 src/measure.c  | 19 +++++++++++++------
 src/measure.h  |  1 +
 3 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/src/meas_lat.c b/src/meas_lat.c
index 250510a..398859f 100644
--- a/src/meas_lat.c
+++ b/src/meas_lat.c
@@ -30,6 +30,8 @@ void free_ms(void* obj) {
   free(obj);
 }
 
+struct measure_state ms_transi = {0};
+
 int streq(char* s1, char* s2) {
   return strcmp(s1, s2) == 0;
 }
@@ -118,7 +120,6 @@ void measlat_stop(
 
 int on_receive_measure_packet(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
   struct measlat_ctx* mctx = fdinfo->cat->app_ctx;
-  struct measure_state msbuf = {0};
   struct measure_state* ms;
   ssize_t nread;
   char url[255];
@@ -130,38 +131,36 @@ int on_receive_measure_packet(struct evt_core_ctx* ctx, struct evt_core_fdinfo*
     // Measurement state exists
     ms = assoc_timer->other;
   } else {
-    // Does not exist yet, we create it.
-    ms = &msbuf;
-    ms->fd = fdinfo->fd;
-    measure_state_init (&mctx->mp, ms);
+    // Does not exist yet, we use a tmp stub.
+    ms = &ms_transi;
   }
 
   // 2. Read data in our measurement object
-  if (ms->mp_nin >= mctx->mp.payload_size || strcmp(mctx->transport, "udp") == 0) ms->mp_nin = 0;
-
-  nread = mctx->connectionless && mctx->role == MEASLAT_SERVER ?
+  ms->mp_nin = mctx->connectionless && mctx->role == MEASLAT_SERVER ?
     recvfrom(fdinfo->fd, ms->mp_in, mctx->mp.payload_size, MSG_TRUNC, (struct sockaddr*)&mctx->addr, &mctx->addrlen) :
-    recv(fdinfo->fd, ms->mp_in + ms->mp_nin, mctx->mp.payload_size - ms->mp_nin, 0);
+    recv(fdinfo->fd, ms->mp_in, mctx->mp.payload_size, 0);
 
-  if (nread > 0) ms->mp_nin += nread;
-
-  if (nread == -1 && errno == EAGAIN || nread == 0) {
-    return 1;
-  }
-  if (!streq("udp", mctx->transport) && ms->mp_nin < mctx->mp.payload_size) {
-    printf("Packet has been fragmented, %ld/%ld\n",ms->mp_nin, mctx->mp.payload_size);
-    return 0;
-  }
+  if (ms->mp_nin <= 0) return 1;
 
   // 3. Process data in our measurement object
   measure_parse (&mctx->mp, ms);
 
   // 4. Persist our measurement object if needed
   // It includes starting a timer.
-  if (ms == &msbuf) {
+  if (ms == &ms_transi) {
+    if (ms->mp_in->counter != 1) {
+      if (!(mctx->connectionless && mctx->role == MEASLAT_SERVER))
+        evt_core_rm_fd (ctx, fdinfo->fd);
+      return 1;
+    }
     struct timespec next_tick = {0};
-    measure_next_tick(&mctx->mp, ms, &next_tick);
-    assoc_timer = register_timer (ctx, mctx, ms, &next_tick);
+    struct measure_state ms_new = {0};
+    measure_state_init(&mctx->mp, &ms_new);
+    ms_new.fd = fdinfo->fd;
+    ms_new.mp_nin = ms->mp_nin;
+    memcpy(ms_new.mp_in, ms->mp_in, mctx->mp.payload_size);
+    measure_next_tick(&mctx->mp, &ms_new, &next_tick);
+    assoc_timer = register_timer (ctx, mctx, &ms_new, &next_tick);
   }
 
   // 5. Check if our measurements are done
@@ -267,8 +266,8 @@ int on_socks5_success_measlat(struct evt_core_ctx* ctx, struct evt_core_fdinfo*
   printf("--- Tor socket registered\n");
 
   struct measure_state ms = {0};
-  ms.fd = fdinfo_n.fd;
   measure_state_init (&mctx->mp, &ms);
+  ms.fd = fdinfo_n.fd;
   register_timer (ctx, mctx, &ms, NULL);
   return 1;
 }
@@ -365,8 +364,8 @@ void spawn_udp_client(struct evt_core_ctx* evts) {
   printf("client port is %u\n", myPort);
 
   struct measure_state ms = {0};
-  ms.fd = udp_sock;
   measure_state_init (&mctx->mp, &ms);
+  ms.fd = udp_sock;
 
   register_timer (evts, mctx, &ms, NULL);
 }
@@ -421,8 +420,8 @@ void spawn_tcp_client(struct evt_core_ctx* evts) {
   printf("client port is %u\n", myPort);
 
   struct measure_state ms = {0};
-  ms.fd = tcp_sock;
   measure_state_init (&mctx->mp, &ms);
+  ms.fd = tcp_sock;
 
   register_timer (evts, mctx, &ms, NULL);
 }
@@ -520,6 +519,8 @@ int main(int argc, char** argv) {
 
   // 2. Check and fix parameters
   measure_params_init (&mctx.mp);
+  measure_state_init (&mctx.mp, &ms_transi);
+
   mctx.addrlen = sizeof(mctx.addr);
   if (mctx.transport == NULL) mctx.transport = "udp";
   if (strcmp(mctx.transport, "udp") == 0) mctx.connectionless = 1;
diff --git a/src/measure.c b/src/measure.c
index 2b1a892..0abe76b 100644
--- a/src/measure.c
+++ b/src/measure.c
@@ -8,15 +8,22 @@ void measure_params_init(struct measure_params* mp) {
 }
 
 void measure_state_init(struct measure_params* mp, struct measure_state* ms) {
-  if ((ms->mp_out = malloc(sizeof(char) * mp->payload_size)) == NULL) {
-    perror("payload malloc failed");
-    exit(EXIT_FAILURE);
+  ms->mp_nin = 0;
+  ms->fd = 0;
+
+  if (ms->mp_out == NULL) {
+    if ((ms->mp_out = malloc(sizeof(char) * mp->payload_size)) == NULL) {
+      perror("payload malloc failed");
+      exit(EXIT_FAILURE);
+    }
   }
   memset(ms->mp_out, 0, mp->payload_size);
 
-  if ((ms->mp_in = malloc(sizeof(char) * mp->payload_size)) == NULL) {
-    perror("payload malloc failed");
-    exit(EXIT_FAILURE);
+  if (ms->mp_in == NULL) {
+    if ((ms->mp_in = malloc(sizeof(char) * mp->payload_size)) == NULL) {
+      perror("payload malloc failed");
+      exit(EXIT_FAILURE);
+    }
   }
   memset(ms->mp_in, 0, mp->payload_size);
 
diff --git a/src/measure.h b/src/measure.h
index c56cfec..3ce1fec 100644
--- a/src/measure.h
+++ b/src/measure.h
@@ -20,6 +20,7 @@ struct measure_state {
   int fd;
 };
 
+#pragma pack(1)
 struct measure_packet {
   uint64_t counter;
   uint8_t flag;