feat: Added heartbeats, fixed server being stuck in authentication, added some logging, rewrote unit test to properly test servers

Author: Jaskowicz1

include/rconpp/server.h

@@ -23,11 +23,13 @@
 namespace rconpp {
 
 struct connected_client {
-	struct sockaddr_in sock_info{};
+	sockaddr_in sock_info{};
 	int socket{0};
 	bool connected{false};
 
 	bool authenticated{false};
+
+	time_t last_heartbeat{0};
 };
 
 struct client_command {
@@ -47,9 +49,7 @@ class RCONPP_EXPORT rcon_server {
 #endif
 
 	std::thread accept_connections_runner;
-
-	// time_t is time since epoch in seconds (last time we ran).
-	std::unordered_map<int, time_t> client_socket_to_last_heartbeat;
+	std::mutex connected_clients_mutex;
 
 public:
 	bool online{false};
@@ -87,8 +87,9 @@ class RCONPP_EXPORT rcon_server {
 	 * @brief Disconnect a client from the server.
 	 *
 	 * @param client_socket The socket of the client to disconnect.
+	 * @param remove_after Should remove client from connected_clients after?
 	 */
-	void disconnect_client(const int client_socket);
+	void disconnect_client(int client_socket, bool remove_after = true);
 
 private:
 
@@ -101,29 +102,21 @@ class RCONPP_EXPORT rcon_server {
 	 */
 	bool startup_server();
 
-	/**
-	 * @brief Ask to receive information from the server for a specified ID.
-	 *
-	 * @param id The ID that we should except the server to return, alongside information.
-	 * @param type The type of packet that we should expect.
-	 *
-	 * @return Data given by the server.
-	 */
-	response receive_information(int32_t id, data_type type);
-
 	/**
 	 * @brief Gathers all the packet's content (based on the length returned by `read_packet_length`)
 	 *
 	 * @param client Client to read packet from.
 	 */
-	void read_packet(rconpp::connected_client client);
+	void read_packet(connected_client& client);
 
 	/**
 	 * @brief Sends a heartbeat to a client.
 	 *
 	 * @param client Client to send a heartbeat to.
+	 *
+	 * @returns bool, true is heartbeat was sent, otherwise false.
 	 */
-	void send_heartbeat(rconpp::connected_client client);
+	bool send_heartbeat(connected_client& client);
 };
 
 } // namespace rconpp

include/rconpp/utilities.h

@@ -101,6 +101,6 @@ RCONPP_EXPORT void report_error();
  *
  * @return The size (not length) of the packet.
  */
-RCONPP_EXPORT int read_packet_size(int socket, const std::function<void(const std::string_view log)>& on_log);
+RCONPP_EXPORT int read_packet_size(int socket);
 
 } // namespace rconpp

src/rconpp/client.cpp

@@ -6,6 +6,9 @@ rconpp::rcon_client::rcon_client(const std::string_view addr, const int _port, c
 }
 
 rconpp::rcon_client::~rcon_client() {
+	if (on_log) {
+		on_log("RCON client is shutting down.");
+	}
 	// Set connected to false, meaning no requests can be attempted during shutdown.
 	connected = false;
 
@@ -142,7 +145,12 @@ rconpp::response rconpp::rcon_client::receive_information(int32_t id, rconpp::da
 }
 
 rconpp::packet rconpp::rcon_client::read_packet() {
-	const int packet_size = read_packet_size(static_cast<int>(sock), on_log);
+	const int packet_size = read_packet_size(static_cast<int>(sock));
+
+	if (packet_size == -1) {
+		on_log("Did not receive a packet in time. Did the server send a response?");
+		return {};
+	}
 
 	packet temp_packet{};
 	temp_packet.length = packet_size + 4;
@@ -205,7 +213,7 @@ void rconpp::rcon_client::start(const bool return_after) {
 
 	// The server will send SERVERDATA_AUTH_RESPONSE once it's happy. If it's not -1, the server will have accepted us!
 	// We use the _sync method here to do a blocking call.
-	const response response = send_data_sync(password, 1, data_type::SERVERDATA_AUTH, true);
+	const response response = send_data_sync(password, 1, SERVERDATA_AUTH, true);
 
 	if (!response.server_responded) {
 		on_log("Login data was incorrect. RCON++ will now abort.");

src/rconpp/server.cpp

@@ -7,14 +7,18 @@ rconpp::rcon_server::rcon_server(const std::string_view addr, const int _port, c
 }
 
 rconpp::rcon_server::~rcon_server() {
+	if (on_log) {
+		on_log("RCON server is shutting down.");
+	}
+
 	// Set connected to false, meaning no requests can be attempted during shutdown.
 	online = false;
 
 	terminating.notify_all();
 
 	// Safely disconnect all clients from server.
 	for(const auto& client : connected_clients) {
-		disconnect_client(client.first);
+		disconnect_client(client.first, false);
 	}
 
 #ifdef _WIN32
@@ -26,10 +30,6 @@ rconpp::rcon_server::~rcon_server() {
 	if (accept_connections_runner.joinable()) {
 		accept_connections_runner.join();
 	}
-
-	if (heartbeat_runner.joinable()) {
-		heartbeat_runner.join();
-	}
 }
 
 bool rconpp::rcon_server::startup_server() {
@@ -84,27 +84,32 @@ bool rconpp::rcon_server::startup_server() {
 	return true;
 }
 
-void rconpp::rcon_server::disconnect_client(const int client_socket) {
+void rconpp::rcon_server::disconnect_client(const int client_socket, const bool remove_after /*= true*/) {
 
 #ifdef _WIN32
 	closesocket(client_socket);
 #else
 	close(client_socket);
 #endif
 
+	std::lock_guard guard(connected_clients_mutex);
+
 	connected_clients.at(client_socket).connected = false;
 
 	if (request_handlers.at(client_socket).joinable()) {
 		request_handlers.at(client_socket).join();
 	}
 
-	connected_clients.erase(client_socket);
+	if (remove_after) {
+		connected_clients.erase(client_socket);
+	}
 }
 
-void rconpp::rcon_server::read_packet(rconpp::connected_client client) {
-	const int packet_size = read_packet_size(static_cast<int>(sock), on_log);
+void rconpp::rcon_server::read_packet(connected_client& client) {
+	const int packet_size = read_packet_size(static_cast<int>(client.socket));
 
-	if (packet_size <= MIN_PACKET_SIZE) {
+	// Silently ignore packet size.
+	if (packet_size < MIN_PACKET_SIZE) {
 		return;
 	}
 
@@ -117,22 +122,26 @@ void rconpp::rcon_server::read_packet(rconpp::connected_client client) {
 		return;
 	}
 
+	// Client is talking to us, we don't need to send a heartbeat if we're being talked to.
+	client.last_heartbeat = time(nullptr);
+
 	std::string packet_data(&buffer[8], &buffer[buffer.size()-2]);
 	int id = bit32_to_int(buffer);
 	int type = type_to_int(buffer);
 
-	rconpp::packet packet_to_send{};
+	packet packet_to_send{};
 
 	if (!client.authenticated) {
+		on_log("Client not authenticated, handling authentication.");
 		if (packet_data == password) {
-			packet_to_send = form_packet("", id, rconpp::data_type::SERVERDATA_AUTH_RESPONSE);
+			packet_to_send = form_packet("", id, SERVERDATA_AUTH_RESPONSE);
 			client.authenticated = true;
 		} else {
-			packet_to_send = form_packet("", -1, rconpp::data_type::SERVERDATA_AUTH_RESPONSE);
+			packet_to_send = form_packet("", -1, SERVERDATA_AUTH_RESPONSE);
 		}
 	} else {
-		if (type != rconpp::data_type::SERVERDATA_EXECCOMMAND) {
-			packet_to_send = form_packet("Invalid packet type (" + std::to_string(type) + "). Double check your packets.", id, rconpp::data_type::SERVERDATA_RESPONSE_VALUE);
+		if (type != SERVERDATA_EXECCOMMAND) {
+			packet_to_send = form_packet("Invalid packet type (" + std::to_string(type) + "). Double check your packets.", id, SERVERDATA_RESPONSE_VALUE);
 			on_log("Invalid packet type (" + std::to_string(type) + ") sent by [" + inet_ntoa(client.sock_info.sin_addr) + ":" + std::to_string(ntohs(client.sock_info.sin_port)) + "]. Double check your packets.");
 		} else {
 			on_log("Client [" + std::string(inet_ntoa(client.sock_info.sin_addr)) + ":" + std::to_string(ntohs(client.sock_info.sin_port)) + "] has asked to execute the command: \"" + packet_data + "\"");
@@ -145,7 +154,7 @@ void rconpp::rcon_server::read_packet(rconpp::connected_client client) {
 				 * It's better to just send no information and let clients assume that meant
 				 * the server didn't like the command.
 				 */
-				packet_to_send = form_packet("", id, rconpp::data_type::SERVERDATA_RESPONSE_VALUE);
+				packet_to_send = form_packet("", id, SERVERDATA_RESPONSE_VALUE);
 			} else {
 				client_command command{};
 				command.command = packet_data;
@@ -155,7 +164,7 @@ void rconpp::rcon_server::read_packet(rconpp::connected_client client) {
 
 				on_log("Sending reply \"" + text_to_send + "\" to client [" + std::string(inet_ntoa(client.sock_info.sin_addr)) + ":" + std::to_string(ntohs(client.sock_info.sin_port)) + "].");
 
-				packet_to_send = form_packet(text_to_send, id, rconpp::data_type::SERVERDATA_RESPONSE_VALUE);
+				packet_to_send = form_packet(text_to_send, id, SERVERDATA_RESPONSE_VALUE);
 			}
 		}
 	}
@@ -169,6 +178,21 @@ void rconpp::rcon_server::read_packet(rconpp::connected_client client) {
 	}
 }
 
+bool rconpp::rcon_server::send_heartbeat(connected_client& client) {
+	on_log("Sending heartbeat to client [" + std::string(inet_ntoa(client.sock_info.sin_addr)) + ":" + std::to_string(ntohs(client.sock_info.sin_port)) + "]");
+
+	packet packet_to_send = form_packet("", -1, SERVERDATA_RESPONSE_VALUE);
+	if (send(client.socket, packet_to_send.data.data(), packet_to_send.length, 0) < 0) {
+		on_log("Failed to send a heartbeat to client [" + std::string(inet_ntoa(client.sock_info.sin_addr)) + ":" + std::to_string(ntohs(client.sock_info.sin_port)) + "]");
+		report_error();
+		return false;
+	}
+
+	client.last_heartbeat = time(nullptr);
+
+	return true;
+}
+
 void rconpp::rcon_server::start(bool return_after) {
 	auto block_calling_thread = [this]() {
 		std::mutex thread_mutex;
@@ -184,7 +208,7 @@ void rconpp::rcon_server::start(bool return_after) {
 	on_log("Attempting to startup an RCON server...");
 
 	if (!startup_server()) {
-		on_log("RCON++ is aborting as it failed to initiate server.");
+		on_log("RCON server is aborting as it failed to initiate server.");
 		return;
 	}
 
@@ -194,39 +218,37 @@ void rconpp::rcon_server::start(bool return_after) {
 
 	accept_connections_runner = std::thread([this]() {
 		while (online) {
-			connected_client client{};
-			struct sockaddr_in client_info{};
+			sockaddr_in client_info{};
 
 			socklen_t client_len = sizeof(client_info);
 			int client_socket = accept(sock, reinterpret_cast<sockaddr*>(&client_info), &client_len);
 
 			if (client_socket == -1) {
 				on_log("client with socket: \"" + std::to_string(client_socket) + "\" failed to connect.");
+				report_error();
 				continue;
 			}
 
 			on_log("Client [" + std::string(inet_ntoa(client_info.sin_addr)) + ":" + std::to_string(ntohs(client_info.sin_port)) + "] has connected to the server.");
 
+			connected_client client{};
+
 			client.sock_info = client_info;
 			client.socket = client_socket;
 			client.connected = true;
 
-			std::thread client_thread([this, client]{
+			std::thread client_thread([this, &client]{
 				while (client.connected) {
 					read_packet(client);
 
-					time_t current_time = time(nullptr);
-					if (client_socket_to_last_heartbeat.find(client.socket) == client_socket_to_last_heartbeat.end()) {
-						client_socket_to_last_heartbeat.insert({ client.socket, current_time });
-					} else {
-						time_t last_time = client_socket_to_last_heartbeat.at(client.socket);
+					const time_t current_time = time(nullptr);
 
-						if (current_time - last_time >= HEARTBEAT_TIME)
+					if (client.authenticated) {
+						if (client.last_heartbeat == 0 || current_time - client.last_heartbeat >= HEARTBEAT_TIME)
 						{
-							send_heartbeat(client);
-
-							// We should check if the heartbeat actually got anything, if it does then insert back into the map.
-							// if it failed, we bin that client off and shut this thread down.
+							if (!send_heartbeat(client)) {
+								disconnect_client(client.socket);
+							}
 						}
 					}
 
@@ -239,6 +261,8 @@ void rconpp::rcon_server::start(bool return_after) {
 
 			request_handlers.at(client_socket).detach();
 
+			std::lock_guard guard(connected_clients_mutex);
+
 			connected_clients.insert({ client_socket, client });
 		}
 	});

src/rconpp/utilities.cpp

@@ -46,7 +46,7 @@ void rconpp::report_error() {
 #endif
 }
 
-int rconpp::read_packet_size(int socket, const std::function<void(const std::string_view log)>& on_log) {
+int rconpp::read_packet_size(int socket) {
 	std::vector<char> buffer{};
 	buffer.resize(4);
 
@@ -55,7 +55,6 @@ int rconpp::read_packet_size(int socket, const std::function<void(const std::str
 	 * We simply just want to read that and then return it.
 	 */
 	if (recv(socket, buffer.data(), 4, 0) == -1) {
-		on_log("Did not receive a packet in time. Did the server send a response?");
 		report_error();
 		return -1;
 	}