/* This is a module which is used to relay packets between pairs of hosts.
 * The pairs are learned when packets are first sent.
 * Once a pair is known, the destinations are swapped and source becomes
 * this host.
 * Ports are enabled and disabled via netlink.
 * Requires kernel to be patched to allow spoofing.
 * 2003-10-xx, Don Mahurin, Creation.
 * 2003-02-xx, Don Mahurin, /proc support
 * 2003-05-26, Don Mahurin, ifdef netlink interface
 * 2006-08-31, Don Mahurin, update for latest netfilter
 */
#include <linux/mm.h>
#include <linux/slab.h>
#include <linux/module.h>
#include <linux/skbuff.h>
#include <linux/proc_fs.h>
#include <linux/ip.h>
#include <linux/udp.h>
#include <net/tcp.h>
#include <net/checksum.h>

#include <linux/netfilter_ipv4/ip_tables.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Don Mahurin <dmnf@openright.org>");
MODULE_DESCRIPTION("iptables RELAY target module");

//static spinlock_t relay_spin_lock = SPIN_LOCK_UNLOCKED; 

#ifdef ENABLE_RELAY_NETLINK
#ifdef EXTERNAL_DRIVER
#include "RELAY_nlm.h"
#else
#include <linux/netfilter_ipv4/RELAY_nlm.h>
#endif
#endif

#include <linux/spinlock.h>

static DEFINE_RWLOCK(relay_lock);

struct addr_port
{
	u_int32_t addr;
	u_int16_t port;
};

#define STATE_OFF 0
#define STATE_ON0 1
#define STATE_ON1 2
#define STATE_ON2 3
#define STATE_ON 3
#define STATE_ACTIVE1 4
#define STATE_ACTIVE2 8
#define STATE_ACTIVE 12
#define STATE_ALL 15

typedef struct port_relay
{
	unsigned char state;
	unsigned char fixed;
	struct addr_port addrs[2];
} port_relay;

static struct proc_dir_entry *proc_net_ipt_RELAY = NULL;
static struct proc_dir_entry *proc_net_ipt_RELAY_new = NULL;

#define MAX_PORTS 8192

static u_int16_t relay_port_min;
static u_int16_t relay_port_max;
static int total_port_relays = 0;

static port_relay *port_relays = NULL;

static int num_port_relays = 0;
static int next_port_relay = 0;

/* we store a (pid,port offset) table as a circular array */
typedef struct pid_port_relay
{
	int pid;
	int port_offset;
} pid_port_relay;

#define MAX_PID_PORT_RELAYS 4

static pid_port_relay pid_ports[MAX_PID_PORT_RELAYS];

static int next_pid_port_relay = 0;
static int first_pid_port_relay = 0;

static void set_pid_port_relay (int pid, int offset)
{
	//printk("set %d : %d \n", pid, offset);
	pid_ports[next_pid_port_relay].pid = pid;
	pid_ports[next_pid_port_relay].port_offset = offset;
	next_pid_port_relay++;
	if(next_pid_port_relay >= MAX_PID_PORT_RELAYS)
	{
		next_pid_port_relay = 0;
	}
	/* trim circular array if needed */
	if(next_pid_port_relay == first_pid_port_relay)
	{
		first_pid_port_relay++;
		if(first_pid_port_relay >= MAX_PID_PORT_RELAYS)
			first_pid_port_relay = 0;
	}
}

static int get_pid_port_relay(int pid)
{
	int r = -1;
	int i;
	//printk("get %d\n", pid);
	for(i = first_pid_port_relay; i != next_pid_port_relay;)
	{
		//printk("check %d %d %d\n", i, pid_ports[i].pid, next_pid_port_relay);

		if(pid_ports[i].pid == pid)
		{
			if(i == first_pid_port_relay)
			{
				first_pid_port_relay++;
				if(first_pid_port_relay >= MAX_PID_PORT_RELAYS)
					first_pid_port_relay = 0;
			}
			pid_ports[i].pid = 0;
			r = pid_ports[i].port_offset;
		}
		i++;
		if(i >= MAX_PID_PORT_RELAYS) i = 0;
	}
	return r;
}

static void init_port_relays(int toggle)
{
	int i;
	write_lock_bh(&relay_lock);
	for(i = 0; i < total_port_relays; i++)
	{
		port_relays[i].state = toggle ? STATE_ON0: STATE_OFF;
		port_relays[i].fixed = 0;
	}
	num_port_relays = total_port_relays * toggle;
	for(i = 0; i < MAX_PID_PORT_RELAYS; i++)
	{
		pid_ports[i].pid = 0;
	}
	write_unlock_bh(&relay_lock);
}

