diff --git a/javalib/api/system-current.txt b/javalib/api/system-current.txt
index 1977321..fe9943d 100644
--- a/javalib/api/system-current.txt
+++ b/javalib/api/system-current.txt
@@ -56,7 +56,7 @@
   }
 
   public final class VirtualMachineConfig {
-    method @NonNull public String getApkPath();
+    method @Nullable public String getApkPath();
     method @NonNull public int getDebugLevel();
     method @IntRange(from=0) public long getEncryptedStorageKib();
     method @IntRange(from=0) public int getMemoryMib();
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 8cebc3c..7c7f4b5 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -769,7 +769,7 @@
                 }
             } catch (IOException e) {
                 // If the file already exists, exception is not thrown.
-                throw new VirtualMachineException("failed to create idsig file", e);
+                throw new VirtualMachineException("Failed to create APK signature file", e);
             }
 
             IVirtualizationService service = mVirtualizationService.connect();
@@ -779,33 +779,16 @@
                     createVmPipes();
                 }
 
-                VirtualMachineAppConfig appConfig = getConfig().toVsConfig();
+                VirtualMachineAppConfig appConfig =
+                        getConfig().toVsConfig(mContext.getPackageManager());
                 appConfig.name = mName;
 
-                // Fill the idsig file by hashing the apk
-                service.createOrUpdateIdsigFile(
-                        appConfig.apk, ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_WRITE));
-
-                for (ExtraApkSpec extraApk : mExtraApks) {
-                    service.createOrUpdateIdsigFile(
-                            ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY),
-                            ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_WRITE));
+                try {
+                    createIdSigs(service, appConfig);
+                } catch (FileNotFoundException e) {
+                    throw new VirtualMachineException("Failed to generate APK signature", e);
                 }
 
-                // Re-open idsig file in read-only mode
-                appConfig.idsig = ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_ONLY);
-                appConfig.instanceImage =
-                        ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_WRITE);
-                if (mEncryptedStoreFilePath != null) {
-                    appConfig.encryptedStorageImage =
-                            ParcelFileDescriptor.open(mEncryptedStoreFilePath, MODE_READ_WRITE);
-                }
-                List<ParcelFileDescriptor> extraIdsigs = new ArrayList<>();
-                for (ExtraApkSpec extraApk : mExtraApks) {
-                    extraIdsigs.add(ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_ONLY));
-                }
-                appConfig.extraIdsigs = extraIdsigs;
-
                 android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
                         android.system.virtualizationservice.VirtualMachineConfig.appConfig(
                                 appConfig);
@@ -814,7 +797,7 @@
                 mVirtualMachine.registerCallback(new CallbackTranslator(service));
                 mContext.registerComponentCallbacks(mMemoryManagementCallbacks);
                 mVirtualMachine.start();
-            } catch (IOException | IllegalStateException | ServiceSpecificException e) {
+            } catch (IllegalStateException | ServiceSpecificException e) {
                 throw new VirtualMachineException(e);
             } catch (RemoteException e) {
                 throw e.rethrowAsRuntimeException();
@@ -822,6 +805,32 @@
         }
     }
 
