diff --git a/android/TerminalApp/AndroidManifest.xml b/android/TerminalApp/AndroidManifest.xml
index 7dab58d..726004c 100644
--- a/android/TerminalApp/AndroidManifest.xml
+++ b/android/TerminalApp/AndroidManifest.xml
@@ -54,7 +54,9 @@
             android:label="@string/settings_port_forwarding_title" />
         <activity android:name=".SettingsRecoveryActivity"
             android:label="@string/settings_recovery_title" />
-        <activity android:name=".ErrorActivity" />
+        <activity android:name=".ErrorActivity"
+            android:label="@string/error_title"
+            android:process=":error" />
         <property
             android:name="android.window.PROPERTY_ACTIVITY_EMBEDDING_SPLITS_ENABLED"
             android:value="true" />
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.java b/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.java
index d6521be..d6ca1e6 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.java
@@ -39,6 +39,15 @@
                             NotificationManager.IMPORTANCE_DEFAULT);
             notificationManager.createNotificationChannel(channel);
         }
+
+        if (!(this instanceof ErrorActivity)) {
+            Thread currentThread = Thread.currentThread();
+            if (!(currentThread.getUncaughtExceptionHandler()
+                    instanceof TerminalExceptionHandler)) {
+                currentThread.setUncaughtExceptionHandler(
+                        new TerminalExceptionHandler(getApplicationContext()));
+            }
+        }
     }
 
     @Override
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/CertificateUtils.java b/android/TerminalApp/java/com/android/virtualization/terminal/CertificateUtils.java
index fa5c382..e3d1a67 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/CertificateUtils.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/CertificateUtils.java
@@ -62,9 +62,8 @@
             }
             return ((KeyStore.PrivateKeyEntry) ks.getEntry(ALIAS, null));
         } catch (Exception e) {
-            Log.e(TAG, "cannot generate or get key", e);
+            throw new RuntimeException("cannot generate or get key", e);
         }
-        return null;
     }
 
     private static void createKey()
@@ -95,7 +94,7 @@
                             + end_cert;
             writer.write(output.getBytes());
         } catch (IOException | CertificateEncodingException e) {
-            Log.d(TAG, "cannot write cert", e);
+            throw new RuntimeException("cannot write certs", e);
         }
     }
 }
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java b/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java
index ac05d78..66ab414 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/InstallerService.java
@@ -41,7 +41,6 @@
 import java.net.SocketException;
 import java.net.UnknownHostException;
 import java.nio.file.Path;
-import java.util.Arrays;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
@@ -75,6 +74,7 @@
                         this, /* requestCode= */ 0, intent, PendingIntent.FLAG_IMMUTABLE);
         mNotification =
                 new Notification.Builder(this, this.getPackageName())
+                        .setSilent(true)
                         .setSmallIcon(R.drawable.ic_launcher_foreground)
                         .setContentTitle(getString(R.string.installer_notif_title_text))
                         .setContentText(getString(R.string.installer_notif_desc_text))
@@ -82,7 +82,9 @@
                         .setContentIntent(pendingIntent)
                         .build();
 
-        mExecutorService = Executors.newSingleThreadExecutor();
+        mExecutorService =
+                Executors.newSingleThreadExecutor(
+                        new TerminalThreadFactory(getApplicationContext()));
 
         mConnectivityManager = getSystemService(ConnectivityManager.class);
         Network defaultNetwork = mConnectivityManager.getBoundNetworkForProcess();
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
index 397a546..624d6ca 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
@@ -32,6 +32,7 @@
 import android.os.Bundle;
 import android.os.ConditionVariable;
 import android.os.Environment;
+import android.os.SystemClock;
 import android.provider.Settings;
 import android.util.Log;
 import android.view.KeyEvent;
@@ -67,6 +68,8 @@
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
 import java.util.Map;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 public class MainActivity extends BaseActivity
         implements VmLauncherService.VmLauncherServiceCallback, AccessibilityStateChangeListener {
@@ -74,9 +77,11 @@
     static final String KEY_DISK_SIZE = "disk_size";
     private static final String VM_ADDR = "192.168.0.2";
     private static final int TTYD_PORT = 7681;
+    private static final int TERMINAL_CONNECTION_TIMEOUT_MS = 10_000;
     private static final int REQUEST_CODE_INSTALLER = 0x33;
     private static final int FONT_SIZE_DEFAULT = 13;
 
+    private ExecutorService mExecutorService;
     private InstalledImage mImage;
     private X509Certificate[] mCertificates;
     private PrivateKey mPrivateKey;
@@ -141,6 +146,11 @@
                             updateModifierKeysVisibility();
                             return insets;
                         });
+
+        mExecutorService =
+                Executors.newSingleThreadExecutor(
+                        new TerminalThreadFactory(getApplicationContext()));
+
         // if installer is launched, it will be handled in onActivityResult
         if (!launchInstaller) {
             if (!Environment.isExternalStorageManager()) {
@@ -323,15 +333,12 @@
                         handler.proceed();
                     }
                 });
-        new Thread(
-                        () -> {
-                            waitUntilVmStarts();
-                            runOnUiThread(
-                                    () ->
-                                            mTerminalView.loadUrl(
-                                                    getTerminalServiceUrl().toString()));
-                        })
-                .start();
+        mExecutorService.execute(
+                () -> {
+                    // TODO(b/376793781): Remove polling
+                    waitUntilVmStarts();
+                    runOnUiThread(() -> mTerminalView.loadUrl(getTerminalServiceUrl().toString()));
+                });
     }
 
     private static void waitUntilVmStarts() {
@@ -341,17 +348,33 @@
         } catch (UnknownHostException e) {
             // this can never happen.
         }
