From 5361e11395114c4d6ccdb4489866364d906cc74a Mon Sep 17 00:00:00 2001
From: Jeff Becker <jeff@i2p.rocks>
Date: Tue, 20 Mar 2018 15:14:31 -0400
Subject: [PATCH] fix race

---
 libi2pd/NTCPSession.cpp | 33 +++++++++++++++++++++------------
 libi2pd/NTCPSession.h   |  6 ++++--
 2 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/libi2pd/NTCPSession.cpp b/libi2pd/NTCPSession.cpp
index e3d6c004..de4306e6 100644
--- a/libi2pd/NTCPSession.cpp
+++ b/libi2pd/NTCPSession.cpp
@@ -24,6 +24,12 @@ namespace i2p
 {
 namespace transport
 {
+
+	struct NTCPWork
+	{
+		std::shared_ptr<NTCPSession> session;
+	};
+	
 	NTCPSession::NTCPSession (NTCPServer& server, std::shared_ptr<const i2p::data::RouterInfo> in_RemoteRouter):
 		TransportSession (in_RemoteRouter, NTCP_ESTABLISH_TIMEOUT),
 		m_Server (server), m_Socket (m_Server.GetService ()),
@@ -177,18 +183,20 @@ namespace transport
 				}
 			}
 			// TODO: check for number of pending keys
-			auto s = shared_from_this ();
-			m_Server.Work(s, [s]() -> std::function<void(void)> {
-					if (!s->m_DHKeysPair)
-						s->m_DHKeysPair = transports.GetNextDHKeysPair ();
-					s->CreateAESKey (s->m_Establisher->phase1.pubKey);
-					return std::bind(&NTCPSession::SendPhase2, s);
+			auto work = new NTCPWork{shared_from_this()};
+			m_Server.Work(work->session, [work]() -> std::function<void(void)> {
+					if (!work->session->m_DHKeysPair)
+						work->session->m_DHKeysPair = transports.GetNextDHKeysPair ();
+					work->session->CreateAESKey (work->session->m_Establisher->phase1.pubKey);
+					return std::bind(&NTCPSession::SendPhase2, work->session, work);
 			});
 		}
 	}
 
-	void NTCPSession::SendPhase2 ()
+	void NTCPSession::SendPhase2 (NTCPWork * work)
 	{
+		if(work)
+			delete work;
 		const uint8_t * y = m_DHKeysPair->GetPublicKey ();
 		memcpy (m_Establisher->phase2.pubKey, y, 256);
 		uint8_t xy[512];
@@ -241,16 +249,17 @@ namespace transport
 		}
 		else
 		{
-			auto s = shared_from_this ();
-			m_Server.Work(s, [s]() -> std::function<void(void)> {
-				s->CreateAESKey (s->m_Establisher->phase2.pubKey);
-				return std::bind(&NTCPSession::HandlePhase2, s);
+			auto work = new NTCPWork{shared_from_this()};
+			m_Server.Work(work->session, [work]() -> std::function<void(void)> {
+				work->session->CreateAESKey (work->session->m_Establisher->phase2.pubKey);
+				return std::bind(&NTCPSession::HandlePhase2, work->session, work);
 			});
 		}
 	}
 
-	void NTCPSession::HandlePhase2 ()
+	void NTCPSession::HandlePhase2 (NTCPWork * work)
 	{
+		if(work) delete work;
 		m_Decryption.SetIV (m_Establisher->phase2.pubKey + 240);
 		m_Encryption.SetIV (m_Establisher->phase1.HXxorHI + 16);
 
diff --git a/libi2pd/NTCPSession.h b/libi2pd/NTCPSession.h
index 5335fbdd..e2207eb7 100644
--- a/libi2pd/NTCPSession.h
+++ b/libi2pd/NTCPSession.h
@@ -35,6 +35,8 @@ namespace transport
 		} encrypted;
 	};
 
+	struct NTCPWork;
+
 	const size_t NTCP_MAX_MESSAGE_SIZE = 16384;
 	const size_t NTCP_BUFFER_SIZE = 1028; // fits 1 tunnel data message
 	const int NTCP_CONNECT_TIMEOUT = 5; // 5 seconds
@@ -77,12 +79,12 @@ namespace transport
 			void SendPhase3 ();
 			void HandlePhase1Sent (const boost::system::error_code& ecode,  std::size_t bytes_transferred);
 			void HandlePhase2Received (const boost::system::error_code& ecode, std::size_t bytes_transferred);
-			void HandlePhase2 ();
+			void HandlePhase2 (NTCPWork * work=nullptr);
 			void HandlePhase3Sent (const boost::system::error_code& ecode, std::size_t bytes_transferred, uint32_t tsA);
 			void HandlePhase4Received (const boost::system::error_code& ecode, std::size_t bytes_transferred, uint32_t tsA);
 
 			//server
-			void SendPhase2 ();
+			void SendPhase2 (NTCPWork * work=nullptr);
 			void SendPhase4 (uint32_t tsA, uint32_t tsB);
 			void HandlePhase1Received (const boost::system::error_code& ecode, std::size_t bytes_transferred);
 			void HandlePhase2Sent (const boost::system::error_code& ecode, std::size_t bytes_transferred, uint32_t tsB);