Add VirtualMachineState enum rather than isRunning bool.

Keep track of payload state and enforce lifecycle expectations.

Bug: 199127239
Test: atest VirtualizationTestCases
Change-Id: I04646c96ec3c15854b6ae705e40979d0f2e29457
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index f4ac467..af53dae 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -31,6 +31,7 @@
 import android.system.virtualizationservice.IVirtualizationService;
 import android.system.virtualizationservice.PartitionType;
 import android.system.virtualizationservice.VirtualMachineAppConfig;
+import android.system.virtualizationservice.VirtualMachineState;
 
 import java.io.File;
 import java.io.FileInputStream;
@@ -224,8 +225,16 @@
     /** Returns the current status of this virtual machine. */
     public @NonNull Status getStatus() throws VirtualMachineException {
         try {
-            if (mVirtualMachine != null && mVirtualMachine.isRunning()) {
-                return Status.RUNNING;
+            if (mVirtualMachine != null) {
+                switch (mVirtualMachine.getState()) {
+                    case VirtualMachineState.STARTING:
+                    case VirtualMachineState.STARTED:
+                    case VirtualMachineState.READY:
+                    case VirtualMachineState.FINISHED:
+                        return Status.RUNNING;
+                    case VirtualMachineState.DEAD:
+                        return Status.STOPPED;
+                }
             }
         } catch (RemoteException e) {
             throw new VirtualMachineException(e);
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
index 081580c..3c89cd7 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
@@ -16,13 +16,14 @@
 package android.system.virtualizationservice;
 
 import android.system.virtualizationservice.IVirtualMachineCallback;
+import android.system.virtualizationservice.VirtualMachineState;
 
 interface IVirtualMachine {
     /** Get the CID allocated to the VM. */
     int getCid();
 
-    /** Returns true if the VM is still running, or false if it has exited for any reason. */
-    boolean isRunning();
+    /** Returns the current lifecycle state of the VM. */
+    VirtualMachineState getState();
 
     /**
      * Register a Binder object to get callbacks when the state of the VM changes, such as if it
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
index d081b8d..672c41a 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
@@ -15,6 +15,8 @@
  */
 package android.system.virtualizationservice;
 
+import android.system.virtualizationservice.VirtualMachineState;
+
 /** Information about a running VM, for debug purposes only. */
 parcelable VirtualMachineDebugInfo {
     /** The CID assigned to the VM. */
@@ -35,6 +37,6 @@
      */
     int requesterPid;
 
-    /** Whether the VM is still running. */
-    boolean running;
+    /** The current lifecycle state of the VM. */
+    VirtualMachineState state;
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineState.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineState.aidl
new file mode 100644
index 0000000..621887f
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineState.aidl
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2021 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.
+ */
+package android.system.virtualizationservice;
+
+/**
+ * The lifecycle state of a VM.
+ */
+@Backing(type="int")
+enum VirtualMachineState {
+    /**
+     * The VM is running, but the payload has not yet started.
+     */
+    STARTING = 1,
+    /**
+     * The VM is running and the payload has been started, but it has not yet indicated that it is
+     * ready.
+     */
+    STARTED = 2,
+    /**
+     * The VM payload has indicated that it is ready to serve requests.
+     */
+    READY = 3,
+    /**
+     * The VM payload has finished but the VM itself is still running.
+     */
+    FINISHED = 4,
+    /**
+     * The VM has died.
+     */
+    DEAD = 5,
+}
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 3cf7ca3..f89181f 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -15,24 +15,25 @@
 //! Implementation of the AIDL interface of the VirtualizationService.
 
 use crate::composite::make_composite_image;
-use crate::crosvm::{CrosvmConfig, DiskFile, VmInstance};
+use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmInstance};
 use crate::payload::add_microdroid_images;
 use crate::{Cid, FIRST_GUEST_CID};
 
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::DiskImage::DiskImage;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::IVirtualizationService;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualMachine::{
     BnVirtualMachine, IVirtualMachine,
 };
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualMachineCallback::IVirtualMachineCallback;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+    DiskImage::DiskImage,
+    IVirtualMachineCallback::IVirtualMachineCallback,
+    IVirtualizationService::IVirtualizationService,
+    PartitionType::PartitionType,
     VirtualMachineAppConfig::VirtualMachineAppConfig,
     VirtualMachineConfig::VirtualMachineConfig,
+    VirtualMachineDebugInfo::VirtualMachineDebugInfo,
     VirtualMachineRawConfig::VirtualMachineRawConfig,
+    VirtualMachineState::VirtualMachineState,
 };
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::VirtualMachineDebugInfo::VirtualMachineDebugInfo;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::PartitionType::PartitionType;
 use android_system_virtualizationservice::binder::{
     self, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, Strong, ThreadState,
 };
@@ -267,7 +268,7 @@
                 requesterUid: vm.requester_uid as i32,
                 requesterSid: vm.requester_sid.clone(),
                 requesterPid: vm.requester_debug_pid,
-                running: vm.running(),
+                state: get_state(&vm),
             })
             .collect();
         Ok(cids)
@@ -329,8 +330,8 @@
     }
 }
 
