pvmfw: Handle config version 1.1

This CL is the initial step toward provisioning device assignment
with the VM DTBO in the pvmfw config payload.

Test: TH, and test locally with config 1.1 and config 1.0
Bug: 291191157, Bug: 291190552
Change-Id: I85553e692003b81d91e7ab7845b6d77a96f6da5c
diff --git a/pvmfw/README.md b/pvmfw/README.md
index 386036d..698972a 100644
--- a/pvmfw/README.md
+++ b/pvmfw/README.md
@@ -139,6 +139,10 @@
 |  offset = (SECOND - HEAD)     |
 |  size = (SECOND_END - SECOND) |
 +-------------------------------+
+|           [Entry 2]           | <-- Entry 2 is present since version 1.1
+|  offset = (THIRD - HEAD)      |
+|  size = (THIRD_END - SECOND)  |
++-------------------------------+
 |              ...              |
 +-------------------------------+
 |           [Entry n]           |
@@ -152,6 +156,10 @@
 |        {Second blob: DP}      |
 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+ <-- SECOND_END
 | (Padding to 8-byte alignment) |
++===============================+ <-- THIRD
+|        {Third blob: VM DTBO}  |
++~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+ <-- THIRD_END
+| (Padding to 8-byte alignment) |
 +===============================+
 |              ...              |
 +===============================+ <-- TAIL
@@ -177,6 +185,11 @@
 - entry 1 may point to a [DTBO] to be applied to the pVM device tree. See
   [debug policy][debug_policy] for an example.
 
+In version 1.1, new blob is added.
+
+- entry 2 may point to a [DTBO] that describes VM DTBO for device assignment.
+  pvmfw will provision assigned devices with the VM DTBO.
+
 [header]: src/config.rs
 [DTBO]: https://android.googlesource.com/platform/external/dtc/+/refs/heads/master/Documentation/dt-object-internal.txt
 [debug_policy]: ../docs/debug/README.md#debug-policy
diff --git a/pvmfw/src/config.rs b/pvmfw/src/config.rs
index 8a24347..db27001 100644
--- a/pvmfw/src/config.rs
+++ b/pvmfw/src/config.rs
@@ -18,6 +18,7 @@
 use core::mem;
 use core::ops::Range;
 use core::result;
+use log::info;
 use static_assertions::const_assert_eq;
 use vmbase::util::RangeExt;
 use zerocopy::{FromBytes, LayoutVerified};
