Skip to content

Commit 42c779f

Browse files
committed
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.
1 parent ea18453 commit 42c779f

File tree

3 files changed

+163
-0
lines changed

3 files changed

+163
-0
lines changed

src/compat.h

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,13 +44,22 @@ typedef unsigned int SOCKET;
4444
#define WSAEINVAL EINVAL
4545
#define WSAEALREADY EALREADY
4646
#define WSAEWOULDBLOCK EWOULDBLOCK
47+
#define WSAEAGAIN EAGAIN
4748
#define WSAEMSGSIZE EMSGSIZE
4849
#define WSAEINTR EINTR
4950
#define WSAEINPROGRESS EINPROGRESS
5051
#define WSAEADDRINUSE EADDRINUSE
5152
#define WSAENOTSOCK EBADF
5253
#define INVALID_SOCKET (SOCKET)(~0)
5354
#define SOCKET_ERROR -1
55+
#else
56+
#ifndef WSAEAGAIN
57+
#ifdef EAGAIN
58+
#define WSAEAGAIN EAGAIN
59+
#else
60+
#define WSAEAGAIN WSAEWOULDBLOCK
61+
#endif
62+
#endif
5463
#endif
5564

5665
#ifdef WIN32

src/util/sock.cpp

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
#include <compat.h>
66
#include <logging.h>
7+
#include <threadinterrupt.h>
78
#include <tinyformat.h>
89
#include <util/sock.h>
910
#include <util/system.h>
@@ -12,12 +13,18 @@
1213
#include <codecvt>
1314
#include <cwchar>
1415
#include <locale>
16+
#include <stdexcept>
1517
#include <string>
1618

1719
#ifdef USE_POLL
1820
#include <poll.h>
1921
#endif
2022

23+
static inline bool IOErrorIsPermanent(int err)
24+
{
25+
return err != WSAEAGAIN && err != WSAEINTR && err != WSAEWOULDBLOCK && err != WSAEINPROGRESS;
26+
}
27+
2128
Sock::Sock() : m_socket(INVALID_SOCKET) {}
2229

2330
Sock::Sock(SOCKET s) : m_socket(s) {}
@@ -125,6 +132,124 @@ bool Sock::Wait(std::chrono::milliseconds timeout, Event requested, Event* occur
125132
#endif /* USE_POLL */
126133
}
127134

