From a18044c504bc946e8f4ea8f5046dc88f133b66af Mon Sep 17 00:00:00 2001
From: Paul Kulchenko <paul@kulchenko.com>
Date: Fri, 8 Jul 2022 07:17:25 -0700
Subject: [PATCH] Add OnServerListen hook to configure listen() (#459)

---
 tool/net/help.txt  |  8 +++++++-
 tool/net/redbean.c | 33 ++++++++++++++++++++++++++++-----
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/tool/net/help.txt b/tool/net/help.txt
index 0c2669b7d..852f3b0d7 100644
--- a/tool/net/help.txt
+++ b/tool/net/help.txt
@@ -501,7 +501,7 @@ HOOKS
   OnClientConnection(ip:int,port:int,serverip:int,serverport:int) → bool
           If this function is defined it'll be called from the main process
           each time redbean accepts a new client connection. If it returns
-          true then redbean will close the connection without calling fork.
+          `true`, redbean will close the connection without calling fork.
 
   OnProcessCreate(pid:int,ip:int,port:int,serverip:int,serverport:int)
           If this function is defined it'll be called from the main process
@@ -517,6 +517,12 @@ HOOKS
           each time redbean reaps a child connection process using wait4().
           This won't be called in uniprocess mode.
 
+  OnServerListen(socketdescriptor:int,serverip:int,serverport:int) → bool
+          If this function is defined it'll be called from the main process
+          before redbean starts listening on a port. This hook can be used
+          to modify socket configuration to set `SO_REUSEPORT`, for example.
+          If it returns `true`, redbean will not listen to that ip/port.
+
   OnServerStart()
           If this function is defined it'll be called from the main process
           right before the main event loop starts.
diff --git a/tool/net/redbean.c b/tool/net/redbean.c
index c991ccf66..c852d6f74 100644
--- a/tool/net/redbean.c
+++ b/tool/net/redbean.c
@@ -1080,8 +1080,8 @@ static bool LuaEvalFile(const char *path) {
 }
 
 static bool LuaOnClientConnection(void) {
+  bool dropit = false;
 #ifndef STATIC
-  bool dropit;
   uint32_t ip, serverip;
   uint16_t port, serverport;
   lua_State *L = GL;
@@ -1096,14 +1096,11 @@ static bool LuaOnClientConnection(void) {
     dropit = lua_toboolean(L, -1);
   } else {
     LogLuaError("OnClientConnection", lua_tostring(L, -1));
-    dropit = false;
   }
   lua_pop(L, 1);  // pop result or error
   AssertLuaStackIsAt(L, 0);
-  return dropit;
-#else
-  return false;
 #endif
+  return dropit;
 }
 
 static void LuaOnProcessCreate(int pid) {
@@ -1127,6 +1124,25 @@ static void LuaOnProcessCreate(int pid) {
 #endif
 }
 
+static bool LuaOnServerListen(int fd, uint32_t ip, uint16_t port) {
+  bool nouse = false;
+#ifndef STATIC
+  lua_State *L = GL;
+  lua_getglobal(L, "OnServerListen");
+  lua_pushinteger(L, fd);
+  lua_pushinteger(L, ip);
+  lua_pushinteger(L, port);
+  if (LuaCallWithTrace(L, 3, 1, NULL) == LUA_OK) {
+    nouse = lua_toboolean(L, -1);
+  } else {
+    LogLuaError("OnServerListen", lua_tostring(L, -1));
+  }
+  lua_pop(L, 1);  // pop result or error
+  AssertLuaStackIsAt(L, 0);
+#endif
+  return nouse;
+}
+
 static void LuaOnProcessDestroy(int pid) {
 #ifndef STATIC
   lua_State *L = GL;
@@ -6913,6 +6929,7 @@ static void Listen(void) {
   char ipbuf[16];
   size_t i, j, n;
   uint32_t ip, port, addrsize, *ifs, *ifp;
+  bool hasonserverlisten = IsHookDefined("OnServerListen");
   if (!ports.n) {
     ProgramPort(8080);
   }
@@ -6939,6 +6956,12 @@ static void Listen(void) {
                                         IPPROTO_TCP, true, &timeout)) == -1) {
         DIEF("(srvr) socket: %m");
       }
+      if (hasonserverlisten &&
+          LuaOnServerListen(servers.p[n].fd, ips.p[i], ports.p[j])) {
+        n--;  // skip this server instance
+        continue;
+      }
+
       if (bind(servers.p[n].fd, &servers.p[n].addr,
                sizeof(servers.p[n].addr)) == -1) {
         DIEF("(srvr) bind error: %m: %hhu.%hhu.%hhu.%hhu:%hu", ips.p[i] >> 24,