#include <linux/etherdevice.h>
#include <linux/ethtool.h>
#include <linux/if_arp.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/list.h>
#include <linux/kthread.h>
#include <linux/netdevice.h>
#include <linux/rcupdate.h>
#include <linux/udp.h>
#include <net/dst.h>
#include <net/genetlink.h>
#include <net/sock.h>

#include "quicktun.h"


#define DRV_NAME	"quicktun"
#define DRV_VERSION	"0.1"
#define DRV_DESCRIPTION	"Quicktun tunnel driver"

#define QUICKTUN_MODE_MASK		0x000f
#define QUICKTUN_FLAG_REMOTE_FLOAT	0x0010

#define MIN_MTU 68
#define MAX_MTU 65535


static LIST_HEAD(quicktun_list);


struct addr_struct {
	struct rcu_head		rcu;
	struct sockaddr_in	addr;
};

struct quicktun_struct {
	struct list_head	list;

	struct net_device	*dev;
	unsigned long		flags;

	struct sockaddr_in	local_address;
	struct addr_struct	*remote_address;

	struct socket		*sock;
};

struct send_work {
	struct work_struct	work;
	struct quicktun_struct	*quicktun;
	struct sk_buff		*skb;
};


static void quicktun_send_bh(struct work_struct *work)
{
	struct send_work *send_work = container_of(work, struct send_work, work);
	struct sk_buff *skb = send_work->skb;
	struct quicktun_struct *tun = send_work->quicktun;
	struct addr_struct *remote_addr;
	struct msghdr msg;
	struct kvec vec;
	int err;
	
	rcu_read_lock();

	remote_addr = rcu_dereference(tun->remote_address);

	if (remote_addr->addr.sin_addr.s_addr) {
		err = skb_linearize(skb);
		if (err < 0)
			goto error;

		vec.iov_base = skb->data;
		vec.iov_len = skb->len;

		msg.msg_name = &remote_addr->addr;
		msg.msg_namelen = sizeof(remote_addr->addr);
		msg.msg_control = NULL;
		msg.msg_controllen = 0;
		msg.msg_flags = MSG_DONTWAIT;

		kernel_sendmsg(tun->sock, &msg, &vec, 1, skb->len);

		tun->dev->stats.tx_packets++;
		tun->dev->stats.tx_bytes += skb->len;
	}
	else
		goto error;

	rcu_read_unlock();

	consume_skb(skb);
	kfree(send_work);

	return;

error:
	rcu_read_unlock();

	kfree_skb(skb);
	kfree(send_work);

	tun->dev->stats.tx_errors++;
}


static netdev_tx_t quicktun_net_xmit(struct sk_buff *skb, struct net_device *dev)
{
	struct send_work *work = kmalloc(sizeof(struct send_work), GFP_ATOMIC);
	if (work) {
		INIT_WORK(&work->work, quicktun_send_bh);
		work->quicktun = netdev_priv(dev);
		work->skb = skb;
		schedule_work(&work->work);
	}

	return NETDEV_TX_OK;
}


static void free_addr_struct(struct rcu_head *rcu)
{
	struct addr_struct *addr = container_of(rcu, struct addr_struct, rcu);
	kfree(addr);
}

static void quicktun_net_uninit(struct net_device *dev)
{
	struct quicktun_struct *tun = netdev_priv(dev);
	call_rcu(&rtnl_dereference(tun->remote_address)->rcu, free_addr_struct);
}


static void quicktun_udp_data_ready(struct sock *sk, int bytes)
{
	struct sk_buff *skb;
	unsigned int len;
	struct quicktun_struct *tun = sk->sk_user_data;
	int err;


	read_lock(&sk->sk_callback_lock);

	skb = skb_recv_datagram(sk, 0, 1, &err);
	if (!skb) {
		read_unlock(&sk->sk_callback_lock);
		if (err == -EAGAIN)
			return;
		pr_info("Quicktun: UDP socket error %d", err);
		return;
	}

	skb_orphan(skb);

	err = skb_linearize(skb);
	if (err < 0)
		goto drop;

	__skb_pull(skb, sizeof(struct udphdr));

	len = skb->len;

	switch (tun->flags & QUICKTUN_MODE_MASK) {
	case QUICKTUN_MODE_ETHERNET:
		if (len < ETH_HLEN)
			goto drop;
		skb->protocol = eth_type_trans(skb, tun->dev);
		break;

	case QUICKTUN_MODE_IP:
		if (len < sizeof(struct iphdr))
			goto drop;

		__skb_tunnel_rx(skb, tun->dev);
		skb_reset_network_header(skb);

		switch (skb->data[0] & 0xf0) {
		case 0x40:
			skb->protocol = htons(ETH_P_IP);
			break;
		case 0x60:
			if (len < sizeof(struct ipv6hdr))
				goto drop;
			skb->protocol = htons(ETH_P_IPV6);
			break;
		default:
			goto drop;
		}
		break;

	default:
		goto drop;
	}

	read_unlock(&sk->sk_callback_lock);

	netif_rx(skb);

	tun->dev->stats.rx_packets++;
	tun->dev->stats.rx_bytes += len;
	return;

drop:
	read_unlock(&sk->sk_callback_lock);

	tun->dev->stats.rx_dropped++;
	kfree_skb(skb);
}

