2 files changed, 54 insertions, 15 deletions

diff --git a/Tools/ProtoProxy/Connection.cpp b/Tools/ProtoProxy/Connection.cpp
index a3d3191fc..32dfe321e 100644
--- a/Tools/ProtoProxy/Connection.cpp
+++ b/Tools/ProtoProxy/Connection.cpp
@@ -49,13 +49,25 @@
 	{ \

 		AString ToServer; \

 		m_ClientBuffer.ReadAgain(ToServer); \

-		if (m_ServerState == csUnencrypted) \

+		switch (m_ServerState) \

 		{ \

-			SERVERSEND(ToServer.data(), ToServer.size()); \

-		} \

-		else \

-		{ \

-			SERVERENCRYPTSEND(ToServer.data(), ToServer.size()); \

+			case csUnencrypted: \

+			{ \

+				SERVERSEND(ToServer.data(), ToServer.size()); \

+				break; \

+			} \

+			case csEncryptedUnderstood: \

+			case csEncryptedUnknown: \

+			{ \

+				SERVERENCRYPTSEND(ToServer.data(), ToServer.size()); \

+				break; \

+			} \

+			case csWaitingForEncryption: \

+			{ \

+				Log("Waiting for server encryption, queued %u bytes", ToServer.size()); \

+				m_ServerEncryptionBuffer.append(ToServer.data(), ToServer.size()); \

+				break; \

+			} \

 		} \

 		DebugSleep(50); \

 	}

@@ -64,13 +76,25 @@
 	{ \

 		AString ToClient; \

 		m_ServerBuffer.ReadAgain(ToClient); \

-		if (m_ClientState == csUnencrypted) \

+		switch (m_ClientState) \

 		{ \

-			CLIENTSEND(ToClient.data(), ToClient.size()); \

-		} \

-		else \

-		{ \

-			CLIENTENCRYPTSEND(ToClient.data(), ToClient.size()); \

+			case csUnencrypted: \

+			{ \

+				CLIENTSEND(ToClient.data(), ToClient.size()); \

+				break; \

+			} \

+			case csEncryptedUnderstood: \

+			case csEncryptedUnknown: \

+			{ \

+				CLIENTENCRYPTSEND(ToClient.data(), ToClient.size()); \

+				break; \

+			} \

+			case csWaitingForEncryption: \

+			{ \

+				Log("Waiting for client encryption, queued %u bytes", ToClient.size()); \

+				m_ClientEncryptionBuffer.append(ToClient.data(), ToClient.size()); \

+				break; \

+			} \

 		} \

 		DebugSleep(50); \

 	}

@@ -379,6 +403,7 @@ bool cConnection::RelayFromServer(void)
 	switch (m_ServerState)

 	{

 		case csUnencrypted:

+		case csWaitingForEncryption:

 		{

 			return DecodeServersPackets(Buffer, res);

 		}

@@ -419,6 +444,7 @@ bool cConnection::RelayFromClient(void)
 	switch (m_ClientState)

 	{

 		case csUnencrypted:

+		case csWaitingForEncryption:

 		{

 			return DecodeClientsPackets(Buffer, res);

 		}

@@ -1543,6 +1569,9 @@ bool cConnection::HandleServerEncryptionKeyResponse(void)
 	}

 	Log("Server communication is now encrypted");

 	m_ServerState = csEncryptedUnderstood;

+	DataLog(m_ServerEncryptionBuffer.data(), m_ServerEncryptionBuffer.size(), "Sending the queued data to server (%u bytes):", m_ServerEncryptionBuffer.size());

+	SERVERENCRYPTSEND(m_ServerEncryptionBuffer.data(), m_ServerEncryptionBuffer.size());

+	m_ServerEncryptionBuffer.clear();

 	return true;

 }

 

@@ -2459,6 +2488,7 @@ void cConnection::SendEncryptionKeyResponse(const AString & a_ServerPublicKey, c
 	ToServer.WriteBEShort(EncryptedLength);

 	ToServer.WriteBuf(EncryptedNonce, EncryptedLength);

 	SERVERSEND(ToServer);

+	m_ServerState = csWaitingForEncryption;

 }

 

 

@@ -2507,6 +2537,11 @@ void cConnection::StartClientEncryption(const AString & a_EncKey, const AString
 	Log("Client connection is now encrypted");

 	m_ClientState = csEncryptedUnderstood;

 	

+	// Send the queued data:

+	DataLog(m_ClientEncryptionBuffer.data(), m_ClientEncryptionBuffer.size(), "Sending the queued data to client (%u bytes):", m_ClientEncryptionBuffer.size());

+	CLIENTENCRYPTSEND(m_ClientEncryptionBuffer.data(), m_ClientEncryptionBuffer.size());

+	m_ClientEncryptionBuffer.clear();

+	

 	// Handle all postponed server data

 	DecodeServersPackets(NULL, 0);

 }

diff --git a/Tools/ProtoProxy/Connection.h b/Tools/ProtoProxy/Connection.h
index 7f3a6f8bb..dafc1b36b 100644
--- a/Tools/ProtoProxy/Connection.h
+++ b/Tools/ProtoProxy/Connection.h
@@ -39,9 +39,10 @@ class cConnection
 	

 	enum eConnectionState

 	{

-		csUnencrypted,          // The connection is not encrypted. Packets must be decoded in order to be able to start decryption.

-		csEncryptedUnderstood,  // The communication is encrypted and so far all packets have been understood, so they can be still decoded

-		csEncryptedUnknown,     // The communication is encrypted, but an unknown packet has been received, so packets cannot be decoded anymore

+		csUnencrypted,           // The connection is not encrypted. Packets must be decoded in order to be able to start decryption.

+		csEncryptedUnderstood,   // The communication is encrypted and so far all packets have been understood, so they can be still decoded

+		csEncryptedUnknown,      // The communication is encrypted, but an unknown packet has been received, so packets cannot be decoded anymore

+		csWaitingForEncryption,  // The communication is waiting for the other line to establish encryption

 	};

 	

 	eConnectionState m_ClientState;

@@ -72,6 +73,9 @@ protected:
 	Decryptor m_ClientDecryptor;

 	Encryptor m_ClientEncryptor;

 	

+	AString m_ClientEncryptionBuffer;  // Buffer for the data to be sent to the client once encryption is established

+	AString m_ServerEncryptionBuffer;  // Buffer for the data to be sent to the server once encryption is established

+	

 	/// Set to true when PACKET_PING is received from the client; will cause special parsing for server kick

 	bool m_HasClientPinged;