Add safe wrapper for RpcPreconnectedClient.

Test: atest compos_key_tests MicrodroidHostTestCases MicrodroidTestApp
Change-Id: I8dcc6c9b0465950bfaced03699fa3167dc3dc641
diff --git a/libs/binder_common/rpc_client.rs b/libs/binder_common/rpc_client.rs
index 1aabe84..33fd732 100644
--- a/libs/binder_common/rpc_client.rs
+++ b/libs/binder_common/rpc_client.rs
@@ -17,10 +17,11 @@
 //! Helpers for implementing an RPC Binder client.
 
 use binder::unstable_api::{new_spibinder, AIBinder};
-use binder::{StatusCode, Strong};
+use binder::{FromIBinder, StatusCode, Strong};
+use std::os::{raw, unix::io::RawFd};
 
 /// Connects to a binder RPC server.
-pub fn connect_rpc_binder<T: binder::FromIBinder + ?Sized>(
+pub fn connect_rpc_binder<T: FromIBinder + ?Sized>(
     cid: u32,
     port: u32,
 ) -> Result<Strong<T>, StatusCode> {
@@ -35,3 +36,40 @@
         Err(StatusCode::BAD_VALUE)
     }
 }
+
+type RequestFd<'a> = &'a mut dyn FnMut() -> Option<RawFd>;
+
+/// Connects to a Binder RPC server, using the given callback to get (and take ownership of) file
+/// descriptors already connected to it.
+pub fn connect_preconnected_rpc_binder<T: FromIBinder + ?Sized>(
+    mut request_fd: impl FnMut() -> Option<RawFd>,
+) -> Result<Strong<T>, StatusCode> {
+    // Double reference the factory because trait objects aren't FFI safe.
+    let mut request_fd_ref: RequestFd = &mut request_fd;
+    let param = &mut request_fd_ref as *mut RequestFd as *mut raw::c_void;
+
+    // SAFETY: AIBinder returned by RpcPreconnectedClient has correct reference count, and the
+    // ownership can be safely taken by new_spibinder. RpcPreconnectedClient does not take ownership
+    // of param, only passing it to request_fd_wrapper.
+    let ibinder = unsafe {
+        new_spibinder(binder_rpc_unstable_bindgen::RpcPreconnectedClient(
+            Some(request_fd_wrapper),
+            param,
+        ) as *mut AIBinder)
+    };
+
+    if let Some(ibinder) = ibinder {
+        <T>::try_from(ibinder)
+    } else {
+        Err(StatusCode::BAD_VALUE)
+    }
+}
+
+unsafe extern "C" fn request_fd_wrapper(param: *mut raw::c_void) -> raw::c_int {
+    // SAFETY: This is only ever called by RpcPreconnectedClient, within the lifetime of the
+    // BinderFdFactory reference, with param being a properly aligned non-null pointer to an
+    // initialized instance.
+    let request_fd_ptr = param as *mut RequestFd;
+    let request_fd = request_fd_ptr.as_mut().unwrap();
+    request_fd().unwrap_or(-1)
+}
diff --git a/vmclient/Android.bp b/vmclient/Android.bp
index 8ad5adf..c219198 100644
--- a/vmclient/Android.bp
+++ b/vmclient/Android.bp
@@ -9,7 +9,7 @@
     edition: "2021",
     rustlibs: [
         "android.system.virtualizationservice-rust",
-        "libbinder_rpc_unstable_bindgen",
+        "libbinder_common",
         "libbinder_rs",
         "liblog_rust",
         "libthiserror",
diff --git a/vmclient/src/errors.rs b/vmclient/src/errors.rs
index 43db7f9..231f81f 100644
--- a/vmclient/src/errors.rs
+++ b/vmclient/src/errors.rs
@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use super::DeathReason;
-use android_system_virtualizationservice::binder::StatusCode;
 use thiserror::Error;
 
 /// An error while waiting for a VM to do something.
@@ -32,14 +31,3 @@
     #[error("VM payload finished.")]
     Finished,
 }
-
-/// An error connecting to a VM RPC Binder service.
-#[derive(Clone, Debug, Eq, Error, PartialEq)]
-pub enum ConnectServiceError {
-    /// The RPC binder connection failed.
-    #[error("Vsock connection to RPC binder failed.")]
-    ConnectionFailed,
-    /// The AIDL service type didn't match.
-    #[error("Service type didn't match ({0}).")]
-    WrongServiceType(StatusCode),
-}
diff --git a/vmclient/src/lib.rs b/vmclient/src/lib.rs
index 9b5b8dd..b3bb635 100644
--- a/vmclient/src/lib.rs
+++ b/vmclient/src/lib.rs
@@ -16,12 +16,11 @@
 
 mod death_reason;
 mod errors;