@@ -81,6 +82,7 @@
 impl Header {
     const MAGIC: u32 = u32::from_ne_bytes(*b"pvmf");
     const VERSION_1_0: Version = Version { major: 1, minor: 0 };
+    const VERSION_1_1: Version = Version { major: 1, minor: 1 };
 
     pub fn total_size(&self) -> usize {
         self.total_size as usize
@@ -101,6 +103,7 @@
     pub fn entry_count(&self) -> Result<usize> {
         let last_entry = match self.version {
             Self::VERSION_1_0 => Entry::DebugPolicy,
+            Self::VERSION_1_1 => Entry::VmDtbo,
             v => return Err(Error::UnsupportedVersion(v)),
         };
 
@@ -112,6 +115,7 @@
 pub enum Entry {
     Bcc,
     DebugPolicy,
+    VmDtbo,
     #[allow(non_camel_case_types)] // TODO: Use mem::variant_count once stable.
     _VARIANT_COUNT,
 }
@@ -181,6 +185,8 @@
             return Err(Error::InvalidFlags(header.flags));
         }
 
+        info!("pvmfw config version: {}", header.version);
+
         // Validate that we won't get an invalid alignment in the following due to padding to u64.
         const_assert_eq!(HEADER_SIZE % mem::size_of::<u64>(), 0);
 
@@ -206,6 +212,7 @@
             // `core::marker::Copy` is not implemented for `core::ops::Range<usize>`.
             Self::validated_body_range(Entry::Bcc, &header_entries, &limits)?,
             Self::validated_body_range(Entry::DebugPolicy, &header_entries, &limits)?,
+            Self::validated_body_range(Entry::VmDtbo, &header_entries, &limits)?,
         ];
 
         Ok(Self { body, ranges })
@@ -216,6 +223,11 @@
         // This assumes that the blobs are in-order w.r.t. the entries.
         let bcc_range = self.get_entry_range(Entry::Bcc).ok_or(Error::MissingEntry(Entry::Bcc))?;
         let dp_range = self.get_entry_range(Entry::DebugPolicy);
+        let vm_dtbo_range = self.get_entry_range(Entry::VmDtbo);
+        // TODO(b/291191157): Provision device assignment with this.
+        if let Some(vm_dtbo_range) = vm_dtbo_range {
+            info!("Found VM DTBO at {:?}", vm_dtbo_range);
+        }
         let bcc_start = bcc_range.start;
         let bcc_end = bcc_range.len();
         let (_, rest) = self.body.split_at_mut(bcc_start);
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/Pvmfw.java b/tests/hostside/helper/java/com/android/microdroid/test/host/Pvmfw.java
index 95eaa58..d752108 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/Pvmfw.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/Pvmfw.java
@@ -33,20 +33,24 @@
     private static final int SIZE_8B = 8; // 8 bytes
     private static final int SIZE_4K = 4 << 10; // 4 KiB, PAGE_SIZE
     private static final int BUFFER_SIZE = 1024;
-    private static final int HEADER_SIZE = Integer.BYTES * 8; // Header has 8 integers.
     private static final int HEADER_MAGIC = 0x666d7670;
-    private static final int HEADER_VERSION = getVersion(1, 0);
+    private static final int HEADER_DEFAULT_VERSION = getVersion(1, 0);
     private static final int HEADER_FLAGS = 0;
 
     @NonNull private final File mPvmfwBinFile;
     @NonNull private final File mBccFile;
     @Nullable private final File mDebugPolicyFile;
+    private final int mVersion;
 
     private Pvmfw(
-            @NonNull File pvmfwBinFile, @NonNull File bccFile, @Nullable File debugPolicyFile) {
+            @NonNull File pvmfwBinFile,
+            @NonNull File bccFile,
+            @Nullable File debugPolicyFile,
+            int version) {
         mPvmfwBinFile = Objects.requireNonNull(pvmfwBinFile);
         mBccFile = Objects.requireNonNull(bccFile);
         mDebugPolicyFile = debugPolicyFile;
+        mVersion = version;
     }
 
     /**
@@ -56,17 +60,22 @@
     public void serialize(@NonNull File outFile) throws IOException {
         Objects.requireNonNull(outFile);
 
-        int bccOffset = HEADER_SIZE;
+        int headerSize = alignTo(getHeaderSize(mVersion), SIZE_8B);
+        int bccOffset = headerSize;
         int bccSize = (int) mBccFile.length();
 
         int debugPolicyOffset = alignTo(bccOffset + bccSize, SIZE_8B);
         int debugPolicySize = mDebugPolicyFile == null ? 0 : (int) mDebugPolicyFile.length();
 
         int totalSize = debugPolicyOffset + debugPolicySize;
+        if (hasVmDtbo(mVersion)) {
+            // Add VM DTBO size as well.
+            totalSize += Integer.BYTES * 2;
+        }
 
-        ByteBuffer header = ByteBuffer.allocate(HEADER_SIZE).order(LITTLE_ENDIAN);
+        ByteBuffer header = ByteBuffer.allocate(headerSize).order(LITTLE_ENDIAN);
         header.putInt(HEADER_MAGIC);
-        header.putInt(HEADER_VERSION);
+        header.putInt(mVersion);
         header.putInt(totalSize);
         header.putInt(HEADER_FLAGS);
         header.putInt(bccOffset);
@@ -74,11 +83,18 @@
         header.putInt(debugPolicyOffset);
         header.putInt(debugPolicySize);
 
+        if (hasVmDtbo(mVersion)) {
+            // Add placeholder entry for VM DTBO.
+            // TODO(b/291191157): Add a real DTBO and test.
+            header.putInt(0);
+            header.putInt(0);
+        }
+
         try (FileOutputStream pvmfw = new FileOutputStream(outFile)) {
             appendFile(pvmfw, mPvmfwBinFile);
             padTo(pvmfw, SIZE_4K);
             pvmfw.write(header.array());
-            padTo(pvmfw, HEADER_SIZE);
+            padTo(pvmfw, SIZE_8B);
             appendFile(pvmfw, mBccFile);
             if (mDebugPolicyFile != null) {
                 padTo(pvmfw, SIZE_8B);
@@ -110,6 +126,19 @@
         }
     }
 
+    private static int getHeaderSize(int version) {
+        if (version == getVersion(1, 0)) {
+            return Integer.BYTES * 8; // Header has 8 integers.
+        }
+        return Integer.BYTES * 10; // Default + VM DTBO (offset, size)
+    }
+
+    private static boolean hasVmDtbo(int version) {
+        int major = getMajorVersion(version);
+        int minor = getMinorVersion(version);
+        return major > 1 || (major == 1 && minor >= 1);
+    }
+
     private static int alignTo(int x, int size) {
         return (x + size - 1) & ~(size - 1);
     }
@@ -118,15 +147,25 @@
         return ((major & 0xFFFF) << 16) | (minor & 0xFFFF);
     }
 
+    private static int getMajorVersion(int version) {
+        return (version >> 16) & 0xFFFF;
+    }
+
+    private static int getMinorVersion(int version) {
+        return version & 0xFFFF;
+    }
+
     /** Builder for {@link Pvmfw}. */
     public static final class Builder {
         @NonNull private final File mPvmfwBinFile;
         @NonNull private final File mBccFile;
         @Nullable private File mDebugPolicyFile;
+        private int mVersion;
 
         public Builder(@NonNull File pvmfwBinFile, @NonNull File bccFile) {
             mPvmfwBinFile = Objects.requireNonNull(pvmfwBinFile);
             mBccFile = Objects.requireNonNull(bccFile);
+            mVersion = HEADER_DEFAULT_VERSION;
         }
 
         @NonNull
@@ -136,8 +175,14 @@
         }
 
         @NonNull
+        public Builder setVersion(int major, int minor) {
+            mVersion = getVersion(major, minor);
+            return this;
+        }
+
+        @NonNull
         public Pvmfw build() {
-            return new Pvmfw(mPvmfwBinFile, mBccFile, mDebugPolicyFile);
+            return new Pvmfw(mPvmfwBinFile, mBccFile, mDebugPolicyFile, mVersion);
         }
     }
 }
