From d382600664f43323ca155ca955c001d741a52d43 Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@netherlabs.nl>
Date: Tue, 19 Sep 2017 22:22:57 +0200
Subject: [PATCH] add ProbaRule to dnsdist: match with given probability

This adds a ProbaRule, ProbaRule(1.0) means 'match always', 0.1 '10%'. Useful for TeeAction.
---
 pdns/dnsdist-lua.cc                     |  5 +++++
 pdns/dnsdistdist/dnsrulactions.cc       | 19 ++++++++++++++++++-
 pdns/dnsdistdist/docs/rules-actions.rst |  8 ++++++++
 pdns/dnsrulactions.hh                   | 15 +++++++++++++--
 4 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index 799c10eb1a..5dfc499bb4 100644
--- 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 49292763e2..5e07bc7d21 100644
--- 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 1e32776dac..8aa459b3bb 100644
--- 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 c00452cfa2..ab004c4666 100644
--- 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:
-- 
2.47.2