check for duplicate and missing fragments

vmlinuxer · Apr 13, 2014 · a510e7c · a510e7c
1 parent e1027ff
commit a510e7c
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 7 deletions.
diff --git a/SSU.cpp b/SSU.cpp
@@ -27,6 +27,12 @@ namespace ssu
 	SSUSession::~SSUSession ()
 	{
 		delete m_DHKeysPair;
+		for (auto it: m_IncomleteMessages)
+			if (it.second)
+			{
+				DeleteI2NPMessage (it.second->msg);
+				delete it.second;
+			}			
 	}	
 
 	void SSUSession::CreateAESandMacKey (uint8_t * pubKey, uint8_t * aesKey, uint8_t * macKey)
@@ -673,9 +679,29 @@ namespace ssu
 				auto it = m_IncomleteMessages.find (msgID);
 				if (it != m_IncomleteMessages.end ())
 				{
-					msg = it->second;
-					memcpy (msg->buf + msg->len, buf, fragmentSize);
-					msg->len += fragmentSize;
+					if (fragmentNum == it->second->nextFragmentNum)
+					{
+						// expected fragment
+						msg = it->second->msg;
+						memcpy (msg->buf + msg->len, buf, fragmentSize);
+						msg->len += fragmentSize;
+						it->second->nextFragmentNum++;
+					}	
+					else if (fragmentNum < it->second->nextFragmentNum)
+						// duplicate fragment
+						LogPrint ("Duplicate fragment ", fragmentNum, " of message ", msgID, ". Ignored");	
+					else
+					{
+						// missing fragment
+						LogPrint ("Missing fragments from ", it->second->nextFragmentNum, " to ", fragmentNum - 1, " of message ", msgID);	
+						//TODO
+					}	
+
+					if (isLast)
+					{
+						delete it->second;
+						m_IncomleteMessages.erase (it);
+					}	
 				}
 				else
 					// TODO:
@@ -691,12 +717,10 @@ namespace ssu
 			if (msg)
 			{					
 				if (!fragmentNum && !isLast)
-					m_IncomleteMessages[msgID] = msg;
+					m_IncomleteMessages[msgID] = new IncompleteMessage (msg);
 				if (isLast)
 				{
 					SendMsgAck (msgID);
-					if (fragmentNum > 0)	
-						m_IncomleteMessages.erase (msgID);
 					msg->FromSSU (msgID);
 					if (m_State == eSessionStateEstablished)
 						i2p::HandleI2NPMessage (msg);

diff --git a/SSU.h b/SSU.h
@@ -124,6 +124,14 @@ namespace ssu
 			void HandleTerminationTimer (const boost::system::error_code& ecode);
 
 		private:
+
+			struct IncompleteMessage
+			{
+				I2NPMessage * msg;
+				uint8_t nextFragmentNum;	
+
+				IncompleteMessage (I2NPMessage * m): msg (m), nextFragmentNum (1) {};
+			};
 
 			SSUServer& m_Server;
 			boost::asio::ip::udp::endpoint m_RemoteEndpoint;
@@ -137,7 +145,7 @@ namespace ssu
 			CryptoPP::CBC_Mode<CryptoPP::AES>::Encryption m_Encryption;	
 			CryptoPP::CBC_Mode<CryptoPP::AES>::Decryption m_Decryption;	
 			uint8_t m_SessionKey[32], m_MacKey[32];
-			std::map<uint32_t, I2NPMessage *> m_IncomleteMessages;
+			std::map<uint32_t, IncompleteMessage *> m_IncomleteMessages;
 			std::list<i2p::I2NPMessage *> m_DelayedMessages;
 	};