From: Bozhan Liang <lbz@rivensbane.com>
Date: Fri, 13 Dec 2024 09:13:46 +0000 (+0800)
Subject: Add GeoIPQueryAttribute to LUA
X-Git-Tag: dnsdist-2.0.0-alpha1~142^2~2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=efbcb1d270ebb781a557297a49c10707eda7757d;p=thirdparty%2Fpdns.git

Add GeoIPQueryAttribute to LUA
---

diff --git a/pdns/lua-record.cc b/pdns/lua-record.cc
index 53709db61d..f450cc0a0c 100644
--- a/pdns/lua-record.cc
+++ b/pdns/lua-record.cc
@@ -1370,6 +1370,7 @@ static void setupLuaRecords(LuaContext& lua) // NOLINT(readability-function-cogn
   lua.writeFunction("geoiplookup", [](const string &ip, const GeoIPInterface::GeoIPQueryAttribute attr) {
     return getGeo(ip, attr);
   });
+  lua.writeVariable("GeoIPQueryAttribute", std::unordered_map<std::string, int>{{"ASn", GeoIPInterface::GeoIPQueryAttribute::ASn}, {"City", GeoIPInterface::GeoIPQueryAttribute::City}, {"Continent", GeoIPInterface::GeoIPQueryAttribute::Continent}, {"Country", GeoIPInterface::GeoIPQueryAttribute::Country}, {"Country2", GeoIPInterface::GeoIPQueryAttribute::Country2}, {"Name", GeoIPInterface::GeoIPQueryAttribute::Name}, {"Region", GeoIPInterface::GeoIPQueryAttribute::Region}, {"Location", GeoIPInterface::GeoIPQueryAttribute::Location}});
 
   typedef const boost::variant<string,vector<pair<int,string> > > combovar_t;
 
diff --git a/regression-tests.auth-py/test_LuaRecords.py b/regression-tests.auth-py/test_LuaRecords.py
index 6e08b957d4..3448f22eef 100644
--- a/regression-tests.auth-py/test_LuaRecords.py
+++ b/regression-tests.auth-py/test_LuaRecords.py
@@ -175,6 +175,10 @@ lookmeup         IN           A  192.0.2.5
 dblookup         IN    LUA    A  "dblookup('lookmeup.example.org', pdns.A)[1]"
 
 whitespace       IN    LUA    TXT "'foo" "bar'"
+geoipqueryattribute IN LUA    TXT ("string.format('%d %d %d %d %d %d %d',"
+                                   "GeoIPQueryAttribute.ASn, GeoIPQueryAttribute.City, GeoIPQueryAttribute.Continent,"
+                                   "GeoIPQueryAttribute.Country, GeoIPQueryAttribute.Country2, GeoIPQueryAttribute.Name,"
+                                   "GeoIPQueryAttribute.Region, GeoIPQueryAttribute.Location)")
         """
     }
     _web_rrsets = []
@@ -1177,6 +1181,22 @@ class TestLuaRecords(BaseLuaTest):
         self.assertRcodeEqual(res, dns.rcode.NOERROR)
         self.assertEqual(res.answer, response.answer)
 
+    def testGeoIPQueryAttribute(self):
+        """
+        Test GeoIPQueryAttribute enum
+        """
+        name = 'geoipqueryattribute.example.org.'
+
+        query = dns.message.make_query(name, 'TXT')
+
+        response = dns.message.make_response(query)
+
+        response.answer.append(dns.rrset.from_text(name, 0, dns.rdataclass.IN, dns.rdatatype.TXT, '"0 1 2 3 4 5 6"'))
+
+        res = self.sendUDPQuery(query)
+        self.assertRcodeEqual(res, dns.rcode.NOERROR)
+        self.assertEqual(res.answer, response.answer)
+
 
 class TestLuaRecordsShared(TestLuaRecords):
     _config_template = """