From 31f9ddda2af15539e5f6f7404785739cfdd52a9d Mon Sep 17 00:00:00 2001 From: Archie Jaskowicz Date: Sun, 17 Dec 2023 21:19:23 +0000 Subject: [PATCH] fix: rcon no longer freezes if server sends no response. --- include/rcon.h | 6 +++++- src/main.cpp | 15 ++++++++++++--- src/rcon.cpp | 41 ++++++++++++++++++++++++++++++++--------- 3 files changed, 49 insertions(+), 13 deletions(-) diff --git a/include/rcon.h b/include/rcon.h index 6b8e0fe..ef158a1 100644 --- a/include/rcon.h +++ b/include/rcon.h @@ -27,6 +27,10 @@ enum data_type { struct rcon_packet { unsigned int bytes; unsigned char* data; + + ~rcon_packet() { + delete[] data; + } }; struct rcon_queued_request { @@ -42,7 +46,7 @@ class rcon { const std::string address; const unsigned int port; const std::string password; - unsigned int sock{0}; + int sock{0}; bool connected{false}; std::vector requests_queued{}; diff --git a/src/main.cpp b/src/main.cpp index 13589b7..84832e5 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -105,6 +105,11 @@ int main() { event.reply(dpp::message(data).set_flags(dpp::m_ephemeral)); }); } else if (event.command.get_command_name() == "command") { + if(FDR::config.allow_achievements) { + event.reply(dpp::message("Since achievements are enabled, you are not allowed to run this command!").set_flags(dpp::m_ephemeral)); + return; + } + if (std::find(event.command.member.get_roles().begin(), event.command.member.get_roles().end(), FDR::config.admin_role) == event.command.member.get_roles().end()) { event.reply(dpp::message("You do not have the required role to run this command!").set_flags(dpp::m_ephemeral)); return; @@ -113,6 +118,10 @@ int main() { auto command_to_run = std::get(event.get_parameter("cmd")); rcon_client.send_data("/command " + command_to_run, 3, data_type::SERVERDATA_EXECCOMMAND, [event](const std::string& data) { + if(data.empty()) { + return; + } + event.reply(dpp::message(data).set_flags(dpp::m_ephemeral)); }); } @@ -154,9 +163,9 @@ int main() { }); }, 120); - rcon_client.send_data("Factorio-Discord-Relay (FDR) has loaded!", 999, data_type::SERVERDATA_EXECCOMMAND); - - bot.message_create(dpp::message(FDR::config.msg_channel, "Factorio-Discord-Relay (FDR) has loaded!")); + bot.message_create(dpp::message(FDR::config.msg_channel, "Factorio-Discord-Relay (FDR) has loaded!"), [&rcon_client](const dpp::confirmation_callback_t& callback) { + rcon_client.send_data("Factorio-Discord-Relay (FDR) has loaded!", 999, data_type::SERVERDATA_EXECCOMMAND); + }); }); bot.start(dpp::st_wait); diff --git a/src/rcon.cpp b/src/rcon.cpp index d80fa99..6fd7223 100644 --- a/src/rcon.cpp +++ b/src/rcon.cpp @@ -92,13 +92,14 @@ bool rcon::connect() { // Make it non-blocking. fcntl(sock, F_SETFL, O_NONBLOCK); - // Set a timeout of 4 milliseconds. - struct timeval tv; + // Set a timeout of 4 seconds. + struct timeval tv{}; tv.tv_sec = 4; tv.tv_usec = 0; fd_set fds; FD_ZERO(&fds); FD_SET(sock, &fds); + setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); // Create temp status int status = -1; @@ -110,7 +111,7 @@ bool rcon::connect() { } } - status = select(sock +1, NULL, &fds, NULL, &tv); + status = select(sock +1, nullptr, &fds, nullptr, &tv); fcntl(sock, F_SETFL, fcntl(sock, F_GETFL, 0) & ~O_NONBLOCK); // If status wasn't zero, we successfully connected. @@ -119,7 +120,7 @@ bool rcon::connect() { void rcon::form_packet(unsigned char packet[], const std::string& data, int32_t id, int32_t type) { const char nullbytes[] = {'\x00', '\x00'}; - const int32_t min_size = sizeof(id) + sizeof(type) + sizeof(nullbytes); + const int32_t min_size = sizeof(id) + sizeof(type) + sizeof(nullbytes); // 10 bytes. const int32_t data_size = static_cast(data.size()) + min_size; if(data_size > 4096) { @@ -134,13 +135,10 @@ void rcon::form_packet(unsigned char packet[], const std::string& data, int32_t packet[4] = id; packet[8] = type; - int write_nullbytes_at{0}; - const char* data_chars = data.c_str(); for(int i = 0; i < data_size; i++) { packet[12 + i] = data_chars[i]; - write_nullbytes_at = 13 + i; } } @@ -149,6 +147,11 @@ std::string rcon::receive_information(int32_t id) { // it should really just keep going for a certain amount of seconds. for(int i=0; i < 500; i++) { rcon_packet packet = read_packet(); + + if(packet.bytes == 0) { + return ""; + } + unsigned char* buffer = packet.data; int offset = packet.bytes - HEADER_SIZE + 3; @@ -167,11 +170,25 @@ std::string rcon::receive_information(int32_t id) { rcon_packet rcon::read_packet() { size_t packet_length = read_packet_length(); + + if(packet_length == 0) { + return {0, nullptr}; + } + auto* buffer = new unsigned char[packet_length]{0}; unsigned int bytes = 0; do { - bytes += ::recv(sock, buffer, packet_length - bytes, 0); + size_t recv_bytes = ::recv(sock, buffer, packet_length - bytes, 0); + if(recv_bytes == -1) { + if ((errno != EAGAIN) && (errno != EWOULDBLOCK)) { + std::cout << "Did not receive a packet in time. Did the server send a response?"; + return {0, nullptr}; + } + } + + bytes += recv_bytes; + } while(bytes < packet_length); return {bytes, buffer}; @@ -179,7 +196,13 @@ rcon_packet rcon::read_packet() { const size_t rcon::read_packet_length() { unsigned char* buffer = new unsigned char[4]{0}; - ::recv(sock, buffer, 4, 0); + size_t recv_bytes = ::recv(sock, buffer, 4, 0); + if(recv_bytes == -1) { + if ((errno != EAGAIN) && (errno != EWOULDBLOCK)) { + std::cout << "Did not receive a packet in time. Did the server send a response?"; + return 0; + } + } const size_t len = byte32_to_int(buffer); delete[] buffer; return len;