From ca749691f8ce741dc5a9582992c4315b92c98938 Mon Sep 17 00:00:00 2001
From: Akash Mozumdar <akashmozumdar@gmail.com>
Date: Sat, 1 Sep 2018 04:23:29 -0400
Subject: [PATCH] refactor extension code and add error handling

---
 GUI/extensions.cpp | 18 +++++++++++-------
 GUI/extensions.h   |  2 +-
 GUI/mainwindow.cpp | 11 +++++++----
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/GUI/extensions.cpp b/GUI/extensions.cpp
index 820f0c3..3a08d11 100644
--- a/GUI/extensions.cpp
+++ b/GUI/extensions.cpp
@@ -26,7 +26,7 @@ std::map<int, QString> LoadExtensions()
 	return extensionNames;
 }
 
-std::wstring DispatchSentenceToExtensions(std::wstring sentence, std::unordered_map<std::string, int> miscInfo)
+bool DispatchSentenceToExtensions(std::wstring& sentence, std::unordered_map<std::string, int> miscInfo)
 {
 	wchar_t* sentenceBuffer = (wchar_t*)malloc((sentence.size() + 1) * sizeof(wchar_t));
 	wcscpy(sentenceBuffer, sentence.c_str());
@@ -34,15 +34,19 @@ std::wstring DispatchSentenceToExtensions(std::wstring sentence, std::unordered_
 	InfoForExtension* miscInfoTraverser = miscInfoLinkedList;
 	for (auto& i : miscInfo) miscInfoTraverser = miscInfoTraverser->nextProperty = new InfoForExtension{ i.first.c_str(), i.second, nullptr };
 	extenMutex.lock_shared();
-	for (auto i : extensions)
+	try
 	{
-		wchar_t* prev = sentenceBuffer;
-		sentenceBuffer = i.second(sentenceBuffer, miscInfoLinkedList);
-		if (sentenceBuffer != prev) free((void*)prev);
+		for (auto i : extensions)
+		{
+			wchar_t* prev = sentenceBuffer;
+			sentenceBuffer = i.second(sentenceBuffer, miscInfoLinkedList);
+			if (sentenceBuffer != prev) free((void*)prev);
+		}
 	}
+	catch (...) { sentenceBuffer[0] = 0; }
 	extenMutex.unlock_shared();
 	delete miscInfoLinkedList;
-	std::wstring newSentence = std::wstring(sentenceBuffer);
+	sentence = std::wstring(sentenceBuffer);
 	free((void*)sentenceBuffer);
-	return newSentence;
+	return sentence.size() > 0;
 }
diff --git a/GUI/extensions.h b/GUI/extensions.h
index ee06e33..5cf12cd 100644
--- a/GUI/extensions.h
+++ b/GUI/extensions.h
@@ -5,7 +5,7 @@
 #include <map>
 
 std::map<int, QString> LoadExtensions();
-std::wstring DispatchSentenceToExtensions(std::wstring sentence, std::unordered_map<std::string, int> miscInfo);
+bool DispatchSentenceToExtensions(std::wstring& sentence, std::unordered_map<std::string, int> miscInfo);
 struct InfoForExtension
 {
 	~InfoForExtension() { if (nextProperty) delete nextProperty; };
diff --git a/GUI/mainwindow.cpp b/GUI/mainwindow.cpp
index c017451..bcedc2f 100644
--- a/GUI/mainwindow.cpp
+++ b/GUI/mainwindow.cpp
@@ -84,9 +84,11 @@ void MainWindow::AddThread(TextThread* thread)
 	);
 	thread->RegisterOutputCallBack([&](TextThread* thread, std::wstring output)
 	{
-		output = DispatchSentenceToExtensions(output, GetInfoForExtensions(thread));
-		output += L"\r\n";
-		emit SigThreadOutput(thread, QString::fromStdWString(output));
+		if (DispatchSentenceToExtensions(output, GetInfoForExtensions(thread)))
+		{
+			output += L"\r\n";
+			emit SigThreadOutput(thread, QString::fromStdWString(output));
+		}
 		return output;
 	});
 }
@@ -150,7 +152,8 @@ std::unordered_map<std::string, int> MainWindow::GetInfoForExtensions(TextThread
 	{ "text number", 0 },
 	{ "process id", thread->GetThreadParam().pid },
 	{ "hook address", (int)thread->GetThreadParam().hook },
-	{ "hook address (upper 32 bits)", (int)(thread->GetThreadParam().hook >> 32) }
+	{ "hook address (upper 32 bits)", (int)(thread->GetThreadParam().hook >> 32) },
+	{ "text handle", (int)thread }
 	};
 }