auth: add byteslimit support to lua ifurlup() function

diff --git a/docs/lua-records/functions.rst b/docs/lua-records/functions.rst
index c0ccb7c2e04273fe25962a9d7acca8870f0cd575..c6f9f8f109dc096efa1a3312f50646894abf0b21 100644 (file)
--- a/docs/lua-records/functions.rst
+++ b/docs/lua-records/functions.rst
@@ -86,6 +86,7 @@ Record creation functions
   - ``timeout``: Maximum time in seconds that you allow the check to take (default 2)
   - ``stringmatch``: check ``url`` for this string, only declare 'up' if found
   - ``useragent``: Set the HTTP "User-Agent" header in the requests. By default it is set to "PowerDNS Authoritative Server"
+  - ``byteslimit``: Limit the maximum download size to ``byteslimit`` bytes (default 0 meaning no limit).
 
   An example of a list of address sets:
 

diff --git a/pdns/lua-record.cc b/pdns/lua-record.cc
index 9526a3816fdfe49ad3503f3739c661188c8c0fbc..d221ce92be0b3e0811d8f5010da3dfd99e541c68 100644 (file)
--- a/pdns/lua-record.cc
+++ b/pdns/lua-record.cc
@@ -100,6 +100,10 @@ private:
       if (cd.opts.count("useragent")) {
         useragent = cd.opts.at("useragent");
       }
+      size_t byteslimit = 0;
+      if (cd.opts.count("byteslimit")) {
+        byteslimit = static_cast<size_t>(std::atoi(cd.opts.at("byteslimit").c_str()));
+      }
       MiniCurl mc(useragent);
 
       string content;
@@ -113,10 +117,10 @@ private:
 
       if (cd.opts.count("source")) {
         ComboAddress src(cd.opts.at("source"));
-        content=mc.getURL(cd.url, rem, &src, timeout);
+        content=mc.getURL(cd.url, rem, &src, timeout, false, false, byteslimit);
       }
       else {
-        content=mc.getURL(cd.url, rem, nullptr, timeout);
+        content=mc.getURL(cd.url, rem, nullptr, timeout, false, false, byteslimit);
       }
       if (cd.opts.count("stringmatch") && content.find(cd.opts.at("stringmatch")) == string::npos) {
         throw std::runtime_error(boost::str(boost::format("unable to match content with `%s`") % cd.opts.at("stringmatch")));

diff --git a/pdns/minicurl.cc b/pdns/minicurl.cc
index d5981c4e417d9bf325dfda05b402bb56806d241a..02b5204562fc7f98fa7113c72a7526aab7d29333 100644 (file)
--- a/pdns/minicurl.cc
+++ b/pdns/minicurl.cc
@@ -57,9 +57,23 @@ MiniCurl::~MiniCurl()
 
 size_t MiniCurl::write_callback(char *ptr, size_t size, size_t nmemb, void *userdata)
 {
-  MiniCurl* us = (MiniCurl*)userdata;
-  us->d_data.append(ptr, size*nmemb);
-  return size*nmemb;
+  if (userdata != nullptr) {
+    MiniCurl* us = static_cast<MiniCurl*>(userdata);
+    us->d_data.append(ptr, size * nmemb);
+    return size * nmemb;
+  }
+  return 0;
+}
+
+size_t MiniCurl::progress_callback(void *clientp, curl_off_t dltotal, curl_off_t dlnow, curl_off_t ultotal, curl_off_t ulnow)
+{
+  if (clientp != nullptr) {
+    MiniCurl* us = static_cast<MiniCurl*>(clientp);
+    if (us->d_byteslimit > 0 && static_cast<size_t>(dlnow) > us->d_byteslimit) {
+      return static_cast<size_t>(dlnow);
+    }
+  }
+  return 0;
 }
 
 static string extractHostFromURL(const std::string& url)
@@ -75,7 +89,7 @@ static string extractHostFromURL(const std::string& url)
   return url.substr(pos, endpos-pos);
 }
 
