/* Route lookup benchmark tool using RTM_GETROUTE */

#include <stdio.h>
#include <stddef.h>

#include <string.h>
#include <errno.h>

#include <time.h>
#include <sys/time.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define _GNU_SOURCE
#include <getopt.h>

#include <libmnl/libmnl.h>

#include <linux/if.h>
#include <linux/if_link.h>
#include <linux/rtnetlink.h>
#include <linux/filter.h>

static int usage(void)
{
	printf("usage: route_bench [ -o ] [ -l count ]\n");
	printf("\t\t[ -s src_ip ] [ -a src_ip_stride ] [ -b src_ip_limit ]\n");
	printf("\t\t[ -d dst_ip ] [ -e dst_ip_stride ] [ -f dst_ip_limit ]\n");
	printf("\t\t[ -i iif ] [ -x iif_stride ] [ -y iif_limit ]\n");
	printf("\t\t[ -m mark ] [ -n mark_stride ] [ -p mark_limit ]\n");
	printf ("\t\t[ -t tos ] [ -q tos_stride ] [ -r tos_limit ]\n");
	return -1;
}

struct bench_state {
	unsigned int		flags;
#define FLAG_SRC_ITERATE	0x00000001
#define FLAG_DST_ITERATE	0x00000002
#define FLAG_MARK_ITERATE	0x00000004
#define FLAG_IIF_ITERATE	0x00000008
#define FLAG_TOS_ITERATE	0x00000010
#define FLAG_ITERATE_ONE	0x00000020

	in_addr_t		src_addr;
	in_addr_t		dst_addr;
	unsigned int		mark;
	unsigned int		iif;
	unsigned int		tos;

	in_addr_t		orig_src_addr;
	int			src_addr_stride;
	in_addr_t		src_addr_limit;

	in_addr_t		orig_dst_addr;
	int			dst_addr_stride;
	in_addr_t		dst_addr_limit;

	unsigned int		orig_mark;
	int			mark_stride;
	unsigned int		mark_limit;

	unsigned int		orig_iif;
	int			iif_stride;
	unsigned int		iif_limit;

	unsigned int		orig_tos;
	int			tos_stride;
	unsigned int		tos_limit;
};

static struct bench_state state;

static void init_bench_state(struct bench_state *s)
{
	s->src_addr = s->orig_src_addr;
	s->dst_addr = s->orig_dst_addr;
	s->mark = s->orig_mark;
	s->iif = s->orig_iif;
	s->tos = s->orig_tos;
}

static void advance_bench_state(struct bench_state *s)
{
	if (s->flags & FLAG_SRC_ITERATE) {
		in_addr_t orig = s->src_addr;

		s->src_addr += s->src_addr_stride;
		if (s->src_addr < orig ||
		    s->src_addr > s->src_addr_limit)
			s->src_addr = s->orig_src_addr;
		else if (s->flags & FLAG_ITERATE_ONE)
			return;
	}
	if (s->flags & FLAG_DST_ITERATE) {
		in_addr_t orig = s->dst_addr;

		s->dst_addr += s->dst_addr_stride;
		if (s->dst_addr < orig ||
		    s->dst_addr > s->dst_addr_limit)
			s->dst_addr = s->orig_dst_addr;
		else if (s->flags & FLAG_ITERATE_ONE)
			return;
	}
	if (s->flags & FLAG_MARK_ITERATE) {
		in_addr_t orig = s->mark;

		s->mark += s->mark_stride;
		if (s->mark < orig ||
		    s->mark > s->mark_limit)
			s->mark = s->orig_mark;
		else if (s->flags & FLAG_ITERATE_ONE)
			return;
	}
	if (s->flags & FLAG_IIF_ITERATE) {
		in_addr_t orig = s->iif;

		s->iif += s->iif_stride;
		if (s->iif < orig ||
		    s->iif > s->iif_limit)
			s->iif = s->orig_iif;
		else if (s->flags & FLAG_ITERATE_ONE)
			return;
	}
	if (s->flags & FLAG_TOS_ITERATE) {
		in_addr_t orig = s->tos;

		s->tos += s->tos_stride;
		if (s->tos < orig ||
		    s->tos > s->tos_limit)
			s->tos = s->orig_tos;
		else if (s->flags & FLAG_ITERATE_ONE)
			return;
	}
}

