nlmsg: add new message batching infrastructure

This patch adds the new message batching infrastructure that allows to
store several messages into one single datagram.

This patch includes an example for ctnetlink.

Signed-off-by: Pablo Neira Ayuso <pablo@netfilter.org>

5 files changed, 401 insertions, 1 deletions

diff --git a/examples/netfilter/Makefile.am b/examples/netfilter/Makefile.am
index 96466e0..98458fd 100644
--- a/examples/netfilter/Makefile.am
+++ b/examples/netfilter/Makefile.am
@@ -2,7 +2,8 @@ include $(top_srcdir)/Make_global.am
 
 check_PROGRAMS = nf-queue \
 		 nf-log \
-		 nfct-event
+		 nfct-event \
+		 nfct-create-batch
 
 nf_queue_SOURCES = nf-queue.c
 nf_queue_LDADD = ../../src/libmnl.la
@@ -12,3 +13,6 @@ nf_log_LDADD = ../../src/libmnl.la
 
 nfct_event_SOURCES = nfct-event.c
 nfct_event_LDADD = ../../src/libmnl.la
+
+nfct_create_batch_SOURCES = nfct-create-batch.c
+nfct_create_batch_LDADD = ../../src/libmnl.la
diff --git a/examples/netfilter/nfct-create-batch.c b/examples/netfilter/nfct-create-batch.c
new file mode 100644
index 0000000..36ad8c7
--- /dev/null
+++ b/examples/netfilter/nfct-create-batch.c
@@ -0,0 +1,192 @@
+/*
+ * (C) 2010 by Pablo Neira Ayuso <pablo@netfilter.org>
+ *
+ * This software may be used and distributed according to the terms
+ * of the GNU General Public License, incorporated herein by reference.
+ */
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <arpa/inet.h>
+#include <time.h>
+#include <sys/select.h>
+#include <string.h>
+
+#include <libmnl/libmnl.h>
+#include <linux/netfilter/nfnetlink.h>
+#include <linux/netfilter/nfnetlink_conntrack.h>
+#include <linux/netfilter/nf_conntrack_common.h>
+#include <linux/netfilter/nf_conntrack_tcp.h>
+
+static void put_msg(char *buf, uint16_t i, int seq)
+{
+	struct nlmsghdr *nlh;
+	struct nfgenmsg *nfh;
+	struct nlattr *nest1, *nest2;
+
+	nlh = mnl_nlmsg_put_header(buf);
+	nlh->nlmsg_type = (NFNL_SUBSYS_CTNETLINK << 8) | IPCTNL_MSG_CT_NEW;
+	nlh->nlmsg_flags = NLM_F_REQUEST|NLM_F_CREATE|NLM_F_EXCL|NLM_F_ACK;
+	nlh->nlmsg_seq = seq;
+
+	nfh = mnl_nlmsg_put_extra_header(nlh, sizeof(struct nfgenmsg));
+	nfh->nfgen_family = AF_INET;
+	nfh->version = NFNETLINK_V0;
+	nfh->res_id = 0;
+
+	nest1 = mnl_attr_nest_start(nlh, CTA_TUPLE_ORIG);
+	nest2 = mnl_attr_nest_start(nlh, CTA_TUPLE_IP);
+	mnl_attr_put_u32(nlh, CTA_IP_V4_SRC, inet_addr("1.1.1.1"));
+	mnl_attr_put_u32(nlh, CTA_IP_V4_DST, inet_addr("2.2.2.2"));
+	mnl_attr_nest_end(nlh, nest2);
+
+	nest2 = mnl_attr_nest_start(nlh, CTA_TUPLE_PROTO);
+	mnl_attr_put_u8(nlh, CTA_PROTO_NUM, IPPROTO_TCP);
+	mnl_attr_put_u16(nlh, CTA_PROTO_SRC_PORT, htons(i));
+	mnl_attr_put_u16(nlh, CTA_PROTO_DST_PORT, htons(1025));
+	mnl_attr_nest_end(nlh, nest2);
+	mnl_attr_nest_end(nlh, nest1);
+
+	nest1 = mnl_attr_nest_start(nlh, CTA_TUPLE_REPLY);
+	nest2 = mnl_attr_nest_start(nlh, CTA_TUPLE_IP);
+	mnl_attr_put_u32(nlh, CTA_IP_V4_SRC, inet_addr("2.2.2.2"));
+	mnl_attr_put_u32(nlh, CTA_IP_V4_DST, inet_addr("1.1.1.1"));
+	mnl_attr_nest_end(nlh, nest2);
+
+	nest2 = mnl_attr_nest_start(nlh, CTA_TUPLE_PROTO);
+	mnl_attr_put_u8(nlh, CTA_PROTO_NUM, IPPROTO_TCP);
+	mnl_attr_put_u16(nlh, CTA_PROTO_SRC_PORT, htons(1025));
+	mnl_attr_put_u16(nlh, CTA_PROTO_DST_PORT, htons(i));
+	mnl_attr_nest_end(nlh, nest2);
+	mnl_attr_nest_end(nlh, nest1);
+
+	nest1 = mnl_attr_nest_start(nlh, CTA_PROTOINFO);
+	nest2 = mnl_attr_nest_start(nlh, CTA_PROTOINFO_TCP);
+	mnl_attr_put_u8(nlh, CTA_PROTOINFO_TCP_STATE, TCP_CONNTRACK_SYN_SENT);
+	mnl_attr_nest_end(nlh, nest2);
+	mnl_attr_nest_end(nlh, nest1);
+
+	mnl_attr_put_u32(nlh, CTA_STATUS, htonl(IPS_CONFIRMED));
+	mnl_attr_put_u32(nlh, CTA_TIMEOUT, htonl(1000));
+}
+
+static int cb_err(const struct nlmsghdr *nlh, void *data)
+{
+	struct nlmsgerr *err = (void *)(nlh + 1);
+	if (err->error != 0)
+		printf("message with seq %u has failed: %s\n",
+			nlh->nlmsg_seq, strerror(-err->error));
+	return MNL_CB_OK;
+}
+
+static mnl_cb_t cb_ctl_array[NLMSG_MIN_TYPE] = {
+	[NLMSG_ERROR] = cb_err,
+};
+
+static void
+send_batch(struct mnl_socket *nl, struct mnl_nlmsg_batch *b, int portid)
+{
+	int ret, fd = mnl_socket_get_fd(nl);
+	size_t len = mnl_nlmsg_batch_size(b);
+	char rcv_buf[MNL_SOCKET_BUFFER_SIZE];
+
+	ret = mnl_socket_sendto(nl, mnl_nlmsg_batch_head(b), len);
+	if (ret == -1) {
+		perror("mnl_socket_recvfrom");
+		exit(EXIT_FAILURE);
+	}
+
+	/* receive and digest all the acknowledgments from the kernel. */
+	struct timeval tv = {
+		.tv_sec		= 0,
+		.tv_usec	= 0
+	};
+	fd_set readfds;
+	FD_ZERO(&readfds);
+	FD_SET(fd, &readfds);
+
+	ret = select(fd+1, &readfds, NULL, NULL, &tv);
+	if (ret == -1) {
+		perror("select");
+		exit(EXIT_FAILURE);
+	}
+	while (ret > 0 && FD_ISSET(fd, &readfds)) {
+		ret = mnl_socket_recvfrom(nl, rcv_buf, sizeof(rcv_buf));
+		if (ret == -1) {
+			perror("mnl_socket_recvfrom");
+			exit(EXIT_FAILURE);
+		}
+
+		ret = mnl_cb_run2(rcv_buf, ret, 0, portid,
+				  NULL, NULL, cb_ctl_array,
+				  MNL_ARRAY_SIZE(cb_ctl_array));
+		if (ret == -1) {
+			perror("mnl_cb_run");
+			exit(EXIT_FAILURE);
+		}
+
+		ret = select(fd+1, &readfds, NULL, NULL, &tv);
+		if (ret == -1) {
+			perror("select");
+			exit(EXIT_FAILURE);
+		}
+		FD_ZERO(&readfds);
+		FD_SET(fd, &readfds);
+	}
+}
+
+int main(void)
+{
+	struct mnl_socket *nl;
+	char snd_buf[MNL_SOCKET_BUFFER_SIZE*2];
+	struct mnl_nlmsg_batch *b;
+	int j;
+	unsigned int seq, portid;
+	uint16_t i;
+
+	nl = mnl_socket_open(NETLINK_NETFILTER);
+	if (nl == NULL) {
+		perror("mnl_socket_open");
+		exit(EXIT_FAILURE);
+	}
+	if (mnl_socket_bind(nl, 0, MNL_SOCKET_AUTOPID) < 0) {
+		perror("mnl_socket_bind");
+		exit(EXIT_FAILURE);
+	}
+	portid = mnl_socket_get_portid(nl);
+
+	/* The buffer that we use to batch messages is MNL_SOCKET_BUFFER_SIZE
+	 * multiplied by 2 bytes long, but we limit the batch to half of it
+	 * since the last message that does not fit the batch goes over the
+	 * upper boundary, if you break this rule, expect memory corruptions. */
+	b = mnl_nlmsg_batch_start(snd_buf, MNL_SOCKET_BUFFER_SIZE);
+	if (b == NULL) {
+		perror("mnl_nlmsg_batch_start");
+		exit(EXIT_FAILURE);
+	}
+
+	seq = time(NULL);
+	for (i=1024, j=0; i<65535; i++, j++) {
+		put_msg(mnl_nlmsg_batch_current(b), i, seq+j);
+
+		/* is there room for more messages in this batch?
+		 * if so, continue. */
+		if (mnl_nlmsg_batch_next(b))
+			continue;
+
+		send_batch(nl, b, portid);
+
+		/* this moves the last message that did not fit into the
+		 * batch to the head of it. */
+		mnl_nlmsg_batch_reset(b);
+	}
+
+	/* check if there is any message in the batch not sent yet. */
+	if (!mnl_nlmsg_batch_is_empty(b))
+		send_batch(nl, b, portid);
+
+	mnl_nlmsg_batch_stop(b);
+	mnl_socket_close(nl);
+
+	return 0;
+}
diff --git a/include/libmnl/libmnl.h b/include/libmnl/libmnl.h
index 1232e1a..78fd0c4 100644
--- a/include/libmnl/libmnl.h
+++ b/include/libmnl/libmnl.h
@@ -68,6 +68,17 @@ extern void *mnl_nlmsg_get_payload_tail(const struct nlmsghdr *nlh);
 /* Netlink message printer */
 extern void mnl_nlmsg_fprintf(FILE *fd, const void *data, size_t datalen, size_t extra_header_size);
 
