From 74a2ea87bc4332e35697e42ab944d7f5b7a3698a Mon Sep 17 00:00:00 2001
From: Peter van Dijk <peter.van.dijk@powerdns.com>
Date: Fri, 17 Dec 2021 19:27:27 +0100
Subject: [PATCH] dnsdist: add newThread() function (experimental, subject to
 change)

---
 pdns/dnsdist-lua.cc                  | 39 ++++++++++++++++++++++++++
 regression-tests.dnsdist/test_Lua.py | 42 ++++++++++++++++++++++++++++
 tasks.py                             |  1 +
 3 files changed, 82 insertions(+)
 create mode 100644 regression-tests.dnsdist/test_Lua.py

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index deb547e78e..0fd6bea61c 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -264,6 +264,40 @@ void checkParameterBound(const std::string& parameter, uint64_t value, size_t ma
   }
 }
 
+static void LuaThread(const std::string code)
+{
+  LuaContext l;
+  // submitToMainThread is camelcased, threadmessage is not.
+  // This follows our tradition of hooks we call being lowercased but functions the user can call being camelcased.
+  l.writeFunction("submitToMainThread", [](std::string cmd, std::unordered_map<std::string, std::string> data) {
+    auto lua = g_lua.lock();
+    // maybe offer more than `void`
+    auto func = lua->readVariable<boost::optional<std::function<void(std::string cmd, std::unordered_map<std::string, std::string> data)>>>("threadmessage");
+    if (func) {
+      func.get()(cmd, data);
+    }
+    else {
+      errlog("Lua thread called submitToMainThread but no threadmessage receiver is defined");
+    }
+  });
+
+  // function threadmessage(cmd, data) print("got thread data:", cmd) for k,v in pairs(data) do print(k,v) end end
+
+  for (;;) {
+    try {
+      l.executeCode(code);
+      errlog("Lua thread exited, restarting in 5 seconds");
+    }
+    catch (const std::exception& e) {
+      errlog("Lua thread crashed, restarting in 5 seconds: %s", e.what());
+    }
+    catch (...) {
+      errlog("Lua thread crashed, restarting in 5 seconds");
+    }
+    sleep(5);
+  }
+}
+
 static void setupLuaConfig(LuaContext& luaCtx, bool client, bool configCheck)
 {
   typedef std::unordered_map<std::string, boost::variant<bool, std::string, vector<pair<int, std::string>>, DownstreamState::checkfunc_t>> newserver_t;
@@ -2759,6 +2793,11 @@ static void setupLuaConfig(LuaContext& luaCtx, bool client, bool configCheck)
     }
   });
 #endif /* HAVE_LIBSSL */
+
+  luaCtx.writeFunction("newThread", [](const std::string& code) {
+    std::thread newThread(LuaThread, code);
+    newThread.detach();
+  });
 }
 
 vector<std::function<void(void)>> setupLua(LuaContext& luaCtx, bool client, bool configCheck, const std::string& config)
diff --git a/regression-tests.dnsdist/test_Lua.py b/regression-tests.dnsdist/test_Lua.py
new file mode 100644
index 0000000000..656c287ec9
--- /dev/null
+++ b/regression-tests.dnsdist/test_Lua.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+
+import base64
+import time
+import unittest
+from dnsdisttests import DNSDistTest
+
+class TestLuaThread(DNSDistTest):
+    _consoleKey = DNSDistTest.generateConsoleKey()
+    _consoleKeyB64 = base64.b64encode(_consoleKey).decode('ascii')
+
+    _config_params = ['_consoleKeyB64', '_consolePort']
+    _config_template = """
+    setKey("%s")
+    controlSocket("127.0.0.1:%s")
+
+    counter = 0
+    function threadmessage(cmd, data)
+        print("counter says", cmd, data.i)
+        counter = tonumber(data.i)
+    end
+
+    newThread([==[
+      local socket = require'socket'
+      local i=1
+      while true
+      do
+          socket.sleep(1)
+          submitToMainThread("setCounter", {i=i})
+          i = i + 1
+      end
+    ]==])
+    """
+
+    def testLuaThreadCounter(self):
+        """
+        LuaThread: Test the lua newThread interface
+        """
+        count1 = self.sendConsoleCommand('counter')
+        time.sleep(3)
+        count2 = self.sendConsoleCommand('counter')
+        self.assertTrue(count2 > count1)
diff --git a/tasks.py b/tasks.py
index 15a012b1bb..1b45d7b74f 100644
--- a/tasks.py
+++ b/tasks.py
@@ -209,6 +209,7 @@ def install_dnsdist_test_deps(c): # FIXME: rename this, we do way more than apt-
               libssl-dev \
               libsystemd0 \
               libsodium23 \
+              lua-socket \
               patch \
               protobuf-compiler \
               python3-venv snmpd prometheus')
-- 
2.47.2