add ProbaRule to dnsdist: match with given probability

This adds a ProbaRule, ProbaRule(1.0) means 'match always', 0.1 '10%'. Useful for TeeAction.

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index 799c10eb1acb114a78370b23066c2751298bed08..5dfc499bb43d0f3b7edc04032aadc3854f9e3551 100644 (file)
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -969,6 +969,11 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
       return std::shared_ptr<DNSRule>(new AllRule());
     });
 
+  g_lua.writeFunction("ProbaRule", [](double proba) {
+      return std::shared_ptr<DNSRule>(new ProbaRule(proba));
+    });
+
+  
   g_lua.writeFunction("QNameRule", [](const std::string& qname) {
       return std::shared_ptr<DNSRule>(new QNameRule(DNSName(qname)));
     });

diff --git a/pdns/dnsdistdist/dnsrulactions.cc b/pdns/dnsdistdist/dnsrulactions.cc
index 49292763e2122921f9fa50d2f1e115bd499af414..5e07bc7d21994bf93c56df561f636d69c91c22b4 100644 (file)
--- a/pdns/dnsdistdist/dnsrulactions.cc
+++ b/pdns/dnsdistdist/dnsrulactions.cc
@@ -21,9 +21,24 @@
  */
 #include "dnsrulactions.hh"
 #include <iostream>
+#include <boost/format.hpp>
 
 using namespace std;
 
+bool ProbaRule::matches(const DNSQuestion* dq) const
+{
+  if(d_proba == 1.0)
+    return true;
+  double rnd = 1.0*random() / RAND_MAX;
+  return rnd > (1.0 - d_proba);
+}
+
+string ProbaRule::toString() const 
+{
+  return "match with prob. " + (boost::format("%0.2f") % d_proba).str();
+}
+
+
 TeeAction::TeeAction(const ComboAddress& ca, bool addECS) : d_remote(ca), d_addECS(addECS)
 {
   d_fd=SSocket(d_remote.sin4.sin_family, SOCK_DGRAM, 0);
@@ -39,10 +54,12 @@ TeeAction::~TeeAction()
   d_worker.join();
 }
 
+
 DNSAction::Action TeeAction::operator()(DNSQuestion* dq, string* ruleresult) const 
 {
-  if(dq->tcp) 
+  if(dq->tcp) {
     d_tcpdrops++;
+  }
   else {
     ssize_t res;
     d_queries++;

diff --git a/pdns/dnsdistdist/docs/rules-actions.rst b/pdns/dnsdistdist/docs/rules-actions.rst
index 1e32776dace29a85b5f887becacd8090038d38fb..8aa459b3bbb1169bc49d8ae712656a1f1af923f3 100644 (file)
--- a/pdns/dnsdistdist/docs/rules-actions.rst
+++ b/pdns/dnsdistdist/docs/rules-actions.rst
@@ -414,6 +414,14 @@ These ``DNSRule``\ s be one of the following items:
 
   :param int code: The opcode to match
 
+.. function:: ProbaRule(probability)
+
+  .. versionadded:: 1.3.0
+
+  Matches queries with a given probability. 1.0 means "always"
+
+  :param double probability: Probability of a match
+
 .. function:: QClassRule(qclass)
 
   Matches queries with the specified ``qclass``.

diff --git a/pdns/dnsrulactions.hh b/pdns/dnsrulactions.hh
index c00452cfa2fa18ce93217ce88a7ed58123d6f5d6..ab004c4666923e880ec251dbb8ea182f1f83d1c6 100644 (file)
--- a/pdns/dnsrulactions.hh
+++ b/pdns/dnsrulactions.hh
@@ -760,6 +760,18 @@ public:
 };
 
 
+class ProbaRule : public DNSRule
+{
+public:
+  ProbaRule(double proba) : d_proba(proba)
+  {
+  }
+  bool matches(const DNSQuestion* dq) const override;
+  string toString() const override;
+  double d_proba;
+};
+
+
 class DropAction : public DNSAction
 {
 public:
@@ -834,6 +846,7 @@ public:
   DNSAction::Action operator()(DNSQuestion* dq, string* ruleresult) const override;
   string toString() const override;
   std::unordered_map<string, double> getStats() const override;
+
 private:
   ComboAddress d_remote;
   std::thread d_worker;
@@ -856,8 +869,6 @@ private:
   bool d_addECS{false};
 };
 
-
-
 class PoolAction : public DNSAction
 {
 public: