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/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.