feature: raw request downstream socket support for stream servers. #a290cac

dndx · web-flow · commit 0ee4e89e239c · 2018-04-14T21:03:12.000-07:00
diff --git a/README.md b/README.md
@@ -301,7 +301,13 @@ is ignored and the raw request socket is always returned. Unlike `ngx_http_lua_m
 you can still call output API functions like `ngx.say`, `ngx.print`, and `ngx.flush`
 after acquiring the raw request socket via this function.
 
-Raw request socket returned by this module will contain the following extra method:
+When stream server is in UDP mode, reading from the downstream socket returned by the
+`ngx.req.socket` call will only return the content of a single packet. Therefore
+the reading call will never block and will return `nil, "no more data"` when all the
+data from the datagram has been consumed. However, you may choose to send multiple UDP
+packets back to the client using the downstream socket.
+
+Raw TCP request socket returned by this module will contain the following extra method:
 
 tcpsock:shutdown
 ----------------
diff --git a/src/ngx_stream_lua_output.c b/src/ngx_stream_lua_output.c
@@ -63,6 +63,10 @@ ngx_stream_lua_ngx_echo(lua_State *L, unsigned newline)
         return luaL_error(L, "no request ctx found");
     }
 
+    if (r->connection->type == SOCK_DGRAM) {
+        return luaL_error(L, "API disabled in the current context");
+    }
+
     ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_CONTENT
                                  | NGX_STREAM_LUA_CONTEXT_PREREAD);
 
@@ -468,6 +472,10 @@ ngx_stream_lua_ngx_flush(lua_State *L)
         return luaL_error(L, "no request ctx found");
     }
 
+    if (r->connection->type == SOCK_DGRAM) {
+        return luaL_error(L, "API disabled in the current context");
+    }
+
     ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_CONTENT
                                  | NGX_STREAM_LUA_CONTEXT_PREREAD);
 
@@ -587,6 +595,10 @@ ngx_stream_lua_ngx_eof(lua_State *L)
         return 2;
     }
 
+    if (r->connection->type == SOCK_DGRAM) {
+        return luaL_error(L, "API disabled in the current context");
+    }
+
     ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_CONTENT
                                  | NGX_STREAM_LUA_CONTEXT_PREREAD);
 
diff --git a/src/ngx_stream_lua_socket_tcp.c b/src/ngx_stream_lua_socket_tcp.c
@@ -101,7 +101,6 @@ static int ngx_stream_lua_socket_receiveuntil_iterator(lua_State *L);
 static ngx_int_t ngx_stream_lua_socket_compile_pattern(u_char *data, size_t len,
     ngx_stream_lua_socket_compiled_pattern_t *cp, ngx_log_t *log);
 static int ngx_stream_lua_socket_cleanup_compiled_pattern(lua_State *L);
-static int ngx_stream_lua_req_socket(lua_State *L);
 static void ngx_stream_lua_req_socket_rev_handler(ngx_stream_lua_request_t *r);
 static int ngx_stream_lua_socket_tcp_getreusedtimes(lua_State *L);
 static int ngx_stream_lua_socket_tcp_setkeepalive(lua_State *L);
@@ -368,14 +367,6 @@ ngx_stream_lua_inject_socket_tcp_api(ngx_log_t *log, lua_State *L)
 }
 
 
-void
-ngx_stream_lua_inject_req_socket_api(lua_State *L)
-{
-    lua_pushcfunction(L, ngx_stream_lua_req_socket);
-    lua_setfield(L, -2, "socket");
-}
-
-
 static int
 ngx_stream_lua_socket_tcp(lua_State *L)
 {
@@ -4255,8 +4246,8 @@ ngx_stream_lua_socket_cleanup_compiled_pattern(lua_State *L)
 }
 
 
-static int
-ngx_stream_lua_req_socket(lua_State *L)
+int
+ngx_stream_lua_req_socket_tcp(lua_State *L)
 {
     int                              n, raw;
     ngx_peer_connection_t           *pc;
diff --git a/src/ngx_stream_lua_socket_tcp.h b/src/ngx_stream_lua_socket_tcp.h
@@ -148,8 +148,8 @@ typedef struct {
 
 
 void ngx_stream_lua_inject_socket_tcp_api(ngx_log_t *log, lua_State *L);
-void ngx_stream_lua_inject_req_socket_api(lua_State *L);
 void ngx_stream_lua_cleanup_conn_pools(lua_State *L);
+int ngx_stream_lua_req_socket_tcp(lua_State *L);
 
 
 #endif /* _NGX_STREAM_LUA_SOCKET_TCP_H_INCLUDED_ */
diff --git a/src/ngx_stream_lua_socket_udp.c b/src/ngx_stream_lua_socket_udp.c
@@ -73,6 +73,8 @@ enum {
 
 static char ngx_stream_lua_socket_udp_metatable_key;
 static char ngx_stream_lua_udp_udata_metatable_key;
+static char ngx_stream_lua_socket_udp_raw_req_socket_metatable_key;
+static char ngx_stream_lua_socket_udp_downstream_udata_metatable_key;
 static u_char ngx_stream_lua_socket_udp_buffer[UDP_MAX_DATAGRAM_SIZE];
 
 
@@ -84,7 +86,7 @@ ngx_stream_lua_inject_socket_udp_api(ngx_log_t *log, lua_State *L)
     lua_pushcfunction(L, ngx_stream_lua_socket_udp);
     lua_setfield(L, -2, "udp"); /* ngx socket */
 
-    /* udp socket object metatable */
+    /* udp upstream socket object metatable */
     lua_pushlightuserdata(L, &ngx_stream_lua_socket_udp_metatable_key);
     lua_createtable(L, 0 /* narr */, 6 /* nrec */);
 
@@ -108,14 +110,41 @@ ngx_stream_lua_inject_socket_udp_api(ngx_log_t *log, lua_State *L)
     lua_rawset(L, LUA_REGISTRYINDEX);
     /* }}} */
 
-    /* udp socket object metatable */
+    /* udp downstream socket object metatable */
+    lua_pushlightuserdata(L, &ngx_stream_lua_socket_udp_raw_req_socket_metatable_key);
+    lua_createtable(L, 0 /* narr */, 4 /* 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_socket_udp_settimeout);
+    lua_setfield(L, -2, "settimeout"); /* ngx socket mt */
+
+    lua_pushvalue(L, -1);
+    lua_setfield(L, -2, "__index");
+    lua_rawset(L, LUA_REGISTRYINDEX);
+    /* }}} */
+
+    /* udp upstream socket object metatable */
     lua_pushlightuserdata(L, &ngx_stream_lua_udp_udata_metatable_key);
     lua_createtable(L, 0 /* narr */, 1 /* nrec */); /* metatable */
     lua_pushcfunction(L, ngx_stream_lua_socket_udp_upstream_destroy);
     lua_setfield(L, -2, "__gc");
     lua_rawset(L, LUA_REGISTRYINDEX);
     /* }}} */
 
+    /* udp downstream socket object metatable */
+    lua_pushlightuserdata(L, &ngx_stream_lua_socket_udp_downstream_udata_metatable_key);
+    lua_createtable(L, 0 /* narr */, 1 /* nrec */); /* metatable */
+    /* share the same destructor as upstream */
+    lua_pushcfunction(L, ngx_stream_lua_socket_udp_upstream_destroy);
+    lua_setfield(L, -2, "__gc");
+    lua_rawset(L, LUA_REGISTRYINDEX);
+    /* }}} */
+
     lua_pop(L, 1);
 }
 
@@ -890,7 +919,7 @@ ngx_stream_lua_socket_udp_send(lua_State *L)
 
     dd("sending query %.*s", (int) query.len, query.data);
 
-    n = ngx_send(u->udp_connection.connection, query.data, query.len);
+    n = ngx_udp_send(u->udp_connection.connection, query.data, query.len);
 
     dd("ngx_send returns %d (query len %d)", (int) n, (int) query.len);
 
@@ -988,7 +1017,28 @@ ngx_stream_lua_socket_udp_receive(lua_State *L)
     ngx_log_debug1(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
                    "lua udp socket receive buffer size: %uz", u->recv_buf_size);
 
-    rc = ngx_stream_lua_socket_udp_read(r, u);
+    if (u->raw_downstream) {
+        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),
+                                  u->recv_buf_size);
+            ngx_memcpy(ngx_stream_lua_socket_udp_buffer,
+                       r->connection->buffer->pos, u->received);
+            r->connection->buffer->pos += u->received;
+
+            ngx_stream_lua_socket_udp_handle_success(r, u);
+
+            rc = NGX_OK;
+
+        } else {
+            lua_pushnil(L);
+            lua_pushliteral(L, "no more data");
+            return 2;
+        }
+
+    } else {
+        rc = ngx_stream_lua_socket_udp_read(r, u);
+    }
 
     if (rc == NGX_ERROR) {
         dd("read failed: %d", (int) u->ft_type);
@@ -1104,7 +1154,11 @@ ngx_stream_lua_socket_udp_finalize(ngx_stream_lua_request_t *r,
         u->resolved->ctx = NULL;
     }
 
-    if (u->udp_connection.connection) {
+    /*
+     * do not close if it is a downstream connection as that will
+     * be handled by stream subsystem itself
+     */
+    if (u->udp_connection.connection && !u->raw_downstream) {
         ngx_log_debug0(NGX_LOG_DEBUG_STREAM, r->connection->log, 0,
                        "lua close socket connection");
 
@@ -1617,4 +1671,113 @@ ngx_stream_lua_udp_socket_cleanup(void *data)
     ngx_stream_lua_socket_udp_finalize(u->request, u);
 }
 
+
+int
+ngx_stream_lua_req_socket_udp(lua_State *L)
+{
+    ngx_stream_lua_udp_connection_t                *pc;
+    ngx_stream_lua_srv_conf_t                      *lscf;
+    ngx_connection_t                               *c;
+    ngx_stream_lua_request_t                       *r;
+    ngx_stream_lua_ctx_t                           *ctx;
+
+    ngx_stream_lua_cleanup_t                       *cln;
+    ngx_stream_lua_co_ctx_t                        *coctx;
+
+    ngx_stream_lua_socket_udp_upstream_t           *u;
+
+    r = ngx_stream_lua_get_req(L);
+
+    ctx = ngx_stream_lua_get_module_ctx(r, ngx_stream_lua_module);
+    if (ctx == NULL) {
+        return luaL_error(L, "no ctx found");
+    }
+
+    ngx_stream_lua_check_context(L, ctx, NGX_STREAM_LUA_CONTEXT_CONTENT
+                                 |NGX_STREAM_LUA_CONTEXT_PREREAD);
+
+    c = r->connection;
+
+    if (c->buffered) {
+        lua_pushnil(L);
+        lua_pushliteral(L, "pending data to write");
+        return 2;
+    }
+
+    dd("ctx acquired raw req socket: %d", ctx->acquired_raw_req_socket);
+
+    if (ctx->acquired_raw_req_socket) {
+        lua_pushnil(L);
+        lua_pushliteral(L, "duplicate call");
+        return 2;
+    }
+
+    ctx->acquired_raw_req_socket = 1;
+
+    lua_createtable(L, 3 /* narr */, 1 /* nrec */); /* the object */
+    lua_pushlightuserdata(L, &ngx_stream_lua_socket_udp_raw_req_socket_metatable_key);
+    lua_rawget(L, LUA_REGISTRYINDEX);
+    lua_setmetatable(L, -2);
+
+    u = lua_newuserdata(L, sizeof(ngx_stream_lua_socket_udp_upstream_t));
+    if (u == NULL) {
+        return luaL_error(L, "no memory");
+    }
+
+#if 1
+    lua_pushlightuserdata(L, &ngx_stream_lua_socket_udp_downstream_udata_metatable_key);
+    lua_rawget(L, LUA_REGISTRYINDEX);
+    lua_setmetatable(L, -2);
+#endif
+
+    lua_rawseti(L, 1, SOCKET_CTX_INDEX);
+
+    ngx_memzero(u, sizeof(ngx_stream_lua_socket_udp_upstream_t));
+
+    u->raw_downstream = 1;
+
+    coctx = ctx->cur_co_ctx;
+
+    u->request = r;
+
+    lscf = ngx_stream_lua_get_module_srv_conf(r, ngx_stream_lua_module);
+
+    u->conf = lscf;
+
+    u->read_timeout = u->conf->read_timeout;
+
+    cln = ngx_stream_lua_cleanup_add(r, 0);
+    if (cln == NULL) {
+        u->ft_type |= NGX_STREAM_LUA_SOCKET_FT_ERROR;
+        lua_pushnil(L);
+        lua_pushliteral(L, "no memory");
+        return 2;
+    }
+
+    cln->handler = ngx_stream_lua_socket_udp_cleanup;
+    cln->data = u;
+    u->cleanup = &cln->handler;
+
+    pc = &u->udp_connection;
+    pc->log = *c->log;
+    pc->connection = c;
+
+    dd("setting data to %p", u);
+
+    coctx->data = u;
+    ctx->downstream = u;
+
+    if (c->read->timer_set) {
+        ngx_del_timer(c->read);
+    }
+
+    if (c->write->timer_set) {
+        ngx_del_timer(c->write);
+    }
+
+    lua_settop(L, 1);
+    return 1;
+}
+
+
 /* vi:set ft=c ts=4 sw=4 et fdm=marker: */
diff --git a/src/ngx_stream_lua_socket_udp.h b/src/ngx_stream_lua_socket_udp.h
@@ -42,22 +42,24 @@ struct ngx_stream_lua_socket_udp_upstream_s {
     ngx_stream_lua_request_t                *request;
     ngx_stream_lua_udp_connection_t          udp_connection;
 
-    ngx_msec_t                       read_timeout;
+    ngx_msec_t                               read_timeout;
 
     ngx_stream_upstream_resolved_t          *resolved;
 
-    ngx_uint_t                       ft_type;
-    ngx_err_t                        socket_errno;
-    size_t                           received; /* for receive */
-    size_t                           recv_buf_size;
+    ngx_uint_t                               ft_type;
+    ngx_err_t                                socket_errno;
+    size_t                                   received; /* for receive */
+    size_t                                   recv_buf_size;
 
     ngx_stream_lua_co_ctx_t                 *co_ctx;
 
-    unsigned                         waiting; /* :1 */
+    unsigned                                 waiting:1;
+    unsigned                                 raw_downstream:1;
 };
 
 
 void ngx_stream_lua_inject_socket_udp_api(ngx_log_t *log, lua_State *L);
+int ngx_stream_lua_req_socket_udp(lua_State *L);
 
 
 #endif /* _NGX_STREAM_LUA_SOCKET_UDP_H_INCLUDED_ */
diff --git a/src/ngx_stream_lua_util.c b/src/ngx_stream_lua_util.c
diff --git a/t/139-req-udp-socket.t b/t/139-req-udp-socket.t
