Merge pull request #7594 from rgacogne/dnsdist-set-rules

[thirdparty/pdns.git] / regression-tests.dnsdist / test_EdnsClientSubnet.py

diff --git a/regression-tests.dnsdist/test_EdnsClientSubnet.py b/regression-tests.dnsdist/test_EdnsClientSubnet.py
index d537e86e9b8fd0378bf6d8460320b4439741edd5..6adb863ed012c16a06ec8a2cfcdf4ad1945652f5 100644 (file)
--- a/regression-tests.dnsdist/test_EdnsClientSubnet.py
+++ b/regression-tests.dnsdist/test_EdnsClientSubnet.py
@@ -3,6 +3,7 @@ import dns
 import clientsubnetoption
 import cookiesoption
 from dnsdisttests import DNSDistTest
+from datetime import datetime, timedelta
 
 class TestEdnsClientSubnetNoOverride(DNSDistTest):
     """
@@ -224,7 +225,7 @@ class TestEdnsClientSubnetNoOverride(DNSDistTest):
         query = dns.message.make_query(name, 'A', 'IN', use_edns=True, payload=4096)
         expectedQuery = dns.message.make_query(name, 'A', 'IN', use_edns=True, payload=4096, options=[ecso])
         response = dns.message.make_response(expectedQuery)
-        ecoResponse = cookiesoption.CookiesOption('deadbeef', 'deadbeef')
+        ecoResponse = cookiesoption.CookiesOption(b'deadbeef', b'deadbeef')
         ecsoResponse = clientsubnetoption.ClientSubnetOption('127.0.0.1', 24, scope=24)
         response.use_edns(edns=True, payload=4096, options=[ecoResponse, ecsoResponse])
         expectedResponse = dns.message.make_response(query)
@@ -267,7 +268,7 @@ class TestEdnsClientSubnetNoOverride(DNSDistTest):
         query = dns.message.make_query(name, 'A', 'IN', use_edns=True, payload=4096)
         expectedQuery = dns.message.make_query(name, 'A', 'IN', use_edns=True, payload=4096, options=[ecso])
         response = dns.message.make_response(expectedQuery)
-        ecoResponse = cookiesoption.CookiesOption('deadbeef', 'deadbeef')
+        ecoResponse = cookiesoption.CookiesOption(b'deadbeef', b'deadbeef')
         ecsoResponse = clientsubnetoption.ClientSubnetOption('127.0.0.1', 24, scope=24)
         response.use_edns(edns=True, payload=4096, options=[ecsoResponse, ecoResponse])
         expectedResponse = dns.message.make_response(query)
@@ -310,7 +311,7 @@ class TestEdnsClientSubnetNoOverride(DNSDistTest):
         query = dns.message.make_query(name, 'A', 'IN', use_edns=True, payload=4096)
         expectedQuery = dns.message.make_query(name, 'A', 'IN', use_edns=True, payload=4096, options=[ecso])
         response = dns.message.make_response(expectedQuery)
-        ecoResponse = cookiesoption.CookiesOption('deadbeef', 'deadbeef')
+        ecoResponse = cookiesoption.CookiesOption(b'deadbeef', b'deadbeef')
         ecsoResponse = clientsubnetoption.ClientSubnetOption('127.0.0.1', 24, scope=24)
         response.use_edns(edns=True, payload=4096, options=[ecoResponse, ecsoResponse, ecoResponse])
         expectedResponse = dns.message.make_response(query)
@@ -564,7 +565,7 @@ class TestECSDisabledByRuleOrLua(DNSDistTest):
         dq.useECS = false
         return DNSAction.None, ""
     end
-    addLuaAction("disabledvialua.ecsrules.tests.powerdns.com.", disableECSViaLua)
+    addAction("disabledvialua.ecsrules.tests.powerdns.com.", LuaAction(disableECSViaLua))
     """
 
     def testWithECSNotDisabled(self):
@@ -672,7 +673,7 @@ class TestECSOverrideSetByRuleOrLua(DNSDistTest):
         dq.ecsOverride = true
         return DNSAction.None, ""
     end
-    addLuaAction("overriddenvialua.ecsrules.tests.powerdns.com.", overrideECSViaLua)
+    addAction("overriddenvialua.ecsrules.tests.powerdns.com.", LuaAction(overrideECSViaLua))
     """
 
     def testWithECSOverrideNotSet(self):
@@ -786,7 +787,7 @@ class TestECSPrefixLengthSetByRuleOrLua(DNSDistTest):
         dq.ecsPrefixLength = 32
         return DNSAction.None, ""
     end
-    addLuaAction("overriddenprefixlengthvialua.ecsrules.tests.powerdns.com.", overrideECSPrefixLengthViaLua)
+    addAction("overriddenprefixlengthvialua.ecsrules.tests.powerdns.com.", LuaAction(overrideECSPrefixLengthViaLua))
     """
 
     def testWithECSPrefixLengthNotOverridden(self):