-/// Waits for incoming connections from VM. If a new connection is made, notify the event to the
-/// client via the callback (if registered).
+/// Waits for incoming connections from VM. If a new connection is made, stores the stream in the
+/// corresponding `VmInstance`.
 fn handle_stream_connection_from_vm(state: Arc<Mutex<State>>) -> Result<()> {
     let listener =
         VsockListener::bind_with_cid_port(VMADDR_CID_HOST, VM_STREAM_SERVICE_PORT as u32)?;
@@ -572,10 +573,10 @@
         Ok(self.instance.cid as i32)
     }
 
-    fn isRunning(&self) -> binder::Result<bool> {
+    fn getState(&self) -> binder::Result<VirtualMachineState> {
         // Don't check permission. The owner of the VM might have passed this binder object to
         // others.
-        Ok(self.instance.running())
+        Ok(get_state(&self.instance))
     }
 
     fn registerCallback(
@@ -738,6 +739,20 @@
     option.as_ref().map(|t| t.as_ref())
 }
 
+/// Gets the `VirtualMachineState` of the given `VmInstance`.
+fn get_state(instance: &VmInstance) -> VirtualMachineState {
+    if instance.running() {
+        match instance.payload_state() {
+            PayloadState::Starting => VirtualMachineState::STARTING,
+            PayloadState::Started => VirtualMachineState::STARTED,
+            PayloadState::Ready => VirtualMachineState::READY,
+            PayloadState::Finished => VirtualMachineState::FINISHED,
+        }
+    } else {
+        VirtualMachineState::DEAD
+    }
+}
+
 /// Converts a `&ParcelFileDescriptor` to a `File` by cloning the file.
 fn clone_file(file: &ParcelFileDescriptor) -> Result<File, Status> {
     file.as_ref().try_clone().map_err(|e| {
@@ -789,6 +804,8 @@
         let cid = cid as Cid;
         if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
             info!("VM having CID {} started payload", cid);
+            vm.update_payload_state(PayloadState::Started)
+                .map_err(|e| new_binder_exception(ExceptionCode::ILLEGAL_STATE, e.to_string()))?;
             let stream = vm.stream.lock().unwrap().take();
             vm.callbacks.notify_payload_started(cid, stream);
             Ok(())
@@ -805,6 +822,8 @@
         let cid = cid as Cid;
         if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
             info!("VM having CID {} payload is ready", cid);
+            vm.update_payload_state(PayloadState::Ready)
+                .map_err(|e| new_binder_exception(ExceptionCode::ILLEGAL_STATE, e.to_string()))?;
             vm.callbacks.notify_payload_ready(cid);
             Ok(())
         } else {
@@ -820,6 +839,8 @@
         let cid = cid as Cid;
         if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
             info!("VM having CID {} finished payload", cid);
+            vm.update_payload_state(PayloadState::Finished)
+                .map_err(|e| new_binder_exception(ExceptionCode::ILLEGAL_STATE, e.to_string()))?;
             vm.callbacks.notify_payload_finished(cid, exit_code);
             Ok(())
         } else {
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 5984ff0..abaa955 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -52,6 +52,18 @@
     pub writable: bool,
 }
 
+/// The lifecycle state which the payload in the VM has reported itself to be in.
+///
+/// Note that the order of enum variants is significant; only forward transitions are allowed by
+/// [`VmInstance::update_payload_state`].
+#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
+pub enum PayloadState {
+    Starting,
+    Started,
+    Ready,
+    Finished,
+}
+
 /// Information about a particular instance of a VM which is running.
 #[derive(Debug)]
 pub struct VmInstance {
@@ -76,6 +88,8 @@
     pub callbacks: VirtualMachineCallbacks,
     /// Input/output stream of the payload run in the VM.
     pub stream: Mutex<Option<VsockStream>>,
+    /// The latest lifecycle state which the payload reported itself to be in.
+    payload_state: Mutex<PayloadState>,
 }
 
 impl VmInstance {
@@ -100,6 +114,7 @@
             running: AtomicBool::new(true),
             callbacks: Default::default(),
             stream: Mutex::new(None),
+            payload_state: Mutex::new(PayloadState::Starting),
         }
     }
 
@@ -154,6 +169,24 @@
         self.running.load(Ordering::Acquire)
     }
 
+    /// Returns the last reported state of the VM payload.
+    pub fn payload_state(&self) -> PayloadState {
+        *self.payload_state.lock().unwrap()
+    }
+
+    /// Updates the payload state to the given value, if it is a valid state transition.
+    pub fn update_payload_state(&self, new_state: PayloadState) -> Result<(), Error> {
+        let mut state_locked = self.payload_state.lock().unwrap();
+        // Only allow forward transitions, e.g. from starting to started or finished, not back in
+        // the other direction.
+        if new_state > *state_locked {
+            *state_locked = new_state;
+            Ok(())
+        } else {
+            bail!("Invalid payload state transition from {:?} to {:?}", *state_locked, new_state)
+        }
+    }
+
     /// Kill the crosvm instance.
     pub fn kill(&self) {
         // TODO: Talk to crosvm to shutdown cleanly.