[libmnl PATCH 9/9] examples: add for mnl_nlmsg_batch_reset_buffer

[Ken-ichirou MATSUZAWA <chamaken@gmail.com>]

Signed-off-by: Ken-ichirou MATSUZAWA <chamas@h4.dion.ne.jp>

---
 examples/mmap/min/Makefile.am         |   6 +-
 examples/mmap/min/nfct-create-batch.c | 248 ++++++++++++++++++++++++++++++++++
 2 files changed, 253 insertions(+), 1 deletion(-)
 create mode 100644 examples/mmap/min/nfct-create-batch.c

diff --git a/examples/mmap/min/Makefile.am b/examples/mmap/min/Makefile.am
index 5b172c2..b053515 100644
--- a/examples/mmap/min/Makefile.am
+++ b/examples/mmap/min/Makefile.am
@@ -5,7 +5,8 @@ check_PROGRAMS = nf-log \
 		 nfct-dump \
 		 rtnl-addr-dump \
 		 rtnl-link-dump \
-		 rtnl-route-dump
+		 rtnl-route-dump \
+		 nfct-create-batch
 
 nf_log_SOURCES = nf-log.c
 nf_log_LDADD = ../../../src/libmnl.la
@@ -24,3 +25,6 @@ rtnl_link_dump_LDADD = ../../../src/libmnl.la
 
 rtnl_route_dump_SOURCES = rtnl-route-dump.c
 rtnl_route_dump_LDADD = ../../../src/libmnl.la