-        try {
-            while (!addr.isReachable(10000)) {}
-        } catch (IOException e) {
-            // give up on network error
-            throw new RuntimeException(e);
+
+        long startTime = SystemClock.elapsedRealtime();
+        while (true) {
+            int remainingTime =
+                    TERMINAL_CONNECTION_TIMEOUT_MS
+                            - (int) (SystemClock.elapsedRealtime() - startTime);
+            if (remainingTime <= 0) {
+                throw new RuntimeException("Connection to terminal timedout");
+            }
+            try {
+                // Note: this quits immediately if VM is unreachable.
+                if (addr.isReachable(remainingTime)) {
+                    return;
+                }
+            } catch (IOException e) {
+                // give up on network error
+                throw new RuntimeException(e);
+            }
         }
-        return;
     }
 
     @Override
     protected void onDestroy() {
+        if (mExecutorService != null) {
+            mExecutorService.shutdown();
+        }
+
         getSystemService(AccessibilityManager.class).removeAccessibilityStateChangeListener(this);
         VmLauncherService.stop(this);
         super.onDestroy();
@@ -471,6 +494,7 @@
         Icon icon = Icon.createWithResource(getResources(), R.drawable.ic_launcher_foreground);
         Notification notification =
                 new Notification.Builder(this, this.getPackageName())
+                        .setSilent(true)
                         .setSmallIcon(R.drawable.ic_launcher_foreground)
                         .setContentTitle(
                                 getResources().getString(R.string.service_notification_title))
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/TerminalExceptionHandler.java b/android/TerminalApp/java/com/android/virtualization/terminal/TerminalExceptionHandler.java
new file mode 100644
index 0000000..4ab2b77
--- /dev/null
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/TerminalExceptionHandler.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright 2024 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.virtualization.terminal;
+
+import android.content.Context;
+import android.util.Log;
+
+public class TerminalExceptionHandler implements Thread.UncaughtExceptionHandler {
+    private static final String TAG = "TerminalExceptionHandler";
+
+    private final Context mContext;
+
+    public TerminalExceptionHandler(Context context) {
+        mContext = context;
+    }
+
+    @Override
+    public void uncaughtException(Thread thread, Throwable throwable) {
+        Exception exception;
+        if (throwable instanceof Exception) {
+            exception = (Exception) throwable;
+        } else {
+            exception = new Exception(throwable);
+        }
+        try {
+            ErrorActivity.start(mContext, exception);
+        } catch (Exception ex) {
+            Log.wtf(TAG, "Failed to launch error activity for an exception", exception);
+        }
+
+        thread.getDefaultUncaughtExceptionHandler().uncaughtException(thread, throwable);
+    }
+}
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/TerminalThreadFactory.java b/android/TerminalApp/java/com/android/virtualization/terminal/TerminalThreadFactory.java
new file mode 100644
index 0000000..5ee535d
--- /dev/null
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/TerminalThreadFactory.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2024 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.virtualization.terminal;
+
+import android.content.Context;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
+
+public class TerminalThreadFactory implements ThreadFactory {
+    private final Context mContext;
+
+    public TerminalThreadFactory(Context context) {
+        mContext = context;
+    }
+
+    @Override
+    public Thread newThread(Runnable r) {
+        Thread thread = Executors.defaultThreadFactory().newThread(r);
+        thread.setUncaughtExceptionHandler(new TerminalExceptionHandler(mContext));
+        return thread;
+    }
+}
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.java b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.java
index 6d2c5bd..f262f1f 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/VmLauncherService.java
@@ -20,9 +20,11 @@
 
 import android.app.Notification;
 import android.app.NotificationManager;
+import android.app.PendingIntent;
 import android.app.Service;
 import android.content.Context;
 import android.content.Intent;
+import android.graphics.drawable.Icon;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
@@ -143,8 +145,13 @@
     @Override
     public int onStartCommand(Intent intent, int flags, int startId) {
         if (Objects.equals(intent.getAction(), ACTION_STOP_VM_LAUNCHER_SERVICE)) {
-            // If there is no Debian service or it fails to shutdown, just stop the service.
-            if (mDebianService == null || !mDebianService.shutdownDebian()) {
+
+            if (mDebianService != null && mDebianService.shutdownDebian()) {
+                // During shutdown, change the notification content to indicate that it's closing
+                Notification notification = createNotificationForTerminalClose();
+                getSystemService(NotificationManager.class).notify(this.hashCode(), notification);
+            } else {
+                // If there is no Debian service or it fails to shutdown, just stop the service.
                 stopSelf();
             }
             return START_NOT_STICKY;
@@ -153,7 +160,8 @@
             Log.d(TAG, "VM instance is already started");
             return START_NOT_STICKY;
         }
-        mExecutorService = Executors.newCachedThreadPool();
+        mExecutorService =
+                Executors.newCachedThreadPool(new TerminalThreadFactory(getApplicationContext()));
 
         InstalledImage image = InstalledImage.getDefault(this);
         ConfigJson json = ConfigJson.from(this, image.getConfigPath());
@@ -172,9 +180,7 @@
             android.os.Trace.endSection();
             android.os.Trace.beginAsyncSection("debianBoot", 0);
         } catch (VirtualMachineException e) {
-            Log.e(TAG, "cannot create runner", e);
-            stopSelf();
-            return START_NOT_STICKY;
+            throw new RuntimeException("cannot create runner", e);
         }
         mVirtualMachine = runner.getVm();
         mResultReceiver =
@@ -204,6 +210,32 @@
         return START_NOT_STICKY;
     }
 
+    private Notification createNotificationForTerminalClose() {
+        Intent stopIntent = new Intent();
+        stopIntent.setClass(this, VmLauncherService.class);
+        stopIntent.setAction(VmLauncherService.ACTION_STOP_VM_LAUNCHER_SERVICE);
+        PendingIntent stopPendingIntent =
+                PendingIntent.getService(
+                        this,
+                        0,
+                        stopIntent,
+                        PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
+        Icon icon = Icon.createWithResource(getResources(), R.drawable.ic_launcher_foreground);
+        String stopActionText =
+                getResources().getString(R.string.service_notification_force_quit_action);
+        String stopNotificationTitle =
+                getResources().getString(R.string.service_notification_close_title);
+        return new Notification.Builder(this, this.getPackageName())
+                .setSmallIcon(R.drawable.ic_launcher_foreground)
+                .setContentTitle(stopNotificationTitle)
+                .setOngoing(true)
+                .setSilent(true)
+                .addAction(
+                        new Notification.Action.Builder(icon, stopActionText, stopPendingIntent)
+                                .build())
+                .build();
+    }
+
     private boolean overrideConfigIfNecessary(VirtualMachineCustomImageConfig.Builder builder) {
         boolean changed = false;
         // TODO: check if ANGLE is enabled for the app.
diff --git a/android/TerminalApp/res/values/strings.xml b/android/TerminalApp/res/values/strings.xml
index 20fd95d..d21ded1 100644
--- a/android/TerminalApp/res/values/strings.xml
+++ b/android/TerminalApp/res/values/strings.xml
@@ -153,6 +153,11 @@
     <!-- Notification action button for closing the virtual machine [CHAR LIMIT=20] -->
     <string name="service_notification_quit_action">Close</string>
 
+    <!-- Notification title for foreground service notification during closing [CHAR LIMIT=none] -->
+    <string name="service_notification_close_title">Terminal is closing</string>
+    <!-- Notification action button for force-closing the virtual machine [CHAR LIMIT=30] -->
+    <string name="service_notification_force_quit_action">Force close</string>
+
     <!-- This string is for toast message to notify that VirGL is enabled. [CHAR LIMIT=40] -->
     <string name="virgl_enabled"><xliff:g>VirGL</xliff:g> is enabled</string>
 </resources>
diff --git a/android/virtmgr/src/aidl.rs b/android/virtmgr/src/aidl.rs
index e9074c6..82a5573 100644
--- a/android/virtmgr/src/aidl.rs
+++ b/android/virtmgr/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::{AudioConfig, CrosvmConfig, DiskFile, SharedPathConfig, DisplayConfig, GpuConfig, InputDeviceOption, PayloadState, UsbConfig, VmContext, VmInstance, VmState};
-use crate::debug_config::DebugConfig;
+use crate::debug_config::{DebugConfig, DebugPolicy};
 use crate::dt_overlay::{create_device_tree_overlay, VM_DT_OVERLAY_MAX_SIZE, VM_DT_OVERLAY_PATH};
 use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images, add_microdroid_vendor_image};
 use crate::selinux::{check_tee_service_permission, getfilecon, getprevcon, SeContext};
@@ -319,6 +319,12 @@
         Ok(Vec::from_iter(SUPPORTED_OS_NAMES.iter().cloned()))
     }
 
+    /// Get printable debug policy for testing and debugging
+    fn getDebugPolicy(&self) -> binder::Result<String> {
+        let debug_policy = DebugPolicy::from_host();
+        Ok(format!("{debug_policy:?}"))
+    }
+
     /// Returns whether given feature is enabled
     fn isFeatureEnabled(&self, feature: &str) -> binder::Result<bool> {
         check_manage_access()?;
diff --git a/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index 0c3f6b7..169c3dc 100644
--- a/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/android/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -82,6 +82,11 @@
      */
     String[] getSupportedOSList();
 
+    /**
+     * Get installed debug policy for test and debugging purpose.
+     */
+    String getDebugPolicy();
+
     /** Returns whether given feature is enabled. */
     boolean isFeatureEnabled(in String feature);
 
diff --git a/android/vm/src/main.rs b/android/vm/src/main.rs
index 7bfd957..830d56c 100644
--- a/android/vm/src/main.rs
+++ b/android/vm/src/main.rs
@@ -483,6 +483,9 @@
     let os_list = get_service()?.getSupportedOSList()?;
     println!("Available OS list: {}", serde_json::to_string(&os_list)?);
 
+    let debug_policy = get_service()?.getDebugPolicy()?;
+    println!("Debug policy: {}", debug_policy);
+
     Ok(())
 }
 
diff --git a/build/microdroid/Android.bp b/build/microdroid/Android.bp
index 68b715d..dea0bf3 100644
--- a/build/microdroid/Android.bp
+++ b/build/microdroid/Android.bp
@@ -598,6 +598,12 @@
             src: ":microdroid_kernel_prebuilt-x86_64",
         },
     },
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "16",
+        },
+    ],
     include_descriptors_from_images: [
         ":microdroid_16k_initrd_normal_hashdesc",
         ":microdroid_16k_initrd_debug_hashdesc",
diff --git a/guest/pvmfw/README.md b/guest/pvmfw/README.md
index 8c8314d..766a923 100644
--- a/guest/pvmfw/README.md
+++ b/guest/pvmfw/README.md
@@ -461,6 +461,7 @@
   - `secretkeeper_protection`: pvmfw defers rollback protection to the guest
   - `supports_uefi_boot`: pvmfw boots the VM as a EFI payload (experimental)
   - `trusty_security_vm`: pvmfw skips rollback protection
+- `"com.android.virt.page_size"`: the guest page size in KiB (optional, defaults to 4)
 
 ## Development
 
diff --git a/guest/pvmfw/avb/Android.bp b/guest/pvmfw/avb/Android.bp
index a1ee626..0294322 100644
--- a/guest/pvmfw/avb/Android.bp
+++ b/guest/pvmfw/avb/Android.bp
@@ -37,10 +37,16 @@
         ":test_image_with_one_hashdesc",
         ":test_image_with_non_initrd_hashdesc",
         ":test_image_with_initrd_and_non_initrd_desc",
-        ":test_image_with_prop_desc",
+        ":test_image_with_invalid_page_size",
+        ":test_image_with_negative_page_size",
+        ":test_image_with_overflow_page_size",
+        ":test_image_with_0k_page_size",
+        ":test_image_with_1k_page_size",
+        ":test_image_with_4k_page_size",
+        ":test_image_with_9k_page_size",
+        ":test_image_with_16k_page_size",
         ":test_image_with_service_vm_prop",
         ":test_image_with_unknown_vm_type_prop",
-        ":test_image_with_multiple_props",
         ":test_image_with_duplicated_capability",
         ":test_image_with_rollback_index_5",
         ":test_image_with_multiple_capabilities",
@@ -117,15 +123,113 @@
 }
 
 avb_add_hash_footer {
-    name: "test_image_with_prop_desc",
+    name: "test_image_with_invalid_page_size",
     src: ":unsigned_test_image",
     partition_name: "boot",
     private_key: ":pvmfw_sign_key",
     salt: "2134",
     props: [
         {
-            name: "mock_prop",
-            value: "3333",
+            name: "com.android.virt.page_size",
+            value: "invalid",
+        },
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_negative_page_size",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "-16",
+        },
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_overflow_page_size",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "18014398509481983",
+        },
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_0k_page_size",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "0",
+        },
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_1k_page_size",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "1",
+        },
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_4k_page_size",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "4",
+        },
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_9k_page_size",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "9",
+        },
+    ],
+}
+
+avb_add_hash_footer {
+    name: "test_image_with_16k_page_size",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.page_size",
+            value: "16",
         },
     ],
 }
