/*
 * Copyright (c) 2012 The DragonFly Project.  All rights reserved.
 * 
 * This code is derived from software contributed to The DragonFly Project
 * by Sepherosa Ziehau <sepherosa@gmail.com>
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 * 3. Neither the name of The DragonFly Project nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific, prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
 * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#define _IP_VHL

#include <sys/param.h>
#include <sys/conf.h>
#include <sys/device.h>
#include <sys/in_cksum.h>
#include <sys/kernel.h>
#include <sys/malloc.h>
#include <sys/mbuf.h>
#include <sys/proc.h>
#include <sys/caps.h>
#include <sys/queue.h>
#include <sys/socket.h>
#include <sys/systm.h>
#include <sys/serialize.h>

#include <net/if.h>
#include <net/if_dl.h>
#include <net/if_var.h>
#include <net/ifq_var.h>
#include <net/ethernet.h>
#include <net/netmsg2.h>
#include <net/netisr2.h>

#include <netinet/in.h>
#include <netinet/in_pcb.h>
#include <netinet/ip.h>
#include <netinet/udp_var.h>

#include "pktgen.h"

#define CDEV_NAME	"pktg"

#define PKTGEN_BUFSZ	2048

#ifndef PKTGEN_DEVCNT
#define PKTGEN_DEVCNT	4
#endif

struct pktgen;

struct netmsg_pktgen {
	struct netmsg_base	np_base;
	struct pktgen		*np_pktg;
	struct ifaltq_subque	*np_ifsq;
};

struct pktgen_buf {
	struct netmsg_base	pb_nmsg;	/* MUST BE THE FIRST */
	void			*pb_buf;
	volatile int		pb_done;
	int			pb_inuse;
	struct ifnet		*pb_ifp;
	struct ifaltq_subque	*pb_ifsq;
	int			pb_len;
	int			pb_cpuid;
	struct pktgen		*pb_pktg;
	LIST_ENTRY(pktgen_buf)	pb_link;
};

struct pktgen_pcpu {
	struct callout		pktg_stop;
	LIST_HEAD(, pktgen_buf)	pktg_buflist;
};

struct pktgen {
	uint32_t		pktg_flags;	/* PKTG_F_ */
	int			pktg_refcnt;

	int			pktg_duration;

	int			pktg_datalen;
	struct ifnet		*pktg_ifp;

	int			pktg_pktenq;

	struct sockaddr_in	pktg_src;
	int			pktg_ndst;
	struct sockaddr_in	*pktg_dst;
	uint8_t			pktg_dst_lladdr[ETHER_ADDR_LEN];

	struct pktgen_pcpu	pktg_pcpu[MAXCPU];
};

#define PKTG_F_CONFIG		0x1
#define PKTG_F_RUNNING		0x4
#define PKTG_F_SWITCH_SRCDST	0x8

static int 		pktgen_modevent(module_t, int, void *);

static void		pktgen_buf_free(void *);
static void		pktgen_buf_ref(void *);
static void		pktgen_buf_send(netmsg_t);

static int		pktgen_config(struct pktgen *,
			    const struct pktgen_conf *);
static int		pktgen_start(struct pktgen *, int);
static void		pktgen_free(struct pktgen *);
static void		pktgen_ref(struct pktgen *);
static void		pktgen_pcpu_stop_cb(void *);
static void		pktgen_mbuf(struct pktgen_buf *, struct mbuf *);
static void		pktgen_start_ifsq(struct pktgen *,
			    struct ifaltq_subque *);
static void		pktgen_start_ifsq_handler(netmsg_t);

static d_open_t		pktgen_open;
static d_close_t	pktgen_close;
static d_ioctl_t	pktgen_ioctl;

static struct dev_ops	pktgen_ops = {
	{ CDEV_NAME, 0, D_MPSAFE },
	.d_open =	pktgen_open,
	.d_close =	pktgen_close,
	.d_ioctl =	pktgen_ioctl,
};

static volatile int		pktgen_refcnt;
static struct lwkt_token pktgen_tok = LWKT_TOKEN_INITIALIZER(pktgen_token);

MALLOC_DECLARE(M_PKTGEN);
MALLOC_DEFINE(M_PKTGEN, CDEV_NAME, "Packet generator");

DEV_MODULE(pktgen, pktgen_modevent, NULL);

