/*
 * netfilter module 'ipt_netquota'
 * handles one quota per IP within a larger network.
 *
 * 06/2004 Olaf Rempel <razzor at kopf minus tisch dot de>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * USAGE:
 * $ iptables -A FORWARD -m netquota --nq-name internet --nq-src 192.168.1.0/24 -j ACCEPT
 *
 * quotas can be read via procfs:
 * cat /proc/net/ipt/netquota/matchname
 *
 * quotas can be set via procfs:
 * echo "192.168.1.1 =1000000" > /proc/net/ipt_netquota/matchname
 * sets quota of 192.168.1.1 to 1000000 bytes (of data within IP datagram)
 *
 * echo "192.168.1.1 +1000000" > /proc/net/ipt_netquota/matchname
 * adds 1000000 bytes to quota of 192.168.1.1
 *
 * echo "192.168.1.1 -1000000" > /proc/net/ipt_netquota/matchname
 * decreases quota of 192.168.1.1 by 1000000 bytes
 *
 * NOTE:
 * netquotas can be referenced from different tables/chains.
 * e.g. to limit traffic from and to a proxy using only one quota, add this:
 * $iptables -A INPUT -p tcp --dport 3128 -m netquota --nq-name proxy --nq-src 192.168.1.0/24 -j ACCEPT
 * $iptables -A OUTPUT -p tcp --sport 3128 -m netquota --nq-name proxy --nq-dst 192.168.1.0/24 -j ACCEPT
 *
 * Network and quote name must be the same to allow this.
 *
 * TODO:
 * - module parameter check
 * - read_locks in nq_*_proc?
 */

#include <linux/module.h>
#include <linux/skbuff.h>
#include <linux/proc_fs.h>
#include <linux/spinlock.h>
#include <linux/vmalloc.h>
#include <linux/ctype.h>

#include <linux/ip.h>

#include <linux/netfilter_ipv4/ip_tables.h>
#include <linux/netfilter_ipv4/ipt_netquota.h>

/*
 * DEBUG == 0 -> no debugging
 * DEBUG |= 1 -> module and match insert/remove debugging
 * DEBUG |= 2 -> proc interface write debugging
 * DEBUG |= 4 -> proc interface read debugging
 * DEBUG |= 8 -> netfilter match debugging
 */
#define DEBUG 0

/* default module parameters */
static unsigned int max_num_of_matches = 16;
static unsigned int proc_file_perms = 0644;
static unsigned int max_netmask = 16;

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("Netwide quota match");
MODULE_AUTHOR("Olaf Rempel <razzor at kopf minus tisch dot de>");
MODULE_PARM(max_num_of_matches, "b");
MODULE_PARM_DESC(max_num_of_matches,
		 " maximum number of netquota matches (default: 16)");
MODULE_PARM(proc_file_perms, "i");
MODULE_PARM_DESC(proc_file_perms,
		 " permissions of /proc/net/ipt_netquota/* files (default: 0644)");
MODULE_PARM(max_netmask, "i");
MODULE_PARM_DESC(max_netmask,
		 " maximum possible network mask of each match (default: 16 == 255.255.0.0)");

/* per IP struct */
struct nq_stat {
	u_int64_t quota;
	spinlock_t lock;
};

/* per match struct */
struct nq_match {
	char name[IPT_NETQUOTA_NAME_LEN];
	struct nq_stat *stat;
	u_int32_t network;
	u_int32_t netmask;
	u_int32_t use_count;
	struct proc_dir_entry *proc_entry;
};

/* global match array */
static struct nq_match **nq_match_arr = NULL;
static rwlock_t nq_match_arr_lock = RW_LOCK_UNLOCKED;

/* /proc/net/ipt_netquota/ directory */
static struct proc_dir_entry *proc_net_ipt_netquota = NULL;

/* procfs read function */
static int
nq_read_proc(char *buffer, char **start, off_t offset,
	     int length, int *eof, void *data)
{
	struct nq_match *this_match = (struct nq_match *)data;
	unsigned int index, count = 0, inc = 0;
	u_int32_t address;

	while (length - count > 40) {
		index = offset + inc;
		address = index + this_match->network;

		if ((address & this_match->netmask) != this_match->network) {
			*eof = 1;
			break;
		}

		spin_lock(&this_match->stat[index].lock);
		count += snprintf(buffer + count, length - count,
				  "%u.%u.%u.%u - %llu\n",
				  HIPQUAD(address),
				  this_match->stat[index].quota);
		spin_unlock(&this_match->stat[index].lock);

		inc++;
	}

#if ((DEBUG & 4) == 4)
	printk(KERN_INFO "ipt_netquota: proc_read: offset=%lu length=%d "
	       "inc=%d count=%d *eof= %d\n", offset, length, inc, count, *eof);
#endif
	/* offset will be offset += inc on next call */
	*start = (char *)inc;
	return count;
}

