Merge "vmbase: Improve safety of console::init()" into main
diff --git a/docs/custom_vm.md b/docs/custom_vm.md
index 13ee620..6c51795 100644
--- a/docs/custom_vm.md
+++ b/docs/custom_vm.md
@@ -56,7 +56,13 @@
 
 ## Graphical VMs
 
-To run OSes with graphics support, follow the instruction below.
+To run OSes with graphics support, simply
+`packages/modules/Virtualization/tests/ferrochrome/ferrochrome.sh`. It prepares
+and launches the ChromiumOS, which is the only officially supported guest
+payload. We will be adding more OSes in the future.
+
+If you want to do so by yourself (e.g. boot with your build), follow the
+instruction below.
 
 ### Prepare a guest image
 
diff --git a/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java b/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
index d9e5229..7c18537 100644
--- a/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
+++ b/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
@@ -56,6 +56,7 @@
     private static final Path IMAGE_VERSION_INFO =
             Path.of(EXTERNAL_STORAGE_DIR + "ferrochrome_image_version");
     private static final Path VM_CONFIG_PATH = Path.of(EXTERNAL_STORAGE_DIR + "vm_config.json");
+    private static final int REQUEST_CODE_VMLAUNCHER = 1;
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -109,10 +110,17 @@
                     }
                     updateStatus("Done.");
                     updateStatus("Starting Ferrochrome...");
-                    runOnUiThread(() -> startActivity(intent));
+                    runOnUiThread(() -> startActivityForResult(intent, REQUEST_CODE_VMLAUNCHER));
                 });
     }
 
+    @Override
+    protected void onActivityResult(int requestCode, int resultCode, Intent data) {
+        if (requestCode == REQUEST_CODE_VMLAUNCHER) {
+            finishAndRemoveTask();
+        }
+    }
+
     private void updateStatus(String line) {
         Log.d(TAG, line);
         runOnUiThread(
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachine.java b/java/framework/src/android/system/virtualmachine/VirtualMachine.java
index bca36a4..f5ce102 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachine.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachine.java
@@ -172,6 +172,7 @@
     private ParcelFileDescriptor mTouchSock;
     private ParcelFileDescriptor mKeySock;
     private ParcelFileDescriptor mMouseSock;
+    private ParcelFileDescriptor mSwitchesSock;
 
     /**
      * Status of a virtual machine
@@ -921,6 +922,13 @@
                 m.pfd = pfds[1];
                 inputDevices.add(InputDevice.mouse(m));
             }
+            if (vmConfig.getCustomImageConfig().useSwitches()) {
+                ParcelFileDescriptor[] pfds = ParcelFileDescriptor.createSocketPair();
+                mSwitchesSock = pfds[0];
+                InputDevice.Switches s = new InputDevice.Switches();
+                s.pfd = pfds[1];
+                inputDevices.add(InputDevice.switches(s));
+            }
         }
         rawConfig.inputDevices = inputDevices.toArray(new InputDevice[0]);
 
@@ -1069,6 +1077,25 @@
                         new InputEvent(EV_SYN, SYN_REPORT, 0)));
     }
 
+    /** @hide */
+    public boolean sendLidEvent(boolean close) {
+        if (mSwitchesSock == null) {
+            Log.d(TAG, "mSwitcheSock == null");
+            return false;
+        }
+
+        // from include/uapi/linux/input-event-codes.h in the kernel.
+        short EV_SYN = 0x00;
+        short EV_SW = 0x05;
+        short SW_LID = 0x00;
+        short SYN_REPORT = 0x00;
+        return writeEventsToSock(
+                mSwitchesSock,
+                Arrays.asList(
+                        new InputEvent(EV_SW, SW_LID, close ? 1 : 0),
+                        new InputEvent(EV_SYN, SYN_REPORT, 0)));
+    }
+
     private boolean writeEventsToSock(ParcelFileDescriptor sock, List<InputEvent> evtList) {
         ByteBuffer byteBuffer =
                 ByteBuffer.allocate(8 /* (type: u16 + code: u16 + value: i32) */ * evtList.size());
@@ -1471,6 +1498,38 @@
         }
     }
 
+    /** @hide */
+    public void suspend() throws VirtualMachineException {
+        synchronized (mLock) {
+            if (mVirtualMachine == null) {
+                throw new VirtualMachineException("VM is not running");
+            }
+            try {
+                mVirtualMachine.suspend();
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            }
+        }
+    }
+
+    /** @hide */
+    public void resume() throws VirtualMachineException {
+        synchronized (mLock) {
+            if (mVirtualMachine == null) {
+                throw new VirtualMachineException("VM is not running");
+            }
+            try {
+                mVirtualMachine.resume();
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            }
+        }
+    }
+
     /**
      * Stops this virtual machine, if it is running.
      *
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
index 8d4886a..7fbfb33 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
@@ -36,6 +36,7 @@
     private static final String KEY_TOUCH = "touch";
     private static final String KEY_KEYBOARD = "keyboard";
     private static final String KEY_MOUSE = "mouse";
+    private static final String KEY_SWITCHES = "switches";
     private static final String KEY_NETWORK = "network";
     private static final String KEY_GPU = "gpu";
 
@@ -49,6 +50,7 @@
     private final boolean touch;
     private final boolean keyboard;
     private final boolean mouse;
+    private final boolean switches;
     private final boolean network;
     @Nullable private final GpuConfig gpuConfig;
 
@@ -94,6 +96,10 @@
         return mouse;
     }
 
+    public boolean useSwitches() {
+        return switches;
+    }
+
     public boolean useNetwork() {
         return network;
     }
@@ -110,6 +116,7 @@
             boolean touch,
             boolean keyboard,
             boolean mouse,
+            boolean switches,
             boolean network,
             GpuConfig gpuConfig) {
         this.name = name;
@@ -122,6 +129,7 @@
         this.touch = touch;
         this.keyboard = keyboard;
         this.mouse = mouse;
+        this.switches = switches;
         this.network = network;
         this.gpuConfig = gpuConfig;
     }
@@ -187,6 +195,7 @@
         pb.putBoolean(KEY_TOUCH, touch);
         pb.putBoolean(KEY_KEYBOARD, keyboard);
         pb.putBoolean(KEY_MOUSE, mouse);
+        pb.putBoolean(KEY_SWITCHES, switches);
         pb.putBoolean(KEY_NETWORK, network);
         pb.putPersistableBundle(
                 KEY_GPU,
@@ -247,6 +256,7 @@
         private boolean touch;
         private boolean keyboard;
         private boolean mouse;
+        private boolean switches;
         private boolean network;
         private GpuConfig gpuConfig;
 
@@ -320,6 +330,12 @@
         }
 
         /** @hide */
+        public Builder useSwitches(boolean switches) {
+            this.switches = switches;
+            return this;
+        }
+
+        /** @hide */
         public Builder useNetwork(boolean network) {
             this.network = network;
             return this;
@@ -338,6 +354,7 @@
                     touch,
                     keyboard,
                     mouse,
+                    switches,
                     network,
                     gpuConfig);
         }
