compsvc: Notify VS when we're ready

Once our server socket is listening, notify virtualizationservice that
we're ready.

Modify clients to wait for ready instead of sleeping.

Bug: 186126194
Test: Manually run compos_key_cmd, compos_verify_key
Change-Id: Ibaff4bf22de9de2aa4faff9e81b24c779643b78c
diff --git a/compos/Android.bp b/compos/Android.bp
index f77b2ca..2597c1c 100644
--- a/compos/Android.bp
+++ b/compos/Android.bp
@@ -34,6 +34,7 @@
     rustlibs: [
         "android.hardware.security.keymint-V1-rust",
         "android.system.keystore2-V1-rust",
+        "android.system.virtualmachineservice-rust",
         "authfs_aidl_interface-rust",
         "compos_aidl_interface-rust",
         "libandroid_logger",
@@ -46,6 +47,7 @@
         "liblibc",
         "liblog_rust",
         "libminijail_rust",
+        "libnix",
         "libring",
         "libscopeguard",
     ],
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 6600f6f..c525f40 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -106,10 +106,7 @@
 
         vm.start()?;
 
-        let cid = vm_state.wait_for_start()?;
-
-        // TODO: Use onPayloadReady to avoid this
-        thread::sleep(Duration::from_secs(3));
+        let cid = vm_state.wait_until_ready()?;
 
         Ok(VmInstance { service, vm, cid })
     }
@@ -214,7 +211,7 @@
         self.state_ready.notify_all();
     }
 
