From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Thu, 30 Nov 2023 16:09:20 +0000 (+0100)
Subject: dnsdist: Add `PayloadSizeRule`
X-Git-Tag: dnsdist-1.9.0-alpha4~4^2~6
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=b7ed5e7bf15f93cb745de840828d8800ee5ee455;p=thirdparty%2Fpdns.git

dnsdist: Add `PayloadSizeRule`

Adding the ability to match on the size of the DNS payload.
---

diff --git a/pdns/dnsdist-lua-rules.cc b/pdns/dnsdist-lua-rules.cc
index 9cbeec7529..64784987e7 100644
--- a/pdns/dnsdist-lua-rules.cc
+++ b/pdns/dnsdist-lua-rules.cc
@@ -778,9 +778,13 @@ void setupLuaRules(LuaContext& luaCtx)
 
   luaCtx.writeFunction("LuaFFIPerThreadRule", [](const std::string& code) {
       return std::shared_ptr<DNSRule>(new LuaFFIPerThreadRule(code));
-  });
+    });
 
   luaCtx.writeFunction("ProxyProtocolValueRule", [](uint8_t type, boost::optional<std::string> value) {
       return std::shared_ptr<DNSRule>(new ProxyProtocolValueRule(type, std::move(value)));
     });
+
+  luaCtx.writeFunction("PayloadSizeRule", [](const std::string& comparison, uint16_t size) {
+    return std::shared_ptr<DNSRule>(new PayloadSizeRule(comparison, size));
+    });
 }
diff --git a/pdns/dnsdistdist/dnsdist-rules.hh b/pdns/dnsdistdist/dnsdist-rules.hh
index c9b932c3dd..d0f2fcb70d 100644
--- a/pdns/dnsdistdist/dnsdist-rules.hh
+++ b/pdns/dnsdistdist/dnsdist-rules.hh
@@ -1350,3 +1350,66 @@ private:
   boost::optional<std::string> d_value;
   uint8_t d_type;
 };
+
+class PayloadSizeRule : public DNSRule
+{
+  enum class Comparisons : uint8_t { equal, greater, greaterOrEqual, smaller, smallerOrEqual };
+public:
+  PayloadSizeRule(const std::string& comparison, uint16_t size): d_size(size)
+  {
+    if (comparison == "equal") {
+      d_comparison = Comparisons::equal;
+    }
+    else if (comparison == "greater") {
+      d_comparison = Comparisons::greater;
+    }
+    else if (comparison == "greaterOrEqual") {
+      d_comparison = Comparisons::greaterOrEqual;
+    }
+    else if (comparison == "smaller") {
+      d_comparison = Comparisons::smaller;
+    }
+    else if (comparison == "smallerOrEqual") {
+      d_comparison = Comparisons::smallerOrEqual;
+    }
+    else {
+      throw std::runtime_error("Unsupported comparison '" + comparison + "'");
+    }
+  }
+
+  bool matches(const DNSQuestion* dq) const override
+  {
+    const auto size = dq->getData().size();
+
+    switch (d_comparison) {
+    case Comparisons::equal:
+      return size == d_size;
+    case Comparisons::greater:
+      return size > d_size;
+    case Comparisons::greaterOrEqual:
+      return size >= d_size;
+    case Comparisons::smaller:
+      return size < d_size;
+    case Comparisons::smallerOrEqual:
+      return size <= d_size;
+    default:
+      return false;
+    }
+  }
+
+  string toString() const override
+  {
+    static const std::array<const std::string, 5> comparisonStr{
+      "equal to" ,
+      "greater than",
+      "equal to or greater than",
+      "smaller than",
+      "equal to or smaller than"
+    };
+    return "payload size is " + comparisonStr.at(static_cast<size_t>(d_comparison)) + " " + std::to_string(d_size);
+  }
+
+private:
+  uint16_t d_size;
+  Comparisons d_comparison;
+};
diff --git a/pdns/dnsdistdist/docs/rules-actions.rst b/pdns/dnsdistdist/docs/rules-actions.rst
index b4266f05aa..7bb7a31c07 100644
--- a/pdns/dnsdistdist/docs/rules-actions.rst
+++ b/pdns/dnsdistdist/docs/rules-actions.rst
@@ -726,6 +726,15 @@ These ``DNSRule``\ s be one of the following items:
 
   :param int code: The opcode to match
 
+.. function:: PayloadSizeRule(comparison, size)
+
+  .. versionadded:: 1.9.0
+
+  Matches queries or responses whose DNS payload size is equal, greater, greater or equal, smaller or smaller or equal to the specified size.
+
+  :param str comparison: The comparison operator to use. Supported values are ``equal``, ``greater``, ``greaterOrEqual``, ``smaller`` and ``smallerOrEqual``.
+  :param int size: The size to compare to.w
+
 .. function:: ProbaRule(probability)
 
   Matches queries with a given probability. 1.0 means "always"
diff --git a/regression-tests.dnsdist/test_Size.py b/regression-tests.dnsdist/test_Size.py
new file mode 100644
index 0000000000..5503587ed8
--- /dev/null
+++ b/regression-tests.dnsdist/test_Size.py
@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+import dns
+from dnsdisttests import DNSDistTest
+
+class TestSize(DNSDistTest):
+
+    _payloadSize = 49
+    _config_template = """
+    addAction(PayloadSizeRule("smaller", %d), SpoofAction("192.0.2.1"))
+    addAction(PayloadSizeRule("greater", %d), SpoofAction("192.0.2.2"))
+    addAction(PayloadSizeRule("equal", %d), SpoofAction("192.0.2.3"))
+    newServer{address="127.0.0.1:%d"}
+    """
+    _config_params = ['_payloadSize', '_payloadSize', '_payloadSize', '_testServerPort']
+
+    def testPayloadSize(self):
+        """
+        Size: Check that PayloadSizeRule works
+        """
+        name = 'payload.size.tests.powerdns.com.'
+        query = dns.message.make_query(name, 'A', 'IN')
+        query.flags &= ~dns.flags.RD
+        expectedResponse = dns.message.make_response(query)
+        rrset = dns.rrset.from_text(name,
+                                    60,
+                                    dns.rdataclass.IN,
+                                    dns.rdatatype.A,
+                                    '192.0.2.3')
+        expectedResponse.answer.append(rrset)
+
+        for method in ("sendUDPQuery", "sendTCPQuery"):
+            sender = getattr(self, method)
+            (_, receivedResponse) = sender(query, response=None, useQueue=False)
+            self.assertTrue(receivedResponse)
+            self.assertEqual(receivedResponse, expectedResponse)