binaryninja/websocket/
provider.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Ref};
use crate::string::{BnStrCompatible, BnString};
use crate::websocket::client;
use crate::websocket::client::{CoreWebsocketClient, WebsocketClient};
use binaryninjacore_sys::*;
use std::ffi::{c_char, c_void};
use std::mem::MaybeUninit;
use std::ptr::NonNull;

pub fn register_websocket_provider<W>(name: &str) -> &'static mut W
where
    W: WebsocketProvider,
{
    let name = name.into_bytes_with_nul();
    let provider_uninit = MaybeUninit::uninit();
    // SAFETY: Websocket provider is never freed
    let leaked_provider = Box::leak(Box::new(provider_uninit));
    let result = unsafe {
        BNRegisterWebsocketProvider(
            name.as_ptr() as *const c_char,
            &mut BNWebsocketProviderCallbacks {
                context: leaked_provider as *mut _ as *mut c_void,
                createClient: Some(cb_create_client::<W>),
            },
        )
    };

    let provider_core = unsafe { CoreWebsocketProvider::from_raw(NonNull::new(result).unwrap()) };
    // We now have the core provider so we can actually construct the object.
    leaked_provider.write(W::from_core(provider_core));
    unsafe { leaked_provider.assume_init_mut() }
}

pub trait WebsocketProvider: Sync + Send + Sized {
    type Client: WebsocketClient;

    fn handle(&self) -> CoreWebsocketProvider;

    /// Called to construct this provider object with the given core object.
    fn from_core(core: CoreWebsocketProvider) -> Self;

    /// Create a new instance of the websocket client.
    fn create_client(&self) -> Result<Ref<CoreWebsocketClient>, ()> {
        let client_uninit = MaybeUninit::uninit();
        // SAFETY: Websocket client is freed by cb_destroy_client
        let leaked_client = Box::leak(Box::new(client_uninit));
        let mut callbacks = BNWebsocketClientCallbacks {
            context: leaked_client as *mut _ as *mut c_void,
            connect: Some(client::cb_connect::<Self::Client>),
            destroyClient: Some(client::cb_destroy_client::<Self::Client>),
            disconnect: Some(client::cb_disconnect::<Self::Client>),
            write: Some(client::cb_write::<Self::Client>),
        };
        let client_ptr =
            unsafe { BNInitWebsocketClient(self.handle().handle.as_ptr(), &mut callbacks) };
        // TODO: If possible pass a sensible error back...
        let client_ptr = NonNull::new(client_ptr).ok_or(())?;
        let client_ref = unsafe { CoreWebsocketClient::ref_from_raw(client_ptr) };
        // We now have the core client so we can actually construct the object.
        leaked_client.write(Self::Client::from_core(client_ref.clone()));
        Ok(client_ref)
    }
}

#[derive(Clone, Copy, Hash, PartialEq, Eq)]
#[repr(transparent)]
pub struct CoreWebsocketProvider {
    handle: NonNull<BNWebsocketProvider>,
}

impl CoreWebsocketProvider {
    pub(crate) unsafe fn from_raw(handle: NonNull<BNWebsocketProvider>) -> Self {
        Self { handle }
    }

    pub fn all() -> Array<Self> {
        let mut count = 0;
        let result = unsafe { BNGetWebsocketProviderList(&mut count) };
        assert!(!result.is_null());
        unsafe { Array::new(result, count, ()) }
    }

    pub fn by_name<S: BnStrCompatible>(name: S) -> Option<CoreWebsocketProvider> {
        let name = name.into_bytes_with_nul();
        let result =
            unsafe { BNGetWebsocketProviderByName(name.as_ref().as_ptr() as *const c_char) };
        NonNull::new(result).map(|h| unsafe { Self::from_raw(h) })
    }

    pub fn name(&self) -> BnString {
        let result = unsafe { BNGetWebsocketProviderName(self.handle.as_ptr()) };
        assert!(!result.is_null());
        unsafe { BnString::from_raw(result) }
    }
}

unsafe impl Sync for CoreWebsocketProvider {}
unsafe impl Send for CoreWebsocketProvider {}

impl CoreArrayProvider for CoreWebsocketProvider {
    type Raw = *mut BNWebsocketProvider;
    type Context = ();
    type Wrapped<'a> = Self;
}

unsafe impl CoreArrayProviderInner for CoreWebsocketProvider {
    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
        BNFreeWebsocketProviderList(raw)
    }

    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
        let handle = NonNull::new(*raw).unwrap();
        Self::from_raw(handle)
    }
}

unsafe extern "C" fn cb_create_client<W: WebsocketProvider>(
    ctxt: *mut c_void,
) -> *mut BNWebsocketClient {
    let ctxt: &mut W = &mut *(ctxt as *mut W);
    match ctxt.create_client() {
        Ok(owned_client) => {
            // SAFETY: The caller is assumed to have picked up this ref.
            Ref::into_raw(owned_client).handle.as_ptr()
        }
        Err(_) => std::ptr::null_mut(),
    }
}