use crate::rc::{Ref, RefCountable}; use crate::string::{BnString, IntoCStr}; use binaryninjacore_sys::*; use std::ffi::{c_char, c_void, CStr}; use std::marker::PhantomData; use std::ops::Deref; use std::ptr::NonNull; pub trait WebsocketClientCallback: Sync + Send { /// Receive a notification that the websocket connection has been connected successfully. /// /// Return `false` if you would like to terminate the connection early. fn connected(&self) -> bool; /// Receive a notification that the websocket connection has been terminated. /// /// For implementations, you must call this at the end of the websocket connection lifecycle /// even if you notify the client of an error. fn disconnected(&self); /// Receive an error from the websocket connection. /// /// For implementations, you typically write the data to an interior mutable buffer on the instance. fn error(&self, msg: &str); /// Receive data from the websocket connection. /// /// For implementations, you typically write the data to an interior mutable buffer on the instance. fn read(&self, data: &[u8]) -> bool; } pub trait WebsocketClient: Sync + Send { /// Called to construct this client object with the given core object. fn from_core(core: Ref) -> Self; fn connect(&self, host: &str, headers: I) -> bool where I: IntoIterator; fn write(&self, data: &[u8]) -> bool; fn disconnect(&self) -> bool; } /// Represents a live websocket connection. /// /// This manages the lifetime of the callback, ensuring it outlives the connection. pub struct ActiveConnection<'a, C: WebsocketClientCallback> { pub client: Ref, _callback: PhantomData<&'a mut C>, } impl<'a, C: WebsocketClientCallback> Deref for ActiveConnection<'a, C> { type Target = CoreWebsocketClient; fn deref(&self) -> &Self::Target { &self.client } } impl<'a, C: WebsocketClientCallback> Drop for ActiveConnection<'a, C> { fn drop(&mut self) { self.client.disconnect(); } } /// Implements a websocket client. /// /// To connect, use [`Ref::connect`] which will return an [`ActiveConnection`] /// which manages the lifecycle of the websocket connection. #[repr(transparent)] pub struct CoreWebsocketClient { pub(crate) handle: NonNull, } impl CoreWebsocketClient { pub(crate) unsafe fn ref_from_raw(handle: NonNull) -> Ref { Ref::new(Self { handle }) } #[allow(clippy::mut_from_ref)] pub(crate) unsafe fn as_raw(&self) -> &mut BNWebsocketClient { &mut *self.handle.as_ptr() } /// Call the connect callback function, forward the callback returned value pub fn notify_connected(&self) -> bool { unsafe { BNNotifyWebsocketClientConnect(self.handle.as_ptr()) } } /// Notify the callback function of a disconnect. This must be called at the end of an active /// websocket connection lifecycle, to free resources. /// /// NOTE: This does not actually disconnect, use the [Self::disconnect] function for that. pub fn notify_disconnected(&self) { unsafe { BNNotifyWebsocketClientDisconnect(self.handle.as_ptr()) } } /// Call the error callback function, this is not a terminating request you must use /// [`CoreWebsocketClient::notify_disconnected`] to terminate the connection. pub fn notify_error(&self, msg: &str) { let error = msg.to_cstr(); unsafe { BNNotifyWebsocketClientError(self.handle.as_ptr(), error.as_ptr()) } } /// Call the read callback function, forward the callback returned value. pub fn notify_read(&self, data: &[u8]) -> bool { unsafe { BNNotifyWebsocketClientReadData( self.handle.as_ptr(), data.as_ptr() as *mut _, data.len().try_into().unwrap(), ) } } pub fn write(&self, data: &[u8]) -> bool { let len = u64::try_from(data.len()).unwrap(); unsafe { BNWriteWebsocketClientData(self.as_raw(), data.as_ptr(), len) != 0 } } pub fn disconnect(&self) -> bool { unsafe { BNDisconnectWebsocketClient(self.as_raw()) } } } impl Ref { /// Initializes the web socket connection, returning the [`ActiveConnection`], once dropped the /// connection will be disconnected. /// /// Connect to a given url, asynchronously. The connection will be run in a /// separate thread managed by the websocket provider. /// /// Callbacks will be called **on the thread of the connection**, so be sure /// to ExecuteOnMainThread any long-running or gui operations in the callbacks. /// /// If the connection succeeds, [WebsocketClientCallback::connected] will be called. On normal /// termination, [WebsocketClientCallback::disconnected] will be called. /// /// If the connection succeeds but later fails, [`WebsocketClientCallback::error`] will be called /// and shortly thereafter [`WebsocketClientCallback::disconnected`] will be called. /// /// If [`WebsocketClientCallback::connected`] or [`WebsocketClientCallback::read`] return false, the /// connection will be aborted. /// /// * `host` - Full url with scheme, domain, optionally port, and path /// * `headers` - HTTP header keys and values /// * `callback` - Callbacks for various websocket events #[must_use] pub fn connect<'a, I, C>( self, host: &str, headers: I, callbacks: &'a C, ) -> Option> where I: IntoIterator, C: WebsocketClientCallback, { let url = host.to_cstr(); let (header_keys, header_values): (Vec<_>, Vec<_>) = headers .into_iter() .map(|(k, v)| (k.to_cstr(), v.to_cstr())) .unzip(); let header_keys: Vec<*const c_char> = header_keys.iter().map(|k| k.as_ptr()).collect(); let header_values: Vec<*const c_char> = header_values.iter().map(|v| v.as_ptr()).collect(); // SAFETY: This context will live for as long as the `ActiveConnection` is alive. // SAFETY: Any subsequent call to BNConnectWebsocketClient will write over the context. let mut output_callbacks = BNWebsocketClientOutputCallbacks { context: callbacks as *const C as *mut C as *mut c_void, connectedCallback: Some(cb_connected::), disconnectedCallback: Some(cb_disconnected::), errorCallback: Some(cb_error::), readCallback: Some(cb_read::), }; let success = unsafe { BNConnectWebsocketClient( self.handle.as_ptr(), url.as_ptr(), header_keys.len().try_into().unwrap(), header_keys.as_ptr(), header_values.as_ptr(), &mut output_callbacks, ) }; if success { Some(ActiveConnection { client: self, _callback: PhantomData, }) } else { None } } } unsafe impl Sync for CoreWebsocketClient {} unsafe impl Send for CoreWebsocketClient {} impl ToOwned for CoreWebsocketClient { type Owned = Ref; fn to_owned(&self) -> Self::Owned { unsafe { RefCountable::inc_ref(self) } } } unsafe impl RefCountable for CoreWebsocketClient { unsafe fn inc_ref(handle: &Self) -> Ref { let result = BNNewWebsocketClientReference(handle.as_raw()); unsafe { Self::ref_from_raw(NonNull::new(result).unwrap()) } } unsafe fn dec_ref(handle: &Self) { BNFreeWebsocketClient(handle.as_raw()) } } pub(crate) unsafe extern "C" fn cb_destroy_client(ctxt: *mut c_void) { let _ = Box::from_raw(ctxt as *mut W); } pub(crate) unsafe extern "C" fn cb_connect( ctxt: *mut c_void, host: *const c_char, header_count: u64, header_keys: *const *const c_char, header_values: *const *const c_char, ) -> bool { let ctxt: &mut W = &mut *(ctxt as *mut W); let host = CStr::from_ptr(host); // SAFETY BnString and *mut c_char are transparent let header_count = usize::try_from(header_count).unwrap(); let header_keys = core::slice::from_raw_parts(header_keys as *const BnString, header_count); let header_values = core::slice::from_raw_parts(header_values as *const BnString, header_count); let header_keys_str = header_keys.iter().map(|s| s.to_string_lossy().to_string()); let header_values_str = header_values .iter() .map(|s| s.to_string_lossy().to_string()); let header = header_keys_str.zip(header_values_str); ctxt.connect(&host.to_string_lossy(), header) } pub(crate) unsafe extern "C" fn cb_write( data: *const u8, len: u64, ctxt: *mut c_void, ) -> bool { let ctxt: &mut W = &mut *(ctxt as *mut W); let len = usize::try_from(len).unwrap(); let data = core::slice::from_raw_parts(data, len); ctxt.write(data) } pub(crate) unsafe extern "C" fn cb_disconnect(ctxt: *mut c_void) -> bool { let ctxt: &mut W = &mut *(ctxt as *mut W); ctxt.disconnect() } unsafe extern "C" fn cb_connected(ctxt: *mut c_void) -> bool { let ctxt: &mut W = &mut *(ctxt as *mut W); ctxt.connected() } unsafe extern "C" fn cb_disconnected(ctxt: *mut c_void) { let ctxt: &mut W = &mut *(ctxt as *mut W); ctxt.disconnected() } unsafe extern "C" fn cb_error(msg: *const c_char, ctxt: *mut c_void) { let ctxt: &mut W = &mut *(ctxt as *mut W); let msg = CStr::from_ptr(msg); ctxt.error(&msg.to_string_lossy()) } unsafe extern "C" fn cb_read( data: *mut u8, len: u64, ctxt: *mut c_void, ) -> bool { let ctxt: &mut W = &mut *(ctxt as *mut W); let len = usize::try_from(len).unwrap(); let data = core::slice::from_raw_parts_mut(data, len); ctxt.read(data) }