#define pr_fmt(fmt) "kbench: " fmt

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/inet.h>

#include <net/route.h>
#include <net/ip_fib.h>

#include <linux/timex.h>

/* We can't just use "get_cycles()" as on some platforms, such
 * as sparc64, that gives system cycles rather than cpu clock
 * cycles.
 */

#ifdef CONFIG_SPARC64
static inline unsigned long long kbench_get_tick(void)
{
	unsigned long long t;

	__asm__ __volatile__("rd %%tick, %0" : "=r" (t));
	return t;
}
#elif defined(CONFIG_X86)
static inline unsigned long long kbench_get_tick(void)
{
	unsigned long long t;

	rdtscll(t);

	return t;
}
#elif defined(CONFIG_POWERPC)
static inline unsigned long long kbench_get_tick(void)
{
	return get_cycles();
}
#else
#error Unsupported architecture, please implement kbench_get_tick()
#endif

#undef IP_ROUTE_HAVE_PREALLOC
#undef IP_ROUTE_CYCLE_SAMPLES
#undef IP_FIB_LOOKUP_EXPORTED

#ifdef IP_ROUTE_CYCLE_SAMPLES
extern int ip_route_output_cycles_1;
extern int ip_route_output_cycles_2;
extern int ip_route_output_cycles_3;
extern int ip_route_output_cycles_4;
extern int ip_route_output_cycles_5;
extern int ip_route_output_cycles_6;
extern int ip_route_output_cycles_7;
#endif

#define DEFAULT_WARMUP_COUNT 100000

#define DEFAULT_DST_IP_ADDR	0x4a800001
#define DEFAULT_SRC_IP_ADDR	0x00000000
#define DEFAULT_OIF		0
#define DEFAULT_IIF		0
#define DEFAULT_MARK		0x00000000
#define DEFAULT_TOS		0x00

static int flow_oif = DEFAULT_OIF;
static int flow_iif = DEFAULT_IIF;
static u32 flow_mark = DEFAULT_MARK;
static u32 flow_dst_ip_addr = DEFAULT_DST_IP_ADDR;
static u32 flow_src_ip_addr = DEFAULT_SRC_IP_ADDR;
static int flow_tos = DEFAULT_TOS;

static char dst_string[64];
static char src_string[64];

module_param_string(dst, dst_string, sizeof(dst_string), 0);
module_param_string(src, src_string, sizeof(src_string), 0);

static void __init flow_setup(void)
{
	if (dst_string[0])
		flow_dst_ip_addr = in_aton(dst_string);
	if (src_string[0])
		flow_src_ip_addr = in_aton(src_string);
}

module_param_named(oif, flow_oif, int, 0);
module_param_named(iif, flow_iif, int, 0);
module_param_named(mark, flow_mark, uint, 0);
module_param_named(tos, flow_tos, int, 0);

static int warmup_count = DEFAULT_WARMUP_COUNT;
module_param_named(count, warmup_count, int, 0);

static void flow_init(struct flowi4 *fl4)
{
	memset(fl4, 0, sizeof(*fl4));
	fl4->flowi4_oif = flow_oif;
	fl4->flowi4_iif = flow_iif;
	fl4->flowi4_mark = flow_mark;
	fl4->flowi4_tos = flow_tos;
	fl4->daddr = flow_dst_ip_addr;
	fl4->saddr = flow_src_ip_addr;
}

static struct rtable *route_output(struct net *net, struct flowi4 *fl4)
{
	return ip_route_output_key(net, fl4);
}

static void do_full_output_lookup_bench(void)
{
	unsigned long long t1, t2, tdiff;
	struct rtable *rt;
	struct flowi4 fl4;
	int i;

	rt = NULL;

	for (i = 0; i < warmup_count; i++) {
		flow_init(&fl4);

		rt = route_output(&init_net, &fl4);
		if (IS_ERR(rt))
			break;
		ip_rt_put(rt);
	}
	if (IS_ERR(rt)) {
		pr_info("ip_route_output_key: err=%ld\n", PTR_ERR(rt));
		return;
	}

#ifdef IP_ROUTE_CYCLE_SAMPLES
	ip_route_output_cycles_1 = 0;
	ip_route_output_cycles_2 = 0;
	ip_route_output_cycles_3 = 0;
	ip_route_output_cycles_4 = 0;
	ip_route_output_cycles_5 = 0;
	ip_route_output_cycles_6 = 0;
	ip_route_output_cycles_7 = 0;
#endif

	flow_init(&fl4);

	t1 = kbench_get_tick();
	rt = route_output(&init_net, &fl4);
	t2 = kbench_get_tick();
	if (!IS_ERR(rt))
		ip_rt_put(rt);

	tdiff = t2 - t1;
	pr_info("ip_route_output_key tdiff: %llu\n", tdiff);
#ifdef IP_ROUTE_CYCLE_SAMPLES
	pr_info("ip_route_output_key cycls: [%d %d %d %d %d %d]\n",
		ip_route_output_cycles_1,
		ip_route_output_cycles_2,
		ip_route_output_cycles_3,
		ip_route_output_cycles_4,
		ip_route_output_cycles_5,
		ip_route_output_cycles_6);
#endif
}