diff --git a/pvmfw/avb/tests/api_test.rs b/pvmfw/avb/tests/api_test.rs
index c6f26ac..8683e69 100644
--- a/pvmfw/avb/tests/api_test.rs
+++ b/pvmfw/avb/tests/api_test.rs
@@ -20,7 +20,11 @@
 use avb::{DescriptorError, SlotVerifyError};
 use avb_bindgen::{AvbFooter, AvbVBMetaImageHeader};
 use pvmfw_avb::{verify_payload, Capability, DebugLevel, PvmfwVerifyError, VerifiedBootData};
-use std::{fs, mem::size_of, ptr};
+use std::{
+    fs,
+    mem::{offset_of, size_of},
+    ptr,
+};
 use utils::*;
 
 const TEST_IMG_WITH_ONE_HASHDESC_PATH: &str = "test_image_with_one_hashdesc.img";
@@ -243,32 +247,20 @@
 fn kernel_footer_with_vbmeta_offset_overwritten_fails_verification() -> Result<()> {
     // Arrange.
     let mut kernel = load_latest_signed_kernel()?;
-    let total_len = kernel.len() as u64;
-    let footer = extract_avb_footer(&kernel)?;
-    assert!(footer.vbmeta_offset < total_len);
-    // TODO: use core::mem::offset_of once stable.
-    let footer_addr = ptr::addr_of!(footer) as *const u8;
-    let vbmeta_offset_addr = ptr::addr_of!(footer.vbmeta_offset) as *const u8;
-    let vbmeta_offset_start =
-        // SAFETY:
-        // - both raw pointers `vbmeta_offset_addr` and `footer_addr` are not null;
-        // - they are both derived from the `footer` object;
-        // - the offset is known from the struct definition to be a small positive number of bytes.
-        unsafe { vbmeta_offset_addr.offset_from(footer_addr) };
-    let footer_start = kernel.len() - size_of::<AvbFooter>();
-    let vbmeta_offset_start = footer_start + usize::try_from(vbmeta_offset_start)?;
+    let footer_offset = get_avb_footer_offset(&kernel)?;
+    let vbmeta_offset_offset = footer_offset + offset_of!(AvbFooter, vbmeta_offset);
+    let vbmeta_offset_bytes = vbmeta_offset_offset..(vbmeta_offset_offset + size_of::<u64>());
 
-    let wrong_offsets = [total_len, u64::MAX];
-    for &wrong_offset in wrong_offsets.iter() {
+    let test_values = [kernel.len(), usize::MAX];
+    for value in test_values {
+        let value = u64::try_from(value).unwrap();
         // Act.
-        kernel[vbmeta_offset_start..(vbmeta_offset_start + size_of::<u64>())]
-            .copy_from_slice(&wrong_offset.to_be_bytes());
+        kernel[vbmeta_offset_bytes.clone()].copy_from_slice(&value.to_be_bytes());
+        // footer is unaligned; copy vbmeta_offset to local variable
+        let vbmeta_offset = extract_avb_footer(&kernel)?.vbmeta_offset;
+        assert_eq!(vbmeta_offset, value);
 
         // Assert.
-        let footer = extract_avb_footer(&kernel)?;
-        // footer is unaligned; copy vbmeta_offset to local variable
-        let vbmeta_offset = footer.vbmeta_offset;
-        assert_eq!(wrong_offset, vbmeta_offset);
         assert_payload_verification_with_initrd_fails(
             &kernel,
             &load_latest_initrd_normal()?,
diff --git a/pvmfw/avb/tests/utils.rs b/pvmfw/avb/tests/utils.rs
index cf37fcf..e989579 100644
--- a/pvmfw/avb/tests/utils.rs
+++ b/pvmfw/avb/tests/utils.rs
@@ -72,8 +72,14 @@
     Ok(fs::read(PUBLIC_KEY_RSA4096_PATH)?)
 }
 
+pub fn get_avb_footer_offset(signed_kernel: &[u8]) -> Result<usize> {
+    let offset = signed_kernel.len().checked_sub(size_of::<AvbFooter>());
+
+    offset.ok_or_else(|| anyhow!("Kernel too small to be AVB-signed"))
+}
+
 pub fn extract_avb_footer(kernel: &[u8]) -> Result<AvbFooter> {
-    let footer_start = kernel.len() - size_of::<AvbFooter>();
+    let footer_start = get_avb_footer_offset(kernel)?;
     // SAFETY: The slice is the same size as the struct which only contains simple data types.
     let mut footer = unsafe {
         transmute::<[u8; size_of::<AvbFooter>()], AvbFooter>(kernel[footer_start..].try_into()?)
diff --git a/rialto/Android.bp b/rialto/Android.bp
index 33fe189..b26a1c4 100644
--- a/rialto/Android.bp
+++ b/rialto/Android.bp
@@ -137,6 +137,7 @@
         "liblibc",
         "liblog_rust",
         "libhwtrust",
+        "libhypervisor_props",
         "libservice_vm_comm",
         "libservice_vm_fake_chain",
         "libservice_vm_manager",
diff --git a/rialto/tests/test.rs b/rialto/tests/test.rs
index 0d57301..cf5630f 100644
--- a/rialto/tests/test.rs
+++ b/rialto/tests/test.rs
@@ -26,7 +26,7 @@
 use client_vm_csr::generate_attestation_key_and_csr;
 use coset::{CborSerializable, CoseMac0, CoseSign};
 use hwtrust::{rkp, session::Session};
-use log::info;
+use log::{info, warn};
 use service_vm_comm::{
     ClientVmAttestationParams, Csr, CsrPayload, EcdsaP256KeyPair, GenerateCertificateRequestParams,
     Request, RequestProcessingError, Response, VmType,
@@ -55,10 +55,15 @@
 #[cfg(dice_changes)]
 #[test]
 fn process_requests_in_protected_vm() -> Result<()> {
-    // The test is skipped if the feature flag |dice_changes| is not enabled, because when
-    // the flag is off, the DICE chain is truncated in the pvmfw, and the service VM cannot
-    // verify the chain due to the missing entries in the chain.
-    check_processing_requests(VmType::ProtectedVm)
+    if hypervisor_props::is_protected_vm_supported()? {
+        // The test is skipped if the feature flag |dice_changes| is not enabled, because when
+        // the flag is off, the DICE chain is truncated in the pvmfw, and the service VM cannot
+        // verify the chain due to the missing entries in the chain.
+        check_processing_requests(VmType::ProtectedVm)
+    } else {
+        warn!("pVMs are not supported on device, skipping test");
+        Ok(())
+    }
 }
 
 #[test]
diff --git a/tests/ferrochrome/AndroidTest.xml b/tests/ferrochrome/AndroidTest.xml
index 79cbe72..8053674 100644
--- a/tests/ferrochrome/AndroidTest.xml
+++ b/tests/ferrochrome/AndroidTest.xml
@@ -32,9 +32,9 @@
          It's too large (6.5G+), so this may break further tests. -->
     <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
         <option name="throw-if-cmd-fail" value="false" />
-        <option name="run-command" value="mkdir /data/local/tmp/ferrochrome" />
+        <option name="run-command" value="mkdir /data/local/tmp" />
         <option name="teardown-command" value="pkill vmlauncher" />
-        <option name="teardown-command" value="rm -rf /data/local/tmp/ferrochrome" />
+        <option name="teardown-command" value="rm /data/local/tmp/chromiumos_test_image.bin" />
     </target_preparer>
 
     <test class="com.android.tradefed.testtype.binary.ExecutableHostTest" >
diff --git a/tests/ferrochrome/ferrochrome.sh b/tests/ferrochrome/ferrochrome.sh
index 6814ac5..210548a 100755
--- a/tests/ferrochrome/ferrochrome.sh
+++ b/tests/ferrochrome/ferrochrome.sh
@@ -21,12 +21,13 @@
 
 FECR_GS_URL="https://storage.googleapis.com/chromiumos-image-archive/ferrochrome-public"
 FECR_DEFAULT_VERSION="R127-15916.0.0"
-FECR_DEVICE_DIR="/data/local/tmp/ferrochrome"
+FECR_DEVICE_DIR="/data/local/tmp"
 FECR_CONFIG_PATH="/data/local/tmp/vm_config.json"  # hardcoded at VmLauncherApp
 FECR_CONSOLE_LOG_PATH="/data/data/\${pkg_name}/files/console.log"
 FECR_BOOT_COMPLETED_LOG="Have fun and send patches!"
 FECR_BOOT_TIMEOUT="300" # 5 minutes (300 seconds)
 ACTION_NAME="android.virtualization.VM_LAUNCHER"
+TRY_UNLOCK_MAX=10
 
 fecr_clean_up() {
   trap - INT
@@ -132,6 +133,35 @@
   adb push ${fecr_script_path}/assets/vm_config.json ${FECR_CONFIG_PATH}
 fi
 
+echo "Ensure screen unlocked"
+
+try_unlock=0
+while [[ "${try_unlock}" -le "${TRY_UNLOCK_MAX}" ]]; do
+  screen_state=$(adb shell dumpsys nfc | sed -n 's/^mScreenState=\(.*\)$/\1/p')
+  case "${screen_state}" in
+    "ON_UNLOCKED")
+      break
+      ;;
+    "ON_LOCKED")
+      # Disclaimer: This can unlock phone only if unlock method is swipe (default after FDR)
+      adb shell input keyevent KEYCODE_MENU
+      ;;
+    "OFF_LOCKED"|"OFF_UNLOCKED")
+      adb shell input keyevent KEYCODE_WAKEUP
+      ;;
+    *)
+      echo "Unknown screen state. Continue to boot, but may fail"
+      break
+      ;;
+  esac
+  sleep 1
+  try_unlock=$((try_unlock+1))
+done
+if [[ "${try_unlock}" -gt "${TRY_UNLOCK_MAX}" ]]; then
+  >&2 echo "Failed to unlock screen. Try again after manual unlock"
+  exit 1
+fi
+
 echo "Starting ferrochrome"
 adb shell am start-activity -a ${ACTION_NAME} > /dev/null
 
diff --git a/virtualizationmanager/Android.bp b/virtualizationmanager/Android.bp
index d1ef4de..ada66dd 100644
--- a/virtualizationmanager/Android.bp
+++ b/virtualizationmanager/Android.bp
@@ -70,6 +70,7 @@
         "libvsock",
         "liblibfdt",
         "libfsfdt",
+        "libhypervisor_props",
         // TODO(b/202115393) stabilize the interface
         "packagemanager_aidl-rust",
     ],
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 9df376a..10dafdf 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -465,9 +465,12 @@
         let kernel = maybe_clone_file(&config.kernel)?;
         let initrd = maybe_clone_file(&config.initrd)?;
 
-        // In a protected VM, we require custom kernels to come from a trusted source (b/237054515).
         if config.protectedVm {
+            // In a protected VM, we require custom kernels to come from a trusted source
+            // (b/237054515).
             check_label_for_kernel_files(&kernel, &initrd).or_service_specific_exception(-1)?;
+            // Fail fast with a meaningful error message in case device doesn't support pVMs.
+            check_protected_vm_is_supported()?;
         }
 
         let zero_filler_path = temporary_directory.join("zero.img");
@@ -798,6 +801,9 @@
         InputDevice::Mouse(mouse) => InputDeviceOption::Mouse(clone_file(
             mouse.pfd.as_ref().ok_or(anyhow!("pfd should have value"))?,
         )?),
+        InputDevice::Switches(switches) => InputDeviceOption::Switches(clone_file(
+            switches.pfd.as_ref().ok_or(anyhow!("pfd should have value"))?,
+        )?),
     })
 }
 /// Given the configuration for a disk image, assembles the `DiskFile` to pass to crosvm.
