From 8a49a2d250d02e20296211bce32dab7184f0ae2f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Krzysztof=20Wojta=C5=9B?= <krzysin@gmail.com>
Date: Wed, 25 Apr 2018 11:40:34 +0200
Subject: [PATCH] lua_add_variable - initialize value (add optional arguments)

---
 src/ngx_stream_lua_directive.c | 47 +++++++++++++++++++++++++++--
 src/ngx_stream_lua_module.c    |  2 +-
 t/139-add-variable.t           | 55 ++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+), 4 deletions(-)

diff --git a/src/ngx_stream_lua_directive.c b/src/ngx_stream_lua_directive.c
index e8732ee8..da9ce80e 100644
--- a/src/ngx_stream_lua_directive.c
+++ b/src/ngx_stream_lua_directive.c
@@ -1209,12 +1209,29 @@ ngx_stream_lua_undefined_var(ngx_stream_session_t *s,
 }
 
 
+static ngx_int_t
+ngx_stream_lua_var(ngx_stream_session_t *s,
+    ngx_stream_variable_value_t *v, uintptr_t data)
+{
+    ngx_str_t *str = (ngx_str_t *) data;
+
+    v->not_found = 0;
+    v->valid = 1;
+    v->no_cacheable = 0;
+    v->not_found = 0;
+    v->len = str->len;
+    v->data = str->data;
+
+    return NGX_OK;
+}
+
+
 char *
 ngx_stream_lua_add_variable(ngx_conf_t *cf, ngx_command_t *cmd,
     void *conf)
 {
     ngx_stream_variable_t           *var;
-    ngx_str_t                       *value;
+    ngx_str_t                       *value, *data;
     ngx_int_t                        ret;
 
     value = cf->args->elts;
@@ -1234,8 +1251,32 @@ ngx_stream_lua_add_variable(ngx_conf_t *cf, ngx_command_t *cmd,
         return NGX_CONF_ERROR;
     }
 
-    if (var->get_handler == NULL) {
-        var->get_handler = ngx_stream_lua_undefined_var;
+    if (cf->args->nelts == 2) {
+        if (var->get_handler == NULL) {
+            var->get_handler = ngx_stream_lua_undefined_var;
+        }
+    } else {
+        if (var->get_handler == NULL) {
+            var->get_handler = ngx_stream_lua_var;
+        }
+
+        if (var->data) {
+            data = (ngx_str_t *) var->data;
+        } else {
+            data = ngx_pnalloc(cf->pool, sizeof(ngx_str_t));
+            if (data == NULL) {
+                return NGX_CONF_ERROR;
+            }
+        }
+        data->len = value[2].len;
+
+        data->data = ngx_pnalloc(cf->pool, value[2].len);
+        if (data->data == NULL) {
+            return NGX_CONF_ERROR;
+        }
+        ngx_memcpy(data->data, value[2].data, value[2].len);
+
+        var->data = (uintptr_t) data;
     }
 
     ret = ngx_stream_get_variable_index(cf, value + 1);
diff --git a/src/ngx_stream_lua_module.c b/src/ngx_stream_lua_module.c
index a3b81c50..64505140 100644
--- a/src/ngx_stream_lua_module.c
+++ b/src/ngx_stream_lua_module.c
@@ -382,7 +382,7 @@ static ngx_command_t ngx_stream_lua_cmds[] = {
       NULL },
 
      { ngx_string("lua_add_variable"),
-      NGX_STREAM_MAIN_CONF|NGX_CONF_TAKE1,
+      NGX_STREAM_MAIN_CONF|NGX_CONF_TAKE12,
       ngx_stream_lua_add_variable,
       0,
       0,
diff --git a/t/139-add-variable.t b/t/139-add-variable.t
index 4798d215..48f67a9b 100644
--- a/t/139-add-variable.t
+++ b/t/139-add-variable.t
@@ -88,3 +88,58 @@ bar
 [error]
 --- error_log
 access log: bar
+
+
+
+=== TEST 5: sanity
+--- stream_config
+    lua_add_variable $foo "foo";
+--- stream_server_config
+    content_by_lua_block {
+        ngx.say(ngx.var.foo)
+        ngx.var.foo = "bar"
+        ngx.say(ngx.var.foo)
+    }
+--- stream_response
+foo
+bar
+--- no_error_log
+[warn]
+[error]
+
+
+
+=== TEST 6: works with C code
+--- stream_config
+    lua_add_variable $foo "foo";
+--- stream_server_config
+    preread_by_lua_block {
+        ngx.var.foo = "bar"
+    }
+
+    return $foo\n;
+--- stream_response
+bar
+--- no_error_log
+[warn]
+[error]
+
+
+
+=== TEST 7: multiple add with same name works
+--- stream_config
+    lua_add_variable $foo "bar";
+    lua_add_variable $foo "foo";
+--- stream_server_config
+    preread_by_lua_block {
+        ngx.say(ngx.var.foo)
+        ngx.var.foo = "bar"
+    }
+
+    return $foo\n;
+--- stream_response
+foo
+bar
+--- no_error_log
+[warn]
+[error]