diff options
Diffstat (limited to 'src/util')
| -rw-r--r-- | src/util/sock.cpp | 10 | ||||
| -rw-r--r-- | src/util/sock.h | 5 |
2 files changed, 12 insertions, 3 deletions
diff --git a/src/util/sock.cpp b/src/util/sock.cpp index e13c52a16a..f9ecfef5d4 100644 --- a/src/util/sock.cpp +++ b/src/util/sock.cpp @@ -175,7 +175,8 @@ void Sock::SendComplete(const std::string& data, std::string Sock::RecvUntilTerminator(uint8_t terminator, std::chrono::milliseconds timeout, - CThreadInterrupt& interrupt) const + CThreadInterrupt& interrupt, + size_t max_data) const { const auto deadline = GetTime<std::chrono::milliseconds>() + timeout; std::string data; @@ -190,9 +191,14 @@ std::string Sock::RecvUntilTerminator(uint8_t terminator, // at a time is about 50 times slower. for (;;) { + if (data.size() >= max_data) { + throw std::runtime_error( + strprintf("Received too many bytes without a terminator (%u)", data.size())); + } + char buf[512]; - const ssize_t peek_ret{Recv(buf, sizeof(buf), MSG_PEEK)}; + const ssize_t peek_ret{Recv(buf, std::min(sizeof(buf), max_data - data.size()), MSG_PEEK)}; switch (peek_ret) { case -1: { diff --git a/src/util/sock.h b/src/util/sock.h index ecebb84205..4b0618dcff 100644 --- a/src/util/sock.h +++ b/src/util/sock.h @@ -135,13 +135,16 @@ public: * @param[in] terminator Character up to which to read from the socket. * @param[in] timeout Timeout for the entire operation. * @param[in] interrupt If this is signaled then the operation is canceled. + * @param[in] max_data The maximum amount of data (in bytes) to receive. If this many bytes + * are received and there is still no terminator, then this method will throw an exception. * @return The data that has been read, without the terminating character. * @throws std::runtime_error if the operation cannot be completed. In this case some bytes may * have been consumed from the socket. */ virtual std::string RecvUntilTerminator(uint8_t terminator, std::chrono::milliseconds timeout, - CThreadInterrupt& interrupt) const; + CThreadInterrupt& interrupt, + size_t max_data) const; /** * Check if still connected. |