feat: add method to reset read buf (#47)

Author: spacewander

lib/resty/apisix/stream/xrpc/socket.lua

@@ -39,11 +39,16 @@ int
 ngx_stream_lua_ffi_socket_tcp_get_send_result(ngx_stream_lua_request_t *r,
     ngx_stream_lua_socket_tcp_upstream_t *u, u_char *errbuf,
     size_t *errbuf_size);
+
+void
+ngx_stream_lua_ffi_socket_tcp_reset_read_buf(ngx_stream_lua_request_t *r,
+    ngx_stream_lua_socket_tcp_upstream_t *u);
 ]]
 local socket_tcp_read = C.ngx_stream_lua_ffi_socket_tcp_read_buf
 local socket_tcp_get_read_result = C.ngx_stream_lua_ffi_socket_tcp_get_read_buf_result
 local socket_tcp_move = C.ngx_stream_lua_ffi_socket_tcp_send_from_socket
 local socket_tcp_get_move_result = C.ngx_stream_lua_ffi_socket_tcp_get_send_result
+local socket_tcp_reset_read_buf = C.ngx_stream_lua_ffi_socket_tcp_reset_read_buf
 
 
 local ERR_BUF_SIZE = 256
@@ -191,6 +196,19 @@ local function move(dst, src)
 end
 
 
+-- reset buffer read from methods `read` or `drain`. Should be used when you don't
+-- want to forward some buffers
+local function reset_read_buf(cosocket)
+    local r = get_request()
+    if not r then
+        error("no request found", 2)
+    end
+
+    local u = get_tcp_socket(cosocket)
+    socket_tcp_reset_read_buf(r, u)
+end
+
+
 local function patch_methods(sk)
     local methods = getmetatable(sk).__index
     local copy = tab_clone(methods)
@@ -202,6 +220,7 @@ local function patch_methods(sk)
     copy.read = read
     copy.drain = drain
     copy.move = move
+    copy.reset_read_buf = reset_read_buf
 
     return {__index = copy}
 end

patch/1.19.9/ngx_stream_lua-xrpc.patch

@@ -1,5 +1,5 @@
 diff --git src/ngx_stream_lua_socket_tcp.c src/ngx_stream_lua_socket_tcp.c
-index 7fcfb45..13e942d 100644
+index 7fcfb45..007305b 100644
 --- src/ngx_stream_lua_socket_tcp.c
 +++ src/ngx_stream_lua_socket_tcp.c
 @@ -234,6 +234,41 @@ enum {
@@ -44,7 +44,7 @@ index 7fcfb45..13e942d 100644
 
  static char ngx_stream_lua_raw_req_socket_metatable_key;
  static char ngx_stream_lua_tcp_socket_metatable_key;
-@@ -6005,6 +6040,605 @@ static ngx_int_t ngx_stream_lua_socket_insert_buffer(
+@@ -6005,6 +6040,620 @@ static ngx_int_t ngx_stream_lua_socket_insert_buffer(
  }
 
 
@@ -55,6 +55,10 @@ index 7fcfb45..13e942d 100644
 +    ngx_chain_t             *cl = u->bufs_in;
 +    ngx_chain_t            **ll = NULL;
 +
++    if (cl == NULL) {
++        return;
++    }
++
 +    if (cl->next) {
 +        ll = &cl->next;
 +    }
@@ -77,6 +81,17 @@ index 7fcfb45..13e942d 100644
 +}
 +
 +
++void
++ngx_stream_lua_ffi_socket_tcp_reset_read_buf(ngx_stream_lua_request_t *r,
++    ngx_stream_lua_socket_tcp_upstream_t *u)
++{
++    ngx_stream_lua_ctx_t    *ctx;
++
++    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
++    ngx_stream_lua_ffi_socket_reset_buf(ctx, u);
++}
++
++
 +static int
 +ngx_stream_lua_socket_tcp_dummy_retval_handler(ngx_stream_lua_request_t *r,
 +    ngx_stream_lua_socket_tcp_upstream_t *u, lua_State *L)

t/stream/xrpc/downstream.t

@@ -299,3 +299,76 @@ stream lua tcp socket allocate new new buf of size 144
     }
 --- stream_request
 hello world
+
+
+
+=== TEST 11: read & drain & move & reset_read_buf
+--- stream_config
+    server {
+        listen 1995;
+        content_by_lua_block {
+            local sk = ngx.req.socket(true)
+            local data = sk:receiveany(128)
+            local exp = "rld"
+            if data ~= exp then
+                ngx.log(ngx.ERR, "actual: ", data, ", expected: ", exp)
+            end
+        }
+    }
+--- stream_server_config
+    content_by_lua_block {
+        local ffi = require("ffi")
+        local ds = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        ds:settimeout(5)
+
+        local us = require("resty.apisix.stream.xrpc.socket").upstream.socket()
+        us:settimeout(50)
+        assert(us:connect("127.0.0.1", 1995))
+
+        assert(ds:read(4))
+        assert(ds:drain(4))
+        ds:reset_read_buf()
+        assert(us:move(ds))
+        assert(ds:drain(3))
+        assert(us:move(ds))
+    }
+--- stream_request
+hello world
+
+
+
+=== TEST 12: move & reset_read_buf + read over buffer in the middle
+--- stream_config
+    server {
+        listen 1995;
+        content_by_lua_block {
+            local sk = ngx.req.socket(true)
+            local data = sk:receive(9 * 8)
+            local exp = ("123456789"):rep(8)
+            if data ~= exp then
+                ngx.log(ngx.ERR, "actual: ", data, ", expected: ", exp)
+            end
+        }
+    }
+--- stream_server_config
+    lua_socket_buffer_size 128;
+    content_by_lua_block {
+        local ffi = require("ffi")
+        local sk = require("resty.apisix.stream.xrpc.socket").downstream.socket()
+        sk:settimeout(5)
+
+        local us = require("resty.apisix.stream.xrpc.socket").upstream.socket()
+        us:settimeout(50)
+        assert(us:connect("127.0.0.1", 1995))
+
+        local len = 9 * 8
+        local p = assert(sk:read(len))
+        local len = 9 * 16
+        local p = assert(sk:read(len))
+        sk:reset_read_buf()
+        local len = 9 * 8
+        local p = assert(sk:read(len))
+        assert(us:move(sk))
+    }
+--- stream_request eval
+"123456789" x 32