From: Otto Moerbeek <otto.moerbeek@open-xchange.com>
Date: Mon, 26 Feb 2024 08:45:04 +0000 (+0100)
Subject: Also allocate custom records via unique ptr
X-Git-Tag: auth-5.0.0-alpha0~2^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ee4618f0a79b097b477b1f20a1bd684d23ee2a09;p=thirdparty%2Fpdns.git

Also allocate custom records via unique ptr
---

diff --git a/pdns/recursordist/filterpo.cc b/pdns/recursordist/filterpo.cc
index 1fbf867376..bfcbbc71e6 100644
--- a/pdns/recursordist/filterpo.cc
+++ b/pdns/recursordist/filterpo.cc
@@ -366,6 +366,17 @@ void DNSFilterEngine::assureZones(size_t zone)
   }
 }
 
+static void addCustom(DNSFilterEngine::Policy& existingPol, const DNSFilterEngine::Policy& pol)
+{
+  if (!existingPol.d_custom) {
+    existingPol.d_custom = make_unique<DNSFilterEngine::Policy::CustomData>();
+  }
+  if (pol.d_custom) {
+    existingPol.d_custom->reserve(existingPol.d_custom->size() + pol.d_custom->size());
+    std::move(pol.d_custom->begin(), pol.d_custom->end(), std::back_inserter(*existingPol.d_custom));
+  }
+}
+
 void DNSFilterEngine::Zone::addNameTrigger(std::unordered_map<DNSName, Policy>& map, const DNSName& n, Policy&& pol, bool ignoreDuplicate, PolicyType ptype)
 {
   auto iter = map.find(n);
@@ -387,9 +398,7 @@ void DNSFilterEngine::Zone::addNameTrigger(std::unordered_map<DNSName, Policy>&
       throw std::runtime_error("Adding a " + getTypeToString(ptype) + "-based filter policy of kind " + getKindToString(pol.d_kind) + " but a policy of kind " + getKindToString(existingPol.d_kind) + " already exists for for the following name: " + n.toLogString());
     }
 
-    existingPol.d_custom.reserve(existingPol.d_custom.size() + pol.d_custom.size());
-
-    std::move(pol.d_custom.begin(), pol.d_custom.end(), std::back_inserter(existingPol.d_custom));
+    addCustom(existingPol, pol);
   }
   else {
     auto& qpol = map.insert({n, std::move(pol)}).first->second;
@@ -419,9 +428,7 @@ void DNSFilterEngine::Zone::addNetmaskTrigger(NetmaskTree<Policy>& nmt, const Ne
       throw std::runtime_error("Adding a " + getTypeToString(ptype) + "-based filter policy of kind " + getKindToString(pol.d_kind) + " but a policy of kind " + getKindToString(existingPol.d_kind) + " already exists for the following netmask: " + netmask.toString());
     }
 
-    existingPol.d_custom.reserve(existingPol.d_custom.size() + pol.d_custom.size());
-
-    std::move(pol.d_custom.begin(), pol.d_custom.end(), std::back_inserter(existingPol.d_custom));
+    addCustom(existingPol, pol);
   }
   else {
     pol.d_zoneData = d_zoneData;
@@ -446,18 +453,20 @@ bool DNSFilterEngine::Zone::rmNameTrigger(std::unordered_map<DNSName, Policy>& m
   /* for custom types, we might have more than one type,
      and then we need to remove only the right ones. */
   bool result = false;
-  for (const auto& toRemove : pol.d_custom) {
-    for (auto it = existing.d_custom.begin(); it != existing.d_custom.end(); ++it) {
-      if (**it == *toRemove) {
-        existing.d_custom.erase(it);
-        result = true;
-        break;
+  if (pol.d_custom && existing.d_custom) {
+    for (const auto& toRemove : *pol.d_custom) {
+      for (auto it = existing.d_custom->begin(); it != existing.d_custom->end(); ++it) {
+        if (**it == *toRemove) {
+          existing.d_custom->erase(it);
+          result = true;
+          break;
+        }
       }
     }
   }
 
   // No records left for this trigger?
-  if (existing.d_custom.empty()) {
+  if (existing.customRecordsSize() == 0) {
     map.erase(found);
     return true;
   }
@@ -482,18 +491,20 @@ bool DNSFilterEngine::Zone::rmNetmaskTrigger(NetmaskTree<Policy>& nmt, const Net
      and then we need to remove only the right ones. */
 
   bool result = false;
-  for (const auto& toRemove : pol.d_custom) {
-    for (auto it = existing.d_custom.begin(); it != existing.d_custom.end(); ++it) {
-      if (**it == *toRemove) {
-        existing.d_custom.erase(it);
-        result = true;
-        break;
+  if (pol.d_custom && existing.d_custom) {
+    for (const auto& toRemove : *pol.d_custom) {
+      for (auto it = existing.d_custom->begin(); it != existing.d_custom->end(); ++it) {
+        if (**it == *toRemove) {
+          existing.d_custom->erase(it);
+          result = true;
+          break;
+        }
       }
     }
   }
 
   // No records left for this trigger?
-  if (existing.d_custom.empty()) {
+  if (existing.customRecordsSize() == 0) {
     nmt.erase(netmask);
     return true;
   }
@@ -594,8 +605,11 @@ std::vector<DNSRecord> DNSFilterEngine::Policy::getCustomRecords(const DNSName&
   }
 
   std::vector<DNSRecord> result;
+  if (customRecordsSize() == 0) {
+    return result;
+  }
 
-  for (const auto& custom : d_custom) {
+  for (const auto& custom : *d_custom) {
     if (qtype != QType::ANY && qtype != custom->getType() && custom->getType() != QType::CNAME) {
       continue;
     }
diff --git a/pdns/recursordist/filterpo.hh b/pdns/recursordist/filterpo.hh
index f91858ac8e..d69b799e5d 100644
--- a/pdns/recursordist/filterpo.hh
+++ b/pdns/recursordist/filterpo.hh
@@ -113,15 +113,18 @@ public:
     }
 
     Policy(PolicyKind kind, PolicyType type, int32_t ttl = 0, std::shared_ptr<PolicyZoneData> data = nullptr, const std::vector<std::shared_ptr<const DNSRecordContent>>& custom = {}) :
-      d_custom(custom), d_zoneData(std::move(data)), d_ttl(ttl), d_kind(kind), d_type(type)
+      d_zoneData(std::move(data)), d_custom(nullptr), d_ttl(ttl), d_kind(kind), d_type(type)
     {
+      if (!custom.empty()) {
+        setCustom(custom);
+      }
     }
 
     ~Policy() = default;
 
     Policy(const Policy& rhs) :
-      d_custom(rhs.d_custom),
       d_zoneData(rhs.d_zoneData),
+      d_custom(rhs.d_custom ? make_unique<CustomData>(*rhs.d_custom) : nullptr),
       d_hitdata(rhs.d_hitdata ? make_unique<HitData>(*rhs.d_hitdata) : nullptr),
       d_ttl(rhs.d_ttl),
       d_kind(rhs.d_kind),
@@ -132,7 +135,9 @@ public:
     Policy& operator=(const Policy& rhs)
     {
       if (this != &rhs) {
-        d_custom = rhs.d_custom;
+        if (rhs.d_custom) {
+          d_custom = make_unique<CustomData>(*rhs.d_custom);
+        }
         d_zoneData = rhs.d_zoneData;
         if (rhs.d_hitdata) {
           d_hitdata = make_unique<HitData>(*rhs.d_hitdata);
@@ -232,8 +237,11 @@ public:
     [[nodiscard]] std::vector<DNSRecord> getCustomRecords(const DNSName& qname, uint16_t qtype) const;
     [[nodiscard]] std::vector<DNSRecord> getRecords(const DNSName& qname) const;
 
-    std::vector<std::shared_ptr<const DNSRecordContent>> d_custom;
     std::shared_ptr<PolicyZoneData> d_zoneData{nullptr};
+
+    using CustomData = std::vector<std::shared_ptr<const DNSRecordContent>>;
+    std::unique_ptr<CustomData> d_custom;
+
     struct HitData
     {
       DNSName d_trigger;
@@ -254,6 +262,19 @@ public:
       }
     }
 
+    void setCustom(const CustomData& custom)
+    {
+      d_custom = make_unique<CustomData>(custom);
+    }
+
+    [[nodiscard]] size_t customRecordsSize() const
+    {
+      if (d_custom) {
+        return d_custom->size();
+      }
+      return 0;
+    }
+
     void setHitData(const DNSName& name, const string& hit)
     {
       HitData hitdata{name, hit};
diff --git a/pdns/recursordist/lua-recursor4.cc b/pdns/recursordist/lua-recursor4.cc
index 8062cb9061..1177cf05cb 100644
--- a/pdns/recursordist/lua-recursor4.cc
+++ b/pdns/recursordist/lua-recursor4.cc
@@ -239,19 +239,21 @@ void RecursorLua4::postPrepareContext()
         return result;
       }
 
-      for (const auto& dr : pol.d_custom) {
-        if (!result.empty()) {
-          result += "\n";
+      if (pol.d_custom) {
+        for (const auto& dnsRecord : *pol.d_custom) {
+          if (!result.empty()) {
+            result += "\n";
+          }
+          result += dnsRecord->getZoneRepresentation();
         }
-        result += dr->getZoneRepresentation();
       }
 
       return result;
     },
     [](DNSFilterEngine::Policy& pol, const std::string& content) {
       // Only CNAMES for now, when we ever add a d_custom_type, there will be pain
-      pol.d_custom.clear();
-      pol.d_custom.push_back(DNSRecordContent::make(QType::CNAME, QClass::IN, content));
+      pol.d_custom = make_unique<DNSFilterEngine::Policy::CustomData>();
+      pol.d_custom->push_back(DNSRecordContent::make(QType::CNAME, QClass::IN, content));
     }
   );
   d_lw->registerFunction("getDH", &DNSQuestion::getDH);
diff --git a/pdns/recursordist/rec-lua-conf.cc b/pdns/recursordist/rec-lua-conf.cc
index e862da35f9..bd080360f5 100644
--- a/pdns/recursordist/rec-lua-conf.cc
+++ b/pdns/recursordist/rec-lua-conf.cc
@@ -109,8 +109,11 @@ static void parseRPZParameters(rpzOptions_t& have, std::shared_ptr<DNSFilterEngi
     defpol->d_kind = (DNSFilterEngine::PolicyKind)boost::get<uint32_t>(have["defpol"]);
     defpol->setName(polName);
     if (defpol->d_kind == DNSFilterEngine::PolicyKind::Custom) {
-      defpol->d_custom.push_back(DNSRecordContent::make(QType::CNAME, QClass::IN,
-                                                        boost::get<string>(have["defcontent"])));
+      if (!defpol->d_custom) {
+        defpol->d_custom = make_unique<DNSFilterEngine::Policy::CustomData>();
+      }
+      defpol->d_custom->push_back(DNSRecordContent::make(QType::CNAME, QClass::IN,
+                                                         boost::get<string>(have["defcontent"])));
 
       if (have.count("defttl") != 0) {
         defpol->d_ttl = static_cast<int32_t>(boost::get<uint32_t>(have["defttl"]));
diff --git a/pdns/recursordist/rpzloader.cc b/pdns/recursordist/rpzloader.cc
index 33c9ae27a5..ef04c8baf9 100644
--- a/pdns/recursordist/rpzloader.cc
+++ b/pdns/recursordist/rpzloader.cc
@@ -146,7 +146,10 @@ static void RPZRecordToPolicy(const DNSRecord& dr, std::shared_ptr<DNSFilterEngi
     }
     else {
       pol.d_kind = DNSFilterEngine::PolicyKind::Custom;
-      pol.d_custom.emplace_back(dr.getContent());
+      if (!pol.d_custom) {
+        pol.d_custom = make_unique<DNSFilterEngine::Policy::CustomData>();
+      }
+      pol.d_custom->emplace_back(dr.getContent());
       // cerr<<"Wants custom "<<crcTarget<<" for "<<dr.d_name<<": ";
     }
   }
@@ -157,7 +160,10 @@ static void RPZRecordToPolicy(const DNSRecord& dr, std::shared_ptr<DNSFilterEngi
     }
     else {
       pol.d_kind = DNSFilterEngine::PolicyKind::Custom;
-      pol.d_custom.emplace_back(dr.getContent());
+      if (!pol.d_custom) {
+        pol.d_custom = make_unique<DNSFilterEngine::Policy::CustomData>();
+      }
+      pol.d_custom->emplace_back(dr.getContent());
       // cerr<<"Wants custom "<<dr.d_content->getZoneRepresentation()<<" for "<<dr.d_name<<": ";
     }
   }