Merge "virtualizationservice: Unbind devices on VM exit" into main
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 8c2099f..d775555 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -18,7 +18,7 @@
 use crate::atom::{
     write_vm_booted_stats, write_vm_creation_stats};
 use crate::composite::make_composite_image;
-use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VfioDevice, VmContext, VmInstance, VmState};
+use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmContext, VmInstance, VmState};
 use crate::debug_config::DebugConfig;
 use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images, add_microdroid_vendor_image};
 use crate::selinux::{getfilecon, SeContext};
@@ -510,14 +510,7 @@
                         .or_binder_exception(ExceptionCode::ILLEGAL_ARGUMENT);
                 }
             }
-            let devices = GLOBAL_SERVICE
-                .bindDevicesToVfioDriver(&config.devices)?
-                .into_iter()
-                .map(|x| VfioDevice {
-                    sysfs_path: PathBuf::from(&x.sysfsPath),
-                    dtbo_label: x.dtboLabel,
-                })
-                .collect::<Vec<_>>();
+            let devices = GLOBAL_SERVICE.bindDevicesToVfioDriver(&config.devices)?;
             let dtbo_file = File::from(
                 GLOBAL_SERVICE
                     .getDtboFile()?
diff --git a/virtualizationmanager/src/crosvm.rs b/virtualizationmanager/src/crosvm.rs
index f0c3e4b..4b3478e 100644
--- a/virtualizationmanager/src/crosvm.rs
+++ b/virtualizationmanager/src/crosvm.rs
@@ -47,6 +47,7 @@
     VirtualMachineAppConfig::DebugLevel::DebugLevel
 };
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IGlobalVmContext::IGlobalVmContext;
+use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IBoundDevice::IBoundDevice;
 use binder::Strong;
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
 use tombstoned_client::{TombstonedConnection, DebuggerdDumpType};
@@ -127,11 +128,7 @@
     pub writable: bool,
 }
 
-#[derive(Clone, Debug)]
-pub struct VfioDevice {
-    pub sysfs_path: PathBuf,
-    pub dtbo_label: String,
-}
+type VfioDevice = Strong<dyn IBoundDevice>;
 
 /// The lifecycle state which the payload in the VM has reported itself to be in.
 ///
@@ -412,10 +409,7 @@
             error!("Error removing temporary files from {:?}: {}", self.temporary_directory, e);
         });
 
-        // TODO(b/278008182): clean up assigned devices.
-        for device in vfio_devices.iter() {
-            info!("NOT RELEASING {device:?}");
-        }
+        drop(vfio_devices); // Cleanup devices.
     }
 
     /// Waits until payload is started, or timeout expires. When timeout occurs, kill
@@ -706,7 +700,7 @@
 
 fn vfio_argument_for_platform_device(device: &VfioDevice) -> Result<String, Error> {
     // Check platform device exists
-    let path = device.sysfs_path.canonicalize()?;
+    let path = Path::new(&device.getSysfsPath()?).canonicalize()?;
     if !path.starts_with(SYSFS_PLATFORM_DEVICES_PATH) {
         bail!("{path:?} is not a platform device");
     }
@@ -718,7 +712,7 @@
     }
 
     if let Some(p) = path.to_str() {
-        Ok(format!("--vfio={p},iommu=pkvm-iommu,dt-symbol={0}", device.dtbo_label))
+        Ok(format!("--vfio={p},iommu=pkvm-iommu,dt-symbol={0}", device.getDtboLabel()?))
     } else {
         bail!("invalid path {path:?}");
     }
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 3f8d193..e89ce44 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -32,6 +32,7 @@
         "libavflog",
         "libbinder_rs",
         "libhypervisor_props",
