Skip to content

Commit cee78df

Browse files
tinysectinysec
authored
Add custom WebSocket providers (#139)
1 parent d7e5a09 commit cee78df

12 files changed

Lines changed: 703 additions & 269 deletions

‎Function/BNConnectWebsocketClient.cs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,13 @@ internal static extern bool BNConnectWebsocketClient(
2929
ulong headerCount ,
3030

3131
// const char** headerKeys
32-
string[] headerKeys ,
32+
IntPtr headerKeys ,
3333

3434
// const char** headerValues
35-
string[] headerValues ,
35+
IntPtr headerValues ,
3636

3737
// BNWebsocketClientOutputCallbacks* callbacks
38-
IntPtr callbacks
38+
in BNWebsocketClientOutputCallbacks callbacks
3939
);
4040
}
41-
}
41+
}

‎Function/BNGetWebsocketProviderList.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ internal static partial class NativeMethods
1919
internal static extern IntPtr BNGetWebsocketProviderList(
2020

2121
// uint64_t* count
22-
IntPtr count
22+
out UIntPtr count
2323

2424
);
2525
}
26-
}
26+
}

‎Function/BNInitWebsocketClient.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ internal static extern IntPtr BNInitWebsocketClient(
2222
IntPtr provider ,
2323

2424
// BNWebsocketClientCallbacks* callbacks
25-
IntPtr callbacks
25+
in BNWebsocketClientCallbacks callbacks
2626

2727
);
2828
}
29-
}
29+
}

‎Function/BNRegisterWebsocketProvider.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ internal static extern IntPtr BNRegisterWebsocketProvider(
2222
[MarshalAs(UnmanagedType.LPUTF8Str)] string name ,
2323

2424
// BNWebsocketProviderCallbacks* callbacks
25-
IntPtr callbacks
25+
in BNWebsocketProviderCallbacks callbacks
2626

2727
);
2828
}
29-
}
29+
}

‎Handle/BNWebsocketClient.cs‎

Lines changed: 146 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -1,163 +1,199 @@
11
using System;
22
using System.Collections.Generic;
33
using System.Runtime.InteropServices;
4-
using Microsoft.Win32.SafeHandles;
54

