diff --git a/check_input.cpp b/check_input.cpp
index 9083623..52faa1e 100644
--- a/check_input.cpp
+++ b/check_input.cpp
@@ -8,9 +8,10 @@
 #include "includes/raygui/raygui.h"
 #include "includes/exception_consts.hpp"
 #include "includes/raygui_helpers.hpp"
+#include "includes/display_text.hpp"
 #include "includes/timer.h"
 
-GameType check_server(char* ip_text, char* port_text) {
+GameType check_server(char* ip_text, char* port_text, const int if_mode) {
 	GameType type;
 	std::string addr;
 	uint16_t port;
@@ -33,7 +34,7 @@ GameType check_server(char* ip_text, char* port_text) {
 	/* Create server socket and wait for client to connect */
 	Server* server = new Server(ES_UDP, addr.data(), port);
 	server->create_socket();
-	display_text_centered("Your code is " + code + "\nWaiting for connection...");
+	display_text("Your code is " + code + "\nWaiting for connection...", if_mode);
 	std::string response = "";
 	char* temp_response = NULL;
 	/* Wait for the client to connect. Since recvAll returns a char*, we need to create a temporary variable to check for NULL. 
@@ -44,30 +45,29 @@ GameType check_server(char* ip_text, char* port_text) {
 	response = std::string(temp_response);
 	
 	server->sendAll("U2");
-	display_text_centered("Connection received from " + server->get_peer_addr());
-	Timer timer = timer_init(3); 
-	while (!timer_done(timer)); // Wait for five seconds
+	display_text("Connection received from " + server->get_peer_addr(), if_mode);
 	
 	type.mode = M_SERVER;
 	type.netsock = server;
 	return type;
 }
 
-GameType check_client(char* code_text) {
+GameType check_client(char* code_text, const int if_mode) {
 	GameType type;
 	std::vector<std::string> addr_port;
 	std::string connect_code = std::string(code_text); /* The connect code is a special string, that contains the server address and port. It is given by the server. */
 	try {
 		addr_port = connect_code::decode(connect_code);
+		if (check_ip_ver(addr_port[0].data()) < 0) {
+			throw std::invalid_argument("Invalid code entered.");
+		}
 		Client* client = new Client(ES_UDP, addr_port[0].data(), std::stoi(addr_port[1]));
 		client->create_socket();
 		/* Send a specific message to the server, and wait for the appropriate response, to know that the server is ready */
 		client->sendAll("GG");
-		std::string msg_from_server = client->recvAllNB();
+		std::string msg_from_server = client->recvAll();
 		if (msg_from_server == "U2") {
-			display_text_centered("Connection made");
-			Timer timer = timer_init(3);
-			while (!timer_done(timer));
+			display_text("Connection made", if_mode);
 		} else {
 			throw std::invalid_argument("Server didn't respond with correct message.");
 		}
diff --git a/includes/check_input.hpp b/includes/check_input.hpp
index ccb7e4c..730c6f0 100644
--- a/includes/check_input.hpp
+++ b/includes/check_input.hpp
@@ -17,12 +17,14 @@ typedef struct {
 
 /* This function checks the IP address and port passed to it, and returns a struct,
 that contains information about the game mode, and contains the server socket.
-It assumes that both ip_text and port_text are non-null
+It assumes that both ip_text and port_text are non-null.
+Any errors are printed using the display_text function, with the given if_type.
 TODO - Add better error checking. */
-GameType check_server(char* ip_text, char* port_text);
+GameType check_server(char* ip_text, char* port_text, const int if_type);
 
 /* NOT IMPLEMENTED YET - This function checks the code given to it, and returns a struct
-that contains information about the game mode, and contains the client socket. */
-GameType check_client(char* code);
+that contains information about the game mode, and contains the client socket.
+Any errors are printed using the display_text function, with the given if_type. */
+GameType check_client(char* code, const int if_type);
 
 #endif