add QPS limiting per domain or subnet

diff --git a/pdns/README-dnsdist.md b/pdns/README-dnsdist.md
index 470aeb800b1864cc456b061c792ecaf62d814103..c866fcc3901f09e0a43f0fdc74dfd2b2f81dca70 100644 (file)
--- a/pdns/README-dnsdist.md
+++ b/pdns/README-dnsdist.md
@@ -154,6 +154,33 @@ This is still much in flux, but for now, try:
  * `topQueries(20,2)`: shows the top-20 two-level domain queries (so `topQueries(20,1)` only shows TLDs)
  * `topResponses(20, 2)`: top-20 servfail responses (use ,3 for NXDOMAIN)
 
+
+Per domain or subnet QPS limiting
+---------------------------------
+If certain domains or source addresses are generating onerous amounts of
+traffic, you can put ceilings on the amount of traffic you are willing to
+forward:
+
+```
+> addQPSLimit("h4xorbooter.xyz.", 10)
+> addQPSLimit({"130.161.0.0/16", "145.14.0.0/16"} , 20)
+> addQPSLimit({"nl.", "be."}, 1)
+> showQPSLimits()
+#   Object                                                 Lim   Passed  Blocked
+0   h4xorbooter.xyz.                                        10        0        0
+1   130.161.0.0/16, 145.14.0.0/16                           20        0        0
+2   nl., be.                                                 1        2        8
+```
+
+To delete a limit:
+```
+> deleteQPSLimit(1)
+> showQPSLimits()
+#   Object                                                 Lim   Passed  Blocked
+0   h4xorbooter.xyz.                                        10        0        0
+1   nl., be.                                                 1       16      251
+```
+
 Dynamic load balancing
 ----------------------
 

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index be88f3965f30071c4dd57014efa4d3897e3287de..117001f516c0085f20683c9d01e5386d68006b69 100644 (file)
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -74,6 +74,7 @@ bool g_console;
 NetmaskGroup g_ACL;
 string g_outputBuffer;
 
