From: bert hubert <bert.hubert@netherlabs.nl>
Date: Mon, 22 May 2017 20:59:11 +0000 (+0200)
Subject: Implement a runtime changeable rule that matches IP address for a certain time.
X-Git-Tag: rec-4.1.0-alpha1~61^2~2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=2b58b6257e08695600cb31c1f3763d0fb924a329;p=thirdparty%2Fpdns.git

Implement a runtime changeable rule that matches IP address for a certain time.
This effectively allows (for example) pool selection from Lua, but then cached.

Sample code:

```
newServer({address="192.168.1.20", pool=""})
newServer({address="8.8.8.8", pool="elgoog"})

tisrElGoog=TimedIPSetRule()
tisrRest=TimedIPSetRule()
addAction(tisrElGoog:slice(), PoolAction("elgoog"))
addAction(tisrRest:slice(), PoolAction(""))

elgoogPeople=newNMG()
elgoogPeople:addMask("192.168.1.0/28")

function pickPool(dq)
	if(elgoogPeople:match(dq.remoteaddr)) -- in real life, this would be external
	then
		print("Lua caught query for a googlePerson")
		tisrElGoog:add(dq.remoteaddr, 10)
		return DNSAction.Pool, "elgoog"
	else
		print("Lua caught query for restPerson")
		tisrRest:add(dq.remoteaddr, 60)
		return DNSAction.None, ""
	end
end

addLuaAction(AllRule(), pickPool)
```
---

diff --git a/pdns/dnsdist-lua2.cc b/pdns/dnsdist-lua2.cc
index c25a8780a5..4efe4235b5 100644
--- a/pdns/dnsdist-lua2.cc
+++ b/pdns/dnsdist-lua2.cc
@@ -1265,6 +1265,22 @@ void moreLua(bool client)
       return std::shared_ptr<DNSRule>(new RDRule());
     });
 
+    g_lua.writeFunction("TimedIPSetRule", []() {
+      return std::shared_ptr<TimedIPSetRule>(new TimedIPSetRule());
+    });
+
+    g_lua.registerFunction<void(std::shared_ptr<TimedIPSetRule>::*)()>("clear", [](std::shared_ptr<TimedIPSetRule> tisr) {
+        tisr->clear();
+      });
+
+    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);
+      });
+        
+    g_lua.registerFunction<std::shared_ptr<DNSRule>(std::shared_ptr<TimedIPSetRule>::*)()>("slice", [](std::shared_ptr<TimedIPSetRule> tisr) {
+        return std::dynamic_pointer_cast<DNSRule>(tisr);
+      });
+    
     g_lua.writeFunction("setWHashedPertubation", [](uint32_t pertub) {
         setLuaSideEffect();
         g_hashperturb = pertub;
diff --git a/pdns/dnsrulactions.hh b/pdns/dnsrulactions.hh
index ea87759f16..f852e3bd02 100644
--- a/pdns/dnsrulactions.hh
+++ b/pdns/dnsrulactions.hh
@@ -134,6 +134,122 @@ private:
   bool d_src;
 };
 
+class TimedIPSetRule : public DNSRule, boost::noncopyable
+{
+private:
+  struct IPv6 {
+    IPv6(const ComboAddress& ca)
+    {
+      static_assert(sizeof(*this)==16, "IPv6 struct has wrong size");
+      memcpy((char*)this, ca.sin6.sin6_addr.s6_addr, 16);
+    }
+    bool operator==(const IPv6& rhs) const
+    {
+      return a==rhs.a && b==rhs.b;
+    }
+    uint64_t a, b;
+  };
+
+public:
+  TimedIPSetRule()
+  {
+    pthread_rwlock_init(&d_lock4, 0);
+    pthread_rwlock_init(&d_lock6, 0);
+  }
+  bool matches(const DNSQuestion* dq) const override
+  {
+    if(dq->remote->sin4.sin_family == AF_INET) {
+      ReadLock rl(&d_lock4);
+      auto fnd = d_ip4s.find(dq->remote->sin4.sin_addr.s_addr);
+      if(fnd == d_ip4s.end()) {
+        return false;
+      }
+      return time(0) < fnd->second;
+    } else {
+      ReadLock rl(&d_lock6);
+      auto fnd = d_ip6s.find({*dq->remote});
+      if(fnd == d_ip6s.end()) {
+        return false;
+      }
+      return time(0) < fnd->second;
+    }
+  }
+
+  void add(const ComboAddress& ca, time_t ttd)
+  {
+    // think twice before adding templates here
+    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;
+    }
+    else {
+      WriteLock rl(&d_lock6);
+      auto res=d_ip6s.insert({{ca}, ttd});
+      if(!res.second && res.first->second < ttd)
+        res.first->second = ttd;
+    }
+  }
+
+  void remove(const ComboAddress& ca)
+  {
+    if(ca.sin4.sin_family == AF_INET) {
+      WriteLock rl(&d_lock4);
+      d_ip4s.erase(ca.sin4.sin_addr.s_addr);
+    }
+    else {
+      WriteLock rl(&d_lock6);
+      d_ip6s.erase({ca});
+    }
+  }
+
+  void clear()
+  {
+    {
+      WriteLock rl(&d_lock4);
+      d_ip4s.clear();
+    }
+    WriteLock rl(&d_lock6);
+    d_ip6s.clear();
+  }
+  
+  string toString() const override
+  {
+    time_t now=time(0);
+    uint64_t count = 0;
+    {
+      ReadLock rl(&d_lock4);
+      for(const auto& ip : d_ip4s)
+        if(now < ip.second)
+          ++count;
+    }
+    {
+      ReadLock rl(&d_lock6);
+      for(const auto& ip : d_ip6s)
+        if(now < ip.second)
+          ++count;
+    }
+    
+    return "Src: "+std::to_string(count)+" ips";
+  }
+private:
+  struct IPv6Hash
+  {
+    std::size_t operator()(const IPv6& ip) const
+    {
+      auto ah=std::hash<uint64_t>{}(ip.a);
+      auto bh=std::hash<uint64_t>{}(ip.b);
+      return ah & (bh<<1);
+    }
+  };
+  std::unordered_map<IPv6, time_t, IPv6Hash> d_ip6s;
+  std::unordered_map<uint32_t, time_t> d_ip4s;
+  mutable pthread_rwlock_t d_lock4;
+  mutable pthread_rwlock_t d_lock6;
+};
+
+
 class AllRule : public DNSRule
 {
 public: