diff --git a/src/shared/file_send.c b/src/shared/file_send.c index f064f93..b5224bd 100644 --- a/src/shared/file_send.c +++ b/src/shared/file_send.c @@ -166,6 +166,8 @@ bool file_send_sendfile_with_skip(File* file, int file_descriptor, bool use_meta } struct pollfd pfd = {.fd = file_descriptor, .events = POLLOUT}; int polled = poll(&pfd, 1, timeout); + if (polled < 0 && errno == EINTR) + continue; if (polled <= 0 || (pfd.revents & (POLLERR | POLLHUP | POLLNVAL))) { close(fd); return false; @@ -183,6 +185,7 @@ bool file_send_sendfile_with_skip(File* file, int file_descriptor, bool use_meta return false; } protocol_note_bytes_written((unsigned long long)sent); + protocol_throttle_bytes((size_t)sent); } close(fd); diff --git a/src/shared/protocol.c b/src/shared/protocol.c index 2a3f844..a78fa0a 100644 --- a/src/shared/protocol.c +++ b/src/shared/protocol.c @@ -301,6 +301,14 @@ static ProtocolSession* legacy_session(int read_fd, int write_fd) { return &legacy_io_session; } +/* Pace an out-of-band write that bypassed protocol_send_n_data (the plaintext + * sendfile fast path). The bound/legacy session is resolved exactly as + * send_n_data resolves it, so the same token-bucket state is throttled and the + * TLS and plaintext transports share identical --bwlimit semantics. */ +void protocol_throttle_bytes(size_t bytes) { + bw_throttle_session(legacy_session(-1, -1), bytes); +} + bool send_n_data(int file_descriptor, const void* data, size_t data_size) { return protocol_send_n_data(legacy_session(-1, file_descriptor), data, data_size); } @@ -439,12 +447,18 @@ static bool protocol_receive_n_data_until(ProtocolSession* session, void* data, } ssize_t bytes_received; - if (session->ssl) - bytes_received = SSL_read(session->ssl, (char*)data + total_bytes_received, - data_size - total_bytes_received); - else + if (session->ssl) { + /* SSL_read takes an int length; clamp a >INT_MAX request into chunks + * (mirrors the send path) so the size_t downcast can never truncate into + * a negative/partial read. */ + size_t ssl_chunk = data_size - total_bytes_received > (size_t)INT_MAX + ? (size_t)INT_MAX + : data_size - total_bytes_received; + bytes_received = SSL_read(session->ssl, (char*)data + total_bytes_received, (int)ssl_chunk); + } else { bytes_received = read(fd, (char*)data + total_bytes_received, data_size - total_bytes_received); + } if (bytes_received <= 0) { if (session->ssl) { int ssl_err = SSL_get_error(session->ssl, (int)bytes_received); diff --git a/src/shared/protocol.h b/src/shared/protocol.h index 4348f09..1ecdab3 100644 --- a/src/shared/protocol.h +++ b/src/shared/protocol.h @@ -220,6 +220,12 @@ SSL* io_get_ssl(void); unsigned long long protocol_bytes_written(void); unsigned long long protocol_bytes_read(void); void protocol_note_bytes_written(unsigned long long bytes); +/* Apply --bwlimit pacing to bytes written outside protocol_send_n_data (the + * plaintext zero-copy sendfile fast path). Resolves the bound/legacy session + * exactly as send_n_data does and runs the same token-bucket throttle, so the + * sendfile transport is paced identically to the buffered/TLS paths. A no-op + * when the effective session has no bandwidth limit. */ +void protocol_throttle_bytes(size_t bytes); void protocol_session_init(ProtocolSession* session, int read_fd, int write_fd); /* Transitional bridge for helpers whose signatures still carry only an fd. */ diff --git a/tests/test_protocol.c b/tests/test_protocol.c index 53e59d7..f1b3c64 100644 --- a/tests/test_protocol.c +++ b/tests/test_protocol.c @@ -2,6 +2,7 @@ #include "test_utils.h" #include #include +#include #include #include @@ -669,6 +670,50 @@ static void test_receive_status_keepalive_emits() { close(to_peer[1]); } +/* protocol_throttle_bytes() must apply the same token-bucket pacing as the + * buffered protocol send path, so the plaintext sendfile fast path honors + * --bwlimit exactly like the TLS path. With bwlimit=1 MB/s the initial burst + * is 100 KB (bwlimit/10); pacing 150 KB therefore owes ~50 KB of debt, i.e. a + * ~50 ms sleep. */ +static void test_protocol_throttle_bytes_paces() { + ProtocolSession session; + protocol_session_init(&session, -1, -1); + protocol_session_bind(&session); + protocol_session_set_bwlimit(&session, 1000000ULL); + + struct timespec start; + clock_gettime(CLOCK_MONOTONIC, &start); + protocol_throttle_bytes(150000); + struct timespec now; + clock_gettime(CLOCK_MONOTONIC, &now); + long long elapsed_ms = + (now.tv_sec - start.tv_sec) * 1000LL + (now.tv_nsec - start.tv_nsec) / 1000000LL; + /* Allow for scheduler slack but require the bulk of the expected 50 ms. */ + EXPECT_TRUE(elapsed_ms >= 40); + + protocol_session_unbind(); +} + +/* With no bandwidth limit the primitive must not sleep, however many bytes it + * is handed. */ +static void test_protocol_throttle_bytes_unlimited() { + ProtocolSession session; + protocol_session_init(&session, -1, -1); + protocol_session_bind(&session); + protocol_session_set_bwlimit(&session, 0); + + struct timespec start; + clock_gettime(CLOCK_MONOTONIC, &start); + protocol_throttle_bytes(100000000ULL); + struct timespec now; + clock_gettime(CLOCK_MONOTONIC, &now); + long long elapsed_ms = + (now.tv_sec - start.tv_sec) * 1000LL + (now.tv_nsec - start.tv_nsec) / 1000000LL; + EXPECT_TRUE(elapsed_ms < 50); + + protocol_session_unbind(); +} + void test_protocol() { test_send_receive_n_data(); test_send_receive_n_data_zero(); @@ -698,4 +743,6 @@ void test_protocol() { test_protocol_accounting_release_does_not_underflow(); test_receive_data_charge_follows_owning_session(); test_data_create_starts_uncharged_and_unowned(); + test_protocol_throttle_bytes_paces(); + test_protocol_throttle_bytes_unlimited(); }