+
+nfct_create_batch_SOURCES = nfct-create-batch.c
+nfct_create_batch_LDADD = ../../../src/libmnl.la
diff --git a/examples/mmap/min/nfct-create-batch.c b/examples/mmap/min/nfct-create-batch.c
new file mode 100644
index 0000000..bb8586d
--- /dev/null
+++ b/examples/mmap/min/nfct-create-batch.c
@@ -0,0 +1,248 @@
+/* This example is placed in the public domain. */
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <arpa/inet.h>
+#include <time.h>
+#include <sys/select.h>
+#include <string.h>
+#include <errno.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_ring_socket *nlm, int portid)
+{
+	int ret, fd = mnl_socket_get_fd(nl);
+	size_t len;
+	char rcv_buf[MNL_SOCKET_BUFFER_SIZE];
+	struct nl_mmap_hdr *hdr;
+	void *ptr;
+
+	if (mnl_socket_sendto(nl, NULL, 0) < 0) {
+		perror("mnl_socket_sendto");
+		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)) {
+		hdr = mnl_ring_get_frame(nlm, MNL_RING_RX);
+		while (hdr->nm_status != NL_MMAP_STATUS_UNUSED) {
+			if (hdr->nm_status == NL_MMAP_STATUS_VALID) {
+				ptr = MNL_MMAP_MSGHDR(hdr);
+				len = hdr->nm_len;
+				if (len == 0)
+					goto next;
+			} else if (hdr->nm_status == NL_MMAP_STATUS_COPY) {
+				len = recv(mnl_socket_get_fd(nl),
+					   rcv_buf, sizeof(rcv_buf), MSG_DONTWAIT);
+				if (len <= 0)
+					break;
+				ptr = rcv_buf;
+			} else
+				break;
+
+			ret = mnl_cb_run2(ptr, len, 0, portid,
+					  NULL, NULL, cb_ctl_array,
+					  MNL_ARRAY_SIZE(cb_ctl_array));
+			if (ret == -1) {
+				perror("mnl_cb_run");
+				exit(EXIT_FAILURE);
+			}
+next:
+			hdr->nm_status = NL_MMAP_STATUS_UNUSED;
+			mnl_ring_advance(nlm, MNL_RING_RX);
+			hdr = mnl_ring_get_frame(nlm, MNL_RING_RX);
+		}
+		FD_ZERO(&readfds);
+		FD_SET(fd, &readfds);
+		ret = select(fd+1, &readfds, NULL, NULL, &tv);
+		if (ret == -1) {
+			perror("select");
+			exit(EXIT_FAILURE);
+		}
+	}
+}
+
+int main(void)
+{
+	struct mnl_socket *nl;
+	struct mnl_ring_socket *nlm;
+	void *snd_buf;
+	struct mnl_nlmsg_batch *b;
+	struct nl_mmap_req tx_nlmr, rx_nlmr;
+	struct nl_mmap_hdr *hdr;
+	int j, nmsgs;
+	uint16_t i;
+	unsigned int seq, portid;
+
+	tx_nlmr = (struct nl_mmap_req)MNL_MMAP_DEFAULT_REQ;
+	rx_nlmr = (struct nl_mmap_req){
+		.nm_block_size	= MNL_SOCKET_BUFFER_SIZE * 16,
+		.nm_block_nr	= 64,
+		.nm_frame_size	= 2048,
+		.nm_frame_nr	= 64 * MNL_SOCKET_BUFFER_SIZE * 16 / 2048 };
+
+	nl = mnl_socket_open(NETLINK_NETFILTER);
+	if (nl == NULL) {
+		perror("mnl_socket_open");
+		exit(EXIT_FAILURE);
+	}
+
+	nlm = mnl_ring_map(nl, &tx_nlmr, &rx_nlmr);
+	if (nlm == NULL) {
+		perror("mnl_ring_map");
+		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);
+
+	hdr = mnl_ring_get_frame(nlm, MNL_RING_TX);
+	if (hdr == NULL) {
+		perror("mnl_ring_get_frame - MNL_RING_TX");
+		exit(EXIT_FAILURE);
+	}
+
+	/* The buffer size that we use to batch messages is
+	 * MNL_MMAP_DEFAULT_REQ.nm_frame_size - MNL_MMAP_HDRLEN,
+	 * 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. */
+	snd_buf = (void *)hdr + NL_MMAP_HDRLEN;
+	b = mnl_nlmsg_batch_start(snd_buf, (tx_nlmr.nm_frame_size - NL_MMAP_HDRLEN) / 2);
+	if (b == NULL) {
+		perror("mnl_nlmsg_batch_start");
+		exit(EXIT_FAILURE);
+	}
+
+	seq = time(NULL);
+	for (i=1024, j=0, nmsgs=1; j<65535; i++, j++, nmsgs++) {
+		put_msg(mnl_nlmsg_batch_current(b), i, seq+j);
+
+		/* is there room for more messages in this batch?
+		 * AND batched messages smaller than frame_nr?
+		 * if so, continue.
+		 *
+		 * we can send over frame_nr messages in a single frame,
+		 * but requesting ACK needs tx frame_nr frames.
+		 */
+		if (mnl_nlmsg_batch_next(b) && nmsgs < rx_nlmr.nm_frame_nr)
+			continue;
+
+		hdr->nm_len = mnl_nlmsg_batch_size(b);
+		hdr->nm_status = NL_MMAP_STATUS_VALID;
+		send_batch(nl, nlm, portid);
+
+		nmsgs = 1;
+		mnl_ring_advance(nlm, MNL_RING_TX);
+
+		hdr = mnl_ring_get_frame(nlm, MNL_RING_TX);
+		if (hdr == NULL) {
+			perror("mnl_ring_get_frame - MNL_RING_TX");
+			exit(EXIT_FAILURE);
+		}
+
+		snd_buf = (void *)hdr + NL_MMAP_HDRLEN;
+		mnl_nlmsg_batch_reset_buffer(b, snd_buf, (tx_nlmr.nm_frame_size - NL_MMAP_HDRLEN) / 2);
+	}
+
+	/* check if there is any message in the batch not sent yet. */
+	if (!mnl_nlmsg_batch_is_empty(b)) {
+		hdr->nm_len = mnl_nlmsg_batch_size(b);
+		hdr->nm_status = NL_MMAP_STATUS_VALID;
+		send_batch(nl, nlm, portid);
+	}
+
+	mnl_nlmsg_batch_stop(b);
+	mnl_ring_unmap(nlm);
+	mnl_socket_close(nl);
+
+	return 0;
+}
-- 
1.8.4.rc3