-mod rpc_binder;
 mod sync;
 
 pub use crate::death_reason::DeathReason;
-pub use crate::errors::{ConnectServiceError, VmWaitError};
-use crate::{rpc_binder::VsockFactory, sync::Monitor};
+pub use crate::errors::VmWaitError;
+use crate::sync::Monitor;
 use android_system_virtualizationservice::{
     aidl::android::system::virtualizationservice::{
         DeathReason::DeathReason as AidlDeathReason,
@@ -36,10 +35,12 @@
         ParcelFileDescriptor, Result as BinderResult, StatusCode, Strong,
     },
 };
+use binder_common::rpc_client::connect_preconnected_rpc_binder;
 use log::warn;
 use std::{
     fmt::{self, Debug, Formatter},
     fs::File,
+    os::unix::io::IntoRawFd,
     sync::Arc,
     time::Duration,
 };
@@ -145,12 +146,19 @@
     pub fn connect_service<T: FromIBinder + ?Sized>(
         &self,
         port: u32,
-    ) -> Result<Strong<T>, ConnectServiceError> {
-        let mut vsock_factory = VsockFactory::new(&*self.vm, port);
-
-        let ibinder = vsock_factory.connect_rpc_client()?;
-
-        FromIBinder::try_from(ibinder).map_err(ConnectServiceError::WrongServiceType)
+    ) -> Result<Strong<T>, StatusCode> {
+        connect_preconnected_rpc_binder(|| {
+            match self.vm.connectVsock(port as i32) {
+                Ok(vsock) => {
+                    // Ownership of the fd is transferred to binder
+                    Some(vsock.into_raw_fd())
+                }
+                Err(e) => {
+                    warn!("Vsock connection failed: {}", e);
+                    None
+                }
+            }
+        })
     }
 
     /// Get ramdump
diff --git a/vmclient/src/rpc_binder.rs b/vmclient/src/rpc_binder.rs
deleted file mode 100644
index 7c2992b..0000000
--- a/vmclient/src/rpc_binder.rs
+++ /dev/null
@@ -1,72 +0,0 @@
-// Copyright 2022, The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use crate::errors::ConnectServiceError;
-use android_system_virtualizationservice::{
-    aidl::android::system::virtualizationservice::IVirtualMachine::IVirtualMachine,
-};
-use binder::unstable_api::{new_spibinder, AIBinder};
-use log::warn;
-use std::os::{raw, unix::io::IntoRawFd};
-
-pub struct VsockFactory<'a> {
-    vm: &'a dyn IVirtualMachine,
-    port: u32,
-}
-
-impl<'a> VsockFactory<'a> {
-    pub fn new(vm: &'a dyn IVirtualMachine, port: u32) -> Self {
-        Self { vm, port }
-    }
-
-    pub fn connect_rpc_client(&mut self) -> Result<binder::SpIBinder, ConnectServiceError> {
-        let param = self.as_void_ptr();
-
-        unsafe {
-            // SAFETY: AIBinder returned by RpcPreconnectedClient has correct reference count, and
-            // the ownership can be safely taken by new_spibinder.
-            // RpcPreconnectedClient does not take ownership of param, only passing it to
-            // request_fd.
-            let binder =
-                binder_rpc_unstable_bindgen::RpcPreconnectedClient(Some(Self::request_fd), param)
-                    as *mut AIBinder;
-            new_spibinder(binder).ok_or(ConnectServiceError::ConnectionFailed)
-        }
-    }
-
-    fn as_void_ptr(&mut self) -> *mut raw::c_void {
-        self as *mut _ as *mut raw::c_void
-    }
-
-    fn new_vsock_fd(&self) -> i32 {
-        match self.vm.connectVsock(self.port as i32) {
-            Ok(vsock) => {
-                // Ownership of the fd is transferred to binder
-                vsock.into_raw_fd()
-            }
-            Err(e) => {
-                warn!("Vsock connection failed: {}", e);
-                -1
-            }
-        }
-    }
-
-    unsafe extern "C" fn request_fd(param: *mut raw::c_void) -> raw::c_int {
-        // SAFETY: This is only ever called by RpcPreconnectedClient, within the lifetime of the
-        // VsockFactory, with param taking the value returned by as_void_ptr (so a properly aligned
-        // non-null pointer to an initialized instance).
-        let vsock_factory = param as *mut Self;
-        vsock_factory.as_ref().unwrap().new_vsock_fd()
-    }
-}