+        "liblazy_static",
         "liblibc",
         "liblog_rust",
         "libnix",
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IBoundDevice.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IBoundDevice.aidl
new file mode 100644
index 0000000..4a37bf7
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IBoundDevice.aidl
@@ -0,0 +1,25 @@
+/*
+ * Copyright 2023 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_internal;
+
+/** A device bound to VFIO driver. */
+interface IBoundDevice {
+    /** Path to SysFS node of the device. */
+    String getSysfsPath();
+
+    /** DTBO label of the device. */
+    String getDtboLabel();
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVfioHandler.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVfioHandler.aidl
index 01906cb..2cf4efd 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVfioHandler.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVfioHandler.aidl
@@ -20,16 +20,22 @@
 import android.system.virtualizationservice_internal.AtomVmBooted;
 import android.system.virtualizationservice_internal.AtomVmCreationRequested;
 import android.system.virtualizationservice_internal.AtomVmExited;
+import android.system.virtualizationservice_internal.IBoundDevice;
 import android.system.virtualizationservice_internal.IGlobalVmContext;
 
 /** VFIO related methods which should be done as root. */
 interface IVfioHandler {
+    parcelable VfioDev {
+        String sysfsPath;
+        String dtboLabel;
+    }
     /**
      * Bind given devices to vfio driver.
      *
-     * @param devices paths of sysfs nodes of devices to assign.
+     * @param devices a list of pairs (sysfs path, DTBO node label) for devices.
+     * @return IBoundDevice list representing a VFIO bound devices.
      */
-    void bindDevicesToVfioDriver(in String[] devices);
+    IBoundDevice[] bindDevicesToVfioDriver(in VfioDev[] devices);
 
     /**
      * Store VM DTBO via the file descriptor.
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index a2cb693..dd94526 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -21,13 +21,10 @@
 import android.system.virtualizationservice_internal.AtomVmBooted;
 import android.system.virtualizationservice_internal.AtomVmCreationRequested;
 import android.system.virtualizationservice_internal.AtomVmExited;
+import android.system.virtualizationservice_internal.IBoundDevice;
 import android.system.virtualizationservice_internal.IGlobalVmContext;
 
 interface IVirtualizationServiceInternal {
-    parcelable BoundDevice {
-        String sysfsPath;
-        String dtboLabel;
-    }
     /**
      * Removes the memlock rlimit of the calling process.
      *
@@ -78,9 +75,9 @@
      * Bind given devices to vfio driver.
      *
      * @param devices paths of sysfs nodes of devices to assign.
-     * @return a list of pairs (sysfs path, DTBO node label) for devices.
+     * @return a list of IBoundDevices representing VFIO bound devices.
      */
-    BoundDevice[] bindDevicesToVfioDriver(in String[] devices);
+    IBoundDevice[] bindDevicesToVfioDriver(in String[] devices);
 
     /** Returns a read-only file descriptor of the VM DTBO file. */
     ParcelFileDescriptor getDtboFile();
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 3ac1e60..80d4725 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -28,15 +28,17 @@
     AtomVmBooted::AtomVmBooted,
     AtomVmCreationRequested::AtomVmCreationRequested,
     AtomVmExited::AtomVmExited,
+    IBoundDevice::IBoundDevice,
     IGlobalVmContext::{BnGlobalVmContext, IGlobalVmContext},
-    IVirtualizationServiceInternal::BoundDevice::BoundDevice,
     IVirtualizationServiceInternal::IVirtualizationServiceInternal,
     IVfioHandler::{BpVfioHandler, IVfioHandler},
+    IVfioHandler::VfioDev::VfioDev,
 };
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::VM_TOMBSTONES_SERVICE_PORT;
 use anyhow::{anyhow, ensure, Context, Result};
 use avflog::LogResult;
 use binder::{self, wait_for_interface, BinderFeatures, ExceptionCode, Interface, LazyServiceGuard, Status, Strong, IntoBinderResult};
+use lazy_static::lazy_static;
 use libc::VMADDR_CID_HOST;
 use log::{error, info, warn};
 use rkpd_client::get_rkpd_attestation_key;
@@ -71,6 +73,12 @@
 
 const CHUNK_RECV_MAX_LEN: usize = 1024;
 
+lazy_static! {
+    static ref VFIO_SERVICE: Strong<dyn IVfioHandler> =
+        wait_for_interface(<BpVfioHandler as IVfioHandler>::get_descriptor())
+            .expect("Could not connect to VfioHandler");
+}
+
 fn is_valid_guest_cid(cid: Cid) -> bool {
     (GUEST_CID_MIN..=GUEST_CID_MAX).contains(&cid)
 }
@@ -219,24 +227,26 @@
             .collect::<Vec<_>>())
     }
 
