From 51762a45b39f906fe72db2d3d46c9890cb01bbd8 Mon Sep 17 00:00:00 2001
From: Ondrej Zajicek <santiago@crfreenet.org>
Date: Sat, 21 Feb 2015 14:05:20 +0100
Subject: [PATCH] Allows user data attached to f_trie_node structure.

Thanks to Alexander Chernikov for the patch.
---
 filter/config.Y |  2 +-
 filter/filter.h |  7 ++++---
 filter/trie.c   | 43 +++++++++++++++++++++++++------------------
 nest/rt-table.c |  4 ++--
 4 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/filter/config.Y b/filter/config.Y
index e50e75ca..7eb2c0a3 100644
--- a/filter/config.Y
+++ b/filter/config.Y
@@ -592,7 +592,7 @@ fprefix:
  ;
 
 fprefix_set:
-   fprefix { $$ = f_new_trie(cfg_mem); trie_add_fprefix($$, &($1.val.px)); }
+   fprefix { $$ = f_new_trie(cfg_mem, sizeof(struct f_trie_node)); trie_add_fprefix($$, &($1.val.px)); }
  | fprefix_set ',' fprefix { $$ = $1; trie_add_fprefix($$, &($3.val.px)); }
  ;
 
diff --git a/filter/filter.h b/filter/filter.h
index 3a6b66d9..2b2d23c2 100644
--- a/filter/filter.h
+++ b/filter/filter.h
@@ -80,8 +80,8 @@ struct f_tree *find_tree(struct f_tree *t, struct f_val val);
 int same_tree(struct f_tree *t1, struct f_tree *t2);
 void tree_format(struct f_tree *t, buffer *buf);
 
-struct f_trie *f_new_trie(linpool *lp);
-void trie_add_prefix(struct f_trie *t, ip_addr px, int plen, int l, int h);
+struct f_trie *f_new_trie(linpool *lp, uint node_size);
+void *trie_add_prefix(struct f_trie *t, ip_addr px, int plen, int l, int h);
 int trie_match_prefix(struct f_trie *t, ip_addr px, int plen);
 int trie_same(struct f_trie *t1, struct f_trie *t2);
 void trie_format(struct f_trie *t, buffer *buf);
@@ -204,7 +204,8 @@ struct f_trie
 {
   linpool *lp;
   int zero;
-  struct f_trie_node root;
+  uint node_size;
+  struct f_trie_node root[0];		/* Root trie node follows */
 };
 
 #define NEW_F_VAL struct f_val * val; val = cfg_alloc(sizeof(struct f_val));
diff --git a/filter/trie.c b/filter/trie.c
index 217d72c3..8af9015e 100644
--- a/filter/trie.c
+++ b/filter/trie.c
@@ -21,7 +21,7 @@
  *
  * We use a bitmask (&accept) to represent accepted prefix lengths
  * at a node. As there are 33 prefix lengths (0..32 for IPv4), but
- * there is just one prefix of zero length in the whole trie so we 
+ * there is just one prefix of zero length in the whole trie so we
  * have &zero flag in &f_trie (indicating whether the trie accepts
  * prefix 0.0.0.0/0) as a special case, and &accept bitmask
  * represents accepted prefix lengths from 1 to 32.
@@ -75,23 +75,24 @@
 #include "filter/filter.h"
 
 /**
- * f_new_trie
- *
- * Allocates and returns a new empty trie.
+ * f_new_trie - allocates and returns a new empty trie
+ * @lp: linear pool to allocate items from
+ * @node_size: node size to be used (&f_trie_node and user data)
  */
 struct f_trie *
-f_new_trie(linpool *lp)
+f_new_trie(linpool *lp, uint node_size)
 {
   struct f_trie * ret;
-  ret = lp_allocz(lp, sizeof(struct f_trie));
+  ret = lp_allocz(lp, sizeof(struct f_trie) + node_size);
   ret->lp = lp;
+  ret->node_size = node_size;
   return ret;
 }
 
 static inline struct f_trie_node *
 new_node(struct f_trie *t, int plen, ip_addr paddr, ip_addr pmask, ip_addr amask)
 {
-  struct f_trie_node *n = lp_allocz(t->lp, sizeof(struct f_trie_node));
+  struct f_trie_node *n = lp_allocz(t->lp, t->node_size);
   n->plen = plen;
   n->addr = paddr;
   n->mask = pmask;
@@ -110,15 +111,19 @@ attach_node(struct f_trie_node *parent, struct f_trie_node *child)
  * @t: trie to add to
  * @px: prefix address
  * @plen: prefix length
- * @l: prefix lower bound 
+ * @l: prefix lower bound
  * @h: prefix upper bound
  *
  * Adds prefix (prefix pattern) @px/@plen to trie @t.  @l and @h are lower
  * and upper bounds on accepted prefix lengths, both inclusive.
  * 0 <= l, h <= 32 (128 for IPv6).
+ *
+ * Returns a pointer to the allocated node. The function can return a pointer to
+ * an existing node if @px and @plen are the same. If px/plen == 0/0 (or ::/0),
+ * a pointer to the root node is returned.
  */
 
-void
+void *
 trie_add_prefix(struct f_trie *t, ip_addr px, int plen, int l, int h)
 {
   if (l == 0)
@@ -133,7 +138,7 @@ trie_add_prefix(struct f_trie *t, ip_addr px, int plen, int l, int h)
   ip_addr pmask = ipa_mkmask(plen);
   ip_addr paddr = ipa_and(px, pmask);
   struct f_trie_node *o = NULL;
-  struct f_trie_node *n = &t->root;
+  struct f_trie_node *n = t->root;
 
   while(n)
     {
@@ -156,7 +161,7 @@ trie_add_prefix(struct f_trie *t, ip_addr px, int plen, int l, int h)
 	  attach_node(o, b);
 	  attach_node(b, n);
 	  attach_node(b, a);
-	  return;
+	  return a;
 	}
 
       if (plen < n->plen)
@@ -166,14 +171,14 @@ trie_add_prefix(struct f_trie *t, ip_addr px, int plen, int l, int h)
 	  struct f_trie_node *a = new_node(t, plen, paddr, pmask, amask);
 	  attach_node(o, a);
 	  attach_node(a, n);
-	  return;
+	  return a;
 	}
-	
+
       if (plen == n->plen)
 	{
 	  /* We already found added node in trie. Just update accept mask */
 	  n->accept = ipa_or(n->accept, amask);
-	  return;
+	  return n;
 	}
 
       /* Update accept mask part M2 and go deeper */
@@ -187,10 +192,12 @@ trie_add_prefix(struct f_trie *t, ip_addr px, int plen, int l, int h)
   /* We add new tail node 'a' after node 'o' */
   struct f_trie_node *a = new_node(t, plen, paddr, pmask, amask);
   attach_node(o, a);
+
+  return a;
 }
 
 /**
- * trie_match
+ * trie_match_prefix
  * @t: trie
  * @px: prefix address
  * @plen: prefix length
@@ -209,7 +216,7 @@ trie_match_prefix(struct f_trie *t, ip_addr px, int plen)
     return t->zero;
 
   int plentest = plen - 1;
-  struct f_trie_node *n = &t->root;
+  struct f_trie_node *n = t->root;
 
   while(n)
     {
@@ -261,7 +268,7 @@ trie_node_same(struct f_trie_node *t1, struct f_trie_node *t2)
 int
 trie_same(struct f_trie *t1, struct f_trie *t2)
 {
-  return (t1->zero == t2->zero) && trie_node_same(&t1->root, &t2->root);
+  return (t1->zero == t2->zero) && trie_node_same(t1->root, t2->root);
 }
 
 static void
@@ -291,7 +298,7 @@ trie_format(struct f_trie *t, buffer *buf)
 
   if (t->zero)
     buffer_print(buf, "%I/%d", IPA_NONE, 0);
-  trie_node_format(&t->root, buf);
+  trie_node_format(t->root, buf);
 
   /* Undo last separator */
   if (buf->pos[-1] != '[')
diff --git a/nest/rt-table.c b/nest/rt-table.c
index 125f1d18..dbe0c50d 100644
--- a/nest/rt-table.c
+++ b/nest/rt-table.c
@@ -1988,7 +1988,7 @@ rt_init_hostcache(rtable *tab)
   hc->slab = sl_new(rt_table_pool, sizeof(struct hostentry));
 
   hc->lp = lp_new(rt_table_pool, 1008);
-  hc->trie = f_new_trie(hc->lp);
+  hc->trie = f_new_trie(hc->lp, sizeof(struct f_trie_node));
 
   tab->hostcache = hc;
 }
@@ -2136,7 +2136,7 @@ rt_update_hostcache(rtable *tab)
 
   /* Reset the trie */
   lp_flush(hc->lp);
-  hc->trie = f_new_trie(hc->lp);
+  hc->trie = f_new_trie(hc->lp, sizeof(struct f_trie_node));
 
   WALK_LIST_DELSAFE(n, x, hc->hostentries)
     {