static port_relay *find_port_relay(unsigned short int port)
{
	int i = port - relay_port_min;
	port_relay* rtn;
	
	read_lock_bh(&relay_lock);
	if(i < 0 || i >= total_port_relays || ! (port_relays[i].state & STATE_ON)) rtn = NULL;
	else
		rtn = &port_relays[i];
	read_unlock_bh(&relay_lock);
	return rtn;
}

static void enable_port_relay(u_int16_t port, uint32_t addr1, uint32_t addr2)
{
	int i = port - relay_port_min;
	unsigned char fixed = 0;

	write_lock_bh(&relay_lock);
	if(total_port_relays==0 || i < 0 ||
	port_relays[i].state & STATE_ON)
	{
		write_unlock_bh(&relay_lock);
		return;
	}

	if(addr1) fixed |= 1;
	port_relays[i].addrs[0].addr = addr1;
	port_relays[i].addrs[0].port = 0;
	if(addr2) fixed |= 2;
	port_relays[i].addrs[1].addr = addr2;
	port_relays[i].addrs[1].port = 0;
	port_relays[i].state |= STATE_ON0;
	port_relays[i].fixed = fixed;

	num_port_relays++;

	if(num_port_relays == total_port_relays)
		next_port_relay = total_port_relays;
	else if(i == next_port_relay)
		while(port_relays[++next_port_relay].state & STATE_ON) {}
	write_unlock_bh(&relay_lock);
}

static int new_relay_port(void)
{
	if( num_port_relays < total_port_relays)
	{
		enable_port_relay(relay_port_min + next_port_relay,
			0,0);
		return relay_port_min + next_port_relay;
	}
	else
		return -1;
}

#ifdef ENABLE_RELAY_NETLINK
static void disable_port_relay(unsigned short int port)
{
	int i = port - relay_port_min;
	write_lock_bh(&relay_lock);
	if(i < 0 ||
		! (port_relays[i].state & STATE_ON))
	{
		write_unlock_bh(&relay_lock);
		return;
	}

	port_relays[i].state &= (~(STATE_ON | STATE_ACTIVE));

	if( i < next_port_relay)
		next_port_relay = i;
	
	num_port_relays--;
	write_unlock_bh(&relay_lock);
}
#endif

static int get_relay_host(u_int16_t port, u_int32_t saddr, u_int16_t sport, u_int32_t *daddr, u_int16_t *dport)
{
	port_relay *relay = find_port_relay(port);
	int state_on;
	if(relay == NULL)
	{
		printk("relay not found (%d)\n", port);
		return 0;
	}
	state_on=relay->state & STATE_ON;

	printk("relay %d state (%d,%d): %x\n", port, relay->state, state_on, saddr);
	if(state_on == STATE_ON2)
	{
		if(saddr == relay->addrs[0].addr && 
			sport == relay->addrs[0].port)
		{
			*daddr = relay->addrs[1].addr;
			*dport = relay->addrs[1].port;
			if(! (relay->state & STATE_ACTIVE1))
			{
				write_lock_bh(&relay_lock);
				relay->state |= STATE_ACTIVE1;
				write_unlock_bh(&relay_lock);
			}
			printk("left %d to %d\n", sport, *dport);
			return 1;
		}
		else if(saddr == relay->addrs[1].addr &&
			sport == relay->addrs[1].port)
		{
			*daddr = relay->addrs[0].addr;
			*dport = relay->addrs[0].port;
			if(! (relay->state & STATE_ACTIVE2))
			{
				write_lock_bh(&relay_lock);
				relay->state |= STATE_ACTIVE2;
				write_unlock_bh(&relay_lock);
			}
			printk("right %d to %d\n", sport, *dport);
			return 1;
		}
		else if(relay->fixed)
		{
			//printk("fixed\n");
			return 0;
		}
		else // unknown - reset
		{
			printk("another\n");
			write_lock_bh(&relay_lock);
			state_on = relay->state = STATE_ON0;
			write_unlock_bh(&relay_lock);
		}
	}
	{
		struct addr_port *paddr = relay->addrs;
		if( state_on == STATE_ON0 || saddr != paddr->addr ||
			sport != paddr->port)
		{
			int n = state_on - 1;
			if(relay->fixed & (n+1) && saddr != paddr[n].addr)
				return 0;
			printk("new(%d) %d %d\n", port, sport, relay->state);
			paddr[n].addr = saddr;
			paddr[n].port = sport;
			state_on++;
			write_lock_bh(&relay_lock);
			relay->state = state_on;
			write_unlock_bh(&relay_lock);
		}
		else
		{
			printk("same(%d) %d %d\n", port, sport, relay->state);
		}
	}
	if(state_on == STATE_ON2)
		return get_relay_host(port,saddr,sport,daddr,dport);
	return 0;
}

