handle socket EOF and more gracefully handle panics in Python bindings

assets/udbserver.h

@@ -1,3 +1,3 @@
 #include <stdint.h>
 
-void udbserver(void* handle, uint16_t port, uint64_t start_addr);
+int32_t udbserver(void* handle, uint16_t port, uint64_t start_addr);

bindings/c/README.md

@@ -5,14 +5,19 @@
 `udbserver` provides a simple API:
 
 ```c
-void udbserver(void* handle, uint16_t port, uint64_t start_addr);
+int32_t udbserver(void* handle, uint16_t port, uint64_t start_addr);
 ```
 
 Parameters:
 - `handle`: The raw handle of a Unicorn instance
 - `port`: The port number to listen on
 - `start_addr`: The address at which the debug server will start and wait for connection. If set to `0`, the debug server starts immediately
 
+Return value:
+- `0`: success
+- `-1`: recoverable runtime error
+- `-2`: panic trapped at the FFI boundary
+
 You can call this API inside a Unicorn hook to integrate `udbserver` within other Unicorn-based projects.
 
 ## Installation

bindings/python/udbserver/__init__.py

@@ -1,8 +1,9 @@
-import os
 import ctypes
+import os
 import sys
+from ctypes import c_int, c_uint16, c_uint64, c_void_p
+
 from unicorn import Uc
-from ctypes import c_void_p, c_uint16, c_uint64
 
 _current_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -14,9 +15,17 @@
 _udbserver_lib = ctypes.cdll.LoadLibrary(os.path.join(_current_dir, _library_file))
 
 _udbserver_lib.udbserver.argtypes = [c_void_p, c_uint16, c_uint64]
-_udbserver_lib.udbserver.restype = None
+_udbserver_lib.udbserver.restype = c_int
+
 
 def udbserver(uc: Uc, port: int = 1234, start_addr: int = 0):
     """Start udbserver.
     """
-    _udbserver_lib.udbserver(int(uc._uch.value), port, start_addr)
+    status = _udbserver_lib.udbserver(int(uc._uch.value), port, start_addr)
+    if status == 0:
+        return
+    if status == -1:
+        raise RuntimeError("udbserver failed with a runtime error")
+    if status == -2:
+        raise RuntimeError("udbserver panicked at the FFI boundary")
+    raise RuntimeError(f"udbserver failed with status code {status}")

src/capi.rs

@@ -3,23 +3,38 @@
 use singlyton::SingletonOption;
 use std::borrow::BorrowMut;
 use std::ffi::c_void;
+use std::panic::AssertUnwindSafe;
 use unicorn_engine::{uc_engine, Unicorn};
 
 type uc_handle = *mut c_void;
 
 static UNICORN: SingletonOption<Unicorn<()>> = SingletonOption::new();
 
-#[no_mangle]
-pub extern "C" fn udbserver(handle: uc_handle, port: u16, start_addr: u64) {
+fn start_udbserver(handle: uc_handle, port: u16, start_addr: u64) -> Result<(), String> {
     if UNICORN.is_some() {
-        return;
+        return Ok(());
     }
-    if let Ok(unicorn) = unsafe { Unicorn::from_handle(handle as *mut uc_engine) } {
-        UNICORN.replace(unicorn);
-    } else {
-        panic!("Failed to convert handle to Unicorn");
+    let unicorn = unsafe { Unicorn::from_handle(handle as *mut uc_engine) }.map_err(|error| format!("Failed to convert handle to Unicorn: {error}"))?;
+    UNICORN.replace(unicorn);
+    crate::udbserver(UNICORN.get_mut().borrow_mut(), port, start_addr).map_err(|error| format!("Failed to start udbserver: {error}"))
+}
+
+#[no_mangle]
+pub extern "C" fn udbserver(handle: uc_handle, port: u16, start_addr: u64) -> i32 {
+    let result = std::panic::catch_unwind(AssertUnwindSafe(|| start_udbserver(handle, port, start_addr)));
+    match result {
+        Ok(Ok(())) => 0,
+        Ok(Err(error)) => {
+            eprintln!("{error}");
+            clean();
+            -1
+        }
+        Err(_) => {
+            eprintln!("udbserver panicked");
+            clean();
+            -2
+        }
     }
-    crate::udbserver(UNICORN.get_mut().borrow_mut(), port, start_addr).expect("Failed to start udbserver");
 }
 
 pub fn clean() {

src/lib.rs

@@ -12,6 +12,7 @@ use gdbstub::target::ext::breakpoints::WatchKind;
 use singlyton::SingletonOption;
 use std::any::Any;
 use std::borrow::BorrowMut;
+use std::io::ErrorKind;
 use std::net::{TcpListener, TcpStream};
 use unicorn_engine::unicorn_const::HookType;
 use unicorn_engine::Unicorn;
@@ -94,6 +95,13 @@ pub(crate) fn udbserver_resume<T: 'static>(addr: Option<u64>) -> DynResult<()> {
     udbserver_loop::<T>()
 }
 
+fn is_disconnect_error(error: &std::io::Error) -> bool {
+    matches!(
+        error.kind(),
+        ErrorKind::UnexpectedEof | ErrorKind::ConnectionReset | ErrorKind::ConnectionAborted | ErrorKind::BrokenPipe | ErrorKind::NotConnected
+    )
+}
+
 fn udbserver_loop<T: 'static>() -> DynResult<()> {
     let gdb_any = GDBSTUB.take().unwrap();
     let mut gdb = *gdb_any
@@ -102,7 +110,11 @@ fn udbserver_loop<T: 'static>() -> DynResult<()> {
     loop {
         gdb = match gdb {
             GdbStubStateMachine::Idle(mut gdb_inner) => {
-                let byte = gdb_inner.borrow_conn().read()?;
+                let byte = match gdb_inner.borrow_conn().read() {
+                    Ok(byte) => byte,
+                    Err(error) if is_disconnect_error(&error) => return handle_disconnect(DisconnectReason::Disconnect),
+                    Err(error) => return Err(Box::new(error)),
+                };
                 let mut emu_any = EMU.get_mut();
                 let emu = emu_any.downcast_mut::<emu::Emu<T>>().expect("Failed to downcast EMU");
                 gdb_inner.incoming_data(emu.borrow_mut(), byte)?