+/* Message batch helpers */
+struct mnl_nlmsg_batch;
+extern struct mnl_nlmsg_batch *mnl_nlmsg_batch_start(void *buf, size_t bufsiz);
+extern bool mnl_nlmsg_batch_next(struct mnl_nlmsg_batch *b);
+extern void mnl_nlmsg_batch_stop(struct mnl_nlmsg_batch *b);
+extern size_t mnl_nlmsg_batch_size(struct mnl_nlmsg_batch *b);
+extern void mnl_nlmsg_batch_reset(struct mnl_nlmsg_batch *b);
+extern void *mnl_nlmsg_batch_head(struct mnl_nlmsg_batch *b);
+extern void *mnl_nlmsg_batch_current(struct mnl_nlmsg_batch *b);
+extern bool mnl_nlmsg_batch_is_empty(struct mnl_nlmsg_batch *b);
+
 /*
  * Netlink attributes API
  */
diff --git a/src/libmnl.map b/src/libmnl.map
index 01bb924..3147ae0 100644
--- a/src/libmnl.map
+++ b/src/libmnl.map
@@ -48,6 +48,14 @@ global:
   mnl_nlmsg_put_header;
   mnl_nlmsg_seq_ok;
   mnl_nlmsg_size;
+  mnl_nlmsg_batch_start;
+  mnl_nlmsg_batch_stop;
+  mnl_nlmsg_batch_next;
+  mnl_nlmsg_batch_size;
+  mnl_nlmsg_batch_reset;
+  mnl_nlmsg_batch_current;
+  mnl_nlmsg_batch_head;
+  mnl_nlmsg_batch_is_empty;
   mnl_socket_bind;
   mnl_socket_close;
   mnl_socket_get_fd;
