Stream: connection limiting module.

stream {
    limit_conn_zone $binary_remote_addr zone=perip:1m;
    limit_conn_log_level error;

    server {
        ...
        limit_conn perip 1;
    }
}

auto/modules

@@ -514,6 +514,11 @@ if [ $STREAM = YES ]; then
         STREAM_SRCS="$STREAM_SRCS $STREAM_SSL_SRCS"
     fi
 
+    if [ $STREAM_LIMIT_CONN = YES ]; then
+        modules="$modules $STREAM_LIMIT_CONN_MODULE"
+        STREAM_SRCS="$STREAM_SRCS $STREAM_LIMIT_CONN_SRCS"
+    fi
+
     if [ $STREAM_ACCESS = YES ]; then
         modules="$modules $STREAM_ACCESS_MODULE"
         STREAM_SRCS="$STREAM_SRCS $STREAM_ACCESS_SRCS"

auto/options

@@ -113,6 +113,7 @@ MAIL_SMTP=YES
 
 STREAM=NO
 STREAM_SSL=NO
+STREAM_LIMIT_CONN=YES
 STREAM_ACCESS=YES
 STREAM_UPSTREAM_HASH=YES
 STREAM_UPSTREAM_LEAST_CONN=YES
@@ -283,6 +284,8 @@ use the \"--with-mail_ssl_module\" option instead"
 
         --with-stream)                   STREAM=YES                 ;;
         --with-stream_ssl_module)        STREAM_SSL=YES             ;;
+        --without-stream_limit_conn_module)
+                                         STREAM_LIMIT_CONN=NO       ;;
         --without-stream_access_module)  STREAM_ACCESS=NO           ;;
         --without-stream_upstream_hash_module)
                                          STREAM_UPSTREAM_HASH=NO    ;;
@@ -452,6 +455,7 @@ cat << END
 
   --with-stream                      enable TCP proxy module
   --with-stream_ssl_module           enable ngx_stream_ssl_module
+  --without-stream_limit_conn_module disable ngx_stream_limit_conn_module
   --without-stream_access_module     disable ngx_stream_access_module
   --without-stream_upstream_hash_module
                                      disable ngx_stream_upstream_hash_module

auto/sources

@@ -568,6 +568,9 @@ STREAM_SSL_MODULE="ngx_stream_ssl_module"
 STREAM_SSL_DEPS="src/stream/ngx_stream_ssl_module.h"
 STREAM_SSL_SRCS="src/stream/ngx_stream_ssl_module.c"
 
+STREAM_LIMIT_CONN_MODULE=ngx_stream_limit_conn_module
+STREAM_LIMIT_CONN_SRCS=src/stream/ngx_stream_limit_conn_module.c
+
 STREAM_ACCESS_MODULE=ngx_stream_access_module
 STREAM_ACCESS_SRCS=src/stream/ngx_stream_access_module.c
 

src/stream/ngx_stream.h

@@ -118,6 +118,7 @@ typedef ngx_int_t (*ngx_stream_access_pt)(ngx_stream_session_t *s);
 typedef struct {
     ngx_array_t             servers;     /* ngx_stream_core_srv_conf_t */
     ngx_array_t             listen;      /* ngx_stream_listen_t */
+    ngx_stream_access_pt    limit_conn_handler;
     ngx_stream_access_pt    access_handler;
 } ngx_stream_core_main_conf_t;
 

src/stream/ngx_stream_handler.c

@@ -147,6 +147,15 @@ ngx_stream_init_connection(ngx_connection_t *c)
 
     cmcf = ngx_stream_get_module_main_conf(s, ngx_stream_core_module);
 