@@ -1244,6 +1250,22 @@
     fn setHostConsoleName(&self, ptsname: &str) -> binder::Result<()> {
         self.instance.vm_context.global_context.setHostConsoleName(ptsname)
     }
+
+    fn suspend(&self) -> binder::Result<()> {
+        self.instance
+            .suspend()
+            .with_context(|| format!("Error suspending VM with CID {}", self.instance.cid))
+            .with_log()
+            .or_service_specific_exception(-1)
+    }
+
+    fn resume(&self) -> binder::Result<()> {
+        self.instance
+            .resume()
+            .with_context(|| format!("Error resuming VM with CID {}", self.instance.cid))
+            .with_log()
+            .or_service_specific_exception(-1)
+    }
 }
 
 impl Drop for VirtualMachine {
@@ -1486,6 +1508,17 @@
     Ok(())
 }
 
+fn check_protected_vm_is_supported() -> binder::Result<()> {
+    let is_pvm_supported =
+        hypervisor_props::is_protected_vm_supported().or_service_specific_exception(-1)?;
+    if is_pvm_supported {
+        Ok(())
+    } else {
+        Err(anyhow!("pVM is not supported"))
+            .or_binder_exception(ExceptionCode::UNSUPPORTED_OPERATION)
+    }
+}
+
 fn check_config_features(config: &VirtualMachineConfig) -> binder::Result<()> {
     if !cfg!(vendor_modules) {
         check_no_vendor_modules(config)?;
diff --git a/virtualizationmanager/src/crosvm.rs b/virtualizationmanager/src/crosvm.rs
index 47ef91a..13c018b 100644
--- a/virtualizationmanager/src/crosvm.rs
+++ b/virtualizationmanager/src/crosvm.rs
@@ -208,6 +208,7 @@
     SingleTouch { file: File, width: u32, height: u32, name: Option<String> },
     Keyboard(File),
     Mouse(File),
+    Switches(File),
 }
 
 type VfioDevice = Strong<dyn IBoundDevice>;
@@ -682,6 +683,28 @@
         conn.notify_completion()?;
         Ok(())
     }