diff --git a/src/nlmsg.c b/src/nlmsg.c
index 4f5e503..b2e3853 100644
--- a/src/nlmsg.c
+++ b/src/nlmsg.c
@@ -8,6 +8,7 @@
  */
 #include <stdbool.h>
 #include <stdio.h>
+#include <stdlib.h>
 #include <ctype.h>
 #include <errno.h>
 #include <string.h>
@@ -371,5 +372,189 @@ mnl_nlmsg_fprintf(FILE *fd, const void *data, size_t datalen,
 }
 
 /**
+ * \defgroup batch Netlink message batch helpers
+ *
+ * This library provides helpers to batch several messages into one single
+ * datagram. These helpers do not perform strict memory boundary checkings.
+ *
+ * The following figure represents a Netlink message batch:
+ *
+ *   |<-------------- MNL_SOCKET_BUFFER_SIZE ------------->|
+ *   |<-------------------- batch ------------------>|     |
+ *   |-----------|-----------|-----------|-----------|-----------|
+ *   |<- nlmsg ->|<- nlmsg ->|<- nlmsg ->|<- nlmsg ->|<- nlmsg ->|
+ *   |-----------|-----------|-----------|-----------|-----------|
+ *                                             ^           ^
+ *                                             |           |
+ *                                        message N   message N+1
+ *
+ * To start the batch, you have to call mnl_nlmsg_batch_start() and you can
+ * use mnl_nlmsg_batch_stop() to release it.
+ *
+ * You have to invoke mnl_nlmsg_batch_next() to get room for a new message
+ * in the batch. If this function returns NULL, it means that the last
+ * message that was added (message N+1 in the figure above) does not fit the
+ * batch. Thus, you have to send the batch (which includes until message N)
+ * and, then, you have to call mnl_nlmsg_batch_reset() to re-initialize
+ * the batch (this moves message N+1 to the head of the buffer). For that
+ * reason, the buffer that you have to use to store the batch must be double
+ * of MNL_SOCKET_BUFFER_SIZE to ensure that the last message (message N+1)
+ * that did not fit into the batch is written inside valid memory boundaries.
+ *
+ * @{
+ */
+
+struct mnl_nlmsg_batch {
+	/* the buffer that is used to store the batch. */
+	void *buf;
+	size_t limit;
+	size_t buflen;
+	/* the current netlink message in the batch. */
+	void *cur;
+	bool overflow;
+};
+
+/**
+ * mnl_nlmsg_batch_start - initialize a batch
+ * \param buf pointer to the buffer that will store this batch
+ * \param limit maximum size of the batch (should be MNL_SOCKET_BUFFER_SIZE).
+ *
+ * The buffer that you pass must be double of MNL_SOCKET_BUFFER_SIZE. The
+ * limit must be half of the buffer size, otherwise expect funny memory
+ * corruptions 8-).
+ *
+ * You can allocate the buffer that you use to store the batch in the stack or
+ * the heap, no restrictions in this regard. This function returns NULL on
+ * error.
+ */
+EXPORT_SYMBOL struct mnl_nlmsg_batch *
+mnl_nlmsg_batch_start(void *buf, size_t limit)
+{
+	struct mnl_nlmsg_batch *b;
+
+	b = malloc(sizeof(struct mnl_nlmsg_batch));
+	if (b == NULL)
+		return NULL;
+
+	b->buf = buf;
+	b->limit = limit;
+	b->buflen = 0;
+	b->cur = buf;
+	b->overflow = false;
+
+	return b;
+}
+
+/**
+ * mnl_nlmsg_batch_stop - release a batch
+ * \param b pointer to batch
+ *
+ * This function returns the amount of data that is part of this batch.
+ */
+EXPORT_SYMBOL void
+mnl_nlmsg_batch_stop(struct mnl_nlmsg_batch *b)
+{
+	free(b);
+}
+
+/**
+ * mnl_nlmsg_batch_next - get room for the next message in the batch
+ * \param b pointer to batch
+ *
+ * This function returns a pointer to the beginning of the new Netlink message
+ * in the batch. It returns false if the last message did not fit into the
+ * batch.
+ *
+ * You have to put at least one message in the batch before calling this
+ * function, otherwise your application is likely to crash.
+ */
+EXPORT_SYMBOL bool
+mnl_nlmsg_batch_next(struct mnl_nlmsg_batch *b)
+{
+	struct nlmsghdr *nlh = b->cur;
+
+	if (b->buflen + nlh->nlmsg_len > b->limit) {
+		b->overflow = true;
+		return false;
+	}
+	b->cur = b->buf + b->buflen + nlh->nlmsg_len;
+	b->buflen += nlh->nlmsg_len;
+	return true;
+}
+
+/**
+ * mnl_nlmsg_batch_reset - reset the batch
+ * \param b pointer to batch
+ *
+ * This function allows to reset a batch, so you can reuse it to create a
+ * new one. This function moves the last message which does not fit the
+ * batch to the head of the buffer, if any.
+ */
+EXPORT_SYMBOL void
+mnl_nlmsg_batch_reset(struct mnl_nlmsg_batch *b)
+{
+	if (b->overflow) {
+		struct nlmsghdr *nlh = b->cur;
+		memcpy(b->buf, b->cur, nlh->nlmsg_len);
+		b->buflen = nlh->nlmsg_len;
+		b->cur = b->buf + b->buflen;
+		b->overflow = false;
+	} else {
+		b->buflen = 0;
+		b->cur = b->buf;
+	}
+}
+
+/**
+ * mnl_nlmsg_batch_size - get current size of the batch
+ * \param b pointer to batch
+ *
+ * This function returns the current size of the batch.
+ */
+EXPORT_SYMBOL size_t
+mnl_nlmsg_batch_size(struct mnl_nlmsg_batch *b)
+{
+	return b->buflen;
+}
+
+/**
+ * mnl_nlmsg_batch_head - get head of this batch
+ * \param b pointer to batch
+ *
+ * This function returns a pointer to the head of the batch, which is the
+ * beginning of the buffer that is used.
+ */
+EXPORT_SYMBOL void *
+mnl_nlmsg_batch_head(struct mnl_nlmsg_batch *b)
+{
+	return b->buf;
+}
+
+/**
+ * mnl_nlmsg_batch_current - returns current position in the batch
+ * \param b pointer to batch
+ *
+ * This function returns a pointer to the current position in the buffer
+ * that is used to store the batch.
+ */
+EXPORT_SYMBOL void *
+mnl_nlmsg_batch_current(struct mnl_nlmsg_batch *b)
+{
+	return b->cur;
+}
+
+/**
+ * mnl_nlmsg_batch_is_empty - check if there is any message in the batch
+ * \param b pointer to batch
+ *
+ * This function returns true if the batch is empty.
+ */
+EXPORT_SYMBOL bool
+mnl_nlmsg_batch_is_empty(struct mnl_nlmsg_batch *b)
+{
+	return b->buflen == 0;
+}
+
+/**
  * @}
  */