From c770bcbf9630a07e5420cbbc12814858c3b4a830 Mon Sep 17 00:00:00 2001
From: Jeff Becker <ampernand@gmail.com>
Date: Sat, 3 Sep 2016 10:24:06 -0400
Subject: [PATCH] prevent race condition in datagram destination

clean up style a bit
---
 AddressBook.cpp | 12 ++++++------
 Datagram.cpp    | 20 ++++++++++++++------
 Datagram.h      | 14 +++++++++-----
 Tag.h           |  9 +++++----
 4 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/AddressBook.cpp b/AddressBook.cpp
index 9f1af809..b993f456 100644
--- a/AddressBook.cpp
+++ b/AddressBook.cpp
@@ -543,11 +543,12 @@ namespace client
 		auto dest = i2p::client::context.GetSharedLocalDestination ();
 		if (dest)
 		{
-      
 			auto datagram = dest->GetDatagramDestination ();
-      if(datagram == nullptr) datagram = dest->CreateDatagramDestination();
-			datagram->SetReceiver (std::bind (&AddressBook::HandleLookupResponse, this, 
-				std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5));
+			if (!datagram)
+				datagram = dest->CreateDatagramDestination ();
+			datagram->SetReceiver (std::bind (&AddressBook::HandleLookupResponse, this,
+				std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, std::placeholders::_5),
+				ADDRESS_RESPONSE_DATAGRAM_PORT);
 		}	
 	}
 	
@@ -557,8 +558,7 @@ namespace client
 		if (dest)
 		{
 			auto datagram = dest->GetDatagramDestination ();
-			if (datagram)
-        datagram->ResetReceiver ();
+			if (datagram) datagram->ResetReceiver (ADDRESS_RESPONSE_DATAGRAM_PORT);
 		}	
 	}
 
diff --git a/Datagram.cpp b/Datagram.cpp
index 13a4af96..cdd468b1 100644
--- a/Datagram.cpp
+++ b/Datagram.cpp
@@ -71,18 +71,26 @@ namespace datagram
 				
 		if (verified)
 		{
-			auto it = m_ReceiversByPorts.find (toPort);
-			if (it != m_ReceiversByPorts.end ())
-				it->second (identity, fromPort, toPort, buf + headerLen, len -headerLen);
-			else if (m_Receiver != nullptr)
-				m_Receiver (identity, fromPort, toPort, buf + headerLen, len -headerLen);
+			auto r = FindReceiver(toPort);
+			if(r)
+				r(identity, fromPort, toPort, buf + headerLen, len -headerLen);
 			else
-				LogPrint (eLogWarning, "Receiver for datagram is not set");	
+				LogPrint (eLogWarning, "DatagramDestination: no receiver for port ", toPort);
 		}
 		else
 			LogPrint (eLogWarning, "Datagram signature verification failed");	
 	}
 
+	DatagramDestination::Receiver DatagramDestination::FindReceiver(uint16_t port)
+	{
+		std::lock_guard<std::mutex> lock(m_ReceiversMutex);
+		Receiver r = m_Receiver;
+		auto itr = m_ReceiversByPorts.find(port);
+		if (itr != m_ReceiversByPorts.end())
+			r = itr->second;
+		return r;
+	}
+
 	void DatagramDestination::HandleDataMessagePayload (uint16_t fromPort, uint16_t toPort, const uint8_t * buf, size_t len)
 	{
 		// unzip it
diff --git a/Datagram.h b/Datagram.h
index bb1bc285..887c82d2 100644
--- a/Datagram.h
+++ b/Datagram.h
@@ -80,8 +80,8 @@ namespace datagram
 			void SetReceiver (const Receiver& receiver) { m_Receiver = receiver; };
 			void ResetReceiver () { m_Receiver = nullptr; };
 
-			void SetReceiver (const Receiver& receiver, uint16_t port) { m_ReceiversByPorts[port] = receiver; };
-			void ResetReceiver (uint16_t port) { m_ReceiversByPorts.erase (port); };
+			void SetReceiver (const Receiver& receiver, uint16_t port) { std::lock_guard<std::mutex> lock(m_ReceiversMutex); m_ReceiversByPorts[port] = receiver; };
+			void ResetReceiver (uint16_t port) { std::lock_guard<std::mutex> lock(m_ReceiversMutex); m_ReceiversByPorts.erase (port); };
     
 		private:
       // clean up after next tick
@@ -96,12 +96,16 @@ namespace datagram
 
 			void HandleDatagram (uint16_t fromPort, uint16_t toPort, const uint8_t * buf, size_t len);
 
+			/** find a receiver by port, if none by port is found try default receiever, otherwise returns nullptr */
+			Receiver FindReceiver(uint16_t port);
+			
 		private:
 			i2p::client::ClientDestination * m_Owner;
-      boost::asio::deadline_timer m_CleanupTimer;
+			boost::asio::deadline_timer m_CleanupTimer;
 			Receiver m_Receiver; // default
-      std::mutex m_SessionsMutex;
-      std::map<i2p::data::IdentHash, std::shared_ptr<DatagramSession> > m_Sessions;
+			std::mutex m_SessionsMutex;
+			std::map<i2p::data::IdentHash, std::shared_ptr<DatagramSession> > m_Sessions;
+			std::mutex m_ReceiversMutex;
 			std::map<uint16_t, Receiver> m_ReceiversByPorts;
 
 			i2p::data::GzipInflator m_Inflator;
diff --git a/Tag.h b/Tag.h
index 6e5638a4..4b393c57 100644
--- a/Tag.h
+++ b/Tag.h
@@ -51,10 +51,11 @@ namespace data {
 			}
 
       /** fill with a value */
-      void Fill(uint8_t c) {
-        memset(m_Buf, c, sz);
-      }
-      
+			void Fill(uint8_t c)
+			{
+				memset(m_Buf, c, sz);
+			}
+			
 			std::string ToBase64 () const
 			{
 				char str[sz*2];