feat: patch UDP socket to resolve host by ourselves (#4473)

Signed-off-by: spacewander <spacewanderlzx@gmail.com>

apisix/patch.lua

@@ -15,13 +15,15 @@
 -- limitations under the License.
 --
 local require = require
+require("resty.dns.resolver") -- preload dns resolver to prevent recursive patch
 local ipmatcher = require("resty.ipmatcher")
 local socket = require("socket")
 local unix_socket = require("socket.unix")
 local ssl = require("ssl")
 local get_phase = ngx.get_phase
 local ngx_socket = ngx.socket
 local original_tcp = ngx.socket.tcp
+local original_udp = ngx.socket.udp
 local concat_tab = table.concat
 local new_tab = require("table.new")
 local log = ngx.log
@@ -84,6 +86,43 @@ do
 end
 
 
+local patch_udp_socket
+do
+    local old_udp_sock_setpeername
+
+    local function new_udp_sock_setpeername(sock, host, port)
+        local core_str = require("apisix.core.string")
+        local resolver = require("apisix.core.resolver")
+
+        if host then
+            if core_str.has_prefix(host, "unix:") then
+                return old_udp_sock_setpeername(sock, host)
+            end
+
+            if not ipmatcher.parse_ipv4(host) and not ipmatcher.parse_ipv6(host) then
+                local err
+                host, err = resolver.parse_domain(host)
+                if not host then
+                    return nil, "failed to parse domain: " .. err
+                end
+            end
+        end
+
+        return old_udp_sock_setpeername(sock, host, port)
+    end
+
+
+    function patch_udp_socket(sock)
+        if not old_udp_sock_setpeername then
+            old_udp_sock_setpeername = sock.setpeername
+        end
+
+        sock.setpeername = new_udp_sock_setpeername
+        return sock
+    end
+end
+
+
 local function flatten(args)
     local buf = new_tab(#args, 0)
     for i, v in ipairs(args) do
@@ -278,6 +317,10 @@ function _M.patch()
 
         return luasocket_tcp()
     end
+
+    ngx_socket.udp = function ()
+        return patch_udp_socket(original_udp())
+    end
 end
 
 

t/misc/patch.t

@@ -184,3 +184,44 @@ m
 
 --- no_error_log
 [error]
+
+
+
+=== TEST 6: resolve host by ourselves (UDP)
+--- yaml_config
+apisix:
+  node_listen: 1984
+  enable_resolv_search_opt: true
+--- config
+    location /t {
+        content_by_lua_block {
+            local sock = ngx.socket.udp()
+            local res, err = sock:setpeername("apisix", 80)
+            if not res then
+                ngx.log(ngx.ERR, err)
+            end
+        }
+    }
+
+
+
+=== TEST 7: ensure our patch works with unix socket
+--- stream_server_config
+    content_by_lua_block {
+    }
+--- stream_config
+    server {
+        listen unix:$TEST_NGINX_HTML_DIR/nginx.sock;
+        content_by_lua_block {
+        }
+    }
+--- config
+    location /t {
+        content_by_lua_block {
+            local sock = ngx.socket.udp()
+            local res, err = sock:setpeername("unix:$TEST_NGINX_HTML_DIR/nginx.sock")
+            if not res then
+                ngx.log(ngx.ERR, err)
+            end
+        }
+    }