Fix leak of scripting instances when new windows are created

5 files changed, 56 insertions, 7 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 0b1949d7..5fcebda9 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -5564,6 +5564,7 @@ __attribute__ ((format (printf, 1, 2)))
 		static void SetCurrentAddressCallback(void* ctxt, uint64_t addr);
 		static void SetCurrentSelectionCallback(void* ctxt, uint64_t begin, uint64_t end);
 		static char* CompleteInputCallback(void* ctxt, const char* text, uint64_t state);
+		static void StopCallback(void* ctxt);
 
 		virtual void DestroyInstance();
 
@@ -5576,6 +5577,7 @@ __attribute__ ((format (printf, 1, 2)))
 		virtual void SetCurrentAddress(uint64_t addr);
 		virtual void SetCurrentSelection(uint64_t begin, uint64_t end);
 		virtual std::string CompleteInput(const std::string& text, uint64_t state);
+		virtual void Stop();
 
 		void Output(const std::string& text);
 		void Error(const std::string& text);
@@ -5603,6 +5605,7 @@ __attribute__ ((format (printf, 1, 2)))
 		virtual void SetCurrentAddress(uint64_t addr) override;
 		virtual void SetCurrentSelection(uint64_t begin, uint64_t end) override;
 		virtual std::string CompleteInput(const std::string& text, uint64_t state) override;
+		virtual void Stop() override;
 	};
 
 	class ScriptingProvider: public StaticCoreRefCountObject<BNScriptingProvider>
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 89b506b4..b034eaec 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2313,6 +2313,8 @@ extern "C"
 	{
 		void* context;
 		void (*destroyInstance)(void* ctxt);
+		void (*externalRefTaken)(void* ctxt);
+		void (*externalRefReleased)(void* ctxt);
 		BNScriptingProviderExecuteResult (*executeScriptInput)(void* ctxt, const char* input);
 		void (*cancelScriptInput)(void* ctxt);
 		void (*setCurrentBinaryView)(void* ctxt, BNBinaryView* view);
@@ -2321,6 +2323,7 @@ extern "C"
 		void (*setCurrentAddress)(void* ctxt, uint64_t addr);
 		void (*setCurrentSelection)(void* ctxt, uint64_t begin, uint64_t end);
 		char* (*completeInput)(void* ctxt, const char* text, uint64_t state);
+		void (*stop)(void* ctxt);
 	};
 
 	struct BNScriptingProviderCallbacks
@@ -5285,6 +5288,7 @@ __attribute__ ((format (printf, 1, 2)))
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentSelection(BNScriptingInstance* instance,
 		uint64_t begin, uint64_t end);
 	BINARYNINJACOREAPI char* BNScriptingInstanceCompleteInput(BNScriptingInstance* instance, const char* text, uint64_t state);
+	BINARYNINJACOREAPI void BNStopScriptingInstance(BNScriptingInstance* instance);
 
 	// Main thread actions
 	BINARYNINJACOREAPI void BNRegisterMainThread(BNMainThreadCallbacks* callbacks);
diff --git a/python/scriptingprovider.py b/python/scriptingprovider.py
index e4061ee3..a336d99b 100644
--- a/python/scriptingprovider.py
+++ b/python/scriptingprovider.py
@@ -110,11 +110,14 @@ class ScriptingOutputListener:
 
 
 class ScriptingInstance:
+	_registered_instances = []
+
 	def __init__(self, provider, handle = None):
 		if handle is None:
 			self._cb = core.BNScriptingInstanceCallbacks()
 			self._cb.context = 0
-			self._cb.destroyInstance = self._cb.destroyInstance.__class__(self._destroy_instance)
+			self._cb.externalRefTaken = self._cb.externalRefTaken.__class__(self._external_ref_taken)
+			self._cb.externalRefReleased = self._cb.externalRefReleased.__class__(self._external_ref_released)
 			self._cb.executeScriptInput = self._cb.executeScriptInput.__class__(self._execute_script_input)
 			self._cb.cancelScriptInput = self._cb.cancelScriptInput.__class__(self._cancel_script_input)
 			self._cb.setCurrentBinaryView = self._cb.setCurrentBinaryView.__class__(self._set_current_binary_view)
@@ -124,6 +127,7 @@ class ScriptingInstance:
 			self._cb.setCurrentSelection = self._cb.setCurrentSelection.__class__(self._set_current_selection)
 			self._cb.completeInput = self._cb.completeInput.__class__(self._complete_input)
 			self._cb.completeInput.restype = ctypes.c_void_p
+			self._cb.stop = self._cb.stop.__class__(self._stop)
 			self.handle = core.BNInitScriptingInstance(provider.handle, self._cb)
 			self.delimiters = ' \t\n`~!@#$%^&*()-=+{}\\|;:\'",<>/?'
 		else:
@@ -134,9 +138,15 @@ class ScriptingInstance:
 		if core is not None:
 			core.BNFreeScriptingInstance(self.handle)
 
