From 937f58338e849b407448deb16c23cd3a439ad0e5 Mon Sep 17 00:00:00 2001
From: Wim <wim@42.be>
Date: Sat, 15 May 2021 20:53:44 +0200
Subject: [PATCH] Support multiple ip addresses for dnsdist-resolver lua script

---
 dockerdata/dnsdist-resolver.lua | 115 +++++++++++++++++++++-----------
 1 file changed, 75 insertions(+), 40 deletions(-)

diff --git a/dockerdata/dnsdist-resolver.lua b/dockerdata/dnsdist-resolver.lua
index 965e6b92a6..89e24b319e 100644
--- a/dockerdata/dnsdist-resolver.lua
+++ b/dockerdata/dnsdist-resolver.lua
@@ -1,6 +1,5 @@
 -- testing oneliner:
 -- resolver = require 'dnsdist-resolver' resolver.maintenance() resolver.servers['www.7bits.nl']={pool='blabla'} resolver.maintenance() os.execute('sleep 3') resolver.maintenance() showServers() resolver.servers['www.7bits.nl']=nil resolver.maintenance() os.execute('sleep 3') resolver.maintenance() showServers()
-
 local _M = {}
 
 -- these are the servers we want - somebody should populate it
@@ -16,102 +15,138 @@ _M.verbose = false
 -- key = name
 -- value = {address, serverObject} (should make these named members)
 local ourservers = {}
+local ourcount = {}
 
 local resolverpipe = io.popen('/usr/local/bin/dnsdist-resolver', 'w')
 
 local function tablecopy(t)
     local t2 = {}
-    for k, v in pairs(t)
-    do
+    for k, v in pairs(t) do
         t2[k] = v
     end
     return t2
 end
 
+local function has_value(tab, val)
+    for _, value in ipairs(tab) do
+        if value == val then
+            return true
+        end
+    end
+
+    return false
+end
+
 local function removeServer(name)
+    if _M.verbose then
+        infolog("removing " .. name)
+    end
+
     rmServer(ourservers[name][2])
     ourservers[name] = nil
 end
 
 local function setServer(name, ip)
     -- adds a server or changes its IP
-    local existing = ourservers[name]
-    if existing ~= nil
-    then
+    local existing = ourservers[name .. "#" .. ip]
+    if existing ~= nil then
         if _M.verbose then
-          infolog(string.format("existing[1] [%s] == ip [%s] ??", existing[1], ip))
+            infolog(string.format("existing[1] [%s] == ip [%s] ??", existing[1], ip))
         end
         -- it exists, check IP
-        if existing[1] == ip
-        then
+        if existing[1] == ip then
             -- IP is correct, done!
             return
         else
             -- IP is wrong, drop and re-add it
-            removeServer(name)
+            removeServer(name .. "#" .. ip)
         end
     end
 
     -- it does not exist, let's add it
     local settings = tablecopy(_M.servers[name])
-    settings.name = name
-    -- TODO: we only take the first IP
+    settings.name = name .. "#" .. ip
     settings.address = ip
-    ourservers[name] = {ip, newServer(settings)}
+    ourservers[name .. "#" .. ip] = {ip, newServer(settings)}
 end
 
 function _M.maintenance()
     -- TODO: only do this if the list has changed
     -- TODO: check return values
-    for k, v in pairs(_M.servers)
-    do
-        resolverpipe:write(k..' ')
+    for k in pairs(_M.servers) do
+        resolverpipe:write(k .. ' ')
     end
     resolverpipe:write('\n')
     resolverpipe:flush()
 
     -- TODO: maybe this failure should be quiet for the first X seconds?
     local ret, resout = pcall(loadfile, '/tmp/dnsdist-resolver.out')
-    if not ret
-    then
+    if not ret then
         error(resout)
     end
 
     -- on purpose no pcall, an error here is a bug
     resout = resout()
 
+    local activeservers = {}
     -- check for servers removed by controller
-    for name, v in pairs(ourservers)
-    do
-        if _M.servers[name] == nil
-        then
+    for ourserver in pairs(ourservers) do
+        activeservers[ourserver] = false
+        for server in pairs(_M.servers) do
+            -- use plain because a dash is a special character ..
+            if ourserver:find(server, 1, true) == 1 then
+                activeservers[ourserver] = true
+            end
+        end
+    end
+
+    for name, active in pairs(activeservers) do
+        if active == false then
             removeServer(name)
         end
     end
 
-    for name, ips in pairs(resout)
-    do
+    for name, ips in pairs(resout) do
         if _M.verbose then
-          infolog("name="..name)
-          for _, ip in ipairs(ips)
-          do
-              infolog("  ip="..ip)
-          end
+            infolog("name=" .. name)
+            for _, ip in ipairs(ips) do
+                infolog("  ip=" .. ip)
+            end
+        end
+
+        -- init our current count
+        if ourcount[name] == nil then
+            ourcount[name] = #ips
+        end
+
+        -- increase our current count if necessary
+        if #ips > ourcount[name] then
+            ourcount[name] = #ips
+            if _M.verbose then
+                infolog("increasing count to " .. ourcount[name] .. " for " .. name)
+            end
         end
 
-        if #ips == 0
-        then
-            -- server has left the building
-            if ourservers[name] ~= nil
-            then
-                removeServer(name)
+        -- remove servers when we've lost ips
+        if #ips < ourcount[name] then
+            for ourserver, server in pairs(ourservers) do
+                -- check if we match the prefix and the ip is gone
+                if ourserver:find(name, 1, true) == 1 and has_value(ips, server[1]) == false then
+                    ourcount[name] = #ips
+                    if _M.verbose then
+                        infolog("ip address not found anymore " .. server[1])
+                        infolog("decreasing count to " .. ourcount[name])
+                    end
+                    removeServer(ourserver)
+                end
             end
         else
-            -- it has IPs
-            if _M.servers[name] ~= nil
-            then
-                -- we want this server
-                setServer(name, ips[1])
+            for _, ip in ipairs(ips) do
+                -- it has IPs
+                if _M.servers[name] ~= nil then
+                    -- we want this server
+                    setServer(name, ip)
+                end
             end
         end
     end
-- 
2.47.2