use shared_ptr for NetDb's I2NPMessages

ipslot · May 11, 2015 · 2b797fc · 2b797fc
1 parent 5cd557e
commit 2b797fc
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 17 deletions.
diff --git a/I2NPProtocol.cpp b/I2NPProtocol.cpp
@@ -36,6 +36,11 @@ namespace i2p
 		delete msg;
 	}	
 
+	std::shared_ptr<I2NPMessage> ToSharedI2NPMessage (I2NPMessage * msg)
+	{
+		return std::shared_ptr<I2NPMessage>(msg, DeleteI2NPMessage);
+	}
+
 	static std::atomic<uint32_t> I2NPmsgID(0); // TODO: create class
 	void FillI2NPMessageHeader (I2NPMessage * msg, I2NPMessageType msgType, uint32_t replyMsgID)
 	{

diff --git a/I2NPProtocol.h b/I2NPProtocol.h
@@ -196,6 +196,8 @@ namespace tunnel
 	I2NPMessage * NewI2NPShortMessage ();
 	I2NPMessage * NewI2NPMessage (size_t len);
 	void DeleteI2NPMessage (I2NPMessage * msg);
+	std::shared_ptr<I2NPMessage> ToSharedI2NPMessage (I2NPMessage * msg);
+
 	void FillI2NPMessageHeader (I2NPMessage * msg, I2NPMessageType msgType, uint32_t replyMsgID = 0);
 	void RenewI2NPMessageHeader (I2NPMessage * msg);
 	I2NPMessage * CreateI2NPMessage (I2NPMessageType msgType, const uint8_t * buf, int len, uint32_t replyMsgID = 0);	

diff --git a/NetDb.cpp b/NetDb.cpp
@@ -102,18 +102,18 @@ namespace data
 						{
 							case eI2NPDatabaseStore:	
 								LogPrint ("DatabaseStore");
-								HandleDatabaseStoreMsg (msg);
+								HandleDatabaseStoreMsg (ToSharedI2NPMessage (msg));
 							break;
 							case eI2NPDatabaseSearchReply:
 								LogPrint ("DatabaseSearchReply");
-								HandleDatabaseSearchReplyMsg (msg);
+								HandleDatabaseSearchReplyMsg (ToSharedI2NPMessage (msg));
 							break;
 							case eI2NPDatabaseLookup:
 								LogPrint ("DatabaseLookup");
-								HandleDatabaseLookupMsg (msg);
+								HandleDatabaseLookupMsg (ToSharedI2NPMessage (msg));
 							break;	
 							default: // WTF?
-								LogPrint ("NetDb: unexpected message type ", msg->GetTypeID ());
+								LogPrint (eLogError, "NetDb: unexpected message type ", msg->GetTypeID ());
 								i2p::HandleI2NPMessage (msg);
 						}	
 						if (numMsgs > 100) break;
@@ -469,15 +469,14 @@ namespace data
 		}	
 	}	
 
-	void NetDb::HandleDatabaseStoreMsg (I2NPMessage * m)
+	void NetDb::HandleDatabaseStoreMsg (std::shared_ptr<I2NPMessage> m)
 	{	
 		const uint8_t * buf = m->GetPayload ();
 		size_t len = m->GetSize ();		
 		IdentHash ident (buf + DATABASE_STORE_KEY_OFFSET);
 		if (ident.IsZero ())
 		{
 			LogPrint (eLogError, "Database store with zero ident. Dropped");
-			i2p::DeleteI2NPMessage (m);
 			return;
 		}	
 		uint32_t replyToken = bufbe32toh (buf + DATABASE_STORE_REPLY_TOKEN_OFFSET);
@@ -539,7 +538,6 @@ namespace data
 			if (size > 2048 || size > len - offset)
 			{
 				LogPrint ("Invalid RouterInfo length ", (int)size);
-				i2p::DeleteI2NPMessage (m);
 				return;
 			}	
 			try
@@ -562,10 +560,9 @@ namespace data
 				LogPrint (eLogError, "DatabaseStore: ", ex.what ());
 			}
 		}	
-		i2p::DeleteI2NPMessage (m);
 	}	
 
-	void NetDb::HandleDatabaseSearchReplyMsg (I2NPMessage * msg)
+	void NetDb::HandleDatabaseSearchReplyMsg (std::shared_ptr<I2NPMessage> msg)
 	{
 		uint8_t * buf = msg->GetPayload ();
 		char key[48];
@@ -652,18 +649,15 @@ namespace data
 			else
 				LogPrint ("Bayan");	
 		}	
-
-		i2p::DeleteI2NPMessage (msg);
 	}	
 
-	void NetDb::HandleDatabaseLookupMsg (I2NPMessage * msg)
+	void NetDb::HandleDatabaseLookupMsg (std::shared_ptr<I2NPMessage> msg)
 	{
 		uint8_t * buf = msg->GetPayload ();
 		IdentHash ident (buf);
 		if (ident.IsZero ())
 		{
 			LogPrint (eLogError, "DatabaseLookup for zero ident. Ignored");
-			i2p::DeleteI2NPMessage (msg);
 			return;
 		}	
 		char key[48];
@@ -789,7 +783,6 @@ namespace data
 			else
 				transports.SendMessage (buf+32, replyMsg);
 		}
-		i2p::DeleteI2NPMessage (msg);
 	}	
 
 	void NetDb::Explore (int numDestinations)

diff --git a/NetDb.h b/NetDb.h
@@ -41,9 +41,9 @@ namespace data
 
 			void RequestDestination (const IdentHash& destination, RequestedDestination::RequestComplete requestComplete = nullptr);			
 
-			void HandleDatabaseStoreMsg (I2NPMessage * msg);
-			void HandleDatabaseSearchReplyMsg (I2NPMessage * msg);
-			void HandleDatabaseLookupMsg (I2NPMessage * msg);			
+			void HandleDatabaseStoreMsg (std::shared_ptr<I2NPMessage> msg);
+			void HandleDatabaseSearchReplyMsg (std::shared_ptr<I2NPMessage> msg);
+			void HandleDatabaseLookupMsg (std::shared_ptr<I2NPMessage> msg);			
 
 			std::shared_ptr<const RouterInfo> GetRandomRouter () const;
 			std::shared_ptr<const RouterInfo> GetRandomRouter (std::shared_ptr<const RouterInfo> compatibleWith) const;