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 493496c..9cb6e4d 100644
--- a/android/TerminalApp/res/values/strings.xml
+++ b/android/TerminalApp/res/values/strings.xml
@@ -155,6 +155,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/debian/build.sh b/build/debian/build.sh
index 19894c2..cc38dfd 100755
--- a/build/debian/build.sh
+++ b/build/debian/build.sh
@@ -11,6 +11,7 @@
 	echo "Options:"
 	echo "-h         Print usage and this help message and exit."
 	echo "-a ARCH    Architecture of the image [default is host arch: $(uname -m)]"
+	echo "-k         Build and use our custom kernel [default is cloud kernel]"
 	echo "-r         Release mode build"
 	echo "-w         Save temp work directory [for debugging]"
 }
@@ -22,7 +23,7 @@
 }
 
 parse_options() {
-	while getopts "a:hrw" option; do
+	while getopts "a:hkrw" option; do
 		case ${option} in
 			h)
 				show_help ; exit
@@ -30,6 +31,9 @@
 			a)
 				arch="$OPTARG"
 				;;
+			k)
+				use_custom_kernel=1
+				;;
 			r)
 				mode=release
 				;;
@@ -117,6 +121,33 @@
 			linux-image-generic
 		)
 	fi
+
+	if [[ "$use_custom_kernel" -eq 1 ]]; then
+		packages+=(
+			bc
+			bison
+			debhelper
+			dh-exec
+			flex
+			gcc-12
+			kernel-wedge
+			libelf-dev
+			libpci-dev
+			lz4
+			pahole
+			python3-jinja2
+			python3-docutils
+			quilt
+			rsync
+		)
+		if [[ "$arch" == "aarch64" ]]; then
+			packages+=(
+				gcc-arm-linux-gnueabihf
+				gcc-12-aarch64-linux-gnu
+			)
+		fi
+	fi
+
 	DEBIAN_FRONTEND=noninteractive \
 	apt install --no-install-recommends --assume-yes "${packages[@]}"
 
@@ -195,6 +226,74 @@
 	build_rust_binary_and_copy shutdown_runner
 }
 
+package_custom_kernel() {
+	if [[ "$use_custom_kernel" != 1 ]]; then
+		echo "linux-headers-generic" >> "${config_space}/package_config/AVF"
+		return
+	fi
+
+	# NOTE: 6.1 is the latest LTS kernel for which Debian's kernel build scripts
+	#       work on Python 3.10, the default version on our Ubuntu 22.04 builders.
+	local debian_kver="6.1.119-1"
+	local custom_flavour="avf"
+	local ksrc_base_url="https://deb.debian.org/debian/pool/main/l/linux"
+
+	local debian_ksrc_url="${ksrc_base_url}/linux_${debian_kver}.debian.tar.xz"
+	local orig_ksrc_url="${ksrc_base_url}/linux_${debian_kver%-*}.orig.tar.xz"
+
+	# 1. Grab original kernel source, merge debian patches etc.
+	mkdir -p "${workdir}/kernel/avf-${debian_arch}"
+	pushd "${workdir}/kernel" > /dev/null
+	wget "$orig_ksrc_url"
+	pushd "avf-${debian_arch}" > /dev/null
+	wget "${debian_ksrc_url}" -O - | tar xJ
+	# TODO: Copy our own kernel patches to debian/patches
+	#       and add patch file names in the desired order to debian/patches/series
+	./debian/rules orig
+
+	local abi_kver="$(sed -nE 's;Package: linux-support-(.*);\1;p' debian/control)"
+	local debarch_flavour="${custom_flavour}-${debian_arch}"
+	local abi_flavour="${abi_kver}-${debarch_flavour}"
+
+	# 2. Define our custom flavour and regenerate control file
+	# NOTE: Our flavour extends Debian's `cloud` config on the `none` featureset.
+	cat > debian/config/${debian_arch}/config.${debarch_flavour} <<EOF
+# TODO: Add our custom kernel config to this file
+EOF
+
+	sed -z "s;\[base\]\nflavours:;[base]\nflavours:\n ${debarch_flavour};" \
+	    -i debian/config/${debian_arch}/none/defines
+	cat >> debian/config/${debian_arch}/none/defines <<EOF
+[${debarch_flavour}_image]
+configs:
+ config.cloud
+ ${debian_arch}/config.${debarch_flavour}
+EOF
+	cat >> debian/config/${debian_arch}/defines <<EOF
+[${debarch_flavour}_description]
+hardware: ${arch} AVF
+hardware-long: ${arch} Android Virtualization Framework
+EOF
+	./debian/rules debian/control || true
+
+	# 3. Build the kernel and generate Debian packages
+	./debian/rules source
+	[[ "$arch" == "$(uname -m)" ]] || export $(dpkg-architecture -a $debian_arch)
+	make -j$(nproc) -f debian/rules.gen \
+	     "binary-arch_${debian_arch}_none_${debarch_flavour}"
+
+	# 4. Copy the packages to localdebs and add their names to package_config/AVF
+	popd > /dev/null
+	cp "linux-headers-${abi_flavour}_${debian_kver}_${debian_arch}.deb" \
+	   "linux-image-${abi_flavour}-unsigned_${debian_kver}_${debian_arch}.deb" \
+	   "${debian_cloud_image}/localdebs/"
+	popd > /dev/null
+	cat >> "${config_space}/package_config/AVF" <<EOF
+linux-headers-${abi_flavour}
+linux-image-${abi_flavour}-unsigned
+EOF
+}
+
 run_fai() {
 	local out="${built_image}"
 	make -C "${debian_cloud_image}" "image_bookworm_nocloud_${debian_arch}"
@@ -238,12 +337,14 @@
 arch="$(uname -m)"
 mode=debug
 save_workdir=0
+use_custom_kernel=0
 
 parse_options "$@"
 check_sudo
 install_prerequisites
 download_debian_cloud_image
 copy_android_config
+package_custom_kernel
 run_fai
 fdisk -l "${built_image}"
 images=()
diff --git a/build/debian/build_in_container.sh b/build/debian/build_in_container.sh
index 5028b74..967f5ab 100755
--- a/build/debian/build_in_container.sh
+++ b/build/debian/build_in_container.sh
@@ -6,17 +6,19 @@
   echo "Options:"
   echo "-h         Print usage and this help message and exit."
   echo "-a ARCH    Architecture of the image [default is host arch: $(uname -m)]"
+  echo "-k         Build and use our custom kernel [default is cloud kernel]"
   echo "-r         Release mode build"
   echo "-s         Leave a shell open [default: only if the build fails]"
   echo "-w         Save temp work directory in the container [for debugging]"
 }
 
 arch="$(uname -m)"
