diff mbox series

[3/8] conntrack: accept parameters from stdin

Message ID 20200925124919.9389-4-mikhail.sennikovskii@cloud.ionos.com
State Under Review
Delegated to: Pablo Neira
Headers show
Series Fast bulk transfers of large sets of ct entries | expand

Commit Message

Mikhail Sennikovsky Sept. 25, 2020, 12:49 p.m. UTC
This commit allows accepting multiple setsĀ of ct entry-related
parameters on stdin.
This is useful when one needs to add/update/delete a large
set of ct entries with a single conntrack tool invocation.

Expected syntax is "conntrack [-I|-D|-U] [table] -".
When invoked like that, conntrack expects ct entry parameters
to be passed to the stdin, each line presenting a separate parameter
set.

Signed-off-by: Mikhail Sennikovsky <mikhail.sennikovskii@cloud.ionos.com>
---
 src/conntrack.c | 196 +++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 161 insertions(+), 35 deletions(-)
diff mbox series

Patch

diff --git a/src/conntrack.c b/src/conntrack.c
index a26fa60..5834f2d 100644
--- a/src/conntrack.c
+++ b/src/conntrack.c
@@ -96,15 +96,18 @@  static struct {
 	struct nfct_bitmask *label_modify;
 } tmpl;
 
+int cur_argc;
+char **cur_argv;
+
 static int alloc_tmpl_objects(void)
 {
+	memset(&tmpl, 0, sizeof(tmpl));
+
 	tmpl.ct = nfct_new();
 	tmpl.exptuple = nfct_new();
 	tmpl.mask = nfct_new();
 	tmpl.exp = nfexp_new();
 
-	memset(&tmpl.mark, 0, sizeof(tmpl.mark));
-
 	return tmpl.ct != NULL && tmpl.exptuple != NULL &&
 	       tmpl.mask != NULL && tmpl.exp != NULL;
 }
@@ -685,12 +688,18 @@  static void free_options(void)
 void __attribute__((noreturn))
 exit_error(enum exittype status, const char *msg, ...)
 {
+	int i;
 	va_list args;
 
 	free_options();
 	va_start(args, msg);
 	fprintf(stderr,"%s v%s (conntrack-tools): ", PROGNAME, VERSION);
 	vfprintf(stderr, msg, args);
+	if (cur_argc) {
+		fprintf(stderr, "\nargs:");
+		for (i = 1; i < cur_argc; ++i)
+			fprintf(stderr, " %s", cur_argv[i]);
+	}
 	fprintf(stderr, "\n");
 	va_end(args);
 	if (status == PARAMETER_PROBLEM)
@@ -2317,23 +2326,63 @@  nfct_set_nat_details(const int opt, struct nf_conntrack *ct,
 		nfct_set_attr_u16(ct, ATTR_DNAT_PORT,
 				  ntohs((uint16_t)atoi(port_str)));
 	}
+}
+
+static int line_to_argcv(char *cmd, char *line, char ***pargv, size_t *pargv_size)
+{
+	char *arg;
+	int argc;
+	char **argv = *pargv;
+	size_t argv_size = *pargv_size;
+	int argc_max = argv_size / sizeof(*argv);
+
+#define _ARG_ADD(_arg) do { \
+		if (argc == argc_max) { \
+			argc_max += 20; \
+			argv_size = argc_max * sizeof (argv[0]); \
+			argv = realloc(argv, argv_size); \
+			if (!argv) \
+				exit_error(OTHER_PROBLEM, "out of memory"); \
+		} \
+		argv[argc] = _arg; \
+		++argc; \
+} while (0)
+
+#define _ARG_SEP " \t\n\r"
+	for (argc = 0, arg = strtok (line, _ARG_SEP);
+			arg;
+			arg = strtok (NULL, _ARG_SEP)) {
+		/*
+		 * getopt_long expects argv[0] to be the command name,
+		 * and would always skip it so we need to include it here
+		 */
+		if (!argc && cmd)
+			_ARG_ADD(cmd);
+		_ARG_ADD(arg);
+	}
+
+#undef _ARG_ADD
+#undef _ARG_SEP
 