+    if (cmcf->limit_conn_handler) {
+        rc = cmcf->limit_conn_handler(s);
+
+        if (rc != NGX_DECLINED) {
+            ngx_stream_close_connection(c);
+            return;
+        }
+    }
+
     if (cmcf->access_handler) {
         rc = cmcf->access_handler(s);
 

src/stream/ngx_stream_limit_conn_module.c

@@ -0,0 +1,632 @@
+
+/*
+ * Copyright (C) Igor Sysoev
+ * Copyright (C) Nginx, Inc.
+ */
+
+
+#include <ngx_config.h>
+#include <ngx_core.h>
+#include <ngx_stream.h>
+
+
+typedef struct {
+    u_char                     color;
+    u_char                     len;
+    u_short                    conn;
+    u_char                     data[1];
+} ngx_stream_limit_conn_node_t;
+
+
+typedef struct {
+    ngx_shm_zone_t            *shm_zone;
+    ngx_rbtree_node_t         *node;
+} ngx_stream_limit_conn_cleanup_t;
+
+
+typedef struct {
+    ngx_rbtree_t              *rbtree;
+} ngx_stream_limit_conn_ctx_t;
+
+
+typedef struct {
+    ngx_shm_zone_t            *shm_zone;
+    ngx_uint_t                 conn;
+} ngx_stream_limit_conn_limit_t;
+
+
+typedef struct {
+    ngx_array_t                limits;
+    ngx_uint_t                 log_level;
+} ngx_stream_limit_conn_conf_t;
+
+
+static ngx_rbtree_node_t *ngx_stream_limit_conn_lookup(ngx_rbtree_t *rbtree,
+    ngx_str_t *key, uint32_t hash);
+static void ngx_stream_limit_conn_cleanup(void *data);
+static ngx_inline void ngx_stream_limit_conn_cleanup_all(ngx_pool_t *pool);
+
+static void *ngx_stream_limit_conn_create_conf(ngx_conf_t *cf);
+static char *ngx_stream_limit_conn_merge_conf(ngx_conf_t *cf, void *parent,
+    void *child);
+static char *ngx_stream_limit_conn_zone(ngx_conf_t *cf, ngx_command_t *cmd,
+    void *conf);
+static char *ngx_stream_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd,
+    void *conf);
+static ngx_int_t ngx_stream_limit_conn_init(ngx_conf_t *cf);
+
+
+static ngx_conf_enum_t  ngx_stream_limit_conn_log_levels[] = {
+    { ngx_string("info"), NGX_LOG_INFO },
+    { ngx_string("notice"), NGX_LOG_NOTICE },
+    { ngx_string("warn"), NGX_LOG_WARN },
+    { ngx_string("error"), NGX_LOG_ERR },
+    { ngx_null_string, 0 }
+};
+
+
+static ngx_command_t  ngx_stream_limit_conn_commands[] = {
+
+    { ngx_string("limit_conn_zone"),
+      NGX_STREAM_MAIN_CONF|NGX_CONF_TAKE2,
+      ngx_stream_limit_conn_zone,
+      0,
+      0,
+      NULL },
+
+    { ngx_string("limit_conn"),
+      NGX_STREAM_MAIN_CONF|NGX_STREAM_SRV_CONF|NGX_CONF_TAKE2,
+      ngx_stream_limit_conn,
+      NGX_STREAM_SRV_CONF_OFFSET,
+      0,
+      NULL },
+
+    { ngx_string("limit_conn_log_level"),
+      NGX_STREAM_MAIN_CONF|NGX_STREAM_SRV_CONF|NGX_CONF_TAKE1,
+      ngx_conf_set_enum_slot,
+      NGX_STREAM_SRV_CONF_OFFSET,
+      offsetof(ngx_stream_limit_conn_conf_t, log_level),
+      &ngx_stream_limit_conn_log_levels },
+
+      ngx_null_command
+};
+
+
+static ngx_stream_module_t  ngx_stream_limit_conn_module_ctx = {
+    ngx_stream_limit_conn_init,            /* postconfiguration */
+
+    NULL,                                  /* create main configuration */
+    NULL,                                  /* init main configuration */
+
+    ngx_stream_limit_conn_create_conf,     /* create server configuration */
+    ngx_stream_limit_conn_merge_conf,      /* merge server configuration */
+};
+
+
+ngx_module_t  ngx_stream_limit_conn_module = {
+    NGX_MODULE_V1,
+    &ngx_stream_limit_conn_module_ctx,       /* module context */
+    ngx_stream_limit_conn_commands,          /* module directives */
+    NGX_STREAM_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_stream_limit_conn_handler(ngx_stream_session_t *s)
+{
+    size_t                            n;
+    uint32_t                          hash;
+    ngx_str_t                         key;
+    ngx_uint_t                        i;
+    ngx_slab_pool_t                  *shpool;
+    ngx_rbtree_node_t                *node;
+    ngx_pool_cleanup_t               *cln;
+    struct sockaddr_in               *sin;
+#if (NGX_HAVE_INET6)
+    struct sockaddr_in6              *sin6;
+#endif
+    ngx_stream_limit_conn_ctx_t      *ctx;
+    ngx_stream_limit_conn_node_t     *lc;
+    ngx_stream_limit_conn_conf_t     *lccf;
+    ngx_stream_limit_conn_limit_t    *limits;
+    ngx_stream_limit_conn_cleanup_t  *lccln;
+
+    switch (s->connection->sockaddr->sa_family) {
+
+    case AF_INET:
+        sin = (struct sockaddr_in *) s->connection->sockaddr;
+
+        key.len = sizeof(in_addr_t);
+        key.data = (u_char *) &sin->sin_addr;
+
+        break;
+
+#if (NGX_HAVE_INET6)
+    case AF_INET6:
+        sin6 = (struct sockaddr_in6 *) s->connection->sockaddr;
+
+        key.len = sizeof(struct in6_addr);
+        key.data = sin6->sin6_addr.s6_addr;
+
+        break;
+#endif
+
+    default:
+        return NGX_DECLINED;
+    }
+
+    hash = ngx_crc32_short(key.data, key.len);
+
+    lccf = ngx_stream_get_module_srv_conf(s, ngx_stream_limit_conn_module);
+    limits = lccf->limits.elts;
+
+    for (i = 0; i < lccf->limits.nelts; i++) {
+        ctx = limits[i].shm_zone->data;
+
+        shpool = (ngx_slab_pool_t *) limits[i].shm_zone->shm.addr;
+
+        ngx_shmtx_lock(&shpool->mutex);
+
+        node = ngx_stream_limit_conn_lookup(ctx->rbtree, &key, hash);
+
+        if (node == NULL) {
+
+            n = offsetof(ngx_rbtree_node_t, color)
+                + offsetof(ngx_stream_limit_conn_node_t, data)
+                + key.len;
+
+            node = ngx_slab_alloc_locked(shpool, n);
+
+            if (node == NULL) {
+                ngx_shmtx_unlock(&shpool->mutex);
+                ngx_stream_limit_conn_cleanup_all(s->connection->pool);
+                return NGX_ABORT;
+            }
+
+            lc = (ngx_stream_limit_conn_node_t *) &node->color;
+
+            node->key = hash;
+            lc->len = (u_char) key.len;
+            lc->conn = 1;
+            ngx_memcpy(lc->data, key.data, key.len);
+
+            ngx_rbtree_insert(ctx->rbtree, node);
+
+        } else {
+
+            lc = (ngx_stream_limit_conn_node_t *) &node->color;
+
+            if ((ngx_uint_t) lc->conn >= limits[i].conn) {
+
+                ngx_shmtx_unlock(&shpool->mutex);
+
+                ngx_log_error(lccf->log_level, s->connection->log, 0,
+                              "limiting connections by zone \"%V\"",
+                              &limits[i].shm_zone->shm.name);
+
+                ngx_stream_limit_conn_cleanup_all(s->connection->pool);
+                return NGX_ABORT;
+            }
+
+            lc->conn++;
+        }
+
+        ngx_log_debug2(NGX_LOG_DEBUG_STREAM, s->connection->log, 0,
+                       "limit conn: %08XD %d", node->key, lc->conn);
+
+        ngx_shmtx_unlock(&shpool->mutex);
+
+        cln = ngx_pool_cleanup_add(s->connection->pool,
+                                   sizeof(ngx_stream_limit_conn_cleanup_t));
+        if (cln == NULL) {
+            return NGX_ERROR;
+        }
+
+        cln->handler = ngx_stream_limit_conn_cleanup;
+        lccln = cln->data;
+
+        lccln->shm_zone = limits[i].shm_zone;
+        lccln->node = node;
+    }
+
+    return NGX_DECLINED;
+}
+
+
+static void
+ngx_stream_limit_conn_rbtree_insert_value(ngx_rbtree_node_t *temp,
+    ngx_rbtree_node_t *node, ngx_rbtree_node_t *sentinel)
+{
+    ngx_rbtree_node_t             **p;
+    ngx_stream_limit_conn_node_t   *lcn, *lcnt;
+
+    for ( ;; ) {
+
+        if (node->key < temp->key) {
+
+            p = &temp->left;
+
+        } else if (node->key > temp->key) {
+
+            p = &temp->right;
+
+        } else { /* node->key == temp->key */
+
+            lcn = (ngx_stream_limit_conn_node_t *) &node->color;
+            lcnt = (ngx_stream_limit_conn_node_t *) &temp->color;
+
+            p = (ngx_memn2cmp(lcn->data, lcnt->data, lcn->len, lcnt->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_rbtree_node_t *
+ngx_stream_limit_conn_lookup(ngx_rbtree_t *rbtree, ngx_str_t *key,
+    uint32_t hash)
+{
+    ngx_int_t                      rc;
+    ngx_rbtree_node_t             *node, *sentinel;
+    ngx_stream_limit_conn_node_t  *lcn;
+
+    node = rbtree->root;
+    sentinel = rbtree->sentinel;
+
+    while (node != sentinel) {
+
+        if (hash < node->key) {
+            node = node->left;
+            continue;
+        }
+
+        if (hash > node->key) {
+            node = node->right;
+            continue;
+        }
+
+        /* hash == node->key */
+
+        lcn = (ngx_stream_limit_conn_node_t *) &node->color;
+
+        rc = ngx_memn2cmp(key->data, lcn->data, key->len, (size_t) lcn->len);
+
+        if (rc == 0) {
+            return node;
+        }
+
+        node = (rc < 0) ? node->left : node->right;
+    }
+
+    return NULL;
+}
+
+
+static void
+ngx_stream_limit_conn_cleanup(void *data)
+{
+    ngx_stream_limit_conn_cleanup_t  *lccln = data;
+
+    ngx_slab_pool_t               *shpool;
+    ngx_rbtree_node_t             *node;
+    ngx_stream_limit_conn_ctx_t   *ctx;
+    ngx_stream_limit_conn_node_t  *lc;
+
+    ctx = lccln->shm_zone->data;
+    shpool = (ngx_slab_pool_t *) lccln->shm_zone->shm.addr;
+    node = lccln->node;
+    lc = (ngx_stream_limit_conn_node_t *) &node->color;
+
+    ngx_shmtx_lock(&shpool->mutex);
+
+    ngx_log_debug2(NGX_LOG_DEBUG_STREAM, lccln->shm_zone->shm.log, 0,
+                   "limit conn cleanup: %08XD %d", node->key, lc->conn);
+
+    lc->conn--;
+
+    if (lc->conn == 0) {
+        ngx_rbtree_delete(ctx->rbtree, node);
+        ngx_slab_free_locked(shpool, node);
+    }
+
+    ngx_shmtx_unlock(&shpool->mutex);
+}
+
+
+static ngx_inline void
+ngx_stream_limit_conn_cleanup_all(ngx_pool_t *pool)
+{
+    ngx_pool_cleanup_t  *cln;
+
+    cln = pool->cleanup;
+
+    while (cln && cln->handler == ngx_stream_limit_conn_cleanup) {
+        ngx_stream_limit_conn_cleanup(cln->data);
+        cln = cln->next;
+    }
+
+    pool->cleanup = cln;
+}
+
+
+static ngx_int_t
+ngx_stream_limit_conn_init_zone(ngx_shm_zone_t *shm_zone, void *data)
+{
+    ngx_stream_limit_conn_ctx_t  *octx = data;
+
+    size_t                        len;
+    ngx_slab_pool_t              *shpool;
+    ngx_rbtree_node_t            *sentinel;
+    ngx_stream_limit_conn_ctx_t  *ctx;
+
+    ctx = shm_zone->data;
+
+    if (octx) {
+        ctx->rbtree = octx->rbtree;
+
+        return NGX_OK;
+    }
+
+    shpool = (ngx_slab_pool_t *) shm_zone->shm.addr;
+
+    if (shm_zone->shm.exists) {
+        ctx->rbtree = shpool->data;
+
+        return NGX_OK;
+    }
+
+    ctx->rbtree = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_t));
+    if (ctx->rbtree == NULL) {
+        return NGX_ERROR;
+    }
+
+    shpool->data = ctx->rbtree;
+
+    sentinel = ngx_slab_alloc(shpool, sizeof(ngx_rbtree_node_t));
+    if (sentinel == NULL) {
+        return NGX_ERROR;
+    }
+
+    ngx_rbtree_init(ctx->rbtree, sentinel,
+                    ngx_stream_limit_conn_rbtree_insert_value);
+
+    len = sizeof(" in limit_conn_zone \"\"") + shm_zone->shm.name.len;
+
+    shpool->log_ctx = ngx_slab_alloc(shpool, len);
+    if (shpool->log_ctx == NULL) {
+        return NGX_ERROR;
+    }
+
+    ngx_sprintf(shpool->log_ctx, " in limit_conn_zone \"%V\"%Z",
+                &shm_zone->shm.name);
+
+    return NGX_OK;
+}
+
+
+static void *
+ngx_stream_limit_conn_create_conf(ngx_conf_t *cf)
+{
+    ngx_stream_limit_conn_conf_t  *conf;
+
+    conf = ngx_pcalloc(cf->pool, sizeof(ngx_stream_limit_conn_conf_t));
+    if (conf == NULL) {
+        return NULL;
+    }
+
+    /*
+     * set by ngx_pcalloc():
+     *
+     *     conf->limits.elts = NULL;
+     */
+
+    conf->log_level = NGX_CONF_UNSET_UINT;
+
+    return conf;
+}
+
+
+static char *
+ngx_stream_limit_conn_merge_conf(ngx_conf_t *cf, void *parent, void *child)
+{
+    ngx_stream_limit_conn_conf_t *prev = parent;
+    ngx_stream_limit_conn_conf_t *conf = child;
+
+    if (conf->limits.elts == NULL) {
+        conf->limits = prev->limits;
+    }
+
+    ngx_conf_merge_uint_value(conf->log_level, prev->log_level, NGX_LOG_ERR);
+
+    return NGX_CONF_OK;
+}
+
+
+static char *
+ngx_stream_limit_conn_zone(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
+{
+    u_char                       *p;
+    ssize_t                       size;
+    ngx_str_t                    *value, name, s;
+    ngx_uint_t                    i;
+    ngx_shm_zone_t               *shm_zone;
+    ngx_stream_limit_conn_ctx_t  *ctx;
+
+    value = cf->args->elts;
+
+    ctx = ngx_pcalloc(cf->pool, sizeof(ngx_stream_limit_conn_ctx_t));
+    if (ctx == NULL) {
+        return NGX_CONF_ERROR;
+    }
+
+    size = 0;
+    name.len = 0;
+
+    for (i = 2; i < cf->args->nelts; i++) {
+
+        if (ngx_strncmp(value[i].data, "zone=", 5) == 0) {
+
+            name.data = value[i].data + 5;
+
+            p = (u_char *) ngx_strchr(name.data, ':');
+
+            if (p == NULL) {
+                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                                   "invalid zone size \"%V\"", &value[i]);
+                return NGX_CONF_ERROR;
+            }
+
+            name.len = p - name.data;
+
+            s.data = p + 1;
+            s.len = value[i].data + value[i].len - s.data;
+
+            size = ngx_parse_size(&s);
+
+            if (size == NGX_ERROR) {
+                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                                   "invalid zone size \"%V\"", &value[i]);
+                return NGX_CONF_ERROR;
+            }
+
+            if (size < (ssize_t) (8 * ngx_pagesize)) {
+                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                                   "zone \"%V\" is too small", &value[i]);
+                return NGX_CONF_ERROR;
+            }
+
+            continue;
+        }
+
+        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                           "invalid parameter \"%V\"", &value[i]);
+        return NGX_CONF_ERROR;
+    }
+
+    if (name.len == 0) {
+        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                           "\"%V\" must have \"zone\" parameter",
+                           &cmd->name);
+        return NGX_CONF_ERROR;
+    }
+
+    shm_zone = ngx_shared_memory_add(cf, &name, size,
+                                     &ngx_stream_limit_conn_module);
+    if (shm_zone == NULL) {
+        return NGX_CONF_ERROR;
+    }
+
+    if (shm_zone->data) {
+        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                           "%V \"%V\" is already bound to key "
+                           "\"$binary_remote_addr\"",
+                           &cmd->name, &name);
+        return NGX_CONF_ERROR;
+    }
+
+    if (ngx_strcmp(value[1].data, "$binary_remote_addr") != 0) {
+        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                           "unsupported key \"%V\", use "
+                           "$binary_remote_addr", &value[1]);
+        return NGX_CONF_ERROR;
+    }
+
+    shm_zone->init = ngx_stream_limit_conn_init_zone;
+    shm_zone->data = ctx;
+
+    return NGX_CONF_OK;
+}
+
+
+static char *
+ngx_stream_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
+{
+    ngx_shm_zone_t                 *shm_zone;
+    ngx_stream_limit_conn_conf_t   *lccf = conf;
+    ngx_stream_limit_conn_limit_t  *limit, *limits;
+
+    ngx_str_t   *value;
+    ngx_int_t    n;
+    ngx_uint_t   i;
+
+    value = cf->args->elts;
+
+    shm_zone = ngx_shared_memory_add(cf, &value[1], 0,
+                                     &ngx_stream_limit_conn_module);
+    if (shm_zone == NULL) {
+        return NGX_CONF_ERROR;
+    }
+
+    limits = lccf->limits.elts;
+
+    if (limits == NULL) {
+        if (ngx_array_init(&lccf->limits, cf->pool, 1,
+                           sizeof(ngx_stream_limit_conn_limit_t))
+            != NGX_OK)
+        {
+            return NGX_CONF_ERROR;
+        }
+    }
+
+    for (i = 0; i < lccf->limits.nelts; i++) {
+        if (shm_zone == limits[i].shm_zone) {
+            return "is duplicate";
+        }
+    }
+
+    n = ngx_atoi(value[2].data, value[2].len);
+    if (n <= 0) {
+        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                           "invalid number of connections \"%V\"", &value[2]);
+        return NGX_CONF_ERROR;
+    }
+
+    if (n > 65535) {
+        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
+                           "connection limit must be less 65536");
+        return NGX_CONF_ERROR;
+    }
+
+    limit = ngx_array_push(&lccf->limits);
+    if (limit == NULL) {
+        return NGX_CONF_ERROR;
+    }
+
+    limit->conn = n;
+    limit->shm_zone = shm_zone;
+
+    return NGX_CONF_OK;
+}
+
+
+static ngx_int_t
+ngx_stream_limit_conn_init(ngx_conf_t *cf)
+{
+    ngx_stream_core_main_conf_t  *cmcf;
+
+    cmcf = ngx_stream_conf_get_module_main_conf(cf, ngx_stream_core_module);
+
+    cmcf->limit_conn_handler = ngx_stream_limit_conn_handler;
+
+    return NGX_OK;
+}