用Netfilter模块实现基于令牌桶的每IP地址流量控制

/**  *     2011/11/06 by marywangran  *    这个版本实现只能使所有规则共享一个上限，作为全局管理局部特例  *    很方便，不多的几条iptables规则即可  */ #include <linux/module.h> #include <linux/skbuff.h> #include <linux/spinlock.h> #include <linux/interrupt.h> #include <net/ip.h>  #include <linux/netfilter/x_tables.h> #include <linux/types.h>  MODULE_LICENSE("GPL"); MODULE_AUTHOR("marywangran <marywangran@126.com>"); MODULE_DESCRIPTION("Xtables: rate-limit match"); MODULE_ALIAS("ipt_limit"); MODULE_ALIAS("ip6t_limit");  struct xt_rateinfo {         __u32 type;    //定义类型，1为源地址限速，2为目标地址限速         __u32 burst;    //定义最大流量，统计值 };  struct src_controler {         struct list_head src_list;         int curr;             //当前的IP连接数         int max;            //最大的IP连接数         spinlock_t lock; };  struct src_entry {         struct list_head list;         int type;        //同xt_rateinfo结构体的type         __u32   addr;            unsigned long prev_hit; //上次包到来的时间         unsigned long toks;    //当前拥有的令牌数         spinlock_t lock; };  struct src_controler *src_ctl;  static bool limit_mt(const struct sk_buff *skb, const struct xt_match_param *par) {         const struct xt_rateinfo *r = par->matchinfo;         unsigned long now = jiffies;          struct list_head *lh;         struct src_entry *entry = NULL;         struct src_entry *find_entry;         long tokens;         struct iphdr *iph = ip_hdr(skb);         __u32 this_addr = 0;          if (r->type == 1) {                 this_addr = iph->saddr;         } else {                 this_addr = iph->daddr;         }          spin_lock (&src_ctl->lock);  //操作链表一定加锁，多CPU下防止并发修改，访问         list_for_each(lh, &src_ctl->src_list) {                 find_entry = list_entry(lh, struct src_entry, list);                 if ((this_addr == find_entry->addr) &&                     (find_entry->type == r->type)) {                         entry = find_entry;                         break;                 }         }         spin_unlock (&src_ctl->lock);         if (entry) {                 spin_lock (&src_ctl->lock);                 list_del(&entry->list);                 list_add(&entry->list, &src_ctl->src_list);                 spin_unlock (&src_ctl->lock);         } else {                    if (src_ctl->curr+1 < src_ctl->max) { add_entry:                         entry = kmalloc(sizeof(struct src_entry), GFP_ATOMIC); //必须使用ATOMIC标志，因为有可能在(软)中断中运行，不能睡眠。                         memset(entry, 0, sizeof(struct src_entry));                         entry->addr = this_addr;                         entry->toks = r->burst; //第一次分配令牌时不加倍(加倍理由为防止使用浮点运算)，以防TCP的慢启动增加突发流量，TCP的慢启动实际上很快，指数级的。                         entry->prev_hit = now;                         entry->type = r->type;                         spin_lock_init(&entry->lock);                         spin_lock (&src_ctl->lock);                         list_add(&entry->list, &src_ctl->src_list);                         src_ctl->curr++;     //应该使用atomic_inc进行递增                         spin_unlock (&src_ctl->lock);                 } else {                         entry = list_entry(src_ctl->src_list.prev, struct src_entry, list);                         if (now-entry->now > 1000) {                                 spin_lock (&src_ctl->lock);                                 list_del(&entry->list);                                 src_ctl->curr--;                                 spin_unlock (&src_ctl->lock);                                 vfree(entry);   //解锁后vfree                                 goto add_entry;                         }                         return 1;                 }         }     //以下根据流逝的时间来确定令牌的数量         tokens = min_t(long, (now-entry->prev_hit)*r->burst, r->burst*1000);         tokens += entry->toks;         if (tokens > (long)r->burst*1000)                         tokens = r->burst*1000;                 tokens -= skb->len*1000; //统一增加HZ倍，避免在内核使用浮点数和除法。          if (tokens >= 0) {                 spin_lock (&entry->lock);                 entry->prev_hit = now;                 entry->toks = tokens;    //令牌积累                 spin_unlock (&entry->lock);                 return 0;         }         return 1; }  static bool limit_mt_check(const struct xt_mtchk_param *par) {         struct xt_rateinfo *r = par->matchinfo;         if (r->burst == 0 || r->type == 0) {                 return false;         }         if (r->type != 1 && r->type != 2)                 return false;         return true; }  static void limit_mt_destroy(const struct xt_mtdtor_param *par) {     //TODO }   static struct xt_match limit_mt_reg __read_mostly = {         .name             = "limit",         .revision         = 0,         .family           = NFPROTO_UNSPEC,         .match            = limit_mt,         .checkentry       = limit_mt_check,         .destroy          = limit_mt_destroy,         .matchsize        = sizeof(struct xt_rateinfo),         .me               = THIS_MODULE, };  static int __init limit_mt_init(void) {         src_ctl = kmalloc(sizeof(struct src_controler), GFP_KERNEL); //初始化全局变量，insmod上下文，可以使用KERNEL标志         memset(src_ctl, 0, sizeof(struct src_controler));         INIT_LIST_HEAD(&src_ctl->src_list);    //初始化全局变量的链表         src_ctl->curr = 0;         src_ctl->max = 2000;         spin_lock_init(&src_ctl->lock);         return xt_register_match(&limit_mt_reg); }  static void __exit limit_mt_exit(void) {         struct src_entry *entry = NULL;         struct list_head *lh = NULL, *lh2 = NULL;         xt_unregister_match(&limit_mt_reg);         spin_lock(&src_ctl->lock);         list_for_each_safe(lh, lh2, &src_ctl->src_list) { //一定要用safe宏，因为这是个外部迭代器                 entry = list_entry(lh, struct src_entry, list);                 list_del(&entry->list);                 kfree(entry);         }         spin_unlock(&src_ctl->lock);         kfree(src_ctl);  }  module_init(limit_mt_init); module_exit(limit_mt_exit);