+	*pargv = argv;
+	*pargv_size = argv_size;
+
+	return argc;
 }
 
 int main(int argc, char *argv[])
 {
 	int c, cmd;
-	unsigned int type = 0, event_mask = 0, l4flags = 0, status = 0;
+	unsigned int type = 0, event_mask = 0, l4flags, status = 0;
 	int res = 0, partial;
-	size_t socketbuffersize = 0;
-	int family = AF_UNSPEC;
-	int protonum = 0;
+	size_t socketbuffersize;
+	int family;
+	int protonum;
 	union ct_address ad;
 	unsigned int command = 0;
-
-	/* we release these objects in the exit_error() path. */
-	if (!alloc_tmpl_objects())
-		exit_error(OTHER_PROBLEM, "out of memory");
+	FILE *opts_file = NULL;
+	char **argv_buf = NULL, *getline_buf = NULL;
+	size_t argv_buf_size = 0, getline_buf_size = 0;
 
 	register_tcp();
 	register_udp();
@@ -2348,6 +2397,23 @@  int main(int argc, char *argv[])
 	/* disable explicit missing arguments error output from getopt_long */
 	opterr = 0;
 
+parse_opts:
+
+	options = 0;
+	filter_family = 0;
+	memset(dir2network, 0, sizeof(dir2network));
+	/* all allocate-able objects get freed zero-inited
+	 * at the end of each iteration */
+
+	l4flags = 0;
+	family = AF_UNSPEC;
+	socketbuffersize = 0;
+	protonum = 0;
+
+	/* we release these objects in the exit_error() path. */
+	if (!alloc_tmpl_objects())
+		exit_error(OTHER_PROBLEM, "out of memory");
+
 	while ((c = getopt_long(argc, argv, getopt_str, opts, NULL)) != -1) {
 	switch(c) {
 		/* commands */
@@ -2585,6 +2651,26 @@  int main(int argc, char *argv[])
 					      "`--dst-nat' with `--any-nat'");
 	}
 	cmd = bit2cmd(command);
+
+	if (!opts_file && optind == argc - 1 && !strcmp(argv[optind], "-")) {
+		switch (command) {
+		case CT_CREATE:
+		case EXP_CREATE:
+		case CT_UPDATE:
+		case CT_DELETE:
+		case EXP_DELETE:
+			break;
+		default:
+			exit_error(PARAMETER_PROBLEM, "stdin mode not supported "
+					"for this command!");
+		}
+		if (options)
+			exit_error(PARAMETER_PROBLEM, "no extra options are expected "
+					"with stdin read mode!");
+		opts_file = stdin;
+		goto next_opts;
+	}
+
 	res = generic_opt_check(options, NUMBER_OF_OPT,
 				commands_v_options[cmd], optflags,
 				addr_valid_flags, ADDR_VALID_FLAGS_MAX,
@@ -2670,8 +2756,8 @@  int main(int argc, char *argv[])
 			printf("</conntrack>\n");
 			fflush(stdout);
 		}
-
 		nfct_close(cth);
+		cth = NULL;
 		break;
 
 	case EXP_LIST:
@@ -2681,12 +2767,13 @@  int main(int argc, char *argv[])
 
 		nfexp_callback_register(cth, NFCT_T_ALL, dump_exp_cb, NULL);
 		res = nfexp_query(cth, NFCT_Q_DUMP, &family);
-		nfct_close(cth);
 
 		if (dump_xml_header_done == 0) {
 			printf("</expect>\n");
 			fflush(stdout);
 		}
+		nfct_close(cth);
+		cth = NULL;
 		break;
 
 	case CT_CREATE:
@@ -2702,14 +2789,15 @@  int main(int argc, char *argv[])
 			nfct_set_attr(tmpl.ct, ATTR_CONNLABELS,
 					xnfct_bitmask_clone(tmpl.label_modify));
 
-		cth = nfct_open(CONNTRACK, 0);
-		if (!cth)
-			exit_error(OTHER_PROBLEM, "Can't open handler");
+		if (!cth) {
+			cth = nfct_open(CONNTRACK, 0);
+			if (!cth)
+				exit_error(OTHER_PROBLEM, "Can't open handler");
+		}
 
 		res = nfct_query(cth, NFCT_Q_CREATE, tmpl.ct);
 		if (res != -1)
 			counter++;
-		nfct_close(cth);
 		break;
 
 	case EXP_CREATE:
@@ -2717,18 +2805,21 @@  int main(int argc, char *argv[])
 		nfexp_set_attr(tmpl.exp, ATTR_EXP_EXPECTED, tmpl.exptuple);
 		nfexp_set_attr(tmpl.exp, ATTR_EXP_MASK, tmpl.mask);
 
-		cth = nfct_open(EXPECT, 0);
-		if (!cth)
-			exit_error(OTHER_PROBLEM, "Can't open handler");
+		if (!cth) {
+			cth = nfct_open(EXPECT, 0);
+			if (!cth)
+				exit_error(OTHER_PROBLEM, "Can't open handler");
+		}
 
 		res = nfexp_query(cth, NFCT_Q_CREATE, tmpl.exp);
-		nfct_close(cth);
 		break;
 
 	case CT_UPDATE:
-		cth = nfct_open(CONNTRACK, 0);
+		if (!cth)
+			cth = nfct_open(CONNTRACK, 0);
 		/* internal handler for delete_cb, otherwise we hit EILSEQ */
-		ith = nfct_open(CONNTRACK, 0);
+		if (!ith)
+			ith = nfct_open(CONNTRACK, 0);
 		if (!cth || !ith)
 			exit_error(OTHER_PROBLEM, "Can't open handler");
 
@@ -2737,13 +2828,13 @@  int main(int argc, char *argv[])
 		nfct_callback_register(cth, NFCT_T_ALL, update_cb, tmpl.ct);
 
 		res = nfct_query(cth, NFCT_Q_DUMP, &family);
-		nfct_close(ith);
-		nfct_close(cth);
 		break;
 		
 	case CT_DELETE:
-		cth = nfct_open(CONNTRACK, 0);
-		ith = nfct_open(CONNTRACK, 0);
+		if (!cth)
+			cth = nfct_open(CONNTRACK, 0);
+		if (!ith)
+			ith = nfct_open(CONNTRACK, 0);
 		if (!cth || !ith)
 			exit_error(OTHER_PROBLEM, "Can't open handler");
 
@@ -2768,19 +2859,18 @@  int main(int argc, char *argv[])
 
 		nfct_filter_dump_destroy(filter_dump);
 
-		nfct_close(ith);
-		nfct_close(cth);
 		break;
 
 	case EXP_DELETE:
 		nfexp_set_attr(tmpl.exp, ATTR_EXP_EXPECTED, tmpl.ct);
 
-		cth = nfct_open(EXPECT, 0);
-		if (!cth)
-			exit_error(OTHER_PROBLEM, "Can't open handler");
+		if (!cth) {
+			cth = nfct_open(EXPECT, 0);
+			if (!cth)
+				exit_error(OTHER_PROBLEM, "Can't open handler");
+		}
 
 		res = nfexp_query(cth, NFCT_Q_DESTROY, tmpl.exp);
-		nfct_close(cth);
 		break;
 
 	case CT_GET:
@@ -2791,6 +2881,7 @@  int main(int argc, char *argv[])
 		nfct_callback_register(cth, NFCT_T_ALL, dump_cb, tmpl.ct);
 		res = nfct_query(cth, NFCT_Q_GET, tmpl.ct);
 		nfct_close(cth);
+		cth = NULL;
 		break;
 
 	case EXP_GET:
@@ -2803,6 +2894,7 @@  int main(int argc, char *argv[])
 		nfexp_callback_register(cth, NFCT_T_ALL, dump_exp_cb, NULL);
 		res = nfexp_query(cth, NFCT_Q_GET, tmpl.exp);
 		nfct_close(cth);
+		cth = NULL;
 		break;
 
 	case CT_FLUSH:
@@ -2810,9 +2902,10 @@  int main(int argc, char *argv[])
 		if (!cth)
 			exit_error(OTHER_PROBLEM, "Can't open handler");
 		res = nfct_query(cth, NFCT_Q_FLUSH, &family);
-		nfct_close(cth);
 		fprintf(stderr, "%s v%s (conntrack-tools): ",PROGNAME,VERSION);
 		fprintf(stderr,"connection tracking table has been emptied.\n");
+		nfct_close(cth);
+		cth = NULL;
 		break;
 
 	case EXP_FLUSH:
@@ -2820,9 +2913,10 @@  int main(int argc, char *argv[])
 		if (!cth)
 			exit_error(OTHER_PROBLEM, "Can't open handler");
 		res = nfexp_query(cth, NFCT_Q_FLUSH, &family);
-		nfct_close(cth);
 		fprintf(stderr, "%s v%s (conntrack-tools): ",PROGNAME,VERSION);
 		fprintf(stderr,"expectation table has been emptied.\n");
+		nfct_close(cth);
+		cth = NULL;
 		break;
 
 	case CT_EVENT:
@@ -2894,6 +2988,8 @@  int main(int argc, char *argv[])
 			res = mnl_cb_run(buf, res, 0, 0, event_cb, tmpl.ct);
 		}
 		mnl_socket_close(sock.mnl);
