Nginx中的频控模块示例

简介

陶辉老师《深入理解Nginx》中的示例代码，支持IP+URL级别的频控。

频控以模块的方式嵌入Nginx。采用 红黑树+链表 的方式实现，每当一个IP访问一次URL，红黑树将会插入一个节点，节点包含本次访问时间。

当相同的IP短时间内访问同样的URL时，红黑树就会查找到刚插入的节点，找出上次的访问时间，判断间隔是否够长，间隔太短的会返回 403 Forbidden，间隔够长就允许访问，并把这次访问时间更新到节点中。

链表的作用又是什么？许多情况下客户端访问了某个URL后就再也不会访问了，这些访问生成的红黑树节点，要及时清理掉避免红黑树过大。于是链表就将红黑树的节点按记录的访问时间有序串起来。每当有新的请求到来时，顺便会检查链表中最久远的几个节点，若节点记录的访问时间与现在太遥远，就可以清理掉了。

配置方法

在 http 块中配置，第一个参数是 IP+URL 连续访问的最短间隔，单位是秒。第二个参数是分配给红黑树+链表的字节数。

http {
    ...
    test_slab 10 32768;
    ...
}

编译方法

频控模块的源码有两个文件：config 和 ngx_http_testslab_module.c，放在一个目录中。编译 nginx 的时候，在 configure 阶段使用 --add-module 把模块添加进去：

./configure --add-module=<源码目录的绝对路径>

然后 make & make install就行了

config

NGX_ADDON_SRCS="$NGX_ADDON_SRCS $ngx_addon_dir/ngx_http_testslab_module.c"

ngx_http_testslab_module.c

#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>

typedef struct {
   
    u_char rbtree_node_data;
    ngx_queue_t queue;
    ngx_msec_t last;
    u_short len;
    u_char data[1];
} ngx_http_testslab_node_t;

typedef struct {
   
    ngx_rbtree_t rbtree;
    ngx_rbtree_node_t sentinel;
    ngx_queue_t queue;
} ngx_http_testslab_shm_t;

typedef struct {
   
    ssize_t shmsize;
    ngx_int_t interval;
    ngx_slab_pool_t *shpool;
    ngx_http_testslab_shm_t* sh;
} ngx_http_testslab_conf_t;

static ngx_int_t ngx_http_testslab_init(ngx_conf_t*);
static void *ngx_http_testslab_create_main_conf(ngx_conf_t*);
static char *ngx_http_testslab_createmem(ngx_conf_t*, ngx_command_t*, void*);
static ngx_int_t ngx_http_testslab_handler(ngx_http_request_t*);
static ngx_int_t ngx_http_testslab_lookup(ngx_http_request_t*, ngx_http_testslab_conf_t*, ngx_uint_t, u_char*, size_t);
static ngx_int_t ngx_http_testslab_shm_init(ngx_shm_zone_t*, void*);
static void ngx_http_testslab_rbtree_insert_value(ngx_rbtree_node_t*, ngx_rbtree_node_t*, ngx_rbtree_node_t*);
static void ngx_http_testslab_expire(ngx_http_request_t*, ngx_http_testslab_conf_t*);

static ngx_command_t  ngx_http_testslab_commands[] = {
   
    {
   
        ngx_string("test_slab"),
        // 仅支持在http块下配置test_slab配置项
        // 必须携带2个参数, 前者为两次成功访问同一URL时的最小间隔秒数
        // 后者为共享内存的大小
        NGX_HTTP_MAIN_CONF|NGX_CONF_TAKE2,
        ngx_http_testslab_createmem,
        0,
        0,
        NULL
    },
    ngx_null_command
};

static ngx_http_module_t  ngx_http_testslab_module_ctx =
{
   
    NULL,                               /* preconfiguration */
    ngx_http_testslab_init,             /* postconfiguration */
    ngx_http_testslab_create_main_conf, /* create main configuration */
    NULL,                               /* init main configuration */
    NULL,                               /* create server configuration */
    NULL,                               /* merge server configuration */
    NULL,                               /* create location configuration */
    NULL                                /* merge location configuration */
};

ngx_module_t  ngx_http_testslab_module =
{
   
    NGX_MODULE_V1,
    &ngx_http_testslab_module_ctx,         /* module context */
    ngx_http_testslab_commands,            /* module directives */
    NGX_HTTP_MODULE,                       /* module type */
    NULL,                                  /* init master */
    NULL,                                  /* init module */
    NULL,                                  /* init process */
    NULL,                                  /* init thread */
    NULL,                                  /* exit thread */