/* parse IP address */
static int
parse_ip(char *cp, unsigned int *address)
{
	unsigned int val, retval = 0;
	char *old_cp = cp, digit, parts = 4;

	while (parts--) {
		digit = 3;
		val = 0;

		while (isdigit(*cp) && digit--) {
			val = (val * 10) + (*cp - '0');
			cp++;
		}

		if (val > 255 || digit == 3)
			return old_cp - cp;

		if (*cp == '.') {
			retval = (retval << 8) | (val & 0xFF);
			cp++;
		} else {
			cp++;
			break;
		}
	}

	if (parts == 0 && *cp != 0) {
		*address = (retval << 8) | (val & 0xFF);
		return cp - old_cp;
	}

	return old_cp - cp;
}

/* procfs write function */
static int
nq_write_proc(struct file *file, const char __user * buffer,
	      unsigned long length, void *data)
{
	struct nq_match *this_match = (struct nq_match *)data;
	unsigned long len = (length > 40) ? 40 : length;
	unsigned int address = 0, index;
	char buf[40], *cp, mode;
	u_int64_t quota = 0;

	if (copy_from_user(buf, buffer, len))
		return -EFAULT;

	if (len < 40)
		buf[len] = 0;

	/* parse the IP */
	cp = buf;
	if ((len = parse_ip(cp, &address)) < 0)
		return -len;

	/* get mode (add/subtract/set) */
	cp += len;
	mode = *cp++;

	/* get value */
	quota = simple_strtoull(cp, &cp, 10);

	if ((address & this_match->netmask) != this_match->network) {
#if ((DEBUG & 2) == 2)
		printk(KERN_INFO "ipt_netquota: proc_write: address %u.%u.%u.%u "
		       "is not a member of match '%s'\n", HIPQUAD(address),
		       this_match->name);
#endif
		return cp - buf;
	}
#if ((DEBUG & 2) == 2)
	printk(KERN_INFO "ipt_netquota: proc_write: address=%u.%u.%u.%u "
	       "mode='%c' quota=%llu\n", HIPQUAD(address), mode, quota);
#endif

	index = address - this_match->network;
	spin_lock(&this_match->stat[index].lock);

	switch (mode) {
	/* ADD to quota */
	case '+':
		/* check for overflow */
		if (this_match->stat[index].quota + quota > quota)
			this_match->stat[index].quota += quota;
		break;

	/* SUBTRACT from quota */
	case '-':
		if (this_match->stat[index].quota >= quota) {
			this_match->stat[index].quota -= quota;
		} else {
			this_match->stat[index].quota = 0;
		}
		break;

	/* SET quota */
	case '=':
		this_match->stat[index].quota = quota;
		break;

	default:
		break;
	}

	spin_unlock(&this_match->stat[index].lock);
	return cp - buffer;
}

/* netfilter match hook, called on every packet that hits us */
static int
match(const struct sk_buff *skb,
      const struct net_device *in,
      const struct net_device *out,
      const void *matchinfo, int offset, int *hotdrop)
{
	struct ipt_netquota_info *info = (struct ipt_netquota_info *)matchinfo;
	struct nq_match *this_match;
	u_int32_t address, index;
	u_int16_t datalen;

	if (skb->len < sizeof(struct iphdr))
		return 1;

	/* datalen = ip datagram length - ip header */
	datalen = skb->len - skb->nh.iph->ihl * 4;

	read_lock(&nq_match_arr_lock);
	this_match = nq_match_arr[info->match_index];
	if (this_match == NULL) {
		printk(KERN_ERR "ipt_netquota: no such quota: '%s'\n", info->name);
		read_unlock(&nq_match_arr_lock);
		return 0;
	}

	/* look at source or destination IP? */
	switch (info->mode) {
	case IPT_NETQUOTA_SRC_MATCH:
		address = ntohl(skb->nh.iph->saddr);
		break;

	case IPT_NETQUOTA_DST_MATCH:
		address = ntohl(skb->nh.iph->daddr);
		break;

	default:
		read_unlock(&nq_match_arr_lock);
		return 0;
		break;
	}

	if ((address & this_match->netmask) != this_match->network) {
#if ((DEBUG & 8) == 8)
		printk(KERN_INFO "ipt_netquota: %u.%u.%u.%u is not a member of "
		       "match '%s'\n", HIPQUAD(address), info->name);
#endif
		read_unlock(&nq_match_arr_lock);
		return 0;
	}

	index = address - this_match->network;
	spin_lock(&this_match->stat[index].lock);

	/* check if we can afford this packet */
	if (this_match->stat[index].quota >= datalen) {
		this_match->stat[index].quota -= datalen;
		spin_unlock(&this_match->stat[index].lock);
#if ((DEBUG & 8) == 8)
		printk(KERN_INFO "ipt_netquota: quota of %u.%u.%u.%u reduced "
		       "by %u\n", HIPQUAD(address), datalen);
#endif
		read_unlock(&nq_match_arr_lock);
		return 1;
	}

	/* quota reached - do not allow any packets from now on */
	this_match->stat[index].quota = 0;
	spin_unlock(&this_match->stat[index].lock);

#if ((DEBUG & 8) == 8)
	printk(KERN_INFO "ipt_netquota: quota of %u.%u.%u.%u is now zero\n",
	       HIPQUAD(address));
#endif

	read_unlock(&nq_match_arr_lock);
	return 0;
}

