#include <linux/etherdevice.h>
#include <linux/ethtool.h>
#include <linux/if_arp.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/netdevice.h>
#include <linux/rculist.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 quicktun_struct {
	struct net_device	*dev;
	unsigned long		flags;

	struct sockaddr_in	local_address;
	struct sockaddr_in	remote_address;

	struct socket		*sock;
};


struct quicktun_list_entry {
	struct list_head	list;
	struct rcu_head		rcu;
	struct quicktun_struct	*quicktun;
};


static netdev_tx_t quicktun_net_xmit(struct sk_buff *skb, struct net_device *dev)
{
	struct quicktun_struct *tun = netdev_priv(dev);
	struct msghdr msg;
	struct kvec vec;
	int err;
	
	if (tun->remote_address.sin_addr.s_addr) {
		err = skb_linearize(skb);
		if (err < 0)
			goto drop;

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

		msg.msg_name = &tun->remote_address;
		msg.msg_namelen = sizeof(tun->remote_address);
		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 drop;

	consume_skb(skb);

	return NETDEV_TX_OK;

 drop:
	kfree_skb(skb);

	tun->dev->stats.tx_dropped++;
	return NETDEV_TX_OK;
}

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;


	skb = skb_recv_datagram(sk, 0, 1, &err);
	if (!skb) {
		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;
	}

	netif_rx_ni(skb);

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

 drop:
	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,
};

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


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);
	tun->remote_address.sin_family = AF_INET;
	tun->remote_address.sin_addr.s_addr = 0;
	tun->remote_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->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 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;
	struct quicktun_list_entry *entry;
	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]);

	entry = kmalloc(sizeof(struct quicktun_struct), GFP_KERNEL);
	dev = alloc_netdev(sizeof(struct quicktun_struct), name, quicktun_setup);
	
	if (!entry || !dev) {
		err = -ENOMEM;
		goto err_free;
	}

	dev_net_set(dev, net);

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

	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.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.sin_addr.s_addr = nla_get_be32(info->attrs[QUICKTUN_A_REMOTE_ADDRESS]);

	if (info->attrs[QUICKTUN_A_REMOTE_PORT])
		tun->remote_address.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;

	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_rcu(&entry->list, &quicktun_list);

	rtnl_unlock();

	return 0;

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

 err_free:
	if (dev) free_netdev(dev);
	if (entry) kfree(entry);
	return err;
}

static void quicktun_free_list_entry(struct rcu_head *rcu) {
	struct quicktun_list_entry *entry = container_of(rcu, struct quicktun_list_entry, rcu);
	kfree(entry);
}

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

	list_for_each_entry(entry, &quicktun_list, list) {
		if (entry->quicktun->dev == dev) {
			break;
		}
	}

	list_del_rcu(&entry->list);
	call_rcu(&entry->rcu, quicktun_free_list_entry);

	unregister_netdevice(dev);
	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 struct genl_family quicktun_nl_family = {
	.id		= GENL_ID_GENERATE,
	.name		= DRV_NAME,
	.version	= 1,
	.hdrsize	= 0,
	.maxattr	= QUICKTUN_A_MAX,
};


static struct nla_policy quicktun_nl_policy[__QUICKTUN_A_MAX] = {
	[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,
	},
};

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_list_entry *entry, *next;

	genl_unregister_family(&quicktun_nl_family);

	rtnl_lock();

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

	rtnl_unlock();
}


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