From bac6e8fbaffc63d0a675b60c07d94c8aae292400 Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@netherlabs.nl>
Date: Thu, 14 Jan 2016 11:45:49 +0100
Subject: [PATCH] implement & document exceedQRate(), plus populate dnsdist.*
 with dns types.

---
 pdns/README-dnsdist.md |  1 +
 pdns/dnsdist-lua.cc    |  5 +++++
 pdns/dnsdist-lua2.cc   | 14 +++++++++++---
 pdns/dnsdistconf.lua   |  7 ++++++-
 4 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/pdns/README-dnsdist.md b/pdns/README-dnsdist.md
index e7c9711ce..cfc2116b5 100644
--- a/pdns/README-dnsdist.md
+++ b/pdns/README-dnsdist.md
@@ -862,6 +862,7 @@ instantiate a server with additional parameters
    * `exceedServFails(rate, seconds)`: get set of addresses that exceed `rate` servails/s over `seconds` seconds
    * `exceedNXDOMAINs(rate, seconds)`: get set of addresses that exceed `rate` NXDOMAIN/s over `seconds` seconds
    * `exceedRespByterate(rate, seconds)`: get set of addresses that exeeded `rate` bytes/s answers over `seconds` seconds
+   * `exceedQRate(rate, seconds)`: get set of address that exceed `rate` queries/s over `seconds` seconds
    * `exceedQTypeRate(type, rate, seconds)`: get set of address that exceed `rate` queries/s for queries of type `type` over `seconds` seconds
  * Advanced functions for writing your own policies and hooks
    * ComboAddress related:
diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index 91df22529..c8260f063 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -132,6 +132,11 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
       {"None",(int)DNSAction::Action::None},
       {"Delay", (int)DNSAction::Action::Delay}}
     );
+
+  vector<pair<string, int> > dd;
+  for(const auto& n : QType::names)
+    dd.push_back({n.first, n.second});
+  g_lua.writeVariable("dnsdist", dd);
   
   g_lua.writeFunction("newServer", 
 		      [client](boost::variant<string,newserver_t> pvars, boost::optional<int> qps)
diff --git a/pdns/dnsdist-lua2.cc b/pdns/dnsdist-lua2.cc
index 94bdd1cc3..0cf7095f7 100644
--- a/pdns/dnsdist-lua2.cc
+++ b/pdns/dnsdist-lua2.cc
@@ -164,7 +164,8 @@ void moreLua()
 			   until.tv_sec += actualSeconds; 
 			   for(const auto& capair : m) {
 			     unsigned int count = 0;
-			     if(auto got = slow.lookup(Netmask(capair.first))) {
+                             auto got = slow.lookup(Netmask(capair.first));
+			     if(got) {
 			       if(until < got->second.until) // had a longer policy
 				 continue;
 			       if(now < got->second.until) // don't inherit count on expired entry
@@ -172,7 +173,8 @@ void moreLua()
 			     }
 			     DynBlock db{msg,until};
 			     db.blocks=count;
-                             warnlog("Inserting dynamic block for %s for %d seconds: %s", capair.first.toString(), actualSeconds, msg);
+                             if(!got)
+                               warnlog("Inserting dynamic block for %s for %d seconds: %s", capair.first.toString(), actualSeconds, msg);
 			     slow.insert(Netmask(capair.first)).second=db;
 			   }
 			   g_dynblockNMG.setState(slow);
@@ -204,10 +206,16 @@ void moreLua()
 	  if(q.qtype==type)
 	    counts[q.requestor]++;
 	});
+    });
 
-
+  g_lua.writeFunction("exceedQRate", [](unsigned int rate, int seconds) {
+      setLuaNoSideEffect();
+      return exceedQueryGen(rate, seconds, [](counts_t& counts, const Rings::Query& q) {
+          counts[q.requestor]++;
+	});
     });
 
+
   g_lua.writeFunction("topBandwidth", [](boost::optional<unsigned int> top_) {
       setLuaNoSideEffect();
       auto top = top_.get_value_or(10);
diff --git a/pdns/dnsdistconf.lua b/pdns/dnsdistconf.lua
index dba9a8be9..ca2dd1341 100644
--- a/pdns/dnsdistconf.lua
+++ b/pdns/dnsdistconf.lua
@@ -97,4 +97,9 @@ function splitSetup(servers, remote, qname, qtype, dh)
 	 end
 end
 
--- setServerPolicyLua("splitSetup", splitSetup)
\ No newline at end of file
+-- setServerPolicyLua("splitSetup", splitSetup)
+
+function maintenance()
+	addDynBlocks(exceedQRate(20, 10), "Exceeded query rate", 60)
+end
+
-- 
2.40.0