From 42c779f503eb8437b6232773a4a2472306cc9f3d Mon Sep 17 00:00:00 2001 From: Vasil Dimov Date: Fri, 4 Dec 2020 15:11:25 +0100 Subject: net: extend Sock with methods for robust send & read until terminator Introduce two high level, convenience methods in the `Sock` class: * `SendComplete()`: keep trying to send the specified data until either successfully sent all of it, timeout or interrupted. * `RecvUntilTerminator()`: read until a terminator is encountered (never after it), timeout or interrupted. These will be convenient in the I2P SAM implementation. `SendComplete()` can also be used in the SOCKS5 implementation instead of calling `send()` directly. --- src/util/sock.cpp | 125 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/util/sock.h | 29 +++++++++++++ 2 files changed, 154 insertions(+) (limited to 'src/util') diff --git a/src/util/sock.cpp b/src/util/sock.cpp index d8fd9d4138..36493ec241 100644 --- a/src/util/sock.cpp +++ b/src/util/sock.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -12,12 +13,18 @@ #include #include #include +#include #include #ifdef USE_POLL #include #endif +static inline bool IOErrorIsPermanent(int err) +{ + return err != WSAEAGAIN && err != WSAEINTR && err != WSAEWOULDBLOCK && err != WSAEINPROGRESS; +} + Sock::Sock() : m_socket(INVALID_SOCKET) {} Sock::Sock(SOCKET s) : m_socket(s) {} @@ -125,6 +132,124 @@ bool Sock::Wait(std::chrono::milliseconds timeout, Event requested, Event* occur #endif /* USE_POLL */ } +void Sock::SendComplete(const std::string& data, + std::chrono::milliseconds timeout, + CThreadInterrupt& interrupt) const +{ + const auto deadline = GetTime() + timeout; + size_t sent{0}; + + for (;;) { + const ssize_t ret{Send(data.data() + sent, data.size() - sent, MSG_NOSIGNAL)}; + + if (ret > 0) { + sent += static_cast(ret); + if (sent == data.size()) { + break; + } + } else { + const int err{WSAGetLastError()}; + if (IOErrorIsPermanent(err)) { + throw std::runtime_error(strprintf("send(): %s", NetworkErrorString(err))); + } + } + + const auto now = GetTime(); + + if (now >= deadline) { + throw std::runtime_error(strprintf( + "Send timeout (sent only %u of %u bytes before that)", sent, data.size())); + } + + if (interrupt) { + throw std::runtime_error(strprintf( + "Send interrupted (sent only %u of %u bytes before that)", sent, data.size())); + } + + // Wait for a short while (or the socket to become ready for sending) before retrying + // if nothing was sent. + const auto wait_time = std::min(deadline - now, std::chrono::milliseconds{MAX_WAIT_FOR_IO}); + Wait(wait_time, SEND); + } +} + +std::string Sock::RecvUntilTerminator(uint8_t terminator, + std::chrono::milliseconds timeout, + CThreadInterrupt& interrupt) const +{ + const auto deadline = GetTime() + timeout; + std::string data; + bool terminator_found{false}; + + // We must not consume any bytes past the terminator from the socket. + // One option is to read one byte at a time and check if we have read a terminator. + // However that is very slow. Instead, we peek at what is in the socket and only read + // as many bytes as possible without crossing the terminator. + // Reading 64 MiB of random data with 262526 terminator chars takes 37 seconds to read + // one byte at a time VS 0.71 seconds with the "peek" solution below. Reading one byte + // at a time is about 50 times slower. + + for (;;) { + char buf[512]; + + const ssize_t peek_ret{Recv(buf, sizeof(buf), MSG_PEEK)}; + + switch (peek_ret) { + case -1: { + const int err{WSAGetLastError()}; + if (IOErrorIsPermanent(err)) { + throw std::runtime_error(strprintf("recv(): %s", NetworkErrorString(err))); + } + break; + } + case 0: + throw std::runtime_error("Connection unexpectedly closed by peer"); + default: + auto end = buf + peek_ret; + auto terminator_pos = std::find(buf, end, terminator); + terminator_found = terminator_pos != end; + + const size_t try_len{terminator_found ? terminator_pos - buf + 1 : + static_cast(peek_ret)}; + + const ssize_t read_ret{Recv(buf, try_len, 0)}; + + if (read_ret < 0 || static_cast(read_ret) != try_len) { + throw std::runtime_error( + strprintf("recv() returned %u bytes on attempt to read %u bytes but previous " + "peek claimed %u bytes are available", + read_ret, try_len, peek_ret)); + } + + // Don't include the terminator in the output. + const size_t append_len{terminator_found ? try_len - 1 : try_len}; + + data.append(buf, buf + append_len); + + if (terminator_found) { + return data; + } + } + + const auto now = GetTime(); + + if (now >= deadline) { + throw std::runtime_error(strprintf( + "Receive timeout (received %u bytes without terminator before that)", data.size())); + } + + if (interrupt) { + throw std::runtime_error(strprintf( + "Receive interrupted (received %u bytes without terminator before that)", + data.size())); + } + + // Wait for a short while (or the socket to become ready for reading) before retrying. + const auto wait_time = std::min(deadline - now, std::chrono::milliseconds{MAX_WAIT_FOR_IO}); + Wait(wait_time, RECV); + } +} + #ifdef WIN32 std::string NetworkErrorString(int err) { diff --git a/src/util/sock.h b/src/util/sock.h index 2d9cac14af..209d30def4 100644 --- a/src/util/sock.h +++ b/src/util/sock.h @@ -6,6 +6,7 @@ #define BITCOIN_UTIL_SOCK_H #include +#include #include #include @@ -114,6 +115,34 @@ public: Event requested, Event* occurred = nullptr) const; + /* Higher level, convenience, methods. These may throw. */ + + /** + * Send the given data, retrying on transient errors. + * @param[in] data Data to send. + * @param[in] timeout Timeout for the entire operation. + * @param[in] interrupt If this is signaled then the operation is canceled. + * @throws std::runtime_error if the operation cannot be completed. In this case only some of + * the data will be written to the socket. + */ + virtual void SendComplete(const std::string& data, + std::chrono::milliseconds timeout, + CThreadInterrupt& interrupt) const; + + /** + * Read from socket until a terminator character is encountered. Will never consume bytes past + * the terminator from the socket. + * @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. + * @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; + private: /** * Contained socket. `INVALID_SOCKET` designates the object is empty. -- cgit v1.2.3