LINUXPATH = /lib/modules/`uname -r`/build CURDIR = $(shell pwd) KBUILD_OUTPUT = $(CURDIR) CROSS_COMPILE = ARCH =  obj-m                           += xt_limit.o  all: limit  limit:         $(MAKE) -C $(LINUXPATH) M=$(CURDIR) modules         @echo "*********************************************"         @echo "*  The MODULE is OK!!"         @echo "*********************************************" .PHONY: clean clean:         rm -rf *.o *.ko *.mod.c *.symvers *.mod.o .*.cmd  ../common/*.o .tmp_versions 

/**  *    2011/11/06 by marywangran  *    修改自iptables-1.4.12/extensions/libxt_limit.c  */  #include <stdio.h> #include <string.h> #include <stdlib.h> #include <xtables.h> #include <linux/netfilter/x_tables.h> #include <linux/netfilter/xt_limit.h>  #define XT_LIMIT_BURST  500000  /**  *    新增轻量级rateinfo结构，对应于内核的等价结构  */ struct xt_rateinfo_new {         __u32 type;         __u32 burst; };  enum {         O_TYPE = 0,         O_BURST, };  static void limit_help(void) {         printf( "limit match options:\n" "--type source[1]|destination[2]        define source limit or destination limit\n" "--limit-burst rate                     rate to match in a burst, default %u\n", XT_LIMIT_BURST); }  static const struct xt_option_entry limit_opts[] = {         {.name = "type", .id = O_TYPE, .type = XTTYPE_STRING},         {.name = "limit-burst", .id = O_BURST, .type = XTTYPE_UINT32,          .flags = XTOPT_PUT, XTOPT_POINTER(struct xt_rateinfo_new, burst),          .min = 0, .max = 1073741824}, //1024*1024*1024         XTOPT_TABLEEND, };  static int parse_rate(const char *rate, uint32_t *val) {         uint32_t r;         r = atoi(rate);         if (!r)                 return 0;         if (r != 1 && r != 2) //1为限制源地址，2为限制目的地址                 return 0;         *val = r;         return 1; }  static void limit_init(struct xt_entry_match *m) {         struct xt_rateinfo_new *r = (struct xt_rateinfo_new *)m->data;         parse_rate("1", &r->type);         r->burst = XT_LIMIT_BURST; }  static void limit_parse(struct xt_option_call *cb) {         struct xt_rateinfo_new *r = cb->data;         xtables_option_parse(cb);         switch (cb->entry->id) {         case O_TYPE:                 if (!parse_rate(cb->arg, &r->type))                         xtables_error(PARAMETER_PROBLEM,                                    "bad rate \"%s\"'", cb->arg);                 break;         }         if (cb->invert)                 xtables_error(PARAMETER_PROBLEM,                            "limit does not support invert"); }  static void print_rate(uint32_t period) {         printf(" %u", period); }  static void limit_print(const void *ip, const struct xt_entry_match *match, int numeric) {         const struct xt_rateinfo_new *r = (const void *)match->data;         printf(" type: avg"); print_rate(r->type);         printf(" burst %u", r->burst); }  static void limit_save(const void *ip, const struct xt_entry_match *match) {         const struct xt_rateinfo_new *r = (const void *)match->data;          printf(" --type"); print_rate(r->type);         if (r->burst != XT_LIMIT_BURST)                 printf(" --limit-burst %u", r->burst); }  static struct xtables_match limit_match = {         .family         = NFPROTO_UNSPEC,         .name           = "limit",         .version        = XTABLES_VERSION,         .size           = XT_ALIGN(sizeof(struct xt_rateinfo_new)),         .help           = limit_help,         .init           = limit_init,         .x6_parse       = limit_parse,         .print          = limit_print,         .save           = limit_save,         .x6_options     = limit_opts, };  void _init(void) {         xtables_register_match(&limit_match); }

安装：

说明：以上的代码都是基于现有的Linux内核以及iptables源码修改的，框架虽在，然逻辑已全非，如果需保留原来的limit功能，请自行实现新的文件。

 本文转自 dog250 51CTO博客，原文链接:http://blog.51cto.com/dog250/1270919