static int
pktgen_modevent(module_t mod, int type, void *data)
{
	int error = 0, i;

	switch (type) {
	case MOD_LOAD:
		for (i = 0; i < PKTGEN_DEVCNT; ++i) {
			make_dev(&pktgen_ops, 0, UID_ROOT, GID_WHEEL, 0600,
			    CDEV_NAME"%d", i);
		}
		break;

	case MOD_UNLOAD:
		if (pktgen_refcnt > 0)
			return EBUSY;
		dev_ops_remove_all(&pktgen_ops);
		break;

	default:
		error = EOPNOTSUPP;
		break;
	}
	return error;
}

static int
pktgen_open(struct dev_open_args *ap)
{
	cdev_t dev = ap->a_head.a_dev;
	struct pktgen *pktg;
	int error, i;

	error = priv_check_cred(ap->a_cred, PRIV_ROOT, 0);
	if (error)
		return error;

	lwkt_gettoken(&pktgen_tok);

	if (dev->si_drv1 != NULL) {
		lwkt_reltoken(&pktgen_tok);
		return EBUSY;
	}

	pktg = kmalloc(sizeof(*pktg), M_PKTGEN, M_ZERO | M_WAITOK);
	for (i = 0; i < ncpus; ++i) {
		struct pktgen_pcpu *p = &pktg->pktg_pcpu[i];

		callout_init_mp(&p->pktg_stop);
		LIST_INIT(&p->pktg_buflist);
	}

	dev->si_drv1 = pktg;
	pktg->pktg_refcnt = 1;

	atomic_add_int(&pktgen_refcnt, 1);

	lwkt_reltoken(&pktgen_tok);
	return 0;
}

static int
pktgen_close(struct dev_close_args *ap)
{
	cdev_t dev = ap->a_head.a_dev;
	struct pktgen *pktg = dev->si_drv1;

	lwkt_gettoken(&pktgen_tok);
	dev->si_drv1 = NULL;
	lwkt_reltoken(&pktgen_tok);

	pktgen_free(pktg);

	return 0;
}

static int
pktgen_ioctl(struct dev_ioctl_args *ap __unused)
{
	cdev_t dev = ap->a_head.a_dev;
	caddr_t data = ap->a_data;
	struct pktgen *pktg = dev->si_drv1;
	int error;

	lwkt_gettoken(&pktgen_tok);

	switch (ap->a_cmd) {
	case PKTGENSTART:
		error = pktgen_start(pktg, 0);
		break;

	case PKTGENMQSTART:
		error = pktgen_start(pktg, 1);
		break;

	case PKTGENSCONF:
		error = pktgen_config(pktg, (const struct pktgen_conf *)data);
		break;

	default:
		error = EOPNOTSUPP;
		break;
	}

	lwkt_reltoken(&pktgen_tok);
	return error;
}

