Skip to content

Commit 2db2d43

Browse files
tinysectinysec
committed
Add custom flow graphs
1 parent 9f623a0 commit 2db2d43

6 files changed

Lines changed: 303 additions & 51 deletions

File tree

‎Function/BNCreateCustomFlowGraph.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ internal static partial class NativeMethods
1818
internal static extern IntPtr BNCreateCustomFlowGraph(
1919

2020
// BNCustomFlowGraph* callbacks
21-
IntPtr callbacks
21+
in BNCustomFlowGraph callbacks
2222
);
2323
}
24-
}
24+
}

‎Handle/BNFlowGraph.cs‎

Lines changed: 52 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,25 @@ namespace BinaryNinja
1010
public abstract class AbstractFlowGraph<T_SELF> : AbstractSafeHandle<T_SELF>
1111
where T_SELF: AbstractFlowGraph<T_SELF>
1212
{
13+
private static readonly NativeDelegates.BNCustomFlowGraphEvent
14+
emptyLayoutCompletion =
15+
AbstractFlowGraph<T_SELF>.IgnoreLayoutCompletion;
16+
17+
private sealed class LayoutCompletionContext
18+
{
19+
private readonly Action callback;
20+
21+
internal LayoutCompletionContext(Action callback)
22+
{
23+
this.callback = callback;
24+
}
25+
26+
internal void Invoke(IntPtr context)
27+
{
28+
this.callback();
29+
}
30+
}
31+
1332
internal AbstractFlowGraph(IntPtr handle , bool owner)
1433
: base(handle , owner)
1534
{
@@ -417,42 +436,51 @@ public void RemoveRenderLayer(RenderLayer layer)
417436

418437
public FlowGraphLayoutRequest StartLayout(Action? callback)
419438
{
420-
IntPtr argCallback = IntPtr.Zero;
421-
422-
if (null != callback)
439+
NativeDelegates.BNCustomFlowGraphEvent callbackAdapter;
440+
if (null == callback)
423441
{
424-
Action<IntPtr> callbackAdapter = ( _ctx => { callback(); } );
425-
426-
// Layout completes asynchronously after BNStartFlowGraphLayout returns, so the
427-
// completion callback fires after this method has exited. PinCallback roots the
428-
// adapter for the process lifetime; otherwise the local would be GC-eligible before
429-
// the layout worker invokes it and the callback would crash.
430-
argCallback = UnsafeUtils.PinCallback(callbackAdapter);
442+
callbackAdapter = AbstractFlowGraph<T_SELF>.emptyLayoutCompletion;
431443
}
444+
else
445+
{
446+
LayoutCompletionContext callbackContext =
447+
new LayoutCompletionContext(callback);
448+
callbackAdapter = callbackContext.Invoke;
449+
}
450+
451+
// The request roots this delegate until completion or cancellation.
452+
IntPtr argCallback = Marshal.GetFunctionPointerForDelegate<
453+
NativeDelegates.BNCustomFlowGraphEvent
454+
>(callbackAdapter);
432455

433456
return new FlowGraphLayoutRequest(
434457
NativeMethods.BNStartFlowGraphLayout(
435458
this.handle ,
436459
IntPtr.Zero ,
437460
argCallback
438-
) ,
439-
true
461+
) ,
462+
true,
463+
callbackAdapter
440464
);
441465
}
442466

467+
private static void IgnoreLayoutCompletion(IntPtr context)
468+
{
469+
}
470+
443471
public bool IsNodeValid(FlowGraphNode node)
444472
{
445473
return NativeMethods.BNIsNodeValidForFlowGraph(this.handle, node.DangerousGetHandle());
446474
}
447475
}
448476

449477
// x
450-
public sealed class FlowGraph : AbstractFlowGraph<FlowGraph>
478+
public partial class FlowGraph : AbstractFlowGraph<FlowGraph>
451479
{
452480
public FlowGraph()
453-
:this( NativeMethods.BNCreateFlowGraph() , true)
481+
: base(IntPtr.Zero, true)
454482
{
455-
483+
this.InitializeCustomGraph();
456484
}
457485

458486
internal FlowGraph(IntPtr handle , bool owner)
@@ -469,7 +497,7 @@ internal FlowGraph(IntPtr handle , bool owner)
469497
return null;
470498
}
471499

472-
return new FlowGraph(handle, true);
500+
return new CoreFlowGraph(handle, true);
473501
}
474502

