Provide a generic reassembler interface.add/generic-reassembler-interface

 * Shall be used for stream based protocols e.g. Kerberos, QUIC, etc.

Signed-off-by: lns <matzeton@googlemail.com>
Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

7 files changed, 359 insertions, 174 deletions

diff --git a/src/include/ndpi_api.h.in b/src/include/ndpi_api.h.in
index 9064ac627..a9c39313c 100644
--- a/src/include/ndpi_api.h.in
+++ b/src/include/ndpi_api.h.in
@@ -1823,6 +1823,21 @@ extern "C" {
 				char *out, u_int out_len,
 				u_int8_t use_json);
 
+  /* ******************************* */
+
+  void ndpi_reassemble_set_buffer_len(struct ndpi_reasm * const reasm,
+                                      u_int64_t buffer_len);
+  int ndpi_reassemble(struct ndpi_reasm * const reasm, uint8_t const * const frag,
+                      uint64_t frag_len, uint64_t frag_offset);
+  int ndpi_reassemble_payload(struct ndpi_reasm * const reasm,
+                              struct ndpi_packet_struct * packet);
+  void ndpi_reassemble_swap_payload(struct ndpi_packet_struct * packet,
+                                    struct ndpi_reasm const * reasm,
+                                    u_int8_t const ** const original_payload,
+                                    u_int16_t * original_payload_packet_len);
+  int ndpi_reassemble_in_progress(struct ndpi_reasm * const reasm);
+  int ndpi_reassemble_is_complete(struct ndpi_reasm * const reasm);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/include/ndpi_typedefs.h b/src/include/ndpi_typedefs.h
index efea443d6..ee4eabbc2 100644
--- a/src/include/ndpi_typedefs.h
+++ b/src/include/ndpi_typedefs.h
@@ -647,6 +647,14 @@ struct ndpi_lru_cache {
   struct ndpi_lru_cache_entry *entries;
 };
 
