lua_add_variable - initialize value (add optional optional arguments)

openresty · Apr 25, 2018 · 26be5ce · 26be5ce
1 parent 479d520
commit 26be5ce
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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]