Fixes to require, loadlib and UdpServer

docs/src/includes/_globals.md

@@ -2459,8 +2459,9 @@ The callback signature is nil on_data(string response, string error)
 
 #### [UdpServer](#UdpServer) udp_listen(string host, int port, function cb)
 
-Start an UDP server on specified address and run callback when data arrives. Return a string from the callback to reply. Requires unsafe mode.
-The server will be closed once the handle is released.
+Start an UDP server on specified address and run callback when data arrives. Set port to 0 to use a random ephemeral port. Return a string from the callback to reply.
+The server will be closed lazily by garbage collection once the handle is released, or immediately by calling close(). Requires unsafe mode.
+<br/>The callback signature is optional<string> on_message(string msg, string src)
 
 ### udp_send
 

docs/src/includes/_types.md

@@ -1057,6 +1057,11 @@ tuple<[Vec2](#Vec2), [Vec2](#Vec2), [Vec2](#Vec2)> | [split()](https://git
 
 Type | Name | Description
 ---- | ---- | -----------
+nil | [close()](https://github.com/spelunky-fyi/overlunky/search?l=Lua&q=close) | Closes the server.
+bool | [open()](https://github.com/spelunky-fyi/overlunky/search?l=Lua&q=open) | Returns true if the port was opened successfully and the server hasn't been closed yet.
+string | [error()](https://github.com/spelunky-fyi/overlunky/search?l=Lua&q=error) | Returns a string explaining the last error, at least if open() returned false.
+string | [host](https://github.com/spelunky-fyi/overlunky/search?l=Lua&q=host) | 
+int | [port](https://github.com/spelunky-fyi/overlunky/search?l=Lua&q=port) | 
 
 ### Vec2
 

examples/udp.lua

@@ -0,0 +1,82 @@
+meta = {
+    name = "UDP echo example",
+    description = "Opens UDP server on a random port and echoes sent messages, closes the server on 'gg'",
+    author = "Dregu",
+    unsafe = true
+}
+
+options = {
+    host = "127.0.0.1",
+    port = 0
+}
+
+function init()
+    deinit()
+
+    server = {
+        -- Save messages to queue instead of printing directly, cause that doesn't work from the server thread
+        queue = {},
+
+        -- Open server on an ephemeral random port, push messages to queue and echo back to sender
+        socket = udp_listen(options.host, options.port, function(msg, src)
+            msg = msg:gsub("%s*$", "")
+            table.insert(server.queue, { msg = msg, src = src })
+            if msg == "gg" then
+                return "bye!\n"
+            else
+                return "echo: " .. msg .. "\n"
+            end
+        end)
+    }
+
+    -- If port was opened successfully, start checking the message queue
+    if server.socket:open() then
+        print(F "Listening on {server.socket.port}, please send some UDP datagrams or 'gg' to close")
+        server.inter = set_global_interval(function()
+            for _, msg in pairs(server.queue) do
+                print(F "Received: '{msg.msg}' from {msg.src}")
+                if msg.msg == "gg" then
+                    server.socket:close()
+                    clear_callback()
+                    print("Server is now closed, have a nice day")
+                end
+            end
+            server.queue = {}
+        end, 1)
+    else
+        print(F "Failed to open server: {server.socket:error()}")
+    end
+end
+
+function deinit()
+    if server then
+        server.socket:close()
+        if server.inter then clear_callback(server.inter) end
+        server = nil
+    end
+end
+
+set_callback(init, ON.LOAD)
+set_callback(init, ON.SCRIPT_ENABLE)
+set_callback(deinit, ON.SCRIPT_DISABLE)
+
+register_option_callback("x", nil, function(ctx)
+    options.host = ctx:win_input_text("Host", options.host)
+    options.port = ctx:win_input_int("Port", options.port)
+    if options.port < 0 then options.port = 0 end
+    if options.port > 65535 then options.port = 65535 end
+    if ctx:win_button("Start server") then init() end
+    ctx:win_inline()
+    if ctx:win_button("Stop server") then deinit() end
+    if server then
+        ctx:win_text(server.socket:error())
+    end
+    if server and server.socket:open() then
+        ctx:win_text(F "Listening on {server.socket.host}:{server.socket.port}\nTry sending something with udp_send:")
+        if ctx:win_button("Send 'Hello World!'") then udp_send(server.socket.host, server.socket.port, "Hello World!") end
+        ctx:win_inline()
+        if ctx:win_button("Send 'gg'") then udp_send(server.socket.host, server.socket.port, "gg") end
+    else
+        ctx:win_text("Stopped")
+    end
+end)

src/game_api/script/lua_require.cpp

@@ -25,24 +25,30 @@ void register_custom_require(sol::state& lua)
     lua.add_package_loader(custom_loader);
 
     lua["__require"] = lua["require"];
+    lua["__loadlib"] = lua["package"]["loadlib"];
 
     /// Custom implementation to trick Lua into allowing to `require 'lib.module'` more than once given it was called from a different source
     lua["require"] = custom_require;
+
+    lua["package"]["loadlib"] = custom_loadlib;
 }
 sol::object custom_require(std::string path)
 {
+    static sol::state& lua = get_lua_vm();
+
+    if (path == "io" || path == "os" || path == "math" || path == "string" || path == "table" || path == "coroutine" || path == "package")
+        return lua[path];
+
     // Turn module into a real path
     {
-        if (path.ends_with(".lua"))
+        if (path.ends_with(".lua") || path.ends_with(".dll"))
         {
             path = path.substr(0, path.size() - 4);
         }
         std::replace(path.begin(), path.end(), '.', '/');
         std::replace(path.begin(), path.end(), ':', '.');
     }
 
-    static sol::state& lua = get_lua_vm();
-
     // Could be preloaded by some unsafe script, which can only be fetched by unsafe scripts
     auto backend = LuaBackend::get_calling_backend();
     const bool unsafe = backend->get_unsafe();
@@ -174,25 +180,112 @@ return info.short_src, info.source
 
     return std::move(res).value_or(sol::nil);
 }
+sol::object custom_loadlib(std::string path, std::string func)
+{
+    static sol::state& lua = get_lua_vm();
+    auto backend = LuaBackend::get_calling_backend();
+    const bool unsafe = backend->get_unsafe();
+
+    // Walk up the stack until we find an _ENV that is not global, then grab the source from that stack index
+    auto [short_source, source] = []() -> std::pair<std::string_view, std::string_view>
+    {
+        return lua.safe_script(R"(
+-- Not available in Lua 5.2+
+local getfenv = getfenv or function(f)
+    f = (type(f) == 'function' and f or debug.getinfo(f + 1, 'f').func)
+    local name, val
+    local up = 0
+    repeat
+        up = up + 1
+        name, val = debug.getupvalue(f, up)
+    until name == '_ENV' or name == nil
+    return val
+end
+
+local env
+local up = 1
+repeat
+    up = up + 1
+    env = getfenv(up)
+until env ~= _G and env ~= nil
+
+local info = debug.getinfo(up)
+return info.short_src, info.source
+)");
+    }();
+
+    if (short_source.starts_with("[string"))
+    {
+        source = backend->get_root();
+    }
+    else
+    {
+        auto last_slash = source.find_last_of("/\\");
+        if (last_slash != std::string::npos)
+        {
+            source = source.substr(0, last_slash);
+        }
+        if (source.starts_with('@'))
+        {
+            source = source.substr(1);
+        }
+    }
+
+    namespace fs = std::filesystem;
+
+    auto loadlib_if_exists = [&](fs::path _path, std::string _func) -> std::optional<sol::object>
+    {
+        if (unsafe && std::filesystem::exists(_path))
+            return lua["__loadlib"](_path.string(), _func);
+        return std::nullopt;
+    };
+
+    /// Valid unsafe options:
+    // path/to/calling/script/path/to/lib.dll
+    // path/to/calling/mod/path/to/lib.dll
+    // path/to/lib.dll
+
+    auto res = loadlib_if_exists(std::filesystem::path(source) / path, func);
+
+    if (!res && source != backend->get_root())
+    {
+        res = loadlib_if_exists(std::filesystem::path(backend->get_root()) / path, func);
+    }
+    if (!res)
+    {
+        res = loadlib_if_exists(std::filesystem::path(path), func);
+    }
+
+    return std::move(res).value_or(sol::nil);
+}
 int custom_loader(lua_State* L)
 {
     std::string path = sol::stack::get<std::string>(L, 1);
     std::replace(path.begin(), path.end(), '.', '/');
     std::replace(path.begin(), path.end(), ':', '.');
     auto backend = LuaBackend::get_calling_backend();
 
-    auto try_load = [&](std::string& _path, std::string_view ext)
+    auto try_load = [&](std::string _path, std::string_view ext)
     {
+        if (ext == ".dll")
+        {
+            auto root = std::string(backend->get_root());
+            auto module = _path.substr(root.size() + 1);
+            auto func = "luaopen_" + module;
+            std::replace(func.begin(), func.end(), '/', '_');
+            std::replace(func.begin(), func.end(), '.', '_');
+            sol::stack::push(L, backend->lua["__loadlib"](_path, func));
+            return true;
+        }
         _path += ext;
-        const auto res = luaL_loadfilex(L, _path.c_str(), "bt");
+        const auto res = luaL_loadfile(L, _path.c_str());
         if (res == LUA_OK)
         {
             backend->lua.push();
-            // The first up-value is always the _ENV of the chunk
+            //  The first up-value is always the _ENV of the chunk
             lua_setupvalue(L, -2, 1);
             return true;
         }
-        _path = _path.substr(0, _path.size() - 1);
         return false;
     };
 

src/game_api/script/lua_require.hpp

@@ -11,4 +11,5 @@ class state;
 // This implementation makes the loaded chunk inherit the env from the loading chunk
 void register_custom_require(sol::state& lua);
 sol::object custom_require(std::string path);
+sol::object custom_loadlib(std::string path, std::string func);
 int custom_loader(struct lua_State* L);

src/game_api/script/usertypes/socket_lua.cpp

@@ -13,15 +13,15 @@ namespace NSocket
 void register_usertypes(sol::state& lua)
 {
     lua.new_usertype<UdpServer>(
-        "UdpServer",
-        sol::no_constructor);
-    /// Start an UDP server on specified address and run callback when data arrives. Return a string from the callback to reply. Requires unsafe mode.
-    /// The server will be closed once the handle is released.
-    lua["udp_listen"] = [](std::string host, in_port_t port, sol::function cb) -> UdpServer*
+        "UdpServer", sol::no_constructor, "close", &UdpServer::close, "open", &UdpServer::open, "error", &UdpServer::error, "host", sol::readonly(&UdpServer::host), "port", sol::readonly(&UdpServer::port));
+
+    /// Start an UDP server on specified address and run callback when data arrives. Set port to 0 to use a random ephemeral port. Return a string from the callback to reply.
+    /// The server will be closed lazily by garbage collection once the handle is released, or immediately by calling close(). Requires unsafe mode.
+    /// <br/>The callback signature is optional<string> on_message(string msg, string src)
+    // lua["udp_listen"] = [](std::string host, in_port_t port, sol::function cb) -> UdpServer*
+    lua["udp_listen"] = [](std::string host, in_port_t port, sol::function cb)
     {
-        // TODO: change the return to std::unique_ptr after fixing the dead lock with the destructor
-        UdpServer* server = new UdpServer(std::move(host), std::move(port), make_safe_cb<UdpServer::SocketCb>(std::move(cb)));
-        return server;
+        return std::unique_ptr<UdpServer>{new UdpServer(std::move(host), std::move(port), make_safe_cb<UdpServer::SocketCb>(std::move(cb)))};
     };
 
     /// Send data to specified UDP address. Requires unsafe mode.

src/game_api/socket.cpp

@@ -2,6 +2,7 @@
 
 #include <Windows.h>             // for GetModuleHandleA, GetProcAddress
 #include <algorithm>             // for max
+#include <chrono>                // for chrono
 #include <detours.h>             // for DetourAttach, DetourTransactionBegin
 #include <exception>             // for exception
 #include <new>                   // for operator new
@@ -56,39 +57,57 @@ void dump_network()
 void udp_data(sockpp::udp_socket socket, UdpServer* server)
 {
     ssize_t n;
-    char buf[500];
+    static char buf[32768];
     sockpp::inet_address src;
-    while (server->kill_thr.test(std::memory_order_acquire) && (n = socket.recv_from(buf, sizeof(buf), &src)) > 0)
+    while (server->kill_thr.test(std::memory_order_acquire) && socket.is_open())
     {
-        std::optional<std::string> ret = server->cb(std::string(buf, n));
-        if (ret)
+        while ((n = socket.recv_from(buf, sizeof(buf), &src)) > 0)
         {
-            socket.send_to(ret.value(), src);
+            std::optional<std::string> ret = server->cb(std::string(buf, n), src.to_string());
+            if (ret)
+            {
+                socket.send_to(ret.value(), src);
+            }
         }
+        std::this_thread::sleep_for(std::chrono::microseconds(1));
     }
 }
 
 UdpServer::UdpServer(std::string host_, in_port_t port_, std::function<SocketCb> cb_)
     : host(host_), port(port_), cb(cb_)
 {
-    sock.bind(sockpp::inet_address(host, port));
-    kill_thr.test_and_set();
-    thr = std::thread(udp_data, std::move(sock), this);
+    if (sock.bind(sockpp::inet_address(host, port)))
+    {
+        const auto addr = (sockpp::inet_address)sock.address();
+        port = addr.port();
+        is_opened = true;
+        sock.set_non_blocking();
+        kill_thr.test_and_set();
+        thr = std::thread(udp_data, std::move(sock), this);
+    }
 }
-void UdpServer::clear() // TODO: fix and expose: this and the destructor causes deadlock
+void UdpServer::close()
 {
+    is_closed = true;
     kill_thr.clear(std::memory_order_release);
-    thr.join();
+    sock.close();
+    sock.shutdown();
+    if (thr.joinable())
+        thr.join();
+}
+bool UdpServer::open()
+{
+    return is_opened && !is_closed && sock.last_error() == 0;
+}
+std::string UdpServer::error()
+{
+    auto err = sock.error_str(sock.last_error());
+    return err.substr(0, err.size() - 2);
 }
 UdpServer::~UdpServer()
 {
-    if (thr.joinable())
-    {
-        kill_thr.clear(std::memory_order_release);
-        thr.join();
-    }
+    close();
 }
-
 bool http_get(const char* sURL, std::string& out, std::string& err)
 {
     const int BUFFER_SIZE = 32768;