From ae27ef11e414d05bead012e49b67d34c6e203609 Mon Sep 17 00:00:00 2001 From: Tyler Wilding Date: Sun, 2 Jun 2024 12:53:33 -0400 Subject: [PATCH] nrepl: cleanup socket handing, resolve spam edge-case --- common/repl/nrepl/ReplServer.cpp | 89 +++++++++++++++++++++----------- common/repl/nrepl/ReplServer.h | 1 + 2 files changed, 61 insertions(+), 29 deletions(-) diff --git a/common/repl/nrepl/ReplServer.cpp b/common/repl/nrepl/ReplServer.cpp index 267d5fa872..b2274a4f5f 100644 --- a/common/repl/nrepl/ReplServer.cpp +++ b/common/repl/nrepl/ReplServer.cpp @@ -27,7 +27,18 @@ ReplServer::~ReplServer() { void ReplServer::post_init() { // Add the listening socket to our set of sockets - lg::info("[nREPL:{}:{}] awaiting connections", tcp_port, listening_socket); + lg::debug("[nREPL:{}:{}] awaiting connections", tcp_port, listening_socket); +} + +void ReplServer::error_response(int socket, const std::string& error) { + std::string msg = fmt::format("[ERROR]: {}", error); + auto resp = write_to_socket(socket, msg.c_str(), msg.size()); + if (resp == -1) { + lg::warn("[nREPL:{}] Client Disconnected: {}", tcp_port, address_to_string(addr), + ntohs(addr.sin_port), socket); + close_socket(socket); + client_sockets.erase(socket); + } } void ReplServer::ping_response(int socket) { @@ -48,7 +59,6 @@ std::optional ReplServer::get_msg() { // Add the server's main listening socket (where we accept clients from) FD_SET(listening_socket, &read_sockets); - int max_sd = listening_socket; for (const int& sock : client_sockets) { if (sock > max_sd) { @@ -60,12 +70,11 @@ std::optional ReplServer::get_msg() { } // Wait for activity on _something_, with a timeout so we don't get stuck here on exit. - struct timeval timeout; - timeout.tv_sec = 0; - timeout.tv_usec = 100000; + struct timeval timeout = {0, 100000}; auto activity = select(max_sd + 1, &read_sockets, NULL, NULL, &timeout); - - if (activity < 0) { // TODO - || error! + if (activity < 0 && errno != EINTR) { + lg::error("[nREPL:{}] select error, returned: {}, errno: {}", tcp_port, activity, + strerror(errno)); return std::nullopt; } @@ -74,44 +83,46 @@ std::optional ReplServer::get_msg() { socklen_t addr_len = sizeof(addr); auto new_socket = accept_socket(listening_socket, (sockaddr*)&addr, &addr_len); if (new_socket < 0) { - // TODO - handle error + if (new_socket != -1) { + lg::error("[nREPL:{}] accept error, returned: {}, errono: {}", tcp_port, new_socket, + strerror(errno)); + } } else { lg::info("[nREPL:{}]: New socket connection: {}:{}:{}", tcp_port, address_to_string(addr), ntohs(addr.sin_port), new_socket); - // Say hello ping_response(new_socket); // Track the new socket if ((int)client_sockets.size() < max_clients) { client_sockets.insert(new_socket); } else { - // TODO - Respond with NO + // Respond with NO and close the socket + lg::warn("[nREPL:{}]: Maximum clients reached. Rejecting connection.", tcp_port); + error_response(new_socket, "Maximum clients reached. Rejecting connection."); + close_socket(new_socket); } } } - // otherwise (and no matter what) check all the clients to see if they have sent us anything - // else its some IO operation on some other socket - // - // RACE - the first client wins - - // TODO - there are ways to do this with iterators but, couldn't figure it out! - std::vector sockets_to_scan(client_sockets.begin(), client_sockets.end()); - for (const int& sock : sockets_to_scan) { + // Check all clients for activity + for (auto it = client_sockets.begin(); it != client_sockets.end();) { + int sock = *it; if (FD_ISSET(sock, &read_sockets)) { // Attempt to read a header - // TODO - should this be in a loop? auto req_bytes = read_from_socket(sock, header_buffer.data(), header_buffer.size()); - if (req_bytes == 0) { - // Socket disconnected + if (req_bytes <= 0) { // TODO - add a queue of messages in the REPL::Wrapper so we can print _BEFORE_ the prompt // is output - lg::warn("[nREPL:{}] Client Disconnected: {}", tcp_port, address_to_string(addr), - ntohs(addr.sin_port), sock); - + if (req_bytes == 0) { + lg::warn("[nREPL:{}] Client Disconnected: {}", tcp_port, address_to_string(addr)); + } else { + lg::warn("[nREPL:{}] Error reading from socket on {}: {}", tcp_port, + address_to_string(addr), strerror(errno)); + } // Cleanup the socket and remove it from our set close_socket(sock); - client_sockets.erase(sock); + it = client_sockets.erase(it); // Erase and move to the next element + continue; } else { // Otherwise, process the message auto* header = (ReplServerHeader*)(header_buffer.data()); @@ -119,7 +130,12 @@ std::optional ReplServer::get_msg() { int expected_size = header->length; int got = 0; int tries = 0; + bool skip_to_next_socket = false; while (got < expected_size) { + if (want_exit_callback()) { + lg::warn("[nREPL:{}] Terminating nREPL early", tcp_port); + return std::nullopt; + } tries++; if (tries > 100) { break; @@ -131,11 +147,25 @@ std::optional ReplServer::get_msg() { tcp_port, got, expected_size, buffer.size()); return std::nullopt; } - auto x = read_from_socket(sock, buffer.data() + got, expected_size - got); - if (want_exit_callback()) { - return std::nullopt; + auto bytes_read = read_from_socket(sock, buffer.data() + got, expected_size - got); + if (bytes_read <= 0) { + if (bytes_read == 0) { + lg::warn("[nREPL:{}] Client Disconnected: {}", tcp_port, address_to_string(addr)); + } else { + lg::warn("[nREPL:{}] Error reading from socket on {}: {}", tcp_port, + address_to_string(addr), strerror(errno)); + } + close_socket(sock); + it = client_sockets.erase(it); // Erase and move to the next element + skip_to_next_socket = true; + break; } - got += x > 0 ? x : 0; + got += bytes_read; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + if (skip_to_next_socket) { + continue; } switch (header->type) { @@ -149,6 +179,7 @@ std::optional ReplServer::get_msg() { } } } + ++it; } return std::nullopt; } diff --git a/common/repl/nrepl/ReplServer.h b/common/repl/nrepl/ReplServer.h index 4cd5e09ba0..9372d5f22f 100644 --- a/common/repl/nrepl/ReplServer.h +++ b/common/repl/nrepl/ReplServer.h @@ -27,5 +27,6 @@ class ReplServer : public XSocketServer { fd_set read_sockets; std::set client_sockets = {}; + void error_response(int socket, const std::string& error); void ping_response(int socket); };