135+
void Sock::SendComplete(const std::string& data,
136+
std::chrono::milliseconds timeout,
137+
CThreadInterrupt& interrupt) const
138+
{
139+
const auto deadline = GetTime<std::chrono::milliseconds>() + timeout;
140+
size_t sent{0};
141+
142+
for (;;) {
143+
const ssize_t ret{Send(data.data() + sent, data.size() - sent, MSG_NOSIGNAL)};
144+
145+
if (ret > 0) {
146+
sent += static_cast<size_t>(ret);
147+
if (sent == data.size()) {
148+
break;
149+
}
150+
} else {
151+
const int err{WSAGetLastError()};
152+
if (IOErrorIsPermanent(err)) {
153+
throw std::runtime_error(strprintf("send(): %s", NetworkErrorString(err)));
154+
}
155+
}
156+
157+
const auto now = GetTime<std::chrono::milliseconds>();
158+
159+
if (now >= deadline) {
160+
throw std::runtime_error(strprintf(
161+
"Send timeout (sent only %u of %u bytes before that)", sent, data.size()));
162+
}
163+
164+
if (interrupt) {
165+
throw std::runtime_error(strprintf(
166+
"Send interrupted (sent only %u of %u bytes before that)", sent, data.size()));
167+
}
168+
169+
// Wait for a short while (or the socket to become ready for sending) before retrying
170+
// if nothing was sent.
171+
const auto wait_time = std::min(deadline - now, std::chrono::milliseconds{MAX_WAIT_FOR_IO});
172+
Wait(wait_time, SEND);
173+
}
174+
}
175+
176+
std::string Sock::RecvUntilTerminator(uint8_t terminator,
177+
std::chrono::milliseconds timeout,
178+
CThreadInterrupt& interrupt) const
179+
{
180+
const auto deadline = GetTime<std::chrono::milliseconds>() + timeout;
181+
std::string data;
182+
bool terminator_found{false};
183+
184+
// We must not consume any bytes past the terminator from the socket.
185+
// One option is to read one byte at a time and check if we have read a terminator.
186+
// However that is very slow. Instead, we peek at what is in the socket and only read
187+
// as many bytes as possible without crossing the terminator.
188+
// Reading 64 MiB of random data with 262526 terminator chars takes 37 seconds to read
189+
// one byte at a time VS 0.71 seconds with the "peek" solution below. Reading one byte
190+
// at a time is about 50 times slower.
191+
192+
for (;;) {
193+
char buf[512];
194+
195+
const ssize_t peek_ret{Recv(buf, sizeof(buf), MSG_PEEK)};
196+
197+
switch (peek_ret) {
198+
case -1: {
199+
const int err{WSAGetLastError()};
200+
if (IOErrorIsPermanent(err)) {
201+
throw std::runtime_error(strprintf("recv(): %s", NetworkErrorString(err)));
202+
}
203+
break;
204+
}
205+
case 0:
206+
throw std::runtime_error("Connection unexpectedly closed by peer");
207+
default:
208+
auto end = buf + peek_ret;
209+
auto terminator_pos = std::find(buf, end, terminator);
210+
terminator_found = terminator_pos != end;
211+
212+
const size_t try_len{terminator_found ? terminator_pos - buf + 1 :
213+
static_cast<size_t>(peek_ret)};
214+
215+
const ssize_t read_ret{Recv(buf, try_len, 0)};
216+
217+
if (read_ret < 0 || static_cast<size_t>(read_ret) != try_len) {
218+
throw std::runtime_error(
219+
strprintf("recv() returned %u bytes on attempt to read %u bytes but previous "
220+
"peek claimed %u bytes are available",
221+
read_ret, try_len, peek_ret));
222+
}
223+
224+
// Don't include the terminator in the output.
225+
const size_t append_len{terminator_found ? try_len - 1 : try_len};
226+
227+
data.append(buf, buf + append_len);
228+
229+
if (terminator_found) {
230+
return data;
231+
}
232+
}
233+
234+
const auto now = GetTime<std::chrono::milliseconds>();
235+
236+
if (now >= deadline) {
237+
throw std::runtime_error(strprintf(
238+
"Receive timeout (received %u bytes without terminator before that)", data.size()));
239+
}
240+
241+
if (interrupt) {
242+
throw std::runtime_error(strprintf(
243+
"Receive interrupted (received %u bytes without terminator before that)",
244+
data.size()));
245+
}
246+
247+
// Wait for a short while (or the socket to become ready for reading) before retrying.
248+
const auto wait_time = std::min(deadline - now, std::chrono::milliseconds{MAX_WAIT_FOR_IO});
249+
Wait(wait_time, RECV);
250+
}
251+
}
252+
128253
#ifdef WIN32
129254
std::string NetworkErrorString(int err)
130255
{

src/util/sock.h

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#define BITCOIN_UTIL_SOCK_H
77

88
#include <compat.h>
9+
#include <threadinterrupt.h>
910
#include <util/time.h>
1011

1112
#include <chrono>
@@ -114,6 +115,34 @@ class Sock
114115
Event requested,
115116
Event* occurred = nullptr) const;
116117

118+
/* Higher level, convenience, methods. These may throw. */
119+
120+
/**
121+
* Send the given data, retrying on transient errors.
122+
* @param[in] data Data to send.
123+
* @param[in] timeout Timeout for the entire operation.
124+
* @param[in] interrupt If this is signaled then the operation is canceled.
125+
* @throws std::runtime_error if the operation cannot be completed. In this case only some of
126+
* the data will be written to the socket.
127+
*/
128+
virtual void SendComplete(const std::string& data,
129+
std::chrono::milliseconds timeout,
130+
CThreadInterrupt& interrupt) const;
131+
132+
/**
133+
* Read from socket until a terminator character is encountered. Will never consume bytes past
134+
* the terminator from the socket.
135+
* @param[in] terminator Character up to which to read from the socket.
136+
* @param[in] timeout Timeout for the entire operation.
137+
* @param[in] interrupt If this is signaled then the operation is canceled.
138+
* @return The data that has been read, without the terminating character.
139+
* @throws std::runtime_error if the operation cannot be completed. In this case some bytes may
140+
* have been consumed from the socket.
141+
*/
142+
virtual std::string RecvUntilTerminator(uint8_t terminator,
143+
std::chrono::milliseconds timeout,
144+
CThreadInterrupt& interrupt) const;
145+
117146
private:
118147
/**
119148
* Contained socket. `INVALID_SOCKET` designates the object is empty.

0 commit comments

Comments
 (0)