Merge "rpc_binder: Refactor users of RpcServer to use new API"
diff --git a/authfs/fd_server/src/main.rs b/authfs/fd_server/src/main.rs
index f1fffdd..21d0e64 100644
--- a/authfs/fd_server/src/main.rs
+++ b/authfs/fd_server/src/main.rs
@@ -29,7 +29,7 @@
 use clap::Parser;
 use log::debug;
 use nix::sys::stat::{umask, Mode};
-use rpcbinder::run_vsock_rpc_server;
+use rpcbinder::RpcServer;
 use std::collections::BTreeMap;
 use std::fs::File;
 use std::os::unix::io::{FromRawFd, OwnedFd};
@@ -135,18 +135,14 @@
     let old_umask = umask(Mode::empty());
     debug!("Setting umask to 0 (old: {:03o})", old_umask.bits());
 
-    let service = FdService::new_binder(fd_pool).as_binder();
     debug!("fd_server is starting as a rpc service.");
-    let retval = run_vsock_rpc_server(service, RPC_SERVICE_PORT, || {
-        debug!("fd_server is ready");
-        // Close the ready-fd if we were given one to signal our readiness.
-        drop(ready_fd.take());
-    });
+    let service = FdService::new_binder(fd_pool).as_binder();
+    let server = RpcServer::new_vsock(service, RPC_SERVICE_PORT)?;
+    debug!("fd_server is ready");
 
-    if retval {
-        debug!("RPC server has shut down gracefully");
-        Ok(())
-    } else {
-        bail!("Premature termination of RPC server");
-    }
+    // Close the ready-fd if we were given one to signal our readiness.
+    drop(ready_fd.take());
+
+    server.join();
+    Ok(())
 }
diff --git a/authfs/service/src/main.rs b/authfs/service/src/main.rs
index 671c06a..e710f07 100644
--- a/authfs/service/src/main.rs
+++ b/authfs/service/src/main.rs
@@ -24,7 +24,7 @@
 
 use anyhow::{bail, Result};
 use log::*;
-use rpcbinder::run_init_unix_domain_rpc_server;
+use rpcbinder::RpcServer;
 use std::ffi::OsString;
 use std::fs::{create_dir, read_dir, remove_dir_all, remove_file};
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -117,15 +117,11 @@
 
     let service = AuthFsService::new_binder(debuggable).as_binder();
     debug!("{} is starting as a rpc service.", AUTHFS_SERVICE_SOCKET_NAME);
-    let retval = run_init_unix_domain_rpc_server(service, AUTHFS_SERVICE_SOCKET_NAME, || {
-        info!("The RPC server '{}' is running.", AUTHFS_SERVICE_SOCKET_NAME);
-    });
-    if retval {
-        info!("The RPC server at '{}' has shut down gracefully.", AUTHFS_SERVICE_SOCKET_NAME);
-        Ok(())
-    } else {
-        bail!("Premature termination of the RPC server '{}'.", AUTHFS_SERVICE_SOCKET_NAME)
-    }
+    let server = RpcServer::new_init_unix_domain(service, AUTHFS_SERVICE_SOCKET_NAME)?;
+    info!("The RPC server '{}' is running.", AUTHFS_SERVICE_SOCKET_NAME);
+    server.join();
+    info!("The RPC server at '{}' has shut down gracefully.", AUTHFS_SERVICE_SOCKET_NAME);
+    Ok(())
 }
 
 fn main() {
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index a4e3903..206dd4b 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -22,10 +22,10 @@
 mod compsvc;
 mod fsverity;
 
-use anyhow::{bail, Result};
+use anyhow::Result;
 use compos_common::COMPOS_VSOCK_PORT;
 use log::{debug, error};
-use rpcbinder::run_vsock_rpc_server;
+use rpcbinder::RpcServer;
 use std::panic;
 use vm_payload_bindgen::AVmPayload_notifyPayloadReady;
 
@@ -45,16 +45,11 @@
         error!("{}", panic_info);
     }));
 
-    let service = compsvc::new_binder()?.as_binder();
     debug!("compsvc is starting as a rpc service.");
+    let service = compsvc::new_binder()?.as_binder();
+    let server = RpcServer::new_vsock(service, COMPOS_VSOCK_PORT)?;
     // SAFETY: Invokes a method from the bindgen library `vm_payload_bindgen`.