-void MiniCurl::setupURL(const std::string& str, const ComboAddress* rem, const ComboAddress* src, int timeout, bool fastopen, bool verify)
+void MiniCurl::setupURL(const std::string& str, const ComboAddress* rem, const ComboAddress* src, int timeout, size_t byteslimit, bool fastopen, bool verify)
 {
   if(rem) {
     struct curl_slist *hostlist = nullptr; // THIS SHOULD BE FREED
@@ -112,6 +126,15 @@ void MiniCurl::setupURL(const std::string& str, const ComboAddress* rem, const C
   curl_easy_setopt(d_curl, CURLOPT_URL, str.c_str());
   curl_easy_setopt(d_curl, CURLOPT_WRITEFUNCTION, write_callback);
   curl_easy_setopt(d_curl, CURLOPT_WRITEDATA, this);
+
+  d_byteslimit = byteslimit;
+  if (d_byteslimit > 0) {
+    /* enable progress meter */
+    curl_easy_setopt(d_curl, CURLOPT_NOPROGRESS, 0L);
+    curl_easy_setopt(d_curl, CURLOPT_XFERINFOFUNCTION, progress_callback);
+    curl_easy_setopt(d_curl, CURLOPT_XFERINFODATA, this);
+  }
+
   curl_easy_setopt(d_curl, CURLOPT_TIMEOUT, static_cast<long>(timeout));
 #if defined(CURL_AT_LEAST_VERSION)
 #if CURL_AT_LEAST_VERSION(7, 49, 0) && defined(__linux__)
@@ -122,24 +145,24 @@ void MiniCurl::setupURL(const std::string& str, const ComboAddress* rem, const C
   d_data.clear();
 }
 
-std::string MiniCurl::getURL(const std::string& str, const ComboAddress* rem, const ComboAddress* src, int timeout, bool fastopen, bool verify)
+std::string MiniCurl::getURL(const std::string& str, const ComboAddress* rem, const ComboAddress* src, int timeout, bool fastopen, bool verify, size_t byteslimit)
 {
-  setupURL(str, rem, src, timeout, fastopen, verify);
+  setupURL(str, rem, src, timeout, byteslimit, fastopen, verify);
   auto res = curl_easy_perform(d_curl);
   long http_code = 0;
   curl_easy_getinfo(d_curl, CURLINFO_RESPONSE_CODE, &http_code);
 
-  if(res != CURLE_OK || http_code != 200)  {
+  if ((res != CURLE_OK && res != CURLE_ABORTED_BY_CALLBACK) || http_code != 200)  {
     throw std::runtime_error("Unable to retrieve URL ("+std::to_string(http_code)+"): "+string(curl_easy_strerror(res)));
   }
-  std::string ret=d_data;
+  std::string ret = d_data;
   d_data.clear();
   return ret;
 }
 
 std::string MiniCurl::postURL(const std::string& str, const std::string& postdata, MiniCurlHeaders& headers, int timeout, bool fastopen, bool verify)
 {
-  setupURL(str, nullptr, nullptr, timeout, fastopen, verify);
+  setupURL(str, nullptr, nullptr, timeout, 0, fastopen, verify);
   setHeaders(headers);
   curl_easy_setopt(d_curl, CURLOPT_POSTFIELDSIZE, postdata.size());
   curl_easy_setopt(d_curl, CURLOPT_POSTFIELDS, postdata.c_str());

diff --git a/pdns/minicurl.hh b/pdns/minicurl.hh
index b59656ae71d8d2f8307488850bc46238c66ec245..c5bbf3316e3590844f7bf9f282a54953102442d0 100644 (file)
--- a/pdns/minicurl.hh
+++ b/pdns/minicurl.hh
@@ -38,14 +38,16 @@ public:
   MiniCurl(const string& useragent="MiniCurl/0.0");
   ~MiniCurl();
   MiniCurl& operator=(const MiniCurl&) = delete;
-  std::string getURL(const std::string& str, const ComboAddress* rem=nullptr, const ComboAddress* src=nullptr, int timeout = 2, bool fastopen = false, bool verify = false);
+  std::string getURL(const std::string& str, const ComboAddress* rem=nullptr, const ComboAddress* src=nullptr, int timeout = 2, bool fastopen = false, bool verify = false, size_t byteslimit = 0);
   std::string postURL(const std::string& str, const std::string& postdata, MiniCurlHeaders& headers, int timeout = 2, bool fastopen = false, bool verify = false);
 private:
   CURL *d_curl;
   static size_t write_callback(char *ptr, size_t size, size_t nmemb, void *userdata);
+  static size_t progress_callback(void *clientp, curl_off_t dltotal, curl_off_t dlnow, curl_off_t ultotal, curl_off_t ulnow);
   std::string d_data;
+  size_t d_byteslimit;
   struct curl_slist* d_header_list = nullptr;
-  void setupURL(const std::string& str, const ComboAddress* rem, const ComboAddress* src, int timeout, bool fastopen, bool verify);
+  void setupURL(const std::string& str, const ComboAddress* rem, const ComboAddress* src, int timeout, size_t byteslimit, bool fastopen, bool verify);
   void setHeaders(const MiniCurlHeaders& headers);
   void clearHeaders();
 };