Merge "[instanceId] allocate instance ID for the Service VM" into main
diff --git a/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 41d244d..068d8f9 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -38,8 +38,6 @@
         "lz4",
         "sign_virt_apex",
         "simg2img",
-        "dtdiff",
-        "dtc", // for dtdiff
     ],
     // java_test_host doesn't have data_native_libs but jni_libs can be used to put
     // native modules under ./lib directory.
@@ -54,6 +52,5 @@
         "liblp",
         "libsparse",
         "libz",
-        "libfdt", // for dtc
     ],
 }
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
index 14cc0ae..007f38c 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
@@ -35,6 +35,7 @@
 import com.android.tradefed.util.RunUtil;
 
 import org.json.JSONArray;
+import org.json.JSONObject;
 
 import java.io.File;
 import java.io.IOException;
@@ -65,6 +66,17 @@
     protected static final Set<String> SUPPORTED_GKI_VERSIONS =
             Collections.unmodifiableSet(new HashSet(Arrays.asList("android14-6.1")));
 
+    /* Keep this sync with AssignableDevice.aidl */
+    public static final class AssignableDevice {
+        public final String node;
+        public final String dtbo_label;
+
+        public AssignableDevice(String node, String dtbo_label) {
+            this.node = node;
+            this.dtbo_label = dtbo_label;
+        }
+    }
+
     public static void prepareVirtualizationTestSetup(ITestDevice androidDevice)
             throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(androidDevice);
@@ -186,18 +198,26 @@
         return pathLine.substring("package:".length());
     }
 
-    public List<String> parseStringArrayFieldsFromVmInfo(String header) throws Exception {
+    public String parseFieldFromVmInfo(String header) throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
         String result = android.run("/apex/com.android.virt/bin/vm", "info");
-        List<String> ret = new ArrayList<>();
         for (String line : result.split("\n")) {
             if (!line.startsWith(header)) continue;
 
-            JSONArray jsonArray = new JSONArray(line.substring(header.length()));
+            return line.substring(header.length());
+        }
+        return "";
+    }
+
+    public List<String> parseStringArrayFieldsFromVmInfo(String header) throws Exception {
+        String field = parseFieldFromVmInfo(header);
+
+        List<String> ret = new ArrayList<>();
+        if (!field.isEmpty()) {
+            JSONArray jsonArray = new JSONArray(field);
             for (int i = 0; i < jsonArray.length(); i++) {
                 ret.add(jsonArray.getString(i));
             }
-            break;
         }
         return ret;
     }
@@ -208,8 +228,20 @@
         return result.contains("enabled");
     }
 
