From 01b8149e8682f12200898dd160ce9f1bc7c1f9c8 Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@powerdns.com>
Date: Tue, 23 May 2017 11:52:34 +0200
Subject: [PATCH] reduce memory usage of TimedIPSet, add cleanup(), add
 autocomplete

---
 pdns/dnsdist-console.cc |  1 +
 pdns/dnsdist-lua2.cc    |  4 ++++
 pdns/dnsrulactions.hh   | 41 +++++++++++++++++++++++++++++++++++------
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/pdns/dnsdist-console.cc b/pdns/dnsdist-console.cc
index 56ef0c419b..6a1e15465c 100644
--- a/pdns/dnsdist-console.cc
+++ b/pdns/dnsdist-console.cc
@@ -404,6 +404,7 @@ const std::vector<ConsoleKeyword> g_consoleKeywords{
   { "SpoofAction", true, "{ip, ...} ", "forge a response with the specified IPv4 (for an A query) or IPv6 (for an AAAA). If you specify multiple addresses, all that match the query type (A, AAAA or ANY) will get spoofed in" },
   { "TCAction", true, "", "create answer to query with TC and RD bits set, to move to TCP" },
   { "testCrypto", true, "", "test of the crypto all works" },
+  { "TimedIPSetRule", true, "", "Create a rule which matches a set of IP addresses which expire"}, 
   { "topBandwidth", true, "top", "show top-`top` clients that consume the most bandwidth over length of ringbuffer" },
   { "topCacheHitResponseRule", true, "", "move the last cache hit response rule to the first position" },
   { "topClients", true, "n", "show top-`n` clients sending the most queries over length of ringbuffer" },
diff --git a/pdns/dnsdist-lua2.cc b/pdns/dnsdist-lua2.cc
index 4efe4235b5..20360e9725 100644
--- a/pdns/dnsdist-lua2.cc
+++ b/pdns/dnsdist-lua2.cc
@@ -1273,6 +1273,10 @@ void moreLua(bool client)
         tisr->clear();
       });
 
+    g_lua.registerFunction<void(std::shared_ptr<TimedIPSetRule>::*)()>("cleanup", [](std::shared_ptr<TimedIPSetRule> tisr) {
+        tisr->cleanup();
+      });
+
     g_lua.registerFunction<void(std::shared_ptr<TimedIPSetRule>::*)(const ComboAddress& ca, int t)>("add", [](std::shared_ptr<TimedIPSetRule> tisr, const ComboAddress& ca, int t) {
         tisr->add(ca, time(0)+t);
       });
diff --git a/pdns/dnsrulactions.hh b/pdns/dnsrulactions.hh
index f852e3bd02..4987602b40 100644
--- a/pdns/dnsrulactions.hh
+++ b/pdns/dnsrulactions.hh
@@ -181,14 +181,14 @@ public:
     if(ca.sin4.sin_family == AF_INET) {
       WriteLock rl(&d_lock4);
       auto res=d_ip4s.insert({ca.sin4.sin_addr.s_addr, ttd});
-      if(!res.second && res.first->second < ttd)
-        res.first->second = ttd;
+      if(!res.second && (time_t)res.first->second < ttd)
+        res.first->second = (uint32_t)ttd;
     }
     else {
       WriteLock rl(&d_lock6);
       auto res=d_ip6s.insert({{ca}, ttd});
-      if(!res.second && res.first->second < ttd)
-        res.first->second = ttd;
+      if(!res.second && (time_t)res.first->second < ttd)
+        res.first->second = (uint32_t)ttd;
     }
   }
 
@@ -213,6 +213,35 @@ public:
     WriteLock rl(&d_lock6);
     d_ip6s.clear();
   }
+
+  void cleanup()
+  {
+    time_t now=time(0);
+    {
+      WriteLock rl(&d_lock4);
+
+      for(auto iter = d_ip4s.begin(); iter != d_ip4s.end(); ) {
+	if(iter->second < now)
+	  iter=d_ip4s.erase(iter);
+	else
+	  ++iter;
+      }
+	    
+    }
+
+    {
+      WriteLock rl(&d_lock6);
+
+      for(auto iter = d_ip6s.begin(); iter != d_ip6s.end(); ) {
+	if(iter->second < now)
+	  iter=d_ip6s.erase(iter);
+	else
+	  ++iter;
+      }
+	    
+    }
+
+  }
   
   string toString() const override
   {
@@ -243,8 +272,8 @@ private:
       return ah & (bh<<1);
     }
   };
-  std::unordered_map<IPv6, time_t, IPv6Hash> d_ip6s;
-  std::unordered_map<uint32_t, time_t> d_ip4s;
+  std::unordered_map<IPv6, uint32_t, IPv6Hash> d_ip6s;
+  std::unordered_map<uint32_t, uint32_t> d_ip4s;
   mutable pthread_rwlock_t d_lock4;
   mutable pthread_rwlock_t d_lock6;
 };
-- 
2.47.2