1 files changed, 365 insertions, 1 deletions

diff --git a/src/netlink.c b/src/netlink.c
index 8ef14011..37981542 100644
--- a/src/netlink.c
+++ b/src/netlink.c
@@ -15,6 +15,8 @@
 #include <netlink/netfilter/nft_rule.h>
 #include <netlink/netfilter/nft_expr.h>
 #include <netlink/netfilter/nft_data.h>
+#include <netlink/netfilter/nft_setelem.h>
+#include <netlink/netfilter/nft_set.h>
 #include <linux/netfilter/nf_tables.h>
 
 #include <nftables.h>
@@ -30,10 +32,13 @@ static struct nl_sock *nf_sock;
 
 static void __init netlink_open_sock(void)
 {
-	// FIXME: should be done dynamically by nft_set and based on set members
 	nlmsg_set_default_size(65536);
 	nf_sock = nl_socket_alloc();
+	if (nf_sock == NULL)
+		memory_allocation_error();
+
 	nfnl_connect(nf_sock);
+	nl_socket_set_nonblocking(nf_sock);
 }
 
 static void __exit netlink_close_sock(void)
@@ -123,6 +128,42 @@ struct nfnl_nft_expr *alloc_nft_expr(int (*init)(struct nfnl_nft_expr *))
 	return nle;
 }
 