-    let retval = run_vsock_rpc_server(service, COMPOS_VSOCK_PORT, || unsafe {
-        AVmPayload_notifyPayloadReady();
-    });
-    if retval {
-        debug!("RPC server has shut down gracefully");
-        Ok(())
-    } else {
-        bail!("Premature termination of RPC server");
-    }
+    unsafe { AVmPayload_notifyPayloadReady() };
+    server.join();
+    Ok(())
 }
diff --git a/microdroid/vm_payload/src/vm_payload_service.rs b/microdroid/vm_payload/src/vm_payload_service.rs
index 88484cc..874b7e1 100644
--- a/microdroid/vm_payload/src/vm_payload_service.rs
+++ b/microdroid/vm_payload/src/vm_payload_service.rs
@@ -20,7 +20,7 @@
 use binder::{Strong, unstable_api::{AIBinder, new_spibinder}};
 use lazy_static::lazy_static;
 use log::{error, info, Level};
-use rpcbinder::{get_unix_domain_rpc_interface, run_vsock_rpc_server};
+use rpcbinder::{get_unix_domain_rpc_interface, RpcServer};
 use std::ffi::CString;
 use std::os::raw::{c_char, c_void};
 use std::ptr;
@@ -105,11 +105,19 @@
     // safely be taken by new_spibinder.
     let service = new_spibinder(service);
     if let Some(service) = service {
-        run_vsock_rpc_server(service, port, || {
-            if let Some(on_ready) = on_ready {
-                on_ready(param);
+        match RpcServer::new_vsock(service, port) {
+            Ok(server) => {
+                if let Some(on_ready) = on_ready {
+                    on_ready(param);
+                }
+                server.join();
+                true
             }
-        })
+            Err(err) => {
+                error!("Failed to start RpcServer: {:?}", err);
+                false
+            }
+        }
     } else {
         error!("Failed to convert the given service from AIBinder to SpIBinder.");
         false
diff --git a/microdroid_manager/src/vm_payload_service.rs b/microdroid_manager/src/vm_payload_service.rs
index fcfc79d..98b9f2b 100644
--- a/microdroid_manager/src/vm_payload_service.rs
+++ b/microdroid_manager/src/vm_payload_service.rs
@@ -18,15 +18,12 @@
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
     BnVmPayloadService, IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME};
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
-use anyhow::{bail, Result};
+use anyhow::Result;
 use binder::{Interface, BinderFeatures, ExceptionCode, Status, Strong};
 use log::{error, info};
 use openssl::hkdf::hkdf;
 use openssl::md::Md;
-use rpcbinder::run_init_unix_domain_rpc_server;
-use std::sync::mpsc;
-use std::thread;
-use std::time::Duration;
+use rpcbinder::RpcServer;
 
 /// Implementation of `IVmPayloadService`.
 struct VmPayloadService {
@@ -101,29 +98,16 @@
         VmPayloadService::new(allow_restricted_apis, vm_service, dice),
         BinderFeatures::default(),
     );
-    let (sender, receiver) = mpsc::channel();
-    thread::spawn(move || {
-        let retval = run_init_unix_domain_rpc_server(
-            vm_payload_binder.as_binder(),
-            VM_PAYLOAD_SERVICE_SOCKET_NAME,
-            || {
-                sender.send(()).unwrap();
-            },
-        );
-        if retval {
-            info!(
-                "The RPC server at '{}' has shut down gracefully.",
-                VM_PAYLOAD_SERVICE_SOCKET_NAME
-            );
-        } else {
-            error!("Premature termination of the RPC server '{}'.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
-        }
+
+    let server = RpcServer::new_init_unix_domain(
+        vm_payload_binder.as_binder(),
+        VM_PAYLOAD_SERVICE_SOCKET_NAME,
+    )?;
+    info!("The RPC server '{}' is running.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
+
+    // Move server reference into a background thread and run it forever.
+    std::thread::spawn(move || {
+        server.join();
     });
-    match receiver.recv_timeout(Duration::from_millis(200)) {
-        Ok(()) => {
-            info!("The RPC server '{}' is running.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
-            Ok(())
-        }
-        _ => bail!("Failed to register service '{}'", VM_PAYLOAD_SERVICE_SOCKET_NAME),
-    }
+    Ok(())
 }