[Rust] Ensure proper lifetime management of `WebsocketClientCallback` objects

1 files changed, 77 insertions, 25 deletions

diff --git a/rust/tests/websocket.rs b/rust/tests/websocket.rs
index 2c25a810..19aacf4c 100644
--- a/rust/tests/websocket.rs
+++ b/rust/tests/websocket.rs
@@ -4,6 +4,8 @@ use binaryninja::websocket::{
     register_websocket_provider, CoreWebsocketClient, CoreWebsocketProvider, WebsocketClient,
     WebsocketClientCallback, WebsocketProvider,
 };
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::RwLock;
 
 struct MyWebsocketProvider {
     core: CoreWebsocketProvider,
@@ -55,27 +57,50 @@ impl WebsocketClient for MyWebsocketClient {
 
 #[derive(Default)]
 struct MyClientCallbacks {
-    data_read: Vec<u8>,
-    did_disconnect: bool,
-    did_error: bool,
+    data_read: RwLock<Vec<u8>>,
+    did_disconnect: AtomicBool,
+    did_error: AtomicBool,
 }
 
 impl WebsocketClientCallback for MyClientCallbacks {
-    fn connected(&mut self) -> bool {
+    fn connected(&self) -> bool {
         true
     }
 
-    fn disconnected(&mut self) {
-        self.did_disconnect = true;
+    fn disconnected(&self) {
+        self.did_disconnect.store(true, Ordering::Relaxed);
     }
 
-    fn error(&mut self, msg: &str) {
+    fn error(&self, msg: &str) {
         assert_eq!(msg, "error");
-        self.did_error = true;
+        self.did_error.store(true, Ordering::Relaxed);
     }
 
-    fn read(&mut self, data: &[u8]) -> bool {
-        self.data_read.extend_from_slice(data);
+    fn read(&self, data: &[u8]) -> bool {
+        self.data_read.write().unwrap().extend_from_slice(data);
+        true
+    }
+}
+
+#[derive(Default)]
+struct LifetimeCallbacks {
+    data: RwLock<Vec<u8>>,
+}
+
+impl WebsocketClientCallback for LifetimeCallbacks {
+    fn connected(&self) -> bool {
+        true
+    }
+
+    fn disconnected(&self) {}
+
+    fn error(&self, _msg: &str) {}
+
+    fn read(&self, data: &[u8]) -> bool {
+        if data == "sent: ".as_bytes() || data == "\n".as_bytes() {
+            return true;
+        }
+        assert_eq!(data, &self.data.read().unwrap()[..]);
         true
     }
 }
@@ -86,12 +111,12 @@ fn reg_websocket_provider() {
     let provider = register_websocket_provider::<MyWebsocketProvider>("RustWebsocketProvider");
     let client = provider.create_client().unwrap();
     let mut callback = MyClientCallbacks::default();
-    let success = client.initialize_connection(
+    let connection = client.connect(
         "url",
         [("header".to_string(), "value".to_string())],
         &mut callback,
     );
-    assert!(success, "Failed to initialize connection!");
+    assert!(connection.is_some(), "Failed to initialize connection!");
 }
 
 #[test]
@@ -101,24 +126,51 @@ fn listen_websocket_provider() {
 
     let client = provider.create_client().unwrap();
     let mut callback = MyClientCallbacks::default();
-    client.initialize_connection(
-        "url",
-        [("header".to_string(), "value".to_string())],
-        &mut callback,
-    );
+    let connection = client
+        .connect(
+            "url",
+            [("header".to_string(), "value".to_string())],
+            &callback,
+        )
+        .expect("Failed to initialize connection!");
 
-    assert!(client.write("test1".as_bytes()));
-    assert!(client.write("test2".as_bytes()));
+    assert!(connection.write("test1".as_bytes()));
+    assert!(connection.write("test2".as_bytes()));
 
-    client.notify_error("error");
-    client.disconnect();
-    drop(client);
+    connection.notify_error("error");
+    connection.disconnect();
 
     assert_eq!(
-        &callback.data_read[..],
+        &callback.data_read.read().unwrap()[..],
         "sent: test1\nsent: test2\n".as_bytes()
     );
     // If we disconnected that means the error callback was not notified.
-    assert!(!callback.did_disconnect);
-    assert!(callback.did_error);
+    assert!(!callback.did_disconnect.load(Ordering::Relaxed));
+    assert!(callback.did_error.load(Ordering::Relaxed));
+}
+
+#[test]
+fn correct_websocket_client_lifetime() {
+    let _session = Session::new().expect("Failed to initialize session");
+    let provider = register_websocket_provider::<MyWebsocketProvider>("RustWebsocketProvider2");
+
+    let client = provider.create_client().unwrap();
+    let callback = LifetimeCallbacks::default();
+    let connection = client
+        .connect(
+            "url",
+            [("header".to_string(), "value".to_string())],
+            &callback,
+        )
+        .expect("Failed to initialize connection!");
+
+    println!("{:?}", callback.data);
+    callback
+        .data
+        .write()
+        .unwrap()
+        .extend(vec![0x55, 0x55, 0x55, 0x55, 0x55]);
+
+    assert!(connection.write(&[0x55, 0x55, 0x55, 0x55, 0x55]));
+    assert!(connection.write(&[0x55, 0x55, 0x55, 0x55, 0x55]));
 }