Add support for TS_HTTP_REQUEST_CLIENT/TS_HTTP_RESPONSE_CLIENT for lua plugin (#9682)

shukitchan · web-flow · commit 671c3609b8cb · 2023-06-05T15:33:09.000-07:00
* Add support for TS_HTTP_REQUEST_CLIENT/TS_HTTP_RESPONSE_CLIENT

* fix clang format error

* fix unused variables problem

* add examples and fix unused variables errors

* add documentation changes and fix clang format errors

* add gold tests
diff --git a/doc/admin-guide/plugins/lua.en.rst b/doc/admin-guide/plugins/lua.en.rst
@@ -647,6 +647,8 @@ Hook point constants
     TS_LUA_HOOK_SEND_RESPONSE_HDR
     TS_LUA_REQUEST_TRANSFORM
     TS_LUA_RESPONSE_TRANSFORM
+    TS_LUA_REQUEST_CLIENT
+    TS_LUA_RESPONSE_CLIENT
 
 These constants are usually used in ts.hook method call.
 
@@ -694,6 +696,12 @@ Additional Information:
 | TS_HTTP_RESPONSE      | TS_LUA_RESPONSE_TRANSFORM |     YES              |    YES             |    YES               |
 | _TRANSFORM_HOOK       |                           |                      |                    |                      |
 +-----------------------+---------------------------+----------------------+--------------------+----------------------+
+| TS_HTTP_REQUEST       | TS_LUA_REQUEST_CLIENT     |     YES              |    NO              |    YES               |
+| _CLIENT_HOOK          |                           |                      |                    |                      |
++-----------------------+---------------------------+----------------------+--------------------+----------------------+
+| TS_HTTP_RESPONSE      | TS_LUA_RESPONSE_CLIENT    |     YES              |    YES             |    YES               |
+| _CLIENT_HOOK          |                           |                      |                    |                      |
++-----------------------+---------------------------+----------------------+--------------------+----------------------+
 | TS_HTTP_TXN           | TS_LUA_HOOK_TXN_CLOSE     |     YES              |    YES             |    YES               |
 | _CLOSE_HOOK           |                           |                      |                    |                      |
 +-----------------------+---------------------------+----------------------+--------------------+----------------------+
diff --git a/plugins/lua/example/test_request_client.lua b/plugins/lua/example/test_request_client.lua
@@ -0,0 +1,35 @@
+--  Licensed to the Apache Software Foundation (ASF) under one
+--  or more contributor license agreements.  See the NOTICE file
+--  distributed with this work for additional information
+--  regarding copyright ownership.  The ASF licenses this file
+--  to you under the Apache License, Version 2.0 (the
+--  "License"); you may not use this file except in compliance
+--  with the License.  You may obtain a copy of the License at
+--
+--  http://www.apache.org/licenses/LICENSE-2.0
+--
+--  Unless required by applicable law or agreed to in writing, software
+--  distributed under the License is distributed on an "AS IS" BASIS,
+--  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+--  See the License for the specific language governing permissions and
+--  limitations under the License.
+
+
+function request_client(data, eos)
+  ts.debug('request_client')
+  ts.ctx['reqbody'] = ts.ctx['reqbody'] .. data
+  ts.debug("req transform got " .. string.len(data) .. "bytes, eos=" .. eos)
+  if (eos == 1) then
+    ts.debug('End of Stream and the reqbody is ... ')
+    ts.debug(ts.ctx['reqbody'])
+  end
+end
+
+function do_global_read_request()
+  ts.debug('do_global_read_request')
+  if (ts.client_request.get_method() == 'POST') then
+    ts.debug('post')
+    ts.ctx['reqbody'] = ''
+    ts.hook(TS_LUA_REQUEST_CLIENT, request_client)
+  end
+end
diff --git a/plugins/lua/example/test_response_client.lua b/plugins/lua/example/test_response_client.lua
@@ -0,0 +1,25 @@
+--  Licensed to the Apache Software Foundation (ASF) under one
+--  or more contributor license agreements.  See the NOTICE file
+--  distributed with this work for additional information
+--  regarding copyright ownership.  The ASF licenses this file
+--  to you under the Apache License, Version 2.0 (the
+--  "License"); you may not use this file except in compliance
+--  with the License.  You may obtain a copy of the License at
+--
+--  http://www.apache.org/licenses/LICENSE-2.0
+--
+--  Unless required by applicable law or agreed to in writing, software
+--  distributed under the License is distributed on an "AS IS" BASIS,
+--  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+--  See the License for the specific language governing permissions and
+--  limitations under the License.
+
+
+function response_client(data, eos)
+    ts.debug("testing")
+    ts.debug(data)
+end
+
+function do_global_read_request()
+    ts.hook(TS_LUA_RESPONSE_CLIENT, response_client)
+end
diff --git a/plugins/lua/ts_lua_hook.c b/plugins/lua/ts_lua_hook.c
@@ -35,6 +35,8 @@ typedef enum {
   TS_LUA_HOOK_TXN_CLOSE,
   TS_LUA_REQUEST_TRANSFORM,
   TS_LUA_RESPONSE_TRANSFORM,
+  TS_LUA_REQUEST_CLIENT,
+  TS_LUA_RESPONSE_CLIENT,
   TS_LUA_HOOK_LAST
 } TSLuaHookID;
 
@@ -52,6 +54,8 @@ char *ts_lua_hook_id_string[] = {"TS_LUA_HOOK_DUMMY",
                                  "TS_LUA_HOOK_TXN_CLOSE",
                                  "TS_LUA_REQUEST_TRANSFORM",
                                  "TS_LUA_RESPONSE_TRANSFORM",
+                                 "TS_LUA_REQUEST_CLIENT",
+                                 "TS_LUA_RESPONSE_CLIENT",
                                  "TS_LUA_HOOK_LAST"};
 
 static int ts_lua_add_hook(lua_State *L);
@@ -240,6 +244,20 @@ ts_lua_add_hook(lua_State *L)
     }
     break;
 
