From: Charles-Henri Bruyand Date: Tue, 6 Dec 2022 11:20:26 +0000 (+0100) Subject: auth: add byteslimit support to lua ifurlup() function X-Git-Tag: dnsdist-1.8.0-rc1~200^2 X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=refs%2Fpull%2F12277%2Fhead;p=thirdparty%2Fpdns.git auth: add byteslimit support to lua ifurlup() function --- diff --git a/docs/lua-records/functions.rst b/docs/lua-records/functions.rst index c0ccb7c2e0..c6f9f8f109 100644 --- 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 9526a3816f..d221ce92be 100644 --- 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(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 d5981c4e41..02b5204562 100644 --- 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(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(clientp); + if (us->d_byteslimit > 0 && static_cast(dlnow) > us->d_byteslimit) { + return static_cast(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(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 b59656ae71..c5bbf3316e 100644 --- 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(); };