STUN: improve demultiplexing of DTLS packets (#2153)

Keep demultiplexing STUN/RTP/RTCP packets after DTLS ones.

We might end up processing the session a little longer, because we will
process the STUN/RTP/RTCP packets after the DTLS handshake.

4 files changed, 78 insertions, 51 deletions

diff --git a/src/lib/ndpi_main.c b/src/lib/ndpi_main.c
index ff6bf9d77..f7dff3e90 100644
--- a/src/lib/ndpi_main.c
+++ b/src/lib/ndpi_main.c
@@ -9889,15 +9889,10 @@ int ndpi_get_lru_cache_ttl(struct ndpi_detection_module_struct *ndpi_struct,
 */
 u_int8_t ndpi_extra_dissection_possible(struct ndpi_detection_module_struct *ndpi_str,
 					struct ndpi_flow_struct *flow) {
-#ifdef NDPI_ENABLE_DEBUG_MESSAGES
-  u_int16_t proto =
-    flow->detected_protocol_stack[1] ? flow->detected_protocol_stack[1] : flow->detected_protocol_stack[0];
-
-  NDPI_LOG_DBG2(ndpi_str, "[DEBUG] %s(%u.%u): %u\n", __FUNCTION__,
+  NDPI_LOG_DBG2(ndpi_str, "Protos (%u.%u): %d\n",
 	 flow->detected_protocol_stack[0],
 	 flow->detected_protocol_stack[1],
-	 proto);
-#endif
+	 !!flow->extra_packets_func);
 
   if(!flow->extra_packets_func)
     return(0);
diff --git a/src/lib/ndpi_private.h b/src/lib/ndpi_private.h
index 58260e6a1..cf549c00c 100644
--- a/src/lib/ndpi_private.h
+++ b/src/lib/ndpi_private.h
@@ -298,7 +298,7 @@ void processCertificateElements(struct ndpi_detection_module_struct *ndpi_struct
                                 struct ndpi_flow_struct *flow,
                                 u_int16_t p_offset, u_int16_t certificate_len);
 void switch_to_tls(struct ndpi_detection_module_struct *ndpi_struct,
-			  struct ndpi_flow_struct *flow);
+			  struct ndpi_flow_struct *flow, int first_dtls_pkt);
 int is_dtls(const u_int8_t *buf, u_int32_t buf_len, u_int32_t *block_len);
 void switch_extra_dissection_to_tls(struct ndpi_detection_module_struct *ndpi_struct,
 				    struct ndpi_flow_struct *flow);
diff --git a/src/lib/protocols/stun.c b/src/lib/protocols/stun.c
index 7cff65530..589c599d2 100644
--- a/src/lib/protocols/stun.c
+++ b/src/lib/protocols/stun.c
@@ -335,6 +335,15 @@ static int keep_extra_dissection(struct ndpi_detection_module_struct *ndpi_struc
   return 0;
 }
 
+static u_int32_t __get_master(struct ndpi_flow_struct *flow) {
+
+  if(flow->detected_protocol_stack[1] != NDPI_PROTOCOL_UNKNOWN)
+    return flow->detected_protocol_stack[1];
+  if(flow->detected_protocol_stack[0] != NDPI_PROTOCOL_UNKNOWN)
+    return flow->detected_protocol_stack[0];
+  return NDPI_PROTOCOL_STUN;
+}
+
 static int stun_search_again(struct ndpi_detection_module_struct *ndpi_struct,
                              struct ndpi_flow_struct *flow)
 {
@@ -343,6 +352,7 @@ static int stun_search_again(struct ndpi_detection_module_struct *ndpi_struct,
   u_int8_t first_byte;
   u_int16_t app_proto = NDPI_PROTOCOL_UNKNOWN;
   u_int32_t unused;
+  int first_dtls_pkt = 0;
 
   NDPI_LOG_DBG2(ndpi_struct, "Packet counter %d protos %d/%d\n", flow->packet_counter,
                 flow->detected_protocol_stack[0], flow->detected_protocol_stack[1]);
@@ -370,24 +380,50 @@ static int stun_search_again(struct ndpi_detection_module_struct *ndpi_struct,
     NDPI_LOG_DBG(ndpi_struct, "DROP or ZRTP range. Unexpected\n");
   } else if(first_byte <= 63) {
     NDPI_LOG_DBG(ndpi_struct, "DTLS\n");
-    if(is_dtls(packet->payload, packet->payload_packet_len, &unused) &&
-       flow->detected_protocol_stack[1] == NDPI_PROTOCOL_UNKNOWN /* No previous subclassification */) {
+
+    if(ndpi_struct->opportunistic_tls_stun_enabled &&
+       is_dtls(packet->payload, packet->payload_packet_len, &unused)) {
+
+      /* Process this DTLS packet via TLS/DTLS code but keep using STUN dissection.
+         This way we can keep demultiplexing DTLS/STUN/RTP */
+
       /* Switching to TLS dissector is tricky, because we are calling one dissector
          from another one, and that is not a common operation...
          Additionally:
-         * at that point protocol stack is already set to STUN
+         * at that point protocol stack is already set to STUN or STUN/XXX
          * we have room for only two protocols in flow->detected_protocol_stack[] so
            we can't have something like STUN/DTLS/SNAPCHAT_CALL
-         * the easiest (!?) solution is to remove STUN, and let TLS dissector to set both
-           master (i.e. DTLS) and subprotocol (if any) */
-      if(ndpi_struct->opportunistic_tls_stun_enabled) {
-        /* TODO: right way? It is a bit scary... do we need to reset something else too? */
-        ndpi_reset_detected_protocol(ndpi_struct, flow);
-        ndpi_int_change_category(ndpi_struct, flow, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED);
-
-        flow->stun.maybe_dtls = 1;
-        NDPI_LOG_DBG(ndpi_struct, "Switch to TLS\n");
-        switch_to_tls(ndpi_struct, flow);
+         * the easiest (!?) solution is to remove everything, and let the TLS dissector
+	   to set both master (i.e. DTLS) and subprotocol (if any) */
+
+      if(packet->tcp) {
+        /* TODO: TLS code assumes that DTLS is only over UDP */
+        NDPI_LOG_DBG(ndpi_struct, "Ignoring DTLS over TCP\n");
+      } else {
+        if(flow->tls_quic.certificate_processed == 1) {
+          NDPI_LOG_DBG(ndpi_struct, "Interesting DTLS stuff already processed. Ignoring\n");
+        } else {
+          if(flow->stun.maybe_dtls == 0) {
+            /* First DTLS packet of the flow */
+            first_dtls_pkt = 1;
+
+            /* TODO: right way? It is a bit scary... do we need to reset something else too? */
+            ndpi_reset_detected_protocol(ndpi_struct, flow);
+            ndpi_int_change_category(ndpi_struct, flow, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED);
+
+            /* Give room for DTLS handshake, where we might have
+               retransmissions and fragments */
+            flow->max_extra_packets_to_check += 10;
+            flow->stun.maybe_dtls = 1;
+	  }
+          NDPI_LOG_DBG(ndpi_struct, "Switch to TLS (%d/%d)\n",
+                       flow->detected_protocol_stack[0], flow->detected_protocol_stack[1]);
+
+	  switch_to_tls(ndpi_struct, flow, first_dtls_pkt);
+
+	  NDPI_LOG_DBG(ndpi_struct, "(%d/%d)\n",
+                       flow->detected_protocol_stack[0], flow->detected_protocol_stack[1]);
+        }
       }
     }
   } else if(first_byte <= 127) {
@@ -402,14 +438,19 @@ static int stun_search_again(struct ndpi_detection_module_struct *ndpi_struct,
       rtp_get_stream_type(packet->payload[1] & 0x7F, &flow->flow_multimedia_type);
 
       if(flow->detected_protocol_stack[1] != NDPI_PROTOCOL_UNKNOWN) {
-        /* STUN/SUBPROTO -> SUBPROTO/RTP */
-        ndpi_set_detected_protocol(ndpi_struct, flow,
-                                   NDPI_PROTOCOL_RTP, flow->detected_protocol_stack[0],
-                                   NDPI_CONFIDENCE_DPI);
+        if(flow->detected_protocol_stack[1] == NDPI_PROTOCOL_DTLS) {
+          /* Keep DTLS/SUBPROTO since we already wrote to flow->protos.tls_quic */
+        } else {
+          /* STUN/SUBPROTO -> SUBPROTO/RTP */
+          ndpi_set_detected_protocol(ndpi_struct, flow,
+                                     NDPI_PROTOCOL_RTP, flow->detected_protocol_stack[0],
+                                     NDPI_CONFIDENCE_DPI);
+        }
       } else {
-        /* STUN -> STUN/RTP */
+        /* STUN -> STUN/RTP, or
+           DTLS -> DTLS/RTP */
         ndpi_set_detected_protocol(ndpi_struct, flow,
-                                   NDPI_PROTOCOL_RTP, NDPI_PROTOCOL_STUN,
+                                   NDPI_PROTOCOL_RTP, __get_master(flow),
                                    NDPI_CONFIDENCE_DPI);
       }
       return 0; /* Stop */
@@ -522,7 +563,7 @@ static void ndpi_int_stun_add_connection(struct ndpi_detection_module_struct *nd
   if(flow->detected_protocol_stack[0] == NDPI_PROTOCOL_UNKNOWN ||
      app_proto != NDPI_PROTOCOL_UNKNOWN) {
     NDPI_LOG_DBG(ndpi_struct, "Setting %d\n", app_proto);
-    ndpi_set_detected_protocol(ndpi_struct, flow, app_proto, NDPI_PROTOCOL_STUN, confidence);
+    ndpi_set_detected_protocol(ndpi_struct, flow, app_proto, __get_master(flow), confidence);
   }
 
   /* This is quite complex. We want extra dissection for:
diff --git a/src/lib/protocols/tls.c b/src/lib/protocols/tls.c
index 44c36cb6d..b86d51f54 100644
--- a/src/lib/protocols/tls.c
+++ b/src/lib/protocols/tls.c
@@ -1231,18 +1231,6 @@ static int ndpi_search_tls_udp(struct ndpi_detection_module_struct *ndpi_struct,
     const u_int8_t *block = (const u_int8_t *)&p[processed];
 
     if(!is_dtls(block, p_len, &block_len)) {
-      if(processed == 0 && /* First block */
-         flow->stun.maybe_dtls == 1) {
-	/* Sometimes STUN packets are interleaved with TLS ones. Ignore STUN ones
-	   since we already are after STUN dissection and we are interested only on
-	   TLS stuff right now */
-#ifdef DEBUG_TLS
-        printf("Probably a stun packet. Keep going with TLS on next packets\n");
-#endif
-        /* Note that we can immediately "return" because, being the first block,
-	   we don't need to restore packet->payload and packet->payload_packet_len */
-        return(1); /* Keep working */
-      }
       no_dtls = 1;
       break;
     }
@@ -1332,6 +1320,7 @@ static int ndpi_search_tls_udp(struct ndpi_detection_module_struct *ndpi_struct,
 #endif
       change_cipher_found = 1;
       processed += block_len + 13;
+      flow->tls_quic.certificate_processed = 1; /* Fake, to avoid extra dissection */
       break;
     } else {
 #ifdef DEBUG_TLS
@@ -1359,7 +1348,6 @@ static int ndpi_search_tls_udp(struct ndpi_detection_module_struct *ndpi_struct,
 
   if(no_dtls || change_cipher_found || flow->tls_quic.certificate_processed) {
     NDPI_EXCLUDE_PROTO_EXT(ndpi_struct, flow, NDPI_PROTOCOL_DTLS);
-    flow->extra_packets_func = NULL;
     return(0); /* That's all */
   } else {
     return(1); /* Keep working */
@@ -1401,19 +1389,21 @@ void switch_extra_dissection_to_tls(struct ndpi_detection_module_struct *ndpi_st
 /* **************************************** */
 
 void switch_to_tls(struct ndpi_detection_module_struct *ndpi_struct,
-		   struct ndpi_flow_struct *flow)
+		   struct ndpi_flow_struct *flow, int first_time)
 {
 #ifdef DEBUG_TLS
   printf("Switching to TLS\n");
 #endif
 
-  /* Reset reassemblers */
-  if(flow->tls_quic.message[0].buffer)
-    ndpi_free(flow->tls_quic.message[0].buffer);
-  memset(&flow->tls_quic.message[0], '\0', sizeof(flow->tls_quic.message[0]));
-  if(flow->tls_quic.message[1].buffer)
-    ndpi_free(flow->tls_quic.message[1].buffer);
-  memset(&flow->tls_quic.message[1], '\0', sizeof(flow->tls_quic.message[1]));
+  if(first_time) {
+    /* Reset reassemblers */
+    if(flow->tls_quic.message[0].buffer)
+      ndpi_free(flow->tls_quic.message[0].buffer);
+    memset(&flow->tls_quic.message[0], '\0', sizeof(flow->tls_quic.message[0]));
+    if(flow->tls_quic.message[1].buffer)
+      ndpi_free(flow->tls_quic.message[1].buffer);
+    memset(&flow->tls_quic.message[1], '\0', sizeof(flow->tls_quic.message[1]));
+  }
 
   ndpi_search_tls_wrapper(ndpi_struct, flow);
 }
@@ -1498,8 +1488,9 @@ static void ndpi_int_tls_add_connection(struct ndpi_detection_module_struct *ndp
   protocol = __get_master(ndpi_struct, flow);
 
   ndpi_set_detected_protocol(ndpi_struct, flow, protocol, protocol, NDPI_CONFIDENCE_DPI);
-
-  tlsInitExtraPacketProcessing(ndpi_struct, flow);
+  /* We don't want to ovewrite STUN extra dissection, if enabled */
+  if(!flow->extra_packets_func)
+    tlsInitExtraPacketProcessing(ndpi_struct, flow);
 }
 
 /* **************************************** */