Merge pull request #13564 from rgacogne/ddist-payload-size-rule

dnsdist: Add `PayloadSizeRule` and `TCResponseAction`

pdns/dnsdist-console.cc

@@ -799,6 +799,7 @@ const std::vector<ConsoleKeyword> g_consoleKeywords{
   { "TagRule", true, "name [, value]", "matches if the tag named 'name' is present, with the given 'value' matching if any" },
   { "TCAction", true, "", "create answer to query with TC and RD bits set, to move to TCP" },
   { "TCPRule", true, "[tcp]", "Matches question received over TCP if tcp is true, over UDP otherwise" },
+  { "TCResponseAction", true, "", "truncate a response" },
   { "TeeAction", true, "remote [, addECS [, local]]", "send copy of query to remote, optionally adding ECS info, optionally set local address" },
   { "testCrypto", true, "", "test of the crypto all works" },
   { "TimedIPSetRule", true, "", "Create a rule which matches a set of IP addresses which expire"},

pdns/dnsdist-lua-actions.cc

@@ -513,6 +513,19 @@ class TCAction : public DNSAction
   }
 };
 
+class TCResponseAction : public DNSResponseAction
+{
+public:
+  DNSResponseAction::Action operator()(DNSResponse* dnsResponse, std::string* ruleresult) const override
+  {
+    return Action::Truncate;
+  }
+  [[nodiscard]] std::string toString() const override
+  {
+    return "tc=1 answer";
+  }
+};
+
 class LuaAction : public DNSAction
 {
 public:
@@ -2549,6 +2562,10 @@ void setupLuaActions(LuaContext& luaCtx)
       return std::shared_ptr<DNSAction>(new TCAction);
     });
 
+  luaCtx.writeFunction("TCResponseAction", []() {
+      return std::shared_ptr<DNSResponseAction>(new TCResponseAction);
+    });
+
   luaCtx.writeFunction("SetDisableValidationAction", []() {
       return std::shared_ptr<DNSAction>(new SetDisableValidationAction);
     });

pdns/dnsdist-lua-rules.cc

@@ -786,9 +786,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));
+    });
 }

pdns/dnsdist-lua-vars.cc

@@ -51,6 +51,7 @@ void setupLuaVars(LuaContext& luaCtx)
       {"Drop",         (int)DNSResponseAction::Action::Drop         },
       {"HeaderModify", (int)DNSResponseAction::Action::HeaderModify },
       {"ServFail",     (int)DNSResponseAction::Action::ServFail     },
+      {"Truncate",     (int)DNSResponseAction::Action::Truncate     },
       {"None",         (int)DNSResponseAction::Action::None         }
     });
 

pdns/dnsdist.cc

@@ -539,6 +539,18 @@ static bool applyRulesToResponse(const std::vector<DNSDistResponseRuleAction>& r
         });
         return true;
         break;
+      case DNSResponseAction::Action::Truncate:
+        if (!dr.overTCP()) {
+          dnsdist::PacketMangling::editDNSHeaderFromPacket(dr.getMutableData(), [](dnsheader& header) {
+            header.tc = true;
+            header.qr = true;
+            return true;
+          });
+          truncateTC(dr.getMutableData(), dr.getMaximumSize(), dr.ids.qname.wirelength());
+          ++dnsdist::metrics::g_stats.ruleTruncated;
+          return true;
+        }
+        break;
         /* non-terminal actions follow */
       case DNSResponseAction::Action::Delay:
         pdns::checked_stoi_into(dr.ids.delayMsec, ruleresult); // sorry

pdns/dnsdist.hh