@@ -159,24 +263,6 @@
 }
 
 avb_add_hash_footer {
-    name: "test_image_with_multiple_props",
-    src: ":unsigned_test_image",
-    partition_name: "boot",
-    private_key: ":pvmfw_sign_key",
-    salt: "2133",
-    props: [
-        {
-            name: "com.android.virt.cap",
-            value: "remote_attest",
-        },
-        {
-            name: "another_vm_type",
-            value: "foo_vm",
-        },
-    ],
-}
-
-avb_add_hash_footer {
     name: "test_image_with_duplicated_capability",
     src: ":unsigned_test_image",
     partition_name: "boot",
diff --git a/guest/pvmfw/avb/src/error.rs b/guest/pvmfw/avb/src/error.rs
index 2e1950a..1307e15 100644
--- a/guest/pvmfw/avb/src/error.rs
+++ b/guest/pvmfw/avb/src/error.rs
@@ -28,6 +28,8 @@
     InvalidDescriptors(DescriptorError),
     /// Unknown vbmeta property.
     UnknownVbmetaProperty,
+    /// VBMeta has invalid page_size property.
+    InvalidPageSize,
 }
 
 impl From<SlotVerifyError<'_>> for PvmfwVerifyError {
@@ -51,6 +53,7 @@
                 write!(f, "VBMeta has invalid descriptors. Error: {:?}", e)
             }
             Self::UnknownVbmetaProperty => write!(f, "Unknown vbmeta property"),
+            Self::InvalidPageSize => write!(f, "Invalid page_size property"),
         }
     }
 }
diff --git a/guest/pvmfw/avb/src/verify.rs b/guest/pvmfw/avb/src/verify.rs
index a073502..8810696 100644
--- a/guest/pvmfw/avb/src/verify.rs
+++ b/guest/pvmfw/avb/src/verify.rs
@@ -17,12 +17,12 @@
 use crate::ops::{Ops, Payload};
 use crate::partition::PartitionName;
 use crate::PvmfwVerifyError;
-use alloc::vec;
 use alloc::vec::Vec;
 use avb::{
-    Descriptor, DescriptorError, DescriptorResult, HashDescriptor, PartitionData,
-    PropertyDescriptor, SlotVerifyError, SlotVerifyNoDataResult, VbmetaData,
+    Descriptor, DescriptorError, DescriptorResult, HashDescriptor, PartitionData, SlotVerifyError,
+    SlotVerifyNoDataResult, VbmetaData,
 };
+use core::str;
 
 // We use this for the rollback_index field if SlotVerifyData has empty rollback_indexes
 const DEFAULT_ROLLBACK_INDEX: u64 = 0;
@@ -45,6 +45,8 @@
     pub capabilities: Vec<Capability>,
     /// Rollback index of kernel.
     pub rollback_index: u64,
+    /// Page size of kernel, if present.
+    pub page_size: Option<usize>,
 }
 
 impl VerifiedBootData<'_> {
@@ -91,14 +93,14 @@
 
     /// Returns the capabilities indicated in `descriptor`, or error if the descriptor has
     /// unexpected contents.
-    fn get_capabilities(descriptor: &PropertyDescriptor) -> Result<Vec<Self>, PvmfwVerifyError> {
-        if descriptor.key != Self::KEY {
-            return Err(PvmfwVerifyError::UnknownVbmetaProperty);
-        }
+    fn get_capabilities(vbmeta_data: &VbmetaData) -> Result<Vec<Self>, PvmfwVerifyError> {
+        let Some(value) = vbmeta_data.get_property_value(Self::KEY) else {
+            return Ok(Vec::new());
+        };
 
         let mut res = Vec::new();
 
-        for v in descriptor.value.split(|b| *b == Self::SEPARATOR) {
+        for v in value.split(|b| *b == Self::SEPARATOR) {
             let cap = match v {
                 Self::REMOTE_ATTEST => Self::RemoteAttest,
                 Self::TRUSTY_SECURITY_VM => Self::TrustySecurityVm,
@@ -153,30 +155,6 @@
     }
 }
 
-/// Verifies that the vbmeta contains at most one property descriptor and it indicates the
-/// vm type is service VM.
-fn verify_property_and_get_capabilities(
-    descriptors: &[Descriptor],
-) -> Result<Vec<Capability>, PvmfwVerifyError> {
-    let mut iter = descriptors.iter().filter_map(|d| match d {
-        Descriptor::Property(p) => Some(p),
-        _ => None,
-    });
-
-    let descriptor = match iter.next() {
-        // No property descriptors -> no capabilities.
-        None => return Ok(vec![]),
-        Some(d) => d,
-    };
-
-    // Multiple property descriptors -> error.
-    if iter.next().is_some() {
-        return Err(DescriptorError::InvalidContents.into());
-    }
-
-    Capability::get_capabilities(descriptor)
-}
-
 /// Hash descriptors extracted from a vbmeta image.
 ///
 /// We always have a kernel hash descriptor and may have initrd normal or debug descriptors.
@@ -243,6 +221,23 @@
     Ok(digest)
 }
 
+/// Returns the indicated payload page size, if present.
+fn read_page_size(vbmeta_data: &VbmetaData) -> Result<Option<usize>, PvmfwVerifyError> {
+    let Some(property) = vbmeta_data.get_property_value("com.android.virt.page_size") else {
+        return Ok(None);
+    };
+    let size = str::from_utf8(property)
+        .or(Err(PvmfwVerifyError::InvalidPageSize))?
+        .parse::<usize>()
+        .or(Err(PvmfwVerifyError::InvalidPageSize))?
+        .checked_mul(1024)
+        // TODO(stable(unsigned_is_multiple_of)): use .is_multiple_of()
+        .filter(|sz| sz % (4 << 10) == 0 && *sz != 0)
+        .ok_or(PvmfwVerifyError::InvalidPageSize)?;
+
+    Ok(Some(size))
+}
+
 /// Verifies the given initrd partition, and checks that the resulting contents looks like expected.
 fn verify_initrd(
     ops: &mut Ops,
@@ -278,7 +273,8 @@
     verify_vbmeta_is_from_kernel_partition(vbmeta_image)?;
     let descriptors = vbmeta_image.descriptors()?;
     let hash_descriptors = HashDescriptors::get(&descriptors)?;
-    let capabilities = verify_property_and_get_capabilities(&descriptors)?;
+    let capabilities = Capability::get_capabilities(vbmeta_image)?;
+    let page_size = read_page_size(vbmeta_image)?;
 
     if initrd.is_none() {
         hash_descriptors.verify_no_initrd()?;
@@ -289,6 +285,7 @@
             public_key: trusted_public_key,
             capabilities,
             rollback_index,
+            page_size,
         });
     }
 
@@ -309,5 +306,6 @@
         public_key: trusted_public_key,
         capabilities,
         rollback_index,
+        page_size,
     })
 }
diff --git a/guest/pvmfw/avb/tests/api_test.rs b/guest/pvmfw/avb/tests/api_test.rs
index 430c4b3..0ed0279 100644
--- a/guest/pvmfw/avb/tests/api_test.rs
+++ b/guest/pvmfw/avb/tests/api_test.rs
@@ -28,11 +28,17 @@
 use utils::*;
 
 const TEST_IMG_WITH_ONE_HASHDESC_PATH: &str = "test_image_with_one_hashdesc.img";
+const TEST_IMG_WITH_INVALID_PAGE_SIZE_PATH: &str = "test_image_with_invalid_page_size.img";
+const TEST_IMG_WITH_NEGATIVE_PAGE_SIZE_PATH: &str = "test_image_with_negative_page_size.img";
+const TEST_IMG_WITH_OVERFLOW_PAGE_SIZE_PATH: &str = "test_image_with_overflow_page_size.img";
+const TEST_IMG_WITH_0K_PAGE_SIZE_PATH: &str = "test_image_with_0k_page_size.img";
+const TEST_IMG_WITH_1K_PAGE_SIZE_PATH: &str = "test_image_with_1k_page_size.img";
+const TEST_IMG_WITH_4K_PAGE_SIZE_PATH: &str = "test_image_with_4k_page_size.img";
+const TEST_IMG_WITH_9K_PAGE_SIZE_PATH: &str = "test_image_with_9k_page_size.img";
+const TEST_IMG_WITH_16K_PAGE_SIZE_PATH: &str = "test_image_with_16k_page_size.img";
 const TEST_IMG_WITH_ROLLBACK_INDEX_5: &str = "test_image_with_rollback_index_5.img";
-const TEST_IMG_WITH_PROP_DESC_PATH: &str = "test_image_with_prop_desc.img";
 const TEST_IMG_WITH_SERVICE_VM_PROP_PATH: &str = "test_image_with_service_vm_prop.img";
 const TEST_IMG_WITH_UNKNOWN_VM_TYPE_PROP_PATH: &str = "test_image_with_unknown_vm_type_prop.img";
