From 4f21535e4066511f0e6dfa224e6409f5fd4f97a8 Mon Sep 17 00:00:00 2001
From: "fu.lin" <fu.lin10@huawei.com>
Date: Thu, 19 Aug 2021 16:29:49 +0800
Subject: [PATCH 40/54] criu: switch to nftables api

usage: criu --use-nft

Signed-off-by: fu.lin <fu.lin10@huawei.com>
---
 criu/config.c             |  1 +
 criu/cr-dump.c            |  4 ++--
 criu/cr-restore.c         |  6 ++++--
 criu/crtools.c            |  1 +
 criu/include/cr_options.h |  1 +
 criu/include/net.h        |  7 +++++--
 criu/include/netfilter.h  |  7 +++++--
 criu/include/sk-inet.h    |  2 +-
 criu/kerndat.c            |  3 ++-
 criu/net.c                | 37 +++++++++++++++++++++++++++----------
 criu/netfilter.c          | 14 +++++++++++---
 criu/nftables.c           | 29 +++++++++++++++++++----------
 criu/sk-tcp.c             | 14 +++++++-------
 13 files changed, 86 insertions(+), 40 deletions(-)

diff --git a/criu/config.c b/criu/config.c
index 9268cd1..676b424 100644
--- a/criu/config.c
+++ b/criu/config.c
@@ -554,6 +554,7 @@ int parse_options(int argc, char **argv, bool *usage_error,
 		BOOL_OPT("ignore-special-dump", &opts.ignore_special_dump),
 		BOOL_OPT("file-locks-repair", &opts.file_locks_repair),
 		{"reserve-ports", 		required_argument,	0, 'P'	},
+		BOOL_OPT("use-nft", &opts.use_nft),
 		{ },
 	};
 
diff --git a/criu/cr-dump.c b/criu/cr-dump.c
index 695c98a..d54b7a7 100644
--- a/criu/cr-dump.c
+++ b/criu/cr-dump.c
@@ -1782,7 +1782,7 @@ static int cr_dump_finish(int ret)
 	 *    start rollback procedure and cleanup everything.
 	 */
 	if (ret || post_dump_ret || opts.final_state == TASK_ALIVE) {
-		network_unlock();
+		network_unlock(opts.tree_id);
 		delete_link_remaps();
 		clean_cr_time_mounts();
 	}
@@ -1931,7 +1931,7 @@ int cr_dump_tasks(pid_t pid)
 	if (collect_pstree_ids())
 		goto err;
 
-	if (network_lock())
+	if (network_lock(opts.tree_id))
 		goto err;
 
 	if (collect_file_locks())
diff --git a/criu/cr-restore.c b/criu/cr-restore.c
index 4fd29a5..945d984 100644
--- a/criu/cr-restore.c
+++ b/criu/cr-restore.c
@@ -2499,7 +2499,7 @@ skip_ns_bouncing:
 		goto out_kill;
 
 	/* Unlock network before disabling repair mode on sockets */
