Nginx中的頻控模塊示例

簡介

陶輝老師《深入理解Nginx》中的示例代碼辟癌,支持IP+URL級別的頻控氏身。

頻控以模塊的方式嵌入Nginx巍棱。采用 紅黑樹+鏈表 的方式實(shí)現(xiàn),每當(dāng)一個(gè)IP訪問一次URL蛋欣,紅黑樹將會插入一個(gè)節(jié)點(diǎn)航徙,節(jié)點(diǎn)包含本次訪問時(shí)間。

當(dāng)相同的IP短時(shí)間內(nèi)訪問同樣的URL時(shí)陷虎,紅黑樹就會查找到剛插入的節(jié)點(diǎn)到踏,找出上次的訪問時(shí)間,判斷間隔是否夠長尚猿,間隔太短的會返回 403 Forbidden窝稿,間隔夠長就允許訪問,并把這次訪問時(shí)間更新到節(jié)點(diǎn)中凿掂。

鏈表的作用又是什么伴榔?許多情況下客戶端訪問了某個(gè)URL后就再也不會訪問了,這些訪問生成的紅黑樹節(jié)點(diǎn)庄萎,要及時(shí)清理掉避免紅黑樹過大踪少。于是鏈表就將紅黑樹的節(jié)點(diǎn)按記錄的訪問時(shí)間有序串起來。每當(dāng)有新的請求到來時(shí)糠涛,順便會檢查鏈表中最久遠(yuǎn)的幾個(gè)節(jié)點(diǎn)秉馏,若節(jié)點(diǎn)記錄的訪問時(shí)間與現(xiàn)在太遙遠(yuǎn),就可以清理掉了脱羡。

配置方法

http 塊中配置萝究,第一個(gè)參數(shù)是 IP+URL 連續(xù)訪問的最短間隔,單位是秒锉罐。第二個(gè)參數(shù)是分配給紅黑樹+鏈表的字節(jié)數(shù)帆竹。

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

編譯方法

頻控模塊的源碼有兩個(gè)文件:config 和 ngx_http_testslab_module.c,放在一個(gè)目錄中脓规。編譯 nginx 的時(shí)候栽连,在 configure 階段使用 --add-module 把模塊添加進(jìn)去:

./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配置項(xiàng)
        // 必須攜帶2個(gè)參數(shù), 前者為兩次成功訪問同一URL時(shí)的最小間隔秒數(shù)
        // 后者為共享內(nèi)存的大小
        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);
    // 設(shè)置模塊在NGX_HTTP_PREACCESS_PHASE階段介入請求的處理
    h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
    if (h == NULL) {
        return NGX_ERROR;
    }
    // 設(shè)置請求的處理方法
    *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參數(shù)錯(cuò)誤, 返回NGX_DECLINED繼續(xù)執(zhí)行下一個(gè)http handler
    if (conf->interval == -1)
        return rc;
    // 以客戶端IP地址(r->connection->addr_text中已經(jīng)保存了解析出的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碼作為紅黑樹的關(guān)鍵字來提高效率
    hash = ngx_crc32_short(data, len);
    // 多進(jìn)程同時(shí)操作同一共享內(nèi)存, 需要加鎖
    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內(nèi)存中分配配置結(jié)構(gòu)體
    conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_testslab_conf_t));
    if (conf == NULL) {
        return NULL;
    }
    // interval初始化為-1, 同時(shí)用于判斷是否未開啟模塊的限速功能
    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執(zhí)行完成后的shm_zone->data
    ngx_http_testslab_conf_t  *oconf = data;
    size_t                     len;
    // shm_zone->data存放著本次初始化cycle時(shí)創(chuàng)建的ngx_http_testslab_conf_t配置結(jié)構(gòu)體
    conf = (ngx_http_testslab_conf_t  *)shm_zone->data;
    // 判斷是否為reload配置項(xiàng)后導(dǎo)致的初始化共享內(nèi)存
    if (oconf) {
        // 本次初始化的共享內(nèi)存不是新創(chuàng)建的
        // 此時(shí), data成員里就是上次創(chuàng)建的ngx_http_testslab_conf_t
        // 將sh和shpool指針指向舊的共享內(nèi)存即可
        conf->sh = oconf->sh;
        conf->shpool = oconf->shpool;
        return NGX_OK;
    }
    // shm.addr里放著共享內(nèi)存首地址:ngx_slab_pool_t結(jié)構(gòu)體
    conf->shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
    // slab共享內(nèi)存中每一次分配的內(nèi)存都用于存放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);
    // 初始化按訪問時(shí)間排序的鏈表
    ngx_queue_init(&conf->sh->queue);
    // slab操作共享內(nèi)存出現(xiàn)錯(cuò)誤時(shí), 其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);
    }
}
最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
  • 序言:七十年代末,一起剝皮案震驚了整個(gè)濱河市侨舆,隨后出現(xiàn)的幾起案子秒紧,更是在濱河造成了極大的恐慌,老刑警劉巖挨下,帶你破解...
    沈念sama閱讀 217,509評論 6 504
  • 序言:濱河連續(xù)發(fā)生了三起死亡事件熔恢,死亡現(xiàn)場離奇詭異,居然都是意外死亡臭笆,警方通過查閱死者的電腦和手機(jī)叙淌,發(fā)現(xiàn)死者居然都...
    沈念sama閱讀 92,806評論 3 394
  • 文/潘曉璐 我一進(jìn)店門,熙熙樓的掌柜王于貴愁眉苦臉地迎上來愁铺,“玉大人鹰霍,你說我怎么就攤上這事∫鹇遥” “怎么了茂洒?”我有些...
    開封第一講書人閱讀 163,875評論 0 354
  • 文/不壞的土叔 我叫張陵,是天一觀的道長瓶竭。 經(jīng)常有香客問我督勺,道長,這世上最難降的妖魔是什么在验? 我笑而不...
    開封第一講書人閱讀 58,441評論 1 293
  • 正文 為了忘掉前任玷氏,我火速辦了婚禮,結(jié)果婚禮上腋舌,老公的妹妹穿的比我還像新娘盏触。我一直安慰自己,他們只是感情好块饺,可當(dāng)我...
    茶點(diǎn)故事閱讀 67,488評論 6 392
  • 文/花漫 我一把揭開白布赞辩。 她就那樣靜靜地躺著,像睡著了一般授艰。 火紅的嫁衣襯著肌膚如雪辨嗽。 梳的紋絲不亂的頭發(fā)上,一...
    開封第一講書人閱讀 51,365評論 1 302
  • 那天淮腾,我揣著相機(jī)與錄音糟需,去河邊找鬼屉佳。 笑死,一個(gè)胖子當(dāng)著我的面吹牛洲押,可吹牛的內(nèi)容都是我干的武花。 我是一名探鬼主播,決...
    沈念sama閱讀 40,190評論 3 418
  • 文/蒼蘭香墨 我猛地睜開眼杈帐,長吁一口氣:“原來是場噩夢啊……” “哼体箕!你這毒婦竟也來了?” 一聲冷哼從身側(cè)響起挑童,我...
    開封第一講書人閱讀 39,062評論 0 276
  • 序言:老撾萬榮一對情侶失蹤累铅,失蹤者是張志新(化名)和其女友劉穎,沒想到半個(gè)月后站叼,有當(dāng)?shù)厝嗽跇淞掷锇l(fā)現(xiàn)了一具尸體娃兽,經(jīng)...
    沈念sama閱讀 45,500評論 1 314
  • 正文 獨(dú)居荒郊野嶺守林人離奇死亡,尸身上長有42處帶血的膿包…… 初始之章·張勛 以下內(nèi)容為張勛視角 年9月15日...
    茶點(diǎn)故事閱讀 37,706評論 3 335
  • 正文 我和宋清朗相戀三年大年,在試婚紗的時(shí)候發(fā)現(xiàn)自己被綠了换薄。 大學(xué)時(shí)的朋友給我發(fā)了我未婚夫和他白月光在一起吃飯的照片。...
    茶點(diǎn)故事閱讀 39,834評論 1 347
  • 序言:一個(gè)原本活蹦亂跳的男人離奇死亡翔试,死狀恐怖轻要,靈堂內(nèi)的尸體忽然破棺而出,到底是詐尸還是另有隱情垦缅,我是刑警寧澤冲泥,帶...
    沈念sama閱讀 35,559評論 5 345
  • 正文 年R本政府宣布,位于F島的核電站壁涎,受9級特大地震影響凡恍,放射性物質(zhì)發(fā)生泄漏。R本人自食惡果不足惜怔球,卻給世界環(huán)境...
    茶點(diǎn)故事閱讀 41,167評論 3 328
  • 文/蒙蒙 一嚼酝、第九天 我趴在偏房一處隱蔽的房頂上張望。 院中可真熱鬧竟坛,春花似錦闽巩、人聲如沸。這莊子的主人今日做“春日...
    開封第一講書人閱讀 31,779評論 0 22
  • 文/蒼蘭香墨 我抬頭看了看天上的太陽。三九已至崭歧,卻和暖如春隅很,著一層夾襖步出監(jiān)牢的瞬間,已是汗流浹背率碾。 一陣腳步聲響...
    開封第一講書人閱讀 32,912評論 1 269
  • 我被黑心中介騙來泰國打工叔营, 沒想到剛下飛機(jī)就差點(diǎn)兒被人妖公主榨干…… 1. 我叫王不留屋彪,地道東北人。 一個(gè)月前我還...
    沈念sama閱讀 47,958評論 2 370
  • 正文 我出身青樓绒尊,卻偏偏與公主長得像撼班,于是被迫代替她去往敵國和親。 傳聞我的和親對象是個(gè)殘疾皇子垒酬,可洞房花燭夜當(dāng)晚...
    茶點(diǎn)故事閱讀 44,779評論 2 354

推薦閱讀更多精彩內(nèi)容