2 files changed, 59 insertions, 150 deletions

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;
 }