-    public List<String> getAssignableDevices() throws Exception {
-        return parseStringArrayFieldsFromVmInfo("Assignable devices: ");
+    public List<AssignableDevice> getAssignableDevices() throws Exception {
+        String field = parseFieldFromVmInfo("Assignable devices: ");
+
+        List<AssignableDevice> ret = new ArrayList<>();
+        if (!field.isEmpty()) {
+            JSONArray jsonArray = new JSONArray(field);
+            for (int i = 0; i < jsonArray.length(); i++) {
+                JSONObject jsonObject = jsonArray.getJSONObject(i);
+                ret.add(
+                        new AssignableDevice(
+                                jsonObject.getString("node"), jsonObject.getString("dtbo_label")));
+            }
+        }
+        return ret;
     }
 
     public List<String> getSupportedOSList() throws Exception {
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index a51bebe..06806ec 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -82,6 +82,7 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
+import java.util.Objects;
 
 @RunWith(DeviceJUnit4Parameterized.class)
 @UseParametersRunnerFactory(DeviceJUnit4ClassRunnerWithParameters.RunnerFactory.class)
@@ -1041,66 +1042,48 @@
         // Check for preconditions
         assumeVfioPlatformSupported();
 
-        List<String> devices = getAssignableDevices();
+        List<AssignableDevice> devices = getAssignableDevices();
         assumeFalse("no assignable devices", devices.isEmpty());
 
-        String vmFdtPath = "/sys/firmware/fdt";
-        File testDir = FileUtil.createTempDir("device_assignment_test");
-        File baseFdtFile = new File(testDir, "base_fdt.dtb");
-        File fdtFile = new File(testDir, "fdt.dtb");
-
-        // Generates baseline DT
-        launchWithDeviceAssignment(/* device= */ null);
-        assertThat(mMicrodroidDevice.pullFile(vmFdtPath, baseFdtFile)).isTrue();
-        getAndroidDevice().shutdownMicrodroid(mMicrodroidDevice);
-
-        // Prepares to run dtdiff. It requires dtc.
-        File dtdiff = findTestFile("dtdiff");
-        RunUtil runner = new RunUtil();
-        String separator = System.getProperty("path.separator");
-        String path = dtdiff.getParent() + separator + System.getenv("PATH");
-        runner.setEnvVariable("PATH", path);
+        String dtSysfsPath = "/proc/device-tree/";
 
         // Try assign devices one by one
-        for (String device : devices) {
-            assertThat(device).isNotNull();
-            launchWithDeviceAssignment(device);
-            assertThat(mMicrodroidDevice.pullFile(vmFdtPath, fdtFile)).isTrue();
+        for (AssignableDevice device : devices) {
+            launchWithDeviceAssignment(device.node);
+
+            String dtPath =
+                    new CommandRunner(mMicrodroidDevice)
+                            .run("cat", dtSysfsPath + "__symbols__/" + device.dtbo_label);
+            assertThat(dtPath).isNotEmpty();
+
+            String resolvedDtPath =
+                    new CommandRunner(mMicrodroidDevice)
+                            .run("readlink", "-e", dtSysfsPath + dtPath);
+            assertThat(resolvedDtPath).isNotEmpty();
+
+            String allDevices =
+                    new CommandRunner(mMicrodroidDevice)
+                            .run("readlink", "-e", "/sys/bus/platform/devices/*/of_node");
+            assertThat(allDevices.split("\n")).asList().contains(resolvedDtPath);
+
             getAndroidDevice().shutdownMicrodroid(mMicrodroidDevice);
-
-            CommandResult result =
-                    runner.runTimedCmd(
-                            500,
-                            dtdiff.getAbsolutePath(),
-                            baseFdtFile.getPath(),
-                            fdtFile.getPath());
-
-            assertWithMessage(
-                            "VM's device tree hasn't changed when assigning "
-                                    + device
-                                    + ", result="
-                                    + result)
-                    .that(result.getStatus())
-                    .isNotEqualTo(CommandStatus.SUCCESS);
+            mMicrodroidDevice = null;
         }
-
-        mMicrodroidDevice = null;
     }
 
     private void launchWithDeviceAssignment(String device) throws Exception {
+        Objects.requireNonNull(device);
         final String configPath = "assets/vm_config.json";
 
-        MicrodroidBuilder builder =
+        mMicrodroidDevice =
                 MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(mProtectedVm)
-                        .gki(mGki);
-        if (device != null) {
-            builder.addAssignableDevice(device);
-        }
-        mMicrodroidDevice = builder.build(getAndroidDevice());
+                        .gki(mGki)
+                        .addAssignableDevice(device)
+                        .build(getAndroidDevice());
 
         assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT)).isTrue();
         assertThat(mMicrodroidDevice.enableAdbRoot()).isTrue();
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index a72f724..b8c6315 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -1521,11 +1521,8 @@
 }
 
 fn is_secretkeeper_supported() -> bool {
-    // TODO(b/327526008): Session establishment wth secretkeeper is failing.
-    // Re-enable this when fixed.
-    let _sk_supported = binder::is_declared(SECRETKEEPER_IDENTIFIER)
-        .expect("Could not check for declared Secretkeeper interface");
-    false
+    binder::is_declared(SECRETKEEPER_IDENTIFIER)
+        .expect("Could not check for declared Secretkeeper interface")
 }
 
 impl VirtualMachineService {
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/AssignableDevice.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/AssignableDevice.aidl
index 014d78c..20114d7 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/AssignableDevice.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/AssignableDevice.aidl
@@ -21,6 +21,6 @@
     /** Path to SysFS node of the device. */
     String node;
 
-    /** Kind of the device. */
-    String kind;
+    /** DTBO label. */
+    String dtbo_label;
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index c6575c8..16975ee 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -108,4 +108,12 @@
      * @param instanceId The ID for the VM.
      */
     void removeVmInstance(in byte[64] instanceId);
+
+    /**
+     * Notification that ownership of a VM has been claimed by the caller.  Note that no permission
+     * checks (with respect to the previous owner) are performed.
+     *
+     * @param instanceId The ID for the VM.
+     */
+    void claimVmInstance(in byte[64] instanceId);
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 05f3cf6..c6150b2 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -372,7 +372,7 @@
         Ok(get_assignable_devices()?
             .device
             .into_iter()
-            .map(|x| AssignableDevice { node: x.sysfs_path, kind: x.kind })
+            .map(|x| AssignableDevice { node: x.sysfs_path, dtbo_label: x.dtbo_label })
             .collect::<Vec<_>>())
     }
 
