From a2249f0a82ca29215b039e15e7755be9f6012735 Mon Sep 17 00:00:00 2001
From: orignal <i2porignal@yandex.ru>
Date: Wed, 24 Jan 2024 13:52:13 -0500
Subject: [PATCH] identify server tunnel session but from ant to ports

---
 libi2pd_client/UDPTunnel.cpp | 57 +++++++++++++++++++++++++-----------
 libi2pd_client/UDPTunnel.h   |  7 +++--
 2 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/libi2pd_client/UDPTunnel.cpp b/libi2pd_client/UDPTunnel.cpp
index 1e4b3d7c..3fcf8b35 100644
--- a/libi2pd_client/UDPTunnel.cpp
+++ b/libi2pd_client/UDPTunnel.cpp
@@ -1,5 +1,5 @@
 /*
-* Copyright (c) 2013-2022, The PurpleI2P Project
+* Copyright (c) 2013-2024, The PurpleI2P Project
 *
 * This file is part of Purple i2pd project and licensed under BSD3
 *
@@ -27,8 +27,17 @@ namespace client
 		m_LastSession->LastActivity = i2p::util::GetMillisecondsSinceEpoch();
 	}
 
-	void I2PUDPServerTunnel::HandleRecvFromI2PRaw (uint16_t, uint16_t, const uint8_t * buf, size_t len)
+	void I2PUDPServerTunnel::HandleRecvFromI2PRaw (uint16_t fromPort, uint16_t toPort, const uint8_t * buf, size_t len)
 	{
+		if (m_LastSession && (fromPort != m_LastSession->RemotePort || toPort != m_LastSession->LocalPort))
+		{
+			std::lock_guard<std::mutex> lock(m_SessionsMutex);
+			auto it = m_Sessions.find (GetSessionIndex (fromPort, toPort));
+			if (it != m_Sessions.end ())
+				m_LastSession = it->second;
+			else	
+				m_LastSession = nullptr;
+		}
 		if (m_LastSession)
 		{
 			m_LastSession->IPSocket.send_to(boost::asio::buffer(buf, len), m_RemoteEndpoint);
@@ -41,11 +50,12 @@ namespace client
 		std::lock_guard<std::mutex> lock(m_SessionsMutex);
 		uint64_t now = i2p::util::GetMillisecondsSinceEpoch();
 		auto itr = m_Sessions.begin();
-		while(itr != m_Sessions.end()) {
-			if(now - (*itr)->LastActivity >= delta )
+		while(itr != m_Sessions.end()) 
+		{
+			if(now - itr->second->LastActivity >= delta )
 				itr = m_Sessions.erase(itr);
 			else
-				++itr;
+				itr++;
 		}
 	}
 
@@ -66,15 +76,25 @@ namespace client
 	UDPSessionPtr I2PUDPServerTunnel::ObtainUDPSession(const i2p::data::IdentityEx& from, uint16_t localPort, uint16_t remotePort)
 	{
 		auto ih = from.GetIdentHash();
-		for (auto & s : m_Sessions )
+		auto idx = GetSessionIndex (remotePort, localPort);
 		{
-			if (s->Identity.GetLL()[0] == ih.GetLL()[0] && remotePort == s->RemotePort)
+			std::lock_guard<std::mutex> lock(m_SessionsMutex);
+			auto it = m_Sessions.find (idx);
+			if (it != m_Sessions.end ())
 			{
-				/** found existing session */
-				LogPrint(eLogDebug, "UDPServer: Found session ", s->IPSocket.local_endpoint(), " ", ih.ToBase32());
-				return s;
-			}
-		}
+				if (it->second->Identity.GetLL()[0] == ih.GetLL()[0])
+				{
+					LogPrint(eLogDebug, "UDPServer: Found session ", it->second->IPSocket.local_endpoint(), " ", ih.ToBase32());
+					return it->second;
+				}
+				else
+				{
+					LogPrint(eLogWarning, "UDPServer: Session with from ", remotePort, " and to ", localPort, " ports already exists. But from differend address. Removed");
+					m_Sessions.erase (it);
+				}	
+			}	
+		}	
+		
 		boost::asio::ip::address addr;
 		/** create new udp session */
 		if(m_IsUniqueLocal && m_LocalAddress.is_loopback())
@@ -84,10 +104,12 @@ namespace client
 		}
 		else
 			addr = m_LocalAddress;
-		boost::asio::ip::udp::endpoint ep(addr, 0);
-		m_Sessions.push_back(std::make_shared<UDPSession>(ep, m_LocalDest, m_RemoteEndpoint, ih, localPort, remotePort));
-		auto & back = m_Sessions.back();
-		return back;
+
+		auto s = std::make_shared<UDPSession>(boost::asio::ip::udp::endpoint(addr, 0), 
+			m_LocalDest, m_RemoteEndpoint, ih, localPort, remotePort);
+		std::lock_guard<std::mutex> lock(m_SessionsMutex);
+		m_Sessions.emplace (idx, s);
+		return s;
 	}
 
 	UDPSession::UDPSession(boost::asio::ip::udp::endpoint localEndpoint,
@@ -175,8 +197,9 @@ namespace client
 		std::vector<std::shared_ptr<DatagramSessionInfo> > sessions;
 		std::lock_guard<std::mutex> lock (m_SessionsMutex);
 
-		for (UDPSessionPtr s: m_Sessions)
+		for (auto it: m_Sessions)
 		{
+			auto s = it.second;
 			if (!s->m_Destination) continue;
 			auto info = s->m_Destination->GetInfoForRemote (s->Identity);
 			if (!info) continue;
diff --git a/libi2pd_client/UDPTunnel.h b/libi2pd_client/UDPTunnel.h
index 3233c3f0..6dff41ea 100644
--- a/libi2pd_client/UDPTunnel.h
+++ b/libi2pd_client/UDPTunnel.h
@@ -1,5 +1,5 @@
 /*
-* Copyright (c) 2013-2022, The PurpleI2P Project
+* Copyright (c) 2013-2024, The PurpleI2P Project
 *
 * This file is part of Purple i2pd project and licensed under BSD3
 *
@@ -104,7 +104,8 @@ namespace client
 			void HandleRecvFromI2P (const i2p::data::IdentityEx& from, uint16_t fromPort, uint16_t toPort, const uint8_t * buf, size_t len);
 			void HandleRecvFromI2PRaw (uint16_t fromPort, uint16_t toPort, const uint8_t * buf, size_t len);
 			UDPSessionPtr ObtainUDPSession (const i2p::data::IdentityEx& from, uint16_t localPort, uint16_t remotePort);
-
+			uint32_t GetSessionIndex (uint16_t fromPort, uint16_t toPort) const { return ((uint32_t)fromPort << 16) + toPort; }
+			
 		private:
 
 			bool m_IsUniqueLocal;
@@ -112,7 +113,7 @@ namespace client
 			boost::asio::ip::address m_LocalAddress;
 			boost::asio::ip::udp::endpoint m_RemoteEndpoint;
 			std::mutex m_SessionsMutex;
-			std::vector<UDPSessionPtr> m_Sessions;
+			std::unordered_map<uint32_t, UDPSessionPtr> m_Sessions; // (from port, to port)->session
 			std::shared_ptr<i2p::client::ClientDestination> m_LocalDest;
 			UDPSessionPtr m_LastSession;
 			bool m_Gzip;