Fix cryptography memory leaks :P

src/Components/Modules/Node.cpp

@@ -2,6 +2,8 @@
 
 namespace Components
 {
+	Utils::Cryptography::ECDSA::Key Node::SignatureKey;
+
 	std::vector<Node::NodeEntry> Node::Nodes;
 	std::vector<Node::DediEntry> Node::Dedis;
 
@@ -396,6 +398,9 @@ namespace Components
 		// ZoneBuilder doesn't require node stuff
 		if (ZoneBuilder::IsEnabled()) return;
 
+		// Generate our ECDSA key
+		Node::SignatureKey = Utils::Cryptography::ECDSA::GenerateKey(512);
+
 		Dvar::OnInit([] ()
 		{
 			Node::Dedis.clear();
@@ -570,6 +575,8 @@ namespace Components
 
 	Node::~Node()
 	{
+		Node::SignatureKey.Free();
+
 		Node::StoreNodes(true);
 		Node::Nodes.clear();
 		Node::Dedis.clear();

src/Components/Modules/Node.hpp

@@ -44,6 +44,7 @@ namespace Components
 		struct NodeEntry
 		{
 			Network::Address address;
+			Utils::Cryptography::ECDSA::Key publicKey;
 			EntryState state;
 			int lastTime; // Last time we heard anything from the server itself
 			int lastHeartbeat; // Last time we got a heartbeat from it
@@ -84,6 +85,8 @@ namespace Components
 		};
 #pragma pack(pop)
 
+		static Utils::Cryptography::ECDSA::Key SignatureKey;
+
 		static std::vector<NodeEntry> Nodes;
 		static std::vector<DediEntry> Dedis;
 

src/Utils/Cryptography.cpp

@@ -30,7 +30,6 @@ namespace Utils
 			ECDSA::Key key;
 
 			register_prng(&sprng_desc);
-			register_hash(&sha1_desc);
 
 			ltc_mp = ltm_desc;
 
@@ -41,7 +40,7 @@ namespace Utils
 
 		std::string ECDSA::SignMessage(Key key, std::string message)
 		{
-			uint8_t buffer[0x200]; // Default size is 512
+			uint8_t buffer[512];
 			DWORD length = sizeof(buffer);
 
 			register_prng(&sprng_desc);
@@ -81,7 +80,7 @@ namespace Utils
 
 		std::string RSA::SignMessage(RSA::Key key, std::string message)
 		{
-			uint8_t buffer[0x200]; // Default size is 512
+			uint8_t buffer[512];
 			DWORD length = sizeof(buffer);
 
 			register_prng(&sprng_desc);

src/Utils/Cryptography.hpp

@@ -14,21 +14,33 @@ namespace Utils
 			class Key
 			{
 			public:
-				Key() { ZeroMemory(&this->KeyStorage, sizeof(this->KeyStorage)); };
-				Key(ecc_key* key) : Key(*key) {};
-				Key(ecc_key key) : KeyStorage(key) {};
-				Key(const Key& obj) : KeyStorage(obj.KeyStorage) {};
+				Key() : KeyStorage(new ecc_key)
+				{
+					ZeroMemory(this->KeyStorage.get(), sizeof(*this->KeyStorage.get()));
+				};
+				Key(ecc_key* key) : Key() { std::memmove(this->KeyStorage.get(), key, sizeof(*key)); };
+				Key(ecc_key key) : Key(&key) {};
+				~Key() 
+				{
+					if (this->KeyStorage.use_count() <= 1)
+					{
+						this->Free();
+					}
+				};
 
-				~Key() {}
+				bool IsValid()
+				{
+					return (!Utils::MemIsSet(this->KeyStorage.get(), 0, sizeof(*this->KeyStorage.get())));
+				}
 
 				ecc_key* GetKeyPtr()
 				{
-					return &this->KeyStorage;
+					return this->KeyStorage.get();
 				}
 
 				std::string GetPublicKey()
 				{
-					uint8_t buffer[0x1000] = { 0 };
+					uint8_t buffer[512] = { 0 };
 					DWORD length = sizeof(buffer);
 
 					if (ecc_ansi_x963_export(this->GetKeyPtr(), buffer, &length) == CRYPT_OK)
@@ -39,8 +51,28 @@ namespace Utils
 					return "";
 				}
 
+				void Set(std::string pubKeyBuffer)
+				{
+					this->Free();
+
+					if (ecc_ansi_x963_import(reinterpret_cast<const uint8_t*>(pubKeyBuffer.data()), pubKeyBuffer.size(), this->KeyStorage.get()) != CRYPT_OK)
+					{
+						ZeroMemory(this->KeyStorage.get(), sizeof(*this->KeyStorage.get()));
+					}
+				}
+
+				void Free()
+				{
+					if (this->IsValid())
+					{
+						ecc_free(this->KeyStorage.get());
+					}
+
+					ZeroMemory(this->KeyStorage.get(), sizeof(*this->KeyStorage.get()));
+				}
+
 			private:
-				ecc_key KeyStorage;
+				std::shared_ptr<ecc_key> KeyStorage;
 			};
 
 			static Key GenerateKey(int bits);
@@ -54,20 +86,42 @@ namespace Utils
 			class Key
 			{
 			public:
-				Key() { ZeroMemory(&this->KeyStorage, sizeof(this->KeyStorage)); };
-				Key(rsa_key* key) : Key(*key) {};
-				Key(rsa_key key) : KeyStorage(key) {};
-				Key(const Key& obj) : KeyStorage(obj.KeyStorage) {};
-
-				~Key() {}
+				Key() : KeyStorage(new rsa_key)
+				{
+					ZeroMemory(this->KeyStorage.get(), sizeof(*this->KeyStorage.get()));
+				};
+				Key(rsa_key* key) : Key() { std::memmove(this->KeyStorage.get(), key, sizeof(*key)); };
+				Key(rsa_key key) : Key(&key) {};
+				~Key()
+				{
+					if (this->KeyStorage.use_count() <= 1)
+					{
+						this->Free();
+					}
+				};
 
 				rsa_key* GetKeyPtr()
 				{
-					return &this->KeyStorage;
+					return this->KeyStorage.get();
+				}
+
+				bool IsValid()
+				{
+					return (!Utils::MemIsSet(this->KeyStorage.get(), 0, sizeof(*this->KeyStorage.get())));
+				}
+
+				void Free()
+				{
+					if (this->IsValid())
+					{
+						rsa_free(this->KeyStorage.get());
+					}
+
+					ZeroMemory(this->KeyStorage.get(), sizeof(*this->KeyStorage.get()));
 				}
 
 			private:
-				rsa_key KeyStorage;
+				std::shared_ptr<rsa_key> KeyStorage;
 			};
 
 			static Key GenerateKey(int bits);

src/Utils/Utils.cpp

@@ -30,6 +30,22 @@ namespace Utils
 		return (strstr(haystack.data(), needle.data()) == (haystack.data() + haystack.size() - needle.size()));
 	}
 
+	// Complementary function for memset, which checks if a memory is set
+	bool MemIsSet(void* mem, char chr, size_t length) 
+	{
+		char* memArr = reinterpret_cast<char*>(mem);
+
+		for (size_t i = 0; i < length; i++)
+		{
+			if (memArr[i] != chr)
+			{
+				return false;
+			}
+		}
+
+		return true;
+	}
+
 	std::vector<std::string> Explode(const std::string& str, char delim)
 	{
 		std::vector<std::string> result;

src/Utils/Utils.hpp

@@ -20,6 +20,8 @@ namespace Utils
 	void WriteFile(std::string file, std::string data);
 	std::string ReadFile(std::string file);
 
+	bool MemIsSet(void* mem, char chr, size_t length);
+
 	class InfoString
 	{
 	public: