From 1da9e926d159e142902237a26db9ca33df0c160b Mon Sep 17 00:00:00 2001
From: Kilemonn <kylegonzalez2255@gmail.com>
Date: Sun, 14 Jul 2024 22:43:28 +0900
Subject: [PATCH] Add TCPSocket constructor from a SocketAddress. Add tests.

---
 src/socket/TCPSocket.cpp       | 28 +++++++++++++++++++++++++++-
 src/socket/TCPSocket.h         |  1 +
 tests/socket/TCPSocketTest.cpp | 30 ++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/socket/TCPSocket.cpp b/src/socket/TCPSocket.cpp
index 78be46f..20bff79 100644
--- a/src/socket/TCPSocket.cpp
+++ b/src/socket/TCPSocket.cpp
@@ -24,7 +24,33 @@ namespace kt
 		this->serverAddress = acceptedAddress;
 	}
 
-	TCPSocket::TCPSocket(const kt::TCPSocket& socket)
+    TCPSocket::TCPSocket(const kt::SocketAddress address)
+    {
+		std::optional<std::string> resolvedHostname = kt::getAddress(address);
+		this->hostname = resolvedHostname.value_or("");
+		this->port = kt::getPortNumber(address);
+		this->protocolVersion = kt::getInternetProtocolVersion(address);
+
+		addrinfo hints = kt::createTcpHints();
+		this->socketDescriptor = socket(address.address.sa_family, hints.ai_socktype, hints.ai_protocol);
+		if (isInvalidSocket(this->socketDescriptor))
+		{
+			throw kt::SocketException("Unable to construct socket to provided addresses with hostname [" + this->hostname + ":" + std::to_string(this->port) + "] " + getErrorCode());
+		}
+
+		int connectionResult = connect(this->socketDescriptor, &address.address, sizeof(address));
+		if (connectionResult == 0)
+		{
+			this->serverAddress = address;
+			this->protocolVersion = static_cast<kt::InternetProtocolVersion>(address.address.sa_family);
+		}
+		else
+		{
+			throw kt::SocketException("Unable to connect to provided address with hostname [" + this->hostname + ":" + std::to_string(this->port) + "] " + getErrorCode());
+		}
+    }
+
+    TCPSocket::TCPSocket(const kt::TCPSocket& socket)
 	{
 		this->socketDescriptor = socket.socketDescriptor;
 		this->hostname = socket.hostname;
diff --git a/src/socket/TCPSocket.h b/src/socket/TCPSocket.h
index e7688f0..f162ca6 100644
--- a/src/socket/TCPSocket.h
+++ b/src/socket/TCPSocket.h
@@ -56,6 +56,7 @@ namespace kt
 			TCPSocket() = delete;
 			TCPSocket(const std::string&, const unsigned short&, const kt::InternetProtocolVersion = kt::InternetProtocolVersion::Any);
 			TCPSocket(const SOCKET&, const std::string&, const unsigned short&, const kt::InternetProtocolVersion, const kt::SocketAddress&);
+			TCPSocket(const kt::SocketAddress);
 
 			TCPSocket(const kt::TCPSocket&);
 			kt::TCPSocket& operator=(const kt::TCPSocket&);
diff --git a/tests/socket/TCPSocketTest.cpp b/tests/socket/TCPSocketTest.cpp
index f910167..696d463 100644
--- a/tests/socket/TCPSocketTest.cpp
+++ b/tests/socket/TCPSocketTest.cpp
@@ -72,6 +72,36 @@ namespace kt
             TCPSocket failedSocket(LOCALHOST, serverSocket.getPort() + 1);
         }, SocketException);
     }
+
+    // Ensure we can construct and connect to a server from the address it is listening on
+    TEST_F(TCPSocketTest, TCPConstructor_FromAddress)
+    {
+        // Accept current incoming connection
+        TCPSocket server = serverSocket.acceptTCPConnection();
+
+        TCPSocket fromAddress(serverSocket.getSocketAddress());
+        TCPSocket acceptedFromAddress = serverSocket.acceptTCPConnection();
+
+        std::string sentFromAddress = "sentFromAddress";
+        ASSERT_TRUE(fromAddress.send(sentFromAddress).first);
+        ASSERT_TRUE(acceptedFromAddress.ready());
+        ASSERT_EQ(sentFromAddress, acceptedFromAddress.receiveAmount(sentFromAddress.size()));
+
+        std::string sentFromServer = "sentFromServer";
+        ASSERT_TRUE(acceptedFromAddress.send(sentFromServer).first);
+        ASSERT_TRUE(fromAddress.ready());
+        ASSERT_EQ(sentFromServer, fromAddress.receiveAmount(sentFromServer.size()));
+    }
+
+    // Ensure we throw a SocketException if we cannot construct a TCP socket from the provided SocketAddress
+    TEST_F(TCPSocketTest, TCPConstructor_FromEmptyAddress)
+    {
+        // Accept the incoming connection to make sure the server is ready
+        TCPSocket server = serverSocket.acceptTCPConnection();
+
+        kt::SocketAddress address{};
+        ASSERT_THROW(TCPSocket fromEmptyAddress(address), SocketException);
+    }
     
     /*
      * Ensure that a Socket created from the copy constructor is still able to send and receive from the copied socket.