From 178dedf78cb74f1385d409596959a21ed261b21c Mon Sep 17 00:00:00 2001
From: orignal <i2porignal@yandex.ru>
Date: Fri, 30 Dec 2016 17:53:54 -0500
Subject: [PATCH] store relay session directly

---
 SSU.cpp        | 9 +++++++--
 SSU.h          | 4 ++--
 SSUSession.cpp | 2 +-
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/SSU.cpp b/SSU.cpp
index 7950dc45..9f74185f 100644
--- a/SSU.cpp
+++ b/SSU.cpp
@@ -154,7 +154,7 @@ namespace transport
 		}	
 	}	
 	
-	void SSUServer::AddRelay (uint32_t tag, const boost::asio::ip::udp::endpoint& relay)
+	void SSUServer::AddRelay (uint32_t tag, std::shared_ptr<SSUSession> relay)
 	{
 		m_Relays[tag] = relay;
 	}	
@@ -168,7 +168,12 @@ namespace transport
 	{
 		auto it = m_Relays.find (tag);
 		if (it != m_Relays.end ())
-			return FindSession (it->second);
+		{	
+			if (it->second->GetState () == eSessionStateEstablished)
+				return it->second;
+			else
+				m_Relays.erase (it);	
+		}	
 		return nullptr;
 	}
 
diff --git a/SSU.h b/SSU.h
index 8753ff77..adbed739 100644
--- a/SSU.h
+++ b/SSU.h
@@ -57,7 +57,7 @@ namespace transport
 			boost::asio::io_service& GetServiceV6 () { return m_ServiceV6; };
 			const boost::asio::ip::udp::endpoint& GetEndpoint () const { return m_Endpoint; };			
 			void Send (const uint8_t * buf, size_t len, const boost::asio::ip::udp::endpoint& to);
-			void AddRelay (uint32_t tag, const boost::asio::ip::udp::endpoint& relay);
+			void AddRelay (uint32_t tag, std::shared_ptr<SSUSession> relay);
 			void RemoveRelay (uint32_t tag);
 			std::shared_ptr<SSUSession> FindRelaySession (uint32_t tag);
 
@@ -120,7 +120,7 @@ namespace transport
 				m_TerminationTimer, m_TerminationTimerV6;
 			std::list<boost::asio::ip::udp::endpoint> m_Introducers; // introducers we are connected to
 			std::map<boost::asio::ip::udp::endpoint, std::shared_ptr<SSUSession> > m_Sessions, m_SessionsV6;
-			std::map<uint32_t, boost::asio::ip::udp::endpoint> m_Relays; // we are introducer
+			std::map<uint32_t, std::shared_ptr<SSUSession> > m_Relays; // we are introducer
 			std::map<uint32_t, PeerTest> m_PeerTests; // nonce -> creation time in milliseconds
 
 		public:
diff --git a/SSUSession.cpp b/SSUSession.cpp
index e67ff1de..a74c6646 100644
--- a/SSUSession.cpp
+++ b/SSUSession.cpp
@@ -462,7 +462,7 @@ namespace transport
 		{
 			RAND_bytes((uint8_t *)&m_SentRelayTag, 4);
 			if (!m_SentRelayTag) m_SentRelayTag = 1;
-			m_Server.AddRelay (m_SentRelayTag, m_RemoteEndpoint);
+			m_Server.AddRelay (m_SentRelayTag, shared_from_this ());
 		}
 		htobe32buf (payload, m_SentRelayTag); 
 		payload += 4; // relay tag