virtualizationservice: Unbind devices on VM exit

When configured to assign VFIO devices to a VM, virtualizationmanager
will bind them to the VFIO driver before starting the VM. This change
rebinds the default driver after VM exit, so the devices can be used by
the host again.

Test: Assign a device to VM, then check `<device_sysfs>/driver` post exit
Bug: 278008519
Change-Id: I48eee0b0dc1bba40a1c0c543d74f47ec4fee6d78
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> {