-	network_unlock();
+	network_unlock(vpid(init));
 	network_status |= NETWORK_UNLOCK;
 
 	/*
@@ -2567,6 +2567,8 @@ skip_ns_bouncing:
 	if (ret != 0)
 		pr_err("Post-resume script ret code %d\n", ret);
 
+	network_delete_set(vpid(init));
+
 	if (!opts.restore_detach && !opts.exec_cmd)
 		wait(NULL);
 
@@ -2712,7 +2714,7 @@ err:
 				pr_err("collect inet sk cinfo fail");
 		}
 		if ((network_status & NETWORK_UNLOCK) == 0)
-			network_unlock();
+			network_unlock(vpid(root_item));
 	}
 
 	return ret;
diff --git a/criu/crtools.c b/criu/crtools.c
index 888714c..aaae44e 100644
--- a/criu/crtools.c
+++ b/criu/crtools.c
@@ -459,6 +459,7 @@ usage:
 "  --ignore-special-dump Ignore special task tid page dump\n"
 "  --file-locks-repair    Use repair mode to dump and restore file locks\n"
 "  --reserve-ports       Reserve src ports in kernel\n"
+"  --use-nft             Use nft API instead of iptables cmd in network locking"
 "\n"
 "Check options:\n"
 "  Without options, \"criu check\" checks availability of absolutely required\n"
diff --git a/criu/include/cr_options.h b/criu/include/cr_options.h
index 85d852f..62fbbb1 100644
--- a/criu/include/cr_options.h
+++ b/criu/include/cr_options.h
@@ -188,6 +188,7 @@ struct cr_options {
 	char			*share_dst_ports;
 	char			*share_src_ports;
 	int			reserve_ports;
+	int			use_nft;
 };
 
 extern struct cr_options opts;
diff --git a/criu/include/net.h b/criu/include/net.h
index bda0ff3..9daea8d 100644
--- a/criu/include/net.h
+++ b/criu/include/net.h
@@ -30,8 +30,11 @@ struct veth_pair {
 
 extern int collect_net_namespaces(bool for_dump);
 
-extern int network_lock(void);
-extern void network_unlock(void);
+extern int network_prepare(pid_t tree_id);
+extern void network_delete_rule(pid_t tree_id);
+extern void network_delete_set(pid_t tree_id);
+extern int network_lock(pid_t tree_id);
+extern void network_unlock(pid_t tree_id);
 extern int network_lock_internal(void);
 
 extern struct ns_desc net_ns_desc;
diff --git a/criu/include/netfilter.h b/criu/include/netfilter.h
index 35ef262..c92762c 100644
--- a/criu/include/netfilter.h
+++ b/criu/include/netfilter.h
@@ -1,9 +1,12 @@
 #ifndef __CR_NETFILTER_H__
 #define __CR_NETFILTER_H__
 
+#include <sys/types.h>
+#include <stdbool.h>
+
 struct inet_sk_desc;
-extern int nf_lock_connection(struct inet_sk_desc *);
-extern int nf_unlock_connection(struct inet_sk_desc *);
+extern int nf_lock_connection(struct inet_sk_desc *, pid_t, bool);
+extern int nf_unlock_connection(struct inet_sk_desc *, bool);
 
 struct inet_sk_info;
 extern int nf_unlock_connection_info(struct inet_sk_info *);
diff --git a/criu/include/sk-inet.h b/criu/include/sk-inet.h
index 4181fbe..88e0881 100644
--- a/criu/include/sk-inet.h
+++ b/criu/include/sk-inet.h
@@ -81,7 +81,7 @@ static inline void tcp_repair_off(int fd)
 
 extern void tcp_locked_conn_add(struct inet_sk_info *);
 extern void rst_unlock_tcp_connections(void);
-extern void cpt_unlock_tcp_connections(void);
+extern void cpt_unlock_tcp_connections(bool);
 
 extern void read_reserved_ports(char *path);
 extern void write_reserved_ports(char *path);
diff --git a/criu/kerndat.c b/criu/kerndat.c
index b2c47c5..59570a2 100644
--- a/criu/kerndat.c
+++ b/criu/kerndat.c
@@ -1065,7 +1065,8 @@ int kerndat_init(void)
 	memset(&kdat, 0, sizeof(kdat));
 
 	preload_socket_modules();
-	preload_netfilter_modules();
+	if (!opts.use_nft)
+		preload_netfilter_modules();
 
 	if (check_pagemap()) {
 		pr_err("check_pagemap failed when initializing kerndat.\n");
diff --git a/criu/net.c b/criu/net.c
index 19329cf..2bd6f77 100644
--- a/criu/net.c
+++ b/criu/net.c
@@ -45,6 +45,7 @@
 #include "util.h"
 #include "external.h"
 #include "fdstore.h"
+#include "nftables.h"
 
 #include "protobuf.h"
 #include "images/netdev.pb-c.h"
@@ -2868,9 +2869,13 @@ int network_lock_internal(void)
 		return -1;
 
 
-	ret |= iptables_restore(false, conf, sizeof(conf) - 1);
-	if (kdat.ipv6)
-		ret |= iptables_restore(true, conf, sizeof(conf) - 1);
+	if (opts.use_nft)
+		ret = nft_lock();
+	else {
+		ret |= iptables_restore(false, conf, sizeof(conf) - 1);
+		if (kdat.ipv6)
+			ret |= iptables_restore(true, conf, sizeof(conf) - 1);
+	}
 
 	if (ret)
 		pr_err("Locking network failed: iptables-restore returned %d. "
@@ -2897,9 +2902,13 @@ static int network_unlock_internal(void)
 	if (switch_ns(root_item->pid->real, &net_ns_desc, &nsret))
 		return -1;
 
-	ret |= iptables_restore(false, conf, sizeof(conf) - 1);
-	if (kdat.ipv6)
-		ret |= iptables_restore(true, conf, sizeof(conf) - 1);
+	if (opts.use_nft)
+		ret = nft_unlock();
+	else {
+		ret |= iptables_restore(false, conf, sizeof(conf) - 1);
+		if (kdat.ipv6)
+			ret |= iptables_restore(true, conf, sizeof(conf) - 1);
+	}
 
 	if (restore_ns(nsret, &net_ns_desc))
 		ret = -1;
@@ -2907,10 +2916,13 @@ static int network_unlock_internal(void)
 	return ret;
 }
 
-int network_lock(void)
+int network_lock(pid_t tree_id)
 {
 	pr_info("Lock network\n");
 
+	if (opts.use_nft && opts.tcp_established_ok && network_prepare(tree_id) < 0)
+		return -1;
+
 	/* Each connection will be locked on dump */
 	if  (!(root_ns_mask & CLONE_NEWNET))
 		return 0;