-const TEST_IMG_WITH_MULTIPLE_PROPS_PATH: &str = "test_image_with_multiple_props.img";
 const TEST_IMG_WITH_DUPLICATED_CAP_PATH: &str = "test_image_with_duplicated_capability.img";
 const TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH: &str = "test_image_with_non_initrd_hashdesc.img";
 const TEST_IMG_WITH_INITRD_AND_NON_INITRD_DESC_PATH: &str =
@@ -51,6 +57,7 @@
         &load_latest_initrd_normal()?,
         b"initrd_normal",
         DebugLevel::None,
+        None,
     )
 }
 
@@ -63,6 +70,7 @@
         salt,
         expected_rollback_index,
         vec![Capability::TrustySecurityVm],
+        None,
     )
 }
 
@@ -72,6 +80,7 @@
         &load_latest_initrd_debug()?,
         b"initrd_debug",
         DebugLevel::Full,
+        None,
     )
 }
 
@@ -93,6 +102,7 @@
         public_key: &public_key,
         capabilities: vec![],
         rollback_index: 0,
+        page_size: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -137,6 +147,7 @@
         public_key: &public_key,
         capabilities: vec![Capability::RemoteAttest],
         rollback_index: 0,
+        page_size: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -154,16 +165,6 @@
 }
 
 #[test]
-fn payload_with_multiple_props_fails_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_fails(
-        &fs::read(TEST_IMG_WITH_MULTIPLE_PROPS_PATH)?,
-        /* initrd= */ None,
-        &load_trusted_public_key()?,
-        PvmfwVerifyError::InvalidDescriptors(DescriptorError::InvalidContents),
-    )
-}
-
-#[test]
 fn payload_with_duplicated_capability_fails_verification_with_no_initrd() -> Result<()> {
     assert_payload_verification_fails(
         &fs::read(TEST_IMG_WITH_DUPLICATED_CAP_PATH)?,
@@ -174,16 +175,6 @@
 }
 
 #[test]
-fn payload_with_prop_descriptor_fails_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_fails(
-        &fs::read(TEST_IMG_WITH_PROP_DESC_PATH)?,
-        /* initrd= */ None,
-        &load_trusted_public_key()?,
-        PvmfwVerifyError::UnknownVbmetaProperty,
-    )
-}
-
-#[test]
 fn payload_expecting_initrd_fails_verification_with_no_initrd() -> Result<()> {
     assert_payload_verification_fails(
         &load_latest_signed_kernel()?,
@@ -257,6 +248,60 @@
 }
 
 #[test]
+fn kernel_has_expected_page_size_invalid() {
+    let kernel = fs::read(TEST_IMG_WITH_INVALID_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Err(PvmfwVerifyError::InvalidPageSize));
+}
+
+#[test]
+fn kernel_has_expected_page_size_negative() {
+    let kernel = fs::read(TEST_IMG_WITH_NEGATIVE_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Err(PvmfwVerifyError::InvalidPageSize));
+}
+
+#[test]
+fn kernel_has_expected_page_size_overflow() {
+    let kernel = fs::read(TEST_IMG_WITH_OVERFLOW_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Err(PvmfwVerifyError::InvalidPageSize));
+}
+
+#[test]
+fn kernel_has_expected_page_size_none() {
+    let kernel = fs::read(TEST_IMG_WITH_ONE_HASHDESC_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Ok(None));
+}
+
+#[test]
+fn kernel_has_expected_page_size_0k() {
+    let kernel = fs::read(TEST_IMG_WITH_0K_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Err(PvmfwVerifyError::InvalidPageSize));
+}
+
+#[test]
+fn kernel_has_expected_page_size_1k() {
+    let kernel = fs::read(TEST_IMG_WITH_1K_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Err(PvmfwVerifyError::InvalidPageSize));
+}
+
+#[test]
+fn kernel_has_expected_page_size_4k() {
+    let kernel = fs::read(TEST_IMG_WITH_4K_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Ok(Some(4usize << 10)));
+}
+
+#[test]
+fn kernel_has_expected_page_size_9k() {
+    let kernel = fs::read(TEST_IMG_WITH_9K_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Err(PvmfwVerifyError::InvalidPageSize));
+}
+
+#[test]
+fn kernel_has_expected_page_size_16k() {
+    let kernel = fs::read(TEST_IMG_WITH_16K_PAGE_SIZE_PATH).unwrap();
+    assert_eq!(read_page_size(&kernel), Ok(Some(16usize << 10)));
+}
+
+#[test]
 fn kernel_footer_with_vbmeta_offset_overwritten_fails_verification() -> Result<()> {
     // Arrange.
     let mut kernel = load_latest_signed_kernel()?;
@@ -412,6 +457,7 @@
         public_key: &public_key,
         capabilities: vec![],
         rollback_index: 5,
+        page_size: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
     Ok(())
diff --git a/guest/pvmfw/avb/tests/utils.rs b/guest/pvmfw/avb/tests/utils.rs
index 61bfbf2..79552b5 100644
--- a/guest/pvmfw/avb/tests/utils.rs
+++ b/guest/pvmfw/avb/tests/utils.rs
@@ -114,6 +114,7 @@
     initrd: &[u8],
     initrd_salt: &[u8],
     expected_debug_level: DebugLevel,
+    page_size: Option<usize>,
 ) -> Result<()> {
     let public_key = load_trusted_public_key()?;
     let kernel = load_latest_signed_kernel()?;
@@ -133,6 +134,7 @@
         public_key: &public_key,
         capabilities,
         rollback_index: if cfg!(llpvm_changes) { 1 } else { 0 },
+        page_size,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -144,6 +146,7 @@
     salt: &[u8],
     expected_rollback_index: u64,
     capabilities: Vec<Capability>,
+    page_size: Option<usize>,
 ) -> Result<()> {
     let public_key = load_trusted_public_key()?;
     let verified_boot_data = verify_payload(
@@ -163,12 +166,23 @@
         public_key: &public_key,
         capabilities,
         rollback_index: expected_rollback_index,
+        page_size,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
     Ok(())
 }
 
+pub fn read_page_size(kernel: &[u8]) -> Result<Option<usize>, PvmfwVerifyError> {
+    let public_key = load_trusted_public_key().unwrap();
+    let verified_boot_data = verify_payload(
+        kernel,
+        None, // initrd
+        &public_key,
+    )?;
+    Ok(verified_boot_data.page_size)
+}
+
 pub fn hash(inputs: &[&[u8]]) -> Digest {
     let mut digester = sha::Sha256::new();
     inputs.iter().for_each(|input| digester.update(input));
diff --git a/guest/pvmfw/src/dice.rs b/guest/pvmfw/src/dice.rs
index b597309..6694881 100644
--- a/guest/pvmfw/src/dice.rs
+++ b/guest/pvmfw/src/dice.rs
@@ -200,6 +200,7 @@
         public_key: b"public key",
         capabilities: vec![],
         rollback_index: 42,
+        page_size: None,
     };
     const HASH: Hash = *b"sixtyfourbyteslongsentencearerarebutletsgiveitatrycantbethathard";
 
diff --git a/guest/pvmfw/src/entry.rs b/guest/pvmfw/src/entry.rs
index 343c2fc..7c46515 100644
--- a/guest/pvmfw/src/entry.rs
+++ b/guest/pvmfw/src/entry.rs
@@ -74,22 +74,36 @@
 configure_heap!(SIZE_128KB);
 limit_stack_size!(SIZE_4KB * 12);
 
+#[derive(Debug)]
+enum NextStage {
+    LinuxBoot(usize),
+    LinuxBootWithUart(usize),
+}
+
 /// Entry point for pVM firmware.
 pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
-    // Limitations in this function:
-    // - can't access non-pvmfw memory (only statically-mapped memory)
-    // - can't access MMIO (except the console, already configured by vmbase)
+    let fdt_address = fdt_address.try_into().unwrap();
+    let payload_start = payload_start.try_into().unwrap();
+    let payload_size = payload_size.try_into().unwrap();
 
-    match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
-        Ok((entry, bcc, keep_uart)) => {
-            jump_to_payload(fdt_address, entry.try_into().unwrap(), bcc, keep_uart)
-        }
-        Err(e) => {
-            const REBOOT_REASON_CONSOLE: usize = 1;
-            console_writeln!(REBOOT_REASON_CONSOLE, "{}", e.as_avf_reboot_string());
-            reboot()
-        }
-    }
+    let reboot_reason = match main_wrapper(fdt_address, payload_start, payload_size) {
+        Err(r) => r,
+        Ok((next_stage, bcc)) => match next_stage {
+            NextStage::LinuxBootWithUart(ep) => jump_to_payload(fdt_address, ep, bcc),
+            NextStage::LinuxBoot(ep) => {
+                if let Err(e) = unshare_uart() {
+                    error!("Failed to unmap UART: {e}");
+                    RebootReason::InternalError
+                } else {
+                    jump_to_payload(fdt_address, ep, bcc)
+                }
+            }
+        },
+    };
+
+    const REBOOT_REASON_CONSOLE: usize = 1;
+    console_writeln!(REBOOT_REASON_CONSOLE, "{}", reboot_reason.as_avf_reboot_string());
+    reboot()
 
     // if we reach this point and return, vmbase::entry::rust_entry() will call power::shutdown().
 }
