From 993dc72ce664612496e4e2c441b37d4bd704db30 Mon Sep 17 00:00:00 2001
From: orignal <i2porignal@yandex.ru>
Date: Sun, 15 Jan 2023 22:50:54 -0500
Subject: [PATCH] use separate pointer to second fragment of incomplete message

---
 libi2pd/SSU2Session.cpp | 71 +++++++++++++++++++++++++++++------------
 libi2pd/SSU2Session.h   |  6 ++--
 2 files changed, 55 insertions(+), 22 deletions(-)

diff --git a/libi2pd/SSU2Session.cpp b/libi2pd/SSU2Session.cpp
index 79e736fe..b3c539bc 100644
--- a/libi2pd/SSU2Session.cpp
+++ b/libi2pd/SSU2Session.cpp
@@ -32,7 +32,55 @@ namespace transport
 		nextFragmentNum++;
 	}
 
+	bool SSU2IncompleteMessage::ConcatOutOfSequenceFragments ()
+	{
+		bool isLast = false;
+		if (nextFragmentNum == 1)
+		{
+			if (secondFragment)
+			{
+				AttachNextFragment (secondFragment->buf, secondFragment->len);
+				isLast = secondFragment->isLast;
+				secondFragment = nullptr;
+			}	
+			else
+				return false;
+			if (isLast) return true;
+		}
+		// might be more
+		if (outOfSequenceFragments)
+		{	
+			for (auto it = outOfSequenceFragments->begin (); it != outOfSequenceFragments->end ();)
+				if (it->first == nextFragmentNum)
+				{
+					AttachNextFragment (it->second->buf, it->second->len);
+					isLast = it->second->isLast;
+					it = outOfSequenceFragments->erase (it);
+				}
+				else
+					break;
+		}	
+		return isLast;
+	}	
 
+	void SSU2IncompleteMessage::AddOutOfSequenceFragment (int fragmentNum, 
+		std::shared_ptr<SSU2IncompleteMessage::Fragment> fragment)
+	{	
+		if (!fragmentNum || !fragment) return; // fragment 0 nun allowed
+		if (fragmentNum < nextFragmentNum) return; // already processed
+		if (fragmentNum == 1)
+		{	
+			if (!secondFragment) secondFragment = fragment;
+		}	
+		else
+		{
+			if (!outOfSequenceFragments) 
+				outOfSequenceFragments.reset (new std::map<int, std::shared_ptr<Fragment> >());
+			outOfSequenceFragments->emplace (fragmentNum, fragment);
+		}	
+		lastFragmentInsertTime = i2p::util::GetSecondsSinceEpoch ();
+	}
+		
 	SSU2Session::SSU2Session (SSU2Server& server, std::shared_ptr<const i2p::data::RouterInfo> in_RemoteRouter,
 		std::shared_ptr<const i2p::data::RouterInfo::Address> addr):
 		TransportSession (in_RemoteRouter, SSU2_CONNECT_TIMEOUT),
@@ -1733,7 +1781,7 @@ namespace transport
 		m->msg = msg;
 		m->nextFragmentNum = 1;
 		m->lastFragmentInsertTime = i2p::util::GetSecondsSinceEpoch ();
-		if (found && ConcatOutOfSequenceFragments (m))
+		if (found && m->ConcatOutOfSequenceFragments ())
 		{
 			// we have all follow-on fragments already
 			m->msg->FromNTCP2 ();
@@ -1764,7 +1812,7 @@ namespace transport
 				}
 				else
 				{
-					if (ConcatOutOfSequenceFragments (it->second))
+					if (it->second->ConcatOutOfSequenceFragments ())
 					{
 						HandleI2NPMsg (std::move (it->second->msg));
 						m_IncompleteMessages.erase (it);
@@ -1792,24 +1840,7 @@ namespace transport
 		memcpy (fragment->buf, buf + 5, len -5);
 		fragment->len = len - 5;
 		fragment->isLast = isLast;
-		it->second->outOfSequenceFragments.emplace (fragmentNum, fragment);
-		it->second->lastFragmentInsertTime = i2p::util::GetSecondsSinceEpoch ();
-	}
-
-	bool SSU2Session::ConcatOutOfSequenceFragments (std::shared_ptr<SSU2IncompleteMessage> m)
-	{
-		if (!m) return false;
-		bool isLast = false;
-		for (auto it = m->outOfSequenceFragments.begin (); it != m->outOfSequenceFragments.end ();)
-			if (it->first == m->nextFragmentNum)
-			{
-				m->AttachNextFragment (it->second->buf, it->second->len);
-				isLast = it->second->isLast;
-				it = m->outOfSequenceFragments.erase (it);
-			}
-			else
-				break;
-		return isLast;
+		it->second->AddOutOfSequenceFragment (fragmentNum, fragment);
 	}
 
 	void SSU2Session::HandleRelayRequest (const uint8_t * buf, size_t len)
diff --git a/libi2pd/SSU2Session.h b/libi2pd/SSU2Session.h
index b1cc234f..97469ed4 100644
--- a/libi2pd/SSU2Session.h
+++ b/libi2pd/SSU2Session.h
@@ -177,9 +177,12 @@ namespace transport
 		std::shared_ptr<I2NPMessage> msg;
 		int nextFragmentNum;
 		uint32_t lastFragmentInsertTime; // in seconds
-		std::map<int, std::shared_ptr<Fragment> > outOfSequenceFragments;
+		std::shared_ptr<Fragment> secondFragment; // fragment #1
+		std::unique_ptr<std::map<int, std::shared_ptr<Fragment> > > outOfSequenceFragments; // fragments #2 and more
 
 		void AttachNextFragment (const uint8_t * fragment, size_t fragmentSize);
+		bool ConcatOutOfSequenceFragments (); // true if message complete
+		void AddOutOfSequenceFragment (int fragmentNum, std::shared_ptr<Fragment> fragment);
 	};
 
 	struct SSU2SentPacket
@@ -306,7 +309,6 @@ namespace transport
 			bool UpdateReceivePacketNum (uint32_t packetNum); // for Ack, returns false if duplicate
 			void HandleFirstFragment (const uint8_t * buf, size_t len);
 			void HandleFollowOnFragment (const uint8_t * buf, size_t len);
-			bool ConcatOutOfSequenceFragments (std::shared_ptr<SSU2IncompleteMessage> m); // true if message complete
 			void HandleRelayRequest (const uint8_t * buf, size_t len);
 			void HandleRelayIntro (const uint8_t * buf, size_t len, int attempts = 0);
 			void HandleRelayResponse (const uint8_t * buf, size_t len);