From 9089c2b35ad19d7b80b23e6583adb2662ca57b83 Mon Sep 17 00:00:00 2001
From: Dennis Felsing <dennis@felsin9.de>
Date: Thu, 13 Dec 2018 16:24:37 +0100
Subject: [PATCH] Don't decompress packets from unknown IPs

---
 src/engine/shared/network.cpp        |  6 ++++-
 src/engine/shared/network.h          |  2 +-
 src/engine/shared/network_client.cpp |  2 +-
 src/engine/shared/network_server.cpp | 39 ++++++++++++++++------------
 4 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/src/engine/shared/network.cpp b/src/engine/shared/network.cpp
index b563f1da4..b5789c609 100644
--- a/src/engine/shared/network.cpp
+++ b/src/engine/shared/network.cpp
@@ -175,7 +175,7 @@ void CNetBase::SendPacket(NETSOCKET Socket, NETADDR *pAddr, CNetPacketConstruct
 }
 
 // TODO: rename this function
-int CNetBase::UnpackPacket(unsigned char *pBuffer, int Size, CNetPacketConstruct *pPacket)
+int CNetBase::UnpackPacket(unsigned char *pBuffer, int Size, CNetPacketConstruct *pPacket, bool Decompress)
 {
 	// check the size
 	if(Size < NET_PACKETHEADERSIZE || Size > NET_MAX_PACKETSIZE)
@@ -225,6 +225,10 @@ int CNetBase::UnpackPacket(unsigned char *pBuffer, int Size, CNetPacketConstruct
 	{
 		if(pPacket->m_Flags&NET_PACKETFLAG_COMPRESSION)
 		{
+			if(!Decompress)
+			{
+				return -1;
+			}
 			// Don't allow compressed control packets.
 			if(pPacket->m_Flags&NET_PACKETFLAG_CONTROL)
 			{
diff --git a/src/engine/shared/network.h b/src/engine/shared/network.h
index 824a4c0d3..9b6de2b11 100644
--- a/src/engine/shared/network.h
+++ b/src/engine/shared/network.h
@@ -469,7 +469,7 @@ public:
 	static void SendPacketConnless(NETSOCKET Socket, NETADDR *pAddr, const void *pData, int DataSize, bool Extended, unsigned char aExtra[4]);
 	static void SendPacket(NETSOCKET Socket, NETADDR *pAddr, CNetPacketConstruct *pPacket, SECURITY_TOKEN SecurityToken);
 
-	static int UnpackPacket(unsigned char *pBuffer, int Size, CNetPacketConstruct *pPacket);
+	static int UnpackPacket(unsigned char *pBuffer, int Size, CNetPacketConstruct *pPacket, bool Decompress);
 
 	// The backroom is ack-NET_MAX_SEQUENCE/2. Used for knowing if we acked a packet or not
 	static int IsSeqInBackroom(int Seq, int Ack);
diff --git a/src/engine/shared/network_client.cpp b/src/engine/shared/network_client.cpp
index f955ef5a8..ff2e70749 100644
--- a/src/engine/shared/network_client.cpp
+++ b/src/engine/shared/network_client.cpp
@@ -70,7 +70,7 @@ int CNetClient::Recv(CNetChunk *pChunk)
 		if(Bytes <= 0)
 			break;
 
-		if(CNetBase::UnpackPacket(m_RecvUnpacker.m_aBuffer, Bytes, &m_RecvUnpacker.m_Data) == 0)
+		if(CNetBase::UnpackPacket(m_RecvUnpacker.m_aBuffer, Bytes, &m_RecvUnpacker.m_Data, /* decompress = */ true) == 0)
 		{
 			if(m_RecvUnpacker.m_Data.m_Flags&NET_PACKETFLAG_CONNLESS)
 			{
diff --git a/src/engine/shared/network_server.cpp b/src/engine/shared/network_server.cpp
index 705a19ad1..345b786ed 100644
--- a/src/engine/shared/network_server.cpp
+++ b/src/engine/shared/network_server.cpp
@@ -615,9 +615,9 @@ int CNetServer::Recv(CNetChunk *pChunk)
 			continue;
 		}
 
-		if(CNetBase::UnpackPacket(m_RecvUnpacker.m_aBuffer, Bytes, &m_RecvUnpacker.m_Data) == 0)
+		if(m_RecvUnpacker.m_Data.m_Flags&NET_PACKETFLAG_CONNLESS)
 		{
-			if(m_RecvUnpacker.m_Data.m_Flags&NET_PACKETFLAG_CONNLESS)
+			if(CNetBase::UnpackPacket(m_RecvUnpacker.m_aBuffer, Bytes, &m_RecvUnpacker.m_Data, /* decompress = */ false) == 0)
 			{
 				pChunk->m_Flags = NETSENDFLAG_CONNLESS;
 				pChunk->m_ClientID = -1;
@@ -631,21 +631,24 @@ int CNetServer::Recv(CNetChunk *pChunk)
 				}
 				return 1;
 			}
-			else
+		}
+		else
+		{
+			// drop invalid ctrl packets
+			if (m_RecvUnpacker.m_Data.m_Flags&NET_PACKETFLAG_CONTROL &&
+					m_RecvUnpacker.m_Data.m_DataSize == 0)
+				continue;
+
+			// normal packet, find matching slot
+			int Slot = GetClientSlot(Addr);
+
+			if (Slot != -1)
 			{
-				// drop invalid ctrl packets
-				if (m_RecvUnpacker.m_Data.m_Flags&NET_PACKETFLAG_CONTROL &&
-						m_RecvUnpacker.m_Data.m_DataSize == 0)
-					continue;
+				// found
 
-				// normal packet, find matching slot
-				int Slot = GetClientSlot(Addr);
-
-				if (Slot != -1)
+				// control
+				if(CNetBase::UnpackPacket(m_RecvUnpacker.m_aBuffer, Bytes, &m_RecvUnpacker.m_Data, /* decompress = */ true) == 0)
 				{
-					// found
-
-					// control
 					if(m_RecvUnpacker.m_Data.m_Flags&NET_PACKETFLAG_CONTROL)
 						OnConnCtrlMsg(Addr, Slot, m_RecvUnpacker.m_Data.m_aChunkData[0], m_RecvUnpacker.m_Data);
 
@@ -655,10 +658,12 @@ int CNetServer::Recv(CNetChunk *pChunk)
 							m_RecvUnpacker.Start(&Addr, &m_aSlots[Slot].m_Connection, Slot);
 					}
 				}
-				else
+			}
+			else
+			{
+				// not found, client that wants to connect
+				if(CNetBase::UnpackPacket(m_RecvUnpacker.m_aBuffer, Bytes, &m_RecvUnpacker.m_Data, /* decompress = */ false) == 0)
 				{
-					// not found, client that wants to connect
-
 					if(IsDDNetControlMsg(&m_RecvUnpacker.m_Data))
 						// got ddnet control msg
 						OnTokenCtrlMsg(Addr, m_RecvUnpacker.m_Data.m_aChunkData[0], m_RecvUnpacker.m_Data);