@@ -102,7 +116,7 @@
     fdt: usize,
     payload: usize,
     payload_size: usize,
-) -> Result<(usize, Range<usize>, bool), RebootReason> {
+) -> Result<(NextStage, Range<usize>), RebootReason> {
     // Limitations in this function:
     // - only access MMIO once (and while) it has been mapped and configured
     // - only perform logging once the logger has been initialized
@@ -122,13 +136,7 @@
 
     let config_entries = appended.get_entries();
 
-    let slices = memory::MemorySlices::new(
-        fdt,
-        payload,
-        payload_size,
-        config_entries.vm_dtbo,
-        config_entries.vm_ref_dt,
-    )?;
+    let slices = memory::MemorySlices::new(fdt, payload, payload_size)?;
 
     // This wrapper allows main() to be blissfully ignorant of platform details.
     let (next_bcc, debuggable_payload) = crate::main(
@@ -137,6 +145,8 @@
         slices.ramdisk,
         config_entries.bcc,
         config_entries.debug_policy,
+        config_entries.vm_dtbo,
+        config_entries.vm_ref_dt,
     )?;
     // Keep UART MMIO_GUARD-ed for debuggable payloads, to enable earlycon.
     let keep_uart = cfg!(debuggable_vms_improvements) && debuggable_payload;
@@ -150,14 +160,20 @@
     })?;
     unshare_all_memory();
 
-    Ok((slices.kernel.as_ptr() as usize, next_bcc, keep_uart))
+    let next_stage = select_next_stage(slices.kernel, keep_uart);
+
+    Ok((next_stage, next_bcc))
 }
 