-    fn bindDevicesToVfioDriver(&self, devices: &[String]) -> binder::Result<Vec<BoundDevice>> {
+    fn bindDevicesToVfioDriver(
+        &self,
+        devices: &[String],
+    ) -> binder::Result<Vec<Strong<dyn IBoundDevice>>> {
         check_use_custom_virtual_machine()?;
 
-        let vfio_service: Strong<dyn IVfioHandler> =
-            wait_for_interface(<BpVfioHandler as IVfioHandler>::get_descriptor())?;
-        vfio_service.bindDevicesToVfioDriver(devices)?;
-
-        Ok(get_assignable_devices()?
+        let devices = get_assignable_devices()?
             .device
             .into_iter()
             .filter_map(|x| {
                 if devices.contains(&x.sysfs_path) {
-                    Some(BoundDevice { sysfsPath: x.sysfs_path, dtboLabel: x.dtbo_label })
+                    Some(VfioDev { sysfsPath: x.sysfs_path, dtboLabel: x.dtbo_label })
                 } else {
+                    warn!("device {} is not assignable", x.sysfs_path);
                     None
                 }
             })
-            .collect::<Vec<_>>())
+            .collect::<Vec<VfioDev>>();
+
+        VFIO_SERVICE.bindDevicesToVfioDriver(devices.as_slice())
     }
 
     fn getDtboFile(&self) -> binder::Result<ParcelFileDescriptor> {
@@ -424,10 +434,7 @@
 
             // Open a write-only file descriptor for vfio_handler.
             let write_fd = File::create(&path).context("Failed to create VM DTBO file")?;
-
-            let vfio_service: Strong<dyn IVfioHandler> =
-                wait_for_interface(<BpVfioHandler as IVfioHandler>::get_descriptor())?;
-            vfio_service.writeVmDtbo(&ParcelFileDescriptor::new(write_fd))?;
+            VFIO_SERVICE.writeVmDtbo(&ParcelFileDescriptor::new(write_fd))?;
 
             // Open read-only. This FD will be cached and returned to clients.
             let read_fd = File::open(&path).context("Failed to open VM DTBO file")?;
diff --git a/virtualizationservice/vfio_handler/src/aidl.rs b/virtualizationservice/vfio_handler/src/aidl.rs
index 63f19c6..c0967af 100644
--- a/virtualizationservice/vfio_handler/src/aidl.rs
+++ b/virtualizationservice/vfio_handler/src/aidl.rs
@@ -15,10 +15,13 @@
 //! Implementation of the AIDL interface of the VirtualizationService.
 
 use anyhow::{anyhow, Context};
+use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IBoundDevice::{IBoundDevice, BnBoundDevice};
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IVfioHandler::IVfioHandler;
+use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IVfioHandler::VfioDev::VfioDev;
 use android_system_virtualizationservice_internal::binder::ParcelFileDescriptor;
-use binder::{self, ExceptionCode, Interface, IntoBinderResult};
+use binder::{self, BinderFeatures, ExceptionCode, Interface, IntoBinderResult, Strong};
 use lazy_static::lazy_static;
+use log::error;
 use std::fs::{read_link, write, File};
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::mem::size_of;
@@ -30,6 +33,38 @@
     FromBytes,
 };
 
+// Device bound to VFIO driver.
+struct BoundDevice {
+    sysfs_path: String,
+    dtbo_label: String,
+}
+
+impl Interface for BoundDevice {}
+
+impl IBoundDevice for BoundDevice {
+    fn getSysfsPath(&self) -> binder::Result<String> {
+        Ok(self.sysfs_path.clone())
+    }
+
+    fn getDtboLabel(&self) -> binder::Result<String> {
+        Ok(self.dtbo_label.clone())
+    }
+}
+
+impl Drop for BoundDevice {
+    fn drop(&mut self) {
+        unbind_device(Path::new(&self.sysfs_path)).unwrap_or_else(|e| {
+            error!("did not restore {} driver: {}", self.sysfs_path, e);
+        });
+    }
+}
+
+impl BoundDevice {
+    fn new_binder(sysfs_path: String, dtbo_label: String) -> Strong<dyn IBoundDevice> {
+        BnBoundDevice::new_binder(BoundDevice { sysfs_path, dtbo_label }, BinderFeatures::default())
+    }
+}
+
 #[derive(Debug, Default)]
 pub struct VfioHandler {}
 
@@ -42,14 +77,22 @@
 impl Interface for VfioHandler {}
 
 impl IVfioHandler for VfioHandler {
-    fn bindDevicesToVfioDriver(&self, devices: &[String]) -> binder::Result<()> {
+    fn bindDevicesToVfioDriver(
+        &self,
+        devices: &[VfioDev],
+    ) -> binder::Result<Vec<Strong<dyn IBoundDevice>>> {
         // permission check is already done by IVirtualizationServiceInternal.
         if !*IS_VFIO_SUPPORTED {
             return Err(anyhow!("VFIO-platform not supported"))
                 .or_binder_exception(ExceptionCode::UNSUPPORTED_OPERATION);
         }
-        devices.iter().try_for_each(|x| bind_device(Path::new(x)))?;
-        Ok(())
+        devices
+            .iter()
+            .map(|d| {
+                bind_device(Path::new(&d.sysfsPath))?;
+                Ok(BoundDevice::new_binder(d.sysfsPath.clone(), d.dtboLabel.clone()))
+            })
+            .collect::<binder::Result<Vec<_>>>()
     }
 
     fn writeVmDtbo(&self, dtbo_fd: &ParcelFileDescriptor) -> binder::Result<()> {
@@ -79,6 +122,11 @@
 const VFIO_PLATFORM_DRIVER_PATH: &str = "/sys/bus/platform/drivers/vfio-platform";
 const SYSFS_PLATFORM_DRIVERS_PROBE_PATH: &str = "/sys/bus/platform/drivers_probe";
 const DT_TABLE_MAGIC: u32 = 0xd7b7ab1e;
+const VFIO_PLATFORM_DRIVER_NAME: &str = "vfio-platform";
+// To remove the override and match the device driver by "compatible" string again,
+// driver_override file must be cleared. Writing an empty string (same as
+// `echo -n "" > driver_override`) won't' clear the file, so append a newline char.
+const DEFAULT_DRIVER: &str = "\n";
 
 /// The structure of DT table header in dtbo.img.
 /// https://source.android.com/docs/core/architecture/dto/partitions
@@ -146,18 +194,15 @@
     group.to_str()?.parse().ok()
 }
 
-fn is_bound_to_vfio_driver(path: &Path) -> bool {
-    let Ok(driver_path) = read_link(path.join("driver")) else {
-        return false;
-    };
-    let Some(driver) = driver_path.file_name() else {
-        return false;
-    };
-    driver.to_str().unwrap_or("") == "vfio-platform"
+fn current_driver(path: &Path) -> Option<String> {
+    let driver_path = read_link(path.join("driver")).ok()?;
+    let bound_driver = driver_path.file_name()?;
+    bound_driver.to_str().map(str::to_string)
 }
 
-fn bind_vfio_driver(path: &Path) -> binder::Result<()> {
-    if is_bound_to_vfio_driver(path) {
+// Try to bind device driver by writing its name to driver_override and triggering driver probe.
+fn try_bind_driver(path: &Path, driver: &str) -> binder::Result<()> {
+    if Some(driver) == current_driver(path).as_deref() {
         // already bound
         return Ok(());
     }
@@ -177,10 +222,13 @@
             .with_context(|| format!("could not unbind {device_str}"))
             .or_service_specific_exception(-1)?;
     }
+    if path.join("driver").exists() {
+        return Err(anyhow!("could not unbind {device_str}")).or_service_specific_exception(-1);
+    }
 
-    // bind to VFIO
-    write(path.join("driver_override"), b"vfio-platform")
-        .with_context(|| format!("could not bind {device_str} to vfio-platform"))
+    // bind to new driver
+    write(path.join("driver_override"), driver.as_bytes())
+        .with_context(|| format!("could not bind {device_str} to '{driver}' driver"))
         .or_service_specific_exception(-1)?;
 
     write(SYSFS_PLATFORM_DRIVERS_PROBE_PATH, device_str.as_bytes())
@@ -188,13 +236,9 @@
         .or_service_specific_exception(-1)?;
 
     // final check
-    if !is_bound_to_vfio_driver(path) {
-        return Err(anyhow!("{path:?} still not bound to vfio driver"))
-            .or_service_specific_exception(-1);
-    }
-
-    if get_device_iommu_group(path).is_none() {
-        return Err(anyhow!("can't get iommu group for {path:?}"))
+    let new_driver = current_driver(path);
+    if new_driver.is_none() || Some(driver) != new_driver.as_deref() && driver != DEFAULT_DRIVER {
+        return Err(anyhow!("{path:?} still not bound to '{driver}' driver"))
             .or_service_specific_exception(-1);
     }
 
@@ -208,7 +252,29 @@
         .or_binder_exception(ExceptionCode::ILLEGAL_ARGUMENT)?;
 
     check_platform_device(&path)?;
-    bind_vfio_driver(&path)
+    try_bind_driver(&path, VFIO_PLATFORM_DRIVER_NAME)?;
+
+    if get_device_iommu_group(&path).is_none() {
+        Err(anyhow!("can't get iommu group for {path:?}")).or_service_specific_exception(-1)
+    } else {
+        Ok(())
+    }
+}
+
+fn unbind_device(path: &Path) -> binder::Result<()> {
+    let path = path
+        .canonicalize()
+        .with_context(|| format!("can't canonicalize {path:?}"))
+        .or_binder_exception(ExceptionCode::ILLEGAL_ARGUMENT)?;
+
+    check_platform_device(&path)?;
+    try_bind_driver(&path, DEFAULT_DRIVER)?;
+
+    if Some(VFIO_PLATFORM_DRIVER_NAME) == current_driver(&path).as_deref() {
+        Err(anyhow!("{path:?} still bound to vfio driver")).or_service_specific_exception(-1)
+    } else {
+        Ok(())
+    }
 }
 
 fn get_dtbo_img_path() -> binder::Result<PathBuf> {