feat: add UDP reqsock peek

README.md

@@ -413,6 +413,7 @@ reqsock:peek
 Peeks into the [preread](https://nginx.org/en/docs/stream/stream_processing.html#preread_phase)
 buffer that contains downstream data sent by the client without consuming them.
 That is, data returned by this API will still be forwarded upstream in later phases.
+It works for both TCP and UDP streams with identical semantics.
 
 This function takes a single required argument, `size`, which is the number of bytes to be peeked.
 Repeated calls to this function always returns data from the beginning of the preread buffer.

src/ngx_stream_lua_socket_udp.c

@@ -41,6 +41,8 @@ static int ngx_stream_lua_socket_udp_setpeername(lua_State *L);
 static int ngx_stream_lua_socket_udp_send(lua_State *L);
 static int ngx_stream_lua_socket_udp_receive(lua_State *L);
 static int ngx_stream_lua_socket_udp_settimeout(lua_State *L);
+static int ngx_stream_lua_req_socket_udp_peek(lua_State *L);
+static ngx_int_t ngx_stream_lua_req_socket_udp_peek_resume(ngx_stream_lua_request_t *r);
 static void ngx_stream_lua_socket_udp_finalize(ngx_stream_lua_request_t *r,
     ngx_stream_lua_socket_udp_upstream_t *u);
 static int ngx_stream_lua_socket_udp_upstream_destroy(lua_State *L);
@@ -140,14 +142,17 @@ ngx_stream_lua_inject_socket_udp_api(ngx_log_t *log, lua_State *L)
     /* udp downstream socket object metatable */
     lua_pushlightuserdata(L, ngx_stream_lua_lightudata_mask(
                           socket_udp_raw_req_socket_metatable_key));
-    lua_createtable(L, 0 /* narr */, 4 /* nrec */);
+    lua_createtable(L, 0 /* narr */, 5 /* nrec */);
 
     lua_pushcfunction(L, ngx_stream_lua_socket_udp_send);
     lua_setfield(L, -2, "send");
 
     lua_pushcfunction(L, ngx_stream_lua_socket_udp_receive);
     lua_setfield(L, -2, "receive");
 
+    lua_pushcfunction(L, ngx_stream_lua_req_socket_udp_peek);
+    lua_setfield(L, -2, "peek");
+
     lua_pushcfunction(L, ngx_stream_lua_socket_udp_settimeout);
     lua_setfield(L, -2, "settimeout"); /* ngx socket mt */
 
@@ -1114,6 +1119,8 @@ ngx_stream_lua_socket_udp_receive(lua_State *L)
                    "lua udp socket receive buffer size: %uz", u->recv_buf_size);
 
     if (u->raw_downstream) {
+        u->read_consumed = 1;
+
         if (ngx_buf_size(r->connection->buffer) > 0) {
             /* we still have unread data */
             u->received = ngx_min((size_t) ngx_buf_size(r->connection->buffer),
@@ -1233,6 +1240,174 @@ ngx_stream_lua_socket_udp_settimeout(lua_State *L)
 }
 
 
+static int
+ngx_stream_lua_req_socket_udp_peek(lua_State *L)
+{
+    ngx_stream_lua_request_t            *r;
+    ngx_connection_t                    *c;
+    ngx_stream_lua_ctx_t                *ctx;
+    ngx_stream_lua_loc_conf_t           *llcf;
+    ngx_stream_lua_co_ctx_t             *coctx;
+    int                                  n;
+    lua_Integer                          bytes;
+    size_t                               size;
+
+    ngx_stream_lua_socket_udp_upstream_t        *u;
+
+    r = ngx_stream_lua_get_req(L);
+    if (r == NULL) {
+        return luaL_error(L, "no request found");
+    }
+
+    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
+    ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_PREREAD);
+
+    n = lua_gettop(L);
+    if (n != 2) {
+        return luaL_error(L, "expecting 2 arguments (including the object), but got %d", n);
+    }
+
+    luaL_checktype(L, 1, LUA_TTABLE);
+
+    lua_rawgeti(L, 1, SOCKET_CTX_INDEX);
+    u = lua_touserdata(L, -1);
+
+    if (u == NULL) {
+        llcf = ngx_stream_lua_get_module_loc_conf(r, ngx_stream_lua_module);
+
+        if (llcf->log_socket_errors) {
+            ngx_log_error(NGX_LOG_ERR, r->connection->log, 0,
+                          "attempt to peek data on a closed socket: u:%p", u);
+        }
+
+        lua_pushnil(L);
+        lua_pushliteral(L, "closed");
+        return 2;
+    }
+
+    if (u->read_consumed) {
+        return luaL_error(L, "attempt to peek on a consumed socket");
+    }
+
+    if (u->waiting) {
+        lua_pushnil(L);
+        lua_pushliteral(L, "socket busy");
+        return 2;
+    }
+
+    c = u->udp_connection.connection;
+
+    if (u->request != r) {
+        return luaL_error(L, "bad request");
+    }
+
+    if (!lua_isnumber(L, 2)) {
+        return luaL_error(L, "argument must be a number");
+    }
+
+    bytes = lua_tointeger(L, 2);
+    if (bytes < 0) {
+        return luaL_argerror(L, 2, "bytes can not be negative");
+    }
+
+    if (bytes == 0) {
+        lua_pushliteral(L, "");
+        return 1;
+    }
+
+    u->length = (size_t) bytes;
+
+    if (c->buffer != NULL) {
+        size = c->buffer->last - c->buffer->pos;
+
+        if (size >= u->length) {
+            lua_pushlstring(L, (char *) c->buffer->pos, u->length);
+            return 1;
+        }
+    }
+
+    coctx = ctx->cur_co_ctx;
+
+    ngx_stream_lua_cleanup_pending_operation(coctx);
+    coctx->cleanup = ngx_stream_lua_udp_socket_cleanup;
+    coctx->data = u;
+
+    ctx->downstream = u;
+    ctx->resume_handler = ngx_stream_lua_req_socket_udp_peek_resume;
+    ctx->peek_needs_more_data = 1;
+    u->co_ctx = coctx;
+    u->waiting = 1;
+
+    return lua_yield(L, 0);
+}
+
+
+static ngx_int_t
+ngx_stream_lua_req_socket_udp_peek_resume(ngx_stream_lua_request_t *r)
+{
+    lua_State                           *vm;
+    ngx_int_t                            rc;
+    ngx_uint_t                           nreqs;
+    ngx_connection_t                    *c;
+    ngx_stream_lua_ctx_t                *ctx;
+    size_t                               size;
+
+    ngx_stream_lua_socket_udp_upstream_t        *u;
+
+    ngx_log_debug0(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
+                   "stream lua udp socket resuming peek");
+
+    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
+    if (ctx == NULL) {
+        return NGX_ERROR;
+    }
+
+    u = ctx->downstream;
+    c = r->connection;
+    vm = ngx_stream_lua_get_lua_vm(r, ctx);
+    nreqs = c->requests;
+
+    size = c->buffer->last - c->buffer->pos;
+
+    if (size < u->length) {
+        ngx_log_debug0(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
+                       "lua peek does not have enough data, returning NGX_AGAIN");
+
+        return ngx_stream_lua_run_posted_threads(c, vm, r, ctx, nreqs);
+    }
+
+    ctx->resume_handler = ngx_stream_lua_wev_handler;
+    r->connection->read->handler = ngx_stream_lua_request_handler;
+
+    lua_pushlstring(u->co_ctx->co, (char *) c->buffer->pos, u->length);
+
+    u->co_ctx->cleanup = NULL;
+    ctx->cur_co_ctx = u->co_ctx;
+    u->co_ctx = NULL;
+    ctx->peek_needs_more_data = 0;
+    u->waiting = 0;
+
+    ngx_log_debug0(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
+                   "lua udp operation done, resuming lua thread");
+
+    rc = ngx_stream_lua_run_thread(vm, r, ctx, 1);
+
+    ngx_log_debug1(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
+                   "lua run thread returned %d", rc);
+
+    if (rc == NGX_AGAIN) {
+        return ngx_stream_lua_run_posted_threads(c, vm, r, ctx, nreqs);
+    }
+
+    if (rc == NGX_DONE) {
+        ngx_stream_lua_finalize_request(r, NGX_DONE);
+        return ngx_stream_lua_run_posted_threads(c, vm, r, ctx, nreqs);
+    }
+
+    return rc;
+}
+
+
 static void
 ngx_stream_lua_socket_udp_finalize(ngx_stream_lua_request_t *r,
     ngx_stream_lua_socket_udp_upstream_t *u)

src/ngx_stream_lua_socket_udp.h

@@ -62,12 +62,14 @@ struct ngx_stream_lua_socket_udp_upstream_s {
     ngx_err_t                                socket_errno;
     size_t                                   received; /* for receive */
     size_t                                   recv_buf_size;
+    size_t                                   length; /* for peek */
 
     ngx_stream_lua_co_ctx_t                 *co_ctx;
 
     unsigned                                 waiting:1;
 
     unsigned                                 raw_downstream:1;
+    unsigned                                 read_consumed:1;
 };
 
 

t/023-preread/req-socket-udp.t

@@ -0,0 +1,119 @@
+use Test::Nginx::Socket::Lua::Stream;
+
+repeat_each(2);
+
+plan tests => repeat_each() * (blocks() * 3 - 2);
+
+no_long_string();
+#no_diff();
+#log_level 'warn';
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: udp peek preserves datagram
+--- stream_server_config
+    set $peek "";
+    preread_by_lua_block {
+        local sock = assert(ngx.req.socket())
+        local data = assert(sock:peek(8))
+        ngx.var.peek = ngx.encode_base64(data)
+        ngx.say(ngx.var.peek)
+    }
+    proxy_pass 127.0.0.1:$TEST_NGINX_RAND_PORT_1 udp;
+--- stream_config
+server {
+    listen 127.0.0.1:$TEST_NGINX_RAND_PORT_1 udp;
+
+    content_by_lua_block {
+        local sock = assert(ngx.req.socket())
+        local data = assert(sock:receive())
+        ngx.log(ngx.DEBUG, "upstream received: ", data)
+        ngx.say("done")
+    }
+}
+--- stream_request chop
+hello world
+--- stream_response
+aGVsbG8gd28=
+done
+--- error_log
+upstream received: hello world
+--- no_error_log
+[error]
+
+=== TEST 2: double peek
+--- stream_server_config
+    preread_by_lua_block {
+        local sock = assert(ngx.req.socket())
+        local data = assert(sock:peek(4))
+        ngx.say(ngx.encode_base64(data))
+        data = assert(sock:peek(4))
+        ngx.say(ngx.encode_base64(data))
+    }
+    return done;
+--- stream_request chop
+hello world
+--- stream_response
+aGVsbA==
+aGVsbA==
+--- no_error_log
+[error]
+
+=== TEST 3: peek after receive
+--- stream_server_config
+    preread_by_lua_block {
+        local sock = assert(ngx.req.socket())
+        local data = assert(sock:receive())
+        ngx.say("received: ", data)
+        ngx.flush(true)
+        sock:peek(1)
+    }
+    return done;
+--- stream_request chop
+hello world
+--- stream_response
+received: hello world
+--- error_log
+attempt to peek on a consumed socket
+--- no_error_log
+[warn]
+
+=== TEST 4: peek timed out
+--- stream_server_config
+    preread_timeout 100ms;
+    preread_by_lua_block {
+        local sock = assert(ngx.req.socket())
+        local data = assert(sock:peek(5))
+        ngx.say("received: ", data)
+        ngx.flush(true)
+        sock:peek(12)
+    }
+    return done;
+--- stream_request chop
+hello world
+--- stream_response
+received: hello
+--- error_log
+finalize stream session: 200
+--- no_error_log
+[warn]
+
+=== TEST 5: preread buffer full
+--- stream_server_config
+    preread_buffer_size 10;
+    preread_by_lua_block {
+        local sock = assert(ngx.req.socket())
+        local data = assert(sock:peek(5))
+        ngx.say("received: ", data)
+        sock:peek(11)
+    }
+    return done;
+--- stream_request chop
+hello world
+--- error_log
+preread buffer full while prereading client data
+finalize stream session: 400
+--- no_error_log
+[warn]