+		nfct_close(cth);
+		cth = NULL;
 		break;
 
 	case EXP_EVENT:
@@ -2922,6 +3018,7 @@  int main(int argc, char *argv[])
 		nfexp_callback_register(cth, NFCT_T_ALL, event_exp_cb, NULL);
 		res = nfexp_catch(cth);
 		nfct_close(cth);
+		cth = NULL;
 		break;
 	case CT_COUNT:
 		/* If we fail with netlink, fall back to /proc to ensure
@@ -2966,6 +3063,7 @@  try_proc_count:
 		nfexp_callback_register(cth, NFCT_T_ALL, count_exp_cb, NULL);
 		res = nfexp_query(cth, NFCT_Q_DUMP, &family);
 		nfct_close(cth);
+		cth = NULL;
 		printf("%d\n", counter);
 		break;
 	case CT_STATS:
@@ -3024,10 +3122,35 @@  try_proc:
 		exit_error(OTHER_PROBLEM, "Operation failed: %s",
 			   err2str(errno, command));
 
+next_opts:
 	free_tmpl_objects();
 	free_options();
-	if (labelmap)
+	if (labelmap) {
 		nfct_labelmap_destroy(labelmap);
+		labelmap = NULL;
+	}
+
+	if (opts_file) {
+		while ((res = getline(&getline_buf, &getline_buf_size, opts_file)) >= 0) {
+			if (!res)
+				continue;
+			argc = line_to_argcv(argv[0], getline_buf, &argv_buf, &argv_buf_size);
+			if (!argc)
+				continue;
+			argv = argv_buf;
+
+			cur_argc = argc;
+			cur_argv = argv;
+
+			optind = 0;
+			goto parse_opts;
+		}
+	}
+
+	if (ith)
+		nfct_close(ith);
+	if (cth)
+		nfct_close(cth);
 
 	if (command && exit_msg[cmd][0]) {
 		fprintf(stderr, "%s v%s (conntrack-tools): ",PROGNAME,VERSION);
@@ -3036,5 +3159,8 @@  try_proc:
 			return EXIT_FAILURE;
 	}
 
+	free(argv_buf);
+	free(getline_buf);
+
 	return EXIT_SUCCESS;
 }