4 files changed, 51 insertions, 29 deletions

diff --git a/src/OSSupport/Network.h b/src/OSSupport/Network.h
index 3ed9885ec..b9ca377cb 100644
--- a/src/OSSupport/Network.h
+++ b/src/OSSupport/Network.h
@@ -139,8 +139,15 @@ public:
 		// Force a virtual destructor for all descendants:
 		virtual ~cListenCallbacks() {}
 
-		/** Called when the TCP server created with Listen() accepts an incoming connection.
-		Provides the newly created Link that can be used for communication. */
+		/** Called when the TCP server created with Listen() receives a new incoming connection.
+		Returns the link callbacks that the server should use for the newly created link.
+		If a nullptr is returned, the connection is dropped immediately;
+		otherwise a new cTCPLink instance is created and OnAccepted() is called. */
+		virtual cTCPLink::cCallbacksPtr OnIncomingConnection(const AString & a_RemoteIPAddress, UInt16 a_RemotePort) = 0;
+
+		/** Called when the TCP server created with Listen() creates a new link for an incoming connection.
+		Provides the newly created Link that can be used for communication.
+		Called right after a successful OnIncomingConnection(). */
 		virtual void OnAccepted(cTCPLink & a_Link) = 0;
 
 		/** Called when the socket fails to listen on the specified port. */
@@ -180,7 +187,7 @@ public:
 	Implemented in TCPLinkImpl.cpp. */
 	static bool Connect(
 		const AString & a_Host,
-		const UInt16 a_Port,
+		UInt16 a_Port,
 		cConnectCallbacksPtr a_ConnectCallbacks,
 		cTCPLink::cCallbacksPtr a_LinkCallbacks
 	);
@@ -192,9 +199,8 @@ public:
 	Returns a cServerHandle that can be used to query the operation status and close the server.
 	Implemented in ServerHandleImpl.cpp. */
 	static cServerHandlePtr Listen(
-		const UInt16 a_Port,
-		cListenCallbacksPtr a_ListenCallbacks,
-		cTCPLink::cCallbacksPtr a_LinkCallbacks
+		UInt16 a_Port,
+		cListenCallbacksPtr a_ListenCallbacks
 	);
 
 
diff --git a/src/OSSupport/ServerHandleImpl.cpp b/src/OSSupport/ServerHandleImpl.cpp
index d81d4ba46..82cbecef2 100644
--- a/src/OSSupport/ServerHandleImpl.cpp
+++ b/src/OSSupport/ServerHandleImpl.cpp
@@ -31,12 +31,8 @@ static bool IsValidSocket(evutil_socket_t a_Socket)
 ////////////////////////////////////////////////////////////////////////////////
 // cServerHandleImpl:
 
-cServerHandleImpl::cServerHandleImpl(
-	cNetwork::cListenCallbacksPtr a_ListenCallbacks,
-	cTCPLink::cCallbacksPtr a_LinkCallbacks
-):
+cServerHandleImpl::cServerHandleImpl(cNetwork::cListenCallbacksPtr a_ListenCallbacks):
 	m_ListenCallbacks(a_ListenCallbacks),
-	m_LinkCallbacks(a_LinkCallbacks),
 	m_ConnListener(nullptr),
 	m_SecondaryConnListener(nullptr),
 	m_IsListening(false),
@@ -92,11 +88,10 @@ void cServerHandleImpl::Close(void)
 
 cServerHandleImplPtr cServerHandleImpl::Listen(
 	UInt16 a_Port,
-	cNetwork::cListenCallbacksPtr a_ListenCallbacks,
-	cTCPLink::cCallbacksPtr a_LinkCallbacks
+	cNetwork::cListenCallbacksPtr a_ListenCallbacks
 )
 {
-	cServerHandleImplPtr res = cServerHandleImplPtr{new cServerHandleImpl(a_ListenCallbacks, a_LinkCallbacks)};
+	cServerHandleImplPtr res = cServerHandleImplPtr{new cServerHandleImpl(a_ListenCallbacks)};
 	if (res->Listen(a_Port))
 	{
 		cNetworkSingleton::Get().AddServer(res);
@@ -253,8 +248,37 @@ void cServerHandleImpl::Callback(evconnlistener * a_Listener, evutil_socket_t a_
 	cServerHandleImpl * Self = reinterpret_cast<cServerHandleImpl *>(a_Self);
 	ASSERT(Self != nullptr);
 
+	// Get the textual IP address and port number out of a_Addr:
+	char IPAddress[128];
+	evutil_inet_ntop(a_Addr->sa_family, a_Addr->sa_data, IPAddress, ARRAYCOUNT(IPAddress));
+	UInt16 Port = 0;
+	switch (a_Addr->sa_family)
+	{
+		case AF_INET:
+		{
+			sockaddr_in * sin = reinterpret_cast<sockaddr_in *>(a_Addr);
+			Port = ntohs(sin->sin_port);
+			break;
+		}
+		case AF_INET6:
+		{
+			sockaddr_in6 * sin6 = reinterpret_cast<sockaddr_in6 *>(a_Addr);
+			Port = ntohs(sin6->sin6_port);
+			break;
+		}
+	}
+
+	// Call the OnIncomingConnection callback to get the link callbacks to use:
+	cTCPLink::cCallbacksPtr LinkCallbacks = Self->m_ListenCallbacks->OnIncomingConnection(IPAddress, Port);
+	if (LinkCallbacks == nullptr)
+	{
+		// Drop the connection:
+		evutil_closesocket(a_Socket);
+		return;
+	}
+
 	// Create a new cTCPLink for the incoming connection:
-	cTCPLinkImplPtr Link = std::make_shared<cTCPLinkImpl>(a_Socket, Self->m_LinkCallbacks, Self, a_Addr, static_cast<socklen_t>(a_Len));
+	cTCPLinkImplPtr Link = std::make_shared<cTCPLinkImpl>(a_Socket, LinkCallbacks, Self, a_Addr, static_cast<socklen_t>(a_Len));
 	{
 		cCSLock Lock(Self->m_CS);
 		Self->m_Connections.push_back(Link);
@@ -289,12 +313,11 @@ void cServerHandleImpl::RemoveLink(const cTCPLinkImpl * a_Link)
 // cNetwork API:
 
 cServerHandlePtr cNetwork::Listen(
-	const UInt16 a_Port,
-	cNetwork::cListenCallbacksPtr a_ListenCallbacks,
-	cTCPLink::cCallbacksPtr a_LinkCallbacks
+	UInt16 a_Port,
+	cNetwork::cListenCallbacksPtr a_ListenCallbacks
 )
 {
-	return cServerHandleImpl::Listen(a_Port, a_ListenCallbacks, a_LinkCallbacks);
+	return cServerHandleImpl::Listen(a_Port, a_ListenCallbacks);
 }
 
 
diff --git a/src/OSSupport/ServerHandleImpl.h b/src/OSSupport/ServerHandleImpl.h
index b325a0f37..33ff787f2 100644
--- a/src/OSSupport/ServerHandleImpl.h
+++ b/src/OSSupport/ServerHandleImpl.h
@@ -46,8 +46,7 @@ public:
 	Always returns a server instance; in the event of a failure, the instance holds the error details. Use IsListening() to query success. */
 	static cServerHandleImplPtr Listen(
 		UInt16 a_Port,
-		cNetwork::cListenCallbacksPtr a_ListenCallbacks,
-		cTCPLink::cCallbacksPtr a_LinkCallbacks
+		cNetwork::cListenCallbacksPtr a_ListenCallbacks
 	);
 
 	// cServerHandle overrides:
@@ -58,9 +57,6 @@ protected:
 	/** The callbacks used to notify about incoming connections. */
 	cNetwork::cListenCallbacksPtr m_ListenCallbacks;
 
-	/** The callbacks used to create new cTCPLink instances for incoming connections. */
-	cTCPLink::cCallbacksPtr m_LinkCallbacks;
-
 	/** The LibEvent handle representing the main listening socket. */
 	evconnlistener * m_ConnListener;
 
@@ -86,10 +82,7 @@ protected:
 
 	/** Creates a new instance with the specified callbacks.
 	Initializes the internals, but doesn't start listening yet. */
-	cServerHandleImpl(
-		cNetwork::cListenCallbacksPtr a_ListenCallbacks,
-		cTCPLink::cCallbacksPtr a_LinkCallbacks
-	);
+	cServerHandleImpl(cNetwork::cListenCallbacksPtr a_ListenCallbacks);
 
 	/** Starts listening on the specified port.
 	Returns true if successful, false on failure. On failure, sets m_ErrorCode and m_ErrorMsg. */
diff --git a/src/OSSupport/TCPLinkImpl.cpp b/src/OSSupport/TCPLinkImpl.cpp
index f87f68280..6f937646f 100644
--- a/src/OSSupport/TCPLinkImpl.cpp
+++ b/src/OSSupport/TCPLinkImpl.cpp
@@ -304,7 +304,7 @@ void cTCPLinkImpl::UpdateRemoteAddress(void)
 
 bool cNetwork::Connect(
 	const AString & a_Host,
-	const UInt16 a_Port,
+	UInt16 a_Port,
 	cNetwork::cConnectCallbacksPtr a_ConnectCallbacks,
 	cTCPLink::cCallbacksPtr a_LinkCallbacks
 )