static int do_bench(int count, in_addr_t src_addr, in_addr_t dst_addr,
		    unsigned int mark, unsigned int iif)
{
	char send_buf[MNL_SOCKET_BUFFER_SIZE];
	unsigned int min, sec, frac, tmp;
	struct sock_filter insns = { .code = BPF_RET | BPF_K,
				     .k = 0 };
	struct sock_fprog filter = { .len = 1, .filter = &insns, };
	struct bench_state *s = &state;
	struct timeval start_time;
	struct timeval end_time;
	struct mnl_socket *nl;
	struct nlmsghdr *nlh;
	unsigned int portid;
	struct rtmsg *rtm;
	int i, err;

	init_bench_state(s);

	nl = mnl_socket_open(NETLINK_ROUTE);
	if (!nl) {
		perror("mnl_socket_open");
		return -1;
	}

	if (mnl_socket_bind(nl, 0, MNL_SOCKET_AUTOPID) < 0) {
		perror("mnl_socket_bind");
		return -1;
	}

	portid = mnl_socket_get_portid(nl);

	err = setsockopt(mnl_socket_get_fd(nl), SOL_SOCKET,
			 SO_ATTACH_FILTER, &filter, sizeof(filter));
	if (err) {
		perror("setsockopt");
		return -1;
	}

	gettimeofday(&start_time, NULL);
	for (i = 0; i < count; i++) {
		unsigned int seq;

		nlh = mnl_nlmsg_put_header(send_buf);
		nlh->nlmsg_type = RTM_GETROUTE;
		nlh->nlmsg_flags = NLM_F_REQUEST;
		nlh->nlmsg_seq = seq = time(NULL);

		rtm = mnl_nlmsg_put_extra_header(nlh, sizeof(struct rtmsg));
		rtm->rtm_family = AF_INET;
		rtm->rtm_tos = s->tos;

		if (s->src_addr)
			mnl_attr_put_u32(nlh, RTA_SRC, s->src_addr);
		if (s->dst_addr)
			mnl_attr_put_u32(nlh, RTA_DST, s->dst_addr);

		if (s->iif)
			mnl_attr_put_u32(nlh, RTA_IIF, s->iif);

#define RTA_MARK	(RTA_TABLE + 1)

		if (s->mark)
			mnl_attr_put_u32(nlh, RTA_MARK, s->mark);


		if (mnl_socket_sendto(nl, nlh, nlh->nlmsg_len) < 0) {
			perror("mnl_socket_sendto");
			return -1;
		}

		/* No need to do a receive, as the socket filter rejects
		 * all packets.
		 */

		advance_bench_state(s);
	}
	gettimeofday(&end_time, NULL);

	sec = end_time.tv_sec - start_time.tv_sec;
	frac = end_time.tv_usec - start_time.tv_usec;

	tmp = frac % 1000;
	frac = (frac * 1000) / 1000000;
	if (tmp >= 500)
		frac++;
	if (frac >= 1000) {
		sec++;
		frac -= 1000;
	}
	
	min = sec / 60;
	sec = sec % 60;

	printf("Result: %um%u.%03us\n", min, sec, frac);

	mnl_socket_close(nl);

	return 0;
}

int main(int argc, char **argv, char **envp)
{
	struct bench_state *s = &state;
	int count, ret;

	count = 1000;

	s->flags = 0;
	s->orig_mark = s->orig_iif = s->orig_tos = 0;
	s->orig_src_addr = s->orig_dst_addr = INADDR_ANY;

	while ((ret = getopt(argc, argv, "l:s:d:i:m:t:a:b:e:f:x:y:n:p:q:r:oh")) >= 0) {
		switch (ret) {
		case 'l':
			sscanf(optarg, "%d", &count);
			break;
		case 's':
			s->orig_src_addr = inet_addr(optarg);
			if (s->orig_src_addr == INADDR_NONE)
				return usage();
			break;
		case 'd':
			s->orig_dst_addr = inet_addr(optarg);
			if (s->orig_dst_addr == INADDR_NONE)
				return usage();
			break;
		case 'i':
			sscanf(optarg, "%u", &s->orig_iif);
			break;
		case 'm':
			sscanf(optarg, "%u", &s->orig_mark);
			break;
		case 't':
			sscanf(optarg, "%u", &s->orig_tos);
			break;
		case 'o':
			s->flags |= FLAG_ITERATE_ONE;
			break;
		case 'a':
			sscanf(optarg, "%u", &s->src_addr_stride);
			s->flags |= FLAG_SRC_ITERATE;
			break;
		case 'b':
			s->src_addr_limit = inet_addr(optarg);
			if (s->src_addr_limit == INADDR_NONE)
				return usage();
			if (!s->src_addr_stride)
				s->src_addr_stride = 1;
			s->flags |= FLAG_SRC_ITERATE;
			break;
		case 'e':
			sscanf(optarg, "%u", &s->dst_addr_stride);
			s->flags |= FLAG_SRC_ITERATE;
			break;
		case 'f':
			s->dst_addr_limit = inet_addr(optarg);
			if (s->dst_addr_limit == INADDR_NONE)
				return usage();
			if (!s->dst_addr_stride)
				s->dst_addr_stride = 1;
			s->flags |= FLAG_DST_ITERATE;
			break;
		case 'x':
			sscanf(optarg, "%u", &s->iif_stride);
			s->flags |= FLAG_IIF_ITERATE;
			break;
		case 'y':
			sscanf(optarg, "%u", &s->iif_limit);
			if (!s->iif_stride)
				s->iif_stride = 1;
			s->flags |= FLAG_IIF_ITERATE;
			break;
		case 'n':
			sscanf(optarg, "%u", &s->mark_stride);
			s->flags |= FLAG_MARK_ITERATE;
			break;
		case 'p':
			sscanf(optarg, "%u", &s->mark_limit);
			if (!s->mark_stride)
				s->mark_stride = 1;
			s->flags |= FLAG_MARK_ITERATE;
			break;
		case 'q':
			sscanf(optarg, "%u", &s->tos_stride);
			s->flags |= FLAG_TOS_ITERATE;
			break;
		case 'r':
			sscanf(optarg, "%u", &s->tos_limit);
			if (!s->tos_stride)
				s->tos_stride = 1;
			s->flags |= FLAG_TOS_ITERATE;
			break;
		case '?':
		case 'h':
			return usage();
		}
	}

	printf("Bench: count(%d) saddr[0x%08x] daddr[0x%08x] mark[0x%x] iif[0x%x]\n",
	       count, s->orig_src_addr, s->orig_dst_addr,
	       s->orig_mark, s->orig_iif);

	return do_bench(count, s->orig_src_addr, s->orig_dst_addr,
			s->orig_mark, s->orig_iif);
}