expr: ct: Add support for setting the mark

This patch adds userspace support for setting properties of tracked connections.
Currently, the connection mark is supported. This can be used to implemented the
same functionality as iptables -j CONNMARK --save-mark.

Signed-off-by: Kristian Evensen <kristian.evensen@gmail.com>
Signed-off-by: Pablo Neira Ayuso <pablo@netfilter.org>

3 files changed, 104 insertions, 26 deletions

diff --git a/include/libnftables/expr.h b/include/libnftables/expr.h
index 25455e4..653bbb0 100644
--- a/include/libnftables/expr.h
+++ b/include/libnftables/expr.h
@@ -124,6 +124,7 @@ enum {
 	NFT_EXPR_CT_DREG	= NFT_RULE_EXPR_ATTR_BASE,
 	NFT_EXPR_CT_KEY,
 	NFT_EXPR_CT_DIR,
+	NFT_EXPR_CT_SREG,
 };
 
 enum {
diff --git a/include/linux/netfilter/nf_tables.h b/include/linux/netfilter/nf_tables.h
index fc0f669..6a22a37 100644
--- a/include/linux/netfilter/nf_tables.h
+++ b/include/linux/netfilter/nf_tables.h
@@ -530,12 +530,14 @@ enum nft_ct_keys {
  * @NFTA_CT_DREG: destination register (NLA_U32)
  * @NFTA_CT_KEY: conntrack data item to load (NLA_U32: nft_ct_keys)
  * @NFTA_CT_DIRECTION: direction in case of directional keys (NLA_U8)
+ * @NFTA_CT_SREG: source register (NLA_U32)
  */
 enum nft_ct_attributes {
 	NFTA_CT_UNSPEC,
 	NFTA_CT_DREG,
 	NFTA_CT_KEY,
 	NFTA_CT_DIRECTION,
+	NFTA_CT_SREG,
 	__NFTA_CT_MAX
 };
 #define NFTA_CT_MAX		(__NFTA_CT_MAX - 1)
diff --git a/src/expr/ct.c b/src/expr/ct.c
index 46e3cef..49d8495 100644
--- a/src/expr/ct.c
+++ b/src/expr/ct.c
@@ -24,7 +24,8 @@
 
 struct nft_expr_ct {
 	enum nft_ct_keys        key;
-	uint32_t		dreg;	/* enum nft_registers */
+	enum nft_registers	dreg;
+	enum nft_registers	sreg;
 	uint8_t			dir;
 };
 
@@ -51,6 +52,9 @@ nft_rule_expr_ct_set(struct nft_rule_expr *e, uint16_t type,
 	case NFT_EXPR_CT_DREG:
 		ct->dreg = *((uint32_t *)data);
 		break;
+	case NFT_EXPR_CT_SREG:
+		ct->sreg = *((uint32_t *)data);
+		break;
 	default:
 		return -1;
 	}
@@ -73,6 +77,9 @@ nft_rule_expr_ct_get(const struct nft_rule_expr *e, uint16_t type,
 	case NFT_EXPR_CT_DREG:
 		*data_len = sizeof(ct->dreg);
 		return &ct->dreg;
+	case NFT_EXPR_CT_SREG:
+		*data_len = sizeof(ct->sreg);
+		return &ct->sreg;
 	}
 	return NULL;
 }
@@ -88,6 +95,7 @@ static int nft_rule_expr_ct_cb(const struct nlattr *attr, void *data)
 	switch(type) {
 	case NFTA_CT_KEY:
 	case NFTA_CT_DREG:
+	case NFTA_CT_SREG:
 		if (mnl_attr_validate(attr, MNL_TYPE_U32) < 0) {
 			perror("mnl_attr_validate");
 			return MNL_CB_ERROR;
@@ -116,6 +124,8 @@ nft_rule_expr_ct_build(struct nlmsghdr *nlh, struct nft_rule_expr *e)
 		mnl_attr_put_u32(nlh, NFTA_CT_DREG, htonl(ct->dreg));
 	if (e->flags & (1 << NFT_EXPR_CT_DIR))
 		mnl_attr_put_u8(nlh, NFTA_CT_DIRECTION, ct->dir);
+	if (e->flags & (1 << NFT_EXPR_CT_SREG))
+		mnl_attr_put_u32(nlh, NFTA_CT_SREG, htonl(ct->sreg));
 }
 
 static int
@@ -131,13 +141,17 @@ nft_rule_expr_ct_parse(struct nft_rule_expr *e, struct nlattr *attr)
 		ct->key = ntohl(mnl_attr_get_u32(tb[NFTA_CT_KEY]));
 		e->flags |= (1 << NFT_EXPR_CT_KEY);
 	}
+	if (tb[NFTA_CT_DIRECTION]) {
+		ct->dir = mnl_attr_get_u8(tb[NFTA_CT_DIRECTION]);
+		e->flags |= (1 << NFT_EXPR_CT_DIR);
+	}
 	if (tb[NFTA_CT_DREG]) {
 		ct->dreg = ntohl(mnl_attr_get_u32(tb[NFTA_CT_DREG]));
 		e->flags |= (1 << NFT_EXPR_CT_DREG);
 	}
-	if (tb[NFTA_CT_DIRECTION]) {
-		ct->dir = mnl_attr_get_u8(tb[NFTA_CT_DIRECTION]);
-		e->flags |= (1 << NFT_EXPR_CT_DIR);
+	if (tb[NFTA_CT_SREG]) {
+		ct->sreg = ntohl(mnl_attr_get_u32(tb[NFTA_CT_SREG]));
+		e->flags |= (1 << NFT_EXPR_CT_SREG);
 	}
 
 	return 0;
@@ -186,10 +200,19 @@ static int nft_rule_expr_ct_json_parse(struct nft_rule_expr *e, json_t *root)
 	uint8_t dir;
 	int key;
 
-	if (nft_jansson_parse_reg(root, "dreg", NFT_TYPE_U32, &reg) < 0)
-		return -1;
+	if (nft_jansson_node_exist(root, "dreg")) {
+		if (nft_jansson_parse_reg(root, "dreg", NFT_TYPE_U32, &reg) < 0)
+			return -1;
+
+		nft_rule_expr_set_u32(e, NFT_EXPR_CT_DREG, reg);
+	}
+
+	if (nft_jansson_node_exist(root, "sreg")) {
+		if (nft_jansson_parse_reg(root, "sreg", NFT_TYPE_U32, &reg) < 0)
+			return -1;
 
-	nft_rule_expr_set_u32(e, NFT_EXPR_CT_DREG, reg);
+		nft_rule_expr_set_u32(e, NFT_EXPR_CT_SREG, reg);
+	}
 
 	if (nft_jansson_node_exist(root, "key")) {
 		key_str = nft_jansson_parse_str(root, "key");
@@ -235,11 +258,16 @@ static int nft_rule_expr_ct_xml_parse(struct nft_rule_expr *e, mxml_node_t *tree
 	uint8_t dir;
 
 	reg = nft_mxml_reg_parse(tree, "dreg", MXML_DESCEND_FIRST);
-	if (reg < 0)
-		return -1;
+	if (reg >= 0) {
+		ct->dreg = reg;
+		e->flags |= (1 << NFT_EXPR_CT_DREG);
+	}
 
-	ct->dreg = reg;
-	e->flags |= (1 << NFT_EXPR_CT_DREG);
+	reg = nft_mxml_reg_parse(tree, "sreg", MXML_DESCEND_FIRST);
+	if (reg >= 0) {
+		ct->sreg = reg;
+		e->flags |= (1 << NFT_EXPR_CT_SREG);
+	}
 
 	key_str = nft_mxml_str_parse(tree, "key", MXML_DESCEND_FIRST,
 				     NFT_XML_MAND);
@@ -274,46 +302,93 @@ err:
 }
 
 static int
-nft_expr_ct_snprintf_json(char *buf, size_t size, struct nft_rule_expr *e)
+nft_rule_expr_ct_snprintf_default(char *buf, size_t size,
+				  struct nft_rule_expr *e)
+{
+	struct nft_expr_ct *ct = nft_expr_data(e);
+
+	if (e->flags & (1 << NFT_EXPR_CT_SREG))
+		return snprintf(buf, size, "set %s with reg %u ",
+				ctkey2str(ct->key), ct->sreg);
+
+	return snprintf(buf, size, "load %s => reg %u dir %u ",
+			ctkey2str(ct->key), ct->dreg, ct->dir);
+}
+
+static int
+nft_rule_expr_ct_snprintf_xml(char *buf, size_t size, struct nft_rule_expr *e)
+{
+	int ret, len = size, offset = 0;
+	struct nft_expr_ct *ct = nft_expr_data(e);
+
+	if (e->flags & (1 << NFT_EXPR_CT_DREG)) {
+		ret = snprintf(buf+offset, len, "<dreg>%u</dreg>", ct->dreg);
+		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
+	}
+
+	if (e->flags & (1 << NFT_EXPR_CT_SREG)) {
+		ret = snprintf(buf+offset, len, "<sreg>%u</sreg>", ct->sreg);
+		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
+	}
+
+	if (e->flags & (1 << NFT_EXPR_CT_KEY)) {
+		ret = snprintf(buf+offset, len, "<key>%s</key>",
+			       ctkey2str(ct->key));
+		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
+	}
+
+	if (e->flags & (1 << NFT_EXPR_CT_DIR)) {
+		ret = snprintf(buf+offset, len, "<dir>%u</dir>", ct->dir);
+		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
+	}
+
+	return offset;
+}
+
+static int
+nft_rule_expr_ct_snprintf_json(char *buf, size_t size, struct nft_rule_expr *e)
 {
 	int ret, len = size, offset = 0;
 	struct nft_expr_ct *ct = nft_expr_data(e);
 
-	ret = snprintf(buf, len, "\"dreg\":%u", ct->dreg);
-	SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
+	if (e->flags & (1 << NFT_EXPR_CT_DREG)) {
+		ret = snprintf(buf+offset, len, "\"dreg\":%u,", ct->dreg);
+		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
+	}
+
+	if (e->flags & (1 << NFT_EXPR_CT_SREG)) {
+		ret = snprintf(buf+offset, len, "\"sreg:\":%u,", ct->sreg);
+		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
+	}
 
 	if (e->flags & (1 << NFT_EXPR_CT_KEY)) {
-		ret = snprintf(buf+offset, len, ",\"key\":\"%s\"",
+		ret = snprintf(buf+offset, len, "\"key\":\"%s\",",
 						ctkey2str(ct->key));
 		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
 	}
 
 	if (e->flags & (1 << NFT_EXPR_CT_DIR)) {
-		ret = snprintf(buf+offset, len, ",\"dir\":%u", ct->dir);
+		ret = snprintf(buf+offset, len, "\"dir\":%u,", ct->dir);
 		SNPRINTF_BUFFER_SIZE(ret, size, len, offset);
 	}
 
-	return offset;
+	/* Remove the last separator characther  */
+	buf[offset-1] = '\0';
 
+	return offset-1;
 }
 
 static int
 nft_rule_expr_ct_snprintf(char *buf, size_t len, uint32_t type,
 			    uint32_t flags, struct nft_rule_expr *e)
 {
-	struct nft_expr_ct *ct = nft_expr_data(e);
-
 	switch(type) {
 	case NFT_OUTPUT_DEFAULT:
-		return snprintf(buf, len, "load %s => reg %u dir %u ",
-				ctkey2str(ct->key), ct->dreg, ct->dir);
+		return nft_rule_expr_ct_snprintf_default(buf, len, e);
 	case NFT_OUTPUT_XML:
-		return snprintf(buf, len, "<dreg>%u</dreg>"
-					  "<key>%s</key>"
-					  "<dir>%u</dir>",
-				ct->dreg, ctkey2str(ct->key), ct->dir);
+		return nft_rule_expr_ct_snprintf_xml(buf, len, e);
 	case NFT_OUTPUT_JSON:
-		return nft_expr_ct_snprintf_json(buf, len, e);
+		return nft_rule_expr_ct_snprintf_json(buf, len, e);
 	default:
 		break;
 	}