+kernel_flag=
 release_flag=
 save_workdir_flag=
 shell_condition="||"
 
-while getopts "a:rsw" option; do
+while getopts "a:hkrsw" option; do
   case ${option} in
     a)
       arch="$OPTARG"
@@ -24,6 +26,9 @@
     h)
       show_help ; exit
       ;;
+    k)
+      kernel_flag="-k"
+      ;;
     r)
       release_flag="-r"
       ;;
@@ -53,4 +58,4 @@
   -v "$ANDROID_BUILD_TOP/packages/modules/Virtualization:/root/Virtualization" \
   --workdir /root/Virtualization/build/debian \
   ubuntu:22.04 \
-  bash -c "/root/Virtualization/build/debian/build.sh -a $arch $release_flag $save_workdir_flag $shell_condition bash"
+  bash -c "/root/Virtualization/build/debian/build.sh -a $arch $release_flag $kernel_flag $save_workdir_flag $shell_condition bash"
diff --git a/build/debian/fai_config/package_config/AVF b/build/debian/fai_config/package_config/AVF
index 2e55e90..a91c354 100644
--- a/build/debian/fai_config/package_config/AVF
+++ b/build/debian/fai_config/package_config/AVF
@@ -1,5 +1,4 @@
 PACKAGES install
 
 bpfcc-tools
-linux-headers-generic
 procps
diff --git a/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh b/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh
index 22ac595..a935591 100644
--- a/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh
+++ b/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh
@@ -5,7 +5,7 @@
 cd "${KOKORO_ARTIFACTS_DIR}/git/avf/build/debian/"
 sudo losetup -D
 grep vmx /proc/cpuinfo || true
-sudo ./build.sh -r -a x86_64
+sudo ./build.sh -a x86_64 -k -r
 sudo mv images.tar.gz ${KOKORO_ARTIFACTS_DIR} || true
 mkdir -p ${KOKORO_ARTIFACTS_DIR}/logs
 sudo cp -r /var/log/fai/* ${KOKORO_ARTIFACTS_DIR}/logs || true
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/libs/libavf/include/android/virtualization.h b/libs/libavf/include/android/virtualization.h
index f33ee75..ac46fca 100644
--- a/libs/libavf/include/android/virtualization.h
+++ b/libs/libavf/include/android/virtualization.h
@@ -18,6 +18,8 @@
 #include <stdbool.h>
 #include <stdint.h>
 #include <stdlib.h>
+#include <sys/cdefs.h>
+#include <time.h>
 
 __BEGIN_DECLS
 
@@ -40,9 +42,9 @@
  * VM by calling {@link AVirtualMachine_createRaw} or releasing it by calling
  * {@link AVirtualMachineRawConfig_destroy}.
  *
- * \return A new virtual machine raw config object.
+ * \return A new virtual machine raw config object. On failure (such as out of memory), it aborts.
  */
