From 0e7604973cf415606eeeb7aa31a8e6f1ac8a8ab3 Mon Sep 17 00:00:00 2001
From: bert hubert <bert.hubert@netherlabs.nl>
Date: Tue, 1 Dec 2015 17:39:59 +0100
Subject: [PATCH] move our RPZ blocking to the most GLORIOUS NetmaskTree
 (thanks Aki!)

---
 pdns/filterpo.cc | 25 +++++++++----------------
 pdns/filterpo.hh |  8 +++++---
 2 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/pdns/filterpo.cc b/pdns/filterpo.cc
index a466455346..c50f7e55cd 100644
--- a/pdns/filterpo.cc
+++ b/pdns/filterpo.cc
@@ -64,11 +64,9 @@ DNSFilterEngine::Policy DNSFilterEngine::getQueryPolicy(const DNSName& qname, co
       return pol;
     }
     
-    for(const auto& qa : z.qpolAddr) {
-      if(qa.first.match(ca)) {
-	//	cerr<<"Had a hit on the IP address ("<<ca.toString()<<") of the client"<<endl;
-	return qa.second;
-      }
+    if(auto fnd=z.qpolAddr.lookup(ca)) {
+      //	cerr<<"Had a hit on the IP address ("<<ca.toString()<<") of the client"<<endl;
+      return fnd->second;
     }
   }
 
@@ -90,12 +88,8 @@ DNSFilterEngine::Policy DNSFilterEngine::getPostPolicy(const vector<DNSRecord>&
       continue;
 
     for(const auto& z : d_zones) {
-      for(const auto& qa : z.postpolAddr) {
-	if(qa.first.match(ca)) {
-	  //	  cerr<<"Had a hit on IP address in answer"<<endl;
-	  return qa.second;
-	}
-      }
+      if(auto fnd=z.postpolAddr.lookup(ca))
+	return fnd->second;
     }
   }
   return Policy{PolicyKind::NoAction};
@@ -105,19 +99,18 @@ void DNSFilterEngine::assureZones(int zone)
 {
   if((int)d_zones.size() <= zone)
     d_zones.resize(zone+1);
-
 }
 
 void DNSFilterEngine::addClientTrigger(const Netmask& nm, Policy pol, int zone)
 {
   assureZones(zone);
-  d_zones[zone].qpolAddr.push_back({nm,pol});
+  d_zones[zone].qpolAddr.insert(nm).second=pol;
 }
 
 void DNSFilterEngine::addResponseTrigger(const Netmask& nm, Policy pol, int zone)
 {
   assureZones(zone);
-  d_zones[zone].postpolAddr.push_back({nm,pol});
+  d_zones[zone].postpolAddr.insert(nm).second=pol;
 }
 
 void DNSFilterEngine::addQNameTrigger(const DNSName& n, Policy pol, int zone)
@@ -137,7 +130,7 @@ bool DNSFilterEngine::rmClientTrigger(const Netmask& nm, Policy pol, int zone)
   assureZones(zone);
 
   auto& qpols = d_zones[zone].qpolAddr;
-  qpols.erase(remove(qpols.begin(), qpols.end(),pair<Netmask,Policy>(nm,pol)), qpols.end());
+  qpols.erase(nm);
   return true;
 }
 
@@ -145,7 +138,7 @@ bool DNSFilterEngine::rmResponseTrigger(const Netmask& nm, Policy pol, int zone)
 {
   assureZones(zone);
   auto& postpols = d_zones[zone].postpolAddr;
-  postpols.erase(remove(postpols.begin(), postpols.end(),pair<Netmask,Policy>(nm,pol)), postpols.end());
+  postpols.erase(nm);  
   return true;
 }
 
diff --git a/pdns/filterpo.hh b/pdns/filterpo.hh
index 8633eb338e..95810eb879 100644
--- a/pdns/filterpo.hh
+++ b/pdns/filterpo.hh
@@ -34,7 +34,6 @@
    Verbatim domain names
    Wildcard versions (*.domain.com does NOT match domain.com)
    Netmasks (IPv4 and IPv6)
-
    Finally, triggers are grouped in different zones. The "first" zone that has a match
    is consulted. Then within that zone, rules again have precedences. 
 */
@@ -72,13 +71,16 @@ public:
   Policy getProcessingPolicy(const DNSName& qname) const;
   Policy getPostPolicy(const vector<DNSRecord>& records) const;
 
+  size_t size() {
+    return d_zones.size();
+  }
 private:
   void assureZones(int zone);
   struct Zone {
     std::map<DNSName, Policy> qpolName;
-    std::vector<pair<Netmask, Policy>> qpolAddr;
+    NetmaskTree<Policy> qpolAddr;
     std::map<DNSName, Policy> propolName;
-    std::vector<pair<Netmask, Policy>> postpolAddr;
+    NetmaskTree<Policy> postpolAddr;
   };
   vector<Zone> d_zones;
 
-- 
2.47.2