Add Lua function to pick records via name hash

This adds a Lua function to return a record based on a weighted hash
of the DNS record name.  One use case here is to consistently return
the same IP address for a particular cache server based on what
subdomain is requesting the data.

diff --git a/.github/actions/spell-check/expect.txt b/.github/actions/spell-check/expect.txt
index e88447ab9c12a89b7695fbd4e3f970360cdefc1c..e8db2f2ba8abd38b9ac2a4dbc43a9d053ac46e61 100644 (file)
--- a/.github/actions/spell-check/expect.txt
+++ b/.github/actions/spell-check/expect.txt
@@ -1008,6 +1008,7 @@ pickhashed
 pickrandom
 pickrandomsample
 pickwhashed
+picknamehashed
 pickwrandom
 piddir
 pidfile

diff --git a/docs/lua-records/functions.rst b/docs/lua-records/functions.rst
index 6902d761ad2c2e8c2782301f04894f96091228f9..a2aec80da4390bc6b338785ed10f877a81ad5197 100644 (file)
--- a/docs/lua-records/functions.rst
+++ b/docs/lua-records/functions.rst
@@ -228,6 +228,29 @@ Record creation functions
                                             "        {100, "198.51.100.5"}             "
                                             "})                                        ")
 
+.. function:: picknamehashed(values)
+
+  Based on the hash of the DNS record name, returns a string from the list
+  supplied, as weighted by the various ``weight`` parameters.
+  Performs no uptime checking.
+
+  :param values: table of weight, string (such as IPv4 or IPv6 address).
+
+  This allows basic persistent load balancing across a number of backends.  It means that
+  test.example.com will always resolve to the same IP, but test2.example.com may go elsewhere.
+
+  This works similar to round-robin load balanacing, but has the advantage of making traffic
+  for the same domain always end up on the same server which can help cache hit rates.
+
+  This function also works for CNAME or TXT records.
+
+  An example::
+
+    mydomain.example.com    IN    LUA    A ("picknamehashed({                             "
+                                            "        {15,  "192.0.2.1"},               "
+                                            "        {100, "198.51.100.5"}             "
+                                            "})                                        ")
+
 
 .. function:: pickwrandom(values)
 

diff --git a/pdns/lua-record.cc b/pdns/lua-record.cc
index 343781b0631c3e0780b937771d9a38154a149fd2..c84c2f870d0724e46f70c511301ea84e19f67093 100644 (file)
--- a/pdns/lua-record.cc
+++ b/pdns/lua-record.cc
@@ -382,6 +382,30 @@ static T pickWeightedHashed(const ComboAddress& bestwho, vector< pair<int, T> >&
   return p->second;
 }
 
+template <typename T>
+static T pickWeightedNameHashed(const DNSName& dnsname, vector< pair<int, T> >& items)
+{
+  if (items.empty()) {
+    throw std::invalid_argument("The items list cannot be empty");
+  }
+  int sum=0;
+  vector< pair<int, T> > pick;
+  pick.reserve(items.size());
+
+  for(auto& i : items) {
+    sum += i.first;
+    pick.push_back({sum, i.second});
+  }
+
+  if (sum == 0) {
+    throw std::invalid_argument("The sum of items cannot be zero");
+  }
+
+  int r = dnsname.hash() % sum;
+  auto p = upper_bound(pick.begin(), pick.end(), r, [](int rarg, const typename decltype(pick)::value_type& a) { return rarg < a.first; });
+  return p->second;
+}
+
 template <typename T>
 static vector<T> pickRandomSample(int n, const vector<T>& items)
 {
@@ -996,6 +1020,20 @@ static void setupLuaRecords(LuaContext& lua) // NOLINT(readability-function-cogn
       return pickWeightedHashed<string>(s_lua_record_ctx->bestwho, items);
     });
 
+  /*
+   * Based on the hash of the record name, return an IP address from the list
+   * supplied, as weighted by the various `weight` parameters
+   * @example picknamehashed({ {15, '1.2.3.4'}, {50, '5.4.3.2'} })
+   */
+  lua.writeFunction("picknamehashed", [](std::unordered_map<int, wiplist_t > ips) {
+      vector< pair<int, string> > items;
+
+      items.reserve(ips.size());
+      for(auto& i : ips)
+        items.emplace_back(atoi(i.second[1].c_str()), i.second[2]);
+
+      return pickWeightedNameHashed<string>(s_lua_record_ctx->qname, items);
+    });
 
   lua.writeFunction("pickclosest", [](const iplist_t& ips) {
       vector<ComboAddress> conv = convComboAddressList(ips);

diff --git a/regression-tests.auth-py/test_LuaRecords.py b/regression-tests.auth-py/test_LuaRecords.py
index 3b1cc721dbeeecb427d7cef8f6ee6daed09bb7be..1164af67d515e0ae7cd46669422cf0845450fb90 100644 (file)
--- a/regression-tests.auth-py/test_LuaRecords.py
+++ b/regression-tests.auth-py/test_LuaRecords.py
@@ -68,6 +68,7 @@ hashed.example.org.          3600 IN LUA  A     "pickhashed({{ '1.2.3.4', '4.3.2
 hashed-v6.example.org.       3600 IN LUA  AAAA  "pickhashed({{ '2001:db8:a0b:12f0::1', 'fe80::2a1:9bff:fe9b:f268' }})"
 hashed-txt.example.org.      3600 IN LUA  TXT   "pickhashed({{ 'bob', 'alice' }})"
 whashed.example.org.         3600 IN LUA  A     "pickwhashed({{ {{15, '1.2.3.4'}}, {{42, '4.3.2.1'}} }})"
+*.namehashed.example.org.    3600 IN LUA  A     "picknamehashed({{ {{15, '1.2.3.4'}}, {{42, '4.3.2.1'}} }})"
 whashed-txt.example.org.     3600 IN LUA  TXT   "pickwhashed({{ {{15, 'bob'}}, {{42, 'alice'}} }})"
 rand.example.org.            3600 IN LUA  A     "pickrandom({{'{prefix}.101', '{prefix}.102'}})"
 rand-txt.example.org.        3600 IN LUA  TXT   "pickrandom({{ 'bob', 'alice' }})"
@@ -773,6 +774,7 @@ createforward6.example.org.                 3600 IN NS   ns2.example.org.
         self.assertRcodeEqual(res, dns.rcode.SERVFAIL)
         self.assertAnswerEmpty(res)
 
+
     def testWHashed(self):
         """
         Basic pickwhashed() test with a set of A records
@@ -790,6 +792,33 @@ createforward6.example.org.                 3600 IN NS   ns2.example.org.
             self.assertRcodeEqual(res, dns.rcode.NOERROR)
             self.assertRRsetInAnswer(res, first.answer[0])
 
+
+    def testNamehashed(self):
+        """
+        Basic picknamehashed() test with a set of A records
+        As the name is hashed, we should always get the same IP back for the same record name.
+        """
+
+        queries = [
+            {
+                'expected': dns.rrset.from_text('test.namehashed.example.org.', 0,
+                                       dns.rdataclass.IN, 'A',
+                                       '1.2.3.4'),
+                'query': dns.message.make_query('test.namehashed.example.org', 'A')
+            },
+            {
+                'expected': dns.rrset.from_text('test2.namehashed.example.org.', 0,
+                                       dns.rdataclass.IN, 'A',
+                                       '4.3.2.1'),
+                'query': dns.message.make_query('test2.namehashed.example.org', 'A')
+            }
+        ]
+        for query in queries :
+            res = self.sendUDPQuery(query['query'])
+            self.assertRcodeEqual(res, dns.rcode.NOERROR)
+            self.assertRRsetInAnswer(res, query['expected'])
+
+
     def testWHashedTxt(self):
         """
         Basic pickwhashed() test with a set of TXT records