static int quicktun_net_change_mtu(struct net_device *dev, int new_mtu)
{
	if (new_mtu < MIN_MTU || new_mtu + dev->hard_header_len > MAX_MTU)
		return -EINVAL;
	dev->mtu = new_mtu;
	return 0;
}


static const struct net_device_ops quicktun_ethernet_netdev_ops = {
	.ndo_start_xmit		= quicktun_net_xmit,
	.ndo_change_mtu		= quicktun_net_change_mtu,
	.ndo_set_mac_address	= eth_mac_addr,
	.ndo_validate_addr	= eth_validate_addr,
	.ndo_uninit		= quicktun_net_uninit,
};

static const struct net_device_ops quicktun_ip_netdev_ops = {
	.ndo_start_xmit		= quicktun_net_xmit,
	.ndo_change_mtu		= quicktun_net_change_mtu,
	.ndo_uninit		= quicktun_net_uninit,
};


static void quicktun_get_drvinfo(struct net_device *dev, struct ethtool_drvinfo *info)
{
	struct quicktun_struct *tun = netdev_priv(dev);

	strcpy(info->driver, DRV_NAME);
	strcpy(info->version, DRV_VERSION);
	strcpy(info->fw_version, "N/A");

	switch (tun->flags & QUICKTUN_MODE_MASK) {
	case QUICKTUN_MODE_ETHERNET:
		strcpy(info->bus_info, "ethernet");
		break;
	case QUICKTUN_MODE_IP:
		strcpy(info->bus_info, "ip");
		break;
	}
}

static u32 quicktun_always_on(struct net_device *dev)
{
	return 1;
}


static const struct ethtool_ops quicktun_ethtool_ops = {
	.get_drvinfo	= quicktun_get_drvinfo,
	.get_link	= quicktun_always_on,
};


static void quicktun_setup(struct net_device *dev)
{
	struct quicktun_struct *tun = netdev_priv(dev);

	tun->local_address.sin_family = AF_INET;
	tun->local_address.sin_addr.s_addr = 0;
	tun->local_address.sin_port = htons(QUICKTUN_DEFAULT_PORT);

	dev->ethtool_ops = &quicktun_ethtool_ops;
}

static void quicktun_net_init(struct net_device *dev)
{
	struct quicktun_struct *tun = netdev_priv(dev);

	switch (tun->flags & QUICKTUN_MODE_MASK) {
	case QUICKTUN_MODE_ETHERNET:
		dev->netdev_ops = &quicktun_ethernet_netdev_ops;

		ether_setup(dev);

		random_ether_addr(dev->dev_addr);
		break;

	case QUICKTUN_MODE_IP:
		dev->netdev_ops = &quicktun_ip_netdev_ops;

		dev->destructor = free_netdev;
		dev->hard_header_len = 0;
		dev->addr_len = 0;
		dev->mtu = 1500;

		dev->type = ARPHRD_NONE;
		dev->flags = IFF_POINTOPOINT | IFF_NOARP | IFF_MULTICAST;
		break;
	}
}

static int quicktun_socket_init(struct net_device *dev)
{
	int err;
	struct quicktun_struct *tun = netdev_priv(dev);


	err = sock_create_kern(AF_INET, SOCK_DGRAM, 0, &tun->sock);
	if (err < 0)
		return err;

	tun->sock->sk->sk_data_ready = quicktun_udp_data_ready;
	tun->sock->sk->sk_user_data = tun;

	err = kernel_bind(tun->sock, (struct sockaddr*)&tun->local_address, sizeof(tun->local_address));
	if (err < 0)
		goto error;

	return 0;
error:
	sock_release(tun->sock);
	return err;
}


static struct genl_family quicktun_nl_family = {
	.id		= GENL_ID_GENERATE,
	.name		= DRV_NAME,
	.version	= 1,
	.hdrsize	= 0,
	.maxattr	= QUICKTUN_A_MAX,
};