+
+    /// Suspends the VM
+    pub fn suspend(&self) -> Result<(), Error> {
+        match vm_control::client::handle_request(
+            &VmRequest::SuspendVcpus,
+            &self.crosvm_control_socket_path,
+        ) {
+            Ok(VmResponse::Ok) => Ok(()),
+            e => bail!("Failed to suspend VM: {e:?}"),
+        }
+    }
+
+    /// Resumes the suspended VM
+    pub fn resume(&self) -> Result<(), Error> {
+        match vm_control::client::handle_request(
+            &VmRequest::ResumeVcpus,
+            &self.crosvm_control_socket_path,
+        ) {
+            Ok(VmResponse::Ok) => Ok(()),
+            e => bail!("Failed to resume: {e:?}"),
+        }
+    }
 }
 
 impl Rss {
@@ -1115,6 +1138,9 @@
                     height,
                     name.as_ref().map_or("".into(), |n| format!(",name={}", n))
                 ),
+                InputDeviceOption::Switches(file) => {
+                    format!("switches[path={}]", add_preserved_fd(&mut preserved_fds, file))
+                }
             });
         }
     }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
index d4001c8..9d1d5d5 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
@@ -50,4 +50,10 @@
 
     /** Set the name of the peer end (ptsname) of the host console. */
     void setHostConsoleName(in @utf8InCpp String pathname);