-AVirtualMachineRawConfig* AVirtualMachineRawConfig_create();
+AVirtualMachineRawConfig* _Nonnull AVirtualMachineRawConfig_create() __INTRODUCED_IN(36);
 
 /**
  * Destroy a virtual machine config object.
@@ -52,61 +54,69 @@
  * `AVirtualMachineRawConfig_destroy` does nothing if `config` is null. A destroyed config object
  * must not be reused.
  */
-void AVirtualMachineRawConfig_destroy(AVirtualMachineRawConfig* config);
+void AVirtualMachineRawConfig_destroy(AVirtualMachineRawConfig* _Nullable config)
+        __INTRODUCED_IN(36);
 
 /**
  * Set a name of a virtual machine.
  *
  * \param config a virtual machine config object.
- * \param name a pointer to a null-terminated string for the name.
+ * \param name a pointer to a null-terminated, UTF-8 encoded string for the name.
  *
- * \return If successful, it returns 0.
+ * \return If successful, it returns 0. If `name` is not a null-terminated UTF-8 encoded string,
+ *   it returns -EINVAL.
  */
-int AVirtualMachineRawConfig_setName(AVirtualMachineRawConfig* config, const char* name);
+int AVirtualMachineRawConfig_setName(AVirtualMachineRawConfig* _Nonnull config,
+                                     const char* _Nonnull name) __INTRODUCED_IN(36);
 
 /**
  * Set an instance ID of a virtual machine.
  *
  * \param config a virtual machine config object.
  * \param instanceId a pointer to a 64-byte buffer for the instance ID.
+ * \param instanceIdSize the number of bytes in `instanceId`.
  *
- * \return If successful, it returns 0.
+ * \return If successful, it returns 0. If `instanceIdSize` is incorrect, it returns -EINVAL.
  */
-int AVirtualMachineRawConfig_setInstanceId(AVirtualMachineRawConfig* config,
-                                           const int8_t* instanceId);
+int AVirtualMachineRawConfig_setInstanceId(AVirtualMachineRawConfig* _Nonnull config,
+                                           const int8_t* _Nonnull instanceId, size_t instanceIdSize)
+        __INTRODUCED_IN(36);
 
 /**
  * Set a kernel image of a virtual machine.
  *
  * \param config a virtual machine config object.
- * \param fd a readable file descriptor containing the kernel image, or -1 to unset.
- *   `AVirtualMachineRawConfig_setKernel` takes ownership of `fd`.
- *
- * \return If successful, it returns 0.
+ * \param fd a readable, seekable, and sized (i.e. report a valid size using fstat()) file
+ *   descriptor containing the kernel image, or -1 to unset. `AVirtualMachineRawConfig_setKernel`
+ *   takes ownership of `fd`.
  */
-int AVirtualMachineRawConfig_setKernel(AVirtualMachineRawConfig* config, int fd);
+void AVirtualMachineRawConfig_setKernel(AVirtualMachineRawConfig* _Nonnull config, int fd)
+        __INTRODUCED_IN(36);
 
 /**
  * Set an init rd of a virtual machine.
  *
  * \param config a virtual machine config object.
- * \param fd a readable file descriptor containing the kernel image, or -1 to unset.
- *   `AVirtualMachineRawConfig_setInitRd` takes ownership of `fd`.
- *
- * \return If successful, it returns 0.
+ * \param fd a readable, seekable, and sized (i.e. report a valid size using fstat()) file
+ *   descriptor containing the init rd image, or -1 to unset. `AVirtualMachineRawConfig_setInitRd`
+ *   takes ownership of `fd`.
  */
-int AVirtualMachineRawConfig_setInitRd(AVirtualMachineRawConfig* config, int fd);
+void AVirtualMachineRawConfig_setInitRd(AVirtualMachineRawConfig* _Nonnull config, int fd)
+        __INTRODUCED_IN(36);
 
 /**
  * Add a disk for a virtual machine.
  *
  * \param config a virtual machine config object.
- * \param fd a readable file descriptor containing the disk image.
- * `AVirtualMachineRawConfig_addDisk` takes ownership of `fd`.
+ * \param fd a readable, seekable, and sized (i.e. report a valid size using fstat()) file
+ *   descriptor containing the disk. `fd` must be writable if If `writable` is true.
+ *   `AVirtualMachineRawConfig_addDisk` takes ownership of `fd`.
+ * \param writable whether this disk should be writable by the virtual machine.
  *
  * \return If successful, it returns 0. If `fd` is invalid, it returns -EINVAL.
  */