-    fn set_started(&self, cid: i32) {
+    fn set_ready(&self, cid: i32) {
         let mut state = self.mutex.lock().unwrap();
         if state.has_died {
             return;
@@ -224,10 +221,10 @@
         self.state_ready.notify_all();
     }
 
-    fn wait_for_start(&self) -> Result<i32> {
+    fn wait_until_ready(&self) -> Result<i32> {
         let (state, result) = self
             .state_ready
-            .wait_timeout_while(self.mutex.lock().unwrap(), Duration::from_secs(10), |state| {
+            .wait_timeout_while(self.mutex.lock().unwrap(), Duration::from_secs(20), |state| {
                 state.cid.is_none() && !state.has_died
             })
             .unwrap();
@@ -255,7 +252,6 @@
         cid: i32,
         stream: Option<&ParcelFileDescriptor>,
     ) -> BinderResult<()> {
-        self.0.set_started(cid);
         if let Some(pfd) = stream {
             if let Err(e) = start_logging(pfd) {
                 warn!("Can't log vm output: {}", e);
@@ -266,7 +262,7 @@
     }
 
     fn onPayloadReady(&self, cid: i32) -> BinderResult<()> {
-        // TODO: Use this to trigger vsock connection
+        self.0.set_ready(cid);
         log::info!("VM payload ready, cid = {}", cid);
         Ok(())
     }
diff --git a/compos/compos_key_cmd/compos_key_cmd.cpp b/compos/compos_key_cmd/compos_key_cmd.cpp
index 4e51d8d..e168648 100644
--- a/compos/compos_key_cmd/compos_key_cmd.cpp
+++ b/compos/compos_key_cmd/compos_key_cmd.cpp
@@ -126,16 +126,16 @@
         std::thread logger([fd = std::move(stream_fd)]() mutable { copyToLog(std::move(fd)); });
         logger.detach();
 
-        {
-            std::unique_lock lock(mMutex);
-            mStarted = true;
-        }
-        mCv.notify_all();
         return ScopedAStatus::ok();
     }
 
     ::ndk::ScopedAStatus onPayloadReady(int32_t in_cid) override {
         LOG(INFO) << "Payload is ready! cid = " << in_cid;
+        {
+            std::unique_lock lock(mMutex);
+            mReady = true;
+        }
+        mCv.notify_all();
         return ScopedAStatus::ok();
     }
 
@@ -154,16 +154,16 @@
         return ScopedAStatus::ok();
     }
 
-    bool waitForStarted() {
+    bool waitUntilReady() {
         std::unique_lock lock(mMutex);
-        return mCv.wait_for(lock, std::chrono::seconds(10), [this] { return mStarted || mDied; }) &&
+        return mCv.wait_for(lock, std::chrono::seconds(20), [this] { return mReady || mDied; }) &&
                 !mDied;
     }
 
 private:
     std::mutex mMutex;
     std::condition_variable mCv;
-    bool mStarted;
+    bool mReady;
     bool mDied;
 };
 
@@ -263,14 +263,10 @@
         }
         LOG(INFO) << "Started VM";
 
-        if (!mCallback->waitForStarted()) {
+        if (!mCallback->waitUntilReady()) {
             return Error() << "VM Payload failed to start";
         }
 
-        // TODO(b/194677789): Implement a polling loop or find a more reliable
-        // way to detect when the service is listening.
-        sleep(3);
-
         return cid;
     }
 
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index 2b445c1..388e79b 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -22,10 +22,26 @@
 mod fsverity;
 mod signer;
 
-use anyhow::{bail, Result};
-use binder::unstable_api::AsNative;
+use android_system_virtualmachineservice::{
+    aidl::android::system::virtualmachineservice::IVirtualMachineService::{
+        IVirtualMachineService, VM_BINDER_SERVICE_PORT,
+    },
+    binder::Strong,
+};
+use anyhow::{anyhow, bail, Context, Result};
+use binder::{
+    unstable_api::{new_spibinder, AIBinder, AsNative},
+    FromIBinder,
+};
 use compos_common::COMPOS_VSOCK_PORT;
-use log::debug;
+use log::{debug, error};
+use nix::ioctl_read_bad;
+use std::fs::OpenOptions;
+use std::os::raw;
+use std::os::unix::io::AsRawFd;
+
+/// The CID representing the host VM
+const VMADDR_CID_HOST: u32 = 2;
 
 fn main() -> Result<()> {
     let args = clap::App::new("compsvc")
@@ -41,12 +57,19 @@
 
     let mut service = compsvc::new_binder()?.as_binder();
     debug!("compsvc is starting as a rpc service.");
+
+    let mut ready_notifier = ReadyNotifier::new()?;
+
     // SAFETY: Service ownership is transferring to the server and won't be valid afterward.
     // Plus the binder objects are threadsafe.
+    // RunRpcServerCallback does not retain a reference to ready_callback, and only ever
+    // calls it with the param we provide during the lifetime of ready_notifier.
     let retval = unsafe {
-        binder_rpc_unstable_bindgen::RunRpcServer(
+        binder_rpc_unstable_bindgen::RunRpcServerCallback(
             service.as_native_mut() as *mut binder_rpc_unstable_bindgen::AIBinder,
             COMPOS_VSOCK_PORT,
+            Some(ReadyNotifier::ready_callback),
+            ready_notifier.as_void_ptr(),
         )
     };
     if retval {
@@ -56,3 +79,69 @@
         bail!("Premature termination of RPC server");
     }
 }
+
+struct ReadyNotifier {
+    vm_service: Strong<dyn IVirtualMachineService>,
+    local_cid: u32,
+}
+
+impl ReadyNotifier {
+    fn new() -> Result<Self> {
+        Ok(Self { vm_service: Self::get_vm_service()?, local_cid: Self::get_local_cid()? })
+    }
+
+    fn notify(&self) {
+        if let Err(e) = self.vm_service.notifyPayloadReady(self.local_cid as i32) {
+            error!("Unable to notify ready: {}", e);
+        }
+    }
+
+    fn as_void_ptr(&mut self) -> *mut raw::c_void {
+        self as *mut _ as *mut raw::c_void
+    }
+
+    unsafe extern "C" fn ready_callback(param: *mut raw::c_void) {
+        // SAFETY: This is only ever called by RunRpcServerCallback, within the lifetime of the
+        // ReadyNotifier, with param taking the value returned by as_void_ptr (so a properly aligned
+        // non-null pointer to an initialized instance).
+        let ready_notifier = param as *mut Self;
+        ready_notifier.as_ref().unwrap().notify()
+    }
+
+    fn get_vm_service() -> Result<Strong<dyn IVirtualMachineService>> {
+        // SAFETY: AIBinder returned by RpcClient has correct reference count, and the ownership
+        // can be safely taken by new_spibinder.
+        let ibinder = unsafe {
+            new_spibinder(binder_rpc_unstable_bindgen::RpcClient(
+                VMADDR_CID_HOST,
+                VM_BINDER_SERVICE_PORT as u32,
+            ) as *mut AIBinder)
+        }
+        .ok_or_else(|| anyhow!("Failed to connect to IVirtualMachineService"))?;
+
+        FromIBinder::try_from(ibinder).context("Connecting to IVirtualMachineService")
+    }
+
+    // TODO(b/199259751): remove this after VS can check the peer addresses of binder clients
+    fn get_local_cid() -> Result<u32> {
+        let f = OpenOptions::new()
+            .read(true)
+            .write(false)
+            .open("/dev/vsock")
+            .context("Failed to open /dev/vsock")?;
+        let mut cid = 0;
+        // SAFETY: the kernel only modifies the given u32 integer.
+        unsafe { vm_sockets_get_local_cid(f.as_raw_fd(), &mut cid) }
+            .context("Failed to get local CID")?;
+        Ok(cid)
+    }
+}
+
+// TODO(b/199259751): remove this after VS can check the peer addresses of binder clients
+const IOCTL_VM_SOCKETS_GET_LOCAL_CID: usize = 0x7b9;
+ioctl_read_bad!(
+    /// Gets local cid from /dev/vsock
+    vm_sockets_get_local_cid,
+    IOCTL_VM_SOCKETS_GET_LOCAL_CID,
+    u32
+);
diff --git a/virtualizationservice/aidl/Android.bp b/virtualizationservice/aidl/Android.bp
index 571cc5d..30a4b03 100644
--- a/virtualizationservice/aidl/Android.bp
+++ b/virtualizationservice/aidl/Android.bp
@@ -39,7 +39,10 @@
     backend: {
         rust: {
             enabled: true,
-            apex_available: ["com.android.virt"],
+            apex_available: [
+                "com.android.virt",
+                "com.android.compos",
+            ],
         },
     },
 }