From: bert hubert <bert.hubert@netherlabs.nl>
Date: Tue, 1 Sep 2015 14:24:39 +0000 (+0200)
Subject: cleanup rule/action infrastructure to be generic in dnsdist
X-Git-Tag: dnsdist-1.0.0-alpha1~248^2~70
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=6bba426c9cf95502c094d244f72aa8e36fccd5e2;p=pdns

cleanup rule/action infrastructure to be generic in dnsdist
---

diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index dc95f5019..7c458197b 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -34,11 +34,15 @@ private:
   func_t d_func;
 };
 
-std::shared_ptr<DNSRule> makeRule(const boost::variant<string,vector<pair<int, string>> >& var)
+typedef boost::variant<string,vector<pair<int, string>>, std::shared_ptr<DNSRule> > luadnsrule_t;
+std::shared_ptr<DNSRule> makeRule(const luadnsrule_t& var)
 {
+  if(auto src = boost::get<std::shared_ptr<DNSRule>>(&var))
+    return *src;
+  
   SuffixMatchNode smn;
   NetmaskGroup nmg;
-  
+
   auto add=[&](string src) {
     try {
       nmg.addMask(src); // need to try mask first, all masks are domain names!
@@ -332,7 +336,7 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
       }catch(std::exception& e) { g_outputBuffer=e.what(); throw; }
     });
 
-  g_lua.writeFunction("addLuaAction", [](boost::variant<string,vector<pair<int, string>> > var, LuaAction::func_t func) 
+  g_lua.writeFunction("addLuaAction", [](luadnsrule_t var, LuaAction::func_t func) 
 		      {
 			auto rule=makeRule(var);
 			g_rulactions.modify([rule,func](decltype(g_rulactions)::value_type& rulactions){
@@ -342,7 +346,25 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
 		      });
 
 
-  g_lua.writeFunction("addPoolRule", [](boost::variant<string,vector<pair<int, string>> > var, string pool) {
+  g_lua.writeFunction("NoRecurseAction", []() {
+      return std::shared_ptr<DNSAction>(new NoRecurseAction);
+    });
+
+  g_lua.writeFunction("MaxQPSIPRule", [](unsigned int qps) {
+      return std::shared_ptr<DNSRule>(new MaxQPSIPRule(qps));
+    });
+
+
+  g_lua.writeFunction("addAction", [](luadnsrule_t var, std::shared_ptr<DNSAction> ea) 
+		      {
+			auto rule=makeRule(var);
+			g_rulactions.modify([rule, ea](decltype(g_rulactions)::value_type& rulactions){
+			    rulactions.push_back({rule, ea});
+			  });
+		      });
+
+
+  g_lua.writeFunction("addPoolRule", [](luadnsrule_t var, string pool) {
       auto rule=makeRule(var);
 	g_rulactions.modify([rule, pool](decltype(g_rulactions)::value_type& rulactions) {
 	    rulactions.push_back({
@@ -351,7 +373,7 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
 	  });
     });
 
-  g_lua.writeFunction("addNoRecurseRule", [](boost::variant<string,vector<pair<int, string>> > var) {
+  g_lua.writeFunction("addNoRecurseRule", [](luadnsrule_t var) {
       auto rule=makeRule(var);
 	g_rulactions.modify([rule](decltype(g_rulactions)::value_type& rulactions) {
 	    rulactions.push_back({
@@ -361,7 +383,7 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
     });
 
 
-  g_lua.writeFunction("addQPSPoolRule", [](boost::variant<string,vector<pair<int, string>> > var, int limit, string pool) {
+  g_lua.writeFunction("addQPSPoolRule", [](luadnsrule_t var, int limit, string pool) {
       auto rule = makeRule(var);
       g_rulactions.modify([rule, pool,limit](decltype(g_rulactions)::value_type& rulactions) {
 	  rulactions.push_back({
@@ -377,7 +399,7 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
 	});
     });
 
-  g_lua.writeFunction("addQPSLimit", [](boost::variant<string,vector<pair<int, string>> > var, int lim) {
+  g_lua.writeFunction("addQPSLimit", [](luadnsrule_t var, int lim) {
       auto rule = makeRule(var);
       g_rulactions.modify([lim,rule](decltype(g_rulactions)::value_type& rulactions) {
 	  rulactions.push_back({rule, 
@@ -385,7 +407,7 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
 	});
     });
    
-  g_lua.writeFunction("addDelay", [](boost::variant<string,vector<pair<int, string>> > var, int msec) {
+  g_lua.writeFunction("addDelay", [](luadnsrule_t var, int msec) {
       auto rule = makeRule(var);
       g_rulactions.modify([msec,rule](decltype(g_rulactions)::value_type& rulactions) {
 	  rulactions.push_back({rule, 
diff --git a/pdns/dnsrulactions.hh b/pdns/dnsrulactions.hh
index 83234e540..f86397a8e 100644
--- a/pdns/dnsrulactions.hh
+++ b/pdns/dnsrulactions.hh
@@ -1,6 +1,35 @@
 #include "dnsdist.hh"
 #include "dnsname.hh"
 
+class MaxQPSIPRule : public DNSRule
+{
+public:
+  MaxQPSIPRule(unsigned int qps) : d_qps(qps) {}
+
+  bool matches(const ComboAddress& remote, const DNSName& qname, uint16_t qtype, dnsheader* dh, int len) const override
+  {
+    ComboAddress zeroport(remote);
+    zeroport.sin4.sin_port=0;
+    auto iter = d_limits.find(zeroport);
+    if(iter == d_limits.end()) {
+      iter=d_limits.insert({zeroport,QPSLimiter(d_qps, d_qps)}).first;
+    }
+    return !iter->second.check();
+  }
+
+  string toString() const override
+  {
+    return "per IP match for QPS over " + std::to_string(d_qps);
+  }
+
+
+private:
+  mutable std::map<ComboAddress, QPSLimiter> d_limits;
+  unsigned int d_qps;
+
+};
+
+
 class NetmaskGroupRule : public DNSRule
 {
 public:
@@ -90,6 +119,7 @@ public:
   }
 };
 
+
 class QPSAction : public DNSAction
 {
 public:
@@ -215,6 +245,6 @@ public:
   }
   string toString() const override
   {
-    return "rd=0 query";
+    return "set rd=0";
   }
 };