From: Yuri Schaeffer <yuri@nlnetlabs.nl>
Date: Mon, 14 Oct 2013 13:07:36 +0000 (+0000)
Subject: enforce limits on the addrtree structure
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=0dbee5840ca644af6affa7e600e0971f7277e60c;p=thirdparty%2Funbound.git

enforce limits on the addrtree structure


git-svn-id: file:///svn/unbound/branches/edns-subnet@2982 be551aaa-1e26-0410-a405-d3ace91eadb9
---

diff --git a/edns-subnet/addrtree.c b/edns-subnet/addrtree.c
index 95ea61f15..a8e1a517a 100644
--- a/edns-subnet/addrtree.c
+++ b/edns-subnet/addrtree.c
@@ -75,7 +75,8 @@ node_create(struct addrtree *tree, void* elem, addrlen_t scope,
 }
 
 struct addrtree* addrtree_create(addrlen_t max_depth, 
-	void (*delfunc)(void *, void *), size_t (*sizefunc)(void *), void *env)
+	void (*delfunc)(void *, void *), size_t (*sizefunc)(void *), 
+	void *env, unsigned int max_node_count)
 {
 	struct addrtree* tree;
 	log_assert(delfunc != NULL);
@@ -95,46 +96,35 @@ struct addrtree* addrtree_create(addrlen_t max_depth,
 	tree->sizefunc = sizefunc;
 	tree->env = env;
 	tree->node_count = 0;
+	tree->max_node_count = max_node_count;
 	return tree;
 }
 
-/** Add node to LRU list as most recently used. */
-void lru_push(struct addrtree *tree, struct addrnode *node)
-{
-	if (!tree->first) {
-		tree->first = node;
-		tree->last = node;
-		node->prev = NULL;
-	} else {
-		tree->last->next = node;
-		node->prev = tree->last->next;
-		tree->last = node;
-	}
-	node->next = NULL;
-}
-
-/** Remove least recently used from LRU list */
-struct addrnode* lru_popfirst(struct addrtree *tree)
+/** 
+ * Scrub a node clean of elem
+ * @param tree: Tree the node lives in.
+ * @param node: Node to be cleaned
+ */
+static void
+clean_node(struct addrtree* tree, struct addrnode* node)
 {
-	struct addrnode* pop;
-	pop = tree->first;
-	log_assert(first != NULL);
-	if (!pop->next) {
-		tree->first = NULL;
-		tree->last = NULL;
-	} else {
-		tree->first = pop->next;
-		tree->first->prev = NULL;
-	}
-	return pop;
+	if (!node->elem) return;
+	tree->delfunc(tree->env, node->elem);
+	node->elem = NULL;
 }
 
 /** Remove specified node from LRU list */
 void lru_pop(struct addrtree *tree, struct addrnode *node)
 {
 	if (node == tree->first) {
-		(void)lru_popfirst(tree);
-	} else if (node == tree->last) {
+		if (!node->next) { /* it is the last as well */
+			tree->first = NULL;
+			tree->last = NULL;
+		} else {
+			tree->first = node->next;
+			tree->first->prev = NULL;
+		}
+	} else if (node == tree->last) { /* but not the first */
 		tree->last = node->prev;
 		tree->last->next = NULL;
 	} else {
@@ -143,41 +133,26 @@ void lru_pop(struct addrtree *tree, struct addrnode *node)
 	}
 }
 
-/** Move node to the end of LRU list */
-void lru_update(struct addrtree *tree, struct addrnode *node)
-{
-	lru_pop(tree, node);
-	lru_push(tree, node);
-}
-
-/** Size in bytes of this data structure */
-size_t addrtree_size(const struct addrtree* tree)
+/** Add node to LRU list as most recently used. */
+void lru_push(struct addrtree *tree, struct addrnode *node)
 {
-	struct addrnode* n;
-	size_t s;
-
-	if (!tree) return 0;
-	s = sizeof (struct addrnode); /* root always exists but not in LRU */
-	if (tree->root->elem) s += tree->sizefunc(tree->root->elem);
-	for(n = tree->first; n; n = n->next) {
-		s += sizeof (struct addredge) + sizeof (struct addrnode);
-		s += tree->sizefunc(n->elem);
+	if (!tree->first) {
+		tree->first = node;
+		tree->last = node;
+		node->prev = NULL;
+	} else {
+		tree->last->next = node;
+		node->prev = tree->last;
+		tree->last = node;
 	}
-	return s;
+	node->next = NULL;
 }
 
-/** 
- * Scrub a node clean of elem
- * @param tree: Tree the node lives in.
- * @param node: Node to be cleaned
- */
-static void
-clean_node(struct addrtree* tree, struct addrnode* node)
+/** Move node to the end of LRU list */
+void lru_update(struct addrtree *tree, struct addrnode *node)
 {
-	tree->node_count--;
-	if (!node->elem) return;
-	tree->delfunc(tree->env, node->elem);
-	node->elem = NULL;
+	lru_pop(tree, node);
+	lru_push(tree, node);
 }
 
 /** 
@@ -186,7 +161,7 @@ clean_node(struct addrtree* tree, struct addrnode* node)
  * @param tree: Tree the node lives in.
  * @param node: Node to be freed
  */
-static void
+static int
 purge_node(struct addrtree *tree, struct addrnode *node)
 {
 	struct addredge *parent_edge, *child_edge = NULL;
@@ -196,6 +171,7 @@ purge_node(struct addrtree *tree, struct addrnode *node)
 	clean_node(tree, node);
 	parent_edge = node->parent_edge;
 	if (keep || !parent_edge) return;
+	tree->node_count--;
 	index = parent_edge->parent_index;
 	child_edge = node->edge[!node->edge[0]];
 	if (child_edge) {
@@ -205,9 +181,54 @@ purge_node(struct addrtree *tree, struct addrnode *node)
 	parent_edge->parent_node->edge[index] = child_edge;
 	free(parent_edge->str);
 	free(parent_edge);
+	lru_pop(tree, node);
 	free(node);
 }
 
+/** If a limit is set remove old nodes while above that limit */
+void lru_cleanup(struct addrtree *tree)
+{
+	struct addrnode *n, *p;
+	int children;
+	if (tree->max_node_count == 0) return;
+	while (tree->node_count > tree->max_node_count) {
+		n = tree->first;
+		if (!n) break;
+		children = (n->edge[0] != NULL) + (n->edge[1] != NULL);
+		/** Don't remove this node, it is either the root or we can't
+		 * do without it because it has 2 children */
+		if (children == 2 || !n->parent_edge) {
+			lru_update(tree, n);
+			continue;
+		}
+		p = n->parent_edge->parent_node;
+		purge_node(tree, n);
+		/** Since we removed n, n's parent p is eligible for deletion
+		 * if it is not the root node, caries no data and has only 1
+		 * child */
+		children = (p->edge[0] != NULL) + (p->edge[1] != NULL);
+		if (!p->elem && children == 1 && p->parent_edge) {
+			purge_node(tree, p);
+		}
+	}
+}
+
+/** Size in bytes of this data structure */
+size_t addrtree_size(const struct addrtree* tree)
+{
+	struct addrnode* n;
+	size_t s;
+
+	if (!tree) return 0;
+	s = sizeof (struct addrnode); /* root always exists but not in LRU */
+	if (tree->root->elem) s += tree->sizefunc(tree->root->elem);
+	for(n = tree->first; n; n = n->next) {
+		s += sizeof (struct addredge) + sizeof (struct addrnode);
+		s += tree->sizefunc(n->elem);
+	}
+	return s;
+}
+
 /** 
  * Free complete addrtree structure
  * @param tree: Tree to free
@@ -218,7 +239,7 @@ void addrtree_delete(struct addrtree* tree)
 	if (!tree) return;
 	clean_node(tree, tree->root);
 	free(tree->root);
-	while(n = tree->first) {
+	while((n = tree->first)) {
 		tree->first = n->next;
 		clean_node(tree, n);
 		free(n->parent_edge->str);
@@ -326,7 +347,6 @@ addrtree_insert(struct addrtree *tree, const addrkey_t *addr,
 			if (!edge->node->elem || edge->node->ttl >= now)
 				break;
 			purge_node(tree, edge->node);
-			lru_pop(tree, edge->node);
 			edge = node->edge[index];
 		}
 		/* Case 2: New leafnode */
@@ -335,10 +355,12 @@ addrtree_insert(struct addrtree *tree, const addrkey_t *addr,
 			if (!newnode) return;
 			if (!edge_create(newnode, addr, sourcemask, node, index)) {
 				clean_node(tree, newnode);
+				tree->node_count--;
 				free(newnode);
 				return;
 			}
 			lru_push(tree, newnode);
+			lru_cleanup(tree);
 			return;
 		}
 		/* Case 3: Traverse edge */
@@ -357,11 +379,12 @@ addrtree_insert(struct addrtree *tree, const addrkey_t *addr,
 			return;
 		if (!edge_create(newnode, addr, common, node, index)) {
 			clean_node(tree, newnode);
+			tree->node_count--;
 			free(newnode);
 			return;
 		}
 		lru_push(tree, newnode);
-		/** connect existing child to our new node */
+		/* connect existing child to our new node */
 		index = getbit(edge->str, edge->len, common);
 		newnode->edge[index] = edge;
 		edge->parent_node = newnode;
@@ -378,11 +401,13 @@ addrtree_insert(struct addrtree *tree, const addrkey_t *addr,
 			newnode = node_create(tree, elem, scope, ttl);
 			if (!edge_create(newnode, addr, sourcemask, node, index^1)) {
 				clean_node(tree, newnode);
+				tree->node_count--;
 				free(newnode);
 				return;
 			}
 			lru_push(tree, newnode);
 		}
+		lru_cleanup(tree);
 		return;
 	}
 }
diff --git a/edns-subnet/addrtree.h b/edns-subnet/addrtree.h
index 8a3969b0e..1a376d0eb 100644
--- a/edns-subnet/addrtree.h
+++ b/edns-subnet/addrtree.h
@@ -41,6 +41,9 @@ struct addrtree {
 	/** Number of elements in the tree (not always equal to number of 
 	 * nodes) */
 	unsigned int node_count;
+	/** Maximum number of allowed nodes, will be enforced by LRU list.
+	 * Excluding the root node, 0 for unlimited */
+	unsigned int max_node_count;
 	/** Maximum prefix length we are willing to cache. */
 	addrlen_t max_depth;
 	/** External function to delete elem. Called as 
@@ -102,7 +105,7 @@ size_t addrtree_size(const struct addrtree* tree);
  */
 struct addrtree* 
 addrtree_create(addrlen_t max_depth, void (*delfunc)(void *, void *), 
-	size_t (*sizefunc)(void *), void *env);
+	size_t (*sizefunc)(void *), void *env, unsigned int max_node_count);
 
 /** 
  * Free tree and all nodes below
diff --git a/edns-subnet/subnetmod.c b/edns-subnet/subnetmod.c
index 9b342d0f4..2ea394b09 100644
--- a/edns-subnet/subnetmod.c
+++ b/edns-subnet/subnetmod.c
@@ -130,11 +130,11 @@ get_tree(struct subnet_msg_cache_data *data, struct edns_data *edns,
 	struct addrtree *tree;
 	if (edns->subnet_addr_fam == EDNSSUBNET_ADDRFAM_IP4) {
 		if (!data->tree4)
-			data->tree4 = addrtree_create(EDNSSUBNET_MAX_SUBNET_IP4, &delfunc, &sizefunc, env);
+			data->tree4 = addrtree_create(EDNSSUBNET_MAX_SUBNET_IP4, &delfunc, &sizefunc, env, 0);
 		tree = data->tree4;
 	} else {
 		if (!data->tree6)
-			data->tree6 = addrtree_create(EDNSSUBNET_MAX_SUBNET_IP6, &delfunc, &sizefunc, env);
+			data->tree6 = addrtree_create(EDNSSUBNET_MAX_SUBNET_IP6, &delfunc, &sizefunc, env, 0);
 		tree = data->tree6;
 	}
 	return tree;
diff --git a/testcode/unitvandergaast.c b/testcode/unitvandergaast.c
index 875467eae..85428939d 100644
--- a/testcode/unitvandergaast.c
+++ b/testcode/unitvandergaast.c
@@ -65,12 +65,12 @@
 	{
 		struct addredge* edge;
 		int i, s, byte;
-		if (indent == 0) printf("-----Tree-----");
+		if (indent == 0) printf("-----Tree-----\n");
 		if (indent > maxdepth) {
 			printf("\n");
 			return;
 		}
-		printf("[node elem:%d]\n", node->elem != NULL);
+		printf("[node elem:%d] (%d)\n", node->elem != NULL, node);
 		for (i = 0; i<2; i++) {
 			if (node->edge[i]) {
 				for (s = 0; s < indent; s++) printf(" ");
@@ -165,7 +165,7 @@ static void consistency_test(void)
 	unit_show_func("edns-subnet/addrtree.h", "Tree consistency check");
 	srand(9195); /* just some value for reproducibility */
 
-	t = addrtree_create(100, &elemfree, NULL, &env);
+	t = addrtree_create(100, &elemfree, NULL, &env, 0);
 	count = t->node_count;
 	unit_assert(count == 0);
 	for (i = 0; i < 1000; i++) {
@@ -181,12 +181,24 @@ static void consistency_test(void)
 	}
 	addrtree_delete(t);
 	unit_show_func("edns-subnet/addrtree.h", "Tree consistency with purge");
-	t = addrtree_create(8, &elemfree, NULL, &env);
+	t = addrtree_create(8, &elemfree, NULL, &env, 0);
+	unit_assert(t->node_count == 0);
+	for (i = 0; i < 1000; i++) {
+		l = randomkey(&k, 128);
+		elem = (struct reply_info *) calloc(1, sizeof(struct reply_info));
+		addrtree_insert(t, k, l, 64, elem, i + 10, i);
+		free(k);
+		unit_assert( !addrtree_inconsistent(t) );
+	}
+	addrtree_delete(t);
+	unit_show_func("edns-subnet/addrtree.h", "Tree consistency with limit");
+	t = addrtree_create(8, &elemfree, NULL, &env, 27);
 	unit_assert(t->node_count == 0);
 	for (i = 0; i < 1000; i++) {
 		l = randomkey(&k, 128);
 		elem = (struct reply_info *) calloc(1, sizeof(struct reply_info));
 		addrtree_insert(t, k, l, 64, elem, i + 10, i);
+		unit_assert( t->node_count <= 27);
 		free(k);
 		unit_assert( !addrtree_inconsistent(t) );
 	}