+
 /* UDP: the grand design. Per socket we listen on for incoming queries there is one thread.
    Then we have a bunch of connected sockets for talking to downstream servers. 
    We send directly to those sockets.
@@ -153,11 +154,22 @@ public:
     return d_passthrough? 0 : d_rate;
   }
 
+  int getPassed() const
+  {
+    return d_passed;
+  }
+  int getBlocked() const
+  {
+    return d_blocked;
+  }
+
   bool check()
   {
     if(d_passthrough)
       return true;
-    d_tokens += 1.0*d_rate * (d_prev.udiffAndSet()/1000000.0);
+    auto delta = d_prev.udiffAndSet();
+  
+    d_tokens += 1.0*d_rate * (delta/1000000.0);
 
     if(d_tokens > d_burst)
       d_tokens = d_burst;
@@ -166,7 +178,11 @@ public:
     if(d_tokens >= 1.0) { // we need this because burst=1 is weird otherwise
       ret=true;
       --d_tokens;
+      d_passed++;
     }
+    else
+      d_blocked++;
+
     return ret; 
   }
 private:
@@ -175,8 +191,12 @@ private:
   unsigned int d_burst;
   double d_tokens;
   StopWatch d_prev;
+  unsigned int d_passed{0};
+  unsigned int d_blocked{0};
 };
 
+vector<pair<boost::variant<SuffixMatchNode,NetmaskGroup>, QPSLimiter> > g_limiters;
+
 struct IDState
 {
   IDState() : origFD(-1) {}
@@ -469,6 +489,24 @@ try
 
     g_rings.queryRing.push_back(qname);
     
+    bool blocked=false;
+    for(auto& lim : g_limiters) {
+      if(auto nmg=boost::get<NetmaskGroup>(&lim.first)) {
+       if(nmg->match(remote) && !lim.second.check()) {
+         blocked=true;
+         break;
+       }
+      }
+      else if(auto smn=boost::get<SuffixMatchNode>(&lim.first)) {
+       if(smn->check(qname) && !lim.second.check()) {
+         blocked=true;
+         break;
+       }
+      }
+    }
+    if(blocked)
+      continue;
+
     if(blockFilter)
     {
       std::lock_guard<std::mutex> lock(g_luamutex);
@@ -931,15 +969,18 @@ void setupLua(bool client)
                            return a->order < b->order;
                          });
                        return ret;
-                      
-
                      } );
 
 
+
+
   g_lua.writeFunction("deleteServer", 
-                     [](std::shared_ptr<DownstreamState> rem)
+                     [](boost::variant<std::shared_ptr<DownstreamState>, int> var)
                      { 
-                       g_dstates.erase(remove(g_dstates.begin(), g_dstates.end(), rem), g_dstates.end());
+                       if(auto* rem = boost::get<shared_ptr<DownstreamState>>(&var))
+                         g_dstates.erase(remove(g_dstates.begin(), g_dstates.end(), *rem), g_dstates.end());
+                       else
+                         g_dstates.erase(g_dstates.begin() + boost::get<int>(var));
                      } );
 
 
@@ -1018,6 +1059,52 @@ void setupLua(bool client)
     });
 
 
+  g_lua.writeFunction("addQPSLimit", [](boost::variant<string,vector<pair<int, string>> > var, int lim) {
+      SuffixMatchNode smn;
+      NetmaskGroup nmg;
+
+      auto add=[&](string src) {
+       try {
+         smn.add(DNSName(src));
+       } catch(...) {
+         nmg.addMask(src);
+       }
+      };
+      if(auto src = boost::get<string>(&var))
+       add(*src);
+      else {
+       for(auto& a : boost::get<vector<pair<int, string>>>(var)) {
+         add(a.second);
+       }
+      }
+      if(nmg.empty())
+       g_limiters.push_back({smn, QPSLimiter(lim, lim)});
+      else
+       g_limiters.push_back({nmg, QPSLimiter(lim, lim)});
+    });
+
+  g_lua.writeFunction("deleteQPSLimit", [](int i) {
+      g_limiters.erase(g_limiters.begin() + i);
+    });
+
+  g_lua.writeFunction("showQPSLimits", []() {
+      boost::format fmt("%-3d %-50s %7d %8d %8d\n");
+      g_outputBuffer += (fmt % "#" % "Object" % "Lim" % "Passed" % "Blocked").str();
+      int num=0;
+      for(const auto& lim : g_limiters) {
+       string name;
+       if(auto nmg=boost::get<NetmaskGroup>(&lim.first)) {
+         name=nmg->toString();
+       }
+       else if(auto smn=boost::get<SuffixMatchNode>(&lim.first)) {
+         name=smn->toString(); 
+       }
+       g_outputBuffer += (fmt % num % name % lim.second.getRate() % lim.second.getPassed() % lim.second.getBlocked()).str();
+       ++num;
+      }
+    });
+
+
   g_lua.writeFunction("getServers", []() {
       vector<pair<int, std::shared_ptr<DownstreamState> > > ret;
       int count=1;

diff --git a/pdns/dnsname.hh b/pdns/dnsname.hh
index 8e03bb03523be0805cd350495c5f4522718db5be..7f1e9d9936a86b30392300d595b3fabb9e139b26 100644 (file)
--- a/pdns/dnsname.hh
+++ b/pdns/dnsname.hh
@@ -71,6 +71,7 @@ struct SuffixMatchNode
   SuffixMatchNode(const std::string& name_="", bool endNode_=false) : name(name_), endNode(endNode_)
   {}
   std::string name;
+  std::string d_human;
   mutable bool endNode;
   mutable std::set<SuffixMatchNode> children;
   bool operator<(const SuffixMatchNode& rhs) const
@@ -80,6 +81,9 @@ struct SuffixMatchNode
 
   void add(const DNSName& name) 
   {
+    if(!d_human.empty())
+      d_human.append(", ");
+    d_human += name.toString();
     add(name.getRawLabels());
   }
 
@@ -103,7 +107,6 @@ struct SuffixMatchNode
     return check(name.getRawLabels());
   }
 
-
   bool check(std::deque<std::string> labels) const
   {
     if(labels.empty()) // optimization
@@ -116,6 +119,10 @@ struct SuffixMatchNode
     labels.pop_back();
     return child->check(labels);
   }
-            
+  
+  std::string toString() const
+  {
+    return d_human;
+  }
 
 };