@@ -267,7 +267,7 @@ public:
 class DNSResponseAction
 {
 public:
-  enum class Action : uint8_t { Allow, Delay, Drop, HeaderModify, ServFail, None };
+  enum class Action : uint8_t { Allow, Delay, Drop, HeaderModify, ServFail, Truncate, None };
   virtual Action operator()(DNSResponse*, string* ruleresult) const =0;
   virtual ~DNSResponseAction()
   {

pdns/dnsdistdist/dnsdist-lua-ffi.hh

@@ -48,11 +48,11 @@ struct dnsdist_ffi_dnsquestion_t
   DNSQuestion* dq{nullptr};
   ComboAddress maskedRemote;
   std::string trailingData;
-  boost::optional<std::string> result{boost::none};
-  boost::optional<std::string> httpPath{boost::none};
-  boost::optional<std::string> httpQueryString{boost::none};
-  boost::optional<std::string> httpHost{boost::none};
-  boost::optional<std::string> httpScheme{boost::none};
+  std::optional<std::string> result{std::nullopt};
+  std::optional<std::string> httpPath{std::nullopt};
+  std::optional<std::string> httpQueryString{std::nullopt};
+  std::optional<std::string> httpHost{std::nullopt};
+  std::optional<std::string> httpScheme{std::nullopt};
   std::unique_ptr<std::vector<dnsdist_ffi_ednsoption_t>> ednsOptionsVect;
   std::unique_ptr<std::vector<dnsdist_ffi_http_header_t>> httpHeadersVect;
   std::unique_ptr<std::vector<dnsdist_ffi_tag_t>> tagsVect;
@@ -78,7 +78,7 @@ struct dnsdist_ffi_dnsresponse_t
   }
 
   DNSResponse* dr{nullptr};
-  boost::optional<std::string> result{boost::none};
+  std::optional<std::string> result{std::nullopt};
 };
 
 // dnsdist_ffi_server_t is a lightuserdata

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;
+};

pdns/dnsdistdist/docs/reference/constants.rst

@@ -154,6 +154,9 @@ All named `QTypes <https://www.iana.org/assignments/dns-parameters/dns-parameter
 DNSResponseAction
 -----------------
 
+.. versionchanged:: 1.9.0
+  The ``DNSResponseAction.Truncate`` value was added.
+
 These constants represent an Action that can be returned from :func:`LuaResponseAction` functions.
 
  * ``DNSResponseAction.Allow``: let the response pass, skipping other rules
@@ -162,3 +165,4 @@ These constants represent an Action that can be returned from :func:`LuaResponse
  * ``DNSResponseAction.HeaderModify``: indicate that the query has been turned into a response
  * ``DNSResponseAction.None``: continue to the next rule
  * ``DNSResponseAction.ServFail``: return a response with a ServFail rcode
+ * ``DNSResponseAction.Truncate``: truncate the response, removing all records from the answer, authority and additional sections if any

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 fits the given comparison.
+
+  :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.
+
 .. function:: ProbaRule(probability)
 
   Matches queries with a given probability. 1.0 means "always"
@@ -1914,6 +1923,13 @@ The following actions exist.
   Before 1.7.0 this action was performed even when the query had been received over TCP, which required the use of :func:`TCPRule` to
   prevent the TC bit from being set over TCP transports.
 
+.. function:: TCResponseAction()
+
+  .. versionadded:: 1.9.0
+
+  Truncate an existing answer, to force the client to TCP. Only applied to answers that will be sent to the client over TCP.
+  In addition to the TC bit being set, all records are removed from the answer, authority and additional sections.
+
 .. function:: TeeAction(remote[, addECS[, local [, addProxyProtocol]]])
 
   .. versionchanged:: 1.8.0

pdns/dnsdistdist/test-dnsdistrules_cc.cc

@@ -157,4 +157,75 @@ BOOST_AUTO_TEST_CASE(test_poolOutstandingRule) {
   BOOST_CHECK_EQUAL(pOR2.matches(&dq), false);
 }
 