static int quicktun_cmd_create_device(struct sk_buff *skb, struct genl_info *info)
{
	unsigned long flags = 0;
	__u16 type = 0;
	struct quicktun_struct *tun;
	struct net *net = genl_info_net(info);
	struct net_device *dev;
	char *name = "qt%d";
	int err;

	if (!capable(CAP_NET_ADMIN))
		return -EPERM;

	if (info->attrs[QUICKTUN_A_MODE])
		type = nla_get_u16(info->attrs[QUICKTUN_A_MODE]);

	switch (type) {
	case QUICKTUN_MODE_ETHERNET:
	case QUICKTUN_MODE_IP:
		flags |= type;
		break;
	default:
		return -EINVAL;
	}

	if (info->attrs[QUICKTUN_A_IFNAME])
		name = nla_data(info->attrs[QUICKTUN_A_IFNAME]);

	dev = alloc_netdev(sizeof(struct quicktun_struct), name, quicktun_setup);
	
	if (!dev)
		return -ENOMEM;

	dev_net_set(dev, net);

	tun = netdev_priv(dev);
	tun->dev = dev;
	tun->flags = flags;

	tun->remote_address = kmalloc(sizeof(struct addr_struct), GFP_KERNEL);
	if (!tun->remote_address) {
		err = -ENOMEM;
		goto err_free_dev;
	}

	tun->remote_address->addr.sin_family = AF_INET;
	tun->remote_address->addr.sin_addr.s_addr = 0;
	tun->remote_address->addr.sin_port = htons(QUICKTUN_DEFAULT_PORT);

	quicktun_net_init(dev);
	
	if (info->attrs[QUICKTUN_A_LOCAL_ADDRESS])
		tun->local_address.sin_addr.s_addr = nla_get_be32(info->attrs[QUICKTUN_A_LOCAL_ADDRESS]);

	/* Set remote port to local port by default */
	if (info->attrs[QUICKTUN_A_LOCAL_PORT])
		tun->remote_address->addr.sin_port = tun->local_address.sin_port = nla_get_be16(info->attrs[QUICKTUN_A_LOCAL_PORT]);

	if (info->attrs[QUICKTUN_A_REMOTE_ADDRESS])
		tun->remote_address->addr.sin_addr.s_addr = nla_get_be32(info->attrs[QUICKTUN_A_REMOTE_ADDRESS]);

	if (info->attrs[QUICKTUN_A_REMOTE_PORT])
		tun->remote_address->addr.sin_port = nla_get_be16(info->attrs[QUICKTUN_A_REMOTE_PORT]);

	if (info->attrs[QUICKTUN_A_REMOTE_FLOAT]) {
		if(nla_get_flag(info->attrs[QUICKTUN_A_REMOTE_FLOAT]))
			tun->flags |= QUICKTUN_FLAG_REMOTE_FLOAT;
		else
			tun->flags &= ~QUICKTUN_FLAG_REMOTE_FLOAT;
	}

	err = quicktun_socket_init(dev);
	if (err < 0)
		goto err_free_dev;

	rtnl_lock();

	if (strchr(dev->name, '%')) {
		err = dev_alloc_name(dev, dev->name);
		if (err < 0)
			goto err_free_sk;
	}

	err = register_netdevice(tun->dev);
	if (err < 0)
		goto err_free_sk;

	list_add_tail(&tun->list, &quicktun_list);

	rtnl_unlock();

	return 0;

err_free_sk:
	rtnl_unlock();
	sock_release(tun->sock);

err_free_dev:
	free_netdev(dev);
	return err;
}


static void __quicktun_delete_device(struct net_device *dev) {
	struct quicktun_struct *tun = netdev_priv(dev);
	struct socket *sock = tun->sock;

	kernel_sock_shutdown(sock, SHUT_RD);

	list_del(&tun->list);

	unregister_netdevice(dev);

	flush_scheduled_work();
	sock_release(sock);
}

static int quicktun_cmd_delete_device(struct sk_buff *skb, struct genl_info *info)
{
	int err;
	struct net *net = genl_info_net(info);
	struct net_device *dev;
	char *name;

	if (!capable(CAP_NET_ADMIN))
		return -EPERM;

	if (!info->attrs[QUICKTUN_A_IFNAME])
		return -EINVAL;

	name = nla_data(info->attrs[QUICKTUN_A_IFNAME]);

	rtnl_lock();

	dev = __dev_get_by_name(net, name);
	if (!dev) {
		err = -EINVAL;
		goto err_unlock;
	}

	if (dev->netdev_ops != &quicktun_ethernet_netdev_ops && dev->netdev_ops != &quicktun_ip_netdev_ops) {
		err = -EINVAL;
		goto err_unlock;
	}

	__quicktun_delete_device(dev);

	rtnl_unlock();

	return 0;

err_unlock:
	rtnl_unlock();

	return err;
}

