Add safe wrapper for RunRpcServerWithFactory.

Test: atest compos_key_tests MicrodroidHostTestCases MicrodroidTestApp
Change-Id: Iaab166d29e0ec08161db50c5e1419283bf304499
diff --git a/libs/binder_common/rpc_server.rs b/libs/binder_common/rpc_server.rs
index 1326755..4261358 100644
--- a/libs/binder_common/rpc_server.rs
+++ b/libs/binder_common/rpc_server.rs
@@ -18,7 +18,7 @@
 
 use binder::unstable_api::AsNative;
 use binder::SpIBinder;
-use std::os::raw;
+use std::{os::raw, ptr::null_mut};
 
 /// Runs a binder RPC server, serving the supplied binder service implementation on the given vsock
 /// port.
@@ -82,3 +82,46 @@
         }
     }
 }
+
+type RpcServerFactoryRef<'a> = &'a mut (dyn FnMut(u32) -> Option<SpIBinder> + Send + Sync);
+
+/// Runs a binder RPC server, using the given factory function to construct a binder service
+/// implementation for each connection.
+///
+/// The current thread is joined to the binder thread pool to handle incoming messages.
+///
+/// Returns true if the server has shutdown normally, false if it failed in some way.
+pub fn run_rpc_server_with_factory(
+    port: u32,
+    mut factory: impl FnMut(u32) -> Option<SpIBinder> + Send + Sync,
+) -> bool {
+    // Double reference the factory because trait objects aren't FFI safe.
+    // NB: The type annotation is necessary to ensure that we have a `dyn` rather than an `impl`.
+    let mut factory_ref: RpcServerFactoryRef = &mut factory;
+    let context = &mut factory_ref as *mut RpcServerFactoryRef as *mut raw::c_void;
+
+    // SAFETY: `factory_wrapper` is only ever called by `RunRpcServerWithFactory`, with context
+    // taking the pointer value above (so a properly aligned non-null pointer to an initialized
+    // `RpcServerFactoryRef`), within the lifetime of `factory_ref` (i.e. no more calls will be made
+    // after `RunRpcServerWithFactory` returns).
+    unsafe {
+        binder_rpc_unstable_bindgen::RunRpcServerWithFactory(Some(factory_wrapper), context, port)
+    }
+}
+
+unsafe extern "C" fn factory_wrapper(
+    cid: u32,
+    context: *mut raw::c_void,
+) -> *mut binder_rpc_unstable_bindgen::AIBinder {
+    // SAFETY: `context` was created from an `&mut RpcServerFactoryRef` by
+    // `run_rpc_server_with_factory`, and we are still within the lifetime of the value it is
+    // pointing to.
+    let factory_ptr = context as *mut RpcServerFactoryRef;
+    let factory = factory_ptr.as_mut().unwrap();
+
+    if let Some(mut service) = factory(cid) {
+        service.as_native_mut() as *mut binder_rpc_unstable_bindgen::AIBinder
+    } else {
+        null_mut()
+    }
+}
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 0c9496a..0a5436b 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -26,7 +26,6 @@
         "libandroid_logger",
         "libanyhow",
         "libbinder_common",
-        "libbinder_rpc_unstable_bindgen",
         "libbinder_rs",
         "libcommand_fds",
         "libdisk",
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index e2e76d5..cc8d8a3 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -19,7 +19,6 @@
 use crate::payload::add_microdroid_images;
 use crate::{Cid, FIRST_GUEST_CID, SYSPROP_LAST_CID};
 use crate::selinux::{SeContext, getfilecon};
-use ::binder::unstable_api::AsNative;
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
     DeathReason::DeathReason,
@@ -36,8 +35,8 @@
     VirtualMachineState::VirtualMachineState,
 };
 use android_system_virtualizationservice::binder::{
-    self, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, StatusCode, Strong,
-    ThreadState,
+    self, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, SpIBinder, Status,
+    StatusCode, Strong, ThreadState,
 };
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
     IVirtualMachineService::{
@@ -46,7 +45,7 @@
     },
 };
 use anyhow::{anyhow, bail, Context, Result};
-use binder_common::{lazy_service::LazyServiceGuard, new_binder_exception};
+use binder_common::{lazy_service::LazyServiceGuard, new_binder_exception, rpc_server::run_rpc_server_with_factory};
 use disk::QcowFile;
 use idsig::{HashAlgorithm, V4Signature};
 use log::{debug, error, info, warn, trace};
@@ -59,10 +58,8 @@
 use std::fs::{create_dir, File, OpenOptions};
 use std::io::{Error, ErrorKind, Write, Read};
 use std::num::NonZeroU32;
-use std::os::raw;
 use std::os::unix::io::{FromRawFd, IntoRawFd};
 use std::path::{Path, PathBuf};
-use std::ptr::null_mut;
 use std::sync::{Arc, Mutex, Weak};
 use tombstoned_client::{TombstonedConnection, DebuggerdDumpType};
 use vmconfig::VmConfig;
@@ -330,28 +327,16 @@
 
         // binder server for vm
         // reference to state (not the state itself) is copied
-        let mut state = service.state.clone();
+        let state = service.state.clone();
         std::thread::spawn(move || {
-            let state_ptr = &mut state as *mut _ as *mut raw::c_void;
-
-            debug!("virtual machine service is starting as an RPC service.");
-            // SAFETY: factory function is only ever called by RunRpcServerWithFactory, within the
-            // lifetime of the state, with context taking the pointer value above (so a properly
-            // aligned non-null pointer to an initialized instance).
-            let retval = unsafe {
-                binder_rpc_unstable_bindgen::RunRpcServerWithFactory(
-                    Some(VirtualMachineService::factory),
-                    state_ptr,
-                    VM_BINDER_SERVICE_PORT as u32,
-                )
-            };
-            if retval {
+            debug!("VirtualMachineService is starting as an RPC service.");
+            if run_rpc_server_with_factory(VM_BINDER_SERVICE_PORT as u32, |cid| {
+                VirtualMachineService::factory(cid, &state)
+            }) {
                 debug!("RPC server has shut down gracefully");
             } else {
-                bail!("Premature termination of RPC server");
+                panic!("Premature termination of RPC server");
             }
-
-            Ok(retval)
         });
         service
     }
@@ -1143,20 +1128,14 @@
 }
 
 impl VirtualMachineService {
-    // SAFETY: Service ownership is held by state, and the binder objects are threadsafe.
-    pub unsafe extern "C" fn factory(
-        cid: Cid,
-        context: *mut raw::c_void,
-    ) -> *mut binder_rpc_unstable_bindgen::AIBinder {
-        let state_ptr = context as *mut Arc<Mutex<State>>;
-        let state = state_ptr.as_ref().unwrap();
+    fn factory(cid: Cid, state: &Arc<Mutex<State>>) -> Option<SpIBinder> {
         if let Some(vm) = state.lock().unwrap().get_vm(cid) {
             let mut vm_service = vm.vm_service.lock().unwrap();
             let service = vm_service.get_or_insert_with(|| Self::new_binder(state.clone(), cid));
-            service.as_binder().as_native_mut() as *mut binder_rpc_unstable_bindgen::AIBinder
+            Some(service.as_binder())
         } else {
             error!("connection from cid={} is not from a guest VM", cid);
-            null_mut()
+            None
         }
     }