简介
陶辉老师《深入理解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
#include
#include
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 */
NULL, /* exit process */
NULL, /* exit master */
NGX_MODULE_V1_PADDING
};
static ngx_int_t
ngx_http_testslab_init(ngx_conf_t *cf)
{
ngx_http_handler_pt *h;
ngx_http_core_main_conf_t *cmcf;
cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);
// 设置模块在NGX_HTTP_PREACCESS_PHASE阶段介入请求的处理
h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
if (h == NULL) {
return NGX_ERROR;
}
// 设置请求的处理方法
*h = ngx_http_testslab_handler;
return NGX_OK;
}
static ngx_int_t
ngx_http_testslab_handler(ngx_http_request_t *r)
{
size_t len;
uint32_t hash;
ngx_int_t rc;
ngx_http_testslab_conf_t *conf;
conf = ngx_http_get_module_main_conf(r, ngx_http_testslab_module);
rc = NGX_DECLINED;
// 如果没有配置test_slab, 或者test_slab参数错误, 返回NGX_DECLINED继续执行下一个http handler
if (conf->interval == -1)
return rc;
// 以客户端IP地址(r->connection->addr_text中已经保存了解析出的IP字符串)
// 和url来识别同一请求
len = r->connection->addr_text.len + r->uri.len;
u_char* data = ngx_palloc(r->pool, len);
ngx_memcpy(data, r->uri.data, r->uri.len);
ngx_memcpy(data+r->uri.len, r->connection->addr_text.data, r->connection->addr_text.len);
// 使用crc32算法将IP+URL字符串生成hash码
// hash码作为红黑树的关键字来提高效率
hash = ngx_crc32_short(data, len);
// 多进程同时操作同一共享内存, 需要加锁
ngx_shmtx_lock(&conf->shpool->mutex);
rc = ngx_http_testslab_lookup(r, conf, hash, data, len);
ngx_shmtx_unlock(&conf->shpool->mutex);
return rc;
}
static char *
ngx_http_testslab_createmem(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
{
ngx_str_t *value;
ngx_shm_zone_t *shm_zone;
ngx_http_testslab_conf_t *mconf = (ngx_http_testslab_conf_t *)conf;
ngx_str_t name = ngx_string("test_slab_shm");
value = cf->args->elts;
mconf->interval = 1000 * ngx_atoi(value[1].data, value[1].len);
if (mconf->interval == NGX_ERROR || mconf->interval == 0) {
mconf->interval = -1;
return "invalid value";
}
mconf->shmsize = ngx_parse_size(&value[2]);
if (mconf->shmsize == (ssize_t) NGX_ERROR || mconf->shmsize == 0) {
mconf->interval = -1;
return "invalid value";
}
shm_zone = ngx_shared_memory_add(cf, &name, mconf->shmsize,
&ngx_http_testslab_module);
if (shm_zone == NULL) {
mconf->interval = -1;
return NGX_CONF_ERROR;
}
shm_zone->init = ngx_http_testslab_shm_init;
shm_zone->data = mconf;
return NGX_CONF_OK;
}
static void *
ngx_http_testslab_create_main_conf(ngx_conf_t *cf)
{
ngx_http_testslab_conf_t *conf;
// 在worker内存中分配配置结构体
conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_testslab_conf_t));
if (conf == NULL) {
return NULL;
}
// interval初始化为-1, 同时用于判断是否未开启模块的限速功能
conf->interval = -1;
conf->shmsize = -1;
return conf;
}
static ngx_int_t
ngx_http_testslab_shm_init(ngx_shm_zone_t *shm_zone, void *data) {
ngx_http_testslab_conf_t *conf;
// data可能为空, 也可能是上次ngx_http_testslab_shm_init执行完成后的shm_zone->data
ngx_http_testslab_conf_t *oconf = data;
size_t len;
// shm_zone->data存放着本次初始化cycle时创建的ngx_http_testslab_conf_t配置结构体
conf = (ngx_http_testslab_conf_t *)shm_zone->data;
// 判断是否为reload配置项后导致的初始化共享内存
if (oconf) {
// 本次初始化的共享内存不是新创建的
// 此时, data成员里就是上次创建的ngx_http_testslab_conf_t
// 将sh和shpool指针指向旧的共享内存即可
conf->sh = oconf->sh;
conf->shpool = oconf->shpool;
return NGX_OK;
}
// shm.addr里放着共享内存首地址:ngx_slab_pool_t结构体
conf->shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
// slab共享内存中每一次分配的内存都用于存放ngx_http_testslab_shm_t
conf->sh = ngx_slab_alloc(conf->shpool, sizeof(ngx_http_testslab_shm_t));
if (conf->sh == NULL) {
return NGX_ERROR;
}
conf->shpool->data = conf->sh;
// 初始化红黑树
ngx_rbtree_init(&conf->sh->rbtree, &conf->sh->sentinel,
ngx_http_testslab_rbtree_insert_value);
// 初始化按访问时间排序的链表
ngx_queue_init(&conf->sh->queue);
// slab操作共享内存出现错误时, 其log输出会将log_ctx字符串作为后缀, 以方便识别
len = sizeof(" in testslab \"\"") + shm_zone->shm.name.len;
conf->shpool->log_ctx = ngx_slab_alloc(conf->shpool, len);
if (conf->shpool->log_ctx == NULL) {
return NGX_ERROR;
}
ngx_sprintf(conf->shpool->log_ctx, " in testslab \"%V\"%Z",
&shm_zone->shm.name);
return NGX_OK;
}
static void
ngx_http_testslab_rbtree_insert_value(ngx_rbtree_node_t *temp, ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
{
ngx_rbtree_node_t **p;
ngx_http_testslab_node_t *lrn, *lrnt;
for ( ;; ) {
if (node->key < temp->key) {
p = &temp->left;
} else if (node->key > temp->key) {
p = &temp->right;
} else {
lrn = (ngx_http_testslab_node_t *) &node->data;
lrnt = (ngx_http_testslab_node_t *) &temp->data;
p = (ngx_memn2cmp(lrn->data, lrnt->data, lrn->len, lrnt->len) < 0) ? &temp->left : &temp->right;
}
if (*p == sentinel) {
break;
}
temp = *p;
}
*p = node;
node->parent = temp;
node->left = sentinel;
node->right = sentinel;
ngx_rbt_red(node);
}
static ngx_int_t
ngx_http_testslab_lookup(
ngx_http_request_t *r, ngx_http_testslab_conf_t *conf, ngx_uint_t hash, u_char* data, size_t len)
{
size_t size;
ngx_int_t rc;
ngx_time_t *tp;
ngx_msec_t now;
ngx_msec_int_t ms;
ngx_rbtree_node_t *node, *sentinel;
ngx_http_testslab_node_t *lr;
tp = ngx_timeofday();
now = (ngx_msec_t) (tp->sec * 1000 + tp->msec);
node = conf->sh->rbtree.root;
sentinel = conf->sh->rbtree.sentinel;
while (node != sentinel) {
if (hash < node->key) {
node = node->left;
continue;
}
if (hash > node->key) {
node = node->right;
continue;
}
lr = (ngx_http_testslab_node_t *) &node->data;
rc = ngx_memn2cmp(data, lr->data, len, (size_t) lr->len);
if (rc == 0) {
ms = (ngx_msec_int_t) (now - lr->last);
if (ms > conf->interval) {
lr->last = now;
ngx_queue_remove(&lr->queue);
ngx_queue_insert_head(&conf->sh->queue, &lr->queue);
return NGX_DECLINED;
} else {
return NGX_HTTP_FORBIDDEN;
}
}
node = (rc < 0) ? node->left : node->right;
}
size = offsetof(ngx_rbtree_node_t, data) + offsetof(ngx_http_testslab_node_t, data) + len;
ngx_http_testslab_expire(r, conf);
node = ngx_slab_alloc_locked(conf->shpool, size);
if (node == NULL) {
return NGX_ERROR;
}
node->key = hash;
lr = (ngx_http_testslab_node_t *) &node->data;
lr->last = now;
lr->len = (u_char) len;
ngx_memcpy(lr->data, data, len);
ngx_rbtree_insert(&conf->sh->rbtree, node);
ngx_queue_insert_head(&conf->sh->queue, &lr->queue);
return NGX_DECLINED;
}
static void
ngx_http_testslab_expire(ngx_http_request_t *r,ngx_http_testslab_conf_t *conf)
{
ngx_time_t *tp;
ngx_msec_t now;
ngx_queue_t *q;
ngx_msec_int_t ms;
ngx_rbtree_node_t *node;
ngx_http_testslab_node_t *lr;
tp = ngx_timeofday();
now = (ngx_msec_t) (tp->sec * 1000 + tp->msec);
while (1) {
if (ngx_queue_empty(&conf->sh->queue)) {
return;
}
q = ngx_queue_last(&conf->sh->queue);
lr = ngx_queue_data(q, ngx_http_testslab_node_t, queue);
node = (ngx_rbtree_node_t *) ((u_char *) lr - offsetof(ngx_rbtree_node_t, data));
ms = (ngx_msec_int_t) (now - lr->last);
if (ms < conf->interval) {
return;
}
ngx_queue_remove(q);
ngx_rbtree_delete(&conf->sh->rbtree, node);
ngx_slab_free_locked(conf->shpool, node);
}
}