@@ -2921,7 +2933,7 @@ int network_lock(void)
 	return network_lock_internal();
 }
 
-void network_unlock(void)
+void network_unlock(pid_t tree_id)
 {
 	pr_info("Unlock network\n");
 
@@ -2930,8 +2942,13 @@ void network_unlock(void)
 		write_reserved_ports(RESERVED_PORTS_PATH);
 	}
 
-	cpt_unlock_tcp_connections();
-	rst_unlock_tcp_connections();
+	if (opts.use_nft && opts.tcp_established_ok)
+		network_delete_rule(tree_id);
+
+	cpt_unlock_tcp_connections(opts.use_nft);
+
+	if (!opts.use_nft)
+		rst_unlock_tcp_connections();
 
 	if (root_ns_mask & CLONE_NEWNET) {
 		run_scripts(ACT_NET_UNLOCK);
diff --git a/criu/netfilter.c b/criu/netfilter.c
index 368651c..b2ec7ed 100644
--- a/criu/netfilter.c
+++ b/criu/netfilter.c
@@ -15,6 +15,8 @@
 #include "sk-inet.h"
 #include "kerndat.h"
 
+#include "nftables.h"
+
 static char buf[512];
 
 /*
@@ -129,13 +131,19 @@ static int nf_connection_switch(struct inet_sk_desc *sk, bool lock)
 	return ret;
 }
 
-int nf_lock_connection(struct inet_sk_desc *sk)
+int nf_lock_connection(struct inet_sk_desc *sk, pid_t tree_id, bool use_nft)
 {
-	return nf_connection_switch(sk, true);
+	if (use_nft)
+		return nft_connection_switch(sk, true, tree_id);
+	else
+		return nf_connection_switch(sk, true);
 }
 
-int nf_unlock_connection(struct inet_sk_desc *sk)
+int nf_unlock_connection(struct inet_sk_desc *sk, bool use_nft)
 {
+	if (use_nft)
+		return 0;
+
 	return nf_connection_switch(sk, false);
 }
 
diff --git a/criu/nftables.c b/criu/nftables.c
index 817f157..739aee4 100644
--- a/criu/nftables.c
+++ b/criu/nftables.c
@@ -653,25 +653,34 @@ int network_prepare(pid_t tree_id)
 	return mnl_common(network_prepare_internal, NULL, &tree_id);
 }
 
-static int network_unprepare_internal(struct mnl_params *params,
-				      batch_func_t _, void *args)
+static int network_delete_rule_internal(struct mnl_params *params,
+					batch_func_t _, void *args)
 {
 	pid_t tree_id = *(pid_t *)args;
 
-	if (nft_rule_common(params, tree_id, false) < 0)
-		return -1;
+	return nft_rule_common(params, tree_id, false);
+}
 
-	if (nft_set_common(params, tree_id, false) < 0)
-		return -1;
+void network_delete_rule(pid_t tree_id)
+{
+	pr_info("unlock network\n");
 
-	return 0;
+	mnl_common(network_delete_rule_internal, NULL, &tree_id);
+}
+
+static int network_delete_set_internal(struct mnl_params *params,
+					batch_func_t _, void *args)
+{
+	pid_t tree_id = *(pid_t *)args;
+
+	return nft_set_common(params, tree_id, false);
 }
 
-void network_unprepare(pid_t tree_id)
+void network_delete_set(pid_t tree_id)
 {
-	pr_info("Unprepare network\n");
+	pr_info("clear nft set\n");
 
-	mnl_common(network_unprepare_internal, NULL, &tree_id);
+	mnl_common(network_delete_set_internal, NULL, &tree_id);
 }
 
 static int add_set_elem_internal(struct nftnl_set *s, void *data, size_t len)
diff --git a/criu/sk-tcp.c b/criu/sk-tcp.c
index 67846c3..a9a9047 100644
--- a/criu/sk-tcp.c
+++ b/criu/sk-tcp.c
@@ -55,7 +55,7 @@ static int tcp_repair_established(int fd, struct inet_sk_desc *sk)
 	}
 
 	if (!(root_ns_mask & CLONE_NEWNET)) {
-		ret = nf_lock_connection(sk);
+		ret = nf_lock_connection(sk, opts.tree_id, opts.use_nft);
 		if (ret < 0)
 			goto err2;
 	}
@@ -70,21 +70,21 @@ static int tcp_repair_established(int fd, struct inet_sk_desc *sk)
 
 err3:
 	if (!(root_ns_mask & CLONE_NEWNET))
-		nf_unlock_connection(sk);
+		nf_unlock_connection(sk, opts.use_nft);
 err2:
 	close(sk->rfd);
 err1:
 	return -1;
 }
 
-static void tcp_unlock_one(struct inet_sk_desc *sk)
+static void tcp_unlock_one(struct inet_sk_desc *sk, bool use_nft)
 {
 	int ret;
 
 	list_del(&sk->rlist);
 
-	if (!(root_ns_mask & CLONE_NEWNET)) {
-		ret = nf_unlock_connection(sk);
+	if (!(root_ns_mask & CLONE_NEWNET) && !use_nft) {
+		ret = nf_unlock_connection(sk, false);
 		if (ret < 0)
 			pr_perror("Failed to unlock TCP connection");
 	}
@@ -101,12 +101,12 @@ static void tcp_unlock_one(struct inet_sk_desc *sk)
 	close(sk->rfd);
 }
 
-void cpt_unlock_tcp_connections(void)
+void cpt_unlock_tcp_connections(bool use_nft)
 {
 	struct inet_sk_desc *sk, *n;
 
 	list_for_each_entry_safe(sk, n, &cpt_tcp_repair_sockets, rlist)
-		tcp_unlock_one(sk);
+		tcp_unlock_one(sk, use_nft);
 }
 
 static int dump_tcp_conn_state(struct inet_sk_desc *sk)
-- 
2.33.1