475503
internal static FlowGraph MustTakeHandle(IntPtr handle)
@@ -479,7 +507,7 @@ internal static FlowGraph MustTakeHandle(IntPtr handle)
479507
throw new ArgumentNullException(nameof(handle));
480508
}
481509

482-
return new FlowGraph(handle, true);
510+
return new CoreFlowGraph(handle, true);
483511
}
484512

485513
internal static FlowGraph? BorrowHandle(IntPtr handle)
@@ -489,7 +517,7 @@ internal static FlowGraph MustTakeHandle(IntPtr handle)
489517
return null;
490518
}
491519

492-
return new FlowGraph(handle, false);
520+
return new CoreFlowGraph(handle, false);
493521
}
494522

495523
internal static FlowGraph MustBorrowHandle(IntPtr handle)
@@ -499,17 +527,17 @@ internal static FlowGraph MustBorrowHandle(IntPtr handle)
499527
throw new ArgumentNullException(nameof(handle));
500528
}
501529

502-
return new FlowGraph(handle, false);
530+
return new CoreFlowGraph(handle, false);
503531
}
504532

505533
/// <summary>
506534
/// Whether the flow graph has pending updates.
507535
/// </summary>
508-
public bool HasUpdates
536+
public virtual bool HasUpdates
509537
{
510538
get
511539
{
512-
return NativeMethods.BNFlowGraphHasUpdates(this.handle);
540+
return false;
513541
}
514542
}
515543

@@ -769,11 +797,9 @@ out ulong arrayLength
769797
/// Produces an updated copy of the graph when it has pending updates, or null.
770798
/// Mirrors Python CoreFlowGraph.update().
771799
/// </summary>
772-
public FlowGraph? Update()
800+
public virtual FlowGraph? Update()
773801
{
774-
return FlowGraph.TakeHandle(
775-
NativeMethods.BNUpdateFlowGraph(this.handle)
776-
);
802+
return null;
777803
}
778804
}
779-
}
805+
}