-int AVirtualMachineRawConfig_addDisk(AVirtualMachineRawConfig* config, int fd);
+int AVirtualMachineRawConfig_addDisk(AVirtualMachineRawConfig* _Nonnull config, int fd,
+                                     bool writable) __INTRODUCED_IN(36);
 
 /**
  * Set how much memory will be given to a virtual machine.
@@ -114,40 +124,41 @@
  * \param config a virtual machine config object.
  * \param memoryMib the amount of RAM to give the virtual machine, in MiB. 0 or negative to use the
  *   default.
- *
- * \return If successful, it returns 0.
  */
-int AVirtualMachineRawConfig_setMemoryMib(AVirtualMachineRawConfig* config, int32_t memoryMib);
+void AVirtualMachineRawConfig_setMemoryMib(AVirtualMachineRawConfig* _Nonnull config,
+                                           int32_t memoryMib) __INTRODUCED_IN(36);
 
 /**
- * Set whether a virtual machine is protected or not.
+ * Set whether the virtual machine's memory will be protected from the host, so the host can't
+ * access its memory.
  *
  * \param config a virtual machine config object.
  * \param protectedVm whether the virtual machine should be protected.
- *
- * \return If successful, it returns 0.
  */
-int AVirtualMachineRawConfig_setProtectedVm(AVirtualMachineRawConfig* config, bool protectedVm);
+void AVirtualMachineRawConfig_setProtectedVm(AVirtualMachineRawConfig* _Nonnull config,
+                                             bool protectedVm) __INTRODUCED_IN(36);
 
 /**
- * Set whether a virtual machine uses memory ballooning or not.
+ * Set whether a virtual machine uses memory ballooning.
  *
  * \param config a virtual machine config object.
  * \param balloon whether the virtual machine should use memory ballooning.
- *
- * \return If successful, it returns 0.
  */
-int AVirtualMachineRawConfig_setBalloon(AVirtualMachineRawConfig* config, bool balloon);
+void AVirtualMachineRawConfig_setBalloon(AVirtualMachineRawConfig* _Nonnull config, bool balloon)
+        __INTRODUCED_IN(36);
 
 /**
  * Set whether to use an alternate, hypervisor-specific authentication method
- * for protected VMs. You don't want to use this.
+ * for protected VMs.
+ *
+ * This option is discouraged. Prefer to use the default authentication method, which is better
+ * tested and integrated into Android. This option must only be used from the vendor partition.
  *
  * \return If successful, it returns 0. It returns `-ENOTSUP` if the hypervisor doesn't have an
  *   alternate auth mode.
  */
-int AVirtualMachineRawConfig_setHypervisorSpecificAuthMethod(AVirtualMachineRawConfig* config,
-                                                             bool enable);
+int AVirtualMachineRawConfig_setHypervisorSpecificAuthMethod(
+        AVirtualMachineRawConfig* _Nonnull config, bool enable) __INTRODUCED_IN(36);
 
 /**
  * Use the specified fd as the backing memfd for a range of the guest
@@ -155,14 +166,15 @@
  *
  * \param config a virtual machine config object.
  * \param fd a memfd
- * \param rangeStart range start IPA
- * \param rangeEnd range end IPA
+ * \param rangeStart range start of guest memory addresses
+ * \param rangeEnd range end of guest memory addresses
  *
  * \return If successful, it returns 0. It returns `-ENOTSUP` if the hypervisor doesn't support
  *   backing memfd.
  */
-int AVirtualMachineRawConfig_addCustomMemoryBackingFile(AVirtualMachineRawConfig* config, int fd,
-                                                        size_t rangeStart, size_t rangeEnd);
+int AVirtualMachineRawConfig_addCustomMemoryBackingFile(AVirtualMachineRawConfig* _Nonnull config,
+                                                        int fd, uint64_t rangeStart,
+                                                        uint64_t rangeEnd) __INTRODUCED_IN(36);
 
 /**
  * Represents a handle on a virtualization service, responsible for managing virtual machines.
@@ -176,8 +188,10 @@
  * The caller takes ownership of the returned service object, and is responsible for releasing it
  * by calling {@link AVirtualizationService_destroy}.
  *
- * \param early set to true when running a service for early virtual machines. See
- *   [`early_vm.md`](../../../../docs/early_vm.md) for more details on early virtual machines.
+ * \param early set to true when running a service for early virtual machines. Early VMs are
+ *   specialized virtual machines that can run even before the `/data` partition is mounted.
+ *   Early VMs must be pre-defined in XML files located at `{partition}/etc/avf/early_vms*.xml`, and
+ *   clients of early VMs must be pre-installed under the same partition.
  * \param service an out parameter that will be set to the service handle.
  *
  * \return
@@ -187,7 +201,8 @@
  *   - If it fails to connect to the spawned `virtmgr`, it leaves `service` untouched and returns
  *     `-ECONNREFUSED`.
  */
