#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/netdevice.h>
#include <linux/rcupdate.h>
#include <linux/udp.h>
#include <net/dst.h>
#include <net/genetlink.h>
#include <net/ip.h>
#include <net/route.h>
#include <net/sock.h>

#include "quicktun.h"
#include "compat.h"


#define DRV_NAME	"quicktun"
#define DRV_VERSION	"1.0"
#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;

	spinlock_t		lock;

	struct net_device	*dev;
	unsigned long		flags;

	struct sockaddr_in	local_address;
	struct addr_struct	*remote_address;

	struct socket		*sock;
};


static netdev_tx_t quicktun_net_xmit(struct sk_buff *skb, struct net_device *dev)
{
	struct quicktun_struct *tun = netdev_priv(dev);
	struct addr_struct *remote_addr;
	struct rtable *rt = NULL;
	struct iphdr *iph;
	struct udphdr *udph;
	__wsum  csum;
	unsigned int len;
	int err;

	rcu_read_lock();

	remote_addr = rcu_dereference(tun->remote_address);

	if (likely(remote_addr->addr.sin_addr.s_addr)) {
		rt = ip_route_output_ports(dev_net(dev), NULL, remote_addr->addr.sin_addr.s_addr, tun->local_address.sin_addr.s_addr,
					   remote_addr->addr.sin_port, tun->local_address.sin_port, IPPROTO_UDP, 0, 0);
		if (!rt) {
			tun->dev->stats.tx_carrier_errors++;
			goto error;
		}

		if (RTABLE_DST(rt)->dev == dev) {
			tun->dev->stats.collisions++;
			goto error;
		}

		if (skb_headroom(skb) < (LL_RESERVED_SPACE(RTABLE_DST(rt)->dev) + sizeof(struct iphdr) + sizeof(struct udphdr))
		    || skb_shared(skb) || (skb_cloned(skb) && !skb_clone_writable(skb, 0))) {
			struct sk_buff *new_skb = skb_realloc_headroom(skb, LL_RESERVED_SPACE(RTABLE_DST(rt)->dev) + sizeof(struct iphdr) + sizeof(struct udphdr));
			if (!new_skb) {
				tun->dev->stats.tx_dropped++;
				goto error;
			}

			dev_kfree_skb(skb);
			skb = new_skb;
		}

		len = skb->len;

		skb_push(skb, sizeof(struct iphdr) + sizeof(struct udphdr));
		skb_reset_network_header(skb);

		skb_dst_drop(skb);
		skb_dst_set(skb, RTABLE_DST(rt));

		iph           = (struct iphdr *)skb->data;
		iph->version  = 4;
		iph->ihl      = sizeof(struct iphdr)>>2;
		iph->frag_off = 0;
		iph->protocol = IPPROTO_UDP;
		iph->tos      = 0;
		iph->saddr    = rt->rt_src;
		iph->daddr    = rt->rt_dst;
		iph->tot_len  = htons(len + sizeof(struct iphdr) + sizeof(struct udphdr));
		iph->ttl      = 64;
		iph->check    = 0;

		udph          = (struct udphdr *)(skb->data + sizeof(struct iphdr));
		udph->source  = tun->local_address.sin_port;
		udph->dest    = remote_addr->addr.sin_port;
		udph->len     = htons(len + sizeof(struct udphdr));
		udph->check   = 0;

		ip_select_ident(ip_hdr(skb), RTABLE_DST(rt), NULL);

		csum = csum_partial(udph, ntohs(udph->len), 0);
		udph->check = csum_tcpudp_magic(iph->saddr, iph->daddr, ntohs(udph->len), IPPROTO_UDP, csum);
		skb->ip_summed = CHECKSUM_NONE;

		nf_reset(skb);
		skb->mark = 0;
		err = ip_local_out(skb);

		if (likely(net_xmit_eval(err) == 0)) {
			tun->dev->stats.tx_bytes += len;
			tun->dev->stats.tx_packets++;
		} else {
			tun->dev->stats.tx_errors++;
			tun->dev->stats.tx_aborted_errors++;
		}
	}
	else {
		tun->dev->stats.tx_carrier_errors++;
		goto error;
	}

	rcu_read_unlock();

	return NETDEV_TX_OK;

error:
	rcu_read_unlock();
	dev_kfree_skb(skb);

	if (rt)
		ip_rt_put(rt);

	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(&rcu_dereference_protected(tun->remote_address, 1)->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;
	struct addr_struct *remote_addr;
	int err;


	read_lock(&sk->sk_callback_lock);

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

	skb_orphan(skb);
	nf_reset(skb);

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

	rcu_read_lock();

	remote_addr = rcu_dereference(tun->remote_address);

	if (unlikely((remote_addr->addr.sin_addr.s_addr != ip_hdr(skb)->saddr) || (remote_addr->addr.sin_port != udp_hdr(skb)->source))) {
		rcu_read_unlock();

		if (ACCESS_ONCE(tun->flags) & QUICKTUN_FLAG_REMOTE_FLOAT) {
			struct addr_struct *addr = kmalloc(sizeof(struct addr_struct), GFP_KERNEL);
			if (!addr)
				goto drop;

			addr->addr.sin_addr.s_addr = ip_hdr(skb)->saddr;
			addr->addr.sin_port = udp_hdr(skb)->source;

			spin_lock(&tun->lock);

			if (tun->flags & QUICKTUN_FLAG_REMOTE_FLOAT) {
				remote_addr = rcu_dereference_protected(tun->remote_address, lockdep_is_held(&tun->lock));
				rcu_assign_pointer(tun->remote_address, addr);
				call_rcu(&remote_addr->rcu, free_addr_struct);
			}

			spin_unlock(&tun->lock);
		}
		else
			goto drop;
	}
	else
		rcu_read_unlock();

	__skb_pull(skb, sizeof(struct udphdr));

	len = skb->len;

	__skb_tunnel_rx(skb, tun->dev);

	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_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_ni(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);

	spin_lock_init(&tun->lock);

	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);

		dev->hard_header_len = LL_MAX_HEADER + sizeof(struct iphdr) + sizeof(struct udphdr) + ETH_HLEN;
		dev->mtu = ETH_DATA_LEN - sizeof(struct iphdr) - sizeof(struct udphdr) - ETH_HLEN;

		random_ether_addr(dev->dev_addr);
		break;

	case QUICKTUN_MODE_IP:
		dev->netdev_ops = &quicktun_ip_netdev_ops;

		dev->destructor = free_netdev;
		dev->addr_len = 4;
		dev->hard_header_len = LL_MAX_HEADER + sizeof(struct iphdr) + sizeof(struct udphdr);
		dev->mtu = ETH_DATA_LEN - sizeof(struct iphdr) - sizeof(struct udphdr);

		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)
{
	__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]);

	if (type != QUICKTUN_MODE_ETHERNET && type != QUICKTUN_MODE_IP)
		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 = type;

	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_u8(info->attrs[QUICKTUN_A_REMOTE_FLOAT]))
			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);

	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 int quicktun_cmd_change_device(struct sk_buff *skb, struct genl_info *info)
{
	int err;
	struct net *net = genl_info_net(info);
	struct addr_struct *old_addr, *addr;
	struct net_device *dev;
	struct quicktun_struct *tun;
	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;
	}

	tun = netdev_priv(dev);

	addr = kmalloc(sizeof(struct addr_struct), GFP_KERNEL);
	if (!addr) {
		err = -ENOMEM;
		goto err_unlock;
	}

	spin_lock(&tun->lock);

	old_addr = rcu_dereference_protected(tun->remote_address, lockdep_is_held(&tun->lock));

	addr->addr = old_addr->addr;

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

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

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

	rcu_assign_pointer(tun->remote_address, addr);
	call_rcu(&old_addr->rcu, free_addr_struct);

	spin_unlock(&tun->lock);

	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;

	rcu_read_lock();

	remote_addr = rcu_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);
	nla_put_u8(skb, QUICKTUN_A_REMOTE_FLOAT, (tun->flags & QUICKTUN_FLAG_REMOTE_FLOAT) ? 1 : 0);

	rcu_read_unlock();

	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_U8 },
};

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_CHANGE_DEVICE,
		.doit = quicktun_cmd_change_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");