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.
This commit is contained in:
Vasil Dimov 2020-12-04 15:11:25 +01:00
parent ea1845315a
commit 42c779f503
No known key found for this signature in database
GPG Key ID: 54DF06F64B55CBBF
3 changed files with 163 additions and 0 deletions

View File

@ -44,6 +44,7 @@ typedef unsigned int SOCKET;
#define WSAEINVAL EINVAL
#define WSAEALREADY EALREADY
#define WSAEWOULDBLOCK EWOULDBLOCK
#define WSAEAGAIN EAGAIN
#define WSAEMSGSIZE EMSGSIZE
#define WSAEINTR EINTR
#define WSAEINPROGRESS EINPROGRESS
@ -51,6 +52,14 @@ typedef unsigned int SOCKET;
#define WSAENOTSOCK EBADF
#define INVALID_SOCKET (SOCKET)(~0)
#define SOCKET_ERROR -1
#else
#ifndef WSAEAGAIN
#ifdef EAGAIN
#define WSAEAGAIN EAGAIN
#else
#define WSAEAGAIN WSAEWOULDBLOCK
#endif
#endif
#endif
#ifdef WIN32

View File

@ -4,6 +4,7 @@
#include <compat.h>
#include <logging.h>
#include <threadinterrupt.h>
#include <tinyformat.h>
#include <util/sock.h>
#include <util/system.h>
@ -12,12 +13,18 @@
#include <codecvt>
#include <cwchar>
#include <locale>
#include <stdexcept>
#include <string>
#ifdef USE_POLL
#include <poll.h>
#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<std::chrono::milliseconds>() + 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<size_t>(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<std::chrono::milliseconds>();
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<std::chrono::milliseconds>() + 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<size_t>(peek_ret)};
const ssize_t read_ret{Recv(buf, try_len, 0)};
if (read_ret < 0 || static_cast<size_t>(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<std::chrono::milliseconds>();
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)
{

View File

@ -6,6 +6,7 @@
#define BITCOIN_UTIL_SOCK_H
#include <compat.h>
#include <threadinterrupt.h>
#include <util/time.h>
#include <chrono>
@ -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.