From bc202c07eb51ba96eb27aa1ddaf5db6d73448678 Mon Sep 17 00:00:00 2001
From: Quentin Dufour <quentin@deuxfleurs.fr>
Date: Mon, 12 Aug 2019 16:38:19 +0200
Subject: [PATCH] Learn to count!

---
 src/packet.c | 12 ++++++------
 src/proxy.c  |  4 ++++
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/packet.c b/src/packet.c
index e96fe66..1042343 100644
--- a/src/packet.c
+++ b/src/packet.c
@@ -93,14 +93,13 @@ enum FD_STATE write_packet_to_tcp(struct evt_core_fdinfo* fdinfo, struct buffer_
 enum FD_STATE write_packet_to_udp(struct evt_core_fdinfo* fdinfo, struct buffer_packet* bp, struct udp_target* udp_t) {
   ssize_t nwrite;
   union abstract_packet* ap = (union abstract_packet*) (&bp->ip + bp->awrite);
-  union abstract_packet* end = buffer_free_ap(bp);
   if (bp->mode != BP_WRITING) return FDS_ERR;
 
-  while (ap != end) {
+  do {
     if (ap->fmt.headers.cmd != CMD_UDP_ENCAPSULATED) continue;
 
     size_t bytes_to_send;
-    size_t pkt_header_size = sizeof(ap->fmt.headers);
+    size_t pkt_header_size = sizeof(ap->fmt.headers) + sizeof(ap->fmt.content.udp_encapsulated) - sizeof(ap->fmt.content.udp_encapsulated.payload);
     struct sockaddr* addr = NULL;
     socklen_t addrlen = 0;
     if (udp_t->set) {
@@ -110,7 +109,7 @@ enum FD_STATE write_packet_to_udp(struct evt_core_fdinfo* fdinfo, struct buffer_
 
     bytes_to_send = ap->fmt.headers.size - pkt_header_size;
     nwrite = sendto(fdinfo->fd,
-                  &(ap->fmt.content.udp_encapsulated),
+                  &(ap->fmt.content.udp_encapsulated.payload),
                   bytes_to_send,
                   0,
                   addr,
@@ -119,7 +118,8 @@ enum FD_STATE write_packet_to_udp(struct evt_core_fdinfo* fdinfo, struct buffer_
     if (nwrite == -1 && errno == EAGAIN) return FDS_AGAIN;
     if (nwrite != bytes_to_send) return FDS_ERR;
     bp->awrite += nwrite;
-  }
+
+  } while((ap = ap_next(ap)) != NULL);
 
   bp->mode = BP_READING;
   bp->aread = 0;
@@ -137,7 +137,7 @@ enum FD_STATE read_packet_from_udp (struct evt_core_fdinfo* fdinfo, struct buffe
     return FDS_ERR;
   }
 
-  size_t pkt_header_size = sizeof(ap->fmt.headers);
+  size_t pkt_header_size = sizeof(ap->fmt.headers) + sizeof(ap->fmt.content.udp_encapsulated) - sizeof(ap->fmt.content.udp_encapsulated.payload);
   size_t udp_packet_size = sizeof(bp->ip) - pkt_header_size;
   socklen_t addrlen = sizeof(struct sockaddr_in);
 
diff --git a/src/proxy.c b/src/proxy.c
index 39c4c44..1d0e169 100644
--- a/src/proxy.c
+++ b/src/proxy.c
@@ -127,18 +127,22 @@ int main_on_udp_write (struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo)
   int write_res = FDS_READY;
 
   // 1. Get current write buffer OR a buffer from the waiting queue OR leave
+  if (ctx->verbose > 1) fprintf(stderr, "  [proxy] Find write buffer\n");
   if ((bp = get_write_buffer(&app_ctx->br, fdinfo)) == NULL) return 1;
 
   // 2. Write buffer
+  if (ctx->verbose > 1) fprintf(stderr, "  [proxy] Write UDP packet\n");
   write_res = write_packet_to_udp(fdinfo, bp, fdinfo->other);
   if (write_res == FDS_ERR) goto co_error;
   if (write_res == FDS_AGAIN) return 1;
 
   // 3. Notify helpers
+  if (ctx->verbose > 1) fprintf(stderr, "  [proxy] Notify traffic capture\n");
   traffic_capture_notify (&app_ctx->cap, bp, "out");
 
   // 4. A whole packet has been written
   // Release the buffer and notify
+  if (ctx->verbose > 1) fprintf(stderr, "  [proxy] Release buffer and notify\n");
   mv_buffer_wtof(&app_ctx->br, fdinfo);
   notify_read(ctx, &app_ctx->br);