static void do_full_input_lookup_bench(void)
{
	unsigned long long t1, t2, tdiff;
	struct net_device *dev;
	struct sk_buff *skb;
	int err, i;

	skb = alloc_skb(4096, GFP_KERNEL);
	if (!skb) {
		pr_info("Cannot alloc SKB for test\n");
		return;
	}
	skb_reset_mac_header(skb);
	skb_reset_network_header(skb);
	ip_hdr(skb)->protocol = IPPROTO_ICMP;
	skb_reserve(skb, MAX_HEADER + sizeof(struct iphdr));

	dev = __dev_get_by_index(&init_net, flow_iif);
	if (dev == NULL) {
		pr_info("Input device does not exist\n");
		goto out_free;
	}
	skb->protocol = htons(ETH_P_IP);
	skb->dev = dev;
	skb->mark = flow_mark;
	local_bh_disable();
	err = 0;
	for (i = 0; i < warmup_count; i++) {
		err = ip_route_input(skb, flow_dst_ip_addr, flow_src_ip_addr, flow_tos, dev);
		if (err)
			break;
		skb_dst_drop(skb);
	}
	local_bh_enable();

	if (err) {
		pr_info("Input route lookup fails with err=%d\n", err);
		goto out_free;
	}

	local_bh_disable();
	t1 = kbench_get_tick();
	err = ip_route_input(skb, flow_dst_ip_addr, flow_src_ip_addr, flow_tos, dev);
	t2 = kbench_get_tick();
	local_bh_enable();

	if (err) {
		pr_info("Input route lookup fails with err=%d\n", err);
		goto out_free;
	}

	skb_dst_drop(skb);

	tdiff = t2 - t1;
	pr_info("ip_route_input tdiff: %llu\n", tdiff);

out_free:
	kfree_skb(skb);
}

static void do_full_lookup_bench(void)
{
	if (!flow_iif)
		do_full_output_lookup_bench();
	else
		do_full_input_lookup_bench();
}

#ifdef IP_ROUTE_HAVE_PREALLOC
static void do_full_lookup_prealloc_bench(void)
{
	unsigned long long t1, t2, tdiff;
	struct rtable *rt, rt_stack;
	struct flowi4 fl4;
	int err, i;

	err = 0;

	for (i = 0; i < warmup_count; i++) {
		flow_init(&fl4);

		rt = ip_route_output_flow_prealloc(&init_net, &fl4, NULL, &rt_stack.dst);
		if (IS_ERR(rt)) {
			err = PTR_ERR(rt);
			break;
		}
		ip_rt_put(rt);
	}
	if (err) {
		pr_info("ip_route_output_key prealloc: err=%d\n", err);
		return;
	}

#ifdef IP_ROUTE_CYCLE_SAMPLES
	ip_route_output_cycles_1 = 0;
	ip_route_output_cycles_2 = 0;
	ip_route_output_cycles_3 = 0;
	ip_route_output_cycles_4 = 0;
	ip_route_output_cycles_5 = 0;
	ip_route_output_cycles_6 = 0;
	ip_route_output_cycles_7 = 0;
#endif

	flow_init(&fl4);

	t1 = kbench_get_tick();
	rt = ip_route_output_flow_prealloc(&init_net, &fl4, NULL, &rt_stack.dst);
	t2 = kbench_get_tick();
	if (!IS_ERR(rt))
		ip_rt_put(rt);

	tdiff = t2 - t1;
	pr_info("ip_route_output_key prealloc tdiff: %llu\n", tdiff);
#ifdef IP_ROUTE_CYCLE_SAMPLES
	pr_info("ip_route_output_key prealloc cycls: [%d %d %d %d %d %d %d]\n",
		ip_route_output_cycles_1,
		ip_route_output_cycles_2,
		ip_route_output_cycles_3,
		ip_route_output_cycles_4,
		ip_route_output_cycles_5,
		ip_route_output_cycles_6,
		ip_route_output_cycles_7);
#endif
}
#endif