-fn jump_to_payload(fdt_address: u64, payload_start: u64, bcc: Range<usize>, keep_uart: bool) -> ! {
-    if !keep_uart {
-        unshare_uart().unwrap();
+fn select_next_stage(kernel: &[u8], keep_uart: bool) -> NextStage {
+    if keep_uart {
+        NextStage::LinuxBootWithUart(kernel.as_ptr() as _)
+    } else {
+        NextStage::LinuxBoot(kernel.as_ptr() as _)
     }
+}
 
+fn jump_to_payload(fdt_address: usize, payload_start: usize, bcc: Range<usize>) -> ! {
     deactivate_dynamic_page_tables();
 
     const ASM_STP_ALIGN: usize = size_of::<u64>() * 2;
@@ -296,8 +312,8 @@
             eh_stack = in(reg) u64::try_from(eh_stack.start.0).unwrap(),
             eh_stack_end = in(reg) u64::try_from(eh_stack.end.0).unwrap(),
             dcache_line_size = in(reg) u64::try_from(min_dcache_line_size()).unwrap(),
-            in("x0") fdt_address,
-            in("x30") payload_start,
+            in("x0") u64::try_from(fdt_address).unwrap(),
+            in("x30") u64::try_from(payload_start).unwrap(),
             options(noreturn),
         );
     };
diff --git a/guest/pvmfw/src/fdt.rs b/guest/pvmfw/src/fdt.rs
index 027f163..bfbd2e6 100644
--- a/guest/pvmfw/src/fdt.rs
+++ b/guest/pvmfw/src/fdt.rs
@@ -16,7 +16,6 @@
 
 use crate::bootargs::BootArgsIterator;
 use crate::device_assignment::{self, DeviceAssignmentInfo, VmDtbo};
-use crate::helpers::GUEST_PAGE_SIZE;
 use crate::Box;
 use crate::RebootReason;
 use alloc::collections::BTreeMap;
@@ -83,7 +82,7 @@
 
 /// Extract from /config the address range containing the pre-loaded kernel. Absence of /config is
 /// not an error.
-fn read_kernel_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
+pub fn read_kernel_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
     let addr = cstr!("kernel-address");
     let size = cstr!("kernel-size");
 
@@ -101,7 +100,7 @@
 
 /// Extract from /chosen the address range containing the pre-loaded ramdisk. Absence is not an
 /// error as there can be initrd-less VM.
-fn read_initrd_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
+pub fn read_initrd_range_from(fdt: &Fdt) -> libfdt::Result<Option<Range<usize>>> {
     let start = cstr!("linux,initrd-start");
     let end = cstr!("linux,initrd-end");
 
@@ -147,7 +146,10 @@
 /// Reads and validates the memory range in the DT.
 ///
 /// Only one memory range is expected with the crosvm setup for now.
-fn read_and_validate_memory_range(fdt: &Fdt) -> Result<Range<usize>, RebootReason> {
+fn read_and_validate_memory_range(
+    fdt: &Fdt,
+    guest_page_size: usize,
+) -> Result<Range<usize>, RebootReason> {
     let mut memory = fdt.memory().map_err(|e| {
         error!("Failed to read memory range from DT: {e}");
         RebootReason::InvalidFdt
@@ -169,8 +171,8 @@
     }
 
     let size = range.len();
-    if size % GUEST_PAGE_SIZE != 0 {
-        error!("Memory size {:#x} is not a multiple of page size {:#x}", size, GUEST_PAGE_SIZE);
+    if size % guest_page_size != 0 {
+        error!("Memory size {:#x} is not a multiple of page size {:#x}", size, guest_page_size);
         return Err(RebootReason::InvalidFdt);
     }
 
@@ -854,16 +856,17 @@
 fn validate_swiotlb_info(
     swiotlb_info: &SwiotlbInfo,
     memory: &Range<usize>,
+    guest_page_size: usize,
 ) -> Result<(), RebootReason> {
     let size = swiotlb_info.size;
     let align = swiotlb_info.align;
 
-    if size == 0 || (size % GUEST_PAGE_SIZE) != 0 {
+    if size == 0 || (size % guest_page_size) != 0 {
         error!("Invalid swiotlb size {:#x}", size);
         return Err(RebootReason::InvalidFdt);
     }
 
-    if let Some(align) = align.filter(|&a| a % GUEST_PAGE_SIZE != 0) {
+    if let Some(align) = align.filter(|&a| a % guest_page_size != 0) {
         error!("Invalid swiotlb alignment {:#x}", align);
         return Err(RebootReason::InvalidFdt);
     }
@@ -989,7 +992,6 @@
 
 #[derive(Debug)]
 pub struct DeviceTreeInfo {
-    pub kernel_range: Option<Range<usize>>,
     pub initrd_range: Option<Range<usize>>,
     pub memory_range: Range<usize>,
     bootargs: Option<CString>,
@@ -1015,15 +1017,11 @@
 }
 
 pub fn sanitize_device_tree(
-    fdt: &mut [u8],
+    fdt: &mut Fdt,
     vm_dtbo: Option<&mut [u8]>,
     vm_ref_dt: Option<&[u8]>,
+    guest_page_size: usize,
 ) -> Result<DeviceTreeInfo, RebootReason> {
-    let fdt = Fdt::from_mut_slice(fdt).map_err(|e| {
-        error!("Failed to load FDT: {e}");
-        RebootReason::InvalidFdt
-    })?;
-
     let vm_dtbo = match vm_dtbo {
         Some(vm_dtbo) => Some(VmDtbo::from_mut_slice(vm_dtbo).map_err(|e| {
             error!("Failed to load VM DTBO: {e}");
@@ -1032,7 +1030,7 @@
         None => None,
     };
 
-    let info = parse_device_tree(fdt, vm_dtbo.as_deref())?;
+    let info = parse_device_tree(fdt, vm_dtbo.as_deref(), guest_page_size)?;
 
     fdt.clone_from(FDT_TEMPLATE).map_err(|e| {
         error!("Failed to instantiate FDT from the template DT: {e}");
@@ -1085,18 +1083,17 @@
     Ok(info)
 }
 
-fn parse_device_tree(fdt: &Fdt, vm_dtbo: Option<&VmDtbo>) -> Result<DeviceTreeInfo, RebootReason> {
-    let kernel_range = read_kernel_range_from(fdt).map_err(|e| {
-        error!("Failed to read kernel range from DT: {e}");
-        RebootReason::InvalidFdt
-    })?;
-
+fn parse_device_tree(
+    fdt: &Fdt,
+    vm_dtbo: Option<&VmDtbo>,
+    guest_page_size: usize,
+) -> Result<DeviceTreeInfo, RebootReason> {
     let initrd_range = read_initrd_range_from(fdt).map_err(|e| {
         error!("Failed to read initrd range from DT: {e}");
         RebootReason::InvalidFdt
     })?;
 
-    let memory_range = read_and_validate_memory_range(fdt)?;
+    let memory_range = read_and_validate_memory_range(fdt, guest_page_size)?;
 
     let bootargs = read_bootargs_from(fdt).map_err(|e| {
         error!("Failed to read bootargs from DT: {e}");
@@ -1149,7 +1146,7 @@
             error!("Swiotlb info missing from DT");
             RebootReason::InvalidFdt
         })?;
-    validate_swiotlb_info(&swiotlb_info, &memory_range)?;
+    validate_swiotlb_info(&swiotlb_info, &memory_range, guest_page_size)?;
 
     let device_assignment = match vm_dtbo {
         Some(vm_dtbo) => {
@@ -1194,7 +1191,6 @@
     })?;
 
     Ok(DeviceTreeInfo {
-        kernel_range,
         initrd_range,
         memory_range,
         bootargs,
diff --git a/guest/pvmfw/src/helpers.rs b/guest/pvmfw/src/helpers.rs
deleted file mode 100644
index 0552640..0000000
--- a/guest/pvmfw/src/helpers.rs
+++ /dev/null
@@ -1,19 +0,0 @@
-// Copyright 2022, 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.
-
-//! Miscellaneous helper functions.
-
-use vmbase::memory::SIZE_4KB;
-
-pub const GUEST_PAGE_SIZE: usize = SIZE_4KB;
diff --git a/guest/pvmfw/src/main.rs b/guest/pvmfw/src/main.rs
index bde03ff..d04db06 100644
--- a/guest/pvmfw/src/main.rs
+++ b/guest/pvmfw/src/main.rs
@@ -28,18 +28,15 @@
 mod exceptions;
 mod fdt;
 mod gpt;
-mod helpers;
 mod instance;
 mod memory;
+mod rollback;
 
 use crate::bcc::Bcc;
 use crate::dice::PartialInputs;
 use crate::entry::RebootReason;
-use crate::fdt::modify_for_next_stage;
-use crate::helpers::GUEST_PAGE_SIZE;
-use crate::instance::EntryBody;
-use crate::instance::Error as InstanceError;
-use crate::instance::{get_recorded_entry, record_instance_entry};
+use crate::fdt::{modify_for_next_stage, sanitize_device_tree};
+use crate::rollback::perform_rollback_protection;
 use alloc::borrow::Cow;
 use alloc::boxed::Box;
 use bssl_avf::Digester;
@@ -49,26 +46,25 @@
 use libfdt::{Fdt, FdtNode};
 use log::{debug, error, info, trace, warn};
 use pvmfw_avb::verify_payload;
-use pvmfw_avb::Capability;
 use pvmfw_avb::DebugLevel;
 use pvmfw_embedded_key::PUBLIC_KEY;
 use vmbase::fdt::pci::{PciError, PciInfo};
 use vmbase::heap;
-use vmbase::memory::flush;
+use vmbase::memory::{flush, init_shared_pool, SIZE_4KB};
 use vmbase::rand;
 use vmbase::virtio::pci;
 
-const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
-
 fn main(
-    fdt: &mut Fdt,
+    untrusted_fdt: &mut Fdt,
     signed_kernel: &[u8],
     ramdisk: Option<&[u8]>,
     current_bcc_handover: &[u8],
     mut debug_policy: Option<&[u8]>,
+    vm_dtbo: Option<&mut [u8]>,
+    vm_ref_dt: Option<&[u8]>,
 ) -> Result<(Range<usize>, bool), RebootReason> {
     info!("pVM firmware");
-    debug!("FDT: {:?}", fdt.as_ptr());
+    debug!("FDT: {:?}", untrusted_fdt.as_ptr());
     debug!("Signed kernel: {:?} ({:#x} bytes)", signed_kernel.as_ptr(), signed_kernel.len());
     debug!("AVB public key: addr={:?}, size={:#x} ({1})", PUBLIC_KEY.as_ptr(), PUBLIC_KEY.len());
     if let Some(rd) = ramdisk {
@@ -97,14 +93,6 @@
         debug_policy = None;
     }
 
-    // Set up PCI bus for VirtIO devices.
-    let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
-    debug!("PCI: {:#x?}", pci_info);
-    let mut pci_root = pci::initialize(pci_info).map_err(|e| {
-        error!("Failed to initialize PCI: {e}");
-        RebootReason::InternalError
-    })?;
-
     let verified_boot_data = verify_payload(signed_kernel, ramdisk, PUBLIC_KEY).map_err(|e| {
         error!("Failed to verify the payload: {e}");
         RebootReason::PayloadVerificationError
@@ -115,7 +103,23 @@
         info!("Please disregard any previous libavb ERROR about initrd_normal.");
     }
 
-    let next_bcc = heap::aligned_boxed_slice(NEXT_BCC_SIZE, GUEST_PAGE_SIZE).ok_or_else(|| {
+    let guest_page_size = verified_boot_data.page_size.unwrap_or(SIZE_4KB);
+    let fdt_info = sanitize_device_tree(untrusted_fdt, vm_dtbo, vm_ref_dt, guest_page_size)?;
+    let fdt = untrusted_fdt; // DT has now been sanitized.
+    let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
+    debug!("PCI: {:#x?}", pci_info);
+    // Set up PCI bus for VirtIO devices.
+    let mut pci_root = pci::initialize(pci_info).map_err(|e| {
+        error!("Failed to initialize PCI: {e}");
+        RebootReason::InternalError
+    })?;
+    init_shared_pool(fdt_info.swiotlb_info.fixed_range()).map_err(|e| {
+        error!("Failed to initialize shared pool: {e}");
+        RebootReason::InternalError
+    })?;
+
+    let next_bcc_size = guest_page_size;
+    let next_bcc = heap::aligned_boxed_slice(next_bcc_size, guest_page_size).ok_or_else(|| {
         error!("Failed to allocate the next-stage BCC");
         RebootReason::InternalError
     })?;
@@ -128,65 +132,14 @@
     })?;
 
     let instance_hash = if cfg!(llpvm_changes) { Some(salt_from_instance_id(fdt)?) } else { None };
-    let defer_rollback_protection = should_defer_rollback_protection(fdt)?
-        && verified_boot_data.has_capability(Capability::SecretkeeperProtection);
-    let (new_instance, salt) = if defer_rollback_protection {
-        info!("Guest OS is capable of Secretkeeper protection, deferring rollback protection");
-        // rollback_index of the image is used as security_version and is expected to be > 0 to
-        // discourage implicit allocation.
-        if verified_boot_data.rollback_index == 0 {
-            error!("Expected positive rollback_index, found 0");
-            return Err(RebootReason::InvalidPayload);
-        };
-        (false, instance_hash.unwrap())
-    } else if verified_boot_data.has_capability(Capability::RemoteAttest) {
-        info!("Service VM capable of remote attestation detected, performing version checks");
-        if service_vm_version::VERSION != verified_boot_data.rollback_index {
-            // For RKP VM, we only boot if the version in the AVB footer of its kernel matches
-            // the one embedded in pvmfw at build time.
-            // This prevents the pvmfw from booting a roll backed RKP VM.
-            error!(
-                "Service VM version mismatch: expected {}, found {}",
-                service_vm_version::VERSION,
-                verified_boot_data.rollback_index
-            );
-            return Err(RebootReason::InvalidPayload);
-        }
-        (false, instance_hash.unwrap())
-    } else if verified_boot_data.has_capability(Capability::TrustySecurityVm) {
-        // The rollback protection of Trusty VMs are handled by AuthMgr, so we don't need to
-        // handle it here.
-        info!("Trusty Security VM detected");
-        (false, instance_hash.unwrap())
-    } else {
-        info!("Fallback to instance.img based rollback checks");
-        let (recorded_entry, mut instance_img, header_index) =
-            get_recorded_entry(&mut pci_root, cdi_seal).map_err(|e| {
-                error!("Failed to get entry from instance.img: {e}");
-                RebootReason::InternalError
-            })?;
-        let (new_instance, salt) = if let Some(entry) = recorded_entry {
-            check_dice_measurements_match_entry(&dice_inputs, &entry)?;
-            let salt = instance_hash.unwrap_or(entry.salt);
-            (false, salt)
-        } else {
-            // New instance!
-            let salt = instance_hash.map_or_else(rand::random_array, Ok).map_err(|e| {
-                error!("Failed to generated instance.img salt: {e}");
-                RebootReason::InternalError
-            })?;
-
-            let entry = EntryBody::new(&dice_inputs, &salt);
-            record_instance_entry(&entry, cdi_seal, &mut instance_img, header_index).map_err(
-                |e| {
-                    error!("Failed to get recorded entry in instance.img: {e}");
-                    RebootReason::InternalError
-                },
-            )?;
-            (true, salt)
-        };
-        (new_instance, salt)
-    };
+    let (new_instance, salt, defer_rollback_protection) = perform_rollback_protection(
+        fdt,
+        &verified_boot_data,
+        &dice_inputs,
+        &mut pci_root,
+        cdi_seal,
+        instance_hash,
+    )?;
     trace!("Got salt for instance: {salt:x?}");
 
     let new_bcc_handover = if cfg!(dice_changes) {
@@ -257,36 +210,6 @@
     Ok((bcc_range, debuggable))
 }
 
-fn check_dice_measurements_match_entry(
-    dice_inputs: &PartialInputs,
-    entry: &EntryBody,
-) -> Result<(), RebootReason> {
-    ensure_dice_measurements_match_entry(dice_inputs, entry).map_err(|e| {
-        error!(
-            "Dice measurements do not match recorded entry. \
-        This may be because of update: {e}"
-        );
-        RebootReason::InternalError
-    })?;
-
-    Ok(())
-}
-
-fn ensure_dice_measurements_match_entry(
-    dice_inputs: &PartialInputs,
-    entry: &EntryBody,
-) -> Result<(), InstanceError> {
-    if entry.code_hash != dice_inputs.code_hash {
-        Err(InstanceError::RecordedCodeHashMismatch)
-    } else if entry.auth_hash != dice_inputs.auth_hash {
-        Err(InstanceError::RecordedAuthHashMismatch)
-    } else if entry.mode() != dice_inputs.mode {
-        Err(InstanceError::RecordedDiceModeMismatch)
-    } else {
-        Ok(())
-    }
-}
-
 // Get the "salt" which is one of the input for DICE derivation.
 // This provides differentiation of secrets for different VM instances with same payloads.
 fn salt_from_instance_id(fdt: &Fdt) -> Result<Hidden, RebootReason> {
@@ -314,18 +237,6 @@
     })
 }
 
-fn should_defer_rollback_protection(fdt: &Fdt) -> Result<bool, RebootReason> {
-    let node = avf_untrusted_node(fdt)?;
-    let defer_rbp = node
-        .getprop(cstr!("defer-rollback-protection"))
-        .map_err(|e| {
-            error!("Failed to get defer-rollback-protection property in DT: {e}");
-            RebootReason::InvalidFdt
-        })?
-        .is_some();
-    Ok(defer_rbp)
-}
-
 fn avf_untrusted_node(fdt: &Fdt) -> Result<FdtNode, RebootReason> {
     let node = fdt.node(cstr!("/avf/untrusted")).map_err(|e| {
         error!("Failed to get /avf/untrusted node: {e}");
diff --git a/guest/pvmfw/src/memory.rs b/guest/pvmfw/src/memory.rs
index 35bfd3a..d2f63b5 100644
--- a/guest/pvmfw/src/memory.rs
+++ b/guest/pvmfw/src/memory.rs
@@ -15,7 +15,7 @@
 //! Low-level allocation and tracking of main memory.
 
 use crate::entry::RebootReason;
-use crate::fdt;
+use crate::fdt::{read_initrd_range_from, read_kernel_range_from};
 use core::num::NonZeroUsize;
 use core::slice;
 use log::debug;
@@ -24,7 +24,7 @@
 use log::warn;
 use vmbase::{
     layout::crosvm,
-    memory::{init_shared_pool, map_data, map_rodata, resize_available_memory},
+    memory::{map_data, map_rodata, resize_available_memory},
 };
 
 pub(crate) struct MemorySlices<'a> {
@@ -34,13 +34,7 @@
 }
 
 impl<'a> MemorySlices<'a> {
-    pub fn new(
-        fdt: usize,
-        kernel: usize,
-        kernel_size: usize,
-        vm_dtbo: Option<&mut [u8]>,
-        vm_ref_dt: Option<&[u8]>,
-    ) -> Result<Self, RebootReason> {
+    pub fn new(fdt: usize, kernel: usize, kernel_size: usize) -> Result<Self, RebootReason> {
         let fdt_size = NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap();
         // TODO - Only map the FDT as read-only, until we modify it right before jump_to_payload()
         // e.g. by generating a DTBO for a template DT in main() and, on return, re-map DT as RW,
@@ -51,44 +45,39 @@
         })?;
 
         // SAFETY: map_data validated the range to be in main memory, mapped, and not overlap.
-        let fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, fdt_size.into()) };
-
-        let info = fdt::sanitize_device_tree(fdt, vm_dtbo, vm_ref_dt)?;
-        let fdt = libfdt::Fdt::from_mut_slice(fdt).map_err(|e| {
-            error!("Failed to load sanitized FDT: {e}");
+        let untrusted_fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, fdt_size.into()) };
+        let untrusted_fdt = libfdt::Fdt::from_mut_slice(untrusted_fdt).map_err(|e| {
+            error!("Failed to load input FDT: {e}");
             RebootReason::InvalidFdt
         })?;
-        debug!("Fdt passed validation!");
 
-        let memory_range = info.memory_range;
+        let memory_range = untrusted_fdt.first_memory_range().map_err(|e| {
+            error!("Failed to read memory range from DT: {e}");
+            RebootReason::InvalidFdt
+        })?;
         debug!("Resizing MemoryTracker to range {memory_range:#x?}");
         resize_available_memory(&memory_range).map_err(|e| {
             error!("Failed to use memory range value from DT: {memory_range:#x?}: {e}");
             RebootReason::InvalidFdt
         })?;
 
-        init_shared_pool(info.swiotlb_info.fixed_range()).map_err(|e| {
-            error!("Failed to initialize shared pool: {e}");
-            RebootReason::InternalError
+        let kernel_range = read_kernel_range_from(untrusted_fdt).map_err(|e| {
+            error!("Failed to read kernel range: {e}");
+            RebootReason::InvalidFdt
         })?;
-
-        let (kernel_start, kernel_size) = if let Some(r) = info.kernel_range {
-            let size = r.len().try_into().map_err(|_| {
-                error!("Invalid kernel size: {:#x}", r.len());
-                RebootReason::InternalError
-            })?;
-            (r.start, size)
+        let (kernel_start, kernel_size) = if let Some(r) = kernel_range {
+            (r.start, r.len())
         } else if cfg!(feature = "legacy") {
             warn!("Failed to find the kernel range in the DT; falling back to legacy ABI");
-            let size = NonZeroUsize::new(kernel_size).ok_or_else(|| {
-                error!("Invalid kernel size: {kernel_size:#x}");
-                RebootReason::InvalidPayload
-            })?;
-            (kernel, size)
+            (kernel, kernel_size)
         } else {
             error!("Failed to locate the kernel from the DT");
             return Err(RebootReason::InvalidPayload);
         };
+        let kernel_size = kernel_size.try_into().map_err(|_| {
+            error!("Invalid kernel size: {kernel_size:#x}");
+            RebootReason::InvalidPayload
+        })?;
 
         map_rodata(kernel_start, kernel_size).map_err(|e| {
             error!("Failed to map kernel range: {e}");
@@ -99,7 +88,11 @@
         // SAFETY: map_rodata validated the range to be in main memory, mapped, and not overlap.
         let kernel = unsafe { slice::from_raw_parts(kernel, kernel_size.into()) };
 
-        let ramdisk = if let Some(r) = info.initrd_range {
+        let initrd_range = read_initrd_range_from(untrusted_fdt).map_err(|e| {
+            error!("Failed to read initrd range: {e}");
+            RebootReason::InvalidFdt
+        })?;
+        let ramdisk = if let Some(r) = initrd_range {
             debug!("Located ramdisk at {r:?}");
             let ramdisk_size = r.len().try_into().map_err(|_| {
                 error!("Invalid ramdisk size: {:#x}", r.len());
@@ -118,6 +111,6 @@
             None
         };
 
-        Ok(Self { fdt, kernel, ramdisk })
+        Ok(Self { fdt: untrusted_fdt, kernel, ramdisk })
     }
 }
diff --git a/guest/pvmfw/src/rollback.rs b/guest/pvmfw/src/rollback.rs
new file mode 100644
index 0000000..bc16332
--- /dev/null
+++ b/guest/pvmfw/src/rollback.rs
@@ -0,0 +1,159 @@
+// Copyright 2024, 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.
+
+//! Support for guest-specific rollback protection (RBP).
+
+use crate::dice::PartialInputs;
+use crate::entry::RebootReason;
+use crate::instance::EntryBody;
+use crate::instance::Error as InstanceError;
+use crate::instance::{get_recorded_entry, record_instance_entry};
+use cstr::cstr;
+use diced_open_dice::Hidden;
+use libfdt::{Fdt, FdtNode};
+use log::{error, info};
+use pvmfw_avb::Capability;
+use pvmfw_avb::VerifiedBootData;
+use virtio_drivers::transport::pci::bus::PciRoot;
+use vmbase::rand;
+
+/// Performs RBP based on the input payload, current DICE chain, and host-controlled platform.
+///
+/// On success, returns a tuple containing:
+/// - `new_instance`: true if a new entry was created using the legacy instance.img solution;
+/// - `salt`: the salt representing the instance, to be used during DICE derivation;
+/// - `defer_rollback_protection`: if RBP is being deferred.
+pub fn perform_rollback_protection(
+    fdt: &Fdt,
+    verified_boot_data: &VerifiedBootData,
+    dice_inputs: &PartialInputs,
+    pci_root: &mut PciRoot,
+    cdi_seal: &[u8],
+    instance_hash: Option<Hidden>,
+) -> Result<(bool, Hidden, bool), RebootReason> {
+    let defer_rollback_protection = should_defer_rollback_protection(fdt)?
+        && verified_boot_data.has_capability(Capability::SecretkeeperProtection);
+    let (new_instance, salt) = if defer_rollback_protection {
+        info!("Guest OS is capable of Secretkeeper protection, deferring rollback protection");
+        // rollback_index of the image is used as security_version and is expected to be > 0 to
+        // discourage implicit allocation.
+        if verified_boot_data.rollback_index == 0 {
+            error!("Expected positive rollback_index, found 0");
+            return Err(RebootReason::InvalidPayload);
+        };
+        (false, instance_hash.unwrap())
+    } else if verified_boot_data.has_capability(Capability::RemoteAttest) {
+        info!("Service VM capable of remote attestation detected, performing version checks");
+        if service_vm_version::VERSION != verified_boot_data.rollback_index {
+            // For RKP VM, we only boot if the version in the AVB footer of its kernel matches
+            // the one embedded in pvmfw at build time.
+            // This prevents the pvmfw from booting a roll backed RKP VM.
+            error!(
+                "Service VM version mismatch: expected {}, found {}",
+                service_vm_version::VERSION,
+                verified_boot_data.rollback_index
+            );
+            return Err(RebootReason::InvalidPayload);
+        }
+        (false, instance_hash.unwrap())
+    } else if verified_boot_data.has_capability(Capability::TrustySecurityVm) {
+        // The rollback protection of Trusty VMs are handled by AuthMgr, so we don't need to
+        // handle it here.
+        info!("Trusty Security VM detected");
+        (false, instance_hash.unwrap())
+    } else {
+        info!("Fallback to instance.img based rollback checks");
+        let (recorded_entry, mut instance_img, header_index) =
+            get_recorded_entry(pci_root, cdi_seal).map_err(|e| {
+                error!("Failed to get entry from instance.img: {e}");
+                RebootReason::InternalError
+            })?;
+        let (new_instance, salt) = if let Some(entry) = recorded_entry {
+            check_dice_measurements_match_entry(dice_inputs, &entry)?;
+            let salt = instance_hash.unwrap_or(entry.salt);
+            (false, salt)
+        } else {
+            // New instance!
+            let salt = instance_hash.map_or_else(rand::random_array, Ok).map_err(|e| {
+                error!("Failed to generated instance.img salt: {e}");
+                RebootReason::InternalError
+            })?;
+
+            let entry = EntryBody::new(dice_inputs, &salt);
+            record_instance_entry(&entry, cdi_seal, &mut instance_img, header_index).map_err(
+                |e| {
+                    error!("Failed to get recorded entry in instance.img: {e}");
+                    RebootReason::InternalError
+                },
+            )?;
+            (true, salt)
+        };
+        (new_instance, salt)
+    };
+
+    Ok((new_instance, salt, defer_rollback_protection))
+}
+
+fn check_dice_measurements_match_entry(
+    dice_inputs: &PartialInputs,
+    entry: &EntryBody,
+) -> Result<(), RebootReason> {
+    ensure_dice_measurements_match_entry(dice_inputs, entry).map_err(|e| {
+        error!(
+            "Dice measurements do not match recorded entry. \
+        This may be because of update: {e}"
+        );
+        RebootReason::InternalError
+    })?;
+
+    Ok(())
+}
+
+fn ensure_dice_measurements_match_entry(
+    dice_inputs: &PartialInputs,
+    entry: &EntryBody,
+) -> Result<(), InstanceError> {
+    if entry.code_hash != dice_inputs.code_hash {
+        Err(InstanceError::RecordedCodeHashMismatch)
+    } else if entry.auth_hash != dice_inputs.auth_hash {
+        Err(InstanceError::RecordedAuthHashMismatch)
+    } else if entry.mode() != dice_inputs.mode {
+        Err(InstanceError::RecordedDiceModeMismatch)
+    } else {
+        Ok(())
+    }
+}
+
+fn should_defer_rollback_protection(fdt: &Fdt) -> Result<bool, RebootReason> {
+    let node = avf_untrusted_node(fdt)?;
+    let defer_rbp = node
+        .getprop(cstr!("defer-rollback-protection"))
+        .map_err(|e| {
+            error!("Failed to get defer-rollback-protection property in DT: {e}");
+            RebootReason::InvalidFdt
+        })?
+        .is_some();
+    Ok(defer_rbp)
+}
+
+fn avf_untrusted_node(fdt: &Fdt) -> Result<FdtNode, RebootReason> {
+    let node = fdt.node(cstr!("/avf/untrusted")).map_err(|e| {
+        error!("Failed to get /avf/untrusted node: {e}");
+        RebootReason::InvalidFdt
+    })?;
+    node.ok_or_else(|| {
+        error!("/avf/untrusted node is missing in DT");
+        RebootReason::InvalidFdt
+    })
+}
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index 8314f43..01af51c 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -15,11 +15,6 @@
 
 java_defaults {
     name: "MicrodroidTestAppsDefaults",
-    test_suites: [
-        "cts",
-        "vts",
-        "general-tests",
-    ],
     static_libs: [
         "com.android.microdroid.testservice-java",
         "com.android.microdroid.test.vmshare_service-java",
@@ -64,17 +59,60 @@
     min_sdk_version: "33",
 }
 
+DATA = [
+    ":MicrodroidTestAppUpdated",
+    ":MicrodroidVmShareApp",
+    ":test_microdroid_vendor_image",
+    ":test_microdroid_vendor_image_unsigned",
+]
+
 android_test {
     name: "MicrodroidTestApp",
     defaults: ["MicrodroidVersionsTestAppDefaults"],
     manifest: "AndroidManifestV5.xml",
-    // Defined in ../vmshareapp/Android.bp
-    data: [
-        ":MicrodroidTestAppUpdated",
-        ":MicrodroidVmShareApp",
-        ":test_microdroid_vendor_image",
-        ":test_microdroid_vendor_image_unsigned",
-    ],
+    test_suites: ["general-tests"],
+    test_config: "AndroidTest.xml",
+    data: DATA,
+}
+
+android_test {
+    name: "MicrodroidTestApp.CTS",
+    defaults: ["MicrodroidVersionsTestAppDefaults"],
+    manifest: "AndroidManifestV5.xml",
+    test_suites: ["cts"],
+    test_config: ":MicrodroidTestApp.CTS.config",
+    data: DATA,
+}
+
+android_test {
+    name: "MicrodroidTestApp.VTS",
+    defaults: ["MicrodroidVersionsTestAppDefaults"],
+    manifest: "AndroidManifestV5.xml",
+    test_suites: ["vts"],
+    test_config: ":MicrodroidTestApp.VTS.config",
+    data: DATA,
+}
+
+genrule {
+    name: "MicrodroidTestApp.CTS.config",
+    srcs: ["AndroidTest.xml"],
+    out: ["out.xml"],
+    cmd: "sed " +
+        "-e 's/<!-- PLACEHOLDER_FOR_ANNOTATION -->/" +
+        "<option name=\"include-annotation\" value=\"com.android.compatibility.common.util.CddTest\" \\/>/' " +
+        "-e 's/MicrodroidTestApp.apk/MicrodroidTestApp.CTS.apk/' " +
+        "$(in) > $(out)",
+}
+
+genrule {
+    name: "MicrodroidTestApp.VTS.config",
+    srcs: ["AndroidTest.xml"],
+    out: ["out.xml"],
+    cmd: "sed " +
+        "-e 's/<!-- PLACEHOLDER_FOR_ANNOTATION -->/" +
+        "<option name=\"include-annotation\" value=\"com.android.compatibility.common.util.VsrTest\" \\/>/' " +
+        "-e 's/MicrodroidTestApp.apk/MicrodroidTestApp.VTS.apk/' " +
+        "$(in) > $(out)",
 }
 
 android_test_helper_app {
diff --git a/tests/testapk/AndroidTest.xml b/tests/testapk/AndroidTest.xml
index e490da4..221c25c 100644
--- a/tests/testapk/AndroidTest.xml
+++ b/tests/testapk/AndroidTest.xml
@@ -43,4 +43,6 @@
     <!-- Controller that will skip the module if a native bridge situation is detected -->
     <!-- For example: module wants to run arm and device is x86 -->
     <object type="module_controller" class="com.android.tradefed.testtype.suite.module.NativeBridgeModuleController" />
+
+    <!-- PLACEHOLDER_FOR_ANNOTATION -->
 </configuration>
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 b1485e3..97a5e78 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -83,7 +83,6 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.function.ThrowingRunnable;
@@ -1931,30 +1930,27 @@
         assertThat(checkVmOutputIsRedirectedToLogcat(true)).isTrue();
     }
 
-    private boolean setSystemProperties(String name, String value) {
+    private boolean isDebugPolicyEnabled(String entry) {
         Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
         UiAutomation uiAutomation = instrumentation.getUiAutomation();
-        String cmd = "setprop " + name + " " + (value.isEmpty() ? "\"\"" : value);
-        return runInShellWithStderr(TAG, uiAutomation, cmd).trim().isEmpty();
+        String cmd = "/apex/com.android.virt/bin/vm info";
+        String output = runInShellWithStderr(TAG, uiAutomation, cmd).trim();
+        for (String line : output.split("\\v")) {
+            if (line.matches("^.*Debug policy.*" + entry + ": true.*$")) {
+                return true;
+            }
+        }
+        return false;
     }
 
     @Test
-    @Ignore("b/372874464")
     public void outputIsNotRedirectedToLogcatIfNotDebuggable() throws Exception {
         assumeSupportedDevice();
 
-        // Disable debug policy to ensure no log output.
-        String sysprop = "hypervisor.virtualizationmanager.debug_policy.path";
-        String old = SystemProperties.get(sysprop);
-        assumeTrue(
-                "Can't disable debug policy. Perhapse user build?",
-                setSystemProperties(sysprop, ""));
+        // Debug policy shouldn't enable log
+        assumeFalse(isDebugPolicyEnabled("log"));
 
-        try {
-            assertThat(checkVmOutputIsRedirectedToLogcat(false)).isFalse();
-        } finally {
-            assertThat(setSystemProperties(sysprop, old)).isTrue();
-        }
+        assertThat(checkVmOutputIsRedirectedToLogcat(false)).isFalse();
     }
 
     @Test