/* netfilter checkentry hook, called to insert a new match */
static int
checkentry(const char *tablename,
	   const struct ipt_ip *ip,
	   void *matchinfo, unsigned int matchinfosize, unsigned int hook_mask)
{
	struct ipt_netquota_info *info = (struct ipt_netquota_info *)matchinfo;
	struct nq_match *new_match;
	unsigned int i, unused = -1, host_count;

	if (matchinfosize != IPT_ALIGN(sizeof(struct ipt_netquota_info)))
		return 0;

	if (!info->name)
		return 0;

	/* find quota name in match array */
	write_lock(&nq_match_arr_lock);
	for (i = 0; i < max_num_of_matches; i++) {

		/* no match / match deleted */
		if (nq_match_arr[i] == NULL) {
			/* remember first unsused position */
			if (unused == -1)
				unused = i;
			continue;
		}

		/* different name */
		if (strncmp
		    (info->name, nq_match_arr[i]->name,
		     IPT_NETQUOTA_NAME_LEN)) {
			continue;
		}

		/* different parameters (only check network/netmask) */
		if (info->network != nq_match_arr[i]->network ||
		    info->netmask != nq_match_arr[i]->netmask) {

			printk(KERN_ERR "ipt_netquota: quota '%s' already exists "
			       "with different network.\n",
			       nq_match_arr[i]->name);

			write_unlock(&nq_match_arr_lock);
			return 0;
		}

		/* ok, this must be the same match, increase reference count */
		nq_match_arr[i]->use_count++;
#if ((DEBUG & 1) == 1)
		printk(KERN_INFO "ipt_netquota: incremented use_count of "
		       "quota '%s' to %d\n", nq_match_arr[i]->name,
		       nq_match_arr[i]->use_count);
#endif
		write_unlock(&nq_match_arr_lock);
		return 1;
	}

	/* no free index found */
	if (unused == -1) {
		printk(KERN_ERR "ipt_netquota: All matches in use. "
		       "(try increasing max_num_of_matches parameter)\n");
		write_unlock(&nq_match_arr_lock);
		return 0;
	}

	/* check netmask */
	host_count = INADDR_BROADCAST - info->netmask + 1;
	if (host_count > (1 << (32 - max_netmask))) {
		printk(KERN_ERR "ipt_netquota: Requested network mask is too big. "
		       "(try increasing max_netmask parameter)\n");
		write_unlock(&nq_match_arr_lock);
		return 0;
	}
#if ((DEBUG & 1) == 1)
	printk(KERN_INFO
	       "ipt_netquota: trying to create new match (alloc %d + %d bytes)\n",
	       sizeof(struct nq_match), host_count * sizeof(struct nq_stat));
#endif

	/* alloc mem for a new netquota */
	new_match = vmalloc(sizeof(struct nq_match));
	if (new_match == NULL) {
		write_unlock(&nq_match_arr_lock);
		return -ENOMEM;
	}

	/* alloc mem for stat counters */
	new_match->stat = vmalloc(host_count * sizeof(struct nq_stat));
	if (new_match->stat == NULL) {
		write_unlock(&nq_match_arr_lock);
		vfree(new_match);
		return -ENOMEM;
	}

	/* populate match struct */
	strncpy(new_match->name, info->name, IPT_NETQUOTA_NAME_LEN);
	new_match->network = info->network;
	new_match->netmask = info->netmask;
	new_match->use_count = 1;

	/* create proc entry */
	new_match->proc_entry =
	    create_proc_entry(new_match->name, proc_file_perms,
			      proc_net_ipt_netquota);
	if (new_match->proc_entry == NULL) {
		write_unlock(&nq_match_arr_lock);
		vfree(new_match->stat);
		vfree(new_match);
		return -ENOMEM;
	}

	/* populate proc struct */
	new_match->proc_entry->owner = THIS_MODULE;
	new_match->proc_entry->read_proc = nq_read_proc;
	new_match->proc_entry->write_proc = nq_write_proc;
	new_match->proc_entry->data = new_match;

	/* init quotas & spinlocks */
	for (i = 0; i < host_count; i++) {
		new_match->stat[i].quota = 0;
		new_match->stat[i].lock = SPIN_LOCK_UNLOCKED;
	}

	/* finally insert new netquota into match array */
	nq_match_arr[unused] = new_match;
	info->match_index = unused;

#if ((DEBUG & 1) == 1)
	printk(KERN_INFO "ipt_netquota: new match '%s' created\n",
	       new_match->name);
#endif
	write_unlock(&nq_match_arr_lock);
	return 1;
}

