Nginx 의 주파수 제어 모듈 예제

간단 한 소개
도 휘 선생님 의 에서 예시 코드 를 깊이 이해 하고 IP + URL 등급 의 주파수 제 어 를 지원 합 니 다.
주파수 제 어 는 모듈 로 Nginx 를 삽입 합 니 다.빨간색 과 검은색 트 리 + 링크 방식 으로 이 루어 집 니 다. IP 가 URL 에 한 번 방문 할 때마다 빨간색 과 검은색 트 리 는 노드 를 삽입 하고 노드 는 이번 방문 시간 을 포함 합 니 다.
같은 IP 가 짧 은 시간 에 같은 URL 에 접근 할 때 빨간색 과 검은색 트 리 는 방금 삽 입 된 노드 를 찾 아 마지막 방문 시간 을 찾 아 간격 이 얼마나 긴 지 판단 하고 간격 이 너무 짧 은 것 은 403 Forbidden 으로 돌아 가 며 간격 이 길 면 접근 을 허용 하고 이번 방문 시간 을 노드 에 업데이트 합 니 다.
링크 의 역할 은 또 무엇 입 니까?많은 상황 에서 클 라 이언 트 가 특정한 URL 을 방문 한 후에 다 시 는 방문 하지 않 을 것 이다. 이런 방문 으로 생 성 된 빨 간 검 은 나무 노드 는 빨 간 검 은 나무 가 너무 크 지 않도록 신속하게 제거 해 야 한다.그래서 링크 는 빨 간 검 은 나무의 노드 를 기 록 된 방문 시간 에 따라 질서 있 게 연결 합 니 다.새로운 요청 이 올 때마다 링크 에서 가장 오래된 몇 개의 노드 를 검사 합 니 다. 노드 에 기 록 된 방문 시간 이 지금 과 너무 멀 면 정리 할 수 있 습 니 다.
설정 방법http 블록 에 첫 번 째 인 자 는 IP + URL 연속 접근 의 가장 짧 은 간격 이 고 단 위 는 초 입 니 다.두 번 째 매개 변 수 는 빨 간 검 은 나무 + 링크 에 분 배 된 바이트 수 입 니 다.

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

컴 파일 방법
주파수 제어 모듈 의 원본 코드 는 두 개의 파일 이 있 습 니 다: config 와 ngxhttp_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);
    }
}