binaryninja/websocket/
client.rs

1use crate::rc::{Ref, RefCountable};
2use crate::string::{BnString, IntoCStr};
3use binaryninjacore_sys::*;
4use std::ffi::{c_char, c_void, CStr};
5use std::ptr::NonNull;
6
7pub trait WebsocketClientCallback: Sync + Send {
8    fn connected(&mut self) -> bool;
9
10    fn disconnected(&mut self);
11
12    fn error(&mut self, msg: &str);
13
14    fn read(&mut self, data: &[u8]) -> bool;
15}
16
17pub trait WebsocketClient: Sync + Send {
18    /// Called to construct this client object with the given core object.
19    fn from_core(core: Ref<CoreWebsocketClient>) -> Self;
20
21    fn connect<I>(&self, host: &str, headers: I) -> bool
22    where
23        I: IntoIterator<Item = (String, String)>;
24
25    fn write(&self, data: &[u8]) -> bool;
26
27    fn disconnect(&self) -> bool;
28}
29
30/// Implements a websocket client.
31#[repr(transparent)]
32pub struct CoreWebsocketClient {
33    pub(crate) handle: NonNull<BNWebsocketClient>,
34}
35
36impl CoreWebsocketClient {
37    pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNWebsocketClient>) -> Ref<Self> {
38        Ref::new(Self { handle })
39    }
40
41    #[allow(clippy::mut_from_ref)]
42    pub(crate) unsafe fn as_raw(&self) -> &mut BNWebsocketClient {
43        &mut *self.handle.as_ptr()
44    }
45
46    /// Initializes the web socket connection.
47    ///
48    /// Connect to a given url, asynchronously. The connection will be run in a
49    /// separate thread managed by the websocket provider.
50    ///
51    /// Callbacks will be called **on the thread of the connection**, so be sure
52    /// to ExecuteOnMainThread any long-running or gui operations in the callbacks.
53    ///
54    /// If the connection succeeds, [WebsocketClientCallback::connected] will be called. On normal
55    /// termination, [WebsocketClientCallback::disconnected] will be called.
56    ///
57    /// If the connection succeeds, but later fails, [WebsocketClientCallback::disconnected] will not
58    /// be called, and [WebsocketClientCallback::error] will be called instead.
59    ///
60    /// If the connection fails, neither [WebsocketClientCallback::connected] nor
61    /// [WebsocketClientCallback::disconnected] will be called, and [WebsocketClientCallback::error]
62    /// will be called instead.
63    ///
64    /// If [WebsocketClientCallback::connected] or [WebsocketClientCallback::read] return false, the
65    /// connection will be aborted.
66    ///
67    /// * `host` - Full url with scheme, domain, optionally port, and path
68    /// * `headers` - HTTP header keys and values
69    /// * `callback` - Callbacks for various websocket events
70    pub fn initialize_connection<I, C>(&self, host: &str, headers: I, callbacks: &mut C) -> bool
71    where
72        I: IntoIterator<Item = (String, String)>,
73        C: WebsocketClientCallback,
74    {
75        let url = host.to_cstr();
76        let (header_keys, header_values): (Vec<_>, Vec<_>) = headers
77            .into_iter()
78            .map(|(k, v)| (k.to_cstr(), v.to_cstr()))
79            .unzip();
80        let header_keys: Vec<*const c_char> = header_keys.iter().map(|k| k.as_ptr()).collect();
81        let header_values: Vec<*const c_char> = header_values.iter().map(|v| v.as_ptr()).collect();
82        // SAFETY: This context will only be live for the duration of BNConnectWebsocketClient
83        // SAFETY: Any subsequent call to BNConnectWebsocketClient will write over the context.
84        let mut output_callbacks = BNWebsocketClientOutputCallbacks {
85            context: callbacks as *mut C as *mut c_void,
86            connectedCallback: Some(cb_connected::<C>),
87            disconnectedCallback: Some(cb_disconnected::<C>),
88            errorCallback: Some(cb_error::<C>),
89            readCallback: Some(cb_read::<C>),
90        };
91        unsafe {
92            BNConnectWebsocketClient(
93                self.handle.as_ptr(),
94                url.as_ptr(),
95                header_keys.len().try_into().unwrap(),
96                header_keys.as_ptr(),
97                header_values.as_ptr(),
98                &mut output_callbacks,
99            )
100        }
101    }
102
103    /// Call the connect callback function, forward the callback returned value
104    pub fn notify_connected(&self) -> bool {
105        unsafe { BNNotifyWebsocketClientConnect(self.handle.as_ptr()) }
106    }
107
108    /// Notify the callback function of a disconnect,
109    ///
110    /// NOTE: This does not actually disconnect, use the [Self::disconnect] function for that.
111    pub fn notify_disconnected(&self) {
112        unsafe { BNNotifyWebsocketClientDisconnect(self.handle.as_ptr()) }
113    }
114
115    /// Call the error callback function
116    pub fn notify_error(&self, msg: &str) {
117        let error = msg.to_cstr();
118        unsafe { BNNotifyWebsocketClientError(self.handle.as_ptr(), error.as_ptr()) }
119    }
120
121    /// Call the read callback function, forward the callback returned value
122    pub fn notify_read(&self, data: &[u8]) -> bool {
123        unsafe {
124            BNNotifyWebsocketClientReadData(
125                self.handle.as_ptr(),
126                data.as_ptr() as *mut _,
127                data.len().try_into().unwrap(),
128            )
129        }
130    }
131
132    pub fn write(&self, data: &[u8]) -> bool {
133        let len = u64::try_from(data.len()).unwrap();
134        unsafe { BNWriteWebsocketClientData(self.as_raw(), data.as_ptr(), len) != 0 }
135    }
136
137    pub fn disconnect(&self) -> bool {
138        unsafe { BNDisconnectWebsocketClient(self.as_raw()) }
139    }
140}
141
142unsafe impl Sync for CoreWebsocketClient {}
143unsafe impl Send for CoreWebsocketClient {}
144
145impl ToOwned for CoreWebsocketClient {
146    type Owned = Ref<Self>;
147
148    fn to_owned(&self) -> Self::Owned {
149        unsafe { RefCountable::inc_ref(self) }
150    }
151}
152
153unsafe impl RefCountable for CoreWebsocketClient {
154    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
155        let result = BNNewWebsocketClientReference(handle.as_raw());
156        unsafe { Self::ref_from_raw(NonNull::new(result).unwrap()) }
157    }
158
159    unsafe fn dec_ref(handle: &Self) {
160        BNFreeWebsocketClient(handle.as_raw())
161    }
162}
163
164pub(crate) unsafe extern "C" fn cb_destroy_client<W: WebsocketClient>(ctxt: *mut c_void) {
165    let _ = Box::from_raw(ctxt as *mut W);
166}
167
168pub(crate) unsafe extern "C" fn cb_connect<W: WebsocketClient>(
169    ctxt: *mut c_void,
170    host: *const c_char,
171    header_count: u64,
172    header_keys: *const *const c_char,
173    header_values: *const *const c_char,
174) -> bool {
175    let ctxt: &mut W = &mut *(ctxt as *mut W);
176    let host = CStr::from_ptr(host);
177    // SAFETY BnString and *mut c_char are transparent
178    let header_count = usize::try_from(header_count).unwrap();
179    let header_keys = core::slice::from_raw_parts(header_keys as *const BnString, header_count);
180    let header_values = core::slice::from_raw_parts(header_values as *const BnString, header_count);
181    let header_keys_str = header_keys.iter().map(|s| s.to_string_lossy().to_string());
182    let header_values_str = header_values
183        .iter()
184        .map(|s| s.to_string_lossy().to_string());
185    let header = header_keys_str.zip(header_values_str);
186    ctxt.connect(&host.to_string_lossy(), header)
187}
188
189pub(crate) unsafe extern "C" fn cb_write<W: WebsocketClient>(
190    data: *const u8,
191    len: u64,
192    ctxt: *mut c_void,
193) -> bool {
194    let ctxt: &mut W = &mut *(ctxt as *mut W);
195    let len = usize::try_from(len).unwrap();
196    let data = core::slice::from_raw_parts(data, len);
197    ctxt.write(data)
198}
199
200pub(crate) unsafe extern "C" fn cb_disconnect<W: WebsocketClient>(ctxt: *mut c_void) -> bool {
201    let ctxt: &mut W = &mut *(ctxt as *mut W);
202    ctxt.disconnect()
203}
204
205unsafe extern "C" fn cb_connected<W: WebsocketClientCallback>(ctxt: *mut c_void) -> bool {
206    let ctxt: &mut W = &mut *(ctxt as *mut W);
207    ctxt.connected()
208}
209
210unsafe extern "C" fn cb_disconnected<W: WebsocketClientCallback>(ctxt: *mut c_void) {
211    let ctxt: &mut W = &mut *(ctxt as *mut W);
212    ctxt.disconnected()
213}
214
215unsafe extern "C" fn cb_error<W: WebsocketClientCallback>(msg: *const c_char, ctxt: *mut c_void) {
216    let ctxt: &mut W = &mut *(ctxt as *mut W);
217    let msg = CStr::from_ptr(msg);
218    ctxt.error(&msg.to_string_lossy())
219}
220
221unsafe extern "C" fn cb_read<W: WebsocketClientCallback>(
222    data: *mut u8,
223    len: u64,
224    ctxt: *mut c_void,
225) -> bool {
226    let ctxt: &mut W = &mut *(ctxt as *mut W);
227    let len = usize::try_from(len).unwrap();
228    let data = core::slice::from_raw_parts_mut(data, len);
229    ctxt.read(data)
230}
binaryninja/websocket/client.rs

binaryninja/websocket/
client.rs
