s4-dnsserver: Build a dns name tree for correct enumeration

The result of EnumRecords/EnumRecords2 RPC calls, is a list of
dns records that are one level below in the name hierarchy starting
from the search name. This patch builds a tree of names to get
the list of records one level below the search names and correctly
count the number of child records for each of those.

Signed-off-by: Andrew Tridgell <tridge@samba.org>

3 files changed, 274 insertions, 122 deletions

diff --git a/source4/rpc_server/dnsserver/dcerpc_dnsserver.c b/source4/rpc_server/dnsserver/dcerpc_dnsserver.c
index 93f315b857..fd0c97795a 100644
--- a/source4/rpc_server/dnsserver/dcerpc_dnsserver.c
+++ b/source4/rpc_server/dnsserver/dcerpc_dnsserver.c
@@ -1429,7 +1429,7 @@ static WERROR dnsserver_enumerate_root_records(struct dnsserver_state *dsstate,
 	for (i=0; i<res->count; i++) {
 		status = dns_fill_records_array(tmp_ctx, NULL, record_type,
 						select_flag, NULL,
-						res->msgs[i], recs,
+						res->msgs[i], 0, recs,
 						&add_names, &add_count);
 		if (!W_ERROR_IS_OK(status)) {
 			talloc_free(tmp_ctx);
@@ -1457,7 +1457,7 @@ static WERROR dnsserver_enumerate_root_records(struct dnsserver_state *dsstate,
 			}
 			status = dns_fill_records_array(tmp_ctx, NULL, DNS_TYPE_A,
 							select_flag, rname,
-							res->msgs[0], recs,
+							res->msgs[0], 0, recs,
 							NULL, NULL);
 			talloc_free(rname);
 			talloc_free(res);
@@ -1487,16 +1487,16 @@ static WERROR dnsserver_enumerate_records(struct dnsserver_state *dsstate,
 					struct DNS_RPC_RECORDS_ARRAY **buffer)
 {
 	TALLOC_CTX *tmp_ctx;
-	char *name, *branch_name;
+	char *name;
 	const char * const attrs[] = { "name", "dnsRecord", NULL };
 	struct ldb_result *res;
 	struct DNS_RPC_RECORDS_ARRAY *recs;
 	char **add_names = NULL;
-	const char *ptr;
 	char *rname;
 	int add_count = 0;
 	int i, ret, len;
 	WERROR status;
+	struct dns_tree *tree, *base, *node;
 
 	tmp_ctx = talloc_new(mem_ctx);
 	W_ERROR_HAVE_NO_MEMORY(tmp_ctx);
@@ -1530,81 +1530,73 @@ static WERROR dnsserver_enumerate_records(struct dnsserver_state *dsstate,
 	ldb_qsort(res->msgs, res->count, sizeof(struct ldb_message *), name,
 			(ldb_qsort_cmp_fn_t)dns_name_compare);
 
-	/* Add the parent record with blank name */
-	ptr = ldb_msg_find_attr_as_string(res->msgs[0], "name", NULL);
-	if (strcmp(ptr, name) == 0 || strcmp(ptr, "@") == 0) {
-		/* parent record found */
-		if (select_flag & DNS_RPC_VIEW_ONLY_CHILDREN) {
-			status = WERR_OK;
-		} else {
-			status = dns_fill_records_array(tmp_ctx, z, record_type,
-							select_flag, NULL,
-							res->msgs[0], recs,
-							&add_names, &add_count);
-		}
-		i = 1;
+	/* Build a tree of name components from dns name */
+	if (strcmp(name, z->name) == 0) {
+		tree = dns_build_tree(tmp_ctx, "@", res);
 	} else {
-		/* parent record not in the search */
-		if (select_flag & DNS_RPC_VIEW_ONLY_CHILDREN) {
-			status = WERR_OK;
-		} else {
-			status = dns_fill_records_array(tmp_ctx, z, record_type,
-							select_flag, NULL,
-							NULL, recs,
-							&add_names, &add_count);
-		}
-		i = 0;
+		tree = dns_build_tree(tmp_ctx, name, res);
 	}
+	W_ERROR_HAVE_NO_MEMORY_AND_FREE(tree, tmp_ctx);
 
-	if (!W_ERROR_IS_OK(status)) {
-		talloc_free(tmp_ctx);
-		return status;
+	/* Find the parent record in the tree */
+	base = tree;
+	while (base->level != -1) {
+		base = base->children[0];
+	}
+
+	/* Add the parent record with blank name */
+	if (!(select_flag & DNS_RPC_VIEW_ONLY_CHILDREN)) {
+		status = dns_fill_records_array(tmp_ctx, z, record_type,
+						select_flag, NULL,
+						base->data, 0,
+						recs, &add_names, &add_count);
+		if (!W_ERROR_IS_OK(status)) {
+			talloc_free(tmp_ctx);
+			return status;
+		}
 	}
 
 	/* Add all the children records */
 	if (!(select_flag & DNS_RPC_VIEW_NO_CHILDREN)) {
-		for ( ; i<res->count; i++) {
-			char *name2;
-			const char *tmp_str;
-
-			ptr = ldb_msg_find_attr_as_string(res->msgs[i], "name", NULL);
-			name2 = dns_split_node_name(tmp_ctx, ptr, name);
-			tmp_str = strrchr(name2, '.');
-			if (tmp_str == NULL) {
-				branch_name = talloc_strdup(tmp_ctx, name2);
-			} else {
-				/* Skip '.' */
-				branch_name = talloc_strdup(tmp_ctx, &tmp_str[1]);
-			}
-			talloc_free(name2);
+		for (i=0; i<base->num_children; i++) {
+			node = base->children[i];
 
 			status = dns_fill_records_array(tmp_ctx, z, record_type,
-							select_flag, branch_name,
-							res->msgs[i], recs,
-							&add_names, &add_count);
+							select_flag, node->name,
+							node->data, node->num_children,
+							recs, &add_names, &add_count);
 			if (!W_ERROR_IS_OK(status)) {
 				talloc_free(tmp_ctx);
 				return status;
 			}
-
-			talloc_free(branch_name);
 		}
 	}
-	talloc_free(res);
 
+	talloc_free(res);
+	talloc_free(tree);
 	talloc_free(name);
 
 	/* Add any additional records */
 	if (select_flag & DNS_RPC_VIEW_ADDITIONAL_DATA) {
 		for (i=0; i<add_count; i++) {
-			name = dns_split_node_name(tmp_ctx, add_names[i], z->name);
-			ret = ldb_search(dsstate->samdb, tmp_ctx, &res, z->zone_dn,
-					LDB_SCOPE_ONELEVEL, attrs,
-					"(&(objectClass=dnsNode)(name=%s))", name);
-			talloc_free(name);
-			if (ret != LDB_SUCCESS || res->count == 0) {
-				talloc_free(res);
-				continue;
+			struct dnsserver_zone *z2;
+
+			/* Search all the available zones for additional name */
+			for (z2 = dsstate->zones; z2; z2 = z2->next) {
+				name = dns_split_node_name(tmp_ctx, add_names[i], z2->name);
+				ret = ldb_search(dsstate->samdb, tmp_ctx, &res, z2->zone_dn,
+						LDB_SCOPE_ONELEVEL, attrs,
+						"(&(objectClass=dnsNode)(name=%s))", name);
+				talloc_free(name);
+				if (ret != LDB_SUCCESS) {
+					continue;
+				}
+				if (res->count == 1) {
+					break;
+				} else {
+					talloc_free(res);
+					continue;
+				}
 			}
 
 			len = strlen(add_names[i]);
@@ -1615,7 +1607,7 @@ static WERROR dnsserver_enumerate_records(struct dnsserver_state *dsstate,
 			}
 			status = dns_fill_records_array(tmp_ctx, NULL, DNS_TYPE_A,
 							select_flag, rname,
-							res->msgs[0], recs,
+							res->msgs[0], 0, recs,
 							NULL, NULL);
 			talloc_free(rname);
 			talloc_free(res);
diff --git a/source4/rpc_server/dnsserver/dnsdata.c b/source4/rpc_server/dnsserver/dnsdata.c
index e1b7f356ff..2dacda7f49 100644
--- a/source4/rpc_server/dnsserver/dnsdata.c
+++ b/source4/rpc_server/dnsserver/dnsdata.c
@@ -462,6 +462,187 @@ struct dnsp_DnssrvRpcRecord *dns_to_dnsp_copy(TALLOC_CTX *mem_ctx, struct DNS_RP
 }
 
 
+/* Intialize tree with given name as the root */
+static struct dns_tree *dns_tree_init(TALLOC_CTX *mem_ctx, const char *name, void *data)
+{
+	struct dns_tree *tree;
+
+	tree = talloc_zero(mem_ctx, struct dns_tree);
+	if (tree == NULL) {
+		return NULL;
+	}
+
+	tree->name = talloc_strdup(tree, name);
+	if (tree->name == NULL) {
+		talloc_free(tree);
+		return NULL;
+	}
+
+	tree->data = data;
+
+	return tree;
+}
+
+
+/* Add a child one level below */
+static struct dns_tree *dns_tree_add(struct dns_tree *tree, const char *name, void *data)
+{
+	struct dns_tree *node;
+
+	node = talloc_zero(tree, struct dns_tree);
+	if (node == NULL) {
+		return NULL;
+	}
+
+	node->name = talloc_strdup(tree, name);
+	if (node->name == NULL) {
+		talloc_free(node);
+		return NULL;
+	}
+	node->level = tree->level + 1;
+	node->num_children = 0;
+	node->children = NULL;
+	node->data = data;
+
+	if (tree->num_children == 0) {
+		tree->children = talloc_zero(tree, struct dns_tree *);
+	} else {
+		tree->children = talloc_realloc(tree, tree->children, struct dns_tree *,
+						tree->num_children+1);
+	}
+	if (tree->children == NULL) {
+		talloc_free(node);
+		return NULL;
+	}
+	tree->children[tree->num_children] = node;
+	tree->num_children++;
+
+	return node;
+}
+
+/* Find a node that matches the name components */
+static struct dns_tree *dns_tree_find(struct dns_tree *tree, int ncount, char **nlist, int *match_count)
+{
+	struct dns_tree *node, *next;
+	int i, j, start;
+
+	*match_count = -1;
+
+	if (strcmp(tree->name, "@") == 0) {
+		start = 0;
+	} else {
+		if (strcmp(tree->name, nlist[ncount-1]) != 0) {
+			return NULL;
+		}
+		start = 1;
+		*match_count = 0;
+	}
+
+	node = tree;
+	for (i=start; i<ncount; i++) {
+		if (node->num_children == 0) {
+			break;
+		}
+		next = NULL;
+		for (j=0; j<node->num_children; j++) {
+			if (strcmp(nlist[(ncount-1)-i], node->children[j]->name) == 0) {
+				next = node->children[j];
+				*match_count = i;
+				break;
+			}
+		}
+		if (next == NULL) {
+			break;
+		} else {
+			node = next;
+		}
+	}
+
+	return node;
+}
+
+/* Build a 2-level tree for resulting dns names */
+struct dns_tree *dns_build_tree(TALLOC_CTX *mem_ctx, const char *name, struct ldb_result *res)
+{
+	struct dns_tree *root, *base, *tree, *node;
+	const char *ptr;
+	int rootcount, ncount;
+	char **nlist;
+	int i, level, match_count;
+
+	rootcount = dns_split_name_components(mem_ctx, name, &nlist);
+	if (rootcount <= 0) {
+		return NULL;
+	}
+
+	root = dns_tree_init(mem_ctx, nlist[rootcount-1], NULL);
+	if (root == NULL) {
+		return NULL;
+	}
+
+	tree = root;
+	for (i=rootcount-2; i>=0; i--) {
+		tree = dns_tree_add(tree, nlist[i], NULL);
+		if (tree == NULL) {
+			goto failed;
+		}
+	}
+
+	base = tree;
+
+	/* Add all names in the result in a tree */
+	for (i=0; i<res->count; i++) {
+		ptr = ldb_msg_find_attr_as_string(res->msgs[i], "name", NULL);
+
+		if (strcmp(ptr, "@") == 0) {
+			base->data = res->msgs[i];
+			continue;
+		} else if (strcmp(ptr, name) == 0) {
+			base->data = res->msgs[i];
+			continue;
+		}
+
+		ncount = dns_split_name_components(root, ptr, &nlist);
+		if (ncount < 0) {
+			goto failed;
+		}
+
+		/* Find matching node */
+		tree = dns_tree_find(root, ncount, nlist, &match_count);
+		if (tree == NULL) {
+			goto failed;
+		}
+
+		/* Add missing name components */
+		for (level=match_count+1; level<ncount; level++) {
+			if (tree->level == rootcount+1) {
+				break;
+			}
+			if (level == ncount-1) {
+				node = dns_tree_add(tree, nlist[(ncount-1)-level], res->msgs[i]);
+			} else {
+				node = dns_tree_add(tree, nlist[(ncount-1)-level], NULL);
+			}
+			if (node == NULL) {
+				goto failed;
+			}
+			tree = node;
+		}
+
+		talloc_free(nlist);
+	}
+
+	/* Mark the base record, so it can be found easily */
+	base->level = -1;
+
+	return root;
+
+failed:
+	talloc_free(root);
+	return NULL;
+}
+
+
 static void _dns_add_name(TALLOC_CTX *mem_ctx, const char *name, char ***add_names, int *add_count)
 {
 	int i;
@@ -529,88 +710,53 @@ WERROR dns_fill_records_array(TALLOC_CTX *mem_ctx,
 				unsigned int select_flag,
 				const char *branch_name,
 				struct ldb_message *msg,
+				int num_children,
 				struct DNS_RPC_RECORDS_ARRAY *recs,
 				char ***add_names,
 				int *add_count)
 {
-	const char *nodename;
 	struct ldb_message_element *el;
+	const char *ptr;
 	int i, j;
-	bool found, node_is_rootzone;
+	bool found;
 
-	/* Check if we already have created record for the branch */
-	found = false;
-	if (branch_name == NULL) {
-		i = 0;
-		if (recs->count > 0) {
-			found = true;
-		}
+	if (recs->count == 0) {
+		recs->rec = talloc_zero(recs, struct DNS_RPC_RECORDS);
 	} else {
-		for (i=0; i<recs->count; i++) {
-			if (strcmp(branch_name, recs->rec[i].dnsNodeName.str) == 0) {
-				found = true;
-				break;
-			}
-		}
+		recs->rec = talloc_realloc(recs, recs->rec, struct DNS_RPC_RECORDS, recs->count+1);
 	}
+	if (recs->rec == NULL) {
+		return WERR_NOMEM;
+	}
+	i = recs->count;
+	recs->rec[i].wLength = 0;
+	recs->rec[i].wRecordCount = 0;
+	recs->rec[i].dwChildCount = num_children;
 
-	/* If not, add empty record */
-	if (!found) {
-		if (recs->count == 0) {
-			recs->rec = talloc_zero(recs, struct DNS_RPC_RECORDS);
-		} else {
-			recs->rec = talloc_realloc(recs, recs->rec, struct DNS_RPC_RECORDS, recs->count+1);
-		}
-		if (recs->rec == NULL) {
-			return WERR_NOMEM;
-		}
-		i = recs->count;
-		recs->rec[i].wLength = 0;
-		recs->rec[i].wRecordCount = 0;
-		recs->rec[i].dwChildCount = 0;
-
-		/* The base records returned with empty name */
-		/* Children records returned with names */
-		if (branch_name == NULL) {
-			recs->rec[i].dnsNodeName.str = talloc_strdup(recs, "");
-			recs->rec[i].dnsNodeName.len = 0;
-		} else {
-			recs->rec[i].dnsNodeName.str = talloc_strdup(recs, branch_name);
-			recs->rec[i].dnsNodeName.len = strlen(branch_name);
-		}
-		recs->rec[i].records = talloc_zero_array(recs, struct DNS_RPC_RECORD, 0);
-		recs->count++;
+	/* The base records returned with empty name */
+	/* Children records returned with names */
+	if (branch_name == NULL) {
+		recs->rec[i].dnsNodeName.str = talloc_strdup(recs, "");
+		recs->rec[i].dnsNodeName.len = 0;
+	} else {
+		recs->rec[i].dnsNodeName.str = talloc_strdup(recs, branch_name);
+		recs->rec[i].dnsNodeName.len = strlen(branch_name);
 	}
+	recs->rec[i].records = talloc_zero_array(recs, struct DNS_RPC_RECORD, 0);
+	recs->count++;
 
 	/* Allow empty records */
 	if (msg == NULL) {
 		return WERR_OK;
 	}
 
-	nodename = ldb_msg_find_attr_as_string(msg, "name", NULL);
-
-	if (strcmp(nodename, "@") == 0) {
-		node_is_rootzone = true;
-	} else {
-		node_is_rootzone = false;
-
-		/* child record */
-		if (branch_name != NULL) {
-			if (branch_name[strlen(branch_name)-1] != '.'
-				&& strcmp(nodename, branch_name) != 0) {
-				recs->rec[i].dwChildCount++;
-				return WERR_OK;
-			}
-		}
-	}
-
+	ptr = ldb_msg_find_attr_as_string(msg, "name", NULL);
 	el = ldb_msg_find_element(msg, "dnsRecord");
 	if (el == NULL || el->values == 0) {
-		DEBUG(0, ("dnsserver: Missing dnsRecord for %s\n", ldb_dn_get_linearized(msg->dn)));
 		return WERR_OK;
 	}
 
-	/* branch level record */
+	/* Add RR records */
 	for (j=0; j<el->num_values; j++) {
 		struct dnsp_DnssrvRpcRecord dnsp_rec;
 		struct DNS_RPC_RECORD *dns_rec;
@@ -661,8 +807,12 @@ WERROR dns_fill_records_array(TALLOC_CTX *mem_ctx,
 				dnsp_to_dns_copy(recs, &dnsp_rec, dns_rec);
 
 				/* Fix record flags */
-				if (node_is_rootzone) {
-					dns_rec->dwFlags |= (DNS_RPC_FLAG_ZONE_ROOT | DNS_RPC_FLAG_AUTH_ZONE_ROOT);
+				if (strcmp(ptr, "@") == 0) {
+					dns_rec->dwFlags |= DNS_RPC_FLAG_ZONE_ROOT;
+
+					if (dnsp_rec.rank == DNS_RANK_ZONE) {
+						dns_rec->dwFlags |= DNS_RPC_FLAG_AUTH_ZONE_ROOT;
+					}
 				}
 
 				if (dns_rec->dwFlags == DNS_RANK_NS_GLUE) {
diff --git a/source4/rpc_server/dnsserver/dnsserver.h b/source4/rpc_server/dnsserver/dnsserver.h
index 5fc13c8704..818dc60b62 100644
--- a/source4/rpc_server/dnsserver/dnsserver.h
+++ b/source4/rpc_server/dnsserver/dnsserver.h
@@ -151,6 +151,14 @@ struct dnsserver_zone {
 };
 
 
+struct dns_tree {
+	const char *name;
+	int level;
+	unsigned int num_children;
+	struct dns_tree **children;
+	void *data;
+};
+
 /* Data structure manipulation functions from dnsdata.c */
 
 struct IP4_ARRAY *ip4_array_copy(TALLOC_CTX *mem_ctx, struct IP4_ARRAY *ip4);
@@ -169,10 +177,12 @@ void dnsp_to_dns_copy(TALLOC_CTX *mem_ctx, struct dnsp_DnssrvRpcRecord *dnsp,
 			struct DNS_RPC_RECORD *dns);
 struct dnsp_DnssrvRpcRecord *dns_to_dnsp_copy(TALLOC_CTX *mem_ctx, struct DNS_RPC_RECORD *dns);
 
+struct dns_tree *dns_build_tree(TALLOC_CTX *mem_ctx, const char *name, struct ldb_result *res);
 WERROR dns_fill_records_array(TALLOC_CTX *mem_ctx, struct dnsserver_zone *z,
 			enum dns_record_type record_type,
 			unsigned int select_flag, const char *zone_name,
-			struct ldb_message *msg, struct DNS_RPC_RECORDS_ARRAY *recs,
+			struct ldb_message *msg, int num_children,
+			struct DNS_RPC_RECORDS_ARRAY *recs,
 			char ***add_names, int *add_count);