binaryninja/websocket/
provider.rs

1use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Ref};
2use crate::string::{BnString, IntoCStr};
3use crate::websocket::client;
4use crate::websocket::client::{CoreWebsocketClient, WebsocketClient};
5use binaryninjacore_sys::*;
6use std::ffi::c_void;
7use std::mem::MaybeUninit;
8use std::ptr::NonNull;
9
10pub fn register_websocket_provider<W>(name: &str) -> &'static mut W
11where
12    W: WebsocketProvider,
13{
14    let name = name.to_cstr();
15    let provider_uninit = MaybeUninit::uninit();
16    // SAFETY: Websocket provider is never freed
17    let leaked_provider = Box::leak(Box::new(provider_uninit));
18    let result = unsafe {
19        BNRegisterWebsocketProvider(
20            name.as_ptr(),
21            &mut BNWebsocketProviderCallbacks {
22                context: leaked_provider as *mut _ as *mut c_void,
23                createClient: Some(cb_create_client::<W>),
24            },
25        )
26    };
27
28    let provider_core = unsafe { CoreWebsocketProvider::from_raw(NonNull::new(result).unwrap()) };
29    // We now have the core provider so we can actually construct the object.
30    leaked_provider.write(W::from_core(provider_core));
31    unsafe { leaked_provider.assume_init_mut() }
32}
33
34pub trait WebsocketProvider: Sync + Send + Sized {
35    type Client: WebsocketClient;
36
37    fn handle(&self) -> CoreWebsocketProvider;
38
39    /// Called to construct this provider object with the given core object.
40    fn from_core(core: CoreWebsocketProvider) -> Self;
41
42    /// Create a new instance of the websocket client.
43    fn create_client(&self) -> Result<Ref<CoreWebsocketClient>, ()> {
44        let client_uninit = MaybeUninit::uninit();
45        // SAFETY: Websocket client is freed by cb_destroy_client
46        let leaked_client = Box::leak(Box::new(client_uninit));
47        let mut callbacks = BNWebsocketClientCallbacks {
48            context: leaked_client as *mut _ as *mut c_void,
49            connect: Some(client::cb_connect::<Self::Client>),
50            destroyClient: Some(client::cb_destroy_client::<Self::Client>),
51            disconnect: Some(client::cb_disconnect::<Self::Client>),
52            write: Some(client::cb_write::<Self::Client>),
53        };
54        let client_ptr =
55            unsafe { BNInitWebsocketClient(self.handle().handle.as_ptr(), &mut callbacks) };
56        // TODO: If possible pass a sensible error back...
57        let client_ptr = NonNull::new(client_ptr).ok_or(())?;
58        let client_ref = unsafe { CoreWebsocketClient::ref_from_raw(client_ptr) };
59        // We now have the core client so we can actually construct the object.
60        leaked_client.write(Self::Client::from_core(client_ref.clone()));
61        Ok(client_ref)
62    }
63}
64
65#[derive(Clone, Copy, Hash, PartialEq, Eq)]
66#[repr(transparent)]
67pub struct CoreWebsocketProvider {
68    handle: NonNull<BNWebsocketProvider>,
69}
70
71impl CoreWebsocketProvider {
72    pub(crate) unsafe fn from_raw(handle: NonNull<BNWebsocketProvider>) -> Self {
73        Self { handle }
74    }
75
76    pub fn all() -> Array<Self> {
77        let mut count = 0;
78        let result = unsafe { BNGetWebsocketProviderList(&mut count) };
79        assert!(!result.is_null());
80        unsafe { Array::new(result, count, ()) }
81    }
82
83    pub fn by_name(name: &str) -> Option<CoreWebsocketProvider> {
84        let name = name.to_cstr();
85        let result = unsafe { BNGetWebsocketProviderByName(name.as_ptr()) };
86        NonNull::new(result).map(|h| unsafe { Self::from_raw(h) })
87    }
88
89    pub fn name(&self) -> String {
90        let result = unsafe { BNGetWebsocketProviderName(self.handle.as_ptr()) };
91        assert!(!result.is_null());
92        unsafe { BnString::into_string(result) }
93    }
94}
95
96unsafe impl Sync for CoreWebsocketProvider {}
97unsafe impl Send for CoreWebsocketProvider {}
98
99impl CoreArrayProvider for CoreWebsocketProvider {
100    type Raw = *mut BNWebsocketProvider;
101    type Context = ();
102    type Wrapped<'a> = Self;
103}
104
105unsafe impl CoreArrayProviderInner for CoreWebsocketProvider {
106    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
107        BNFreeWebsocketProviderList(raw)
108    }
109
110    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
111        let handle = NonNull::new(*raw).unwrap();
112        Self::from_raw(handle)
113    }
114}
115
116unsafe extern "C" fn cb_create_client<W: WebsocketProvider>(
117    ctxt: *mut c_void,
118) -> *mut BNWebsocketClient {
119    let ctxt: &mut W = &mut *(ctxt as *mut W);
120    match ctxt.create_client() {
121        Ok(owned_client) => {
122            // SAFETY: The caller is assumed to have picked up this ref.
123            Ref::into_raw(owned_client).handle.as_ptr()
124        }
125        Err(_) => std::ptr::null_mut(),
126    }
127}