From c58dfe7e2762cdeab0829d0727677610f222d971 Mon Sep 17 00:00:00 2001 From: aaurizon <750048+aaurizon@users.noreply.github.com> Date: Tue, 26 Aug 2025 22:53:16 +0200 Subject: [PATCH 1/2] feat: add udp reqsock peek --- README.md | 1 + src/ngx_stream_lua_socket_udp.c | 177 +++++++++++++++++++++++++++++++- src/ngx_stream_lua_socket_udp.h | 2 + t/023-preread/req-socket-udp.t | 119 +++++++++++++++++++++ 4 files changed, 298 insertions(+), 1 deletion(-) create mode 100644 t/023-preread/req-socket-udp.t diff --git a/README.md b/README.md index 3cdd5060..bcc3e337 100644 --- a/README.md +++ b/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. diff --git a/src/ngx_stream_lua_socket_udp.c b/src/ngx_stream_lua_socket_udp.c index d2612f59..7bddd6cc 100644 --- a/src/ngx_stream_lua_socket_udp.c +++ b/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,7 +142,7 @@ 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"); @@ -148,6 +150,9 @@ ngx_stream_lua_inject_socket_udp_api(ngx_log_t *log, lua_State *L) 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_coctx_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) diff --git a/src/ngx_stream_lua_socket_udp.h b/src/ngx_stream_lua_socket_udp.h index 8f5b9e95..89e7c4f7 100644 --- a/src/ngx_stream_lua_socket_udp.h +++ b/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; }; diff --git a/t/023-preread/req-socket-udp.t b/t/023-preread/req-socket-udp.t new file mode 100644 index 00000000..09d13cb5 --- /dev/null +++ b/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] From 36aaa3f486b9406776f6829719ad6995d120e383 Mon Sep 17 00:00:00 2001 From: aaurizon <750048+aaurizon@users.noreply.github.com> Date: Wed, 27 Aug 2025 00:12:34 +0200 Subject: [PATCH 2/2] fix: use UDP socket cleanup for peek --- src/ngx_stream_lua_socket_udp.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ngx_stream_lua_socket_udp.c b/src/ngx_stream_lua_socket_udp.c index 7bddd6cc..1b5a4dc4 100644 --- a/src/ngx_stream_lua_socket_udp.c +++ b/src/ngx_stream_lua_socket_udp.c @@ -1329,7 +1329,7 @@ ngx_stream_lua_req_socket_udp_peek(lua_State *L) coctx = ctx->cur_co_ctx; ngx_stream_lua_cleanup_pending_operation(coctx); - coctx->cleanup = ngx_stream_lua_coctx_cleanup; + coctx->cleanup = ngx_stream_lua_udp_socket_cleanup; coctx->data = u; ctx->downstream = u;