-int AVirtualizationService_create(AVirtualizationService** service, bool early);
+int AVirtualizationService_create(AVirtualizationService* _Null_unspecified* _Nonnull service,
+                                  bool early) __INTRODUCED_IN(36);
 
 /**
  * Destroy a VirtualizationService object.
@@ -197,7 +212,7 @@
  *
  * \param service a handle on a virtualization service.
  */
-void AVirtualizationService_destroy(AVirtualizationService* service);
+void AVirtualizationService_destroy(AVirtualizationService* _Nullable service) __INTRODUCED_IN(36);
 
 /**
  * Represents a handle on a virtual machine.
@@ -208,55 +223,55 @@
  * The reason why a virtual machine stopped.
  * @see AVirtualMachine_waitForStop
  */
-enum StopReason : int32_t {
+enum AVirtualMachineStopReason : int32_t {
     /**
      * VirtualizationService died.
      */
-    VIRTUALIZATION_SERVICE_DIED = 1,
+    AVIRTUAL_MACHINE_VIRTUALIZATION_SERVICE_DIED = 1,
     /**
      * There was an error waiting for the virtual machine.
      */
-    INFRASTRUCTURE_ERROR = 2,
+    AVIRTUAL_MACHINE_INFRASTRUCTURE_ERROR = 2,
     /**
      * The virtual machine was killed.
      */
-    KILLED = 3,
+    AVIRTUAL_MACHINE_KILLED = 3,
     /**
      * The virtual machine stopped for an unknown reason.
      */
-    UNKNOWN = 4,
+    AVIRTUAL_MACHINE_UNKNOWN = 4,
     /**
      * The virtual machine requested to shut down.
      */
-    SHUTDOWN = 5,
+    AVIRTUAL_MACHINE_SHUTDOWN = 5,
     /**
      * crosvm had an error starting the virtual machine.
      */
-    START_FAILED = 6,
+    AVIRTUAL_MACHINE_START_FAILED = 6,
     /**
      * The virtual machine requested to reboot, possibly as the result of a kernel panic.
      */
-    REBOOT = 7,
+    AVIRTUAL_MACHINE_REBOOT = 7,
     /**
      * The virtual machine or crosvm crashed.
      */
-    CRASH = 8,
+    AVIRTUAL_MACHINE_CRASH = 8,
     /**
      * The pVM firmware failed to verify the VM because the public key doesn't match.
      */
-    PVM_FIRMWARE_PUBLIC_KEY_MISMATCH = 9,
+    AVIRTUAL_MACHINE_PVM_FIRMWARE_PUBLIC_KEY_MISMATCH = 9,
     /**
      * The pVM firmware failed to verify the VM because the instance image changed.
      */
-    PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED = 10,
+    AVIRTUAL_MACHINE_PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED = 10,
     /**
      * The virtual machine was killed due to hangup.
      */
-    HANGUP = 11,
+    AVIRTUAL_MACHINE_HANGUP = 11,
     /**
      * VirtualizationService sent a stop reason which was not recognised by the client library.
      */
-    UNRECOGNISED = 0,
+    AVIRTUAL_MACHINE_UNRECOGNISED = 0,
 };
 
 /**
@@ -282,45 +297,67 @@
  * \return If successful, it sets `vm` and returns 0. Otherwise, it leaves `vm` untouched and
  *   returns `-EIO`.
  */
-int AVirtualMachine_createRaw(const AVirtualizationService* service,
-                              AVirtualMachineRawConfig* config, int consoleOutFd, int consoleInFd,
-                              int logFd, AVirtualMachine** vm);
+int AVirtualMachine_createRaw(const AVirtualizationService* _Nonnull service,
+                              AVirtualMachineRawConfig* _Nonnull config, int consoleOutFd,
+                              int consoleInFd, int logFd,
+                              AVirtualMachine* _Null_unspecified* _Nonnull vm) __INTRODUCED_IN(36);
 
 /**
- * Start a virtual machine.
+ * Start a virtual machine. `AVirtualMachine_start` is synchronous and blocks until the virtual
+ * machine is initialized and free to start executing code, or until an error happens.
  *
  * \param vm a handle on a virtual machine.
  *
  * \return If successful, it returns 0. Otherwise, it returns `-EIO`.
  */
-int AVirtualMachine_start(AVirtualMachine* vm);
+int AVirtualMachine_start(AVirtualMachine* _Nonnull vm) __INTRODUCED_IN(36);
 
 /**
- * Stop a virtual machine.
+ * Stop a virtual machine. Stopping a virtual machine is like pulling the plug on a real computer;
+ * the machine halts immediately. Software running on the virtual machine is not notified of the
+ * event, the instance might be left in an inconsistent state.
+ *
+ * For a graceful shutdown, you could request the virtual machine to exit itself, and wait for the
+ * virtual machine to stop by `AVirtualMachine_waitForStop`.
+ *
+ * A stopped virtual machine can be re-started by calling `AVirtualMachine_start`.
+ *
+ * `AVirtualMachine_stop` stops a virtual machine by sending a signal to the process. This operation
+ * is synchronous and `AVirtualMachine_stop` may block.
  *
  * \param vm a handle on a virtual machine.
  *
  * \return If successful, it returns 0. Otherwise, it returns `-EIO`.
  */
-int AVirtualMachine_stop(AVirtualMachine* vm);
+int AVirtualMachine_stop(AVirtualMachine* _Nonnull vm) __INTRODUCED_IN(36);
 
 /**
- * Wait until a virtual machine stops.
+ * Wait until a virtual machine stops or the given timeout elapses.
  *
  * \param vm a handle on a virtual machine.
+ * \param timeout the timeout, or null to wait indefinitely.
+ * \param reason An out parameter that will be set to the reason why the virtual machine stopped.
  *
- * \return The reason why the virtual machine stopped.
+ * \return
+ *   - If the virtual machine stops within the timeout (or indefinitely if `timeout` is null), it
+ *     sets `reason` and returns true.
+ *   - If the timeout expired, it returns `false`.
  */
-enum StopReason AVirtualMachine_waitForStop(AVirtualMachine* vm);
+bool AVirtualMachine_waitForStop(AVirtualMachine* _Nonnull vm,
+                                 const struct timespec* _Nullable timeout,
+                                 enum AVirtualMachineStopReason* _Nonnull reason)
+        __INTRODUCED_IN(36);
 
 /**
- * Destroy a virtual machine.
+ * Destroy a virtual machine object. If the virtual machine is still running,
+ * `AVirtualMachine_destroy` first stops the virtual machine by sending a signal to the process.
+ * This operation is synchronous and `AVirtualMachine_destroy` may block.
  *
  * `AVirtualMachine_destroy` does nothing if `vm` is null. A destroyed virtual machine must not be
  * reused.
  *
  * \param vm a handle on a virtual machine.
  */
-void AVirtualMachine_destroy(AVirtualMachine* vm);
+void AVirtualMachine_destroy(AVirtualMachine* _Nullable vm) __INTRODUCED_IN(36);
 
 __END_DECLS
diff --git a/libs/libavf/libavf.map.txt b/libs/libavf/libavf.map.txt
index ecb4cc9..05a5b35 100644
--- a/libs/libavf/libavf.map.txt
+++ b/libs/libavf/libavf.map.txt
@@ -1,4 +1,4 @@
-LIBAVF {
+LIBAVF { # introduced=36
   global:
     AVirtualMachineRawConfig_create; # apex llndk
     AVirtualMachineRawConfig_destroy; # apex llndk
diff --git a/libs/libavf/src/lib.rs b/libs/libavf/src/lib.rs
index 0a8f891..1d7861f 100644
--- a/libs/libavf/src/lib.rs
+++ b/libs/libavf/src/lib.rs
@@ -19,6 +19,7 @@
 use std::os::fd::FromRawFd;
 use std::os::raw::{c_char, c_int};
 use std::ptr;
+use std::time::Duration;
 
 use android_system_virtualizationservice::{
     aidl::android::system::virtualizationservice::{
@@ -28,7 +29,8 @@
     },
     binder::{ParcelFileDescriptor, Strong},
 };
-use avf_bindgen::StopReason;
+use avf_bindgen::AVirtualMachineStopReason;
+use libc::timespec;
 use vmclient::{DeathReason, VirtualizationService, VmInstance};
 
 /// Create a new virtual machine config object with no properties.
@@ -70,8 +72,14 @@
     // AVirtualMachineRawConfig_create. It's the only reference to the object.
     let config = unsafe { &mut *config };
     // SAFETY: `name` is assumed to be a pointer to a valid C string.
-    config.name = unsafe { CStr::from_ptr(name) }.to_string_lossy().into_owned();
-    0
+    let name = unsafe { CStr::from_ptr(name) };
+    match name.to_str() {
+        Ok(name) => {
+            config.name = name.to_owned();
+            0
+        }
+        Err(_) => -libc::EINVAL,
+    }
 }
 
 /// Set an instance ID of a virtual machine.
@@ -83,7 +91,12 @@
 pub unsafe extern "C" fn AVirtualMachineRawConfig_setInstanceId(
     config: *mut VirtualMachineRawConfig,
     instance_id: *const u8,
+    instance_id_size: usize,
 ) -> c_int {
+    if instance_id_size != 64 {
+        return -libc::EINVAL;
+    }
+
     // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
     // AVirtualMachineRawConfig_create. It's the only reference to the object.
     let config = unsafe { &mut *config };
@@ -91,7 +104,7 @@
     // is assumed to be a valid object returned by AVirtuaMachineConfig_create.
     // Both never overlap.
     unsafe {
-        ptr::copy_nonoverlapping(instance_id, config.instanceId.as_mut_ptr(), 64);
+        ptr::copy_nonoverlapping(instance_id, config.instanceId.as_mut_ptr(), instance_id_size);
     }
     0
 }
@@ -106,13 +119,12 @@
 pub unsafe extern "C" fn AVirtualMachineRawConfig_setKernel(
     config: *mut VirtualMachineRawConfig,
     fd: c_int,
-) -> c_int {
+) {
     let file = get_file_from_fd(fd);
     // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
     // AVirtualMachineRawConfig_create. It's the only reference to the object.
     let config = unsafe { &mut *config };
     config.kernel = file.map(ParcelFileDescriptor::new);
-    0
 }
 
 /// Set an init rd of a virtual machine.
@@ -125,13 +137,12 @@
 pub unsafe extern "C" fn AVirtualMachineRawConfig_setInitRd(
     config: *mut VirtualMachineRawConfig,
     fd: c_int,
-) -> c_int {
+) {
     let file = get_file_from_fd(fd);
     // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
     // AVirtualMachineRawConfig_create. It's the only reference to the object.
     let config = unsafe { &mut *config };
     config.initrd = file.map(ParcelFileDescriptor::new);
-    0
 }
 
 /// Add a disk for a virtual machine.
@@ -172,12 +183,11 @@
 pub unsafe extern "C" fn AVirtualMachineRawConfig_setMemoryMib(
     config: *mut VirtualMachineRawConfig,
     memory_mib: i32,
-) -> c_int {
+) {
     // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
     // AVirtualMachineRawConfig_create. It's the only reference to the object.
     let config = unsafe { &mut *config };
     config.memoryMib = memory_mib;
-    0
 }
 
 /// Set whether a virtual machine is protected or not.
@@ -188,12 +198,11 @@
 pub unsafe extern "C" fn AVirtualMachineRawConfig_setProtectedVm(
     config: *mut VirtualMachineRawConfig,
     protected_vm: bool,
-) -> c_int {
+) {
     // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
     // AVirtualMachineRawConfig_create. It's the only reference to the object.
     let config = unsafe { &mut *config };
     config.protectedVm = protected_vm;
-    0
 }
 
 /// Set whether a virtual machine uses memory ballooning or not.
@@ -204,12 +213,11 @@
 pub unsafe extern "C" fn AVirtualMachineRawConfig_setBalloon(
     config: *mut VirtualMachineRawConfig,
     balloon: bool,
-) -> c_int {
+) {
     // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
     // AVirtualMachineRawConfig_create. It's the only reference to the object.
     let config = unsafe { &mut *config };
     config.noBalloon = !balloon;
-    0
 }
 
 /// NOT IMPLEMENTED.
@@ -232,8 +240,8 @@
 pub extern "C" fn AVirtualMachineRawConfig_addCustomMemoryBackingFile(
     _config: *mut VirtualMachineRawConfig,
     _fd: c_int,
-    _range_start: usize,
-    _range_end: usize,
+    _range_start: u64,
+    _range_end: u64,
 ) -> c_int {
     -libc::ENOTSUP
 }
@@ -360,31 +368,64 @@
     }
 }
 