65
namespace BinaryNinja
76
{
8-
/// <summary>
9-
/// Represents a WebSocket client instance created by a WebSocket provider.
10-
/// The caller takes ownership of the handle returned from the provider's CreateClient method.
11-
/// </summary>
12-
public sealed class WebsocketClient : AbstractSafeHandle<WebsocketClient>
7+
/// <summary>Represents a WebSocket client created by a provider.</summary>
8+
public abstract partial class WebsocketClient :
9+
AbstractSafeHandle<WebsocketClient>
1310
{
14-
/// <summary>
15-
/// Initializes a new WebsocketClient wrapper around an existing native handle.
16-
/// </summary>
17-
/// <param name="handle">The native pointer to the BNWebsocketClient object.</param>
18-
/// <param name="owner">True if this wrapper owns the handle and should free it on dispose.</param>
19-
internal WebsocketClient(IntPtr handle, bool owner)
11+
private WebsocketConnectedDelegate? connectedHandler;
12+
private WebsocketDisconnectedDelegate? disconnectedHandler;
13+
private WebsocketErrorDelegate? errorHandler;
14+
private WebsocketDataDelegate? dataHandler;
15+
private NativeDelegates.BNWebsocketConnected? connectedCallback;
16+
private NativeDelegates.BNWebsocketDisconnected? disconnectedCallback;
17+
private NativeDelegates.BNWebsocketError? errorCallback;
18+
private NativeDelegates.BNWebsocketRead? readCallback;
19+
private bool connectionStarted;
20+
21+
private protected WebsocketClient(IntPtr handle, bool owner)
2022
: base(handle, owner)
2123
{
2224
}
2325

24-
/// <summary>
25-
/// Creates an owned reference by incrementing the native reference count.
26-
/// Returns null if the handle is zero.
27-
/// </summary>
28-
/// <param name="handle">The native BNWebsocketClient pointer.</param>
29-
/// <returns>A new owned WebsocketClient, or null if the handle is zero.</returns>
30-
internal static WebsocketClient? NewFromHandle(IntPtr handle)
26+
internal static WebsocketClient? TakeHandle(IntPtr handle)
3127
{
32-
if (handle == IntPtr.Zero)
33-
{
34-
return null;
35-
}
36-
37-
return new WebsocketClient(
38-
NativeMethods.BNNewWebsocketClientReference(handle),
39-
true
40-
);
28+
return IntPtr.Zero == handle
29+
? null
30+
: new CoreWebsocketClient(handle);
4131
}
4232

43-
/// <summary>
44-
/// Creates an owned reference by incrementing the native reference count.
45-
/// Throws if the handle is zero.
46-
/// </summary>
47-
/// <param name="handle">The native BNWebsocketClient pointer.</param>
48-
/// <returns>A new owned WebsocketClient.</returns>
49-
internal static WebsocketClient MustNewFromHandle(IntPtr handle)
33+
/// <summary>Starts an asynchronous connection.</summary>
34+
public bool Connect(
35+
string url,
36+
IReadOnlyDictionary<string, string>? headers = null,
37+
WebsocketConnectedDelegate? onConnected = null,
38+
WebsocketDisconnectedDelegate? onDisconnected = null,
39+
WebsocketErrorDelegate? onError = null,
40+
WebsocketDataDelegate? onData = null
41+
)
5042
{
51-
if (handle == IntPtr.Zero)
43+
if (this.connectionStarted)
5244
{
53-
throw new ArgumentNullException(nameof(handle));
45+
throw new InvalidOperationException(
46+
"A WebSocket client cannot connect more than once."
47+
);
5448
}
5549

56-
return new WebsocketClient(
57-
NativeMethods.BNNewWebsocketClientReference(handle),
58-
true
50+
this.connectionStarted = true;
51+
this.connectedHandler = onConnected;
52+
this.disconnectedHandler = onDisconnected;
53+
this.errorHandler = onError;
54+
this.dataHandler = onData;
55+
this.connectedCallback = new NativeDelegates.BNWebsocketConnected(
56+
this.InvokeConnected
57+
);
58+
this.disconnectedCallback =
59+
new NativeDelegates.BNWebsocketDisconnected(
60+
this.InvokeDisconnected
61+
);
62+
this.errorCallback = new NativeDelegates.BNWebsocketError(
63+
this.InvokeError
64+
);
65+
this.readCallback = new NativeDelegates.BNWebsocketRead(
66+
this.InvokeRead
67+
);
68+
69+
string[] keys;
70+
string[] values;
71+
CopyHeaders(headers, out keys, out values);
72+
BNWebsocketClientOutputCallbacks callbacks =
73+
new BNWebsocketClientOutputCallbacks();
74+
callbacks.context = IntPtr.Zero;
75+
callbacks.connectedCallback = Pointer(this.connectedCallback);
76+
callbacks.disconnectedCallback = Pointer(
77+
this.disconnectedCallback
78+
);
79+
callbacks.errorCallback = Pointer(this.errorCallback);
80+
callbacks.readCallback = Pointer(this.readCallback);
81+
82+
using ScopedAllocator allocator = new ScopedAllocator();
83+
return NativeMethods.BNConnectWebsocketClient(
84+
this.handle,
85+
url,
86+
(ulong)keys.Length,
87+
allocator.AllocUtf8StringArray(keys),
88+
allocator.AllocUtf8StringArray(values),
89+
in callbacks
5990
);
6091
}
6192

62-
/// <summary>
63-
/// Takes ownership of an existing handle without incrementing the reference count.
64-
/// Returns null if the handle is zero.
65-
/// </summary>
66-
/// <param name="handle">The native BNWebsocketClient pointer.</param>
67-
/// <returns>A new owned WebsocketClient, or null if the handle is zero.</returns>
68-
internal static WebsocketClient? TakeHandle(IntPtr handle)
93+
/// <summary>Writes data to this client.</summary>
94+
public bool Write(byte[] data)
6995
{
70-
if (handle == IntPtr.Zero)
96+
if (null == data)
7197
{
72-
return null;
98+
throw new ArgumentNullException(nameof(data));
7399
}
74100

75-
return new WebsocketClient(handle, true);
76-
}
77-
78-
/// <summary>
79-
/// Takes ownership of an existing handle without incrementing the reference count.
80-
/// Throws if the handle is zero.
81-
/// </summary>
82-
/// <param name="handle">The native BNWebsocketClient pointer.</param>
83-
/// <returns>A new owned WebsocketClient.</returns>
84-
internal static WebsocketClient MustTakeHandle(IntPtr handle)
85-
{
86-
if (handle == IntPtr.Zero)
101+
GCHandle pin = GCHandle.Alloc(data, GCHandleType.Pinned);
102+
try
87103
{
88-
throw new ArgumentNullException(nameof(handle));
104+
return 0 != NativeMethods.BNWriteWebsocketClientData(
105+
this.handle,
106+
0 == data.Length ? IntPtr.Zero : pin.AddrOfPinnedObject(),
107+
(ulong)data.Length
108+
);
89109
}
90-
91-
return new WebsocketClient(handle, true);
92-
}
93-
94-
/// <summary>
95-
/// Borrows a native handle without taking ownership. Returns null if the handle is zero.
96-
/// </summary>
97-
/// <param name="handle">The native BNWebsocketClient pointer.</param>
98-
/// <returns>A new WebsocketClient that will not free the handle on dispose.</returns>
99-
internal static WebsocketClient? BorrowHandle(IntPtr handle)
100-
{
101-
if (handle == IntPtr.Zero)
110+
finally
102111
{
103-
return null;
112+
pin.Free();
104113
}
105-
106-
return new WebsocketClient(handle, false);
107114
}
108115

109-
/// <summary>
110-
/// Borrows a native handle without taking ownership. Throws if the handle is zero.
111-
/// </summary>
112-
/// <param name="handle">The native BNWebsocketClient pointer.</param>
113-
/// <returns>A new WebsocketClient that will not free the handle on dispose.</returns>
114-
internal static WebsocketClient MustBorrowHandle(IntPtr handle)
116+
/// <summary>Writes a raw buffer and returns the core result.</summary>
117+
public ulong WriteData(IntPtr data, ulong length)
115118
{
116-
if (handle == IntPtr.Zero)
117-
{
118-
throw new ArgumentNullException(nameof(handle));
119-
}
119+
return NativeMethods.BNWriteWebsocketClientData(
120+
this.handle,
121+
data,
122+
length
123+
);
124+
}
120125

121-
return new WebsocketClient(handle, false);
126+
/// <summary>Disconnects this client.</summary>
127+
public bool Disconnect()
128+
{
129+
return NativeMethods.BNDisconnectWebsocketClient(this.handle);
122130
}
123131

124-
/// <summary>
125-
/// Releases the native BNWebsocketClient handle when this instance is disposed or finalized.
126-
/// </summary>
127-
/// <returns>True if the handle was successfully released.</returns>
128132
protected override bool ReleaseHandle()
129133
{
130134
if (!this.IsInvalid)
131135
{
132-
// Free the native client handle and mark it invalid to prevent double-free.
133136
NativeMethods.BNFreeWebsocketClient(this.handle);
134137
this.SetHandleAsInvalid();
135138
}
136139

137140
return true;
138141
}
139142

140-
/// <summary>
141-
/// Disconnects this WebSocket client, closing the underlying connection.
142-
/// </summary>
143-
/// <returns>True if the disconnect succeeded; false on failure.</returns>
144-
public bool Disconnect()
143+
private static void CopyHeaders(
144+
IReadOnlyDictionary<string, string>? headers,
145+
out string[] keys,
146+
out string[] values
147+
)
145148
{
146-
// Delegate to the native disconnect API.
147-
return NativeMethods.BNDisconnectWebsocketClient(this.handle);
149+
if (null == headers || 0 == headers.Count)
150+
{
151+
keys = Array.Empty<string>();
152+
values = Array.Empty<string>();
153+
return;
154+
}
155+
156+
keys = new string[headers.Count];
157+
values = new string[headers.Count];
158+
int index = 0;
159+
foreach (KeyValuePair<string, string> header in headers)
160+
{
161+
keys[index] = header.Key;
162+
values[index] = header.Value;
163+
index++;
164+
}
148165
}
149166

150-
/// <summary>
151-
/// Writes raw byte data to the WebSocket connection.
152-
/// The caller is responsible for ensuring the pointer is valid for the given length.
153-
/// </summary>
154-
/// <param name="data">Pointer to the byte buffer to send.</param>
155-
/// <param name="length">Number of bytes to send from the buffer.</param>
156-
/// <returns>The number of bytes actually written; zero on failure.</returns>
157-
public ulong WriteData(IntPtr data, ulong length)
167+
private static IntPtr Pointer<TDelegate>(TDelegate callback)
168+
where TDelegate : Delegate
169+
{
170+
return Marshal.GetFunctionPointerForDelegate(callback);
171+
}
172+
173+
private sealed class CoreWebsocketClient : WebsocketClient
158174
{
159-
// Forward the raw data pointer and byte count to the native write API.
160-
return NativeMethods.BNWriteWebsocketClientData(this.handle, data, length);
175+
internal CoreWebsocketClient(IntPtr handle)
176+
: base(handle, true)
177+
{
178+
}
179+
180+
protected override bool PerformConnect(
181+
string host,
182+
IReadOnlyDictionary<string, string> headers
183+
)
184+
{
185+
return false;
186+
}
187+
188+
protected override bool PerformWrite(byte[] data)
189+
{
190+
return false;
191+
}
192+
193+
protected override bool PerformDisconnect()
194+
{
195+
return false;
196+
}
161197
}
162198
}
163199
}

0 commit comments

Comments
 (0)