static int
pktgen_config(struct pktgen *pktg, const struct pktgen_conf *conf)
{
	const struct sockaddr_in *sin;
	struct sockaddr_in *dst = NULL;
	const struct sockaddr *sa;
	struct ifnet *ifp;
	size_t dst_size;
	int i, error, pktenq;

	if (pktg->pktg_flags & (PKTG_F_RUNNING | PKTG_F_CONFIG))
		return EBUSY;

	if (conf->pc_datalen <= 0 ||
	    conf->pc_datalen > ETHERMTU - sizeof(struct udpiphdr))
		return EINVAL;
	if (conf->pc_duration <= 0)
		return EINVAL;

	sin = &conf->pc_src;
	if (sin->sin_family != AF_INET)
		return EPROTONOSUPPORT;
	if (sin->sin_port == 0)
		return EINVAL;

	if (conf->pc_ndst <= 0)
		return EINVAL;
	dst_size = conf->pc_ndst * sizeof(struct sockaddr_in);

	dst = kmalloc(dst_size, M_PKTGEN, M_WAITOK | M_NULLOK);
	if (dst == NULL)
		return ENOMEM;

	error = copyin(conf->pc_dst, dst, dst_size);
	if (error)
		goto failed;

	for (i = 0; i < conf->pc_ndst; ++i) {
		sin = &dst[i];
		if (sin->sin_family != AF_INET) {
			error = EPROTONOSUPPORT;
			goto failed;
		}
		if (sin->sin_port == 0) {
			error = EINVAL;
			goto failed;
		}
	}

	/*
	 * XXX NOT MPSAFE.
	 * Make sure that pc_ifname is not gone.
	 */
	ifnet_lock();
	ifp = ifunit(conf->pc_ifname);
	ifnet_unlock();
	if (ifp == NULL) {
		error = ENXIO;
		goto failed;
	}

	pktenq = conf->pc_pktenq;
	if (pktenq < 0 || pktenq > ifp->if_snd.altq_maxlen) {
		error = ENOBUFS;
		goto failed;
	} else if (pktenq == 0) {
		pktenq = (ifp->if_snd.altq_maxlen * 3) / 4;
	}

	sa = &conf->pc_dst_lladdr;
	if (sa->sa_family != AF_LINK) {
		error = EPROTONOSUPPORT;
		goto failed;
	}
	if (sa->sa_len != ETHER_ADDR_LEN) {
		error = EPROTONOSUPPORT;
		goto failed;
	}
	if (ETHER_IS_MULTICAST(sa->sa_data) ||
	    bcmp(sa->sa_data, ifp->if_broadcastaddr, ifp->if_addrlen) == 0) {
		error = EADDRNOTAVAIL;
		goto failed;
	}

	/*
	 * Accept the config
	 */
	pktg->pktg_flags |= PKTG_F_CONFIG;

	if (conf->pc_flags & PKTGEN_FLAG_SWITCH_SRCDST)
		pktg->pktg_flags |= PKTG_F_SWITCH_SRCDST;
	pktg->pktg_duration = conf->pc_duration;
	pktg->pktg_datalen = conf->pc_datalen;
	pktg->pktg_pktenq = pktenq;
	pktg->pktg_ifp = ifp;
	pktg->pktg_src = conf->pc_src;
	pktg->pktg_ndst = conf->pc_ndst;
	KKASSERT(pktg->pktg_dst == NULL);
	pktg->pktg_dst = dst;
	bcopy(sa->sa_data, pktg->pktg_dst_lladdr, ETHER_ADDR_LEN);

	return 0;

failed:
	if (dst != NULL)
		kfree(dst, M_PKTGEN);
	return error;
}

static void
pktgen_start_ifsq(struct pktgen *pktg, struct ifaltq_subque *ifsq)
{
	struct netmsg_pktgen *np;

	np = kmalloc(sizeof(*np), M_LWKTMSG, M_WAITOK);
	netmsg_init(&np->np_base, NULL, &netisr_afree_rport, 0,
	    pktgen_start_ifsq_handler);
	np->np_pktg = pktg;
	np->np_ifsq = ifsq;

	lwkt_sendmsg(netisr_cpuport(ifsq_get_cpuid(ifsq)), &np->np_base.lmsg);
}

static int
pktgen_start(struct pktgen *pktg, int mq)
{
	struct ifaltq *ifq;

	if ((pktg->pktg_flags & PKTG_F_CONFIG) == 0)
		return EINVAL;
	if (pktg->pktg_flags & PKTG_F_RUNNING)
		return EBUSY;
	pktg->pktg_flags |= PKTG_F_RUNNING;

	ifq = &pktg->pktg_ifp->if_snd;
	if (!mq) {
		pktgen_ref(pktg);
		pktgen_start_ifsq(pktg, ifq_get_subq_default(ifq));
	} else {
		int i;

		for (i = 0; i < ifq->altq_subq_cnt; ++i)
			pktgen_ref(pktg);
		for (i = 0; i < ifq->altq_subq_cnt; ++i)
			pktgen_start_ifsq(pktg, ifq_get_subq(ifq, i));
	}
	return 0;
}

