IPv6: add support for IPv6 risk exceptions (#2122)

1 files changed, 39 insertions, 3 deletions

diff --git a/src/lib/ndpi_utils.c b/src/lib/ndpi_utils.c
index 8bd7393d3..bd28d82c4 100644
--- a/src/lib/ndpi_utils.c
+++ b/src/lib/ndpi_utils.c
@@ -2332,11 +2332,11 @@ static u_int64_t ndpi_host_ip_risk_ptree_match(struct ndpi_detection_module_stru
   ndpi_prefix_t prefix;
   ndpi_patricia_node_t *node;
 
-  if(!ndpi_str->protocols_ptree)
+  if(!ndpi_str->ip_risk_mask_ptree)
     return((u_int64_t)-1);
 
   /* Make sure all in network byte order otherwise compares wont work */
-  ndpi_fill_prefix_v4(&prefix, pin, 32, ((ndpi_patricia_tree_t *) ndpi_str->protocols_ptree)->maxbits);
+  ndpi_fill_prefix_v4(&prefix, pin, 32, ((ndpi_patricia_tree_t *) ndpi_str->ip_risk_mask_ptree)->maxbits);
   node = ndpi_patricia_search_best(ndpi_str->ip_risk_mask_ptree, &prefix);
 
   if(node)
@@ -2347,6 +2347,26 @@ static u_int64_t ndpi_host_ip_risk_ptree_match(struct ndpi_detection_module_stru
 
 /* ********************************************************************************* */
 
+static u_int64_t ndpi_host_ip_risk_ptree_match6(struct ndpi_detection_module_struct *ndpi_str,
+					        struct in6_addr *pin6) {
+  ndpi_prefix_t prefix;
+  ndpi_patricia_node_t *node;
+
+  if(!ndpi_str->ip_risk_mask_ptree6)
+    return((u_int64_t)-1);
+
+  /* Make sure all in network byte order otherwise compares wont work */
+  ndpi_fill_prefix_v6(&prefix, pin6, 128, ((ndpi_patricia_tree_t *) ndpi_str->ip_risk_mask_ptree6)->maxbits);
+  node = ndpi_patricia_search_best(ndpi_str->ip_risk_mask_ptree6, &prefix);
+
+  if(node)
+    return(node->value.u.uv64);
+  else
+    return((u_int64_t)-1);
+}
+
+/* ********************************************************************************* */
+
 /* Check isuerDN exception */
 u_int8_t ndpi_check_issuerdn_risk_exception(struct ndpi_detection_module_struct *ndpi_str,
 					    char *issuerDN) {
@@ -2413,6 +2433,20 @@ static u_int8_t ndpi_check_ipv4_exception(struct ndpi_detection_module_struct *n
 
 /* ********************************************************************************* */
 
+static u_int8_t ndpi_check_ipv6_exception(struct ndpi_detection_module_struct *ndpi_str,
+					  struct ndpi_flow_struct *flow,
+					  struct in6_addr *addr) {
+  u_int64_t r;
+
+  r = ndpi_host_ip_risk_ptree_match6(ndpi_str, addr);
+
+  if(flow) flow->risk_mask &= r;
+
+  return((r != (u_int64_t)-1) ? 1 : 0);
+}
+
+/* ********************************************************************************* */
+
 void ndpi_handle_risk_exceptions(struct ndpi_detection_module_struct *ndpi_str,
 				 struct ndpi_flow_struct *flow) {
   if(flow->risk == 0) return; /* Nothing to do */
@@ -2449,11 +2483,13 @@ void ndpi_handle_risk_exceptions(struct ndpi_detection_module_struct *ndpi_str,
     }
   }
 
-  /* TODO: add IPv6 support */
   if(!flow->ip_risk_mask_evaluated) {
     if(flow->is_ipv6 == 0) {
       ndpi_check_ipv4_exception(ndpi_str, flow, flow->c_address.v4 /* Client */);
       ndpi_check_ipv4_exception(ndpi_str, flow, flow->s_address.v4 /* Server */);
+    } else {
+      ndpi_check_ipv6_exception(ndpi_str, flow, (struct in6_addr *)&flow->c_address.v6 /* Client */);
+      ndpi_check_ipv6_exception(ndpi_str, flow, (struct in6_addr *)&flow->s_address.v6 /* Server */);
     }
 
     flow->ip_risk_mask_evaluated = 1;