Exemplo n.º 1
0
static ngx_int_t
ngx_http_limit_speed_body_filter(ngx_http_request_t *r, ngx_chain_t *in)
{
    ngx_uint_t                       speed;
    ngx_http_limit_speed_req_ctx_t  *rctx;

    if (ngx_http_limit_speed_get_ctx(r) != NGX_OK) {
        goto done;
    }

    rctx = ngx_http_get_module_ctx(r, ngx_http_limit_speed_module);

    if (rctx == NULL || rctx->ls == NULL || rctx->ls->conn == 0) {
        goto done;
    }

    ngx_log_error(NGX_LOG_DEBUG, r->connection->log, 0,
                  "limit speed rate: %ui, conn: %ui, rctx: %p %p",
                  r->main->limit_rate, rctx->ls->conn, rctx, r);

    speed = rctx->speed;

    r->main->limit_rate = speed / rctx->ls->conn;

done:
    return ngx_http_next_body_filter(r, in);
}
Exemplo n.º 2
0
static ngx_int_t
ngx_http_limit_speed_handler(ngx_http_request_t *r)
{
    size_t                           len, n;
    uint32_t                         hash;
    ngx_int_t                        rc;
    ngx_slab_pool_t                 *shpool;
    ngx_rbtree_node_t               *node, *sentinel;
    ngx_pool_cleanup_t              *cln;
    ngx_http_variable_value_t       *vv;
    ngx_http_limit_speed_ctx_t      *ctx;
    ngx_http_limit_speed_node_t     *ls;
    ngx_http_limit_speed_conf_t     *lscf;
    ngx_http_limit_speed_cleanup_t  *lscln;
    ngx_http_limit_speed_req_ctx_t  *rctx;

    if (r->main->limit_rate) {
        return NGX_DECLINED;
    }

    lscf = ngx_http_get_module_loc_conf(r, ngx_http_limit_speed_module);

    if (lscf->shm_zone == NULL) {
        return NGX_DECLINED;
    }

    ctx = lscf->shm_zone->data;

    vv = ngx_http_get_indexed_variable(r, ctx->index);

    if (vv == NULL || vv->not_found) {
        return NGX_DECLINED;
    }

    len = vv->len;

    if (len == 0) {
        return NGX_DECLINED;
    }

    if (lscf->var_max_len) {
        len = len > lscf->var_max_len ? lscf->var_max_len : len;
    }

    hash = ngx_crc32_short(vv->data, len);

    cln = ngx_pool_cleanup_add(r->pool, sizeof(ngx_http_limit_speed_cleanup_t));
    if (cln == NULL) {
        return NGX_HTTP_INTERNAL_SERVER_ERROR;
    }

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

    ngx_shmtx_lock(&shpool->mutex);

    node = ctx->rbtree->root;
    sentinel = ctx->rbtree->sentinel;

    while (node != sentinel) {

        if (hash < node->key) {
            node = node->left;
            continue;
        }

        if (hash > node->key) {
            node = node->right;
            continue;
        }

        /* hash == node->key */

        do {
            ls = (ngx_http_limit_speed_node_t *) &node->color;

            rc = ngx_memn2cmp(vv->data, ls->data, len, (size_t) ls->len);

            if (rc == 0) {
                ls->conn++;
                goto done;
            }

            node = (rc < 0) ? node->left : node->right;

        } while (node != sentinel && hash == node->key);

        break;
    }

    n = offsetof(ngx_rbtree_node_t, color)
        + offsetof(ngx_http_limit_speed_node_t, data)
        + len;

    node = ngx_slab_alloc_locked(shpool, n);
    if (node == NULL) {
        ngx_shmtx_unlock(&shpool->mutex);
        return NGX_DECLINED;
    }

    ls = (ngx_http_limit_speed_node_t *) &node->color;

    node->key = hash;
    ls->len = (u_char) len;
    ls->conn = 1;
    ngx_memcpy(ls->data, vv->data, len);

    ngx_rbtree_insert(ctx->rbtree, node);

done:

    r->main->limit_rate = lscf->speed / ls->conn;

    ngx_log_debug3(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
                   "limit speed zone: %08XD conn=%d, speed=%d",
                   node->key, ls->conn, r->main->limit_rate);

    ngx_shmtx_unlock(&shpool->mutex);

    cln->handler = ngx_http_limit_speed_cleanup;
    lscln = cln->data;

    lscln->shm_zone = lscf->shm_zone;
    lscln->node = node;

    if (ngx_http_limit_speed_get_ctx(r) != NGX_OK) {
        return NGX_DECLINED;
    }

    rctx = ngx_http_get_module_ctx(r, ngx_http_limit_speed_module);
    rctx->speed = lscf->speed;
    rctx->ls = ls;

    return NGX_DECLINED;
}