diff --git a/tests/hostside/java/com/android/microdroid/test/PvmfwImgTest.java b/tests/hostside/java/com/android/microdroid/test/PvmfwImgTest.java
new file mode 100644
index 0000000..30ad9e2
--- /dev/null
+++ b/tests/hostside/java/com/android/microdroid/test/PvmfwImgTest.java
@@ -0,0 +1,168 @@
+/*
+ * 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 com.android.microdroid.test;
+
+import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assume.assumeTrue;
+import static org.junit.Assume.assumeFalse;
+import static org.junit.Assert.assertThrows;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
+import com.android.microdroid.test.host.Pvmfw;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.DeviceRuntimeException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.util.FileUtil;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.Objects;
+
+/** Tests pvmfw.img and pvmfw */
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class PvmfwImgTest extends MicrodroidHostTestCaseBase {
+    @NonNull private static final String PVMFW_FILE_NAME = "pvmfw_test.bin";
+    @NonNull private static final String BCC_FILE_NAME = "bcc.dat";
+    @NonNull private static final String PACKAGE_FILE_NAME = "MicrodroidTestApp.apk";
+    @NonNull private static final String PACKAGE_NAME = "com.android.microdroid.test";
+    @NonNull private static final String MICRODROID_DEBUG_FULL = "full";
+    @NonNull private static final String MICRODROID_CONFIG_PATH = "assets/vm_config_apex.json";
+    private static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
+    private static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
+
+    @NonNull private static final String CUSTOM_PVMFW_FILE_PREFIX = "pvmfw";
+    @NonNull private static final String CUSTOM_PVMFW_FILE_SUFFIX = ".bin";
+    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
+    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
+
+    @Nullable private static File mPvmfwBinFileOnHost;
+    @Nullable private static File mBccFileOnHost;
+
+    @Nullable private TestDevice mAndroidDevice;
+    @Nullable private ITestDevice mMicrodroidDevice;
+    @Nullable private File mCustomPvmfwBinFileOnHost;
+
+    @Before
+    public void setUp() throws Exception {
+        mAndroidDevice = (TestDevice) Objects.requireNonNull(getDevice());
+
+        // Check device capabilities
+        assumeDeviceIsCapable(mAndroidDevice);
+        assumeTrue(
+                "Skip if protected VMs are not supported",
+                mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
+        assumeFalse("Test requires setprop for using custom pvmfw and adb root", isUserBuild());
+
+        assumeTrue("Skip if adb root fails", mAndroidDevice.enableAdbRoot());
+
+        // tradefed copies the test artfacts under /tmp when running tests,
+        // so we should *find* the artifacts with the file name.
+        mPvmfwBinFileOnHost =
+                getTestInformation().getDependencyFile(PVMFW_FILE_NAME, /* targetFirst= */ false);
+        mBccFileOnHost =
+                getTestInformation().getDependencyFile(BCC_FILE_NAME, /* targetFirst= */ false);
+
+        // Prepare for system properties for custom pvmfw.img.
+        // File will be prepared later in individual test and then pushed to device
+        // when launching with launchProtectedVmAndWaitForBootCompleted().
+        mCustomPvmfwBinFileOnHost =
+                FileUtil.createTempFile(CUSTOM_PVMFW_FILE_PREFIX, CUSTOM_PVMFW_FILE_SUFFIX);
+        mAndroidDevice.setProperty(CUSTOM_PVMFW_IMG_PATH_PROP, CUSTOM_PVMFW_IMG_PATH);
+
+        // Prepare for launching microdroid
+        mAndroidDevice.installPackage(findTestFile(PACKAGE_FILE_NAME), /* reinstall */ false);
+        prepareVirtualizationTestSetup(mAndroidDevice);
+        mMicrodroidDevice = null;
+    }
+
+    @After
+    public void shutdown() throws Exception {
+        if (!mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true)) {
+            return;
+        }
+        if (mMicrodroidDevice != null) {
+            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
+            mMicrodroidDevice = null;
+        }
+        mAndroidDevice.uninstallPackage(PACKAGE_NAME);
+
+        // Cleanup for custom pvmfw.img
+        mAndroidDevice.setProperty(CUSTOM_PVMFW_IMG_PATH_PROP, "");
+        FileUtil.deleteFile(mCustomPvmfwBinFileOnHost);
+
+        cleanUpVirtualizationTestSetup(mAndroidDevice);
+
+        mAndroidDevice.disableAdbRoot();
+    }
+
+    @Test
+    public void testConfigVersion1_0_boots() throws Exception {
+        Pvmfw pvmfw =
+                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 0).build();
+        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+
+        launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
+    }
+
+    @Test
+    public void testConfigVersion1_1_boots() throws Exception {
+        Pvmfw pvmfw =
+                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 1).build();
+        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+
+        launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
+    }
+
+    @Test
+    public void testInvalidConfigVersion_doesNotBoot() throws Exception {
+        // Disclaimer: Update versions when it becomes valid
+        Pvmfw pvmfw =
+                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 100).build();
+        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+
+        assertThrows(
+                "pvmfw shouldn't boot with invalid version",
+                DeviceRuntimeException.class,
+                () -> launchProtectedVmAndWaitForBootCompleted(BOOT_FAILURE_WAIT_TIME_MS));
+    }
+
+    private ITestDevice launchProtectedVmAndWaitForBootCompleted(long adbTimeoutMs)
+            throws DeviceNotAvailableException {
+        mMicrodroidDevice =
+                MicrodroidBuilder.fromDevicePath(
+                                getPathForPackage(PACKAGE_NAME), MICRODROID_CONFIG_PATH)
+                        .debugLevel(MICRODROID_DEBUG_FULL)
+                        .protectedVm(true)
+                        .addBootFile(mCustomPvmfwBinFileOnHost, PVMFW_FILE_NAME)
+                        .setAdbConnectTimeoutMs(adbTimeoutMs)
+                        .build(mAndroidDevice);
+        assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS)).isTrue();
+        return mMicrodroidDevice;
+    }
+}