-/// Wait until a virtual machine stops.
+fn death_reason_to_stop_reason(death_reason: DeathReason) -> AVirtualMachineStopReason {
+    match death_reason {
+        DeathReason::VirtualizationServiceDied => {
+            AVirtualMachineStopReason::AVIRTUAL_MACHINE_VIRTUALIZATION_SERVICE_DIED
+        }
+        DeathReason::InfrastructureError => {
+            AVirtualMachineStopReason::AVIRTUAL_MACHINE_INFRASTRUCTURE_ERROR
+        }
+        DeathReason::Killed => AVirtualMachineStopReason::AVIRTUAL_MACHINE_KILLED,
+        DeathReason::Unknown => AVirtualMachineStopReason::AVIRTUAL_MACHINE_UNKNOWN,
+        DeathReason::Shutdown => AVirtualMachineStopReason::AVIRTUAL_MACHINE_SHUTDOWN,
+        DeathReason::StartFailed => AVirtualMachineStopReason::AVIRTUAL_MACHINE_START_FAILED,
+        DeathReason::Reboot => AVirtualMachineStopReason::AVIRTUAL_MACHINE_REBOOT,
+        DeathReason::Crash => AVirtualMachineStopReason::AVIRTUAL_MACHINE_CRASH,
+        DeathReason::PvmFirmwarePublicKeyMismatch => {
+            AVirtualMachineStopReason::AVIRTUAL_MACHINE_PVM_FIRMWARE_PUBLIC_KEY_MISMATCH
+        }
+        DeathReason::PvmFirmwareInstanceImageChanged => {
+            AVirtualMachineStopReason::AVIRTUAL_MACHINE_PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED
+        }
+        DeathReason::Hangup => AVirtualMachineStopReason::AVIRTUAL_MACHINE_HANGUP,
+        _ => AVirtualMachineStopReason::AVIRTUAL_MACHINE_UNRECOGNISED,
+    }
+}
+
+/// Wait until a virtual machine stops or the timeout elapses.
 ///
 /// # Safety
