Add category and breed support for custom rules (#2872)

Close #2594

12 files changed, 110 insertions, 10 deletions

diff --git a/example/protos.txt b/example/protos.txt
index 8da9ca0b9..2dc3b3318 100644
--- a/example/protos.txt
+++ b/example/protos.txt
@@ -63,6 +63,11 @@ ipv6:[fe80::76ac:b9ff:fe6c:c124]:12717@CustomProtocolG=2050
 ipv6:[fe80::76ac:b9ff:fe6c:c124]:12718@CustomProtocolH=65535
 ipv6:[fe80::76ac:b9ff:fe6c:c124]:12719@CustomProtocolI=65534
 
+ipv6:[2001:db8::c2fd:b817:5ca8:82dd]:16690@CustomProtocolJ,breed=1
+ipv6:[2001:db8::cba5:51b2:8733:6d9e]:38542@CustomProtocolK,cat=3
+ipv6:[2001:db8::cc14:67e6:fcd:b96d]:37464@CustomProtocolL=3002,cat=5,breed=3
+ipv6:[2003:db8::4]@CustomProtocolM=3003,cat=106,breed=5
+
 #
 # You can use symbolic IP addreses if you want
 #
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_1.txt b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_1.txt
new file mode 100644
index 000000000..4d0ece167
--- /dev/null
+++ b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_1.txt
@@ -0,0 +1 @@
+ipv6:[2001:db8::c2fd:b817:5ca8:82dd]:16690@CustomProtocolJ,breed=1
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_2.txt b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_2.txt
new file mode 100644
index 000000000..4c3624310
--- /dev/null
+++ b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_2.txt
@@ -0,0 +1 @@
+ipv6:[2001:db8::cba5:51b2:8733:6d9e]:38542@CustomProtocolK,cat=3
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_3.txt b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_3.txt
new file mode 100644
index 000000000..9064b8a51
--- /dev/null
+++ b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_3.txt
@@ -0,0 +1 @@
+ipv6:[2001:db8::cc14:67e6:fcd:b96d]:37464@CustomProtocolL=3002,cat=5,breed=3
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_4.txt b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_4.txt
new file mode 100644
index 000000000..01462459e
--- /dev/null
+++ b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_4.txt
@@ -0,0 +1 @@
+ipv6:[2003:db8::4]@CustomProtocolM=3003,cat=106,breed=5
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_5.txt b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_5.txt
new file mode 100644
index 000000000..8207e6f64
--- /dev/null
+++ b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_5.txt
@@ -0,0 +1 @@
+ip:127.0.0.1@CustomProtocolJ,breed=1
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_6.txt b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_6.txt
new file mode 100644
index 000000000..8f70d5de4
--- /dev/null
+++ b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_6.txt
@@ -0,0 +1 @@
+ip:127.0.0.1@CustomProtocolJ,cat=1,breed=1
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_7.txt b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_7.txt
new file mode 100644
index 000000000..c3cd78ba2
--- /dev/null
+++ b/fuzz/corpus/fuzz_filecfg_protocols/custom_cat_7.txt
@@ -0,0 +1 @@
+ip:127.0.0.1/8@CustomProtocolJ=1000,breed=1
diff --git a/fuzz/corpus/fuzz_filecfg_protocols/protos.txt b/fuzz/corpus/fuzz_filecfg_protocols/protos.txt
index 8da9ca0b9..2dc3b3318 100644
--- a/fuzz/corpus/fuzz_filecfg_protocols/protos.txt
+++ b/fuzz/corpus/fuzz_filecfg_protocols/protos.txt
@@ -63,6 +63,11 @@ ipv6:[fe80::76ac:b9ff:fe6c:c124]:12717@CustomProtocolG=2050
 ipv6:[fe80::76ac:b9ff:fe6c:c124]:12718@CustomProtocolH=65535
 ipv6:[fe80::76ac:b9ff:fe6c:c124]:12719@CustomProtocolI=65534
 
+ipv6:[2001:db8::c2fd:b817:5ca8:82dd]:16690@CustomProtocolJ,breed=1
+ipv6:[2001:db8::cba5:51b2:8733:6d9e]:38542@CustomProtocolK,cat=3
+ipv6:[2001:db8::cc14:67e6:fcd:b96d]:37464@CustomProtocolL=3002,cat=5,breed=3
+ipv6:[2003:db8::4]@CustomProtocolM=3003,cat=106,breed=5
+
 #
 # You can use symbolic IP addreses if you want
 #
diff --git a/src/lib/ndpi_main.c b/src/lib/ndpi_main.c
index 9aa62fd0f..48dda1f2b 100644
--- a/src/lib/ndpi_main.c
+++ b/src/lib/ndpi_main.c
@@ -5275,6 +5275,10 @@ static int ndpi_handle_rule(struct ndpi_detection_module_struct *ndpi_str,
   u_int subprotocol_id, i;
   int ret = 0;
 
+  char *additional_params = NULL;
+  ndpi_protocol_category_t category = NDPI_PROTOCOL_CATEGORY_UNSPECIFIED;
+  ndpi_protocol_breed_t breed = NDPI_PROTOCOL_ACCEPTABLE;
+
   at = strrchr(rule, '@');
   if(at == NULL) {
     /* This looks like a mask rule or an invalid rule */
@@ -5336,36 +5340,80 @@ static int ndpi_handle_rule(struct ndpi_detection_module_struct *ndpi_str,
 
   if(def == NULL) {
     ndpi_port_range ports_a[MAX_DEFAULT_PORTS], ports_b[MAX_DEFAULT_PORTS];
-    char *equal = strchr(proto, '=');
     u_int16_t user_proto_id, proto_id;
 
     proto_id = ndpi_str->num_supported_protocols; /* First free id */
     user_proto_id = proto_id; /* By default, external id is equal to the internal one */
 
-    if(equal != NULL) {
-      /* PROTO=VALUE */
+    char *first_comma = strchr(proto, ',');
+    char *proto_name = proto;
+    
+    if(first_comma != NULL) {
+      first_comma[0] = '\0';
+      additional_params = &first_comma[1];
+    }
 
+    char *equal = strchr(proto_name, '=');
+    
+    if(equal != NULL) {
       equal[0] = '\0';
-      user_proto_id = atoi(&equal[1]);
+      char *id_part = &equal[1];
+      
+      const char *errstrp;
+      user_proto_id = ndpi_strtonum(id_part, ndpi_str->num_supported_protocols, 65535, &errstrp, 10);
+      if(errstrp != NULL) {
+        NDPI_LOG_ERR(ndpi_str, "Invalid protocol ID '%s': %s\n", id_part, errstrp);
+        return(-1);
+      }
 
-      NDPI_LOG_DBG(ndpi_str, "***** ADDING MAPPING %s: %u -> %u\n", proto, proto_id, user_proto_id);
+      NDPI_LOG_DBG(ndpi_str, "***** ADDING MAPPING %s: %u -> %u\n", proto_name, proto_id, user_proto_id);
     }
 
     /* TODO */
     if(ndpi_str->num_custom_protocols >= (NDPI_MAX_NUM_CUSTOM_PROTOCOLS - 1)) {
       NDPI_LOG_ERR(ndpi_str, "Too many protocols defined (%u): skipping protocol %s\n",
-		   ndpi_str->num_custom_protocols, proto);
+		   ndpi_str->num_custom_protocols, proto_name);
       return(-2);
     }
 
     ndpi_add_user_proto_id_mapping(ndpi_str, proto_id, user_proto_id);
 
+    /* Parse additional parameters like cat= and breed= */
+    if(additional_params != NULL) {
+      char *param = NULL;
+      char *params_copy = additional_params;
+
+      while((param = strsep(&params_copy, ",")) != NULL) {
+        if(strncmp(param, "cat=", 4) == 0) {
+          char *cat_value = &param[4];
+          const char *errstrp;
+          
+          int cat_id = ndpi_strtonum(cat_value, 1, NDPI_PROTOCOL_NUM_CATEGORIES-1, &errstrp, 10);
+          if(errstrp == NULL) {
+            category = (ndpi_protocol_category_t)cat_id;
+          } else {
+            NDPI_LOG_ERR(ndpi_str, "Invalid category ID '%s': %s\n", cat_value, errstrp);
+          }
+        } else if (strncmp(param, "breed=", 6) == 0) {
+          char *breed_value = &param[6];
+          const char *errstrp;
+          
+          int breed_id = ndpi_strtonum(breed_value, NDPI_PROTOCOL_SAFE, NDPI_PROTOCOL_UNRATED-1, &errstrp, 10);
+          if(errstrp == NULL) {
+            breed = (ndpi_protocol_breed_t)breed_id;
+          } else {
+            NDPI_LOG_ERR(ndpi_str, "Invalid breed ID '%s': %s\n", breed_value, errstrp);
+          }
+        }
+      }
+    }
+
     ndpi_set_proto_defaults(ndpi_str, 1 /* is_cleartext */,
 			    1 /* is_app_protocol */,
-			    NDPI_PROTOCOL_ACCEPTABLE, /* TODO add protocol breed support in rules */
+			    breed,
 			    proto_id,
 			    proto, /* protoName */
-			    NDPI_PROTOCOL_CATEGORY_UNSPECIFIED, /* TODO add protocol category support in rules */
+			    category,
 			    NDPI_PROTOCOL_QOE_CATEGORY_UNSPECIFIED,
 			    ndpi_build_default_ports(ports_a, 0, 0, 0, 0, 0) /* TCP */,
 			    ndpi_build_default_ports(ports_b, 0, 0, 0, 0, 0) /* UDP */,
@@ -5456,14 +5504,14 @@ static int ndpi_handle_rule(struct ndpi_detection_module_struct *ndpi_str,
       if(rc != 0)
 	return(rc);
     } else {
-      ndpi_add_host_url_subprotocol(ndpi_str, value, subprotocol_id, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED,
-				    NDPI_PROTOCOL_ACCEPTABLE, 0);
+      ndpi_add_host_url_subprotocol(ndpi_str, value, subprotocol_id, category, breed, 0);
     }
   }
 
   return(ret);
 }
 
+
 /* ******************************************************************** */
 
 /*
diff --git a/tests/cfgs/default/pcap/custom_breed_cat.pcap b/tests/cfgs/default/pcap/custom_breed_cat.pcap
new file mode 100644
index 000000000..9c3e9d593
--- /dev/null
+++ b/tests/cfgs/default/pcap/custom_breed_cat.pcap
diff --git a/tests/cfgs/default/result/custom_breed_cat.pcap.out b/tests/cfgs/default/result/custom_breed_cat.pcap.out
new file mode 100644
index 000000000..715fa8766
--- /dev/null
+++ b/tests/cfgs/default/result/custom_breed_cat.pcap.out
@@ -0,0 +1,35 @@
+DPI Packets (TCP):	4	(1.00 pkts/flow)
+Confidence Match by custom rule: 4 (flows)
+Num dissector calls: 0 (0.00 diss/flow)
+LRU cache ookla:      0/0/0 (insert/search/found)
+LRU cache bittorrent: 0/0/0 (insert/search/found)
+LRU cache stun:       0/0/0 (insert/search/found)
+LRU cache tls_cert:   0/0/0 (insert/search/found)
+LRU cache mining:     0/0/0 (insert/search/found)
+LRU cache msteams:    0/0/0 (insert/search/found)
+LRU cache fpc_dns:    0/0/0 (insert/search/found)
+Automa host:          0/0 (search/found)
+Automa domain:        0/0 (search/found)
+Automa tls cert:      0/0 (search/found)
+Automa risk mask:     0/0 (search/found)
+Automa common alpns:  0/0 (search/found)
+Patricia risk mask:   0/0 (search/found)
+Patricia risk mask IPv6: 0/0 (search/found)
+Patricia risk:        0/0 (search/found)
+Patricia risk IPv6:   0/0 (search/found)
+Patricia protocols:   0/0 (search/found)
+Patricia protocols IPv6: 5/4 (search/found)
+
+CustomProtocolJ	1	94	1
+CustomProtocolK	1	94	1
+CustomProtocolL	1	94	1
+CustomProtocolM	1	94	1
+
+Acceptable                       2 188           2            
+Unsafe                           1 94            1            
+Dangerous                        1 94            1            
+
+	1	TCP [2001:db8::1]:33408 -> [2001:db8::c2fd:b817:5ca8:82dd]:16690 [proto: 467/CustomProtocolJ][IP: 467/CustomProtocolJ][ClearText][Confidence: Match by custom rule][FPC: 0/Unknown, Confidence: Unknown][DPI packets: 1][1 pkts/94 bytes -> 0 pkts/0 bytes][Goodput ratio: 0/0][< 1 sec][Risk: ** Unidirectional Traffic **][Risk Score: 10][Risk Info: No server to client traffic][TCP Fingerprint: 2_64_65476_dc3a5db5296b/Unknown][Plen Bins: 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
+	2	TCP [2001:db8::2]:33408 -> [2001:db8::cba5:51b2:8733:6d9e]:38542 [proto: 468/CustomProtocolK][IP: 468/CustomProtocolK][ClearText][Confidence: Match by custom rule][FPC: 0/Unknown, Confidence: Unknown][DPI packets: 1][cat: Email/3][1 pkts/94 bytes -> 0 pkts/0 bytes][Goodput ratio: 0/0][< 1 sec][Risk: ** Unidirectional Traffic **][Risk Score: 10][Risk Info: No server to client traffic][TCP Fingerprint: 2_64_65476_dc3a5db5296b/Unknown][Plen Bins: 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
+	3	TCP [2003:db8::3]:33408 -> [2001:db8::cc14:67e6:fcd:b96d]:37464 [proto: 3002/CustomProtocolL][IP: 3002/CustomProtocolL][ClearText][Confidence: Match by custom rule][FPC: 0/Unknown, Confidence: Unknown][DPI packets: 1][cat: Web/5][1 pkts/94 bytes -> 0 pkts/0 bytes][Goodput ratio: 0/0][< 1 sec][Risk: ** Unidirectional Traffic **][Risk Score: 10][Risk Info: No server to client traffic][TCP Fingerprint: 2_64_65476_dc3a5db5296b/Unknown][Plen Bins: 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
+	4	TCP [2003:db8::4]:33408 -> [2001:db8::7b51:316f:9fe:4940]:42695 [proto: 3003/CustomProtocolM][IP: 3003/CustomProtocolM][ClearText][Confidence: Match by custom rule][FPC: 0/Unknown, Confidence: Unknown][DPI packets: 1][cat: Crypto_Currency/106][1 pkts/94 bytes -> 0 pkts/0 bytes][Goodput ratio: 0/0][< 1 sec][Risk: ** Unidirectional Traffic **][Risk Score: 10][Risk Info: No server to client traffic][TCP Fingerprint: 2_64_65476_dc3a5db5296b/Unknown][Plen Bins: 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]