#ifdef ENABLE_RELAY_NETLINK

static struct sock *relaynl;
static DECLARE_MUTEX(relaynl_sem);

static int
relay_rcv_nl_event(struct notifier_block *this,
                 unsigned long event, void *ptr)
{
	return NOTIFY_DONE;
}

static struct notifier_block relay_nl_notifier = {
	relay_rcv_nl_event,
	NULL,
	0
};

static int
relay_receive_peer(struct relay_peer_msg *pmsg,
                 unsigned char type, unsigned int len)
{
	int status = 0;

	if (len < sizeof(*pmsg))
		return -EINVAL;

	switch (type) {
	case RELAY_NLM_ENABLE_ALL:
		init_port_relays(1);
		break;
	case RELAY_NLM_DISABLE_ALL:
		init_port_relays(0);
		break;
	case RELAY_NLM_ENABLE:
		enable_port_relay(pmsg->port, pmsg->addr1, pmsg->addr2);
		printk("added port %u\n", pmsg->port);
		break;
	case RELAY_NLM_DISABLE:
		disable_port_relay(pmsg->port);
		break;
	default:
		status = -EINVAL;
	}
	return status;
}


#define RCV_SKB_FAIL(err) do { netlink_ack(skb, nlh, (err)); return; } while (0)

static inline void
relay_rcv_skb(struct sk_buff *skb)
{
	int status, type, pid, flags, nlmsglen, skblen;
	struct nlmsghdr *nlh;

	skblen = skb->len;
	if (skblen < sizeof(*nlh))
		return;

	nlh = (struct nlmsghdr *)skb->data;
	nlmsglen = nlh->nlmsg_len;
	if (nlmsglen < sizeof(*nlh) || skblen < nlmsglen)
		return;

	pid = nlh->nlmsg_pid;
	flags = nlh->nlmsg_flags;
	
	if(pid <= 0 || !(flags & NLM_F_REQUEST) || flags & NLM_F_MULTI)
		RCV_SKB_FAIL(-EINVAL);
		
	if (flags & MSG_TRUNC)
		RCV_SKB_FAIL(-ECOMM);
		
	type = nlh->nlmsg_type;
	if (type < NLMSG_NOOP || type >= RELAY_NLM_MAX)
		RCV_SKB_FAIL(-EINVAL);
		
	if (type <= RELAY_NLM_BASE)
		return;
		
	if(!cap_raised(NETLINK_CB(skb).eff_cap, CAP_NET_ADMIN))
		RCV_SKB_FAIL(-EPERM);
	
	status = relay_receive_peer(NLMSG_DATA(nlh), type,
	                          skblen - NLMSG_LENGTH(0));
	if (status < 0)
		RCV_SKB_FAIL(status);
		
	if (flags & NLM_F_ACK)
		netlink_ack(skb, nlh, 0);
        return;
}

static void
relay_rcv_sk(struct sock *sk, int len)
{
	do {
		struct sk_buff *skb;

		if (down_trylock(&relaynl_sem))
			return;
			
		while ((skb = skb_dequeue(&sk->receive_queue)) != NULL) {
			relay_rcv_skb(skb);
			kfree_skb(skb);
		}
		
		up(&relaynl_sem);

	} while (relaynl && relaynl->receive_queue.qlen);
}

#endif

static int
init_or_cleanup(int init)
{
	int status = -ENOMEM;
	
	if (!init)
		goto cleanup;

#ifdef ENABLE_RELAY_NETLINK
	netlink_register_notifier(&relay_nl_notifier);
	relaynl = netlink_kernel_create(NETLINK_FIREWALL, relay_rcv_sk);
	if (relaynl == NULL) {
		printk(KERN_ERR "ip_queue: failed to create netlink socket\n");
		goto cleanup_netlink_notifier;
	}
#endif

	status = 0;

	return status;

cleanup:
	nf_unregister_queue_handler(PF_INET);
//        spin_lock_bh(&relay_spin_lock);
//        spin_lock_bh(&relay_spin_unlock);
	
#ifdef ENABLE_RELAY_NETLINK
	sock_release(relaynl->socket);
	down(&relaynl_sem);
	up(&relaynl_sem);
	
cleanup_netlink_notifier:

	netlink_unregister_notifier(&relay_nl_notifier);
#endif
	return status;
}

