Fix memory leak and use after free bugs in BinaryView and FlowGraph objects

1 files changed, 12 insertions, 0 deletions

diff --git a/flowgraph.cpp b/flowgraph.cpp
index 3de9c90e..9e3466c8 100644
--- a/flowgraph.cpp
+++ b/flowgraph.cpp
@@ -76,6 +76,11 @@ FlowGraph::FlowGraph()
 	callbacks.prepareForLayout = PrepareForLayoutCallback;
 	callbacks.populateNodes = PopulateNodesCallback;
 	callbacks.completeLayout = CompleteLayoutCallback;
+	callbacks.update = UpdateCallback;
+	callbacks.freeObject = FreeObjectCallback;
+	callbacks.externalRefTaken = nullptr;
+	callbacks.externalRefReleased = nullptr;
+	AddRefForRegistration();
 	m_object = BNCreateCustomFlowGraph(&callbacks);
 }
 
@@ -117,6 +122,13 @@ BNFlowGraph* FlowGraph::UpdateCallback(void* ctxt)
 }
 
 
+void FlowGraph::FreeObjectCallback(void* ctxt)
+{
+	FlowGraph* graph = (FlowGraph*)ctxt;
+	graph->ReleaseForRegistration();
+}
+
+
 void FlowGraph::FinishPrepareForLayout()
 {
 	BNFinishPrepareForLayout(m_object);