/* Called when an entry of this type is deleted. */
static void destroy(void *matchinfo, unsigned int matchinfosize)
{
	struct ipt_netquota_info *info = matchinfo;
	unsigned int i, j, host_count;

	if (matchinfosize != IPT_ALIGN(sizeof(struct ipt_netquota_info)))
		return;

	/* find quota name in match array */
	write_lock(&nq_match_arr_lock);
	for (i = 0; i < max_num_of_matches; i++) {

		/* no match / match deleted */
		if (nq_match_arr[i] == NULL) {
			continue;
		}

		/* different name */
		if (strncmp(info->name, nq_match_arr[i]->name,
		     IPT_NETQUOTA_NAME_LEN)) {
			continue;
		}

		/* different parameters (note: only network/netmask check) */
		if (info->network != nq_match_arr[i]->network ||
		    info->netmask != nq_match_arr[i]->netmask) {

			write_unlock(&nq_match_arr_lock);
			return;
		}

		/* ok, must be the same match, decrease reference count */
		nq_match_arr[i]->use_count--;
#if ((DEBUG & 1) == 1)
		printk(KERN_INFO "ipt_netquota: decremented use_count of match "
		       "'%s' to %d\n", nq_match_arr[i]->name,
		       nq_match_arr[i]->use_count);
#endif
		if (nq_match_arr[i]->use_count > 0) {
			write_unlock(&nq_match_arr_lock);
			return;
		}
#if ((DEBUG & 1) == 1)
		printk(KERN_INFO "ipt_netquota: destroy match '%s'\n",
		       nq_match_arr[i]->name);
#endif
		host_count = INADDR_BROADCAST - info->netmask + 1;

		/* first, try to lock all counters
		 * -> all other contexts (if any) in match() must leave (unlock)
		 * -> new contexts cannot enter match() due to global write_lock
		 * -> remove the counter locks
		 * -> we are alone[tm] and can remove this match
		 */
		for (j = 0; j < host_count; j++)
			spin_lock(&nq_match_arr[i]->stat[j].lock);

		for (j = 0; j < host_count; j++)
			spin_unlock(&nq_match_arr[i]->stat[j].lock);

		remove_proc_entry(nq_match_arr[i]->name, proc_net_ipt_netquota);

		vfree(nq_match_arr[i]->stat);
		vfree(nq_match_arr[i]);
		nq_match_arr[i] = NULL;

		write_unlock(&nq_match_arr_lock);
		return;
	}

	printk(KERN_ERR "ipt_netquota: no such quota: '%s'\n",
	       info->name);

	write_unlock(&nq_match_arr_lock);
	return;
}

static struct ipt_match netquota_match = {
	.name = "netquota",
	.match = match,
	.checkentry = checkentry,
	.destroy = destroy,
	.me = THIS_MODULE
};

static int __init init(void)
{
#if ((DEBUG & 1) == 1)
	u_int32_t netmask = ~((1 << (32 - max_netmask)) - 1);
	printk(KERN_INFO "ipt_netquota: module loaded (max. matches: %d  "
	       "max. netmask: %u.%u.%u.%u)\n", max_num_of_matches,
	       HIPQUAD(netmask));
#endif
	/* create match array */
	nq_match_arr = vmalloc(sizeof(struct nq_match) * max_num_of_matches);
	if (nq_match_arr == NULL) {
		return -ENOMEM;
	}

	memset(nq_match_arr, 0, sizeof(struct nq_match) * max_num_of_matches);

	/* create /proc/net/ipt_netquota/ */
	proc_net_ipt_netquota = proc_mkdir("ipt_netquota", proc_net);
	if (proc_net_ipt_netquota == NULL) {
		vfree(nq_match_arr);
		return -ENOMEM;
	}

	if (ipt_register_match(&netquota_match)) {
		vfree(nq_match_arr);
		return -EINVAL;
	}

	return 0;
}

static void __exit fini(void)
{
	ipt_unregister_match(&netquota_match);
	remove_proc_entry("ipt_netquota", proc_net);

	vfree(nq_match_arr);
#if ((DEBUG & 1) == 1)
	printk(KERN_INFO "ipt_netquota: module unloaded\n");
#endif
}

module_init(init);
module_exit(fini);