+    private void createIdSigs(IVirtualizationService service, VirtualMachineAppConfig appConfig)
+            throws RemoteException, FileNotFoundException {
+        // Fill the idsig file by hashing the apk
+        service.createOrUpdateIdsigFile(
+                appConfig.apk, ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_WRITE));
+
+        for (ExtraApkSpec extraApk : mExtraApks) {
+            service.createOrUpdateIdsigFile(
+                    ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY),
+                    ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_WRITE));
+        }
+
+        // Re-open idsig files in read-only mode
+        appConfig.idsig = ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_ONLY);
+        appConfig.instanceImage = ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_WRITE);
+        if (mEncryptedStoreFilePath != null) {
+            appConfig.encryptedStorageImage =
+                    ParcelFileDescriptor.open(mEncryptedStoreFilePath, MODE_READ_WRITE);
+        }
+        List<ParcelFileDescriptor> extraIdsigs = new ArrayList<>();
+        for (ExtraApkSpec extraApk : mExtraApks) {
+            extraIdsigs.add(ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_ONLY));
+        }
+        appConfig.extraIdsigs = extraIdsigs;
+    }
+
     @GuardedBy("mLock")
     private void createVmPipes() throws VirtualMachineException {
         try {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index f5c3cd2..b358f9e 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -29,6 +29,8 @@
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
 import android.content.Context;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
 import android.sysprop.HypervisorProperties;
@@ -58,8 +60,9 @@
     private static final String[] EMPTY_STRING_ARRAY = {};
 
     // These define the schema of the config file persisted on disk.
-    private static final int VERSION = 3;
+    private static final int VERSION = 4;
     private static final String KEY_VERSION = "version";
+    private static final String KEY_PACKAGENAME = "packageName";
     private static final String KEY_APKPATH = "apkPath";
     private static final String KEY_PAYLOADCONFIGPATH = "payloadConfigPath";
     private static final String KEY_PAYLOADBINARYNAME = "payloadBinaryPath";
@@ -94,8 +97,11 @@
      */
     @SystemApi public static final int DEBUG_LEVEL_FULL = 1;
 
+    /** Name of a package whose primary APK contains the VM payload. */
+    @Nullable private final String mPackageName;
+
     /** Absolute path to the APK file containing the VM payload. */
-    @NonNull private final String mApkPath;
+    @Nullable private final String mApkPath;
 
     @DebugLevel private final int mDebugLevel;
 
@@ -129,7 +135,8 @@
     private final boolean mVmOutputCaptured;
 
     private VirtualMachineConfig(
-            @NonNull String apkPath,
+            @Nullable String packageName,
+            @Nullable String apkPath,
             @Nullable String payloadConfigPath,
             @Nullable String payloadBinaryName,
             @DebugLevel int debugLevel,
@@ -139,6 +146,7 @@
             long encryptedStorageKib,
             boolean vmOutputCaptured) {
         // This is only called from Builder.build(); the builder handles parameter validation.
+        mPackageName = packageName;
         mApkPath = apkPath;
         mPayloadConfigPath = payloadConfigPath;
         mPayloadBinaryName = payloadBinaryName;
@@ -191,8 +199,13 @@
                     "Version " + version + " too high; current is " + VERSION);
         }
 
-        Builder builder = new Builder();
-        builder.setApkPath(b.getString(KEY_APKPATH));
+        String packageName = b.getString(KEY_PACKAGENAME);
+        Builder builder = new Builder(packageName);
+
+        String apkPath = b.getString(KEY_APKPATH);
+        if (apkPath != null) {
+            builder.setApkPath(apkPath);
+        }
 
         String payloadConfigPath = b.getString(KEY_PAYLOADCONFIGPATH);
         if (payloadConfigPath == null) {
@@ -234,7 +247,12 @@
     private void serializeOutputStream(@NonNull OutputStream output) throws IOException {
         PersistableBundle b = new PersistableBundle();
         b.putInt(KEY_VERSION, VERSION);
-        b.putString(KEY_APKPATH, mApkPath);
+        if (mPackageName != null) {
+            b.putString(KEY_PACKAGENAME, mPackageName);
+        }
+        if (mApkPath != null) {
+            b.putString(KEY_APKPATH, mApkPath);
+        }
         b.putString(KEY_PAYLOADCONFIGPATH, mPayloadConfigPath);
         b.putString(KEY_PAYLOADBINARYNAME, mPayloadBinaryName);
         b.putInt(KEY_DEBUGLEVEL, mDebugLevel);
@@ -252,12 +270,13 @@
 
     /**
      * Returns the absolute path of the APK which should contain the binary payload that will
-     * execute within the VM.
+     * execute within the VM. Returns null if no specific path has been set, so the primary APK will
+     * be used.
      *
      * @hide
      */
     @SystemApi
-    @NonNull
+    @Nullable
     public String getApkPath() {
         return mApkPath;
     }
@@ -383,7 +402,8 @@
                 && this.mVmOutputCaptured == other.mVmOutputCaptured
                 && Objects.equals(this.mPayloadConfigPath, other.mPayloadConfigPath)
                 && Objects.equals(this.mPayloadBinaryName, other.mPayloadBinaryName)
-                && this.mApkPath.equals(other.mApkPath);
+                && Objects.equals(this.mPackageName, other.mPackageName)
+                && Objects.equals(this.mApkPath, other.mApkPath);
     }
 
     /**
@@ -393,9 +413,28 @@
      * app-owned files and that could be abused to run a VM with software that the calling
      * application doesn't own.
      */
-    VirtualMachineAppConfig toVsConfig() throws FileNotFoundException {
+    VirtualMachineAppConfig toVsConfig(@NonNull PackageManager packageManager)
+            throws VirtualMachineException {
         VirtualMachineAppConfig vsConfig = new VirtualMachineAppConfig();
-        vsConfig.apk = ParcelFileDescriptor.open(new File(mApkPath), MODE_READ_ONLY);
+
+        String apkPath = mApkPath;
+        if (apkPath == null) {
+            try {
+                ApplicationInfo appInfo =
+                        packageManager.getApplicationInfo(
+                                mPackageName, PackageManager.ApplicationInfoFlags.of(0));
+                // This really is the path to the APK, not a directory.
+                apkPath = appInfo.sourceDir;
+            } catch (PackageManager.NameNotFoundException e) {
+                throw new VirtualMachineException("Package not found", e);
+            }
+        }
+
+        try {
+            vsConfig.apk = ParcelFileDescriptor.open(new File(apkPath), MODE_READ_ONLY);
+        } catch (FileNotFoundException e) {
+            throw new VirtualMachineException("Failed to open APK", e);
+        }
         if (mPayloadBinaryName != null) {
             VirtualMachinePayloadConfig payloadConfig = new VirtualMachinePayloadConfig();
             payloadConfig.payloadBinaryName = mPayloadBinaryName;
@@ -428,7 +467,7 @@
      */
     @SystemApi
     public static final class Builder {
-        @Nullable private final Context mContext;
+        @Nullable private final String mPackageName;
         @Nullable private String mApkPath;
         @Nullable private String mPayloadConfigPath;
         @Nullable private String mPayloadBinaryName;
@@ -447,15 +486,15 @@
          */
         @SystemApi
         public Builder(@NonNull Context context) {
-            mContext = requireNonNull(context, "context must not be null");
+            mPackageName = requireNonNull(context, "context must not be null").getPackageName();
         }
 
         /**
-         * Creates a builder with no associated context; {@link #setApkPath} must be called to
-         * specify which APK contains the payload.
+         * Creates a builder for a specific package. If packageName is null, {@link #setApkPath}
+         * must be called to specify the APK containing the payload.
          */
-        private Builder() {
-            mContext = null;
+        private Builder(@Nullable String packageName) {
+            mPackageName = packageName;
         }
 
         /**
@@ -466,14 +505,16 @@
         @SystemApi
         @NonNull
         public VirtualMachineConfig build() {
-            String apkPath;
-            if (mApkPath == null) {
-                if (mContext == null) {
-                    throw new IllegalStateException("apkPath must be specified");
-                }
-                apkPath = mContext.getPackageCodePath();
-            } else {
+            String apkPath = null;
+            String packageName = null;
+
+            if (mApkPath != null) {
                 apkPath = mApkPath;
+            } else if (mPackageName != null) {
+                packageName = mPackageName;
+            } else {
+                // This should never happen, unless we're deserializing a bad config
+                throw new IllegalStateException("apkPath or packageName must be specified");
             }
 
             if (mPayloadBinaryName == null) {
@@ -496,6 +537,7 @@
             }
 
             return new VirtualMachineConfig(
+                    packageName,
                     apkPath,
                     mPayloadConfigPath,
                     mPayloadBinaryName,
diff --git a/pvmfw/avb/src/ops.rs b/pvmfw/avb/src/ops.rs
index 903fecb..e7f0ac7 100644
--- a/pvmfw/avb/src/ops.rs
+++ b/pvmfw/avb/src/ops.rs
@@ -21,7 +21,7 @@
 use crate::utils::{self, as_ref, is_not_null, to_nonnull, write};
 use avb_bindgen::{
     avb_slot_verify, avb_slot_verify_data_free, AvbHashtreeErrorMode, AvbIOResult, AvbOps,
-    AvbSlotVerifyData, AvbSlotVerifyFlags, AvbVBMetaData,
+    AvbPartitionData, AvbSlotVerifyData, AvbSlotVerifyFlags, AvbVBMetaData,
 };
 use core::{
     ffi::{c_char, c_void, CStr},
@@ -229,15 +229,13 @@
 extern "C" fn read_rollback_index(
     _ops: *mut AvbOps,
     _rollback_index_location: usize,
-    _out_rollback_index: *mut u64,
+    out_rollback_index: *mut u64,
 ) -> AvbIOResult {
-    // TODO(b/256148034): Write -1 to out_rollback_index
-    // so that we won't compare the current rollback index with uninitialized number
-    // in avb_slot_verify.
-
-    // Rollback protection is not yet implemented, but
-    // this method is required by `avb_slot_verify()`.
-    AvbIOResult::AVB_IO_RESULT_OK
+    // Rollback protection is not yet implemented, but this method is required by
+    // `avb_slot_verify()`.
+    // We set `out_rollback_index` to 0 to ensure that the default rollback index (0)
+    // is never smaller than it, thus the rollback index check will pass.
+    to_avb_io_result(write(out_rollback_index, 0))
 }
 
 extern "C" fn get_unique_guid_for_partition(
@@ -327,4 +325,15 @@
             unsafe { slice::from_raw_parts(data.vbmeta_images, data.num_vbmeta_images) };
         Ok(vbmeta_images)
     }
+
+    pub(crate) fn loaded_partitions(&self) -> Result<&[AvbPartitionData], AvbSlotVerifyError> {
+        let data = self.as_ref();
+        is_not_null(data.loaded_partitions).map_err(|_| AvbSlotVerifyError::Io)?;
+        // SAFETY: It is safe as the raw pointer `data.loaded_partitions` is a nonnull pointer and
+        // is guaranteed by libavb to point to a valid `AvbPartitionData` array as part of the
+        // `AvbSlotVerifyData` struct.
+        let loaded_partitions =
+            unsafe { slice::from_raw_parts(data.loaded_partitions, data.num_loaded_partitions) };
+        Ok(loaded_partitions)
+    }
 }
diff --git a/pvmfw/avb/src/verify.rs b/pvmfw/avb/src/verify.rs
index d6a0cb2..a062061 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -20,7 +20,7 @@
 use crate::utils::{is_not_null, to_usize, usize_checked_add, write};
 use avb_bindgen::{
     avb_descriptor_foreach, avb_hash_descriptor_validate_and_byteswap, AvbDescriptor,
-    AvbHashDescriptor, AvbVBMetaData,
+    AvbHashDescriptor, AvbPartitionData, AvbVBMetaData,
 };
 use core::{
     ffi::{c_char, c_void},
@@ -135,6 +135,16 @@
     }
 }
 
+fn verify_only_one_vbmeta_exists(
+    vbmeta_images: &[AvbVBMetaData],
+) -> Result<(), AvbSlotVerifyError> {
+    if vbmeta_images.len() == 1 {
+        Ok(())
+    } else {
+        Err(AvbSlotVerifyError::InvalidMetadata)
+    }
+}
+
 fn verify_vbmeta_is_from_kernel_partition(
     vbmeta_image: &AvbVBMetaData,
 ) -> Result<(), AvbSlotVerifyError> {
@@ -144,6 +154,29 @@
     }
 }
 
+fn verify_loaded_partition_has_expected_length(
+    loaded_partitions: &[AvbPartitionData],
+    partition_name: PartitionName,
+    expected_len: usize,
+) -> Result<(), AvbSlotVerifyError> {
+    if loaded_partitions.len() != 1 {
+        // Only one partition should be loaded in each verify result.
+        return Err(AvbSlotVerifyError::Io);
+    }
+    let loaded_partition = loaded_partitions[0];
+    if !PartitionName::try_from(loaded_partition.partition_name as *const c_char)
+        .map_or(false, |p| p == partition_name)
+    {
+        // Only the requested partition should be loaded.
+        return Err(AvbSlotVerifyError::Io);
+    }
+    if loaded_partition.data_size == expected_len {
+        Ok(())
+    } else {
+        Err(AvbSlotVerifyError::Verification)
+    }
+}
+
 /// Verifies the payload (signed kernel + initrd) against the trusted public key.
 pub fn verify_payload(
     kernel: &[u8],
@@ -153,13 +186,12 @@
     let mut payload = Payload::new(kernel, initrd, trusted_public_key);
     let mut ops = Ops::from(&mut payload);
     let kernel_verify_result = ops.verify_partition(PartitionName::Kernel.as_cstr())?;
+
     let vbmeta_images = kernel_verify_result.vbmeta_images()?;
-    if vbmeta_images.len() != 1 {
-        // There can only be one VBMeta.
-        return Err(AvbSlotVerifyError::InvalidMetadata);
-    }
+    verify_only_one_vbmeta_exists(vbmeta_images)?;
     let vbmeta_image = vbmeta_images[0];
     verify_vbmeta_is_from_kernel_partition(&vbmeta_image)?;
+
     if initrd.is_none() {
         verify_vbmeta_has_no_initrd_descriptor(&vbmeta_image)?;
         return Ok(DebugLevel::None);
@@ -167,12 +199,20 @@
     // TODO(b/256148034): Check the vbmeta doesn't have hash descriptors other than
     // boot, initrd_normal, initrd_debug.
 
-    let debug_level = if ops.verify_partition(PartitionName::InitrdNormal.as_cstr()).is_ok() {
-        DebugLevel::None
-    } else if ops.verify_partition(PartitionName::InitrdDebug.as_cstr()).is_ok() {
-        DebugLevel::Full
-    } else {
-        return Err(AvbSlotVerifyError::Verification);
-    };
+    let initrd = initrd.unwrap();
+    let (debug_level, initrd_verify_result, initrd_partition_name) =
+        if let Ok(result) = ops.verify_partition(PartitionName::InitrdNormal.as_cstr()) {
+            (DebugLevel::None, result, PartitionName::InitrdNormal)
+        } else if let Ok(result) = ops.verify_partition(PartitionName::InitrdDebug.as_cstr()) {
+            (DebugLevel::Full, result, PartitionName::InitrdDebug)
+        } else {
+            return Err(AvbSlotVerifyError::Verification);
+        };
+    let loaded_partitions = initrd_verify_result.loaded_partitions()?;
+    verify_loaded_partition_has_expected_length(
+        loaded_partitions,
+        initrd_partition_name,
+        initrd.len(),
+    )?;
     Ok(debug_level)
 }
diff --git a/pvmfw/avb/tests/api_test.rs b/pvmfw/avb/tests/api_test.rs
index 41ead59..0572789 100644
--- a/pvmfw/avb/tests/api_test.rs
+++ b/pvmfw/avb/tests/api_test.rs
@@ -202,6 +202,19 @@
 }
 
 #[test]
+fn extended_initrd_fails_verification() -> Result<()> {
+    let mut initrd = load_latest_initrd_normal()?;
+    initrd.extend(b"androidboot.vbmeta.digest=1111");
+
+    assert_payload_verification_with_initrd_eq(
+        &load_latest_signed_kernel()?,
+        &initrd,
+        &load_trusted_public_key()?,
+        Err(AvbSlotVerifyError::Verification),
+    )
+}
+
+#[test]
 fn tampered_vbmeta_fails_verification() -> Result<()> {
     let mut kernel = load_latest_signed_kernel()?;
     let footer = extract_avb_footer(&kernel)?;
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index e1a2e40..7bd5f08 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -31,6 +31,7 @@
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 
 import android.content.Context;
+import android.content.ContextWrapper;
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
 import android.os.ParcelFileDescriptor.AutoCloseInputStream;
@@ -342,7 +343,7 @@
         VirtualMachineConfig.Builder minimalBuilder = newVmConfigBuilder();
         VirtualMachineConfig minimal = minimalBuilder.setPayloadBinaryName("binary.so").build();
 
-        assertThat(minimal.getApkPath()).isEqualTo(getContext().getPackageCodePath());
+        assertThat(minimal.getApkPath()).isNull();
         assertThat(minimal.getDebugLevel()).isEqualTo(DEBUG_LEVEL_NONE);
         assertThat(minimal.getMemoryMib()).isEqualTo(0);
         assertThat(minimal.getNumCpus()).isEqualTo(1);
@@ -425,13 +426,9 @@
         assertThat(e).hasMessageThat().contains("debug level must be FULL to capture output");
     }
 
-    private VirtualMachineConfig.Builder newBaselineBuilder() {
-        return newVmConfigBuilder().setPayloadBinaryName("binary.so").setApkPath("/apk/path");
-    }
-
     @Test
     @CddTest(requirements = {"9.17/C-1-1"})
-    public void compatibleConfigTests() throws Exception {
+    public void compatibleConfigTests() {
         int maxCpus = Runtime.getRuntime().availableProcessors();
 
         VirtualMachineConfig baseline = newBaselineBuilder().build();
@@ -467,6 +464,31 @@
                 newBaselineBuilder().setDebugLevel(DEBUG_LEVEL_FULL);
         VirtualMachineConfig debuggable = debuggableBuilder.build();
         assertConfigCompatible(debuggable, debuggableBuilder.setVmOutputCaptured(true)).isFalse();
+
+        VirtualMachineConfig currentContextConfig =
+                new VirtualMachineConfig.Builder(getContext())
+                        .setProtectedVm(isProtectedVm())
+                        .setPayloadBinaryName("binary.so")
+                        .build();
+
+        // packageName is not directly exposed by the config, so we have to be a bit creative
+        // to modify it.
+        Context otherContext =
+                new ContextWrapper(getContext()) {
+                    @Override
+                    public String getPackageName() {
+                        return "other.package.name";
+                    }
+                };
+        VirtualMachineConfig.Builder otherContextBuilder =
+                new VirtualMachineConfig.Builder(otherContext)
+                        .setProtectedVm(isProtectedVm())
+                        .setPayloadBinaryName("binary.so");
+        assertConfigCompatible(currentContextConfig, otherContextBuilder).isFalse();
+    }
+
+    private VirtualMachineConfig.Builder newBaselineBuilder() {
+        return newVmConfigBuilder().setPayloadBinaryName("binary.so").setApkPath("/apk/path");
     }
 
     private BooleanSubject assertConfigCompatible(