xtables-restore: Introduce rule counter tokenizer function

The same piece of code appears three times, introduce a function to take
care of tokenizing and error reporting.

Pass buffer pointer via reference so it can be updated to point to after
the counters (if found).

While being at it, drop pointless casting when passing pcnt/bcnt to
add_argv().

Signed-off-by: Phil Sutter <phil@nwl.cc>
Acked-by: Pablo Neira Ayuso <pablo@netfilter.org>

6 files changed, 70 insertions, 91 deletions

diff --git a/iptables/iptables-restore.c b/iptables/iptables-restore.c
index 6bc182bf..3655b3e8 100644
--- a/iptables/iptables-restore.c
+++ b/iptables/iptables-restore.c
@@ -313,44 +313,17 @@ ip46tables_restore_main(struct iptables_restore_cb *cb, int argc, char *argv[])
 			int a;
 			char *pcnt = NULL;
 			char *bcnt = NULL;
-			char *parsestart;
-
-			if (buffer[0] == '[') {
-				/* we have counters in our input */
-				char *ptr = strchr(buffer, ']');
-
-				if (!ptr)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need ]\n",
-						   line);
-
-				pcnt = strtok(buffer+1, ":");
-				if (!pcnt)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need :\n",
-						   line);
-
-				bcnt = strtok(NULL, "]");
-				if (!bcnt)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need ]\n",
-						   line);
-
-				/* start command parsing after counter */
-				parsestart = ptr + 1;
-			} else {
-				/* start command parsing at start of line */
-				parsestart = buffer;
-			}
+			char *parsestart = buffer;
 
 			add_argv(argv[0], 0);
 			add_argv("-t", 0);
 			add_argv(curtable, 0);
 
+			tokenize_rule_counters(&parsestart, &pcnt, &bcnt, line);
 			if (counters && pcnt && bcnt) {
 				add_argv("--set-counters", 0);
-				add_argv((char *) pcnt, 0);
-				add_argv((char *) bcnt, 0);
+				add_argv(pcnt, 0);
+				add_argv(bcnt, 0);
 			}
 
 			add_param_to_argv(parsestart, line);
diff --git a/iptables/iptables-xml.c b/iptables/iptables-xml.c
index 36ad7845..5255e097 100644
--- a/iptables/iptables-xml.c
+++ b/iptables/iptables-xml.c
@@ -644,7 +644,7 @@ iptables_xml_main(int argc, char *argv[])
 			unsigned int a;
 			char *pcnt = NULL;
 			char *bcnt = NULL;
-			char *parsestart;
+			char *parsestart = buffer;
 			char *chain = NULL;
 
 			/* the parser */
@@ -652,34 +652,7 @@ iptables_xml_main(int argc, char *argv[])
 			int quote_open, quoted;
 			char param_buffer[1024];
 