@@ -885,3 +886,83 @@ class TestECSPrefixLengthSetByRuleOrLua(DNSDistTest):
         receivedQuery.id = expectedQuery.id
         self.checkQueryEDNSWithECS(expectedQuery, receivedQuery)
         self.checkResponseNoEDNS(expectedResponse, receivedResponse)
+
+class TestECSPrefixSetByRule(DNSDistTest):
+    """
+    dnsdist is configured to set the EDNS0 Client Subnet
+    option for incoming queries to the actual source IP,
+    but we override it for some queries via SetECSAction().
+    """
+
+    _config_template = """
+    setECSOverride(false)
+    setECSSourcePrefixV4(32)
+    setECSSourcePrefixV6(128)
+    newServer{address="127.0.0.1:%s", useClientSubnet=true}
+    addAction(makeRule("setecsaction.ecsrules.tests.powerdns.com."), SetECSAction("192.0.2.1/32"))
+    """
+
+    def testWithRegularECS(self):
+        """
+        ECS Prefix: not set
+        """
+        name = 'notsetecsaction.ecsrules.tests.powerdns.com.'
+        ecso = clientsubnetoption.ClientSubnetOption('127.0.0.1', 32)
+        query = dns.message.make_query(name, 'A', 'IN')
+        expectedQuery = dns.message.make_query(name, 'A', 'IN', use_edns=True, options=[ecso], payload=512)
+        response = dns.message.make_response(query)
+        response.use_edns(edns=True, payload=4096, options=[ecso])
+        rrset = dns.rrset.from_text(name,
+                                    3600,
+                                    dns.rdataclass.IN,
+                                    dns.rdatatype.A,
+                                    '127.0.0.1')
+        response.answer.append(rrset)
+        expectedResponse = dns.message.make_response(query)
+        expectedResponse.answer.append(rrset)
+
+        (receivedQuery, receivedResponse) = self.sendUDPQuery(query, response)
+        self.assertTrue(receivedQuery)
+        self.assertTrue(receivedResponse)
+        receivedQuery.id = expectedQuery.id
+        self.checkQueryEDNSWithECS(expectedQuery, receivedQuery)
+        self.checkResponseNoEDNS(expectedResponse, receivedResponse)
+
+        (receivedQuery, receivedResponse) = self.sendTCPQuery(query, response)
+        self.assertTrue(receivedQuery)
+        self.assertTrue(receivedResponse)
+        receivedQuery.id = expectedQuery.id
+        self.checkQueryEDNSWithECS(expectedQuery, receivedQuery)
+        self.checkResponseNoEDNS(expectedResponse, receivedResponse)
+
+    def testWithECSSetByRule(self):
+        """
+        ECS Prefix: set with SetECSAction
+        """
+        name = 'setecsaction.ecsrules.tests.powerdns.com.'
+        ecso = clientsubnetoption.ClientSubnetOption('192.0.2.1', 32)
+        query = dns.message.make_query(name, 'A', 'IN')
+        expectedQuery = dns.message.make_query(name, 'A', 'IN', use_edns=True, options=[ecso], payload=512)
+        response = dns.message.make_response(expectedQuery)
+        rrset = dns.rrset.from_text(name,
+                                    3600,
+                                    dns.rdataclass.IN,
+                                    dns.rdatatype.A,
+                                    '127.0.0.1')
+        response.answer.append(rrset)
+        expectedResponse = dns.message.make_response(query)
+        expectedResponse.answer.append(rrset)
+
+        (receivedQuery, receivedResponse) = self.sendUDPQuery(query, response)
+        self.assertTrue(receivedQuery)
+        self.assertTrue(receivedResponse)
+        receivedQuery.id = expectedQuery.id
+        self.checkQueryEDNSWithECS(expectedQuery, receivedQuery)
+        self.checkResponseNoEDNS(expectedResponse, receivedResponse)
+
+        (receivedQuery, receivedResponse) = self.sendTCPQuery(query, response)
+        self.assertTrue(receivedQuery)
+        self.assertTrue(receivedResponse)
+        receivedQuery.id = expectedQuery.id
+        self.checkQueryEDNSWithECS(expectedQuery, receivedQuery)
+        self.checkResponseNoEDNS(expectedResponse, receivedResponse)