+BOOST_AUTO_TEST_CASE(test_payloadSizeRule) {
+  auto dnsQuestion = getDQ();
+
+  {
+    PayloadSizeRule rule("equal", dnsQuestion.getData().size());
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), true);
+    BOOST_CHECK_EQUAL(rule.toString(), "payload size is equal to " + std::to_string(dnsQuestion.getData().size()));
+  }
+
+  {
+    PayloadSizeRule rule("equal", dnsQuestion.getData().size() + 1);
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), false);
+  }
+
+  {
+    PayloadSizeRule rule("greater", dnsQuestion.getData().size());
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), false);
+    BOOST_CHECK_EQUAL(rule.toString(), "payload size is greater than " + std::to_string(dnsQuestion.getData().size()));
+  }
+
+  {
+    PayloadSizeRule rule("greater", dnsQuestion.getData().size() - 1);
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), true);
+  }
+
+  {
+    PayloadSizeRule rule("smaller", dnsQuestion.getData().size());
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), false);
+    BOOST_CHECK_EQUAL(rule.toString(), "payload size is smaller than " + std::to_string(dnsQuestion.getData().size()));
+  }
+
+  {
+    PayloadSizeRule rule("smaller", dnsQuestion.getData().size() + 1);
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), true);
+  }
+
+  {
+    PayloadSizeRule rule("greaterOrEqual", dnsQuestion.getData().size());
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), true);
+    BOOST_CHECK_EQUAL(rule.toString(), "payload size is equal to or greater than " + std::to_string(dnsQuestion.getData().size()));
+  }
+
+  {
+    PayloadSizeRule rule("greaterOrEqual", dnsQuestion.getData().size() - 1);
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), true);
+  }
+
+  {
+    PayloadSizeRule rule("greaterOrEqual", dnsQuestion.getData().size() + 1);
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), false);
+  }
+
+  {
+    PayloadSizeRule rule("smallerOrEqual", dnsQuestion.getData().size());
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), true);
+    BOOST_CHECK_EQUAL(rule.toString(), "payload size is equal to or smaller than " + std::to_string(dnsQuestion.getData().size()));
+  }
+
+  {
+    PayloadSizeRule rule("smallerOrEqual", dnsQuestion.getData().size() + 1);
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), true);
+  }
+
+  {
+    PayloadSizeRule rule("smallerOrEqual", dnsQuestion.getData().size() - 1);
+    BOOST_CHECK_EQUAL(rule.matches(&dnsQuestion), false);
+  }
+
+  BOOST_CHECK_THROW(PayloadSizeRule("invalid", 42U), std::runtime_error);
+}
+
 BOOST_AUTO_TEST_SUITE_END()

regression-tests.dnsdist/test_Size.py

@@ -0,0 +1,79 @@
+#!/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)
+        self.assertEqual(len(query.to_wire()), self._payloadSize)
+
+        for method in ("sendUDPQuery", "sendTCPQuery"):
+            sender = getattr(self, method)
+            (_, receivedResponse) = sender(query, response=None, useQueue=False)
+            self.assertTrue(receivedResponse)
+            self.assertEqual(receivedResponse, expectedResponse)
+
+class TestTruncateOversizedResponse(DNSDistTest):
+
+    _payloadSize = 512
+    _config_template = """
+    addResponseAction(PayloadSizeRule("greater", %d), TCResponseAction())
+    newServer{address="127.0.0.1:%d"}
+    """
+    _config_params = ['_payloadSize', '_testServerPort']
+
+    def testTruncateOversizedResponse(self):
+        """
+        Size: Truncate oversized response
+        """
+        name = 'truncate-oversized.size.tests.powerdns.com.'
+        query = dns.message.make_query(name, 'TXT', 'IN')
+        query.flags &= ~dns.flags.RD
+        expectedResponse = dns.message.make_response(query)
+        expectedResponse.flags |= dns.flags.TC
+
+        backendResponse = dns.message.make_response(query)
+        content = ''
+        for i in range(2):
+            if len(content) > 0:
+                content = content + ' '
+            content = content + 'A' * 255
+        rrset = dns.rrset.from_text(name,
+                                    3600,
+                                    dns.rdataclass.IN,
+                                    dns.rdatatype.TXT,
+                                    content)
+        backendResponse.answer.append(rrset)
+        self.assertGreater(len(backendResponse.to_wire()), self._payloadSize)
+
+        (receivedQuery, receivedResponse) = self.sendUDPQuery(query, response=backendResponse)
+        self.assertTrue(receivedQuery)
+        self.assertTrue(receivedResponse)
+        self.assertEqual(receivedResponse, expectedResponse)
+
+        (receivedQuery, receivedResponse) = self.sendTCPQuery(query, response=backendResponse)
+        self.assertTrue(receivedQuery)
+        self.assertTrue(receivedResponse)
+        self.assertEqual(receivedResponse, backendResponse)