@@ -436,6 +436,28 @@
         }
         Ok(())
     }
+
+    fn claimVmInstance(&self, instance_id: &[u8; 64]) -> binder::Result<()> {
+        let state = &mut *self.state.lock().unwrap();
+        if let Some(sk_state) = &mut state.sk_state {
+            let uid = get_calling_uid();
+            info!(
+                "Claiming a VM's instance_id: {:?}, for uid: {:?}",
+                hex::encode(instance_id),
+                uid
+            );
+
+            let user_id = multiuser_get_user_id(uid);
+            let app_id = multiuser_get_app_id(uid);
+            info!("Recording possible new owner of state for (user_id={user_id}, app_id={app_id})");
+            if let Err(e) = sk_state.add_id(instance_id, user_id, app_id) {
+                error!("Failed to update the instance_id owner: {e:?}");
+            }
+        } else {
+            info!("ignoring claimVmInstance() as no ISecretkeeper");
+        }
+        Ok(())
+    }
 }
 
 impl IVirtualizationMaintenance for VirtualizationServiceInternal {
@@ -476,10 +498,8 @@
     }
 }
 
-// KEEP IN SYNC WITH assignable_devices.xsd
 #[derive(Debug, Deserialize)]
 struct Device {
-    kind: String,
     dtbo_label: String,
     sysfs_path: String,
 }
diff --git a/virtualizationservice/src/maintenance.rs b/virtualizationservice/src/maintenance.rs
index 219df7d..f950db9 100644
--- a/virtualizationservice/src/maintenance.rs
+++ b/virtualizationservice/src/maintenance.rs
@@ -96,7 +96,8 @@
         }
     }
 
-    /// Record a new VM ID.
+    /// Record a new VM ID.  If there is an existing owner (user_id, app_id) for the VM ID,
+    /// it will be replaced.
     pub fn add_id(&mut self, vm_id: &VmId, user_id: u32, app_id: u32) -> Result<()> {
         let user_id: i32 = user_id.try_into().context(format!("user_id {user_id} out of range"))?;
         let app_id: i32 = app_id.try_into().context(format!("app_id {app_id} out of range"))?;
diff --git a/vm/src/main.rs b/vm/src/main.rs
index bc05ec3..b60f2db 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -30,6 +30,7 @@
 use create_idsig::command_create_idsig;
 use create_partition::command_create_partition;
 use run::{command_run, command_run_app, command_run_microdroid};
+use serde::Serialize;
 use std::num::NonZeroU16;
 use std::path::{Path, PathBuf};
 
@@ -402,8 +403,17 @@
         println!("VFIO-platform is not supported.");
     }
 
+    #[derive(Serialize)]
+    struct AssignableDevice {
+        node: String,
+        dtbo_label: String,
+    }
+
     let devices = get_service()?.getAssignableDevices()?;
-    let devices = devices.into_iter().map(|x| x.node).collect::<Vec<_>>();
+    let devices: Vec<_> = devices
+        .into_iter()
+        .map(|device| AssignableDevice { node: device.node, dtbo_label: device.dtbo_label })
+        .collect();
     println!("Assignable devices: {}", serde_json::to_string(&devices)?);
 
     let os_list = get_service()?.getSupportedOSList()?;