Add size limit

Signed-off-by: Otto Moerbeek <otto.moerbeek@open-xchange.com>

diff --git a/pdns/recursordist/lua-recursor4.cc b/pdns/recursordist/lua-recursor4.cc
index a91ca83fd0102071770c1388f653540cc8cc2211..8a278bb5246433c5c7f9f12bd870282c810e124a 100644 (file)
--- a/pdns/recursordist/lua-recursor4.cc
+++ b/pdns/recursordist/lua-recursor4.cc
@@ -523,9 +523,9 @@ void RecursorLua4::postPrepareContext() // NOLINT(readability-function-cognitive
       return {dir, name};
   });
 
-  d_lw->writeFunction("getNSSpeedTable", []() {
+  d_lw->writeFunction("getNSSpeedTable", [](size_t maxSize) {
     std::string ret;
-    auto number = SyncRes::getNSSpeedTable(ret);
+    auto number = SyncRes::getNSSpeedTable(maxSize, ret);
     return std::tuple<std::string, size_t>{ret, number};
   });
 

diff --git a/pdns/recursordist/syncres.cc b/pdns/recursordist/syncres.cc
index 9b09876602fcb6a4fff04e3c63f71b8be7a52076..a2d7350969b8e78b9e4a4f37341239b16f40c598 100644 (file)
--- a/pdns/recursordist/syncres.cc
+++ b/pdns/recursordist/syncres.cc
@@ -290,8 +290,14 @@ public:
     }
   }
 
-  size_t getPB(std::string& ret) const
+  size_t getPB(size_t maxSize, std::string& ret) const
   {
+    auto log = g_slog->withName("syncres")->withValues("maxSize", Logging::Loggable(maxSize));
+    log->info(Logr::Info, "Producing nsspeed dump");
+
+    // A observed average record size is 60;
+    size_t estimate = maxSize == 0 ? size() * 60 : maxSize + 4096; // We may overshoot (will be rolled back)
+
     protozero::pbf_builder<PBNSSpeedDump> full(ret);
     full.add_string(PBNSSpeedDump::required_string_version, getPDNSVersion());
     full.add_string(PBNSSpeedDump::required_string_identity, SyncRes::s_serverID);
@@ -300,11 +306,19 @@ public:
     full.add_string(PBNSSpeedDump::required_string_type, "PBNSSpeedDump");
 
     size_t count = 0;
+    ret.reserve(estimate);
+
     for (const auto& entry : *this) {
       protozero::pbf_builder<PBNSSpeedEntry> message(full, PBNSSpeedDump::repeated_message_nsspeedEntry);
       getEntry(message, &entry);
+      if (ret.size() > maxSize) {
+        message.rollback();
+        log->info(Logr::Info, "Produced nsspeed dump (max size reached)", "size", Logging::Loggable(ret.size()), "count", Logging::Loggable(count));
+        return count;
+      }
       ++count;
     }
+    log->info(Logr::Info, "Produced nsspeed dump", "size", Logging::Loggable(ret.size()), "count", Logging::Loggable(count));
     return count;
   }
 
@@ -426,10 +440,10 @@ public:
 
 static LockGuarded<nsspeeds_t> s_nsSpeeds;
 
-size_t SyncRes::getNSSpeedTable(std::string& ret)
+size_t SyncRes::getNSSpeedTable(size_t maxSize, std::string& ret)
 {
   const auto copy = *s_nsSpeeds.lock();
-  return copy.getPB(ret);
+  return copy.getPB(maxSize, ret);
 }
 
 size_t SyncRes::putIntoNSSpeedTable(const std::string& ret)

diff --git a/pdns/recursordist/syncres.hh b/pdns/recursordist/syncres.hh
index 0a661a71f5bd01840f95793e5d4e292da89915a2..91e800627125a9eff4cbff9a236f0d90142a89aa 100644 (file)
--- a/pdns/recursordist/syncres.hh
+++ b/pdns/recursordist/syncres.hh
@@ -170,7 +170,7 @@ public:
   static uint64_t doDumpSavedParentNSSets(int fileDesc);
   static uint64_t doDumpDoTProbeMap(int fileDesc);
 
-  static size_t getNSSpeedTable(std::string& ret);
+  static size_t getNSSpeedTable(size_t maxSize, std::string& ret);
   static size_t putIntoNSSpeedTable(const std::string& ret);
 
   static int getRootNS(struct timeval now, asyncresolve_t asyncCallback, unsigned int depth, Logr::log_t);