From e0a44fb3bc250197e30c5023fa53fc5c27a5cdcf Mon Sep 17 00:00:00 2001
From: peterbell10 <peterbell10@live.co.uk>
Date: Mon, 17 Jul 2017 00:19:09 +0100
Subject: Handle incomplete packets in cProtocolRecognizer

---
 src/Protocol/ProtocolRecognizer.cpp | 97 +++++++++++++++++++++----------------
 src/Protocol/ProtocolRecognizer.h   |  5 +-
 2 files changed, 58 insertions(+), 44 deletions(-)

(limited to 'src/Protocol')

diff --git a/src/Protocol/ProtocolRecognizer.cpp b/src/Protocol/ProtocolRecognizer.cpp
index 812d81721..187a73a6e 100644
--- a/src/Protocol/ProtocolRecognizer.cpp
+++ b/src/Protocol/ProtocolRecognizer.cpp
@@ -83,46 +83,58 @@ void cProtocolRecognizer::DataReceived(const char * a_Data, size_t a_Size)
 		return;
 	}
 
-	if (m_InPingForUnrecognizedVersion)
+	if (!m_InPingForUnrecognizedVersion)
 	{
-		// We already know the verison; handle it here.
-		UInt32 PacketLen;
-		UInt32 PacketID;
-		if (!m_Buffer.ReadVarInt32(PacketLen))
+		if (TryRecognizeProtocol())
 		{
+			// The protocol has just been recognized, dump the whole m_Buffer contents into it for parsing:
+			AString Dump;
+			m_Buffer.ResetRead();
+			m_Buffer.ReadAll(Dump);
+			m_Protocol->DataReceived(Dump.data(), Dump.size());
 			return;
 		}
-		if (!m_Buffer.ReadVarInt32(PacketID))
+		else
 		{
-			return;
+			m_Buffer.ResetRead();
 		}
-		if ((PacketID != 0x01) || (PacketLen != 9))
+	}
+
+	if (!m_InPingForUnrecognizedVersion)
+	{
+		return;
+	}
+
+	// Handle server list ping packets
+	for (;;)
+	{
+		UInt32 PacketLen;
+		UInt32 PacketID;
+		if (
+			!m_Buffer.ReadVarInt32(PacketLen) ||
+			!m_Buffer.CanReadBytes(PacketLen) ||
+			!m_Buffer.ReadVarInt32(PacketID)
+		)
 		{
-			// Not a Ping packet
-			return;
+			// Not enough data
+			m_Buffer.ResetRead();
+			break;
 		}
 
-		Int64 Data;
-		if (!m_Buffer.ReadBEInt64(Data))
+		if ((PacketID == 0x00) && (PacketLen == 1))  // Request packet
 		{
+			HandlePacketStatusRequest();
+		}
+		else if ((PacketID == 0x01) && (PacketLen == 9))  // Ping packet
+		{
+			HandlePacketStatusPing();
+		}
+		else
+		{
+			m_Client->Kick("Server list ping failed, unrecognized packet");
 			return;
 		}
-
-		cPacketizer Pkt(*this, 0x01);  // Pong packet
-		Pkt.WriteBEInt64(Data);
-		return;
-	}
-
-	if (!TryRecognizeProtocol())
-	{
-		return;
 	}
-
-	// The protocol has just been recognized, dump the whole m_Buffer contents into it for parsing:
-	AString Dump;
-	m_Buffer.ResetRead();
-	m_Buffer.ReadAll(Dump);
-	m_Protocol->DataReceived(Dump.data(), Dump.size());
 }
 
 
@@ -1108,20 +1120,6 @@ bool cProtocolRecognizer::TryRecognizeLengthedProtocol(UInt32 a_PacketLengthRema
 			else
 			{
 				m_InPingForUnrecognizedVersion = true;
-
-				UInt32 PacketLen;
-				UInt32 PacketID;
-				if (!m_Buffer.ReadVarInt32(PacketLen))
-				{
-					return false;
-				}
-				if (!m_Buffer.ReadVarInt32(PacketID))
-				{
-					return false;
-				}
-				ASSERT(PacketID == 0x00);  // Request packet
-				ASSERT(PacketLen == 1);  // No payload except for packet ID
-				SendPingStatusResponse();
 			}
 			return false;
 		}
@@ -1155,7 +1153,7 @@ void cProtocolRecognizer::SendPacket(cPacketizer & a_Pkt)
 
 
 
-void cProtocolRecognizer::SendPingStatusResponse(void)
+void cProtocolRecognizer::HandlePacketStatusRequest(void)
 {
 	cServer * Server = cRoot::Get()->GetServer();
 	AString ServerDescription = Server->GetDescription();
@@ -1200,3 +1198,18 @@ void cProtocolRecognizer::SendPingStatusResponse(void)
 
 
 
+void cProtocolRecognizer::HandlePacketStatusPing()
+{
+	Int64 Timestamp;
+	if (!m_Buffer.ReadBEInt64(Timestamp))
+	{
+		return;
+	}
+
+	cPacketizer Pkt(*this, 0x01);  // Pong packet
+	Pkt.WriteBEInt64(Timestamp);
+}
+
+
+
+
diff --git a/src/Protocol/ProtocolRecognizer.h b/src/Protocol/ProtocolRecognizer.h
index 1a638fb21..32dcca940 100644
--- a/src/Protocol/ProtocolRecognizer.h
+++ b/src/Protocol/ProtocolRecognizer.h
@@ -143,8 +143,6 @@ public:
 
 	virtual void SendData(const char * a_Data, size_t a_Size) override;
 
-	void SendPingStatusResponse(void);
-
 protected:
 	/** The recognized protocol */
 	cProtocol * m_Protocol;
@@ -155,6 +153,9 @@ protected:
 	/** Is a server list ping for an unrecognized version currently occuring? */
 	bool m_InPingForUnrecognizedVersion;
 
+	// Packet handlers while in status state (m_InPingForUnrecognizedVersion == true)
+	void HandlePacketStatusRequest();
+	void HandlePacketStatusPing();
 
 	/** Tries to recognize protocol based on m_Buffer contents; returns true if recognized */
 	bool TryRecognizeProtocol(void);
-- 
cgit v1.2.3