static void
pktgen_start_ifsq_handler(netmsg_t nmsg)
{
	struct netmsg_pktgen *np = (struct netmsg_pktgen *)nmsg;
	struct pktgen *pktg = np->np_pktg;
	struct ifaltq_subque *ifsq = np->np_ifsq;

	struct mbuf *m, *head = NULL, **next;
	struct ifnet *ifp;
	struct pktgen_pcpu *p;
	int cpuid, i, alloc_cnt, keep_cnt;

	u_short ulen, psum;
	int len, ip_len;

	/* Reply ASAP */
	lwkt_replymsg(&np->np_base.lmsg, 0);

	ifp = pktg->pktg_ifp;

	cpuid = ifsq_get_cpuid(ifsq);
	KKASSERT(cpuid == mycpuid);

	p = &pktg->pktg_pcpu[cpuid];

	keep_cnt = pktg->pktg_pktenq;
	alloc_cnt = keep_cnt * 2;

	/*
	 * Prefault enough mbuf into mbuf objcache
	 */
	next = &head;
	for (i = 0; i < alloc_cnt; ++i) {
		MGETHDR(m, M_WAITOK, MT_DATA);
		*next = m;
		next = &m->m_nextpkt;
	}

	for (i = 0; i < alloc_cnt - keep_cnt; ++i) {
		m = head;
		head = m->m_nextpkt;
		m->m_nextpkt = NULL;
		m_freem(m);
	}
	KKASSERT(head != NULL);

	/*
	 * Setup the packets' data
	 */
	ip_len = pktg->pktg_datalen + sizeof(struct udpiphdr);
	len = ip_len + ETHER_HDR_LEN;

	psum = htons((u_short)pktg->pktg_datalen + sizeof(struct udphdr) +
	    IPPROTO_UDP);
	ulen = htons(pktg->pktg_datalen + sizeof(struct udphdr));

	m = head;
	i = 0;
	while (m != NULL) {
		struct mbuf *nextm;
		const struct sockaddr_in *dst;
		struct pktgen_buf *pb;
		struct ip *ip;
		struct udpiphdr *ui;
		struct ether_header *eh;

		pktgen_ref(pktg);

		pb = kmalloc(sizeof(*pb), M_PKTGEN, M_WAITOK | M_ZERO);
		pb->pb_ifp = ifp;
		pb->pb_ifsq = ifsq;
		pb->pb_inuse = 1;
		pb->pb_buf = kmalloc(PKTGEN_BUFSZ, M_PKTGEN, M_WAITOK);
		pb->pb_len = len;
		pb->pb_cpuid = cpuid;
		pb->pb_pktg = pktg;
		netmsg_init(&pb->pb_nmsg, NULL, &netisr_adone_rport, 0,
		    pktgen_buf_send);
		LIST_INSERT_HEAD(&p->pktg_buflist, pb, pb_link);

		dst = &pktg->pktg_dst[i % pktg->pktg_ndst];
		++i;

		m->m_ext.ext_arg = pb;
		m->m_ext.ext_buf = pb->pb_buf;
		m->m_ext.ext_free = pktgen_buf_free;
		m->m_ext.ext_ref = pktgen_buf_ref;
		m->m_ext.ext_size = PKTGEN_BUFSZ;

		m->m_data = m->m_ext.ext_buf;
		m->m_flags |= M_EXT;
		m->m_len = m->m_pkthdr.len = len;

		m->m_data += ETHER_HDR_LEN;
		m->m_len -= ETHER_HDR_LEN;
		m->m_pkthdr.len -= ETHER_HDR_LEN;

		ui = mtod(m, struct udpiphdr *);
		ui->ui_pr = IPPROTO_UDP;
		if (pktg->pktg_flags & PKTG_F_SWITCH_SRCDST) {
			ui->ui_src.s_addr = dst->sin_addr.s_addr;
			ui->ui_dst.s_addr = pktg->pktg_src.sin_addr.s_addr;
			ui->ui_sport = dst->sin_port;
			ui->ui_dport = pktg->pktg_src.sin_port;
		} else {
			ui->ui_src.s_addr = pktg->pktg_src.sin_addr.s_addr;
			ui->ui_dst.s_addr = dst->sin_addr.s_addr;
			ui->ui_sport = pktg->pktg_src.sin_port;
			ui->ui_dport = dst->sin_port;
		}
		ui->ui_ulen = ulen;
		ui->ui_sum = in_pseudo(ui->ui_src.s_addr, ui->ui_dst.s_addr,
		    psum);
		m->m_pkthdr.csum_data = offsetof(struct udphdr, uh_sum);

		ip = (struct ip *)ui;
		ip->ip_len = ip_len;
		ip->ip_ttl = 64;	/* XXX */
		ip->ip_tos = 0;		/* XXX */
		ip->ip_vhl = IP_VHL_BORING;
		ip->ip_off = 0;
		ip->ip_sum = 0;
		ip->ip_id = ip_newid();

		in_delayed_cksum(m);

		ip->ip_len = htons(ip->ip_len);
		ip->ip_sum = in_cksum_hdr(ip);

		m->m_data -= ETHER_HDR_LEN;
		m->m_len += ETHER_HDR_LEN;
		m->m_pkthdr.len += ETHER_HDR_LEN;

		eh = mtod(m, struct ether_header *);
		bcopy(pktg->pktg_dst_lladdr, eh->ether_dhost, ETHER_ADDR_LEN);
		bcopy(IF_LLADDR(ifp), eh->ether_shost, ETHER_ADDR_LEN);
		eh->ether_type = htons(ETHERTYPE_IP);

		nextm = m->m_nextpkt;
		m->m_nextpkt = NULL;

		ifq_dispatch(ifp, m, NULL);

		m = nextm;
	}

	callout_reset(&p->pktg_stop, pktg->pktg_duration * hz,
	    pktgen_pcpu_stop_cb, p);

	pktgen_free(pktg);
}

