Implement nDPI patricia tree API

4 files changed, 130 insertions, 5 deletions

diff --git a/src/include/ndpi_api.h b/src/include/ndpi_api.h
index 3c30f1f69..611ebe261 100644
--- a/src/include/ndpi_api.h
+++ b/src/include/ndpi_api.h
@@ -837,6 +837,12 @@ extern "C" {
 				      char *outbuf, u_int outbuf_len);
   char* ndpi_ssl_version2str(u_int16_t version, u_int8_t *unknown_tls_version);
 
+  /* ptree (trie) API */
+  ndpi_ptree_t* ndpi_ptree_create();
+  int ndpi_ptree_insert(ndpi_ptree_t *tree, const ndpi_ip_addr_t *addr, u_int8_t bits, uint user_data);
+  int ndpi_ptree_match_addr(ndpi_ptree_t *tree, const ndpi_ip_addr_t *addr, uint *user_data);
+  void ndpi_ptree_destroy(ndpi_ptree_t *tree);
+
   /* Serializer */
   int ndpi_init_serializer_ll(ndpi_serializer *serializer, ndpi_serialization_format fmt,
 			      u_int32_t buffer_size);
diff --git a/src/include/ndpi_main.h b/src/include/ndpi_main.h
index 06aa8bc96..7263410de 100644
--- a/src/include/ndpi_main.h
+++ b/src/include/ndpi_main.h
@@ -114,6 +114,7 @@ extern "C" {
 
   extern int ndpi_parse_ip_string(const char *ip_str, ndpi_ip_addr_t *parsed_ip);
   extern char *ndpi_get_ip_string(const ndpi_ip_addr_t * ip, char *buf, u_int buf_len);
+  extern u_int8_t ndpi_is_ipv6(const ndpi_ip_addr_t *ip);
 
   extern char* ndpi_get_proto_by_id(struct ndpi_detection_module_struct *ndpi_mod, u_int id);
   u_int16_t ndpi_get_proto_by_name(struct ndpi_detection_module_struct *ndpi_mod, const char *name);
diff --git a/src/include/ndpi_typedefs.h b/src/include/ndpi_typedefs.h
index 4366df5c1..8eb481f47 100644
--- a/src/include/ndpi_typedefs.h
+++ b/src/include/ndpi_typedefs.h
@@ -1438,4 +1438,8 @@ struct ndpi_analyze_struct {
 #define MAX_SERIES_LEN      512
 #define MIN_SERIES_LEN      8
 
+/* **************************************** */
+
+typedef struct ndpi_ptree ndpi_ptree_t;
+
 #endif /* __NDPI_TYPEDEFS_H__ */
diff --git a/src/lib/ndpi_main.c b/src/lib/ndpi_main.c
index a72917b7f..c665429ef 100644
--- a/src/lib/ndpi_main.c
+++ b/src/lib/ndpi_main.c
@@ -145,6 +145,14 @@ char * ndpi_strdup(const char *s)
 
 /* *********************************************************************************** */
 
+/* Opaque structure defined here */
+struct ndpi_ptree {
+  patricia_tree_t *v4;
+  patricia_tree_t *v6;
+};
+
+/* *********************************************************************************** */
+
 u_int32_t ndpi_detection_get_sizeof_ndpi_flow_struct(void) { return(sizeof(struct ndpi_flow_struct)); }
 
 /* *********************************************************************************** */
@@ -1835,7 +1843,7 @@ static int ac_match_handler(AC_MATCH_t *m, AC_TEXT_t *txt, AC_REP_t *match) {
 
 /* ******************************************************************** */
 
-static int fill_prefix_v4(prefix_t *p, struct in_addr *a, int b, int mb) {
+static int fill_prefix_v4(prefix_t *p, const struct in_addr *a, int b, int mb) {
   do {
     if(b < 0 || b > mb)
       return(-1);
@@ -1852,6 +1860,18 @@ static int fill_prefix_v4(prefix_t *p, struct in_addr *a, int b, int mb) {
 
 /* ******************************************* */
 
+static int fill_prefix_v6(prefix_t *prefix, const struct in6_addr *addr, int bits, int maxbits) {
+  if(bits < 0 || bits > maxbits)
+    return -1;
+
+  memcpy(&prefix->add.sin6, addr, (maxbits + 7) / 8);
+  prefix->family = AF_INET6, prefix->bitlen = bits, prefix->ref_count = 0;
+
+  return 0;
+}
+
+/* ******************************************* */
+
 u_int16_t ndpi_network_ptree_match(struct ndpi_detection_module_struct *ndpi_str,
 				   struct in_addr *pin /* network byte order */) {
   prefix_t prefix;
@@ -5637,14 +5657,23 @@ void ndpi_packet_dst_ip_get(const struct ndpi_packet_struct *packet, ndpi_ip_add
 
 /* ********************************************************************************* */
 
+u_int8_t ndpi_is_ipv6(const ndpi_ip_addr_t *ip) {
+#ifdef NDPI_DETECTION_SUPPORT_IPV6
+  return(ip->ipv6.u6_addr.u6_addr32[1] != 0 ||
+         ip->ipv6.u6_addr.u6_addr32[2] != 0 ||
+         ip->ipv6.u6_addr.u6_addr32[3] != 0);
+#else
+  return(0);
+#endif
+}
+
+/* ********************************************************************************* */
+
 char *ndpi_get_ip_string(const ndpi_ip_addr_t * ip, char *buf, u_int buf_len) {
   const u_int8_t *a = (const u_int8_t *) &ip->ipv4;
 
 #ifdef NDPI_DETECTION_SUPPORT_IPV6
-  if(ip->ipv6.u6_addr.u6_addr32[1] != 0 ||
-     ip->ipv6.u6_addr.u6_addr32[2] != 0 ||
-     ip->ipv6.u6_addr.u6_addr32[3] != 0) {
-
+  if(ndpi_is_ipv6(ip)) {
     if(inet_ntop(AF_INET6, &ip->ipv6.u6_addr, buf, buf_len) == NULL)
       buf[0] = '\0';
 
@@ -6512,3 +6541,88 @@ ndpi_l4_proto_info ndpi_get_l4_proto_info(struct ndpi_detection_module_struct *n
 
   return(ndpi_l4_proto_unknown); /* default */
 }
+
+/* ******************************************************************** */
+
+ndpi_ptree_t* ndpi_ptree_create() {
+  ndpi_ptree_t *tree = (ndpi_ptree_t*) ndpi_malloc(sizeof(ndpi_ptree_t));
+
+  if(tree) {
+    tree->v4 = ndpi_New_Patricia(32);
+    tree->v6 = ndpi_New_Patricia(128);
+
+    if((!tree->v4) || (!tree->v6)) {
+      ndpi_ptree_destroy(tree);
+      return(NULL);
+    }
+  }
+
+  return(tree);
+}
+
+/* ******************************************************************** */
+
+void ndpi_ptree_destroy(ndpi_ptree_t *tree) {
+  if(tree) {
+    if(tree->v4) ndpi_Destroy_Patricia(tree->v4, free_ptree_data);
+    if(tree->v6) ndpi_Destroy_Patricia(tree->v6, free_ptree_data);
+
+    ndpi_free(tree);
+  }
+}
+
+/* ******************************************************************** */
+
+int ndpi_ptree_insert(ndpi_ptree_t *tree, const ndpi_ip_addr_t *addr, u_int8_t bits, uint user_data) {
+  u_int8_t is_v6 = ndpi_is_ipv6(addr);
+  patricia_tree_t *ptree = is_v6 ? tree->v6 : tree->v4;
+  prefix_t prefix;
+  patricia_node_t *node;
+
+  if(bits > ptree->maxbits)
+    return(-1);
+
+  if(is_v6)
+    fill_prefix_v6(&prefix, (const struct in6_addr*)&addr->ipv6, bits, ptree->maxbits);
+  else
+    fill_prefix_v4(&prefix, (const struct in_addr*)&addr->ipv4, bits, ptree->maxbits);
+
+  /* Verify that the node does not already exist */
+  node = ndpi_patricia_search_best(ptree, &prefix);
+
+  if(node && (node->prefix->bitlen == bits))
+    return(-2);
+
+  node = ndpi_patricia_lookup(ptree, &prefix);
+
+  if(node != NULL) {
+    node->value.user_value = user_data;
+    return(0);
+  }
+
+  return(-3);
+}
+
+/* ******************************************************************** */
+
+int ndpi_ptree_match_addr(ndpi_ptree_t *tree, const ndpi_ip_addr_t *addr, uint *user_data) {
+  u_int8_t is_v6 = ndpi_is_ipv6(addr);
+  patricia_tree_t *ptree = is_v6 ? tree->v6 : tree->v4;
+  prefix_t prefix;
+  patricia_node_t *node;
+  int bits = ptree->maxbits;
+
+  if(is_v6)
+    fill_prefix_v6(&prefix, (const struct in6_addr*)&addr->ipv6, bits, ptree->maxbits);
+  else
+    fill_prefix_v4(&prefix, (const struct in_addr*)&addr->ipv4, bits, ptree->maxbits);
+
+  node = ndpi_patricia_search_best(ptree, &prefix);
+
+  if(node) {
+    *user_data = node->value.user_value;
+    return(0);
+  }
+
+  return(-1);
+}