1 files changed, 13 insertions, 16 deletions

diff --git a/src/lib/protocols/wireguard.c b/src/lib/protocols/wireguard.c
index 90d814464..57ad347c3 100644
--- a/src/lib/protocols/wireguard.c
+++ b/src/lib/protocols/wireguard.c
@@ -82,14 +82,14 @@ static void ndpi_search_wireguard(struct ndpi_detection_module_struct *ndpi_stru
    * Note that handshake packets have a slightly different structure, but they are larger.
    */
   if (packet->payload_packet_len < 32) {
-    NDPI_EXCLUDE_PROTO(ndpi_struct, flow);
+    NDPI_EXCLUDE_DISSECTOR(ndpi_struct, flow);
     return;
   }
   /*
    * The next three bytes after the message type are reserved and set to zero.
    */
   if (payload[1] != 0 || payload[2] != 0 || payload[3] != 0) {
-    NDPI_EXCLUDE_PROTO(ndpi_struct, flow);
+    NDPI_EXCLUDE_DISSECTOR(ndpi_struct, flow);
     return;
   }
 
@@ -139,12 +139,13 @@ static void ndpi_search_wireguard(struct ndpi_detection_module_struct *ndpi_stru
       u_int32_t receiver_index = get_u_int32_t(payload, 8);
 
       if (receiver_index == flow->l4.udp.wireguard_peer_index[1 - packet->packet_direction]) {
-        if(packet->payload_packet_len == 100)
+        if(packet->payload_packet_len == 100 &&
+           ndpi_struct->cfg.wireguard_subclassification_by_ip /* TODO: the right option? */)
           ndpi_int_wireguard_add_connection(ndpi_struct, flow, NDPI_PROTOCOL_TUNNELBEAR);
         else
           ndpi_int_wireguard_add_connection(ndpi_struct, flow, NDPI_PROTOCOL_UNKNOWN);
       } else {
-        NDPI_EXCLUDE_PROTO(ndpi_struct, flow);
+        NDPI_EXCLUDE_DISSECTOR(ndpi_struct, flow);
       }
     }
     /* need more packets before deciding */
@@ -160,7 +161,7 @@ static void ndpi_search_wireguard(struct ndpi_detection_module_struct *ndpi_stru
       if (receiver_index == flow->l4.udp.wireguard_peer_index[1 - packet->packet_direction]) {
         ndpi_int_wireguard_add_connection(ndpi_struct, flow, NDPI_PROTOCOL_UNKNOWN);
       } else {
-        NDPI_EXCLUDE_PROTO(ndpi_struct, flow);
+        NDPI_EXCLUDE_DISSECTOR(ndpi_struct, flow);
       }
     }
     /* need more packets before deciding */
@@ -188,23 +189,19 @@ static void ndpi_search_wireguard(struct ndpi_detection_module_struct *ndpi_stru
       if (receiver_index == flow->l4.udp.wireguard_peer_index[packet->packet_direction]) {
         ndpi_int_wireguard_add_connection(ndpi_struct, flow, NDPI_PROTOCOL_UNKNOWN);
       } else {
-        NDPI_EXCLUDE_PROTO(ndpi_struct, flow);
+        NDPI_EXCLUDE_DISSECTOR(ndpi_struct, flow);
       }
     }
     /* need more packets before deciding */
   } else {
-    NDPI_EXCLUDE_PROTO(ndpi_struct, flow);
+    NDPI_EXCLUDE_DISSECTOR(ndpi_struct, flow);
   }
 }
 
-void init_wireguard_dissector(struct ndpi_detection_module_struct *ndpi_struct, u_int32_t *id)
+void init_wireguard_dissector(struct ndpi_detection_module_struct *ndpi_struct)
 {
-  ndpi_set_bitmask_protocol_detection("WireGuard", ndpi_struct, *id,
-				      NDPI_PROTOCOL_WIREGUARD,
-				      ndpi_search_wireguard,
-				      NDPI_SELECTION_BITMASK_PROTOCOL_V4_V6_UDP_WITH_PAYLOAD,
-				      SAVE_DETECTION_BITMASK_AS_UNKNOWN,
-				      ADD_TO_DETECTION_BITMASK);
-
-  *id += 1;
+  register_dissector("WireGuard", ndpi_struct,
+                     ndpi_search_wireguard,
+                     NDPI_SELECTION_BITMASK_PROTOCOL_V4_V6_UDP_WITH_PAYLOAD,
+                     1, NDPI_PROTOCOL_WIREGUARD);
 }