changeset 7595:9606d93aa586

Limit conn: $limit_conn_status variable. The variable takes one of the values: PASSED, REJECTED or REJECTED_DRY_RUN.
author Roman Arutyunyan <arut@nginx.com>
date Mon, 18 Nov 2019 17:48:32 +0300
parents 359b0ea2b067
children b45f052483b8
files src/http/modules/ngx_http_limit_conn_module.c src/http/ngx_http_request.h src/stream/ngx_stream.h src/stream/ngx_stream_limit_conn_module.c
diffstat 4 files changed, 149 insertions(+), 5 deletions(-) [+]
line wrap: on
line diff
--- a/src/http/modules/ngx_http_limit_conn_module.c
+++ b/src/http/modules/ngx_http_limit_conn_module.c
@@ -10,6 +10,11 @@
 #include <ngx_http.h>
 
 
+#define NGX_HTTP_LIMIT_CONN_PASSED            1
+#define NGX_HTTP_LIMIT_CONN_REJECTED          2
+#define NGX_HTTP_LIMIT_CONN_REJECTED_DRY_RUN  3
+
+
 typedef struct {
     u_char                     color;
     u_char                     len;
@@ -49,6 +54,8 @@ static ngx_rbtree_node_t *ngx_http_limit
 static void ngx_http_limit_conn_cleanup(void *data);
 static ngx_inline void ngx_http_limit_conn_cleanup_all(ngx_pool_t *pool);
 
+static ngx_int_t ngx_http_limit_conn_status_variable(ngx_http_request_t *r,
+    ngx_http_variable_value_t *v, uintptr_t data);
 static void *ngx_http_limit_conn_create_conf(ngx_conf_t *cf);
 static char *ngx_http_limit_conn_merge_conf(ngx_conf_t *cf, void *parent,
     void *child);
@@ -56,6 +63,7 @@ static char *ngx_http_limit_conn_zone(ng
     void *conf);
 static char *ngx_http_limit_conn(ngx_conf_t *cf, ngx_command_t *cmd,
     void *conf);
+static ngx_int_t ngx_http_limit_conn_add_variables(ngx_conf_t *cf);
 static ngx_int_t ngx_http_limit_conn_init(ngx_conf_t *cf);
 
 
@@ -115,7 +123,7 @@ static ngx_command_t  ngx_http_limit_con
 
 
 static ngx_http_module_t  ngx_http_limit_conn_module_ctx = {
-    NULL,                                  /* preconfiguration */
+    ngx_http_limit_conn_add_variables,     /* preconfiguration */
     ngx_http_limit_conn_init,              /* postconfiguration */
 
     NULL,                                  /* create main configuration */
@@ -145,6 +153,22 @@ ngx_module_t  ngx_http_limit_conn_module
 };
 
 
+static ngx_http_variable_t  ngx_http_limit_conn_vars[] = {
+
+    { ngx_string("limit_conn_status"), NULL,
+      ngx_http_limit_conn_status_variable, 0, NGX_HTTP_VAR_NOCACHEABLE, 0 },
+
+      ngx_http_null_variable
+};
+
+
+static ngx_str_t  ngx_http_limit_conn_status[] = {
+    ngx_string("PASSED"),
+    ngx_string("REJECTED"),
+    ngx_string("REJECTED_DRY_RUN")
+};
+
+
 static ngx_int_t
 ngx_http_limit_conn_handler(ngx_http_request_t *r)
 {
@@ -161,7 +185,7 @@ ngx_http_limit_conn_handler(ngx_http_req
     ngx_http_limit_conn_limit_t    *limits;
     ngx_http_limit_conn_cleanup_t  *lccln;
 
-    if (r->main->limit_conn_set) {
+    if (r->main->limit_conn_status) {
         return NGX_DECLINED;
     }
 
@@ -187,7 +211,7 @@ ngx_http_limit_conn_handler(ngx_http_req
             continue;
         }
 
-        r->main->limit_conn_set = 1;
+        r->main->limit_conn_status = NGX_HTTP_LIMIT_CONN_PASSED;
 
         hash = ngx_crc32_short(key.data, key.len);
 
@@ -210,9 +234,13 @@ ngx_http_limit_conn_handler(ngx_http_req
                 ngx_http_limit_conn_cleanup_all(r->pool);
 
                 if (lccf->dry_run) {
+                    r->main->limit_conn_status =
+                                          NGX_HTTP_LIMIT_CONN_REJECTED_DRY_RUN;
                     return NGX_DECLINED;
                 }
 
+                r->main->limit_conn_status = NGX_HTTP_LIMIT_CONN_REJECTED;
+
                 return lccf->status_code;
             }
 
@@ -241,9 +269,13 @@ ngx_http_limit_conn_handler(ngx_http_req
                 ngx_http_limit_conn_cleanup_all(r->pool);
 
                 if (lccf->dry_run) {
+                    r->main->limit_conn_status =
+                                          NGX_HTTP_LIMIT_CONN_REJECTED_DRY_RUN;
                     return NGX_DECLINED;
                 }
 
+                r->main->limit_conn_status = NGX_HTTP_LIMIT_CONN_REJECTED;
+
                 return lccf->status_code;
             }
 
@@ -467,6 +499,25 @@ ngx_http_limit_conn_init_zone(ngx_shm_zo
 }
 
 
+static ngx_int_t
+ngx_http_limit_conn_status_variable(ngx_http_request_t *r,
+    ngx_http_variable_value_t *v, uintptr_t data)
+{
+    if (r->main->limit_conn_status == 0) {
+        v->not_found = 1;
+        return NGX_OK;
+    }
+
+    v->valid = 1;
+    v->no_cacheable = 0;
+    v->not_found = 0;
+    v->len = ngx_http_limit_conn_status[r->main->limit_conn_status - 1].len;
+    v->data = ngx_http_limit_conn_status[r->main->limit_conn_status - 1].data;
+
+    return NGX_OK;
+}
+
+
 static void *
 ngx_http_limit_conn_create_conf(ngx_conf_t *cf)
 {
@@ -674,6 +725,25 @@ ngx_http_limit_conn(ngx_conf_t *cf, ngx_
 
 
 static ngx_int_t
+ngx_http_limit_conn_add_variables(ngx_conf_t *cf)
+{
+    ngx_http_variable_t  *var, *v;
+
+    for (v = ngx_http_limit_conn_vars; v->name.len; v++) {
+        var = ngx_http_add_variable(cf, &v->name, v->flags);
+        if (var == NULL) {
+            return NGX_ERROR;
+        }
+
+        var->get_handler = v->get_handler;
+        var->data = v->data;
+    }
+
+    return NGX_OK;
+}
+
+
+static ngx_int_t
 ngx_http_limit_conn_init(ngx_conf_t *cf)
 {
     ngx_http_handler_pt        *h;
--- a/src/http/ngx_http_request.h
+++ b/src/http/ngx_http_request.h
@@ -512,7 +512,7 @@ struct ngx_http_request_s {
      * ngx_http_limit_conn_module and ngx_http_limit_req_module
      * we use the bit fields in the request structure
      */
-    unsigned                          limit_conn_set:1;
+    unsigned                          limit_conn_status:2;
     unsigned                          limit_req_status:3;
 
     unsigned                          limit_rate_set:1;
--- a/src/stream/ngx_stream.h
+++ b/src/stream/ngx_stream.h
@@ -226,6 +226,8 @@ struct ngx_stream_session_s {
     unsigned                       stat_processing:1;
 
     unsigned                       health_check:1;
+
+    unsigned                       limit_conn_status:2;
 };
 
 
--- a/src/stream/ngx_stream_limit_conn_module.c
+++ b/src/stream/ngx_stream_limit_conn_module.c
@@ -10,6 +10,11 @@
 #include <ngx_stream.h>
 
 
+#define NGX_STREAM_LIMIT_CONN_PASSED            1
+#define NGX_STREAM_LIMIT_CONN_REJECTED          2
+#define NGX_STREAM_LIMIT_CONN_REJECTED_DRY_RUN  3
+
+
 typedef struct {
     u_char                       color;
     u_char                       len;
@@ -48,6 +53,8 @@ static ngx_rbtree_node_t *ngx_stream_lim
 static void ngx_stream_limit_conn_cleanup(void *data);
 static ngx_inline void ngx_stream_limit_conn_cleanup_all(ngx_pool_t *pool);
 
+static ngx_int_t ngx_stream_limit_conn_status_variable(ngx_stream_session_t *s,
+    ngx_stream_variable_value_t *v, uintptr_t data);
 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);
@@ -55,6 +62,7 @@ static char *ngx_stream_limit_conn_zone(
     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_add_variables(ngx_conf_t *cf);
 static ngx_int_t ngx_stream_limit_conn_init(ngx_conf_t *cf);
 
 
@@ -102,7 +110,7 @@ static ngx_command_t  ngx_stream_limit_c
 
 
 static ngx_stream_module_t  ngx_stream_limit_conn_module_ctx = {
-    NULL,                                  /* preconfiguration */
+    ngx_stream_limit_conn_add_variables,   /* preconfiguration */
     ngx_stream_limit_conn_init,            /* postconfiguration */
 
     NULL,                                  /* create main configuration */
@@ -129,6 +137,22 @@ ngx_module_t  ngx_stream_limit_conn_modu
 };
 
 
+static ngx_stream_variable_t  ngx_stream_limit_conn_vars[] = {
+
+    { ngx_string("limit_conn_status"), NULL,
+      ngx_stream_limit_conn_status_variable, 0, NGX_STREAM_VAR_NOCACHEABLE, 0 },
+
+      ngx_stream_null_variable
+};
+
+
+static ngx_str_t  ngx_stream_limit_conn_status[] = {
+    ngx_string("PASSED"),
+    ngx_string("REJECTED"),
+    ngx_string("REJECTED_DRY_RUN")
+};
+
+
 static ngx_int_t
 ngx_stream_limit_conn_handler(ngx_stream_session_t *s)
 {
@@ -167,6 +191,8 @@ ngx_stream_limit_conn_handler(ngx_stream
             continue;
         }
 
+        s->limit_conn_status = NGX_STREAM_LIMIT_CONN_PASSED;
+
         hash = ngx_crc32_short(key.data, key.len);
 
         shpool = (ngx_slab_pool_t *) limits[i].shm_zone->shm.addr;
@@ -188,9 +214,13 @@ ngx_stream_limit_conn_handler(ngx_stream
                 ngx_stream_limit_conn_cleanup_all(s->connection->pool);
 
                 if (lccf->dry_run) {
+                    s->limit_conn_status =
+                                        NGX_STREAM_LIMIT_CONN_REJECTED_DRY_RUN;
                     return NGX_DECLINED;
                 }
 
+                s->limit_conn_status = NGX_STREAM_LIMIT_CONN_REJECTED;
+
                 return NGX_STREAM_SERVICE_UNAVAILABLE;
             }
 
@@ -219,9 +249,13 @@ ngx_stream_limit_conn_handler(ngx_stream
                 ngx_stream_limit_conn_cleanup_all(s->connection->pool);
 
                 if (lccf->dry_run) {
+                    s->limit_conn_status =
+                                        NGX_STREAM_LIMIT_CONN_REJECTED_DRY_RUN;
                     return NGX_DECLINED;
                 }
 
+                s->limit_conn_status = NGX_STREAM_LIMIT_CONN_REJECTED;
+
                 return NGX_STREAM_SERVICE_UNAVAILABLE;
             }
 
@@ -446,6 +480,25 @@ ngx_stream_limit_conn_init_zone(ngx_shm_
 }
 
 
+static ngx_int_t
+ngx_stream_limit_conn_status_variable(ngx_stream_session_t *s,
+    ngx_stream_variable_value_t *v, uintptr_t data)
+{
+    if (s->limit_conn_status == 0) {
+        v->not_found = 1;
+        return NGX_OK;
+    }
+
+    v->valid = 1;
+    v->no_cacheable = 0;
+    v->not_found = 0;
+    v->len = ngx_stream_limit_conn_status[s->limit_conn_status - 1].len;
+    v->data = ngx_stream_limit_conn_status[s->limit_conn_status - 1].data;
+
+    return NGX_OK;
+}
+
+
 static void *
 ngx_stream_limit_conn_create_conf(ngx_conf_t *cf)
 {
@@ -650,6 +703,25 @@ ngx_stream_limit_conn(ngx_conf_t *cf, ng
 
 
 static ngx_int_t
+ngx_stream_limit_conn_add_variables(ngx_conf_t *cf)
+{
+    ngx_stream_variable_t  *var, *v;
+
+    for (v = ngx_stream_limit_conn_vars; v->name.len; v++) {
+        var = ngx_stream_add_variable(cf, &v->name, v->flags);
+        if (var == NULL) {
+            return NGX_ERROR;
+        }
+
+        var->get_handler = v->get_handler;
+        var->data = v->data;
+    }
+
+    return NGX_OK;
+}
+
+
+static ngx_int_t
 ngx_stream_limit_conn_init(ngx_conf_t *cf)
 {
     ngx_stream_handler_pt        *h;