-	def _destroy_instance(self, ctxt):
+	def _external_ref_taken(self, ctxt):
 		try:
-			self.perform_destroy_instance()
+			self.__class__._registered_instances.append(self)
+		except:
+			log_error(traceback.format_exc())
+
+	def _external_ref_released(self, ctxt):
+		try:
+			self.__class__._registered_instances.remove(self)
 		except:
 			log_error(traceback.format_exc())
 
@@ -213,9 +223,11 @@ class ScriptingInstance:
 			log_error(traceback.format_exc())
 			return "".encode("utf-8")
 
-	@abc.abstractmethod
-	def perform_destroy_instance(self):
-		raise NotImplementedError
+	def _stop(self, ctxt):
+		try:
+			self.perform_stop()
+		except:
+			log_error(traceback.format_exc())
 
 	@abc.abstractmethod
 	def perform_execute_script_input(self, text):
@@ -249,6 +261,10 @@ class ScriptingInstance:
 	def perform_complete_input(self, text:str, state) -> str:
 		return NotImplemented
 
+	@abc.abstractmethod
+	def perform_stop(self):
+		return NotImplemented
+
 	@property
 	def input_ready_state(self):
 		return core.BNGetScriptingInstanceInputReadyState(self.handle)
@@ -293,6 +309,9 @@ class ScriptingInstance:
 	def complete_input(self, text, state):
 		return core.BNScriptingInstanceCompleteInput(self.handle, text, state)
 
+	def stop(self):
+		core.BNStopScriptingInstance(self.handle)
+
 	def register_output_listener(self, listener):
 		listener._register(self.handle)
 		self.listeners.append(listener)
@@ -728,7 +747,7 @@ from binaryninja import *
 		self.input_ready_state = ScriptingProviderInputReadyState.ReadyForScriptExecution
 
 	@abc.abstractmethod
-	def perform_destroy_instance(self):
+	def perform_stop(self):
 		self.interpreter.end()
 
 	@abc.abstractmethod
diff --git a/scriptingprovider.cpp b/scriptingprovider.cpp
index 4c11ccc4..76cb67c2 100644
--- a/scriptingprovider.cpp
+++ b/scriptingprovider.cpp
@@ -54,6 +54,8 @@ ScriptingInstance::ScriptingInstance(ScriptingProvider* provider)
 	BNScriptingInstanceCallbacks cb;
 	cb.context = this;
 	cb.destroyInstance = DestroyInstanceCallback;
+	cb.externalRefTaken = nullptr;
+	cb.externalRefReleased = nullptr;
 	cb.executeScriptInput = ExecuteScriptInputCallback;
 	cb.cancelScriptInput = CancelScriptInputCallback;
 	cb.setCurrentBinaryView = SetCurrentBinaryViewCallback;
@@ -62,6 +64,7 @@ ScriptingInstance::ScriptingInstance(ScriptingProvider* provider)
 	cb.setCurrentAddress = SetCurrentAddressCallback;
 	cb.setCurrentSelection = SetCurrentSelectionCallback;
 	cb.completeInput = CompleteInputCallback;
+	cb.stop = StopCallback;
 	AddRefForRegistration();
 	m_object = BNInitScriptingInstance(provider->GetObject(), &cb);
 }
@@ -136,6 +139,13 @@ char* ScriptingInstance::CompleteInputCallback(void* ctxt, const char* text, uin
 }
 
 
+void ScriptingInstance::StopCallback(void* ctxt)
+{
+	ScriptingInstance* instance = (ScriptingInstance*)ctxt;
+	instance->Stop();
+}
+
+
 void ScriptingInstance::DestroyInstance()
 {
 	ReleaseForRegistration();
@@ -225,6 +235,11 @@ void ScriptingInstance::SetDelimiters(const std::string& delimiters)
 }
 
 
+void ScriptingInstance::Stop()
+{
+}
+
+
 CoreScriptingInstance::CoreScriptingInstance(BNScriptingInstance* instance): ScriptingInstance(instance)
 {
 }
@@ -281,6 +296,12 @@ std::string CoreScriptingInstance::CompleteInput(const std::string& text, uint64
 }
 
 
+void CoreScriptingInstance::Stop()
+{
+	BNStopScriptingInstance(m_object);
+}
+
+
 ScriptingProvider::ScriptingProvider(const string& name, const string& apiName): m_nameForRegister(name), m_apiNameForRegister(apiName)
 {
 }
diff --git a/ui/scriptingconsole.h b/ui/scriptingconsole.h
index b0e20253..3726514b 100644
--- a/ui/scriptingconsole.h
+++ b/ui/scriptingconsole.h
@@ -189,4 +189,6 @@ public:
 	void moveUpInHistory();
 	void moveDownInHistory();
 	std::vector<std::string> reverseSearch(const QString& text);
+
+	void closing();
 };