Added ability to store custom category file in patricia tree

6 files changed, 95 insertions, 37 deletions

diff --git a/example/ndpiReader.c b/example/ndpiReader.c
index 4dd3f3632..576f6a4f3 100644
--- a/example/ndpiReader.c
+++ b/example/ndpiReader.c
@@ -2218,9 +2218,17 @@ static void setupDetection(u_int16_t thread_id, pcap_t * pcap_handle) {
   if(_protoFilePath != NULL)
     ndpi_load_protocols_file(ndpi_thread_info[thread_id].workflow->ndpi_struct, _protoFilePath);
 
-  if(_customCategoryFilePath)
-    ndpi_load_categories_file(ndpi_thread_info[thread_id].workflow->ndpi_struct, _customCategoryFilePath);
+  if(_customCategoryFilePath) {
+    char *label = strrchr(_customCategoryFilePath, '/');
 
+    if(label != NULL)
+      label = &label[1];
+    else
+      label = _customCategoryFilePath;
+    
+    ndpi_load_categories_file(ndpi_thread_info[thread_id].workflow->ndpi_struct, _customCategoryFilePath, label);
+  }
+  
   if(_riskyDomainFilePath)
     ndpi_load_risk_domain_file(ndpi_thread_info[thread_id].workflow->ndpi_struct, _riskyDomainFilePath);
 
diff --git a/src/include/ndpi_api.h.in b/src/include/ndpi_api.h.in
index 9401731fa..9ad9bd7c2 100644
--- a/src/include/ndpi_api.h.in
+++ b/src/include/ndpi_api.h.in
@@ -109,7 +109,7 @@ extern "C" {
   void * ndpi_flow_malloc(size_t size);
   void   ndpi_flow_free(void *ptr);
   u_int32_t ndpi_get_tot_allocated_memory(void);
-  
+
   /**
    * Search the first occurrence of substring -find- in -s-
    * The search is limited to the first -slen- characters of the string
@@ -756,10 +756,11 @@ extern "C" {
    *
    * @par     ndpi_mod = the detection module
    * @par     path     = the path of the file
+   * @par     user_data = pointer to some user data value
    * @return  0 if the file is loaded correctly;
    *          -1 else
    */
-  int ndpi_load_categories_file(struct ndpi_detection_module_struct *ndpi_str, const char* path);
+  int ndpi_load_categories_file(struct ndpi_detection_module_struct *ndpi_str, const char* path, void *user_data);
 
   /**
    * Read a file and load the list of risky domains
@@ -922,12 +923,16 @@ extern "C" {
   int ndpi_match_string(void *_automa, char *string_to_match);
 
   int ndpi_load_ip_category(struct ndpi_detection_module_struct *ndpi_struct,
-			    const char *ip_address_and_mask, ndpi_protocol_category_t category);
+			    const char *ip_address_and_mask, ndpi_protocol_category_t category,
+			    void *user_data);
   int ndpi_load_hostname_category(struct ndpi_detection_module_struct *ndpi_struct,
 				  const char *name_to_add, ndpi_protocol_category_t category);
   int ndpi_load_category(struct ndpi_detection_module_struct *ndpi_struct,
-			 const char *ip_or_name, ndpi_protocol_category_t category);
+			 const char *ip_or_name, ndpi_protocol_category_t category,
+			 void *user_data);
   int ndpi_enable_loaded_categories(struct ndpi_detection_module_struct *ndpi_struct);
+  void* ndpi_find_ipv4_category_userdata(struct ndpi_detection_module_struct *ndpi_str,
+					 u_int32_t saddr);
   int ndpi_fill_ip_protocol_category(struct ndpi_detection_module_struct *ndpi_struct,
 				     u_int32_t saddr,
 				     u_int32_t daddr,
@@ -986,12 +991,12 @@ extern "C" {
    *
    * @par    ndpi_struct  = the struct created for the protocol detection
    * @par    days         = the number of days threshold for emitting the alert
-   *                        
+   *
    */
   void ndpi_set_tls_cert_expire_days(struct ndpi_detection_module_struct *ndpi_str,
 				     u_int8_t days);
 
-  
+
   /* Utility functions to set ndpi malloc/free/print wrappers */
   void set_ndpi_malloc(void* (*__ndpi_malloc)(size_t size));
   void set_ndpi_free(void  (*__ndpi_free)(void *ptr));
@@ -1090,7 +1095,7 @@ extern "C" {
   u_int64_t ndpi_ntohll(u_int64_t v);
   u_int8_t ndpi_is_valid_protoId(u_int16_t protoId);
   u_int8_t ndpi_is_encrypted_proto(struct ndpi_detection_module_struct *ndpi_str, ndpi_protocol proto);
-  
+
   /* DGA */
   int ndpi_check_dga_name(struct ndpi_detection_module_struct *ndpi_str,
 			  struct ndpi_flow_struct *flow,
@@ -1589,13 +1594,13 @@ extern "C" {
   int ndpi_ses_init(struct ndpi_ses_struct *ses, double alpha, float significance);
   int ndpi_ses_add_value(struct ndpi_ses_struct *ses, const u_int64_t _value, double *forecast, double *confidence_band);
   void ndpi_ses_fitting(double *values, u_int32_t num_values, float *ret_alpha);
-  
+
   /* ******************************* */
 
   int ndpi_des_init(struct ndpi_des_struct *des, double alpha, double beta, float significance);
   int ndpi_des_add_value(struct ndpi_des_struct *des, const u_int64_t _value, double *forecast, double *confidence_band);
   void ndpi_des_fitting(double *values, u_int32_t num_values, float *ret_alpha, float *ret_beta);
-  
+
   /* ******************************* */
 
   int   ndpi_jitter_init(struct ndpi_jitter_struct *hw, u_int16_t num_periods);
@@ -1662,7 +1667,7 @@ extern "C" {
   /* ******************************* */
 
   /*
-   * Finds outliers using Z-score  
+   * Finds outliers using Z-score
    * Z-Score = (Value - Mean) / StdDev
    *
    * @par    values      = pointer to the individual values to be analyzed [in]
@@ -1698,7 +1703,7 @@ extern "C" {
   /* ******************************* */
 
   char* ndpi_get_flow_name(struct ndpi_flow_struct *flow);
-  
+
   /* ******************************* */
 
   ndpi_bitmap* ndpi_bitmap_alloc(void);
@@ -1714,13 +1719,13 @@ extern "C" {
 
   void ndpi_bitmap_and(ndpi_bitmap* a, ndpi_bitmap* b_and);
   void ndpi_bitmap_or(ndpi_bitmap* a, ndpi_bitmap* b_or);
-  
+
   ndpi_bitmap_iterator* ndpi_bitmap_iterator_alloc(ndpi_bitmap* b);
   void ndpi_bitmap_iterator_free(ndpi_bitmap* b);
   bool ndpi_bitmap_iterator_next(ndpi_bitmap_iterator* i, uint32_t *value);
 
   /* ******************************* */
-  
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/include/ndpi_patricia_typedefs.h b/src/include/ndpi_patricia_typedefs.h
index 3b190b2fa..c9f7fddac 100644
--- a/src/include/ndpi_patricia_typedefs.h
+++ b/src/include/ndpi_patricia_typedefs.h
@@ -73,7 +73,7 @@ union ndpi_patricia_node_value_t {
     } uv32;
     
     u_int64_t uv64;
-
+    
     void *user_data;
   } u;
 };
@@ -95,6 +95,7 @@ typedef struct _ndpi_patricia_node_t {
   struct _ndpi_patricia_node_t *l, *r;	/* left and right children */
   struct _ndpi_patricia_node_t *parent;/* may be used */
   void *data;			/* pointer to data */
+  void *custom_user_data;       /* pointer to custom userdata */
   union ndpi_patricia_node_value_t value;
 } ndpi_patricia_node_t;
 
diff --git a/src/include/ndpi_typedefs.h b/src/include/ndpi_typedefs.h
index 109b3a0cc..124347fb0 100644
--- a/src/include/ndpi_typedefs.h
+++ b/src/include/ndpi_typedefs.h
@@ -973,9 +973,10 @@ typedef struct ndpi_proto {
   */
   u_int16_t master_protocol /* e.g. HTTP */, app_protocol /* e.g. FaceBook */;
   ndpi_protocol_category_t category;
+  void *custom_category_userdata;
 } ndpi_protocol;
 
-#define NDPI_PROTOCOL_NULL { NDPI_PROTOCOL_UNKNOWN , NDPI_PROTOCOL_UNKNOWN , NDPI_PROTOCOL_CATEGORY_UNSPECIFIED }
+#define NDPI_PROTOCOL_NULL { NDPI_PROTOCOL_UNKNOWN , NDPI_PROTOCOL_UNKNOWN , NDPI_PROTOCOL_CATEGORY_UNSPECIFIED, NULL }
 
 #define NUM_CUSTOM_CATEGORIES      5
 #define CUSTOM_CATEGORY_LABEL_LEN 32
diff --git a/src/lib/ndpi_main.c b/src/lib/ndpi_main.c
index a9e34f30b..63a5cec89 100644
--- a/src/lib/ndpi_main.c
+++ b/src/lib/ndpi_main.c
@@ -3458,7 +3458,8 @@ int ndpi_handle_rule(struct ndpi_detection_module_struct *ndpi_str, char *rule,
  *  - host and category are separated by a single TAB
  *  - empty lines or lines starting with # are ignored
  */
-int ndpi_load_categories_file(struct ndpi_detection_module_struct *ndpi_str, const char *path) {
+int ndpi_load_categories_file(struct ndpi_detection_module_struct *ndpi_str,
+			      const char *path, void *user_data) {
   char buffer[512], *line, *name, *category, *saveptr;
   FILE *fd;
   int len, num = 0;
@@ -3488,7 +3489,9 @@ int ndpi_load_categories_file(struct ndpi_detection_module_struct *ndpi_str, con
       category = strtok_r(NULL, "\t", &saveptr);
 
       if(category) {
-	int rc = ndpi_load_category(ndpi_str, name, (ndpi_protocol_category_t) atoi(category));
+	int rc = ndpi_load_category(ndpi_str, name,
+				    (ndpi_protocol_category_t) atoi(category),
+				    user_data);
 
 	if(rc >= 0)
 	  num++;
@@ -5368,7 +5371,7 @@ static void ndpi_add_connection_as_zoom(struct ndpi_detection_module_struct *ndp
 
 ndpi_protocol ndpi_detection_giveup(struct ndpi_detection_module_struct *ndpi_str, struct ndpi_flow_struct *flow,
 				    u_int8_t enable_guess, u_int8_t *protocol_was_guessed) {
-  ndpi_protocol ret = {NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED};
+  ndpi_protocol ret = {NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED, NULL};
   u_int16_t guessed_protocol_id = NDPI_PROTOCOL_UNKNOWN, guessed_host_protocol_id = NDPI_PROTOCOL_UNKNOWN;
   
   /* *** We can't access ndpi_str->packet from this function!! *** */
@@ -5563,8 +5566,10 @@ void ndpi_process_extra_packet(struct ndpi_detection_module_struct *ndpi_str, st
 
 /* ********************************************************************************* */
 
-int ndpi_load_ip_category(struct ndpi_detection_module_struct *ndpi_str, const char *ip_address_and_mask,
-			  ndpi_protocol_category_t category) {
+int ndpi_load_ip_category(struct ndpi_detection_module_struct *ndpi_str,
+			  const char *ip_address_and_mask,
+			  ndpi_protocol_category_t category,
+			  void *user_data) {
   ndpi_patricia_node_t *node;
   struct in_addr pin;
   int bits = 32;
@@ -5589,15 +5594,17 @@ int ndpi_load_ip_category(struct ndpi_detection_module_struct *ndpi_str, const c
 
   if((node = add_to_ptree(ndpi_str->custom_categories.ipAddresses_shadow, AF_INET, &pin, bits)) != NULL) {
     node->value.u.uv32.user_value = (u_int16_t)category, node->value.u.uv32.additional_user_value = 0;
+    node->custom_user_data = user_data;
   }
 
+  
   return(0);
 }
 
-
 /* ********************************************************************************* */
 
-int ndpi_load_hostname_category(struct ndpi_detection_module_struct *ndpi_str, const char *name_to_add,
+int ndpi_load_hostname_category(struct ndpi_detection_module_struct *ndpi_str,
+				const char *name_to_add,
 				ndpi_protocol_category_t category) {
 
   if(ndpi_str->custom_categories.hostnames_shadow.ac_automa == NULL)
@@ -5606,7 +5613,8 @@ int ndpi_load_hostname_category(struct ndpi_detection_module_struct *ndpi_str, c
   if(name_to_add == NULL)
     return(-1);
 
-  return ndpi_string_to_automa(ndpi_str,(AC_AUTOMATA_t *)ndpi_str->custom_categories.hostnames_shadow.ac_automa,
+  return ndpi_string_to_automa(ndpi_str,
+			       (AC_AUTOMATA_t *)ndpi_str->custom_categories.hostnames_shadow.ac_automa,
 			       name_to_add,category,category, 0, 0, 1); /* at_end */
 }
 
@@ -5614,14 +5622,20 @@ int ndpi_load_hostname_category(struct ndpi_detection_module_struct *ndpi_str, c
 
 /* Loads an IP or name category */
 int ndpi_load_category(struct ndpi_detection_module_struct *ndpi_struct, const char *ip_or_name,
-		       ndpi_protocol_category_t category) {
+		       ndpi_protocol_category_t category, void *user_data) {
   int rv;
 
   /* Try to load as IP address first */
-  rv = ndpi_load_ip_category(ndpi_struct, ip_or_name, category);
+  rv = ndpi_load_ip_category(ndpi_struct, ip_or_name, category, user_data);
 
   if(rv < 0) {
-    /* IP load failed, load as hostname */
+    /* 
+       IP load failed, load as hostname 
+
+       NOTE: 
+       we cannot add user_data here as with Aho-Corasick this
+       information would not be used
+    */
     rv = ndpi_load_hostname_category(ndpi_struct, ip_or_name, category);
   }
 
@@ -5632,10 +5646,12 @@ int ndpi_load_category(struct ndpi_detection_module_struct *ndpi_struct, const c
 
 int ndpi_enable_loaded_categories(struct ndpi_detection_module_struct *ndpi_str) {
   int i;
-
+  static char *built_in = "built-in";
+  
   /* First add the nDPI known categories matches */
   for(i = 0; category_match[i].string_to_match != NULL; i++)
-    ndpi_load_category(ndpi_str, category_match[i].string_to_match, category_match[i].protocol_category);
+    ndpi_load_category(ndpi_str, category_match[i].string_to_match,
+		       category_match[i].protocol_category, built_in);
 
   /* Free */
   ac_automata_release((AC_AUTOMATA_t *) ndpi_str->custom_categories.hostnames.ac_automa,
@@ -5667,8 +5683,33 @@ int ndpi_enable_loaded_categories(struct ndpi_detection_module_struct *ndpi_str)
 
 /* ********************************************************************************* */
 
-int ndpi_fill_ip_protocol_category(struct ndpi_detection_module_struct *ndpi_str, u_int32_t saddr, u_int32_t daddr,
+/* NOTE u_int32_t is represented in network byte order */
+void* ndpi_find_ipv4_category_userdata(struct ndpi_detection_module_struct *ndpi_str,
+				       u_int32_t saddr) {
+  ndpi_patricia_node_t *node;
+  
+  if(saddr == 0)
+    node = NULL;
+  else {
+    ndpi_prefix_t prefix;
+    
+    ndpi_fill_prefix_v4(&prefix, (struct in_addr *) &saddr, 32,
+			((ndpi_patricia_tree_t *) ndpi_str->protocols_ptree)->maxbits);
+    node = ndpi_patricia_search_best(ndpi_str->custom_categories.ipAddresses, &prefix);
+  }
+
+  return(node ? node->custom_user_data : NULL);
+}
+
+/* ********************************************************************************* */
+
+/* NOTE u_int32_t is represented in network byte order */
+int ndpi_fill_ip_protocol_category(struct ndpi_detection_module_struct *ndpi_str,
+				   u_int32_t saddr, u_int32_t daddr,
 				   ndpi_protocol *ret) {
+
+  ret->custom_category_userdata = NULL;
+  
   if(ndpi_str->custom_categories.categories_loaded) {
     ndpi_prefix_t prefix;
     ndpi_patricia_node_t *node;
@@ -5692,7 +5733,7 @@ int ndpi_fill_ip_protocol_category(struct ndpi_detection_module_struct *ndpi_str
 
     if(node) {
       ret->category = (ndpi_protocol_category_t) node->value.u.uv32.user_value;
-
+      ret->custom_category_userdata = node->custom_user_data;
       return(1);
     }
   }
@@ -5896,7 +5937,7 @@ ndpi_protocol ndpi_detection_process_packet(struct ndpi_detection_module_struct
   struct ndpi_packet_struct *packet = &ndpi_str->packet;
   NDPI_SELECTION_BITMASK_PROTOCOL_SIZE ndpi_selection_packet;
   u_int32_t num_calls = 0;
-  ndpi_protocol ret = { flow->detected_protocol_stack[1], flow->detected_protocol_stack[0], flow->category };
+  ndpi_protocol ret = { flow->detected_protocol_stack[1], flow->detected_protocol_stack[0], flow->category, NULL };
 
   if(ndpi_str->ndpi_log_level >= NDPI_LOG_TRACE)
     NDPI_LOG(flow ? flow->detected_protocol_stack[0] : NDPI_PROTOCOL_UNKNOWN, ndpi_str, NDPI_LOG_TRACE,
@@ -6721,8 +6762,10 @@ ndpi_protocol_category_t ndpi_get_flow_category(struct ndpi_detection_module_str
   return(flow->category);
 }
 
+/* ********************************************************************************* */
+
 void ndpi_get_flow_ndpi_proto(struct ndpi_detection_module_struct *ndpi_str, struct ndpi_flow_struct *flow,
-                struct ndpi_proto * ndpi_proto)
+			      struct ndpi_proto * ndpi_proto)
 {
   ndpi_proto->master_protocol = ndpi_get_flow_masterprotocol(ndpi_str, flow);
   ndpi_proto->app_protocol = ndpi_get_flow_appprotocol(ndpi_str, flow);
@@ -6963,7 +7006,7 @@ ndpi_protocol ndpi_guess_undetected_protocol(struct ndpi_detection_module_struct
 					     u_int32_t dhost /* host byte order */, u_int16_t dport) {
   u_int32_t rc;
   struct in_addr addr;
-  ndpi_protocol ret = {NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED};
+  ndpi_protocol ret = {NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_UNKNOWN, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED, NULL};
   u_int8_t user_defined_proto;
 
 #ifdef BITTORRENT_CACHE_DEBUG
diff --git a/src/lib/protocols/tls.c b/src/lib/protocols/tls.c
index c5142abde..fbb21ef00 100644
--- a/src/lib/protocols/tls.c
+++ b/src/lib/protocols/tls.c
@@ -303,7 +303,7 @@ static void checkTLSSubprotocol(struct ndpi_detection_module_struct *ndpi_struct
 
       if(ndpi_lru_find_cache(ndpi_struct->tls_cert_cache, key,
 			     &cached_proto, 0 /* Don't remove it as it can be used for other connections */)) {
-	ndpi_protocol ret = { NDPI_PROTOCOL_TLS, cached_proto, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED};
+	ndpi_protocol ret = { NDPI_PROTOCOL_TLS, cached_proto, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED, NULL};
 
 	ndpi_set_detected_protocol(ndpi_struct, flow, cached_proto, NDPI_PROTOCOL_TLS, NDPI_CONFIDENCE_DPI_CACHE);
 	flow->category = ndpi_get_proto_category(ndpi_struct, ret);
@@ -638,7 +638,7 @@ static void processCertificateElements(struct ndpi_detection_module_struct *ndpi
       if(rc == 0) {
 	/* Match found */
 	u_int16_t proto_id = (u_int16_t)val;
-	ndpi_protocol ret = { NDPI_PROTOCOL_TLS, proto_id, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED};
+	ndpi_protocol ret = { NDPI_PROTOCOL_TLS, proto_id, NDPI_PROTOCOL_CATEGORY_UNSPECIFIED, NULL};
 
 	ndpi_set_detected_protocol(ndpi_struct, flow, proto_id, NDPI_PROTOCOL_TLS, NDPI_CONFIDENCE_DPI);
 	flow->category = ndpi_get_proto_category(ndpi_struct, ret);