+
+    /** Suspends the VM. */
+    void suspend();
+
+    /** Resumes the suspended VM. */
+    void resume();
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/InputDevice.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/InputDevice.aidl
index 56c5b6d..5a7ed4a 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/InputDevice.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/InputDevice.aidl
@@ -38,8 +38,15 @@
     parcelable Mouse {
         ParcelFileDescriptor pfd;
     }
+
+    // Switches input
+    parcelable Switches {
+        ParcelFileDescriptor pfd;
+    }
+
     SingleTouch singleTouch;
     EvDev evDev;
     Keyboard keyboard;
     Mouse mouse;
+    Switches switches;
 }
diff --git a/vmbase/src/bionic.rs b/vmbase/src/bionic.rs
index 72f86c4..6ea8d60 100644
--- a/vmbase/src/bionic.rs
+++ b/vmbase/src/bionic.rs
@@ -111,7 +111,7 @@
 ///
 /// # Note
 ///
-/// This Rust functions is missing the last argument of its C/C++ counterpart, a va_list.
+/// This Rust function is missing the last argument of its C/C++ counterpart, a va_list.
 #[no_mangle]
 unsafe extern "C" fn async_safe_fatal_va_list(prefix: *const c_char, format: *const c_char) {
     // SAFETY: The caller guaranteed that both strings were valid and NUL-terminated.
diff --git a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
index 0be2e57..e13d2c9 100644
--- a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
+++ b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
@@ -188,6 +188,7 @@
             customImageConfigBuilder.useTouch(true);
             customImageConfigBuilder.useKeyboard(true);
             customImageConfigBuilder.useMouse(true);
+            customImageConfigBuilder.useSwitches(true);
             customImageConfigBuilder.useNetwork(true);
 
             configBuilder.setCustomImageConfig(customImageConfigBuilder.build());
@@ -241,41 +242,34 @@
 
                     @Override
                     public void onPayloadStarted(VirtualMachine vm) {
-                        Log.e(TAG, "payload start");
+                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
+                        // this.
                     }
 
                     @Override
                     public void onPayloadReady(VirtualMachine vm) {
-                        // This check doesn't 100% prevent race condition or UI hang.
-                        // However, it's fine for demo.
-                        if (mService.isShutdown()) {
-                            return;
-                        }
-                        Log.d(TAG, "(Payload is ready. Testing VM service...)");
+                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
+                        // this.
                     }
 
                     @Override
                     public void onPayloadFinished(VirtualMachine vm, int exitCode) {
-                        // This check doesn't 100% prevent race condition, but is fine for demo.
-                        if (!mService.isShutdown()) {
-                            Log.d(
-                                    TAG,
-                                    String.format("(Payload finished. exit code: %d)", exitCode));
-                        }
+                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
+                        // this.
                     }
 
                     @Override
                     public void onError(VirtualMachine vm, int errorCode, String message) {
-                        Log.d(
-                                TAG,
-                                String.format(
-                                        "(Error occurred. code: %d, message: %s)",
-                                        errorCode, message));
+                        Log.e(TAG, "Error from VM. code: " + errorCode + " (" + message + ")");
+                        setResult(RESULT_CANCELED);
+                        finish();
                     }
 
                     @Override
                     public void onStopped(VirtualMachine vm, int reason) {
-                        Log.e(TAG, "vm stop");
+                        Log.d(TAG, "VM stopped. Reason: " + reason);
+                        setResult(RESULT_OK);
+                        finish();
                     }
                 };
 
@@ -422,6 +416,32 @@
     }
 
     @Override
+    protected void onStop() {
+        super.onStop();
+        if (mVirtualMachine != null) {
+            try {
+                mVirtualMachine.sendLidEvent(/* close */ true);
+                mVirtualMachine.suspend();
+            } catch (VirtualMachineException e) {
+                Log.e(TAG, "Failed to suspend VM" + e);
+            }
+        }
+    }
+
+    @Override
+    protected void onRestart() {
+        super.onRestart();
+        if (mVirtualMachine != null) {
+            try {
+                mVirtualMachine.resume();
+                mVirtualMachine.sendLidEvent(/* close */ false);
+            } catch (VirtualMachineException e) {
+                Log.e(TAG, "Failed to resume VM" + e);
+            }
+        }
+    }
+
+    @Override
     protected void onDestroy() {
         super.onDestroy();
         if (mExecutorService != null) {