‎Handle/BNFlowGraphCustom.cs‎

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Runtime.InteropServices;
4+
5+
namespace BinaryNinja
6+
{
7+
public partial class FlowGraph
8+
{
9+
private static readonly object externalReferenceLock = new object();
10+
11+
private static readonly List<FlowGraph> externallyReferencedGraphs =
12+
new List<FlowGraph>();
13+
14+
private NativeDelegates.BNCustomFlowGraphEvent? prepareForLayoutCallback;
15+
16+
private NativeDelegates.BNCustomFlowGraphEvent? populateNodesCallback;
17+
18+
private NativeDelegates.BNCustomFlowGraphEvent? completeLayoutCallback;
19+
20+
private NativeDelegates.BNCustomFlowGraphUpdate? updateCallback;
21+
22+
private NativeDelegates.BNCustomFlowGraphEvent? externalReferenceTakenCallback;
23+
24+
private NativeDelegates.BNCustomFlowGraphEvent? externalReferenceReleasedCallback;
25+
26+
/// <summary>Prepares this graph before the layout worker requests nodes.</summary>
27+
protected virtual void PrepareForLayout()
28+
{
29+
this.FinishPrepareForLayout();
30+
}
31+
32+
/// <summary>Populates this graph on the layout worker thread.</summary>
33+
protected virtual void PopulateNodes()
34+
{
35+
}
36+
37+
/// <summary>Runs after the core completes graph layout.</summary>
38+
protected virtual void CompleteLayout()
39+
{
40+
}
41+
42+
private void InitializeCustomGraph()
43+
{
44+
this.prepareForLayoutCallback = new NativeDelegates.BNCustomFlowGraphEvent(
45+
this.InvokePrepareForLayout
46+
);
47+
this.populateNodesCallback = new NativeDelegates.BNCustomFlowGraphEvent(
48+
this.InvokePopulateNodes
49+
);
50+
this.completeLayoutCallback = new NativeDelegates.BNCustomFlowGraphEvent(
51+
this.InvokeCompleteLayout
52+
);
53+
this.updateCallback = new NativeDelegates.BNCustomFlowGraphUpdate(
54+
this.InvokeUpdate
55+
);
56+
this.externalReferenceTakenCallback =
57+
new NativeDelegates.BNCustomFlowGraphEvent(
58+
this.InvokeExternalReferenceTaken
59+
);
60+
this.externalReferenceReleasedCallback =
61+
new NativeDelegates.BNCustomFlowGraphEvent(
62+
this.InvokeExternalReferenceReleased
63+
);
64+
65+
BNCustomFlowGraph callbacks = new BNCustomFlowGraph();
66+
callbacks.context = IntPtr.Zero;
67+
callbacks.prepareForLayout = Marshal.GetFunctionPointerForDelegate(
68+
this.prepareForLayoutCallback
69+
);
70+
callbacks.populateNodes = Marshal.GetFunctionPointerForDelegate(
71+
this.populateNodesCallback
72+
);
73+
callbacks.completeLayout = Marshal.GetFunctionPointerForDelegate(
74+
this.completeLayoutCallback
75+
);
76+
callbacks.update = Marshal.GetFunctionPointerForDelegate(
77+
this.updateCallback
78+
);
79+
callbacks.freeObject = IntPtr.Zero;
80+
callbacks.externalRefTaken = Marshal.GetFunctionPointerForDelegate(
81+
this.externalReferenceTakenCallback
82+
);
83+
callbacks.externalRefReleased = Marshal.GetFunctionPointerForDelegate(
84+
this.externalReferenceReleasedCallback
85+
);
86+
87+
IntPtr handle = NativeMethods.BNCreateCustomFlowGraph(in callbacks);
88+
if (IntPtr.Zero == handle)
89+
{
90+
throw new InvalidOperationException(
91+
"The core rejected the custom flow graph."
92+
);
93+
}
94+
95+
this.SetHandle(handle);
96+
}
97+
98+
private void InvokePrepareForLayout(IntPtr context)
99+
{
100+
try
101+
{
102+
this.PrepareForLayout();
103+
}
104+
catch (Exception exception)
105+
{
106+
Core.LogError(
107+
"Unhandled exception in FlowGraph.PrepareForLayout: {0}",
108+
exception
109+
);
110+
this.FinishPrepareForLayout();
111+
}
112+
}
113+
114+
private void InvokePopulateNodes(IntPtr context)
115+
{
116+
try
117+
{
118+
this.PopulateNodes();
119+
}
120+
catch (Exception exception)
121+
{
122+
Core.LogError(
123+
"Unhandled exception in FlowGraph.PopulateNodes: {0}",
124+
exception
125+
);
126+
}
127+
}
128+
129+
private void InvokeCompleteLayout(IntPtr context)
130+
{
131+
try
132+
{
133+
this.CompleteLayout();
134+
}
135+
catch (Exception exception)
136+
{
137+
Core.LogError(
138+
"Unhandled exception in FlowGraph.CompleteLayout: {0}",
139+
exception
140+
);
141+
}
142+
}
143+
144+
private IntPtr InvokeUpdate(IntPtr context)
145+
{
146+
try
147+
{
148+
FlowGraph? graph = this.Update();
149+
if (null == graph)
150+
{
151+
return IntPtr.Zero;
152+
}
153+
154+
return NativeMethods.BNNewFlowGraphReference(
155+
graph.DangerousGetHandle()
156+
);
157+
}
158+
catch (Exception exception)
159+
{
160+
Core.LogError(
161+
"Unhandled exception in FlowGraph.Update: {0}",
162+
exception
163+
);
164+
return IntPtr.Zero;
165+
}
166+
}
167+
168+
private void InvokeExternalReferenceTaken(IntPtr context)
169+
{
170+
lock (FlowGraph.externalReferenceLock)
171+
{
172+
FlowGraph.externallyReferencedGraphs.Add(this);
173+
}
174+
}
175+
176+
private void InvokeExternalReferenceReleased(IntPtr context)
177+
{
178+
lock (FlowGraph.externalReferenceLock)
179+
{
180+
FlowGraph.externallyReferencedGraphs.Remove(this);
181+
}
182+
}
183+
184+
private sealed class CoreFlowGraph : FlowGraph
185+
{
186+
internal CoreFlowGraph(IntPtr handle, bool owner)
187+
: base(handle, owner)
188+
{
189+
}
190+
191+
public override bool HasUpdates
192+
{
193+
get
194+
{
195+
if (!this.UpdateQueryMode)
196+
{
197+
return false;
198+
}
199+
200+
return NativeMethods.BNFlowGraphHasUpdates(this.handle);
201+
}
202+
}
203+
204+
public override FlowGraph? Update()
205+
{
206+
return FlowGraph.TakeHandle(
207+
NativeMethods.BNUpdateFlowGraph(this.handle)
208+
);
209+
}
210+
}
211+
}
212+
}

0 commit comments

Comments
 (0)