+  case TS_LUA_REQUEST_CLIENT:
+  case TS_LUA_RESPONSE_CLIENT:
+    if (http_ctx) {
+      connp = TSTransformCreate(ts_lua_client_entry, http_ctx->txnp);
+      ts_lua_create_http_transform_ctx(http_ctx, connp);
+
+      if (entry == TS_LUA_REQUEST_CLIENT) {
+        TSHttpTxnHookAdd(http_ctx->txnp, TS_HTTP_REQUEST_CLIENT_HOOK, connp);
+      } else {
+        TSHttpTxnHookAdd(http_ctx->txnp, TS_HTTP_RESPONSE_CLIENT_HOOK, connp);
+      }
+    }
+    break;
+
   default:
     break;
   }
diff --git a/plugins/lua/ts_lua_transform.c b/plugins/lua/ts_lua_transform.c
@@ -18,8 +18,230 @@
 
 #include "ts_lua_util.h"
 
+static int ts_lua_client_handler(TSCont contp, ts_lua_http_transform_ctx *transform_ctx, TSEvent event, int n);
 static int ts_lua_transform_handler(TSCont contp, ts_lua_http_transform_ctx *transform_ctx, TSEvent event, int n);
 
+int
+ts_lua_client_entry(TSCont contp, TSEvent ev, void *edata)
+{
+  int n, event;
+  TSVIO input_vio;
+  ts_lua_http_transform_ctx *transform_ctx;
+
+  event         = (int)ev;
+  transform_ctx = (ts_lua_http_transform_ctx *)TSContDataGet(contp);
+
+  n = 0;
+
+  switch (event) {
+  case TS_EVENT_ERROR:
+    input_vio = TSVConnWriteVIOGet(contp);
+    TSContCall(TSVIOContGet(input_vio), TS_EVENT_ERROR, input_vio);
+    break;
+
+  case TS_EVENT_VCONN_WRITE_COMPLETE:
+    TSDebug(TS_LUA_DEBUG_TAG, "[%s] received TS_EVENT_VCONN_WRITE_COMPLETE", __FUNCTION__);
+    break;
+
+  case TS_LUA_EVENT_COROUTINE_CONT:
+    n = (intptr_t)edata;
+  /* FALL THROUGH */
+  case TS_EVENT_VCONN_WRITE_READY:
+  default:
+    ts_lua_client_handler(contp, transform_ctx, event, n);
+    break;
+  }
+
+  return 0;
+}
+
+static int
+ts_lua_client_handler(TSCont contp, ts_lua_http_transform_ctx *transform_ctx, TSEvent event, int n)
+{
+  TSVIO input_vio;
+  TSIOBufferReader input_reader;
+  TSIOBufferBlock blk;
+  int64_t toread, towrite, blk_len, upstream_done, input_avail, input_wm_bytes;
+  const char *start;
+  int ret, eos, rc, top, empty_input;
+  ts_lua_coroutine *crt;
+  ts_lua_cont_info *ci;
+
+  lua_State *L;
+  TSMutex mtxp;
+
+  ci  = &transform_ctx->cinfo;
+  crt = &ci->routine;
+
+  mtxp = crt->mctx->mutexp;
+  L    = crt->lua;
+
+  input_vio = TSVConnWriteVIOGet(contp);
+
+  empty_input = 0;
+  if (!TSVIOBufferGet(input_vio)) {
+    TSDebug(TS_LUA_DEBUG_TAG, "[%s] no input VIO and output VIO", __FUNCTION__);
+    empty_input = 1;
+  } else { // input VIO exists
+    input_wm_bytes = TSIOBufferWaterMarkGet(TSVIOBufferGet(input_vio));
+    if (transform_ctx->upstream_watermark_bytes >= 0 && transform_ctx->upstream_watermark_bytes != input_wm_bytes) {
+      TSDebug(TS_LUA_DEBUG_TAG, "[%s] Setting input_vio watermark to %" PRId64 " bytes", __FUNCTION__,
+              transform_ctx->upstream_watermark_bytes);
+      TSIOBufferWaterMarkSet(TSVIOBufferGet(input_vio), transform_ctx->upstream_watermark_bytes);
+    }
+  }
+
+  if (empty_input == 0) {
+    input_reader = TSVIOReaderGet(input_vio);
+  }
+
+  if (!transform_ctx->output.buffer) {
+    transform_ctx->output.buffer = TSIOBufferCreate();
+    transform_ctx->output.reader = TSIOBufferReaderAlloc(transform_ctx->output.buffer);
+
+    transform_ctx->reserved.buffer = TSIOBufferCreate();
+    transform_ctx->reserved.reader = TSIOBufferReaderAlloc(transform_ctx->reserved.buffer);
+
+    if (empty_input == 0) {
+      transform_ctx->upstream_bytes = TSVIONBytesGet(input_vio);
+    } else {
+      transform_ctx->upstream_bytes = 0;
+    }
+
+    transform_ctx->downstream_bytes = INT64_MAX;
+  }
+
+  if (empty_input == 0) {
+    input_avail   = TSIOBufferReaderAvail(input_reader);
+    upstream_done = TSVIONDoneGet(input_vio);
+    toread        = TSVIONTodoGet(input_vio);
+
+    if (toread <= input_avail) { // upstream finished
+      eos = 1;
+    } else {
+      eos = 0;
+    }
+  } else {
+    input_avail   = 0;
+    upstream_done = 0;
+    toread        = 0;
+    eos           = 1;
+  }
+
+  if (input_avail > 0) {
+    // move to the reserved.buffer
+    TSIOBufferCopy(transform_ctx->reserved.buffer, input_reader, input_avail, 0);
+
+    // reset input
+    TSIOBufferReaderConsume(input_reader, input_avail);
+    TSVIONDoneSet(input_vio, upstream_done + input_avail);
+  }
+
+  if (empty_input == 0) {
+    towrite = TSIOBufferReaderAvail(transform_ctx->reserved.reader);
+  } else {
+    towrite = 0;
+  }
+
+  TSMutexLock(mtxp);
+  ts_lua_set_cont_info(L, ci);
+
+  do {
+    if (event == TS_LUA_EVENT_COROUTINE_CONT) {
+      event = 0;
+      goto launch;
+    } else {
+      n = 2;
+    }
+
+    if (towrite == 0 && empty_input == 0) {
+      break;
+    }
+
+    // additional condition for client
+    if (towrite == 0) {
+      break;
+    }
+
+    if (empty_input == 0) {
+      blk   = TSIOBufferReaderStart(transform_ctx->reserved.reader);
+      start = TSIOBufferBlockReadStart(blk, transform_ctx->reserved.reader, &blk_len);
+
+      lua_pushlightuserdata(L, transform_ctx);
+      lua_rawget(L, LUA_GLOBALSINDEX); /* push function */
+
+      if (towrite > blk_len) {
+        lua_pushlstring(L, start, (size_t)blk_len);
+        towrite -= blk_len;
+        TSIOBufferReaderConsume(transform_ctx->reserved.reader, blk_len);
+      } else {
+        lua_pushlstring(L, start, (size_t)towrite);
+        TSIOBufferReaderConsume(transform_ctx->reserved.reader, towrite);
+        towrite = 0;
+      }
+
+      if (!towrite && eos) {
+        lua_pushinteger(L, 1); /* second param, data finished */
+      } else {
+        lua_pushinteger(L, 0); /* second param, data not finish */
+      }
+    } else {
+      lua_pushlightuserdata(L, transform_ctx);
+      lua_rawget(L, LUA_GLOBALSINDEX); /* push function */
+
+      lua_pushlstring(L, "", 0);
+      lua_pushinteger(L, 1); /* second param, data finished */
+    }
+
+  launch:
+    rc  = lua_resume(L, n);
+    top = lua_gettop(L);
+
+    switch (rc) {
+    case LUA_YIELD: // coroutine yield
+      TSMutexUnlock(mtxp);
+      return 0;
+
+    case 0: // coroutine success
+      ret = 0;
+      break;
+
+    default: // coroutine failed
+      TSError("[ts_lua][%s] lua_resume failed: %s", __FUNCTION__, lua_tostring(L, -1));
+      ret = 1;
+      break;
+    }
+
+    lua_pop(L, top);
+
+    if (ret || (eos && !towrite)) { // EOS
+      eos = 1;
+      break;
+    }
+
+  } while (towrite > 0);
+
+  TSMutexUnlock(mtxp);
+
+  if (toread > input_avail) { // upstream not finished.
+    if (eos) {
+      if (empty_input == 0) {
+        TSContCall(TSVIOContGet(input_vio), TS_EVENT_VCONN_EOS, input_vio);
+      }
+    } else {
+      if (empty_input == 0) {
+        TSContCall(TSVIOContGet(input_vio), TS_EVENT_VCONN_WRITE_READY, input_vio);
+      }
+    }
+  } else { // upstream is finished.
+    if (empty_input == 0) {
+      TSContCall(TSVIOContGet(input_vio), TS_EVENT_VCONN_WRITE_COMPLETE, input_vio);
+    }
+  }
+
+  return 0;
+}
+
 int
 ts_lua_transform_entry(TSCont contp, TSEvent ev, void *edata)
 {
diff --git a/plugins/lua/ts_lua_transform.h b/plugins/lua/ts_lua_transform.h
@@ -20,4 +20,5 @@
 
 #include "ts_lua_common.h"
 
+int ts_lua_client_entry(TSCont contp, TSEvent event, void *edata);
 int ts_lua_transform_entry(TSCont contp, TSEvent event, void *edata);
diff --git a/tests/gold_tests/pluginTest/lua/client_hook.lua b/tests/gold_tests/pluginTest/lua/client_hook.lua
@@ -0,0 +1,24 @@
+--  Licensed to the Apache Software Foundation (ASF) under one
+--  or more contributor license agreements.  See the NOTICE file
+--  distributed with this work for additional information
+--  regarding copyright ownership.  The ASF licenses this file
+--  to you under the Apache License, Version 2.0 (the
+--  "License"); you may not use this file except in compliance
+--  with the License.  You may obtain a copy of the License at
+--
+--  http://www.apache.org/licenses/LICENSE-2.0
+--
+--  Unless required by applicable law or agreed to in writing, software
+--  distributed under the License is distributed on an "AS IS" BASIS,
+--  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+--  See the License for the specific language governing permissions and
+--  limitations under the License.
+
+function response_client(data, eos)
+    ts.debug(data)
+end
+
+function do_remap()
+    ts.hook(TS_LUA_RESPONSE_CLIENT, response_client)
+    return 0
+end
diff --git a/tests/gold_tests/pluginTest/lua/lua_client_hook.test.py b/tests/gold_tests/pluginTest/lua/lua_client_hook.test.py

Original file line number	Diff line number	Diff line change
`@@ -20,4 +20,5 @@`
`20`	`20`
`21`	`21`	`#include "ts_lua_common.h"`
`22`	`22`
	`23`	`+int ts_lua_client_entry(TSCont contp, TSEvent event, void *edata);`
`23`	`24`	`int ts_lua_transform_entry(TSCont contp, TSEvent event, void *edata);`