diff options
-rw-r--r-- | src/mongo/util/bufreader.h | 13 | ||||
-rw-r--r-- | src/mongo/util/net/http_client_curl.cpp | 47 |
2 files changed, 50 insertions, 10 deletions
diff --git a/src/mongo/util/bufreader.h b/src/mongo/util/bufreader.h index 8c30070bada..b8d931add9a 100644 --- a/src/mongo/util/bufreader.h +++ b/src/mongo/util/bufreader.h @@ -106,6 +106,11 @@ public: invariant(_pos >= _start); } + /** back up to beginging of buffer */ + void rewindToStart() { + _pos = _start; + } + /** return current position pointer, and advance by len */ const void* skip(unsigned len) { ConstDataRangeCursor cdrc(_pos, _end); @@ -124,6 +129,14 @@ public: s = readCStr().toString(); } + /** + * Return a view of the next len bytes and advance by len. + */ + StringData readBytes(size_t len) { + // Note: the call to skip() includes a check that at least 'len' bytes remain in the buffer. + return StringData(reinterpret_cast<const char*>(skip(len)), len); + } + const void* pos() { return _pos; } diff --git a/src/mongo/util/net/http_client_curl.cpp b/src/mongo/util/net/http_client_curl.cpp index 3e6183d2ebf..82433b4caf0 100644 --- a/src/mongo/util/net/http_client_curl.cpp +++ b/src/mongo/util/net/http_client_curl.cpp @@ -53,6 +53,7 @@ #include "mongo/util/alarm.h" #include "mongo/util/alarm_runner_background_thread.h" #include "mongo/util/assert_util.h" +#include "mongo/util/bufreader.h" #include "mongo/util/concurrency/thread_pool.h" #include "mongo/util/functional.h" #include "mongo/util/net/hostandport.h" @@ -159,20 +160,40 @@ size_t WriteMemoryCallback(void* ptr, size_t size, size_t nmemb, void* data) { */ size_t ReadMemoryCallback(char* buffer, size_t size, size_t nitems, void* instream) { - auto* cdrc = reinterpret_cast<ConstDataRangeCursor*>(instream); + auto* bufReader = reinterpret_cast<BufReader*>(instream); size_t ret = 0; - if (cdrc->length() > 0) { - size_t readSize = std::min(size * nitems, cdrc->length()); - memcpy(buffer, cdrc->data(), readSize); - invariant(cdrc->advanceNoThrow(readSize).isOK()); + if (bufReader->remaining() > 0) { + size_t readSize = + std::min(size * nitems, static_cast<unsigned long>(bufReader->remaining())); + auto buf = bufReader->readBytes(readSize); + memcpy(buffer, buf.rawData(), readSize); ret = readSize; } return ret; } +/** + * Seek into for data to the remote side + */ +size_t SeekMemoryCallback(void* clientp, curl_off_t offset, int origin) { + + // Curl will call this in readrewind but only to reset the stream to the beginning + // In other protocols (like FTP, SSH) or HTTP resumption they may ask for partial buffers which + // we do not support. + if (offset != 0 || origin != SEEK_SET) { + return CURL_SEEKFUNC_CANTSEEK; + } + + auto* bufReader = reinterpret_cast<BufReader*>(clientp); + + bufReader->rewindToStart(); + + return CURL_SEEKFUNC_OK; +} + struct CurlEasyCleanup { void operator()(CURL* handle) { if (handle) { @@ -630,7 +651,7 @@ private: curl_easy_setopt(handle, CURLOPT_CONNECTTIMEOUT, longSeconds(_connectTimeout)); - ConstDataRangeCursor cdrc(cdr); + BufReader bufReader(cdr.data(), cdr.length()); switch (method) { case HttpMethod::kGET: uassert(ErrorCodes::BadValue, @@ -645,16 +666,22 @@ private: curl_easy_setopt(handle, CURLOPT_POST, 1); curl_easy_setopt(handle, CURLOPT_READFUNCTION, ReadMemoryCallback); - curl_easy_setopt(handle, CURLOPT_READDATA, &cdrc); - curl_easy_setopt(handle, CURLOPT_POSTFIELDSIZE, (long)cdrc.length()); + curl_easy_setopt(handle, CURLOPT_READDATA, &bufReader); + curl_easy_setopt(handle, CURLOPT_POSTFIELDSIZE, (long)bufReader.remaining()); + + curl_easy_setopt(handle, CURLOPT_SEEKFUNCTION, SeekMemoryCallback); + curl_easy_setopt(handle, CURLOPT_SEEKDATA, &bufReader); break; case HttpMethod::kPUT: curl_easy_setopt(handle, CURLOPT_POST, 0); curl_easy_setopt(handle, CURLOPT_PUT, 1); curl_easy_setopt(handle, CURLOPT_READFUNCTION, ReadMemoryCallback); - curl_easy_setopt(handle, CURLOPT_READDATA, &cdrc); - curl_easy_setopt(handle, CURLOPT_INFILESIZE_LARGE, (long)cdrc.length()); + curl_easy_setopt(handle, CURLOPT_READDATA, &bufReader); + curl_easy_setopt(handle, CURLOPT_INFILESIZE_LARGE, (long)bufReader.remaining()); + + curl_easy_setopt(handle, CURLOPT_SEEKFUNCTION, SeekMemoryCallback); + curl_easy_setopt(handle, CURLOPT_SEEKDATA, &bufReader); break; default: MONGO_UNREACHABLE; |