Nginx 基於slab內存實現的客戶端訪問限制（IP+URL）

本實例來源於《深入理解Nginx》書中的源碼。網上未查到可複製使用的源碼，因此親自編譯測試並寫此博客，定爲轉載類型，供nginx愛好者學習交流。

 

測試結果：

10秒內第一次訪問結果：HelloWorld

10秒內第2、3..次訪問結果：403 Forbidden

 

源碼：

config:

http {

test_slab 10 10M; #限制10s一次的訪問頻率，slab內存大小爲10M

server {

          location /hello {
            default_type text/html;
            content_by_lua_block {
                ngx.say("HelloWorld")
            }
        }
}

}

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_handler(ngx_http_request_t *r);

static char * ngx_http_testslab_createmem(ngx_conf_t *cf, ngx_conf_t *cmd, void *conf);

static ngx_int_t ngx_http_testslab_init(ngx_conf_t *cf);

static void *ngx_http_testslab_create_main_conf(ngx_conf_t *cf);

static ngx_command_t  ngx_http_testslab_commands[] =
{

    {
        ngx_string("test_slab"),
        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 void *ngx_http_testslab_create_main_conf(ngx_conf_t *cf)
{
    ngx_http_testslab_conf_t *conf;

    conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_testslab_conf_t));

    if (NULL == conf)
    {
        return NULL;
    }

    conf->interval = -1;
    conf->shmsize = -1;

    return conf;
}

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 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);
   }

}

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 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);

    h = ngx_array_push(&cmcf->phases[NGX_HTTP_PREACCESS_PHASE].handlers);
    if (NULL == h)
    {
        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;

    if (conf->interval == -1)
    {
        return rc;
    }

    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);

    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 ngx_int_t
ngx_http_testslab_shm_init(ngx_shm_zone_t *shm_zone, void *data)
{
    ngx_http_testslab_conf_t *conf;

    ngx_http_testslab_conf_t *oconf = data;
    size_t len;

    conf = (ngx_http_testslab_conf_t *)shm_zone->data;

    if (oconf)
    {
        conf->sh = oconf->sh;
        conf->shpool = oconf->shpool;

        return NGX_OK;
    }

    conf->shpool = (ngx_slab_pool_t *)shm_zone->shm.addr;

    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);

    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 char * ngx_http_testslab_createmem(ngx_conf_t *cf, ngx_conf_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 (NULL == shm_zone)
    {
        mconf->interval = -1;
        return NGX_CONF_ERROR;
    }

    shm_zone->init = ngx_http_testslab_shm_init;

    shm_zone->data = mconf;

    return NGX_CONF_OK;
}