static void __put_tunnel_spec(struct sk_buff *skb, struct quicktun_struct *tun)
{
	struct nlattr *nested = nla_nest_start(skb, QUICKTUN_A_TUNNEL_SPEC);
	struct addr_struct *remote_addr;

	remote_addr = rtnl_dereference(tun->remote_address);

	nla_put_string(skb, QUICKTUN_A_IFNAME, tun->dev->name);
	nla_put_u16(skb, QUICKTUN_A_MODE, tun->flags & QUICKTUN_MODE_MASK);
	nla_put_u32(skb, QUICKTUN_A_LOCAL_ADDRESS, tun->local_address.sin_addr.s_addr);
	nla_put_u16(skb, QUICKTUN_A_LOCAL_PORT, tun->local_address.sin_port);
	nla_put_u32(skb, QUICKTUN_A_REMOTE_ADDRESS, remote_addr->addr.sin_addr.s_addr);
	nla_put_u16(skb, QUICKTUN_A_REMOTE_PORT, remote_addr->addr.sin_port);

	if (tun->flags & QUICKTUN_FLAG_REMOTE_FLOAT)
		nla_put_flag(skb, QUICKTUN_A_REMOTE_FLOAT);

	nla_nest_end(skb, nested);
}

static int quicktun_cmd_get_device(struct sk_buff *skb, struct genl_info *info)
{
	struct sk_buff *ret;
	void *skb_head;
	int err;
	char *name = NULL;

	ret = genlmsg_new(NLMSG_GOODSIZE, GFP_KERNEL);
	if (ret == NULL)
		return -ENOMEM;

	skb_head = genlmsg_put(ret, 0, info->snd_seq, &quicktun_nl_family, 0, QUICKTUN_CMD_GET_DEVICE);

	if (skb_head == NULL)
		return -ENOMEM;

	if (info->attrs[QUICKTUN_A_IFNAME])
		name = nla_data(info->attrs[QUICKTUN_A_IFNAME]);

	rtnl_lock();

	if (name) {
		struct net_device *dev;

		dev = __dev_get_by_name(genl_info_net(info), name);
		if (!dev) {
			err = -EINVAL;
			goto err_unlock;
		}

		__put_tunnel_spec(ret, netdev_priv(dev));
	}
	else {
		struct quicktun_struct *entry, *next;

		list_for_each_entry_safe(entry, next, &quicktun_list, list)
			__put_tunnel_spec(ret, entry);
	}

	rtnl_unlock();

	genlmsg_end(ret, skb_head);

	err = genlmsg_unicast(genl_info_net(info), ret, info->snd_pid);
	if (err)
		return err;

	return 0;

err_unlock:
	rtnl_unlock();

	return err;
}


static struct nla_policy quicktun_nl_policy[__QUICKTUN_A_MAX] = {
	[QUICKTUN_A_TUNNEL_SPEC]	= { .type = NLA_NESTED },
	[QUICKTUN_A_IFNAME]		= { .type = NLA_NUL_STRING, .len = IFNAMSIZ - 1 },
	[QUICKTUN_A_MODE]		= { .type = NLA_U16 },
	[QUICKTUN_A_LOCAL_ADDRESS]	= { .type = NLA_U32 },
	[QUICKTUN_A_LOCAL_PORT]		= { .type = NLA_U16 },
	[QUICKTUN_A_REMOTE_ADDRESS]	= { .type = NLA_U32 },
	[QUICKTUN_A_REMOTE_PORT]	= { .type = NLA_U16 },
	[QUICKTUN_A_REMOTE_FLOAT]	= { .type = NLA_FLAG },
};

static struct genl_ops quicktun_nl_ops[] = {
	{
		.cmd = QUICKTUN_CMD_CREATE_DEVICE,
		.doit = quicktun_cmd_create_device,
		.policy = quicktun_nl_policy,
	},
	{
		.cmd = QUICKTUN_CMD_DELETE_DEVICE,
		.doit = quicktun_cmd_delete_device,
		.policy = quicktun_nl_policy,
	},
	{
		.cmd = QUICKTUN_CMD_GET_DEVICE,
		.doit = quicktun_cmd_get_device,
		.policy = quicktun_nl_policy,
	},
};

static int __init quicktun_init(void)
{
	int ret = 0;

	pr_info("%s, %s\n", DRV_DESCRIPTION, DRV_VERSION);

	ret = genl_register_family_with_ops(&quicktun_nl_family, quicktun_nl_ops, ARRAY_SIZE(quicktun_nl_ops));
	return ret;
}

static void quicktun_exit(void)
{
	struct quicktun_struct *entry, *next;

	genl_unregister_family(&quicktun_nl_family);

	rtnl_lock();

	list_for_each_entry_safe(entry, next, &quicktun_list, list) {
		__quicktun_delete_device(entry->dev);
	}

	rtnl_unlock();
}


module_init(quicktun_init);
module_exit(quicktun_exit);
MODULE_DESCRIPTION(DRV_DESCRIPTION);
MODULE_LICENSE("GPL");