-/// `vm` must be a pointer returned by `AVirtualMachine_createRaw`.
+/// `vm` must be a pointer returned by `AVirtualMachine_createRaw`. `timeout` must be a valid
+/// pointer to a `struct timespec` object or null. `reason` must be a valid, non-null pointer to an
+/// AVirtualMachineStopReason object.
 #[no_mangle]
-pub unsafe extern "C" fn AVirtualMachine_waitForStop(vm: *const VmInstance) -> StopReason {
+pub unsafe extern "C" fn AVirtualMachine_waitForStop(
+    vm: *const VmInstance,
+    timeout: *const timespec,
+    reason: *mut AVirtualMachineStopReason,
+) -> bool {
     // SAFETY: `vm` is assumed to be a valid, non-null pointer returned by
     // AVirtualMachine_create. It's the only reference to the object.
     let vm = unsafe { &*vm };
-    match vm.wait_for_death() {
-        DeathReason::VirtualizationServiceDied => StopReason::VIRTUALIZATION_SERVICE_DIED,
-        DeathReason::InfrastructureError => StopReason::INFRASTRUCTURE_ERROR,
-        DeathReason::Killed => StopReason::KILLED,
-        DeathReason::Unknown => StopReason::UNKNOWN,
-        DeathReason::Shutdown => StopReason::SHUTDOWN,
-        DeathReason::StartFailed => StopReason::START_FAILED,
-        DeathReason::Reboot => StopReason::REBOOT,
-        DeathReason::Crash => StopReason::CRASH,
-        DeathReason::PvmFirmwarePublicKeyMismatch => StopReason::PVM_FIRMWARE_PUBLIC_KEY_MISMATCH,
-        DeathReason::PvmFirmwareInstanceImageChanged => {
-            StopReason::PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED
+
+    let death_reason = if timeout.is_null() {
+        vm.wait_for_death()
+    } else {
+        // SAFETY: `timeout` is assumed to be a valid pointer to a `struct timespec` object if
+        // non-null.
+        let timeout = unsafe { &*timeout };
+        let timeout = Duration::new(timeout.tv_sec as u64, timeout.tv_nsec as u32);
+        match vm.wait_for_death_with_timeout(timeout) {
+            Some(death_reason) => death_reason,
+            None => return false,
         }
-        DeathReason::Hangup => StopReason::HANGUP,
-        _ => StopReason::UNRECOGNISED,
-    }
+    };
+
+    // SAFETY: `reason` is assumed to be a valid, non-null pointer to an
+    // AVirtualMachineStopReason object.
+    unsafe { *reason = death_reason_to_stop_reason(death_reason) };
+    true
 }
 
 /// Destroy a virtual machine.
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