-			if (buffer[0] == '[') {
-				/* we have counters in our input */
-				char *ptr = strchr(buffer, ']');
-
-				if (!ptr)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need ]\n",
-						   line);
-
-				pcnt = strtok(buffer + 1, ":");
-				if (!pcnt)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need :\n",
-						   line);
-
-				bcnt = strtok(NULL, "]");
-				if (!bcnt)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need ]\n",
-						   line);
-
-				/* start command parsing after counter */
-				parsestart = ptr + 1;
-			} else {
-				/* start command parsing at start of line */
-				parsestart = buffer;
-			}
-
+			tokenize_rule_counters(&parsestart, &pcnt, &bcnt, line);
 
 			/* This is a 'real' parser crafted in artist mode
 			 * not hacker mode. If the author can live with that
diff --git a/iptables/tests/shell/testcases/ipt-restore/0008-restore-counters_0 b/iptables/tests/shell/testcases/ipt-restore/0008-restore-counters_0
new file mode 100755
index 00000000..5ac70682
--- /dev/null
+++ b/iptables/tests/shell/testcases/ipt-restore/0008-restore-counters_0
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+set -e
+
+DUMP="*filter
+:foo - [23:42]
+[13:37] -A foo -j ACCEPT
+COMMIT
+"
+
+EXPECT=":foo - [0:0]
+[0:0] -A foo -j ACCEPT"
+
+$XT_MULTI iptables-restore <<< "$DUMP"
+diff -u -Z <(echo -e "$EXPECT") <($XT_MULTI iptables-save --counters | grep foo)
+
+# iptables-*-restore ignores custom chain counters :(
+EXPECT=":foo - [0:0]
+[13:37] -A foo -j ACCEPT"
+
+$XT_MULTI iptables-restore --counters <<< "$DUMP"
+diff -u -Z <(echo -e "$EXPECT") <($XT_MULTI iptables-save --counters | grep foo)
diff --git a/iptables/xshared.c b/iptables/xshared.c
index 5e6cd4ae..ba723f59 100644
--- a/iptables/xshared.c
+++ b/iptables/xshared.c
@@ -373,6 +373,43 @@ int parse_counters(const char *string, struct xt_counters *ctr)
 	return ret == 2;
 }
 
+/* Tokenize counters argument of typical iptables-restore format rule.
+ *
+ * If *bufferp contains counters, update *pcntp and *bcntp to point at them,
+ * change bytes after counters in *bufferp to nul-bytes, update *bufferp to
+ * point to after the counters and return true.
+ * If *bufferp does not contain counters, return false.
+ * If syntax is wrong in *bufferp, call xtables_error() and hence exit().
+ * */
+bool tokenize_rule_counters(char **bufferp, char **pcntp, char **bcntp, int line)
+{
+	char *ptr, *buffer = *bufferp, *pcnt, *bcnt;
+
+	if (buffer[0] != '[')
+		return false;
+
+	/* we have counters in our input */
+
+	ptr = strchr(buffer, ']');
+	if (!ptr)
+		xtables_error(PARAMETER_PROBLEM, "Bad line %u: need ]\n", line);
+
+	pcnt = strtok(buffer+1, ":");
+	if (!pcnt)
+		xtables_error(PARAMETER_PROBLEM, "Bad line %u: need :\n", line);
+
+	bcnt = strtok(NULL, "]");
+	if (!bcnt)
+		xtables_error(PARAMETER_PROBLEM, "Bad line %u: need ]\n", line);
+
+	*pcntp = pcnt;
+	*bcntp = bcnt;
+	/* start command parsing after counter */
+	*bufferp = ptr + 1;
+
+	return true;
+}
+
 inline bool xs_has_arg(int argc, char *argv[])
 {
 	return optind < argc &&
diff --git a/iptables/xshared.h b/iptables/xshared.h
index b08c700e..21f4e8fd 100644
--- a/iptables/xshared.h
+++ b/iptables/xshared.h
@@ -151,6 +151,7 @@ extern int xtables_lock_or_exit(int wait, struct timeval *tv);
 int parse_wait_time(int argc, char *argv[]);
 void parse_wait_interval(int argc, char *argv[], struct timeval *wait_interval);
 int parse_counters(const char *string, struct xt_counters *ctr);
+bool tokenize_rule_counters(char **bufferp, char **pcnt, char **bcnt, int line);
 bool xs_has_arg(int argc, char *argv[]);
 
 extern const struct xtables_afinfo *afinfo;
diff --git a/iptables/xtables-restore.c b/iptables/xtables-restore.c
index 24bfce51..4652d631 100644
--- a/iptables/xtables-restore.c
+++ b/iptables/xtables-restore.c
@@ -230,47 +230,20 @@ void xtables_restore_parse(struct nft_handle *h,
 			int a;
 			char *pcnt = NULL;
 			char *bcnt = NULL;
-			char *parsestart;
+			char *parsestart = buffer;
 
 			/* reset the newargv */
 			newargc = 0;
 
-			if (buffer[0] == '[') {
-				/* we have counters in our input */
-				char *ptr = strchr(buffer, ']');
-
-				if (!ptr)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need ]\n",
-						   line);
-
-				pcnt = strtok(buffer+1, ":");
-				if (!pcnt)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need :\n",
-						   line);
-
-				bcnt = strtok(NULL, "]");
-				if (!bcnt)
-					xtables_error(PARAMETER_PROBLEM,
-						   "Bad line %u: need ]\n",
-						   line);
-
-				/* start command parsing after counter */
-				parsestart = ptr + 1;
-			} else {
-				/* start command parsing at start of line */
-				parsestart = buffer;
-			}
-
 			add_argv(xt_params->program_name, 0);
 			add_argv("-t", 0);
 			add_argv(curtable->name, 0);
 
+			tokenize_rule_counters(&parsestart, &pcnt, &bcnt, line);
 			if (counters && pcnt && bcnt) {
 				add_argv("--set-counters", 0);
-				add_argv((char *) pcnt, 0);
-				add_argv((char *) bcnt, 0);
+				add_argv(pcnt, 0);
+				add_argv(bcnt, 0);
 			}
 
 			add_param_to_argv(parsestart, line);