#include	<linux/module.h>
#include	<linux/skbuff.h>
#include	<linux/ip.h>
#include        <linux/in.h>
#include	<asm/uaccess.h>
#include	<linux/mm.h>
#include	<linux/time.h>

/* Include all protocols we supposed to know headers of */
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/icmp.h>

#ifndef CONFIG_NETFILTER
#define CONFIG_NETFILTER
#endif

#include <linux/netfilter_ipv4/ip_tables.h>
#include <linux/netfilter_ipv4/ipt_ACC.h>
#include <linux/netfilter_ipv4/lockhelp.h>

static DECLARE_RWLOCK(acc_lock);

static	struct acc_ctrl		main_table;
static	struct acc_ctrl		back_table;
static	int	hash_size = 4096;
static	int	hash_module = 1021;
static	int	mem_limit = 16;


static void free_table(struct acc_ctrl *table)
{
	struct acc_block	*blk;
	while( (blk=table->heap) ) {
		table->heap=blk->next;
		kfree(blk);
	}
	if(table->hash)
		kfree(table->hash);
	memset(table,0,sizeof(struct acc_ctrl));
}

static void init_table(struct acc_ctrl *tbl)
{
	memset(tbl,0,sizeof(struct acc_ctrl));
	tbl->hash_size=hash_size;
	tbl->hash_module=hash_module;
	if( (tbl->hash=kmalloc(tbl->hash_size,GFP_KERNEL)) == 0 ) {
		printk("ACC: hash allocation failure\n");
	} else
		memset(tbl->hash,0,tbl->hash_size);
	tbl->last_heap = tbl->heap = kmalloc(ACC_HEAP_BLK_SZ,GFP_KERNEL);
	if( tbl->heap != 0) {
		memset(tbl->heap,0,ACC_HEAP_BLK_SZ);
		tbl->blocks_count=1;
		tbl->free_count=ACC_ENTRIES_PER_BLOCK;
	}
}

inline void backup_table(void)
{
	struct acc_ctrl tmp;
	struct	timeval	t;
	free_table(&back_table);
	init_table(&tmp);
	do_gettimeofday(&t);
	WRITE_LOCK(&acc_lock);
	memcpy(&back_table,&main_table,sizeof(struct acc_ctrl));
	memcpy(&main_table,&tmp,sizeof(struct acc_ctrl));
	memcpy(&back_table.time_off,&t,sizeof(struct timeval));
	memcpy(&main_table.time_on,&t,sizeof(struct timeval));
	WRITE_UNLOCK(&acc_lock);
}

inline int add_heap_blk(void)
{
	struct acc_block *blk;
	if( main_table.blocks_count >= mem_limit || (blk=kmalloc(ACC_HEAP_BLK_SZ,GFP_ATOMIC)) == 0)
		return 0;
	memset(blk,0,ACC_HEAP_BLK_SZ);
	if( ! main_table.heap )
		main_table.heap = blk;
	else if( main_table.last_heap )
		main_table.last_heap->next = blk;
	main_table.last_heap = blk;
	// skip  blk->next = NULL because blk is already zeroed
	main_table.free_count = ACC_ENTRIES_PER_BLOCK;
	main_table.blocks_count++;
	return 1;
}

//FIXME: Write a better hash function for use with protoz and ports
inline int hash_func_acc(__u32 src,__u32 dst,__u32 mark)
{
	return (
		( (src % 4073) << 20) + ( (dst % 4079) << 8 ) +
		(mark % 251) ) % main_table.hash_module;
}

inline struct acc_entry *new_entry(__u32 src, __u16 sprt, __u32 dst, __u16 dprt, __u32 mark,__u32 count, __u8 proto, __u8 hooknum)
{
	struct acc_entry        *e;
	if( main_table.free_count == 0 && add_heap_blk() == 0)
		return NULL;
	e = &main_table.last_heap->bl[ACC_ENTRIES_PER_BLOCK-main_table.free_count--];
	e->src=src;
	e->sprt=sprt;
	e->dst=dst;
	e->dprt=dprt;
	e->mark=mark;
	e->count=count;
	e->proto=proto;
	e->hooknum=hooknum;
	// e->next already zeroed by add_heap_blk
	return e;
}

inline __u64 lost(__u32 count) { return main_table.lost += count;}

static __u64 new_packet(__u32 src, __u16 sprt, __u32 dst, __u16 dprt, __u32 mark,__u32 count, __u8 proto, __u8 hooknum)
{
	int	h;
	struct acc_entry	*e;
	if( main_table.hash == 0)
		return lost(count);
	h=hash_func_acc(src,dst,mark);
	for(e = main_table.hash[h]; e ;e=e->next )
		if( e->src == src && e->sprt==sprt && e->dst == dst && e->dprt==dprt && e->mark == mark && e->proto == proto && e->hooknum == hooknum)
				return e->count += count;
		else if( e->next == 0 ) {
			if( (e->next=new_entry(src,sprt, dst, dprt, mark,count,proto,hooknum)) == 0)
				return lost(count);
			return 0;
		}
	if( (main_table.hash[h] = new_entry(src,sprt,dst, dprt,mark,count,proto,hooknum)) == 0)
		return lost(count);
	return 0;
}

static int set_acc_ctl(struct sock *sk, int cmd, void *user, unsigned int len)
{

	struct ip_acc_set_rq req;
	if (!capable(CAP_NET_ADMIN))
		return -EPERM;
	if (cmd != SO_IP_ACC)
		return -EBADF;
	if (len != sizeof(req))
		return -EINVAL;
	if (copy_from_user(&req, user, sizeof(req)) != 0)
		return -EFAULT;
	if ( req.op & IP_ACC_SET_HASH )
		switch ( req.hash_size ) {
			case 1:
				hash_size = 4096;
				hash_module = 1021;
				break;
			case 4:
				hash_size = 16384;
				hash_module = 4093;
				break;
			case 16:
				hash_size = 65536;
				hash_module = 16381;
				break;
			default:
				return -EINVAL;
		}
	if( req.op & IP_ACC_SET_LIMIT ) {
		if( req.mem_limit >= 16 && req.mem_limit <= 1024 )
			mem_limit = req.mem_limit;
		else
			return -EINVAL;
	}
	if( req.op & IP_ACC_SET_FLUSH )
		backup_table();
	return 0;
}

static int get_acc_ctl(struct sock *sk, int cmd, void *user, int *len)
{
	struct acc_block *blk;
	int	nblk;
	struct ip_acc_get_info	info;
	if (!capable(CAP_NET_ADMIN))
		return -EPERM;
	if (cmd == SO_IP_ACC_INFO) {
		if(*len != sizeof(info))
			return -EINVAL;
		info.blocks = back_table.blocks_count;
		memcpy(&info.time_on,&back_table.time_on,sizeof(struct timeval));
		memcpy(&info.time_off,&back_table.time_off,sizeof(struct timeval));
		info.lost = back_table.lost;
		if (copy_to_user(user,&info,sizeof(info)) != 0)
			return -EFAULT;
		return 0;
	}
	else if (cmd == SO_IP_ACC_BLOCK) {
		if(*len != sizeof(struct acc_block))
			return -EINVAL;
		copy_from_user(&nblk,user,sizeof(int));
		if( nblk < 0 || nblk >= back_table.blocks_count )
			return -EINVAL;
		for( blk = back_table.heap; blk ; blk=blk->next,nblk-- ) {
			if( nblk == 0 ) {
				copy_to_user(user,blk,sizeof(struct acc_block));
				return 0;
			}
		}
		return -EINVAL;
	}
	return -EBADF;
}

static unsigned int
target(struct sk_buff **pskb,
       unsigned int hooknum,
       const struct net_device *in,
       const struct net_device *out,
       const void *targinfo,
       void *userinfo)
{ 
	const struct ipt_acc_target_info *accinfo = targinfo;
	struct iphdr *iph = (*pskb)->nh.iph;
	
	struct icmphdr *icmp_hdr;
	//	struct igmphdr *igmp_hdr;
	struct tcphdr *tcp_hdr;
	struct udphdr *udp_hdr;

	__u16 sport=0;
	__u16 dport=0;

	//FIXME: Rewrite this uglu pointers to smart and pretty new one :)
	switch(iph->protocol)
	  {
	  case IPPROTO_ICMP: 
	    icmp_hdr = (*pskb)->h.icmph;
	    sport=icmp_hdr->type;
	    dport=icmp_hdr->code;
	    break;
	  
	  case IPPROTO_IGMP:
	    //FIXME: Add normal IGMP headers processing. ICMP has almost the same, so it works.
	    icmp_hdr = (*pskb)->h.icmph;
	    sport=icmp_hdr->type;
	    dport=icmp_hdr->code;
	    break;

	  case IPPROTO_TCP: 
	    tcp_hdr = (*pskb)->h.th;
	    //We don't do ntohs() here. In userspace we trust. So Don't forget to modify ipacc.c
	
	    sport=tcp_hdr->source;
	    dport=tcp_hdr->dest;	    

	    break;

	  case IPPROTO_UDP: 
	    udp_hdr = (*pskb)->h.uh;
	    sport=udp_hdr->source;
	    dport=udp_hdr->dest;
	    break;

	  default : break; //If we don't know portnumbers for that proto, it's 0 :)

	  }

	
	printk("%d: %d -> %d || %d -> %d \n",hooknum, sport, dport, ntohs(sport), ntohs(dport));
	WRITE_LOCK(&acc_lock);
		new_packet(iph->saddr, sport, iph->daddr, dport ,accinfo->mark,ntohs(iph->tot_len), iph->protocol, hooknum);
	WRITE_UNLOCK(&acc_lock);
	return IPT_CONTINUE;
}

static int
checkentry(const char *tablename,
	   const struct ipt_entry *e,
           void *targinfo,
           unsigned int targinfosize,
           unsigned int hook_mask)
{
	if (targinfosize != IPT_ALIGN(sizeof(struct ipt_acc_target_info))) {
		printk(KERN_WARNING "ACC: targinfosize %u != %Zu\n",
		       targinfosize,
		       IPT_ALIGN(sizeof(struct ipt_acc_target_info)));
		return 0;
	}
	return 1;
}


static struct ipt_target ipt_acc_reg                                           
= { { NULL, NULL }, "ACC", target, checkentry, NULL, THIS_MODULE };            

static struct nf_sockopt_ops so_acc
= { { NULL, NULL }, PF_INET,
    SO_IP_ACC, SO_IP_ACC+1, &set_acc_ctl,
    SO_IP_ACC, SO_IP_ACC+2, &get_acc_ctl,
    0, NULL };


static int __init init(void)
{
	struct	timeval	t;
	if (ipt_register_target(&ipt_acc_reg))
		return -EINVAL;
	if( nf_register_sockopt(&so_acc) ) {
		ipt_unregister_target(&ipt_acc_reg);
		return -EINVAL;
	}
	memset(&back_table,0,sizeof(struct acc_ctrl));
	init_table(&main_table);
	do_gettimeofday(&t);
	memcpy(&main_table.time_on,&t,sizeof(struct timeval));
	return 0;
}

static void __exit fini(void)
{
	free_table(&main_table);
	free_table(&back_table);
	nf_unregister_sockopt(&so_acc);
	ipt_unregister_target(&ipt_acc_reg);
}

module_init(init);
module_exit(fini);