static u_int16_t
cheat_check(u_int32_t oldvalinv, u_int32_t newval, u_int16_t oldcheck)
{
	u_int32_t diffs[] = { oldvalinv, newval };
	return csum_fold(csum_partial((char *)diffs, sizeof(diffs),
				      oldcheck^0xFFFF));
}

static unsigned int
target(struct sk_buff **pskb,
       const struct net_device *in,
       const struct net_device *out,
       unsigned int hooknum,
       const struct xt_target *target,
       const void *targinfo,
       void *userinfo)
{
	struct iphdr *iph = (*pskb)->nh.iph;
	void *ipdata = (void *)iph + iph->ihl * 4;

	u_int32_t destaddr;
	u_int16_t destport;

	if(iph->protocol == IPPROTO_TCP)
	{
		struct tcphdr *tcph = (struct tcphdr *)ipdata;
		if(! get_relay_host(ntohs(tcph->dest), iph->saddr, ntohs(tcph->source), &destaddr, &destport)) return IPT_CONTINUE; 
		destport = htons(destport);

		tcph->check = cheat_check(~iph->daddr,
			destaddr,
			cheat_check(tcph->dest ^ 0xFFFF,
				destport,
				tcph->check));

		/* change source port to relay port */
		tcph->check = cheat_check(tcph->source ^ 0xFFFF,
					tcph->dest,
					tcph->check);

		tcph->check = cheat_check(~iph->saddr,
				iph->daddr, tcph->check);

		tcph->source = tcph->dest;
		tcph->dest = destport;
	}
	else if(iph->protocol == IPPROTO_UDP)
	{
		struct udphdr *udph = (struct udphdr *)ipdata;
		if(! get_relay_host(ntohs(udph->dest), iph->saddr, ntohs(udph->source), &destaddr, &destport)) return IPT_CONTINUE; 
		destport = htons(destport);

		if (udph->check) /* 0 is a special case meaning no checksum */
		{
			udph->check = cheat_check(~iph->daddr,
				destaddr,
				cheat_check(udph->dest ^ 0xFFFF,
					destport,
					udph->check));

			/* change source port to relay port */
			udph->check = cheat_check(udph->source ^ 0xFFFF,
					udph->dest,
					udph->check);

			udph->check = cheat_check(~iph->saddr,
				iph->daddr, udph->check);

			udph->source = udph->dest;
			udph->dest = destport;
		}

	}
	iph->check = cheat_check(~iph->daddr,
		 destaddr, iph->check);

	iph->check = cheat_check(~iph->saddr,
		 iph->daddr, iph->check);

	iph->saddr = iph->daddr;
	iph->daddr = destaddr;

//	(*pskb)->nfcache |= NFC_ALTERED;

	return IPT_CONTINUE;
}

static int get_port_range(const struct ipt_entry *e,
		u_int16_t *start,
		u_int16_t *end)
{
	int p = sizeof(struct ipt_entry);
	while(p < e->target_offset && p < e->next_offset)
	{
		struct ipt_entry_match *m = (struct ipt_entry_match *)((void *)e + p);
		if (strcmp(m->u.kernel.match->name, "tcp") == 0)
		{
			const struct ipt_tcp *tcpinfo = (const struct ipt_tcp *)m->data;
			*start = tcpinfo->dpts[0];
			*end = tcpinfo->dpts[1];
			return 1;
		}
		else if(strcmp(m->u.kernel.match->name, "udp") == 0)
		{
			const struct ipt_udp *udpinfo = (const struct ipt_udp *)m->data;
			*start = udpinfo->dpts[0];
			*end = udpinfo->dpts[1];
			return 1;
		}
		p += sizeof(struct ipt_entry_match);
	}
	return 0;
}

static int
checkentry(const char *tablename,
           const void *e_void,
           const struct xt_target *target,
           void *targinfo,
           unsigned int targinfosize,
           unsigned int hook_mask)
{
	const struct ipt_entry *e = e_void;
	int total;

/*
	if (strcmp(tablename, "mangle") != 0) {
		printk(KERN_WARNING "RELAY: can only be called from \"mangle\" table, not \"%s\"\n", tablename);
		return 0;
	}
*/

	if(!get_port_range(e, &relay_port_min, &relay_port_max))
	{
		printk("problem getting port range\n");
		return 0;
	}

	total = relay_port_max - relay_port_min + 1;
	if(total > MAX_PORTS)
		return 0;

	port_relays = (port_relay *)kmalloc(sizeof(port_relay) * total, GFP_KERNEL);
	if(port_relays == NULL)
		return 0;

	total_port_relays = total;
	init_port_relays(0);
	printk("relay %d %d %d\n", relay_port_min, relay_port_max, total_port_relays);

	return 1;
}