static void
pktgen_mbuf(struct pktgen_buf *pb, struct mbuf *m)
{
	m->m_ext.ext_arg = pb;
	m->m_ext.ext_buf = pb->pb_buf;
	m->m_ext.ext_free = pktgen_buf_free;
	m->m_ext.ext_ref = pktgen_buf_ref;
	m->m_ext.ext_size = PKTGEN_BUFSZ;

	m->m_data = m->m_ext.ext_buf;
	m->m_flags |= M_EXT;
	m->m_len = m->m_pkthdr.len = pb->pb_len;
}

static void
pktgen_buf_send(netmsg_t msg)
{
	struct pktgen_buf *pb = (struct pktgen_buf *)msg;
	struct mbuf *m;

	KKASSERT(&curthread->td_msgport == netisr_cpuport(pb->pb_cpuid));

	crit_enter();
	lwkt_replymsg(&pb->pb_nmsg.lmsg, 0);
	crit_exit();

	MGETHDR(m, M_WAITOK, MT_DATA);
	pktgen_mbuf(pb, m);
	ifq_dispatch(pb->pb_ifp, m, NULL);
}

static void
pktgen_buf_free(void *arg)
{
	struct pktgen_buf *pb = arg;
	struct mbuf *m;

	KKASSERT(pb->pb_inuse > 0);
	if (pb->pb_done) {
		if (atomic_fetchadd_int(&pb->pb_inuse, -1) == 1) {
			struct pktgen *pktg;

			pktg = pb->pb_pktg;
			crit_enter();
			LIST_REMOVE(pb, pb_link);
			crit_exit();
			kfree(pb->pb_buf, M_PKTGEN);
			kfree(pb, M_PKTGEN);

			pktgen_free(pktg);
		}
		return;
	}

	if (&curthread->td_msgport != netisr_cpuport(pb->pb_cpuid)) {
		KKASSERT(pb->pb_cpuid == mycpuid);
		crit_enter();
		KKASSERT(pb->pb_nmsg.lmsg.ms_flags & MSGF_DONE);
		lwkt_sendmsg(netisr_cpuport(pb->pb_cpuid), &pb->pb_nmsg.lmsg);
		crit_exit();
		return;
	}

	MGETHDR(m, M_WAITOK, MT_DATA);
	pktgen_mbuf(pb, m);
	ifsq_enqueue(pb->pb_ifsq, m, NULL);
}

static void
pktgen_buf_ref(void *arg)
{
	struct pktgen_buf *pb = arg;

	panic("%s should never be called\n", __func__);

	KKASSERT(pb->pb_inuse > 0);
	atomic_add_int(&pb->pb_inuse, 1);
}

static void
pktgen_free(struct pktgen *pktg)
{
	KKASSERT(pktg->pktg_refcnt > 0);
	if (atomic_fetchadd_int(&pktg->pktg_refcnt, -1) == 1) {
		int i;

		if (pktg->pktg_dst != NULL)
			kfree(pktg->pktg_dst, M_PKTGEN);

		for (i = 0; i < ncpus; ++i)
			KKASSERT(LIST_EMPTY(&pktg->pktg_pcpu[i].pktg_buflist));
		kfree(pktg, M_PKTGEN);
	}

	KKASSERT(pktgen_refcnt > 0);
	atomic_subtract_int(&pktgen_refcnt, 1);
}

static void
pktgen_pcpu_stop_cb(void *arg)
{
	struct pktgen_pcpu *p = arg;
	struct pktgen_buf *pb;

	crit_enter();
	LIST_FOREACH(pb, &p->pktg_buflist, pb_link)
		pb->pb_done = 1;
	crit_exit();
}

static void
pktgen_ref(struct pktgen *pktg)
{
	atomic_add_int(&pktg->pktg_refcnt, 1);
	atomic_add_int(&pktgen_refcnt, 1);
}