#ifdef IP_FIB_LOOKUP_EXPORTED
static void do_fib_lookup_bench(void)
{
	unsigned long long t1, t2, tdiff;
	struct fib_result res;
	struct flowi4 fl4;
	int err, i;

	flow_init(&fl4);

	for (i = 0; i < warmup_count; i++) {
		struct fib_table *table;

		table = fib_get_table(&init_net, RT_TABLE_MAIN);
		if (!table) {
			pr_info("fib_lookup: No main table.\n");
			return;
		}
		err = fib_table_lookup(table, &fl4, &res, FIB_LOOKUP_NOREF);
		if (err) {
			pr_info("fib_lookup: err=%d\n", err);
			return;
		}
	}

	{
		struct fib_table *table;

		t1 = kbench_get_tick();
		table = fib_get_table(&init_net, RT_TABLE_MAIN);
		if (!table) {
			pr_info("fib_lookup: No main table.\n");
			return;
		}
		err = fib_table_lookup(table, &fl4, &res, FIB_LOOKUP_NOREF);
		t2 = kbench_get_tick();
	}
	tdiff = t2 - t1;
	pr_info("fib_lookup tdiff: %llu\n", tdiff);
}

struct net_route {
	struct net_device	*dev;
	unsigned long		_metrics;
	struct neighbour	*neighbour;
	struct hh_cache		*hh;
	atomic_t		__refcnt;
};

struct ipv4_route {
	struct net_route	base;

	int			rt_genid;

	__be32			rt_dst;
	__be32			rt_src;
	__u16			rt_type;
	__u8			rt_tos;

	int			rt_iif;
	int			rt_oif;
	__u32			rt_mark;

	__be32			rt_gateway;
	__be32			rt_spec_dst;

	u32			rt_peer_genid;
	struct inet_peer	*peer;
	struct fib_info		*fi;
};

static atomic_t foo_id = ATOMIC_INIT(1);

static int new_output_lookup(const struct flowi *fl, struct ipv4_route *rt)
{
	struct fib_table *table = fib_get_table(&init_net, RT_TABLE_MAIN);
	struct fib_result res;
	int err;

	if (!table) {
		pr_info("new_output_lookup: No main table.\n");
		return -ENODEV;
	}
	err = fib_table_lookup(table, fl, &res, FIB_LOOKUP_NOREF);
	if (err)
		return err;

	rt->rt_genid = atomic_read(&foo_id);
	rt->rt_dst = fl->fl4_dst;
	rt->rt_src = fl->fl4_src;
	rt->rt_type = res.type;
	rt->rt_tos = fl->fl4_tos;
	rt->rt_iif = fl->iif;
	rt->rt_oif = fl->oif;
	rt->rt_mark = fl->mark;
	rt->rt_gateway = FIB_RES_GW(res);
	rt->rt_spec_dst = fl->fl4_src;
	rt->rt_peer_genid = 0;
	rt->peer = NULL;
	rt->fi = NULL;

	return 0;
}

static void do_new_lookup_bench(void)
{
	unsigned long long t1, t2, tdiff;
	struct ipv4_route rt;
	struct flowi fl;
	int err, i;

	flow_init(&fl);

	for (i = 0; i < warmup_count; i++) {
		err = new_output_lookup(&fl, &rt);
		if (err) {
			pr_info("new_output_lookup: err=%d\n", err);
			return;
		}
	}
	t1 = kbench_get_tick();
	err = new_output_lookup(&fl, &rt);
	t2 = kbench_get_tick();

	if (err) {
		pr_info("new_output_lookup: err=%d\n", err);
		return;
	}

	tdiff = t2 - t1;
	pr_info("new_output_lookup tdiff: %llu\n", tdiff);
}
#endif

static void do_bench(void)
{
	do_full_lookup_bench();
	do_full_lookup_bench();
	do_full_lookup_bench();
	do_full_lookup_bench();

#ifdef IP_ROUTE_HAVE_PREALLOC
	do_full_lookup_prealloc_bench();
	do_full_lookup_prealloc_bench();
	do_full_lookup_prealloc_bench();
	do_full_lookup_prealloc_bench();
#endif

#ifdef IP_FIB_LOOKUP_EXPORTED
	do_fib_lookup_bench();
	do_fib_lookup_bench();
	do_fib_lookup_bench();
	do_fib_lookup_bench();

	do_new_lookup_bench();
	do_new_lookup_bench();
	do_new_lookup_bench();
	do_new_lookup_bench();
#endif
}

static int __init kbench_init(void)
{
	flow_setup();

	pr_info("flow [IIF(%d),OIF(%d),MARK(0x%08x),D(%pI4),S(%pI4),TOS(0x%02x)]\n",
		flow_iif, flow_oif, flow_mark,
		&flow_dst_ip_addr,
		&flow_src_ip_addr, flow_tos);

#if defined(CONFIG_X86)
	if (!boot_cpu_has(X86_FEATURE_TSC)) {
		pr_err("X86 TSC is required, but is unavailable.\n");
		return -EINVAL;
	}
#endif

	pr_info("sizeof(struct rtable)==%zd\n", sizeof(struct rtable));

	do_bench();

	return -ENODEV;
}

static void __exit kbench_exit(void)
{
}

module_init(kbench_init);
module_exit(kbench_exit);
MODULE_LICENSE("GPL");