From 4e09b39735086b166764925b79dc478408b1133a Mon Sep 17 00:00:00 2001
From: orignal <romakoshelkin@yandex.ru>
Date: Wed, 11 Jun 2014 10:56:20 -0400
Subject: [PATCH] check sequence of fragments

---
 TunnelEndpoint.cpp | 35 ++++++++++++++++++++++++-----------
 TunnelEndpoint.h   |  7 ++++++-
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/TunnelEndpoint.cpp b/TunnelEndpoint.cpp
index fc92ec29..cfbc39ec 100644
--- a/TunnelEndpoint.cpp
+++ b/TunnelEndpoint.cpp
@@ -26,7 +26,8 @@ namespace tunnel
 				
 				bool isFollowOnFragment = flag & 0x80, isLastFragment = true;		
 				uint32_t msgID = 0;
-				TunnelMessageBlock m;
+				int fragmentNum = 0;
+				TunnelMessageBlockEx m;
 				if (!isFollowOnFragment)
 				{	
 					// first fragment
@@ -68,7 +69,7 @@ namespace tunnel
 					// follow on
 					msgID = be32toh (*(uint32_t *)fragment); // MessageID			
 					fragment += 4; 
-					int fragmentNum = (flag >> 1) & 0x3F; // 6 bits
+					fragmentNum = (flag >> 1) & 0x3F; // 6 bits
 					isLastFragment = flag & 0x01;
 					LogPrint ("Follow on fragment ", fragmentNum, " of message ", msgID, isLastFragment ? " last" : " non-last");
 				}	
@@ -101,22 +102,34 @@ namespace tunnel
 					if (msgID) // msgID is presented, assume message is fragmented
 					{
 						if (!isFollowOnFragment) // create new incomlete message
+						{
+							m.nextFragmentNum = 1;
 							m_IncompleteMessages[msgID] = m;
+						}
 						else
 						{
 							auto it = m_IncompleteMessages.find (msgID);
 							if (it != m_IncompleteMessages.end())
 							{
-								I2NPMessage * incompleteMessage = it->second.data; 
-								memcpy (incompleteMessage->buf + incompleteMessage->len, fragment, size); // concatenate fragment
-								incompleteMessage->len += size;
-								// TODO: check fragmentNum sequence
-								if (isLastFragment)
+								if (fragmentNum == it->second.nextFragmentNum)
 								{
-									// message complete
-									HandleNextMessage (it->second);	
-									m_IncompleteMessages.erase (it); 
-								}	
+									I2NPMessage * incompleteMessage = it->second.data; 
+									memcpy (incompleteMessage->buf + incompleteMessage->len, fragment, size); // concatenate fragment
+									incompleteMessage->len += size;
+									if (isLastFragment)
+									{
+										// message complete
+										HandleNextMessage (it->second);	
+										m_IncompleteMessages.erase (it); 
+									}	
+									else
+										it->second.nextFragmentNum++;
+								}
+								else
+								{	
+									LogPrint ("Unexpected fragment ", fragmentNum, " instead ", it->second.nextFragmentNum, " of message ", msgID, ". Discarded");
+									m_IncompleteMessages.erase (it); // TODO: store unexpect fragment for a while
+								}
 							}
 							else
 								LogPrint ("First fragment of message ", msgID, " not found. Discarded");
diff --git a/TunnelEndpoint.h b/TunnelEndpoint.h
index b18f204b..2323e574 100644
--- a/TunnelEndpoint.h
+++ b/TunnelEndpoint.h
@@ -26,7 +26,12 @@ namespace tunnel
 			
 		private:
 		
-			std::map<uint32_t, TunnelMessageBlock> m_IncompleteMessages;
+			struct TunnelMessageBlockEx: public TunnelMessageBlock
+			{
+				uint8_t nextFragmentNum;
+			};				
+
+			std::map<uint32_t, TunnelMessageBlockEx> m_IncompleteMessages;
 			size_t m_NumReceivedBytes;
 	};	
 }