From 00196e975a3cd4156b1d675d192c875535d7e5f0 Mon Sep 17 00:00:00 2001
From: madmaxoft <github@xoft.cz>
Date: Sun, 28 Jul 2013 18:15:19 +0200
Subject: [PATCH] ProtoProxy now properly waits for both sides to establish
 encryption

No more "End of stream" kicks in the client. Data sent while one connection is encrypted and the other is not is buffered and sent when the other link establishes encryption.
---
 Tools/ProtoProxy/Connection.cpp | 59 ++++++++++++++++++++++++++-------
 Tools/ProtoProxy/Connection.h   | 10 ++++--
 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;