From 176c92c26bfc9b4a36b99c58b66e117fbe5ae706 Mon Sep 17 00:00:00 2001
From: Phil Sutter <phil@nwl.cc>
Date: Tue, 7 Jul 2020 11:43:26 +0200
Subject: nft: Introduce a dedicated base chain array

Preparing for sorted chain output, introduce a per-table array holding
base chains indexed by nf_inet_hooks value. Since the latter is ordered
correctly, iterating over the array will return base chains in expected
order.

Signed-off-by: Phil Sutter <phil@nwl.cc>
---
 iptables/nft-cache.c | 34 +++++++++++++++++++++++++++++++++-
 iptables/nft.c       | 12 +++++++++++-
 iptables/nft.h       |  1 +
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/iptables/nft-cache.c b/iptables/nft-cache.c
index f62e5100..bd19b6df 100644
--- a/iptables/nft-cache.c
+++ b/iptables/nft-cache.c
@@ -208,7 +208,24 @@ int nft_cache_add_chain(struct nft_handle *h, const struct builtin_table *t,
 	const char *cname = nftnl_chain_get_str(c, NFTNL_CHAIN_NAME);
 	struct nft_chain *nc = nft_chain_alloc(c);
 
-	list_add_tail(&nc->head, &h->cache->table[t->type].chains->list);
+	if (nftnl_chain_is_set(c, NFTNL_CHAIN_HOOKNUM)) {
+		uint32_t hooknum = nftnl_chain_get_u32(c, NFTNL_CHAIN_HOOKNUM);
+
+		if (hooknum >= NF_INET_NUMHOOKS) {
+			nft_chain_free(nc);
+			return -EINVAL;
+		}
+
+		if (h->cache->table[t->type].base_chains[hooknum]) {
+			nft_chain_free(nc);
+			return -EEXIST;
+		}
+
+		h->cache->table[t->type].base_chains[hooknum] = nc;
+	} else {
+		list_add_tail(&nc->head,
+			      &h->cache->table[t->type].chains->list);
+	}
 	hlist_add_head(&nc->hnode, chain_name_hlist(h, t, cname));
 	return 0;
 }
@@ -609,6 +626,19 @@ static int __flush_set_cache(struct nftnl_set *s, void *data)
 	return 0;
 }
 
+static void flush_base_chain_cache(struct nft_chain **base_chains)
+{
+	int i;
+
+	for (i = 0; i < NF_INET_NUMHOOKS; i++) {
+		if (!base_chains[i])
+			continue;
+		hlist_del(&base_chains[i]->hnode);
+		nft_chain_free(base_chains[i]);
+		base_chains[i] = NULL;
+	}
+}
+
 static int flush_cache(struct nft_handle *h, struct nft_cache *c,
 		       const char *tablename)
 {
@@ -620,6 +650,7 @@ static int flush_cache(struct nft_handle *h, struct nft_cache *c,
 		if (!table)
 			return 0;
 
+		flush_base_chain_cache(c->table[table->type].base_chains);
 		nft_chain_foreach(h, tablename, __flush_chain_cache, NULL);
 
 		if (c->table[table->type].sets)
@@ -632,6 +663,7 @@ static int flush_cache(struct nft_handle *h, struct nft_cache *c,
 		if (h->tables[i].name == NULL)
 			continue;
 
+		flush_base_chain_cache(c->table[i].base_chains);
 		if (c->table[i].chains) {
 			nft_chain_list_free(c->table[i].chains);
 			c->table[i].chains = NULL;
diff --git a/iptables/nft.c b/iptables/nft.c
index 1b740005..4187e691 100644
--- a/iptables/nft.c
+++ b/iptables/nft.c
@@ -2396,12 +2396,22 @@ int nft_chain_foreach(struct nft_handle *h, const char *table,
 	const struct builtin_table *t;
 	struct nft_chain_list *list;
 	struct nft_chain *c, *c_bak;
-	int ret;
+	int i, ret;
 
 	t = nft_table_builtin_find(h, table);
 	if (!t)
 		return -1;
 
+	for (i = 0; i < NF_INET_NUMHOOKS; i++) {
+		c = h->cache->table[t->type].base_chains[i];
+		if (!c)
+			continue;
+
+		ret = cb(c, data);
+		if (ret < 0)
+			return ret;
+	}
+
 	list = h->cache->table[t->type].chains;
 	if (!list)
 		return -1;
diff --git a/iptables/nft.h b/iptables/nft.h
index ac227b4c..1a2506ee 100644
--- a/iptables/nft.h
+++ b/iptables/nft.h
@@ -40,6 +40,7 @@ enum nft_cache_level {
 
 struct nft_cache {
 	struct {
+		struct nft_chain	*base_chains[NF_INET_NUMHOOKS];
 		struct nft_chain_list	*chains;
 		struct nftnl_set_list	*sets;
 		bool			exists;
-- 
cgit v1.2.3