+struct nfnl_nft_set *alloc_nft_set(const struct handle *h)
+{
+	struct nfnl_nft_set *nls;
+
+	nls = nfnl_nft_set_alloc();
+	if (nls == NULL)
+		memory_allocation_error();
+	nfnl_nft_set_set_family(nls, h->family);
+	nfnl_nft_set_set_table(nls, h->table, strlen(h->table) + 1);
+	if (h->set != NULL)
+		nfnl_nft_set_set_name(nls, h->set, strlen(h->set) + 1);
+	return nls;
+}
+
+static struct nfnl_nft_setelem *alloc_nft_setelem(const struct expr *expr)
+{
+	struct nfnl_nft_setelem *nlse;
+
+	nlse = nfnl_nft_setelem_alloc();
+	if (nlse == NULL)
+		memory_allocation_error();
+
+	if (expr->ops->type == EXPR_VALUE || expr->flags & EXPR_F_INTERVAL_END)
+		nfnl_nft_setelem_set_key(nlse, netlink_gen_data(expr));
+	else {
+		assert(expr->ops->type == EXPR_MAPPING);
+		nfnl_nft_setelem_set_key(nlse, netlink_gen_data(expr->left));
+		nfnl_nft_setelem_set_data(nlse, netlink_gen_data(expr->right));
+	}
+
+	if (expr->flags & EXPR_F_INTERVAL_END)
+		nfnl_nft_setelem_set_flags(nlse, NFT_SET_ELEM_INTERVAL_END);
+
+	return nlse;
+}
+
 struct nfnl_nft_data *alloc_nft_data(const void *data, unsigned int len)
 {
 	struct nfnl_nft_data *nld;
@@ -585,3 +626,326 @@ int netlink_flush_table(struct netlink_ctx *ctx, const struct handle *h)
 {
 	return netlink_flush_rules(ctx, h);
 }
+
+static enum nft_data_types dtype_map_to_kernel(const struct datatype *dtype)
+{
+	switch (dtype->type) {
+	case TYPE_VERDICT:
+		return NFT_DATA_VERDICT;
+	default:
+		return dtype->type;
+	}
+}
+
+static const struct datatype *dtype_map_from_kernel(enum nft_data_types type)
+{
+	switch (type) {
+	case NFT_DATA_VERDICT:
+		return &verdict_type;
+	default:
+		return datatype_lookup(type);
+	}
+}
+
+static void add_set_cb(struct nl_object *obj, void *arg)
+{
+	struct nfnl_nft_set *nls = (struct nfnl_nft_set *)obj;
+	struct netlink_ctx *ctx = arg;
+	struct set *set = ctx->set;
+
+#if TRACE
+	netlink_dump_object(OBJ_CAST(nls));
+#endif
+	set->handle.set = xstrdup(nfnl_nft_set_get_name(nls));
+}
+
+static int netlink_add_set_cb(struct nl_msg *msg, void *arg)
+{
+	return nl_msg_parse(msg, add_set_cb, arg);
+}
+
+int netlink_add_set(struct netlink_ctx *ctx, const struct handle *h,
+		    struct set *set)
+{
+	struct nfnl_nft_set *nls;
+	int err;
+
+	nls = alloc_nft_set(h);
+	nfnl_nft_set_set_flags(nls, set->flags);
+	nfnl_nft_set_set_keytype(nls, dtype_map_to_kernel(set->keytype));
+	nfnl_nft_set_set_keylen(nls, set->keylen / BITS_PER_BYTE);
+	if (set->flags & NFT_SET_MAP) {
+		nfnl_nft_set_set_datatype(nls, dtype_map_to_kernel(set->datatype));
+		nfnl_nft_set_set_datalen(nls, set->datalen / BITS_PER_BYTE);
+	}
+#if TRACE
+	netlink_dump_object(OBJ_CAST(nls));
+#endif
+
+	ctx->set = set;
+	netlink_set_callback(netlink_add_set_cb, ctx);
+	err = nfnl_nft_set_add(nf_sock, nls, NLM_F_EXCL | NLM_F_ECHO);
+	if (err == 0)
+		err = nl_recvmsgs_default(nf_sock);
+	netlink_set_callback(NULL, NULL);
+	nfnl_nft_set_put(nls);
+	ctx->set = NULL;
+
+	if (err < 0)
+		netlink_io_error(ctx, NULL, "Could not add set: %s",
+				 nl_geterror(err));
+	return err;
+}
+
+int netlink_delete_set(struct netlink_ctx *ctx, const struct handle *h)
+{
+	struct nfnl_nft_set *nls;
+	int err;
+
+	nls = alloc_nft_set(h);
+	err = nfnl_nft_set_delete(nf_sock, nls, 0);
+	nfnl_nft_set_put(nls);
+
+	if (err < 0)
+		netlink_io_error(ctx, NULL, "Could not delete set: %s",
+				 nl_geterror(err));
+	return err;
+}
+
+static void list_set_cb(struct nl_object *obj, void *arg)
+{
+	struct nfnl_nft_set *nls = (struct nfnl_nft_set *)obj;
+	struct netlink_ctx *ctx = arg;
+	const struct datatype *keytype, *datatype;
+	uint32_t flags;
+	struct set *set;
+#if TRACE
+	netlink_dump_object(obj);
+#endif
+	if (!nfnl_nft_set_test_family(nls) ||
+	    !nfnl_nft_set_test_table(nls) ||
+	    !nfnl_nft_set_test_name(nls) ||
+	    !nfnl_nft_set_test_keytype(nls) ||
+	    !nfnl_nft_set_test_keylen(nls)) {
+		netlink_io_error(ctx, NULL, "Incomplete set received");
+		return;
+	}
+
+	keytype = dtype_map_from_kernel(nfnl_nft_set_get_keytype(nls));
+	if (keytype == NULL) {
+		netlink_io_error(ctx, NULL, "Unknown data type in set key %u",
+				 nfnl_nft_set_get_keytype(nls));
+		return;
+	}
+
+	flags = nfnl_nft_set_get_flags(nls);
+	if (flags & NFT_SET_MAP) {
+		datatype = dtype_map_from_kernel(nfnl_nft_set_get_datatype(nls));
+		if (datatype == NULL) {
+			netlink_io_error(ctx, NULL, "Unknown data type in set key %u",
+					 nfnl_nft_set_get_datatype(nls));
+			return;
+		}
+	} else
+		datatype = NULL;
+
+	set = set_alloc(&internal_location);
+	set->handle.family = nfnl_nft_set_get_family(nls);
+	set->handle.table  = xstrdup(nfnl_nft_set_get_table(nls));
+	set->handle.set    = xstrdup(nfnl_nft_set_get_name(nls));
+	set->keytype       = keytype;
+	set->keylen        = nfnl_nft_set_get_keylen(nls) * BITS_PER_BYTE;
+	set->flags         = flags;
+	set->datatype      = datatype;
+	set->datalen       = nfnl_nft_set_get_datalen(nls) * BITS_PER_BYTE;
+	list_add_tail(&set->list, &ctx->list);
+}
+
+int netlink_list_sets(struct netlink_ctx *ctx, const struct handle *h)
+{
+	struct nl_cache *set_cache;
+	int err;
+
+	err = nfnl_nft_set_alloc_cache(nf_sock, h->family, h->table, &set_cache);
+	if (err < 0)
+		return netlink_io_error(ctx, NULL,
+					"Could not receive sets from kernel: %s",
+					nl_geterror(err));
+
+	nl_cache_foreach(set_cache, list_set_cb, ctx);
+	nl_cache_free(set_cache);
+	return 0;
+}
+
+static int netlink_get_set_cb(struct nl_msg *msg, void *arg)
+{
+	return nl_msg_parse(msg, list_set_cb, arg);
+}
+
+int netlink_get_set(struct netlink_ctx *ctx, const struct handle *h)
+{
+	struct nfnl_nft_set *nls;
+	int err;
+
+	nls = alloc_nft_set(h);
+#if TRACE
+	netlink_dump_object(OBJ_CAST(nls));
+#endif
+	netlink_set_callback(netlink_get_set_cb, ctx);
+	err = nfnl_nft_set_query(nf_sock, nls, 0);
+	if (err == 0)
+		err = nl_recvmsgs_default(nf_sock);
+	netlink_set_callback(NULL, NULL);
+
+	nfnl_nft_set_put(nls);
+
+	if (err < 0)
+		return netlink_io_error(ctx, NULL,
+					"Could not receive set from kernel: %s",
+					nl_geterror(err));
+	return err;
+}
+
+static int alloc_setelem_cache(const struct expr *set, struct nl_cache **res)
+{
+	struct nfnl_nft_setelem *nlse;
+	struct nl_cache *elements;
+	const struct expr *expr;
+	int err;
+
+	err = nl_cache_alloc_name("netfilter/nft_setelem", &elements);
+	if (err < 0)
+		return err;
+	list_for_each_entry(expr, &set->expressions, list) {
+		nlse = alloc_nft_setelem(expr);
+		nl_cache_add(elements, OBJ_CAST(nlse));
+	}
+	*res = elements;
+	return 0;
+}
+
+int netlink_add_setelems(struct netlink_ctx *ctx, const struct handle *h,
+			 const struct expr *expr)
+{
+	struct nfnl_nft_set *nls;
+	struct nl_cache *elements;
+	int err;
+
+	nls = alloc_nft_set(h);
+#if TRACE
+	netlink_dump_object(OBJ_CAST(nls));
+#endif
+	err = alloc_setelem_cache(expr, &elements);
+	if (err < 0)
+		goto out;
+	err = nfnl_nft_setelem_add(nf_sock, nls, elements, 0);
+	if (err < 0)
+		goto out;
+	err = nl_recvmsgs_default(nf_sock);
+out:
+	nfnl_nft_set_put(nls);
+	if (err < 0)
+		netlink_io_error(ctx, NULL, "Could not add set elements: %s",
+				 nl_geterror(err));
+	return err;
+}
+
+int netlink_delete_setelems(struct netlink_ctx *ctx, const struct handle *h,
+			    const struct expr *expr)
+{
+	struct nfnl_nft_set *nls;
+	struct nl_cache *elements;
+	int err;
+
+	nls = alloc_nft_set(h);
+	err = alloc_setelem_cache(expr, &elements);
+	if (err < 0)
+		goto out;
+	err = nfnl_nft_setelem_delete(nf_sock, nls, elements, 0);
+	if (err < 0)
+		goto out;
+	err = nl_recvmsgs_default(nf_sock);
+out:
+	nfnl_nft_set_put(nls);
+	if (err < 0)
+		netlink_io_error(ctx, NULL, "Could not delete set elements: %s",
+				 nl_geterror(err));
+	return err;
+}
+
+static void list_setelem_cb(struct nl_object *obj, void *arg)
+{
+	struct nfnl_nft_setelem *nlse = nl_object_priv(obj);
+	struct nfnl_nft_data *nld;
+	struct netlink_ctx *ctx = arg;
+	struct set *set = ctx->set;
+	struct expr *expr, *data;
+	uint32_t flags;
+#if TRACE
+	netlink_dump_object(obj);
+#endif
+	if (!nfnl_nft_setelem_test_key(nlse)) {
+		netlink_io_error(ctx, NULL, "Incomplete set element received");
+		return;
+	}
+
+	nld   = nfnl_nft_setelem_get_key(nlse);
+	flags = nfnl_nft_setelem_get_flags(nlse);
+
+	expr = netlink_alloc_value(&internal_location, nld);
+	expr->dtype	= set->keytype;
+	expr->byteorder	= set->keytype->byteorder;
+	if (expr->byteorder == BYTEORDER_HOST_ENDIAN)
+		mpz_switch_byteorder(expr->value, expr->len / BITS_PER_BYTE);
+
+	if (flags & NFT_SET_ELEM_INTERVAL_END)
+		expr->flags |= EXPR_F_INTERVAL_END;
+	else if (nfnl_nft_setelem_test_data(nlse)) {
+		nld = nfnl_nft_setelem_get_data(nlse);
+
+		data = netlink_alloc_data(&internal_location, nld,
+					  set->datatype->type == EXPR_VERDICT ?
+					  NFT_REG_VERDICT : NFT_REG_1);
+		data->dtype = set->datatype;
+
+		expr = mapping_expr_alloc(&internal_location, expr, data);
+	}
+
+	compound_expr_add(set->init, expr);
+}
+
+extern void interval_map_decompose(struct expr *set);
+
+int netlink_get_setelems(struct netlink_ctx *ctx, const struct handle *h,
+			 struct set *set)
+{
+	struct nl_cache *elements;
+	struct nfnl_nft_set *nls;
+	int err;
+
+	nls = alloc_nft_set(h);
+#if TRACE
+	netlink_dump_object(OBJ_CAST(nls));
+#endif
+	err = nfnl_nft_setelem_alloc_cache(nf_sock, nls, &elements);
+	if (err < 0)
+		goto out;
+	err = nl_recvmsgs_default(nf_sock);
+	if (err < 0)
+		goto out;
+
+	ctx->set = set;
+	set->init = set_expr_alloc(&internal_location);
+	nl_cache_foreach(elements, list_setelem_cb, ctx);
+	nl_cache_free(elements);
+	ctx->set = NULL;
+
+	if (set->flags & NFT_SET_INTERVAL)
+		interval_map_decompose(set->init);
+out:
+	nfnl_nft_set_put(nls);
+	if (err < 0)
+		netlink_io_error(ctx, NULL, "Could not receive set elements: %s",
+				 nl_geterror(err));
+	return err;
+}