nrepl: cleanup socket handing, resolve spam edge-case

This commit is contained in:
Tyler Wilding
2024-06-02 12:53:33 -04:00
parent eebae4c524
commit ae27ef11e4
2 changed files with 61 additions and 29 deletions
+60 -29
View File
@@ -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<std::string> 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<std::string> 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<std::string> 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<int> 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<std::string> 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<std::string> 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<std::string> ReplServer::get_msg() {
}
}
}
++it;
}
return std::nullopt;
}
+1
View File
@@ -27,5 +27,6 @@ class ReplServer : public XSocketServer {
fd_set read_sockets;
std::set<int> client_sockets = {};
void error_response(int socket, const std::string& error);
void ping_response(int socket);
};