+struct ndpi_reasm {
+  u_int8_t *buf;
+  u_int8_t *buf_bitmap;
+  u_int32_t buf_last_pos;
+  u_int32_t buf_req;
+  u_int32_t buf_len;
+};
+
 /* ************************************************** */
 
 struct ndpi_flow_tcp_struct {
@@ -808,11 +816,6 @@ struct ndpi_flow_udp_struct {
   u_int8_t wireguard_stage;
   u_int32_t wireguard_peer_index[2];
 
-  /* NDPI_PROTOCOL_QUIC */
-  u_int8_t *quic_reasm_buf;
-  u_int8_t *quic_reasm_buf_bitmap;
-  u_int32_t quic_reasm_buf_last_pos;
-
   /* NDPI_PROTOCOL_CSGO */
   u_int8_t csgo_strid[18],csgo_state,csgo_s2;
   u_int32_t csgo_id2;
@@ -1293,6 +1296,8 @@ struct ndpi_flow_struct {
     struct ndpi_flow_udp_struct udp;
   } l4;
 
+  struct ndpi_reasm reassemble;
+
   /* Some protocols calculate the entropy. */
   float entropy;
 
@@ -1330,16 +1335,6 @@ struct ndpi_flow_struct {
     char *nat_ip; /* Via HTTP X-Forwarded-For */
   } http;
 
-  /*
-     Put outside of the union to avoid issues in case the protocol
-     is remapped to somethign pther than Kerberos due to a faulty
-     dissector
-  */
-  struct {
-    char *pktbuf;
-    u_int16_t pktbuf_maxlen, pktbuf_currlen;
-  } kerberos_buf;
-
   struct {
     u_int8_t num_udp_pkts, num_binding_requests;
     u_int16_t num_processed_pkts;
diff --git a/src/lib/ndpi_main.c b/src/lib/ndpi_main.c
index 7929b3b88..8521d5dda 100644
--- a/src/lib/ndpi_main.c
+++ b/src/lib/ndpi_main.c
@@ -4895,9 +4895,6 @@ void ndpi_free_flow_data(struct ndpi_flow_struct* flow) {
     if(flow->http.detected_os)
       ndpi_free(flow->http.detected_os);
 
-    if(flow->kerberos_buf.pktbuf)
-      ndpi_free(flow->kerberos_buf.pktbuf);
-
     if(flow_is_proto(flow, NDPI_PROTOCOL_QUIC) ||
        flow_is_proto(flow, NDPI_PROTOCOL_TLS) ||
        flow_is_proto(flow, NDPI_PROTOCOL_DTLS) ||
@@ -4930,12 +4927,13 @@ void ndpi_free_flow_data(struct ndpi_flow_struct* flow) {
 	ndpi_free(flow->l4.tcp.tls.message[1].buffer);
     }
 
-    if(flow->l4_proto == IPPROTO_UDP) {
-      if(flow->l4.udp.quic_reasm_buf){
-        ndpi_free(flow->l4.udp.quic_reasm_buf);
-        if(flow->l4.udp.quic_reasm_buf_bitmap)
-          ndpi_free(flow->l4.udp.quic_reasm_buf_bitmap);
-      }
+    if (flow->reassemble.buf != NULL)
+    {
+      ndpi_free(flow->reassemble.buf);
+    }
+    if (flow->reassemble.buf_bitmap != NULL)
+    {
+      ndpi_free(flow->reassemble.buf_bitmap);
     }
   }
 }
diff --git a/src/lib/ndpi_reassmble.c b/src/lib/ndpi_reassmble.c
new file mode 100644
index 000000000..50141dacc
--- /dev/null
+++ b/src/lib/ndpi_reassmble.c
@@ -0,0 +1,197 @@
+/*
+ * ndpi_reassmble.c
+ *
+ * Copyright (C) 2022 - ntop.org
+ *
+ * This file is part of nDPI, an open source deep packet inspection
+ * library.
+ *
+ * nDPI is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * nDPI is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with nDPI.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+#include "ndpi_config.h"
+#include "ndpi_api.h"
+
+/* ********************************************************************************* */
+
+static void update_reasm_buf_bitmap(u_int8_t *buffer_bitmap,
+                                    const u_int32_t buffer_bitmap_size,
+                                    const u_int32_t recv_pos,
+                                    const u_int32_t recv_len)
+{
+  if (!recv_len || !buffer_bitmap_size || recv_pos + recv_len > buffer_bitmap_size * 8)
+    return;
+  const u_int32_t start_byte = recv_pos / 8;
+  const u_int32_t end_byte = (recv_pos + recv_len - 1) / 8;
+  const u_int32_t start_bit = recv_pos % 8;
+  const u_int32_t end_bit = (start_bit + recv_len - 1) % 8;
+  if (start_byte == end_byte)
+  {
+    // fill from bit 'start_bit' until bit 'end_bit', both inclusive
+    buffer_bitmap[start_byte] |= (((1U << recv_len) - 1U) << start_bit);
+  } else {
+    u_int32_t i;
+    for (i = start_byte + 1; i <= end_byte - 1; i++)
+    {
+      buffer_bitmap[i] = 0xff; // completely received byte
+    }
+    // fill from bit 'start_bit' until bit 7, both inclusive
+    buffer_bitmap[start_byte] |= ~((1U << start_bit) - 1U);
+    // fill from bit 0 until bit 'end_bit', both inclusive
+    buffer_bitmap[end_byte] |= (1U << (end_bit + 1U)) - 1U;
+  }
+}
+
+/* ********************************************************************************* */
+
+static uint64_t ndpi_reassemble_calculate_buffer_len(uint64_t min_len)
+{
+   return min_len + (8 - min_len % 8);
+}
+
+void ndpi_reassemble_set_buffer_len(struct ndpi_reasm * const reasm,
+                                    uint64_t buffer_len)
+{
+  u_int32_t new_len = ndpi_reassemble_calculate_buffer_len(buffer_len);
+
+  reasm->buf_req = buffer_len;
+
+  if (reasm->buf_len == 0)
+  {
+    reasm->buf_len = new_len;
+  } else if (new_len > reasm->buf_len) {
+    uint8_t * const reasm_buf = (uint8_t *)ndpi_realloc(reasm->buf, reasm->buf_len, new_len);
+    uint8_t * const reasm_buf_bitmap = (uint8_t *)ndpi_realloc(reasm->buf_bitmap, reasm->buf_len, new_len);
+    if (reasm_buf != NULL && reasm_buf_bitmap != NULL) {
+      reasm->buf = reasm_buf;
+      reasm->buf_bitmap = reasm_buf_bitmap;
+
+      memset(reasm->buf_bitmap + reasm->buf_len, 0, new_len - reasm->buf_len);
+      reasm->buf_len = new_len;
+      reasm->buf_last_pos = ndpi_min(reasm->buf_last_pos, new_len);
+    }
+  }
+}
+
+/* ********************************************************************************* */
+
+int ndpi_reassemble(struct ndpi_reasm * const reasm, uint8_t const * const frag,
+                    uint64_t frag_len, uint64_t frag_offset)
+{
+  uint64_t max_reasm_buffer_len;
+  const uint64_t last_pos = frag_offset + frag_len;
+
+  if (reasm->buf == NULL)
+  {
+    if (reasm->buf_len == 0)
+    {
+      ndpi_reassemble_set_buffer_len(reasm,
+                                     4096 /* default: a couple of MTU sized packets */);
+    }
+    max_reasm_buffer_len = reasm->buf_len;
+    reasm->buf = (uint8_t *)ndpi_malloc(max_reasm_buffer_len);
+    reasm->buf_bitmap = (uint8_t *)ndpi_calloc(max_reasm_buffer_len,
+                                               sizeof(uint8_t));
+    if (reasm->buf == NULL || reasm->buf_bitmap == NULL)
+    {
+      return -1;
+    }
+    reasm->buf_last_pos = 0;
+  } else {
+    max_reasm_buffer_len = reasm->buf_len;
+  }
+
+  const uint64_t reasm_buffer_bitmap_len = reasm->buf_len / 8;
+
+  if (last_pos > max_reasm_buffer_len)
+  {
+    return -3;
+  }
+
+  memcpy(&reasm->buf[frag_offset], frag, frag_len);
+  if (last_pos > reasm->buf_last_pos)
+  {
+    reasm->buf_last_pos = last_pos;
+  }
+  update_reasm_buf_bitmap(reasm->buf_bitmap, reasm_buffer_bitmap_len, frag_offset, frag_len);
+
+  return 0;
+}
+
+int ndpi_reassemble_payload(struct ndpi_reasm * const reasm, struct ndpi_packet_struct * packet)
+{
+  return ndpi_reassemble(reasm, packet->payload, packet->payload_packet_len, reasm->buf_last_pos);
+}
+
+/* ********************************************************************************* */
+
+int ndpi_reassemble_in_progress(struct ndpi_reasm * const reasm)
+{
+  return reasm->buf_len > 0 && ndpi_reassemble_is_complete(reasm) == 0;
+}
+
+/* ********************************************************************************* */
+
+int ndpi_reassemble_is_complete(struct ndpi_reasm * const reasm)
+{
+  const u_int32_t complete_bytes = reasm->buf_last_pos / 8;
+  const u_int32_t remaining_bits = reasm->buf_last_pos % 8;
+  u_int32_t i;
+
+  if (reasm->buf_last_pos < reasm->buf_req)
+  {
+    return 0;
+  }
+
+  if (reasm->buf_last_pos > reasm->buf_req)
+  {
+    return 1;
+  }
+
+  for(i = 0; i < complete_bytes; i++)
+  {
+    if (reasm->buf_bitmap[i] != 0xff)
+    {
+      return 0;
+    }
+  }
+
+  if (remaining_bits && reasm->buf_bitmap[complete_bytes] != (1U << (remaining_bits)) - 1)
+  {
+    return 0;
+  }
+
+  return 1;
+}
+
+/* ********************************************************************************* */
+
+void ndpi_reassemble_swap_payload(struct ndpi_packet_struct * packet,
+                                  struct ndpi_reasm const * reasm,
+                                  u_int8_t const ** const original_payload,
+                                  u_int16_t * original_payload_packet_len)
+{
+  if (reasm->buf == packet->payload)
+  {
+    packet->payload = *original_payload;
+    packet->payload_packet_len = *original_payload_packet_len;
+  } else {
+    *original_payload = packet->payload;
+    *original_payload_packet_len = packet->payload_packet_len;
+
+    packet->payload = reasm->buf;
+    packet->payload_packet_len = reasm->buf_req;
+  }
+}
diff --git a/src/lib/protocols/kerberos.c b/src/lib/protocols/kerberos.c
index 92ee7defe..fd35be435 100644
--- a/src/lib/protocols/kerberos.c
+++ b/src/lib/protocols/kerberos.c
@@ -322,30 +322,27 @@ void ndpi_search_kerberos(struct ndpi_detection_module_struct *ndpi_struct,
 #ifdef KERBEROS_DEBUG
   printf("\n[Kerberos] Process packet [len: %u]\n", packet->payload_packet_len);
 #endif
-    
-  if(flow->kerberos_buf.pktbuf != NULL) {
-    u_int missing = flow->kerberos_buf.pktbuf_maxlen - flow->kerberos_buf.pktbuf_currlen;
-
-    if(packet->payload_packet_len <= missing) {
-      memcpy(&flow->kerberos_buf.pktbuf[flow->kerberos_buf.pktbuf_currlen], packet->payload, packet->payload_packet_len);
-      flow->kerberos_buf.pktbuf_currlen += packet->payload_packet_len;
-
-      if(flow->kerberos_buf.pktbuf_currlen == flow->kerberos_buf.pktbuf_maxlen) {
-	original_packet_payload = packet->payload;
-	original_payload_packet_len = packet->payload_packet_len;
-	packet->payload = (u_int8_t *)flow->kerberos_buf.pktbuf;
-	packet->payload_packet_len = flow->kerberos_buf.pktbuf_currlen;
+
+  if (ndpi_reassemble_in_progress(&flow->reassemble) != 0)
+  {
+    if (ndpi_reassemble_payload(&flow->reassemble, packet) != 0)
+    {
+      return;
+    }
+    if (ndpi_reassemble_is_complete(&flow->reassemble) != 0)
+    {
+      ndpi_reassemble_swap_payload(packet, &flow->reassemble, &original_packet_payload,
+                                   &original_payload_packet_len);
 #ifdef KERBEROS_DEBUG
-	printf("[Kerberos] Packet is now full: processing\n");
+      printf("[Kerberos] Packet is now full: processing\n");
 #endif
-      } else {
+    } else {
 #ifdef KERBEROS_DEBUG
-	printf("[Kerberos] Missing %u bytes: skipping\n",
-	       flow->kerberos_buf.pktbuf_maxlen - flow->kerberos_buf.pktbuf_currlen);
+        printf("[Kerberos] Missing %u bytes: skipping\n",
+               flow->reassemble.buf_len - flow->reassemble.buf_last_pos);
 #endif
 
-	return;
-      }
+        return;
     }
   }
 
@@ -356,44 +353,30 @@ void ndpi_search_kerberos(struct ndpi_detection_module_struct *ndpi_struct,
 
     if(packet->tcp) {
       kerberos_len = ntohl(get_u_int32_t(packet->payload, 0)),
-	expected_len = packet->payload_packet_len - 4;
+      expected_len = packet->payload_packet_len - 4;
       base_offset = 4;
-    } else
+    } else {
       base_offset = 0, kerberos_len = expected_len = packet->payload_packet_len;
+    }
 
 #ifdef KERBEROS_DEBUG
     printf("[Kerberos] [Kerberos len: %u][expected_len: %u]\n", kerberos_len, expected_len);
 #endif
 
-    if(kerberos_len < 12000) {
+    if (kerberos_len < 12000 && packet->tcp != NULL &&
+        kerberos_len > expected_len) {
       /*
-	Kerberos packets might be too long for a TCP packet
-	so it could be split across two packets. Instead of
-	rebuilding the stream we use a heuristic approach
-      */
-      if(kerberos_len > expected_len) {
-	if(packet->tcp) {
-	  if(flow->kerberos_buf.pktbuf == NULL) {
-	    flow->kerberos_buf.pktbuf = (char*)ndpi_malloc(kerberos_len+4);
-
-	    if(flow->kerberos_buf.pktbuf != NULL) {
-	      flow->kerberos_buf.pktbuf_maxlen = kerberos_len+4;	      
+       * Kerberos packets might be too long for a TCP packet
+       * so it could be split across two packets. Instead of
+       * rebuilding the stream we use a heuristic approach
+       */
+      ndpi_reassemble_set_buffer_len(&flow->reassemble, kerberos_len + 4);
+      ndpi_reassemble_payload(&flow->reassemble, packet);
 #ifdef KERBEROS_DEBUG
-	      printf("[Kerberos] Allocated %u bytes\n", flow->kerberos_buf.pktbuf_maxlen);
-#endif	      
-	    }
-	  }
-	  
-	  if(flow->kerberos_buf.pktbuf != NULL) {
-	    if(packet->payload_packet_len <= flow->kerberos_buf.pktbuf_maxlen) {
-	      memcpy(flow->kerberos_buf.pktbuf, packet->payload, packet->payload_packet_len);
-	      flow->kerberos_buf.pktbuf_currlen = packet->payload_packet_len;
-	    }
-	  }
-	}
-	
-	return;
-      } else if(kerberos_len == expected_len) {
+      printf("[Kerberos] Allocated %u bytes\n", flow->reassemble.buf_len);
+#endif
+      return;
+    } else if(kerberos_len == expected_len) {
 	if(packet->payload_packet_len > 64) {
 	  u_int16_t koffset, i;
 
@@ -611,12 +594,10 @@ void ndpi_search_kerberos(struct ndpi_detection_module_struct *ndpi_struct,
 		    ndpi_snprintf(flow->protos.kerberos.domain, sizeof(flow->protos.kerberos.domain), "%s", realm_str);
 
 		    /* If necessary we can decode sname */
-		    if(flow->kerberos_buf.pktbuf) {
-			    ndpi_free(flow->kerberos_buf.pktbuf);
-			    packet->payload = original_packet_payload;
-			    packet->payload_packet_len = original_payload_packet_len;
+		    if(ndpi_reassemble_in_progress(&flow->reassemble) != 0) {
+                ndpi_reassemble_swap_payload(packet, &flow->reassemble, &original_packet_payload,
+                                             &original_payload_packet_len);
 		    }
-		    flow->kerberos_buf.pktbuf = NULL;
 		  }
 		}
 	      }
@@ -629,14 +610,13 @@ void ndpi_search_kerberos(struct ndpi_detection_module_struct *ndpi_struct,
 	        flow->extra_packets_func = ndpi_search_kerberos_extra;
 	      }
 
-	      if(flow->kerberos_buf.pktbuf != NULL) {
-		ndpi_free(flow->kerberos_buf.pktbuf);
-		packet->payload = original_packet_payload;
-		packet->payload_packet_len = original_payload_packet_len;
-		flow->kerberos_buf.pktbuf = NULL;
-	      }
-
-	      return;
+          if (ndpi_reassemble_in_progress(&flow->reassemble) != 0)
+          {
+            ndpi_reassemble_swap_payload(packet, &flow->reassemble, &original_packet_payload,
+                                         &original_payload_packet_len);
+            // XXX: Free reassemble data?
+           }
+           return;
 	    } else if(msg_type == 0x0d) /* TGS-REP */ {
 	      NDPI_LOG_DBG(ndpi_struct, "[Kerberos] Processing TGS-REP\n");
 
@@ -659,9 +639,8 @@ void ndpi_search_kerberos(struct ndpi_detection_module_struct *ndpi_struct,
 	    }
 
 	    return;
-	  }
-	}
       }
+	}
     } else {
 #ifdef KERBEROS_DEBUG
       printf("[Kerberos][s/dport: %u/%u] Skipping packet: too long [kerberos_len: %u]\n",
@@ -679,18 +658,10 @@ void ndpi_search_kerberos(struct ndpi_detection_module_struct *ndpi_struct,
 static int ndpi_search_kerberos_extra(struct ndpi_detection_module_struct *ndpi_struct,
 				      struct ndpi_flow_struct *flow)
 {
-  struct ndpi_packet_struct *packet = &ndpi_struct->packet;
-
 #ifdef KERBEROS_DEBUG
   printf("[Kerberos] Extra function\n");
 #endif
 
-  /* Unfortunately, generic "extra function" code doesn't honour protocol bitmask */
-  /* TODO: handle that in ndpi_main.c for all the protocols */
-  if(packet->payload_packet_len == 0 ||
-     packet->tcp_retransmission)
-    return 1;
-
   /* Possibly dissect the reply */
   ndpi_search_kerberos(ndpi_struct, flow);
 
diff --git a/src/lib/protocols/quic.c b/src/lib/protocols/quic.c
index aafeb3397..f93cb5f53 100644
--- a/src/lib/protocols/quic.c
+++ b/src/lib/protocols/quic.c
@@ -1007,72 +1007,6 @@ static uint8_t *decrypt_initial_packet(struct ndpi_detection_module_struct *ndpi
   return NULL;
 }
 
-static void update_reasm_buf_bitmap(u_int8_t *buffer_bitmap,
-				    const u_int32_t buffer_bitmap_size,
-				    const u_int32_t recv_pos,
-				    const u_int32_t recv_len)
-{
-  if (!recv_len || !buffer_bitmap_size || recv_pos + recv_len > buffer_bitmap_size * 8)
-    return;
-  const u_int32_t start_byte = recv_pos / 8;
-  const u_int32_t end_byte = (recv_pos + recv_len - 1) / 8;
-  const u_int32_t start_bit = recv_pos % 8;
-  const u_int32_t end_bit = (start_bit + recv_len - 1) % 8;
-  if (start_byte == end_byte)
-    buffer_bitmap[start_byte] |= (((1U << recv_len) - 1U) << start_bit); // fill from bit 'start_bit' until bit 'end_bit', both inclusive
-  else{
-    u_int32_t i;
-
-    for (i = start_byte + 1; i <= end_byte - 1; i++)
-      buffer_bitmap[i] = 0xff; // completely received byte
-    buffer_bitmap[start_byte] |= ~((1U << start_bit) - 1U); // fill from bit 'start_bit' until bit 7, both inclusive
-    buffer_bitmap[end_byte] |= (1U << (end_bit + 1U)) - 1U; // fill from bit 0 until bit 'end_bit', both inclusive
-  }
-}
-
-static int is_reasm_buf_complete(const u_int8_t *buffer_bitmap,
-                                 const u_int32_t buffer_len)
-{
-  const u_int32_t complete_bytes = buffer_len / 8;
-  const u_int32_t remaining_bits = buffer_len % 8;
-  u_int32_t i;
-
-  for(i = 0; i < complete_bytes; i++)
-    if (buffer_bitmap[i] != 0xff)
-      return 0;
-
-  if (remaining_bits && buffer_bitmap[complete_bytes] != (1U << (remaining_bits)) - 1) 
-    return 0;
-    
-  return 1;
-}
-
-static int __reassemble(struct ndpi_flow_struct *flow, const u_int8_t *frag,
-                        uint64_t frag_len, uint64_t frag_offset,
-                        const u_int8_t **buf, u_int64_t *buf_len)
-{
-  const uint64_t max_quic_reasm_buffer_len = 4096; /* Let's say a couple of full-MTU packets... Must be multiple of 8*/
-  const uint64_t quic_reasm_buffer_bitmap_len = max_quic_reasm_buffer_len / 8;
-  const uint64_t last_pos = frag_offset + frag_len;
-
-  if(!flow->l4.udp.quic_reasm_buf) {
-    flow->l4.udp.quic_reasm_buf = (uint8_t *)ndpi_malloc(max_quic_reasm_buffer_len);
-    flow->l4.udp.quic_reasm_buf_bitmap = (uint8_t *)ndpi_calloc(quic_reasm_buffer_bitmap_len, sizeof(uint8_t));
-    if(!flow->l4.udp.quic_reasm_buf || !flow->l4.udp.quic_reasm_buf_bitmap)
-      return -1; /* Memory error */
-    flow->l4.udp.quic_reasm_buf_last_pos = 0;
-  }
-  if(last_pos > max_quic_reasm_buffer_len)
-    return -3; /* Buffer too small */
-
-  memcpy(&flow->l4.udp.quic_reasm_buf[frag_offset], frag, frag_len);
-  flow->l4.udp.quic_reasm_buf_last_pos = last_pos > flow->l4.udp.quic_reasm_buf_last_pos ? last_pos : flow->l4.udp.quic_reasm_buf_last_pos;
-  update_reasm_buf_bitmap(flow->l4.udp.quic_reasm_buf_bitmap, quic_reasm_buffer_bitmap_len, frag_offset, frag_len); 
-
-  *buf = flow->l4.udp.quic_reasm_buf;
-  *buf_len = flow->l4.udp.quic_reasm_buf_last_pos;
-  return 0;
-}
 static int is_ch_complete(const u_int8_t *buf, uint64_t buf_len)
 {
   uint32_t msg_len;
@@ -1087,17 +1021,16 @@ static int is_ch_complete(const u_int8_t *buf, uint64_t buf_len)
 }
 static int is_ch_reassembler_pending(struct ndpi_flow_struct *flow)
 {
-  return flow->l4.udp.quic_reasm_buf != NULL &&
-    !(is_reasm_buf_complete(flow->l4.udp.quic_reasm_buf_bitmap, flow->l4.udp.quic_reasm_buf_last_pos)
-      && is_ch_complete(flow->l4.udp.quic_reasm_buf, flow->l4.udp.quic_reasm_buf_last_pos));
+  return flow->reassemble.buf != NULL &&
+         !(ndpi_reassemble_is_complete(&flow->reassemble)
+            && is_ch_complete(flow->reassemble.buf, flow->reassemble.buf_last_pos));
 }
 static const uint8_t *get_reassembled_crypto_data(struct ndpi_detection_module_struct *ndpi_struct,
-						  struct ndpi_flow_struct *flow,
-						  const u_int8_t *frag,
-						  uint64_t frag_offset, uint64_t frag_len,
-						  uint64_t *crypto_data_len)
+                                                  struct ndpi_flow_struct *flow,
+                                                  const u_int8_t *frag,
+                                                  uint64_t frag_offset, uint64_t frag_len,
+                                                  uint64_t *crypto_data_len)
 {
-  const u_int8_t *crypto_data;
   int rc;
 
   NDPI_LOG_DBG2(ndpi_struct, "frag %d/%d\n", frag_offset, frag_len);
@@ -1110,18 +1043,23 @@ static const uint8_t *get_reassembled_crypto_data(struct ndpi_detection_module_s
     return frag;
   }
 
-  rc = __reassemble(flow, frag, frag_len, frag_offset,
-                    &crypto_data, crypto_data_len);
+  if (frag_offset == 0 && frag_len >= 4)
+  {
+    ndpi_reassemble_set_buffer_len(&flow->reassemble, (frag[1] << 16) + (frag[2] << 8) + frag[3] + 4);
+  }
+  rc = ndpi_reassemble(&flow->reassemble, frag, frag_len, frag_offset);
   if(rc == 0) {
-    if(is_reasm_buf_complete(flow->l4.udp.quic_reasm_buf_bitmap, *crypto_data_len) &&
-       is_ch_complete(crypto_data, *crypto_data_len)) {
+    if(ndpi_reassemble_is_complete(&flow->reassemble) &&
+      is_ch_complete(flow->reassemble.buf, flow->reassemble.buf_last_pos)) {
       NDPI_LOG_DBG2(ndpi_struct, "Reassembler completed!\n");
-      return crypto_data;
+      *crypto_data_len = flow->reassemble.buf_last_pos;
+      return flow->reassemble.buf;
     }
     NDPI_LOG_DBG2(ndpi_struct, "CH not yet completed\n");
   } else {
     NDPI_LOG_DBG(ndpi_struct, "Reassembler error: %d\n", rc);
   }
+
   return NULL;
 }
 
diff --git a/tests/unit/unit.c b/tests/unit/unit.c
index 13de97a1f..152777be7 100644
--- a/tests/unit/unit.c
+++ b/tests/unit/unit.c
@@ -317,6 +317,76 @@ int serializeProtoUnitTest(void)
 
 /* *********************************************** */
 
+static int reassmbleUnitTest()
+{
+  u_int8_t const * org_payload = NULL;
+  u_int16_t org_payload_len;
+  char const * const packets[] = { "There", " ", "is", " ", "some", " ", "payload", " ", "on", " ", "the", " ", "wire!" };
+  struct ndpi_packet_struct packet = {};
+  struct ndpi_reasm reasm = {};
+  size_t i, j;
+
+  assert(ndpi_reassemble_in_progress(&reasm) == 0);
+  assert(ndpi_reassemble_is_complete(&reasm) != 0);
+  for (i = 0; i < NDPI_ARRAY_LENGTH(packets); ++i)
+  {
+    ndpi_reassemble_set_buffer_len(&reasm, reasm.buf_req + strlen(packets[i]));
+    assert(ndpi_reassemble_is_complete(&reasm) == 0);
+    assert(ndpi_reassemble_in_progress(&reasm) != 0);
+
+    packet.payload = (u_int8_t const *)packets[i];
+    packet.payload_packet_len = strlen(packets[i]);
+
+    assert(ndpi_reassemble_payload(&reasm, &packet) == 0);
+  }
+  assert(ndpi_reassemble_in_progress(&reasm) == 0);
+  assert(ndpi_reassemble_is_complete(&reasm) != 0);
+
+  ndpi_reassemble_swap_payload(&packet, &reasm, &org_payload, &org_payload_len);
+
+  j = 0;
+  for (i = 0; i < NDPI_ARRAY_LENGTH(packets); ++i)
+  {
+    assert(strncmp((char *)&packet.payload[j], packets[i], strlen(packets[i])) == 0);
+    j += strlen(packets[i]);
+  }
+  assert(j == packet.payload_packet_len);
+
+  ndpi_reassemble_swap_payload(&packet, &reasm, &org_payload, &org_payload_len);
+
+  free(reasm.buf);
+  free(reasm.buf_bitmap);
+  memset(&reasm, 0, sizeof(reasm));
+
+  assert(ndpi_reassemble_in_progress(&reasm) == 0);
+  assert(ndpi_reassemble_is_complete(&reasm) != 0);
+  j = 0;
+  for (i = 0; i < NDPI_ARRAY_LENGTH(packets); ++i)
+  {
+    j += strlen(packets[i]);
+  }
+  ndpi_reassemble_set_buffer_len(&reasm, j);
+  assert(ndpi_reassemble_is_complete(&reasm) == 0);
+  for (i = NDPI_ARRAY_LENGTH(packets); i > 0; --i)
+  {
+    assert(ndpi_reassemble_is_complete(&reasm) == 0);
+    assert(ndpi_reassemble_in_progress(&reasm) != 0);
+
+    j -= strlen(packets[i - 1]);
+    assert(ndpi_reassemble(&reasm, (u_int8_t const *)packets[i - 1], strlen(packets[i - 1]), j) == 0);
+  }
+
+  free(reasm.buf);
+  free(reasm.buf_bitmap);
+  memset(&reasm, 0, sizeof(reasm));
+
+  printf("%30s                      OK\n", __FUNCTION__);
+
+  return 0;
+}
+
+/* *********************************************** */
+
 int main(int argc, char **argv) {
 #ifndef WIN32
   int c;
@@ -355,6 +425,7 @@ int main(int argc, char **argv) {
   /* Tests */
   if (serializerUnitTest() != 0) return -1;
   if (serializeProtoUnitTest() != 0) return -1;
+  if (reassmbleUnitTest() != 0) return -1;
 
   return 0;
 }