static struct ipt_target ipt_relay_reg
= {
	.name = "RELAY",
	.target = target,
        .targetsize = 0,
        .table = "mangle",
	.hooks = (1 << NF_IP_PRE_ROUTING) | (1 << NF_IP_FORWARD) | (1 << NF_IP_POST_ROUTING),
	.checkentry = checkentry,
	.me = THIS_MODULE
};

static int proc_read_relay_new(char *page, char **start, off_t offset, int length, int *eof, void *data)
{
	int len;
	int port;

	if(offset)
		return 0;
	port = new_relay_port();
	len = sprintf(page, "%d\n", port);

	return len;
}

static int proc_relay_get_state(char *page, char **start, off_t offset, int length, int *eof, void *data)
{
	int len;
	int i;

	*eof = 1;

	read_lock_bh(&relay_lock);
	i = get_pid_port_relay(current->pid);
	//if(offset)
//		return 0;

	if(i >= total_port_relays || i < 0)
		len = 0;
	else
		len = sprintf(page, "%d\n", port_relays[i].state);
	
	read_unlock_bh(&relay_lock);

	return len;
}

static int my_atoi(const char *buf, int max, int *offset)
{
	int n = 0;
	int i = *offset;
	while(i < max && (buf[i] == '\n' || buf[i] == ' ')) { i++; }
	if( i >= max) { *offset = i; return -1; }
        while(i < max && buf[i] && buf[i] != '\n' && buf[i] != ' ')
	{
		if(buf[i] < '0' || buf[i] > '9')
		{
			n = -1;
			break;
		}
		//printk("b %c %d\n", buf[i],  n);
		n= 10 * n + (buf[i] - '0');
		//printk("c %d\n", n);
		i++;
	}
	*offset = i;
	return n;
}

static int proc_relay_set_state(struct file *file, const char *buffer, unsigned long count, void *data)
{
	int offset = 0;
	int port = my_atoi(buffer, count, &offset);
	int newstate = my_atoi(buffer, count, &offset);
	int i = port - relay_port_min;

	if(i < 0 || i >= total_port_relays)
		return count;

	//printk("set a %d %d\n", port, state;


	write_lock_bh(&relay_lock);
	if(newstate >= 0)
	{
		int maxstate = port_relays[i].state & STATE_ALL;
		if(maxstate & STATE_ON)
		{
			if(!(maxstate & newstate & STATE_ON1))
				maxstate &= ~STATE_ACTIVE1;
			if(!(maxstate & newstate & STATE_ON2))
				maxstate &= ~STATE_ACTIVE2;
		}
		else if(newstate == STATE_ALL || newstate == STATE_ON0 )
			maxstate = STATE_ON0;
		else
			maxstate &= STATE_OFF;
		port_relays[i].state = maxstate & newstate;
	}
	else
		set_pid_port_relay(current->pid, i);
	write_unlock_bh(&relay_lock);


	return count;
}

static int __init init(void)
{
	proc_net_ipt_RELAY = proc_mkdir("ipt_RELAY",proc_net);
	if(!proc_net_ipt_RELAY) return -ENOMEM;

	proc_net_ipt_RELAY_new = create_proc_entry("new", 0644, proc_net_ipt_RELAY);
	proc_net_ipt_RELAY_new->owner = THIS_MODULE;
	proc_net_ipt_RELAY_new->read_proc = proc_read_relay_new;

	proc_net_ipt_RELAY_new = create_proc_entry("state", 0644, proc_net_ipt_RELAY);
	proc_net_ipt_RELAY_new->owner = THIS_MODULE;
	proc_net_ipt_RELAY_new->write_proc = proc_relay_set_state;
	proc_net_ipt_RELAY_new->read_proc = proc_relay_get_state;

	init_or_cleanup(1);
	if (ipt_register_target(&ipt_relay_reg))
		return -EINVAL;

	return 0;
}

static void __exit fini(void)
{
	init_or_cleanup(0);
	ipt_unregister_target(&ipt_relay_reg);
}

module_init(init);
module_exit(fini);