From cfb4e8187940b89acf163a7f4181bf04abd333d0 Mon Sep 17 00:00:00 2001
From: Kristian Evensen <kristian.evensen@gmail.com>
Date: Tue, 1 May 2018 22:44:16 +0200
Subject: [PATCH] NFQUEUE early insert

---
 net/netfilter/nfnetlink_queue.c | 69 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/net/netfilter/nfnetlink_queue.c b/net/netfilter/nfnetlink_queue.c
index c97966298..254866d5c 100644
--- a/net/netfilter/nfnetlink_queue.c
+++ b/net/netfilter/nfnetlink_queue.c
@@ -43,6 +43,9 @@
 
 #if IS_ENABLED(CONFIG_NF_CONNTRACK)
 #include <net/netfilter/nf_conntrack.h>
+#include <net/netfilter/nf_conntrack_core.h>
+#include <net/netfilter/nf_conntrack_l3proto.h>
+#include <net/netfilter/nf_conntrack_l4proto.h>
 #endif
 
 #define NFQNL_QMAX_DEFAULT 1024
@@ -1151,6 +1154,64 @@ static int nfqa_parse_bridge(struct nf_queue_entry *entry,
 	return 0;
 }
 
+#if IS_ENABLED(CONFIG_NF_CONNTRACK)
+static void nfqnl_update_ct(struct net *net, struct sk_buff *skb)
+{
+	struct nf_conn *ct = NULL;
+	enum ip_conntrack_info ctinfo;
+	const struct nf_conntrack_l3proto *l3proto;
+	const struct nf_conntrack_l4proto *l4proto;
+	u_int16_t l3num;
+	unsigned int dataoff;
+	u_int8_t l4num;
+	struct nf_conntrack_tuple tuple;
+	struct nf_conntrack_tuple_hash *h;
+
+	ct = nf_ct_get(skb, &ctinfo);
+	l3num = nf_ct_l3num(ct);
+	l3proto = nf_ct_l3proto_find_get(l3num);
+
+	if (l3proto->get_l4proto(skb, skb_network_offset(skb), &dataoff,
+				 &l4num) <= 0) {
+		pr_info("Failed to get l4 protonum\n");
+		return;
+	}
+
+	l4proto = nf_ct_l4proto_find_get(l3num, l4num);
+
+	if (!nf_ct_get_tuple(skb, skb_network_offset(skb), dataoff, l3num,
+			     l4num, net, &tuple, l3proto, l4proto)) {
+		pr_info("Failed to get tuple\n");
+		return;
+	}
+
+#if IS_ENABLED(CONFIG_NF_CONNTRACK_ZONES)
+	h = nf_conntrack_find_get(net, &(ct->zone), &tuple);
+#else
+	h = nf_conntrack_find_get(net, NULL, &tuple);
+#endif
+
+	if (!h) {
+		pr_info("early insert: tuple %u %pI4:%hu -> %pI4:%hu\n",
+		       tuple.dst.protonum, &(tuple.src.u3.ip), ntohs(tuple.src.u.all),
+	       	       &(tuple.dst.u3.ip), ntohs(tuple.dst.u.all));
+
+		rcu_read_lock();
+		if (!nf_conntrack_hash_check_insert(ct)) {
+			pr_info("early insert confirmed\n");
+		}
+		rcu_read_unlock();
+	} else {
+		pr_info("replace: tuple %u %pI4:%hu -> %pI4:%hu\n",
+		       tuple.dst.protonum, &(tuple.src.u3.ip), ntohs(tuple.src.u.all),
+	       	       &(tuple.dst.u3.ip), ntohs(tuple.dst.u.all));
+		nf_ct_put(ct);
+		ct = nf_ct_tuplehash_to_ctrack(h);
+		nf_ct_set(skb, ct, IP_CT_NEW);
+	}
+}
+#endif
+
 static int nfqnl_recv_verdict(struct net *net, struct sock *ctnl,
 			      struct sk_buff *skb,
 			      const struct nlmsghdr *nlh,
@@ -1213,6 +1274,14 @@ static int nfqnl_recv_verdict(struct net *net, struct sock *ctnl,
 	if (nfqa[NFQA_MARK])
 		entry->skb->mark = ntohl(nla_get_be32(nfqa[NFQA_MARK]));
 
+#if IS_ENABLED(CONFIG_NF_CONNTRACK)
+	nf_ct_get(entry->skb, &ctinfo);
+
+	if (ctinfo == IP_CT_NEW && verdict != NF_STOLEN && verdict != NF_DROP) {
+		nfqnl_update_ct(net, entry->skb);
+	}
+#endif
+
 	nf_reinject(entry, verdict);
 	return 0;
 }
-- 
2.14.1