Merge "Disable dylib variants of nostd bindgen modules" into main
diff --git a/TEST_MAPPING b/TEST_MAPPING
index a822210..fc88a59 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -46,6 +46,9 @@
     },
     {
       "name": "libdice_driver_test"
+    },
+    {
+      "name": "vm_accessor_test"
     }
   ],
   "avf-postsubmit": [
@@ -64,9 +67,6 @@
     {
       "name": "AvfRkpdAppGoogleIntegrationTests",
       "keywords": ["internal"]
-    },
-    {
-      "name": "vm_accessor_test"
     }
   ],
   "ferrochrome-postsubmit": [
diff --git a/android/FerrochromeApp/AndroidManifest.xml b/android/FerrochromeApp/AndroidManifest.xml
index d640c4a..f6d3f6a 100644
--- a/android/FerrochromeApp/AndroidManifest.xml
+++ b/android/FerrochromeApp/AndroidManifest.xml
@@ -12,6 +12,9 @@
         <intent>
             <action android:name="android.virtualization.VM_LAUNCHER" />
         </intent>
+        <intent>
+            <action android:name="android.virtualization.FERROCHROME_DOWNLOADER" />
+        </intent>
     </queries>
     <application
         android:label="Ferrochrome">
diff --git a/android/FerrochromeApp/custom_vm_setup.sh b/android/FerrochromeApp/custom_vm_setup.sh
index a5480ff..4dce0c7 100644
--- a/android/FerrochromeApp/custom_vm_setup.sh
+++ b/android/FerrochromeApp/custom_vm_setup.sh
@@ -7,13 +7,13 @@
 }
 
 function install() {
-  user=$(cmd user get-main-user)
-  src_dir=/data/media/${user}/ferrochrome/
+  src_dir=$(getprop debug.custom_vm_setup.path)
+  src_dir=${src_dir/#\/storage\/emulated\//\/data\/media\/}
   dst_dir=/data/local/tmp/
 
   cat $(find ${src_dir} -name "images.tar.gz*" | sort) | tar xz -C ${dst_dir}
-  cp -u ${src_dir}vm_config.json ${dst_dir}
-  chmod 666 ${dst_dir}*
+  cp -u ${src_dir}/vm_config.json ${dst_dir}
+  chmod 666 ${dst_dir}/*
 
   # increase the size of state.img to the multiple of 4096
   num_blocks=$(du -b -K ${dst_dir}state.img | cut -f 1)
@@ -21,8 +21,8 @@
   additional_blocks=$((( ${required_num_blocks} - ${num_blocks} )))
   dd if=/dev/zero bs=512 count=${additional_blocks} >> ${dst_dir}state.img
 
-  rm ${src_dir}images.tar.gz*
-  rm ${src_dir}vm_config.json
+  rm ${src_dir}/images.tar.gz*
+  rm ${src_dir}/vm_config.json
 }
 
 setprop debug.custom_vm_setup.done false
diff --git a/android/FerrochromeApp/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java b/android/FerrochromeApp/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
index 2df5cab..dba0078 100644
--- a/android/FerrochromeApp/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
+++ b/android/FerrochromeApp/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
@@ -16,14 +16,17 @@
 
 package com.android.virtualization.ferrochrome;
 
+import android.annotation.WorkerThread;
 import android.app.Activity;
 import android.app.ActivityManager;
 import android.content.Intent;
+import android.content.pm.ActivityInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
 import android.os.Bundle;
 import android.os.Environment;
 import android.os.SystemProperties;
+import android.text.TextUtils;
 import android.util.Log;
 import android.view.WindowManager;
 import android.widget.TextView;
@@ -43,12 +46,20 @@
 public class FerrochromeActivity extends Activity {
     private static final String TAG = FerrochromeActivity.class.getName();
     private static final String ACTION_VM_LAUNCHER = "android.virtualization.VM_LAUNCHER";
+    private static final String ACTION_FERROCHROME_DOWNLOAD =
+            "android.virtualization.FERROCHROME_DOWNLOADER";
+    private static final String EXTRA_FERROCHROME_DEST_DIR = "dest_dir";
+    private static final String EXTRA_FERROCHROME_UPDATE_NEEDED = "update_needed";
 
     private static final Path DEST_DIR =
             Path.of(Environment.getExternalStorageDirectory().getPath(), "ferrochrome");
+    private static final String ASSET_DIR = "ferrochrome";
     private static final Path VERSION_FILE = Path.of(DEST_DIR.toString(), "version");
 
     private static final int REQUEST_CODE_VMLAUNCHER = 1;
+    private static final int REQUEST_CODE_FERROCHROME_DOWNLOADER = 2;
+
+    private ResolvedActivity mVmLauncher;
 
     ExecutorService executorService = Executors.newSingleThreadExecutor();
 
@@ -66,25 +77,28 @@
         getWindow().addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON);
 
         // Find VM Launcher
-        Intent intent = new Intent(ACTION_VM_LAUNCHER).setFlags(Intent.FLAG_ACTIVITY_SINGLE_TOP);
-        PackageManager pm = getPackageManager();
-        List<ResolveInfo> resolveInfos =
-                pm.queryIntentActivities(intent, PackageManager.MATCH_DEFAULT_ONLY);
-        if (resolveInfos == null || resolveInfos.size() != 1) {
+        mVmLauncher = ResolvedActivity.resolve(getPackageManager(), ACTION_VM_LAUNCHER);
+        if (mVmLauncher == null) {
             updateStatus("Failed to resolve VM Launcher");
             return;
         }
 
         // Clean up the existing vm launcher process if there is
         ActivityManager am = getSystemService(ActivityManager.class);
-        am.killBackgroundProcesses(resolveInfos.get(0).activityInfo.packageName);
+        am.killBackgroundProcesses(mVmLauncher.activityInfo.packageName);
 
         executorService.execute(
                 () -> {
-                    if (updateImageIfNeeded()) {
-                        updateStatus("Starting Ferrochrome...");
-                        runOnUiThread(
-                                () -> startActivityForResult(intent, REQUEST_CODE_VMLAUNCHER));
+                    if (hasLocalAssets()) {
+                        if (updateImageIfNeeded()) {
+                            updateStatus("Starting Ferrochrome...");
+                            runOnUiThread(
+                                    () ->
+                                            startActivityForResult(
+                                                    mVmLauncher.intent, REQUEST_CODE_VMLAUNCHER));
+                        }
+                    } else {
+                        tryLaunchDownloader();
                     }
                 });
     }
@@ -93,9 +107,55 @@
     protected void onActivityResult(int requestCode, int resultCode, Intent data) {
         if (requestCode == REQUEST_CODE_VMLAUNCHER) {
             finishAndRemoveTask();
+        } else if (requestCode == REQUEST_CODE_FERROCHROME_DOWNLOADER) {
+            String destDir = data.getStringExtra(EXTRA_FERROCHROME_DEST_DIR);
+            boolean updateNeeded =
+                    data.getBooleanExtra(EXTRA_FERROCHROME_UPDATE_NEEDED, /* default= */ true);
+
+            if (resultCode != RESULT_OK || TextUtils.isEmpty(destDir)) {
+                Log.w(
+                        TAG,
+                        "Ferrochrome downloader returned error, code="
+                                + resultCode
+                                + ", dest="
+                                + destDir);
+                updateStatus("User didn't accepted ferrochrome download..");
+                return;
+            }
+
+            Log.w(TAG, "Ferrochrome downloader returned OK");
+
+            if (!updateNeeded) {
+                updateStatus("Starting Ferrochrome...");
+                startActivityForResult(mVmLauncher.intent, REQUEST_CODE_VMLAUNCHER);
+            }
+
+            executorService.execute(
+                    () -> {
+                        if (!extractImages(destDir)) {
+                            updateStatus("Images from downloader looks bad..");
+                            return;
+                        }
+                        updateStatus("Starting Ferrochrome...");
+                        runOnUiThread(
+                                () ->
+                                        startActivityForResult(
+                                                mVmLauncher.intent, REQUEST_CODE_VMLAUNCHER));
+                    });
         }
     }
 
+    @WorkerThread
+    private boolean hasLocalAssets() {
+        try {
+            String[] files = getAssets().list(ASSET_DIR);
+            return files != null && files.length > 0;
+        } catch (IOException e) {
+            return false;
+        }
+    }
+
+    @WorkerThread
     private boolean updateImageIfNeeded() {
         if (!isUpdateNeeded()) {
             Log.d(TAG, "No update needed.");
@@ -107,13 +167,8 @@
                 Files.createDirectory(DEST_DIR);
             }
 
-            String[] files = getAssets().list("ferrochrome");
-            if (files == null || files.length == 0) {
-                updateStatus("ChromeOS image not found. Please go/try-ferrochrome");
-                return false;
-            }
-
             updateStatus("Copying images...");
+            String[] files = getAssets().list("ferrochrome");
             for (String file : files) {
                 updateStatus(file);
                 Path dst = Path.of(DEST_DIR.toString(), file);
@@ -126,7 +181,38 @@
         }
         updateStatus("Done.");
 
+        return extractImages(DEST_DIR.toAbsolutePath().toString());
+    }
+
+    @WorkerThread
+    private void tryLaunchDownloader() {
+        // TODO(jaewan): Add safeguard to check whether ferrochrome downloader is valid.
+        Log.w(TAG, "No built-in assets found. Try again with ferrochrome downloader");
+
+        ResolvedActivity downloader =
+                ResolvedActivity.resolve(getPackageManager(), ACTION_FERROCHROME_DOWNLOAD);
+        if (downloader == null) {
+            Log.d(TAG, "Ferrochrome downloader doesn't exist");
+            updateStatus("ChromeOS image not found. Please go/try-ferrochrome");
+            return;
+        }
+        String pkgName = downloader.activityInfo.packageName;
+        Log.d(TAG, "Resolved Ferrochrome Downloader, pkgName=" + pkgName);
+        updateStatus("Launching Ferrochrome downloader for update");
+
+        // onActivityResult() will handle downloader result.
+        startActivityForResult(downloader.intent, REQUEST_CODE_FERROCHROME_DOWNLOADER);
+    }
+
+    @WorkerThread
+    private boolean extractImages(String destDir) {
         updateStatus("Extracting images...");
+
+        if (TextUtils.isEmpty(destDir)) {
+            throw new RuntimeException("Internal error: destDir shouldn't be null");
+        }
+
+        SystemProperties.set("debug.custom_vm_setup.path", destDir);
         SystemProperties.set("debug.custom_vm_setup.done", "false");
         SystemProperties.set("debug.custom_vm_setup.start", "true");
         while (!SystemProperties.getBoolean("debug.custom_vm_setup.done", false)) {
@@ -143,6 +229,7 @@
         return true;
     }
 
+    @WorkerThread
     private boolean isUpdateNeeded() {
         Path[] pathsToCheck = {DEST_DIR, VERSION_FILE};
         for (Path p : pathsToCheck) {
@@ -188,4 +275,33 @@
                     statusView.append(line + "\n");
                 });
     }
+
+    private static final class ResolvedActivity {
+        public final ActivityInfo activityInfo;
+        public final Intent intent;
+
+        private ResolvedActivity(ActivityInfo activityInfo, Intent intent) {
+            this.activityInfo = activityInfo;
+            this.intent = intent;
+        }
+
+        /* synthetic access */
+        static ResolvedActivity resolve(PackageManager pm, String action) {
+            Intent intent = new Intent(action).setFlags(Intent.FLAG_ACTIVITY_SINGLE_TOP);
+            List<ResolveInfo> resolveInfos =
+                    pm.queryIntentActivities(intent, PackageManager.MATCH_DEFAULT_ONLY);
+            if (resolveInfos == null || resolveInfos.size() != 1) {
+                Log.w(
+                        TAG,
+                        "Failed to resolve activity, action="
+                                + action
+                                + ", resolved="
+                                + resolveInfos);
+                return null;
+            }
+            ActivityInfo activityInfo = resolveInfos.get(0).activityInfo;
+            intent.setClassName(activityInfo.packageName, activityInfo.name);
+            return new ResolvedActivity(activityInfo, intent);
+        }
+    }
 }
diff --git a/android/FerrochromeApp/vm_config.json.template b/android/FerrochromeApp/vm_config.json.template
index 6e024ba..380f016 100644
--- a/android/FerrochromeApp/vm_config.json.template
+++ b/android/FerrochromeApp/vm_config.json.template
@@ -43,7 +43,7 @@
     },
     "audio": {
         "speaker": true,
-         "microphone": true
+        "microphone": true
     },
     "gpu": {
         "backend": "virglrenderer",
@@ -52,5 +52,5 @@
     "display": {
         "scale": "0.77",
         "refresh_rate": "30"
-     }
+    }
 }
diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/VmConfigJson.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/ConfigJson.java
similarity index 98%
rename from android/VmLauncherApp/java/com/android/virtualization/vmlauncher/VmConfigJson.java
rename to android/VmLauncherApp/java/com/android/virtualization/vmlauncher/ConfigJson.java
index 8116743..6d39b46 100644
--- a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/VmConfigJson.java
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/ConfigJson.java
@@ -36,10 +36,10 @@
 import java.util.Arrays;
 
 /** This class and its inner classes model vm_config.json. */
-class VmConfigJson {
+class ConfigJson {
     private static final boolean DEBUG = true;
 
-    private VmConfigJson() {}
+    private ConfigJson() {}
 
     @SerializedName("protected")
     private boolean isProtected;
@@ -64,9 +64,9 @@
     private GpuJson gpu;
 
     /** Parses JSON file at jsonPath */
-    static VmConfigJson from(String jsonPath) {
+    static ConfigJson from(String jsonPath) {
         try (FileReader r = new FileReader(jsonPath)) {
-            return new Gson().fromJson(r, VmConfigJson.class);
+            return new Gson().fromJson(r, ConfigJson.class);
         } catch (Exception e) {
             throw new RuntimeException("Failed to parse " + jsonPath, e);
         }
diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/DisplayProvider.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/DisplayProvider.java
new file mode 100644
index 0000000..6eba709
--- /dev/null
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/DisplayProvider.java
@@ -0,0 +1,210 @@
+/*
+ * Copyright (C) 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.vmlauncher;
+
+import android.crosvm.ICrosvmAndroidDisplayService;
+import android.graphics.PixelFormat;
+import android.os.IBinder;
+import android.os.ParcelFileDescriptor;
+import android.os.RemoteException;
+import android.os.ServiceManager;
+import android.system.virtualizationservice_internal.IVirtualizationServiceInternal;
+import android.util.Log;
+import android.view.SurfaceControl;
+import android.view.SurfaceHolder;
+import android.view.SurfaceView;
+
+import libcore.io.IoBridge;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+/** Presents Android-side surface where VM can use as a display */
+class DisplayProvider {
+    private static final String TAG = MainActivity.TAG;
+    private final SurfaceView mMainView;
+    private final SurfaceView mCursorView;
+    private final IVirtualizationServiceInternal mVirtService;
+    private CursorHandler mCursorHandler;
+
+    DisplayProvider(SurfaceView mainView, SurfaceView cursorView) {
+        mMainView = mainView;
+        mCursorView = cursorView;
+
+        mMainView.setSurfaceLifecycle(SurfaceView.SURFACE_LIFECYCLE_FOLLOWS_ATTACHMENT);
+        mMainView.getHolder().addCallback(new Callback(Callback.SurfaceKind.MAIN));
+
+        mCursorView.setSurfaceLifecycle(SurfaceView.SURFACE_LIFECYCLE_FOLLOWS_ATTACHMENT);
+        mCursorView.getHolder().addCallback(new Callback(Callback.SurfaceKind.CURSOR));
+        mCursorView.getHolder().setFormat(PixelFormat.RGBA_8888);
+        // TODO: do we need this z-order?
+        mCursorView.setZOrderMediaOverlay(true);
+
+        IBinder b = ServiceManager.waitForService("android.system.virtualizationservice");
+        mVirtService = IVirtualizationServiceInternal.Stub.asInterface(b);
+        try {
+            // To ensure that the previous display service is removed.
+            mVirtService.clearDisplayService();
+        } catch (RemoteException e) {
+            throw new RuntimeException("Failed to clear prior display service", e);
+        }
+    }
+
+    void notifyDisplayIsGoingToInvisible() {
+        // When the display is going to be invisible (by putting in the background), save the frame
+        // of the main surface so that we can re-draw it next time the display becomes visible. This
+        // is to save the duration of time where nothing is drawn by VM.
+        try {
+            getDisplayService().saveFrameForSurface(false /* forCursor */);
+        } catch (RemoteException e) {
+            throw new RuntimeException("Failed to save frame for the main surface", e);
+        }
+    }
+
+    private synchronized ICrosvmAndroidDisplayService getDisplayService() {
+        try {
+            IBinder b = mVirtService.waitDisplayService();
+            return ICrosvmAndroidDisplayService.Stub.asInterface(b);
+        } catch (Exception e) {
+            throw new RuntimeException("Error while getting display service", e);
+        }
+    }
+
+    private class Callback implements SurfaceHolder.Callback {
+        enum SurfaceKind {
+            MAIN,
+            CURSOR
+        }
+
+        private final SurfaceKind mSurfaceKind;
+
+        Callback(SurfaceKind kind) {
+            mSurfaceKind = kind;
+        }
+
+        private boolean isForCursor() {
+            return mSurfaceKind == SurfaceKind.CURSOR;
+        }
+
+        @Override
+        public void surfaceCreated(SurfaceHolder holder) {
+            try {
+                getDisplayService().setSurface(holder.getSurface(), isForCursor());
+            } catch (Exception e) {
+                // TODO: don't consume this exception silently. For some unknown reason, setSurface
+                // call above throws IllegalArgumentException and that fails the surface
+                // configuration.
+                Log.e(TAG, "Failed to present surface " + mSurfaceKind + " to VM", e);
+            }
+
+            try {
+                switch (mSurfaceKind) {
+                    case MAIN:
+                        getDisplayService().drawSavedFrameForSurface(isForCursor());
+                        break;
+                    case CURSOR:
+                        ParcelFileDescriptor stream = createNewCursorStream();
+                        getDisplayService().setCursorStream(stream);
+                        break;
+                }
+            } catch (Exception e) {
+                // TODO: don't consume exceptions here too
+                Log.e(TAG, "Failed to configure surface " + mSurfaceKind, e);
+            }
+        }
+
+        @Override
+        public void surfaceChanged(SurfaceHolder holder, int format, int width, int height) {
+            // TODO: support resizeable display. We could actually change the display size that the
+            // VM sees, or keep the size and render it by fitting it in the new surface.
+        }
+
+        @Override
+        public void surfaceDestroyed(SurfaceHolder holder) {
+            try {
+                getDisplayService().removeSurface(isForCursor());
+            } catch (RemoteException e) {
+                throw new RuntimeException("Error while destroying surface for " + mSurfaceKind, e);
+            }
+        }
+    }
+
+    private ParcelFileDescriptor createNewCursorStream() {
+        if (mCursorHandler != null) {
+            mCursorHandler.interrupt();
+        }
+        ParcelFileDescriptor[] pfds;
+        try {
+            pfds = ParcelFileDescriptor.createSocketPair();
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to create socketpair for cursor stream", e);
+        }
+        mCursorHandler = new CursorHandler(pfds[0]);
+        mCursorHandler.start();
+        return pfds[1];
+    }
+
+    /**
+     * Thread reading cursor coordinate from a stream, and updating the position of the cursor
+     * surface accordingly.
+     */
+    private class CursorHandler extends Thread {
+        private final ParcelFileDescriptor mStream;
+        private final SurfaceControl mCursor;
+        private final SurfaceControl.Transaction mTransaction;
+
+        CursorHandler(ParcelFileDescriptor stream) {
+            mStream = stream;
+            mCursor = DisplayProvider.this.mCursorView.getSurfaceControl();
+            mTransaction = new SurfaceControl.Transaction();
+
+            SurfaceControl main = DisplayProvider.this.mMainView.getSurfaceControl();
+            mTransaction.reparent(mCursor, main).apply();
+        }
+
+        @Override
+        public void run() {
+            try {
+                ByteBuffer byteBuffer = ByteBuffer.allocate(8 /* (x: u32, y: u32) */);
+                byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
+                while (true) {
+                    if (Thread.interrupted()) {
+                        Log.d(TAG, "CursorHandler thread interrupted!");
+                        return;
+                    }
+                    byteBuffer.clear();
+                    int bytes =
+                            IoBridge.read(
+                                    mStream.getFileDescriptor(),
+                                    byteBuffer.array(),
+                                    0,
+                                    byteBuffer.array().length);
+                    if (bytes == -1) {
+                        Log.e(TAG, "cannot read from cursor stream, stop the handler");
+                        return;
+                    }
+                    float x = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
+                    float y = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
+                    mTransaction.setPosition(mCursor, x, y).apply();
+                }
+            } catch (IOException e) {
+                Log.e(TAG, "failed to run CursorHandler", e);
+            }
+        }
+    }
+}
diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/InputForwarder.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/InputForwarder.java
new file mode 100644
index 0000000..1be362b
--- /dev/null
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/InputForwarder.java
@@ -0,0 +1,148 @@
+/*
+ * Copyright (C) 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.vmlauncher;
+
+import android.content.Context;
+import android.hardware.input.InputManager;
+import android.os.Handler;
+import android.system.virtualmachine.VirtualMachine;
+import android.system.virtualmachine.VirtualMachineCustomImageConfig;
+import android.util.Log;
+import android.view.InputDevice;
+import android.view.KeyEvent;
+import android.view.View;
+
+/** Forwards input events (touch, mouse, ...) from Android to VM */
+class InputForwarder {
+    private static final String TAG = MainActivity.TAG;
+    private final Context mContext;
+    private final VirtualMachine mVirtualMachine;
+    private InputManager.InputDeviceListener mInputDeviceListener;
+
+    private boolean isTabletMode = false;
+
+    InputForwarder(
+            Context context,
+            VirtualMachine vm,
+            View touchReceiver,
+            View mouseReceiver,
+            View keyReceiver) {
+        mContext = context;
+        mVirtualMachine = vm;
+
+        VirtualMachineCustomImageConfig config = vm.getConfig().getCustomImageConfig();
+        if (config.useTouch()) {
+            setupTouchReceiver(touchReceiver);
+        }
+        if (config.useMouse() || config.useTrackpad()) {
+            setupMouseReceiver(mouseReceiver);
+        }
+        if (config.useKeyboard()) {
+            setupKeyReceiver(keyReceiver);
+        }
+        if (config.useSwitches()) {
+            // Any view's handler is fine.
+            setupTabletModeHandler(touchReceiver.getHandler());
+        }
+    }
+
+    void cleanUp() {
+        if (mInputDeviceListener != null) {
+            InputManager im = mContext.getSystemService(InputManager.class);
+            im.unregisterInputDeviceListener(mInputDeviceListener);
+            mInputDeviceListener = null;
+        }
+    }
+
+    private void setupTouchReceiver(View receiver) {
+        receiver.setOnTouchListener(
+                (v, event) -> {
+                    return mVirtualMachine.sendMultiTouchEvent(event);
+                });
+    }
+
+    private void setupMouseReceiver(View receiver) {
+        receiver.requestUnbufferedDispatch(InputDevice.SOURCE_ANY);
+        receiver.setOnCapturedPointerListener(
+                (v, event) -> {
+                    int eventSource = event.getSource();
+                    if ((eventSource & InputDevice.SOURCE_CLASS_POSITION) != 0) {
+                        return mVirtualMachine.sendTrackpadEvent(event);
+                    }
+                    return mVirtualMachine.sendMouseEvent(event);
+                });
+    }
+
+    private void setupKeyReceiver(View receiver) {
+        receiver.setOnKeyListener(
+                (v, code, event) -> {
+                    // TODO: this is guest-os specific. It shouldn't be handled here.
+                    if (isVolumeKey(code)) {
+                        return false;
+                    }
+                    return mVirtualMachine.sendKeyEvent(event);
+                });
+    }
+
+    private static boolean isVolumeKey(int keyCode) {
+        return keyCode == KeyEvent.KEYCODE_VOLUME_UP
+                || keyCode == KeyEvent.KEYCODE_VOLUME_DOWN
+                || keyCode == KeyEvent.KEYCODE_VOLUME_MUTE;
+    }
+
+    private void setupTabletModeHandler(Handler handler) {
+        InputManager im = mContext.getSystemService(InputManager.class);
+        mInputDeviceListener =
+                new InputManager.InputDeviceListener() {
+                    @Override
+                    public void onInputDeviceAdded(int deviceId) {
+                        setTabletModeConditionally();
+                    }
+
+                    @Override
+                    public void onInputDeviceRemoved(int deviceId) {
+                        setTabletModeConditionally();
+                    }
+
+                    @Override
+                    public void onInputDeviceChanged(int deviceId) {
+                        setTabletModeConditionally();
+                    }
+                };
+        im.registerInputDeviceListener(mInputDeviceListener, handler);
+    }
+
+    private static boolean hasPhysicalKeyboard() {
+        for (int id : InputDevice.getDeviceIds()) {
+            InputDevice d = InputDevice.getDevice(id);
+            if (!d.isVirtual() && d.isEnabled() && d.isFullKeyboard()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    void setTabletModeConditionally() {
+        boolean tabletModeNeeded = !hasPhysicalKeyboard();
+        if (tabletModeNeeded != isTabletMode) {
+            String mode = tabletModeNeeded ? "tablet mode" : "desktop mode";
+            Log.d(TAG, "switching to " + mode);
+            isTabletMode = tabletModeNeeded;
+            mVirtualMachine.sendTabletModeEvent(tabletModeNeeded);
+        }
+    }
+}
diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/Logger.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/Logger.java
new file mode 100644
index 0000000..e1cb285
--- /dev/null
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/Logger.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 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.vmlauncher;
+
+import android.system.virtualmachine.VirtualMachine;
+import android.system.virtualmachine.VirtualMachineConfig;
+import android.system.virtualmachine.VirtualMachineException;
+import android.util.Log;
+
+import libcore.io.Streams;
+
+import java.io.BufferedOutputStream;
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.concurrent.ExecutorService;
+
+/**
+ * Forwards VM's console output to a file on the Android side, and VM's log output to Android logd.
+ */
+class Logger {
+    private Logger() {}
+
+    static void setup(VirtualMachine vm, Path path, ExecutorService executor) {
+        if (vm.getConfig().getDebugLevel() != VirtualMachineConfig.DEBUG_LEVEL_FULL) {
+            return;
+        }
+
+        try {
+            InputStream console = vm.getConsoleOutput();
+            OutputStream file = Files.newOutputStream(path, StandardOpenOption.CREATE);
+            executor.submit(() -> Streams.copy(console, new LineBufferedOutputStream(file)));
+
+            InputStream log = vm.getLogOutput();
+            executor.submit(() -> writeToLogd(log, vm.getName()));
+        } catch (VirtualMachineException | IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static boolean writeToLogd(InputStream input, String vmName) throws IOException {
+        BufferedReader reader = new BufferedReader(new InputStreamReader(input));
+        String line;
+        while ((line = reader.readLine()) != null && !Thread.interrupted()) {
+            Log.d(vmName, line);
+        }
+        // TODO: find out why javac complains when the return type of this method is void. It
+        // (incorrectly?) thinks that IOException should be caught inside the lambda.
+        return true;
+    }
+
+    private static class LineBufferedOutputStream extends BufferedOutputStream {
+        LineBufferedOutputStream(OutputStream out) {
+            super(out);
+        }
+
+        @Override
+        public void write(byte[] buf, int off, int len) throws IOException {
+            super.write(buf, off, len);
+            for (int i = 0; i < len; ++i) {
+                if (buf[off + i] == '\n') {
+                    flush();
+                    break;
+                }
+            }
+        }
+    }
+}
diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java
index 160140a..7c7d7ef 100644
--- a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java
@@ -23,51 +23,35 @@
 import android.content.ClipData;
 import android.content.ClipboardManager;
 import android.content.Intent;
-import android.crosvm.ICrosvmAndroidDisplayService;
-import android.graphics.PixelFormat;
-import android.hardware.input.InputManager;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
-import android.os.RemoteException;
-import android.os.ServiceManager;
-import android.system.virtualizationservice_internal.IVirtualizationServiceInternal;
 import android.system.virtualmachine.VirtualMachine;
-import android.system.virtualmachine.VirtualMachineCallback;
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineException;
-import android.system.virtualmachine.VirtualMachineManager;
 import android.util.Log;
-import android.view.InputDevice;
-import android.view.KeyEvent;
-import android.view.SurfaceControl;
-import android.view.SurfaceHolder;
 import android.view.SurfaceView;
 import android.view.View;
+import android.view.Window;
 import android.view.WindowInsets;
 import android.view.WindowInsetsController;
-import android.view.WindowManager;
 
-import libcore.io.IoBridge;
-
-import java.io.BufferedOutputStream;
-import java.io.BufferedReader;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
-public class MainActivity extends Activity implements InputManager.InputDeviceListener {
-    private static final String TAG = "VmLauncherApp";
-    private static final String VM_NAME = "my_custom_vm";
+public class MainActivity extends Activity {
+    static final String TAG = "VmLauncherApp";
+    // TODO: this path should be from outside of this activity
+    private static final String VM_CONFIG_PATH = "/data/local/tmp/vm_config.json";
 
-    private static final boolean DEBUG = true;
     private static final int RECORD_AUDIO_PERMISSION_REQUEST_CODE = 101;
 
     private static final String ACTION_VM_LAUNCHER = "android.virtualization.VM_LAUNCHER";
@@ -75,81 +59,9 @@
 
     private ExecutorService mExecutorService;
     private VirtualMachine mVirtualMachine;
-    private CursorHandler mCursorHandler;
     private ClipboardManager mClipboardManager;
-
-
-    private static boolean isVolumeKey(int keyCode) {
-        return keyCode == KeyEvent.KEYCODE_VOLUME_UP
-                || keyCode == KeyEvent.KEYCODE_VOLUME_DOWN
-                || keyCode == KeyEvent.KEYCODE_VOLUME_MUTE;
-    }
-
-    @Override
-    public boolean onKeyDown(int keyCode, KeyEvent event) {
-        if (mVirtualMachine == null) {
-            return false;
-        }
-        return !isVolumeKey(keyCode) && mVirtualMachine.sendKeyEvent(event);
-    }
-
-    @Override
-    public boolean onKeyUp(int keyCode, KeyEvent event) {
-        if (mVirtualMachine == null) {
-            return false;
-        }
-        return !isVolumeKey(keyCode) && mVirtualMachine.sendKeyEvent(event);
-    }
-
-    private void registerInputDeviceListener() {
-        InputManager inputManager = getSystemService(InputManager.class);
-        if (inputManager == null) {
-            Log.e(TAG, "failed to registerInputDeviceListener because InputManager is null");
-            return;
-        }
-        inputManager.registerInputDeviceListener(this, null);
-    }
-
-    private void unregisterInputDeviceListener() {
-        InputManager inputManager = getSystemService(InputManager.class);
-        if (inputManager == null) {
-            Log.e(TAG, "failed to unregisterInputDeviceListener because InputManager is null");
-            return;
-        }
-        inputManager.unregisterInputDeviceListener(this);
-    }
-
-    private void setTabletModeConditionally() {
-        if (mVirtualMachine == null) {
-            Log.e(TAG, "failed to setTabletModeConditionally because VirtualMachine is null");
-            return;
-        }
-        for (int id : InputDevice.getDeviceIds()) {
-            InputDevice d = InputDevice.getDevice(id);
-            if (!d.isVirtual() && d.isEnabled() && d.isFullKeyboard()) {
-                Log.d(TAG, "the device has a physical keyboard, turn off tablet mode");
-                mVirtualMachine.sendTabletModeEvent(false);
-                return;
-            }
-        }
-        mVirtualMachine.sendTabletModeEvent(true);
-        Log.d(TAG, "the device doesn't have a physical keyboard, turn on tablet mode");
-    }
-
-    @Override
-    public void onInputDeviceAdded(int deviceId) {
-        setTabletModeConditionally();
-    }
-
-    @Override
-    public void onInputDeviceRemoved(int deviceId) {
-        setTabletModeConditionally();
-    }
-
-    @Override
-    public void onInputDeviceChanged(int deviceId) {
-        setTabletModeConditionally();
-    }
+    private InputForwarder mInputForwarder;
+    private DisplayProvider mDisplayProvider;
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -162,263 +74,85 @@
         }
         checkAndRequestRecordAudioPermission();
         mExecutorService = Executors.newCachedThreadPool();
+
+        ConfigJson json = ConfigJson.from(VM_CONFIG_PATH);
+        VirtualMachineConfig config = json.toConfig(this);
+
+        Runner runner;
         try {
-            // To ensure that the previous display service is removed.
-            IVirtualizationServiceInternal.Stub.asInterface(
-                            ServiceManager.waitForService("android.system.virtualizationservice"))
-                    .clearDisplayService();
-        } catch (RemoteException e) {
-            Log.d(TAG, "failed to clearDisplayService");
-        }
-        getWindow().setDecorFitsSystemWindows(false);
-        setContentView(R.layout.activity_main);
-        VirtualMachineCallback callback =
-                new VirtualMachineCallback() {
-                    // store reference to ExecutorService to avoid race condition
-                    private final ExecutorService mService = mExecutorService;
-
-                    @Override
-                    public void onPayloadStarted(VirtualMachine vm) {
-                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
-                        // this.
-                    }
-
-                    @Override
-                    public void onPayloadReady(VirtualMachine vm) {
-                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
-                        // this.
-                    }
-
-                    @Override
-                    public void onPayloadFinished(VirtualMachine vm, int exitCode) {
-                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
-                        // this.
-                    }
-
-                    @Override
-                    public void onError(VirtualMachine vm, int errorCode, String message) {
-                        Log.e(TAG, "Error from VM. code: " + errorCode + " (" + message + ")");
-                        setResult(RESULT_CANCELED);
-                        finish();
-                    }
-
-                    @Override
-                    public void onStopped(VirtualMachine vm, int reason) {
-                        Log.d(TAG, "VM stopped. Reason: " + reason);
-                        setResult(RESULT_OK);
-                        finish();
-                    }
-                };
-
-        try {
-            VirtualMachineConfig config =
-                    VmConfigJson.from("/data/local/tmp/vm_config.json").toConfig(this);
-            VirtualMachineManager vmm =
-                    getApplication().getSystemService(VirtualMachineManager.class);
-            if (vmm == null) {
-                Log.e(TAG, "vmm is null");
-                return;
-            }
-            mVirtualMachine = vmm.getOrCreate(VM_NAME, config);
-            try {
-                mVirtualMachine.setConfig(config);
-            } catch (VirtualMachineException e) {
-                vmm.delete(VM_NAME);
-                mVirtualMachine = vmm.create(VM_NAME, config);
-                Log.e(TAG, "error for setting VM config", e);
-            }
-
-            Log.d(TAG, "vm start");
-            mVirtualMachine.run();
-            mVirtualMachine.setCallback(Executors.newSingleThreadExecutor(), callback);
-            if (DEBUG) {
-                InputStream console = mVirtualMachine.getConsoleOutput();
-                InputStream log = mVirtualMachine.getLogOutput();
-                OutputStream consoleLogFile =
-                        new LineBufferedOutputStream(
-                                getApplicationContext().openFileOutput("console.log", 0));
-                mExecutorService.execute(new CopyStreamTask("console", console, consoleLogFile));
-                mExecutorService.execute(new Reader("log", log));
-            }
-        } catch (VirtualMachineException | IOException e) {
+            runner = Runner.create(this, config);
+        } catch (VirtualMachineException e) {
             throw new RuntimeException(e);
         }
-
-        SurfaceView surfaceView = findViewById(R.id.surface_view);
-        SurfaceView cursorSurfaceView = findViewById(R.id.cursor_surface_view);
-        cursorSurfaceView.setZOrderMediaOverlay(true);
-        View backgroundTouchView = findViewById(R.id.background_touch_view);
-        backgroundTouchView.setOnTouchListener(
-                (v, event) -> {
-                    if (mVirtualMachine == null) {
-                        return false;
-                    }
-                    return mVirtualMachine.sendMultiTouchEvent(event);
-                });
-        surfaceView.requestUnbufferedDispatch(InputDevice.SOURCE_ANY);
-        surfaceView.setOnCapturedPointerListener(
-                (v, event) -> {
-                    if (mVirtualMachine == null) {
-                        return false;
-                    }
-                    int eventSource = event.getSource();
-                    if ((eventSource & InputDevice.SOURCE_CLASS_POSITION) != 0) {
-                        return mVirtualMachine.sendTrackpadEvent(event);
-                    }
-                    return mVirtualMachine.sendMouseEvent(event);
-                });
-        surfaceView
-                .getHolder()
-                .addCallback(
-                        // TODO(b/331708504): it should be handled in AVF framework.
-                        new SurfaceHolder.Callback() {
-                            @Override
-                            public void surfaceCreated(SurfaceHolder holder) {
-                                Log.d(
-                                        TAG,
-                                        "surface size: "
-                                                + holder.getSurfaceFrame().flattenToString());
-                                Log.d(
-                                        TAG,
-                                        "ICrosvmAndroidDisplayService.setSurface("
-                                                + holder.getSurface()
-                                                + ")");
-                                runWithDisplayService(
-                                        s ->
-                                                s.setSurface(
-                                                        holder.getSurface(),
-                                                        false /* forCursor */));
-                                // TODO  execute the above and the below togther with the same call
-                                // to runWithDisplayService. Currently this doesn't work because
-                                // setSurface somtimes trigger an exception and as a result
-                                // drawSavedFrameForSurface is skipped.
-                                runWithDisplayService(
-                                        s -> s.drawSavedFrameForSurface(false /* forCursor */));
-                            }
-
-                            @Override
-                            public void surfaceChanged(
-                                    SurfaceHolder holder, int format, int width, int height) {
-                                Log.d(
-                                        TAG,
-                                        "surface changed, width: " + width + ", height: " + height);
-                            }
-
-                            @Override
-                            public void surfaceDestroyed(SurfaceHolder holder) {
-                                Log.d(TAG, "ICrosvmAndroidDisplayService.removeSurface()");
-                                runWithDisplayService(
-                                        (service) -> service.removeSurface(false /* forCursor */));
-                            }
+        mVirtualMachine = runner.getVm();
+        runner.getExitStatus()
+                .thenAcceptAsync(
+                        success -> {
+                            setResult(success ? RESULT_OK : RESULT_CANCELED);
+                            finish();
                         });
-        cursorSurfaceView.getHolder().setFormat(PixelFormat.RGBA_8888);
-        cursorSurfaceView
-                .getHolder()
-                .addCallback(
-                        new SurfaceHolder.Callback() {
-                            @Override
-                            public void surfaceCreated(SurfaceHolder holder) {
-                                try {
-                                    ParcelFileDescriptor[] pfds =
-                                            ParcelFileDescriptor.createSocketPair();
-                                    if (mCursorHandler != null) {
-                                        mCursorHandler.interrupt();
-                                    }
-                                    mCursorHandler =
-                                            new CursorHandler(
-                                                    surfaceView.getSurfaceControl(),
-                                                    cursorSurfaceView.getSurfaceControl(),
-                                                    pfds[0]);
-                                    mCursorHandler.start();
-                                    runWithDisplayService(
-                                            (service) -> service.setCursorStream(pfds[1]));
-                                } catch (Exception e) {
-                                    Log.d(TAG, "failed to run cursor stream handler", e);
-                                }
-                                Log.d(
-                                        TAG,
-                                        "ICrosvmAndroidDisplayService.setSurface("
-                                                + holder.getSurface()
-                                                + ")");
-                                runWithDisplayService(
-                                        (service) ->
-                                                service.setSurface(
-                                                        holder.getSurface(), true /* forCursor */));
-                            }
 
-                            @Override
-                            public void surfaceChanged(
-                                    SurfaceHolder holder, int format, int width, int height) {
-                                Log.d(
-                                        TAG,
-                                        "cursor surface changed, width: "
-                                                + width
-                                                + ", height: "
-                                                + height);
-                            }
+        // Setup UI
+        setContentView(R.layout.activity_main);
+        SurfaceView mainView = findViewById(R.id.surface_view);
+        SurfaceView cursorView = findViewById(R.id.cursor_surface_view);
+        View touchView = findViewById(R.id.background_touch_view);
+        makeFullscreen();
 
-                            @Override
-                            public void surfaceDestroyed(SurfaceHolder holder) {
-                                Log.d(TAG, "ICrosvmAndroidDisplayService.removeSurface()");
-                                runWithDisplayService(
-                                        (service) -> service.removeSurface(true /* forCursor */));
-                            }
-                        });
-        getWindow().addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON);
+        // Connect the views to the VM
+        mInputForwarder = new InputForwarder(this, mVirtualMachine, touchView, mainView, mainView);
+        mDisplayProvider = new DisplayProvider(mainView, cursorView);
 
-        // Fullscreen:
-        WindowInsetsController windowInsetsController = surfaceView.getWindowInsetsController();
-        windowInsetsController.setSystemBarsBehavior(
+        Path logPath = getFileStreamPath(mVirtualMachine.getName() + ".log").toPath();
+        Logger.setup(mVirtualMachine, logPath, mExecutorService);
+    }
+
+    private void makeFullscreen() {
+        Window w = getWindow();
+        w.setDecorFitsSystemWindows(false);
+        WindowInsetsController insetsCtrl = w.getInsetsController();
+        insetsCtrl.hide(WindowInsets.Type.systemBars());
+        insetsCtrl.setSystemBarsBehavior(
                 WindowInsetsController.BEHAVIOR_SHOW_TRANSIENT_BARS_BY_SWIPE);
-        windowInsetsController.hide(WindowInsets.Type.systemBars());
-        registerInputDeviceListener();
     }
 
     @Override
     protected void onResume() {
         super.onResume();
-        setTabletModeConditionally();
+        mInputForwarder.setTabletModeConditionally();
     }
 
     @Override
     protected void onPause() {
         super.onPause();
-        runWithDisplayService(s -> s.saveFrameForSurface(false /* forCursor */));
+        mDisplayProvider.notifyDisplayIsGoingToInvisible();
     }
 
     @Override
     protected void onStop() {
         super.onStop();
-        if (mVirtualMachine != null) {
-            try {
-                mVirtualMachine.sendLidEvent(/* close */ true);
-                mVirtualMachine.suspend();
-            } catch (VirtualMachineException e) {
-                Log.e(TAG, "Failed to suspend VM" + e);
-            }
+        try {
+            mVirtualMachine.suspend();
+        } catch (VirtualMachineException e) {
+            Log.e(TAG, "Failed to suspend VM" + e);
         }
     }
 
     @Override
     protected void onRestart() {
         super.onRestart();
-        if (mVirtualMachine != null) {
-            try {
-                mVirtualMachine.resume();
-                mVirtualMachine.sendLidEvent(/* close */ false);
-            } catch (VirtualMachineException e) {
-                Log.e(TAG, "Failed to resume VM" + e);
-            }
+        try {
+            mVirtualMachine.resume();
+        } catch (VirtualMachineException e) {
+            Log.e(TAG, "Failed to resume VM" + e);
         }
     }
 
     @Override
     protected void onDestroy() {
         super.onDestroy();
-        if (mExecutorService != null) {
-            mExecutorService.shutdownNow();
-        }
-        unregisterInputDeviceListener();
+        mExecutorService.shutdownNow();
+        mInputForwarder.cleanUp();
         Log.d(TAG, "destroyed");
     }
 
@@ -571,73 +305,6 @@
         }
     }
 
-    @FunctionalInterface
-    public interface RemoteExceptionCheckedFunction<T> {
-        void apply(T t) throws RemoteException;
-    }
-
-    private void runWithDisplayService(
-            RemoteExceptionCheckedFunction<ICrosvmAndroidDisplayService> func) {
-        IVirtualizationServiceInternal vs =
-                IVirtualizationServiceInternal.Stub.asInterface(
-                        ServiceManager.waitForService("android.system.virtualizationservice"));
-        try {
-            Log.d(TAG, "wait for the display service");
-            ICrosvmAndroidDisplayService service =
-                    ICrosvmAndroidDisplayService.Stub.asInterface(vs.waitDisplayService());
-            assert service != null;
-            func.apply(service);
-            Log.d(TAG, "display service runs successfully");
-        } catch (Exception e) {
-            Log.d(TAG, "error on running display service", e);
-        }
-    }
-
-    static class CursorHandler extends Thread {
-        private final SurfaceControl mCursor;
-        private final ParcelFileDescriptor mStream;
-        private final SurfaceControl.Transaction mTransaction;
-
-        CursorHandler(SurfaceControl main, SurfaceControl cursor, ParcelFileDescriptor stream) {
-            mCursor = cursor;
-            mStream = stream;
-            mTransaction = new SurfaceControl.Transaction();
-
-            mTransaction.reparent(cursor, main).apply();
-        }
-
-        @Override
-        public void run() {
-            Log.d(TAG, "running CursorHandler");
-            try {
-                ByteBuffer byteBuffer = ByteBuffer.allocate(8 /* (x: u32, y: u32) */);
-                byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
-                while (true) {
-                    if (Thread.interrupted()) {
-                        Log.d(TAG, "interrupted: exiting CursorHandler");
-                        return;
-                    }
-                    byteBuffer.clear();
-                    int bytes =
-                            IoBridge.read(
-                                    mStream.getFileDescriptor(),
-                                    byteBuffer.array(),
-                                    0,
-                                    byteBuffer.array().length);
-                    if (bytes == -1) {
-                        Log.e(TAG, "cannot read from cursor stream, stop the handler");
-                        return;
-                    }
-                    float x = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
-                    float y = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
-                    mTransaction.setPosition(mCursor, x, y).apply();
-                }
-            } catch (IOException e) {
-                Log.e(TAG, "failed to run CursorHandler", e);
-            }
-        }
-    }
-
     private void checkAndRequestRecordAudioPermission() {
         if (getApplicationContext().checkSelfPermission(permission.RECORD_AUDIO)
                 != PERMISSION_GRANTED) {
@@ -646,72 +313,4 @@
         }
     }
 
-    /** Reads data from an input stream and posts it to the output data */
-    static class Reader implements Runnable {
-        private final String mName;
-        private final InputStream mStream;
-
-        Reader(String name, InputStream stream) {
-            mName = name;
-            mStream = stream;
-        }
-
-        @Override
-        public void run() {
-            try {
-                BufferedReader reader = new BufferedReader(new InputStreamReader(mStream));
-                String line;
-                while ((line = reader.readLine()) != null && !Thread.interrupted()) {
-                    Log.d(TAG, mName + ": " + line);
-                }
-            } catch (IOException e) {
-                Log.e(TAG, "Exception while posting " + mName + " output: " + e.getMessage());
-            }
-        }
-    }
-
-    private static class CopyStreamTask implements Runnable {
-        private final String mName;
-        private final InputStream mIn;
-        private final OutputStream mOut;
-
-        CopyStreamTask(String name, InputStream in, OutputStream out) {
-            mName = name;
-            mIn = in;
-            mOut = out;
-        }
-
-        @Override
-        public void run() {
-            try {
-                byte[] buffer = new byte[2048];
-                while (!Thread.interrupted()) {
-                    int len = mIn.read(buffer);
-                    if (len < 0) {
-                        break;
-                    }
-                    mOut.write(buffer, 0, len);
-                }
-            } catch (Exception e) {
-                Log.e(TAG, "Exception while posting " + mName, e);
-            }
-        }
-    }
-
-    private static class LineBufferedOutputStream extends BufferedOutputStream {
-        LineBufferedOutputStream(OutputStream out) {
-            super(out);
-        }
-
-        @Override
-        public void write(byte[] buf, int off, int len) throws IOException {
-            super.write(buf, off, len);
-            for (int i = 0; i < len; ++i) {
-                if (buf[off + i] == '\n') {
-                    flush();
-                    break;
-                }
-            }
-        }
-    }
 }
diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/Runner.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/Runner.java
new file mode 100644
index 0000000..f50ec86
--- /dev/null
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/Runner.java
@@ -0,0 +1,111 @@
+/*
+ * Copyright (C) 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.vmlauncher;
+
+import android.content.Context;
+import android.system.virtualmachine.VirtualMachine;
+import android.system.virtualmachine.VirtualMachineCallback;
+import android.system.virtualmachine.VirtualMachineConfig;
+import android.system.virtualmachine.VirtualMachineCustomImageConfig;
+import android.system.virtualmachine.VirtualMachineException;
+import android.system.virtualmachine.VirtualMachineManager;
+import android.util.Log;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ForkJoinPool;
+
+/** Utility class for creating a VM and waiting for it to finish. */
+class Runner {
+    private static final String TAG = MainActivity.TAG;
+    private final VirtualMachine mVirtualMachine;
+    private final Callback mCallback;
+
+    private Runner(VirtualMachine vm, Callback cb) {
+        mVirtualMachine = vm;
+        mCallback = cb;
+    }
+
+    /** Create a virtual machine of the given config, under the given context. */
+    static Runner create(Context context, VirtualMachineConfig config)
+            throws VirtualMachineException {
+        VirtualMachineManager vmm = context.getSystemService(VirtualMachineManager.class);
+        VirtualMachineCustomImageConfig customConfig = config.getCustomImageConfig();
+        if (customConfig == null) {
+            throw new RuntimeException("CustomImageConfig is missing");
+        }
+
+        String name = customConfig.getName();
+        if (name == null || name.isEmpty()) {
+            throw new RuntimeException("Virtual machine's name is missing in the config");
+        }
+
+        VirtualMachine vm = vmm.getOrCreate(name, config);
+        try {
+            vm.setConfig(config);
+        } catch (VirtualMachineException e) {
+            vmm.delete(name);
+            vm = vmm.create(name, config);
+            Log.w(TAG, "Re-creating virtual machine (" + name + ")", e);
+        }
+
+        Callback cb = new Callback();
+        vm.setCallback(ForkJoinPool.commonPool(), cb);
+        vm.run();
+        return new Runner(vm, cb);
+    }
+
+    /** Give access to the underlying VirtualMachine object. */
+    VirtualMachine getVm() {
+        return mVirtualMachine;
+    }
+
+    /** Get future about VM's exit status. */
+    CompletableFuture<Boolean> getExitStatus() {
+        return mCallback.mFinishedSuccessfully;
+    }
+
+    private static class Callback implements VirtualMachineCallback {
+        final CompletableFuture<Boolean> mFinishedSuccessfully = new CompletableFuture<>();
+
+        @Override
+        public void onPayloadStarted(VirtualMachine vm) {
+            // This event is only from Microdroid-based VM. Custom VM shouldn't emit this.
+        }
+
+        @Override
+        public void onPayloadReady(VirtualMachine vm) {
+            // This event is only from Microdroid-based VM. Custom VM shouldn't emit this.
+        }
+
+        @Override
+        public void onPayloadFinished(VirtualMachine vm, int exitCode) {
+            // This event is only from Microdroid-based VM. Custom VM shouldn't emit this.
+        }
+
+        @Override
+        public void onError(VirtualMachine vm, int errorCode, String message) {
+            Log.e(TAG, "Error from VM. code: " + errorCode + " (" + message + ")");
+            mFinishedSuccessfully.complete(false);
+        }
+
+        @Override
+        public void onStopped(VirtualMachine vm, int reason) {
+            Log.d(TAG, "VM stopped. Reason: " + reason);
+            mFinishedSuccessfully.complete(true);
+        }
+    }
+}
diff --git a/android/VmLauncherApp/proguard.flags b/android/VmLauncherApp/proguard.flags
index 5e05ecf..13ec24e 100644
--- a/android/VmLauncherApp/proguard.flags
+++ b/android/VmLauncherApp/proguard.flags
@@ -1,7 +1,7 @@
 # Keep the no-args constructor of the deserialized class
--keepclassmembers class com.android.virtualization.vmlauncher.VmConfigJson {
+-keepclassmembers class com.android.virtualization.vmlauncher.ConfigJson {
   <init>();
 }
--keepclassmembers class com.android.virtualization.vmlauncher.VmConfigJson$* {
+-keepclassmembers class com.android.virtualization.vmlauncher.ConfigJson$* {
   <init>();
 }
diff --git a/docs/custom_vm.md b/docs/custom_vm.md
index cdeddf5..6a1cc00 100644
--- a/docs/custom_vm.md
+++ b/docs/custom_vm.md
@@ -25,19 +25,35 @@
 The `vm` command also has other subcommands for debugging; run
 `/apex/com.android.virt/bin/vm help` for details.
 
-### Running Debian with u-boot
-1. Prepare u-boot binary from `u-boot_crosvm_aarch64` in https://ci.android.com/builds/branches/aosp_u-boot-mainline/grid
-or build it by https://source.android.com/docs/devices/cuttlefish/bootloader-dev#develop-bootloader
-2. Prepare Debian image from https://cloud.debian.org/images/cloud/ (We tested nocloud image)
-3. Copy `u-boot.bin`, Debian image file(like `debian-12-nocloud-arm64.raw`) and `vm_config.json` to `/data/local/tmp`
+### Running Debian
+1. Download an ARM64 image from https://cloud.debian.org/images/cloud/ (We tested nocloud image)
+
+2. Resize the image
+```shell
+truncate -s 20G debian.img
+virt-resize --expand /dev/sda1 <download_image_file> debian.img
+```
+
+3. Copy the image file
+```shell
+tar cfS debian.img.tar debian.img
+adb push debian.img.tar /data/local/tmp/
+adb shell tar xf /data/local/tmp/debian.img.tar -C /data/local/tmp/
+adb shell rm /data/local/tmp/debian.img.tar
+adb shell chmod a+w /data/local/tmp/debian.img
+rm debian.img.tar
+```
+
+Note: we tar and untar to keep the image file sparse.
+
+4. Make the VM config file
 ```shell
 cat > vm_config.json <<EOF
 {
     "name": "debian",
-    "bootloader": "/data/local/tmp/u-boot.bin",
     "disks": [
         {
-            "image": "/data/local/tmp/debian-12-nocloud-arm64.raw",
+            "image": "/data/local/tmp/debian.img",
             "partitions": [],
             "writable": true
         }
@@ -45,14 +61,51 @@
     "protected": false,
     "cpu_topology": "match_host",
     "platform_version": "~1.0",
-    "memory_mib" : 8096
+    "memory_mib": 8096,
+    "debuggable": true,
+    "console_out": true,
+    "connect_console": true,
+    "console_input_device": "ttyS0",
+    "network": true,
+    "input": {
+        "touchscreen": true,
+        "keyboard": true,
+        "mouse": true,
+        "trackpad": true,
+        "switches": true
+    },
+    "audio": {
+        "speaker": true,
+         "microphone": true
+    },
+    "gpu": {
+        "backend": "virglrenderer",
+        "context_types": ["virgl2"]
+    },
+    "display": {
+        "refresh_rate": "30"
+    }
 }
 EOF
-adb push `u-boot.bin` /data/local/tmp
-adb push `debian-12-nocloud-arm64.raw` /data/local/tmp
-adb push vm_config.json /data/local/tmp/vm_config.json
+adb push vm_config.json /data/local/tmp/
 ```
-4. Launch VmLauncherApp(the detail will be explain below)
+
+5. Launch VmLauncherApp(the detail will be explain below)
+
+6. For console, we can refer to `Debugging` section below. (id: root)
+
+7. For graphical shell, you need to install xfce(for now, only xfce is tested)
+```
+apt install task-xfce-desktop
+dpkg --configure -a (if necessary)
+systemctl set-default graphical.target
+
+# need non-root user for graphical shell
+adduser linux
+# optional
+adduser linux sudo
+reboot
+```
 
 ## Graphical VMs
 
@@ -192,16 +245,34 @@
             "writable": true
         }
     ],
+    "protected": false,
+    "cpu_topology": "match_host",
+    "platform_version": "~1.0",
+    "memory_mib": 8096,
+    "debuggable": true,
+    "console_out": true,
+    "connect_console": true,
+    "console_input_device": "hvc0",
+    "network": true,
+    "input": {
+        "touchscreen": true,
+        "keyboard": true,
+        "mouse": true,
+        "trackpad": true,
+        "switches": true
+    },
+    "audio": {
+        "speaker": true,
+        "microphone": true
+    },
     "gpu": {
         "backend": "virglrenderer",
         "context_types": ["virgl2"]
     },
-    "params": "root=/dev/vda3 rootwait noinitrd ro enforcing=0 cros_debug cros_secure",
-    "protected": false,
-    "cpu_topology": "match_host",
-    "platform_version": "~1.0",
-    "memory_mib" : 8096,
-    "console_input_device": "hvc0"
+    "display": {
+        "scale": "0.77",
+        "refresh_rate": "30"
+    }
 }
 ```
 
@@ -235,15 +306,15 @@
 ```
 
 To see console logs only, check
-`/data/data/com{,.google}.android.virtualization.vmlauncher/files/console.log`
+`/data/data/com{,.google}.android.virtualization.vmlauncher/files/${vm_name}.log`
 
 For HSUM enabled devices,
-`/data/user/${current_user_id}/com{,.google}.android.virtualization.vmlauncher/files/console.log`
+`/data/user/${current_user_id}/com{,.google}.android.virtualization.vmlauncher/files/${vm_name}.log`
 
 You can monitor console out as follows
 
 ```shell
-$ adb shell 'su root tail +0 -F /data/user/$(am get-current-user)/com{,.google}.android.virtualization.vmlauncher/files/console.log'
+$ adb shell 'su root tail +0 -F /data/user/$(am get-current-user)/com{,.google}.android.virtualization.vmlauncher/files/${vm_name}.log'
 ```
 
 For ChromiumOS, you can enter to the console via SSH connection. Check your IP
diff --git a/guest/pvmfw/image.ld b/guest/pvmfw/image.ld
index 18bb3ba..fb26806 100644
--- a/guest/pvmfw/image.ld
+++ b/guest/pvmfw/image.ld
@@ -18,5 +18,4 @@
 {
 	image		: ORIGIN = 0x7fc00000, LENGTH = 2M
 	writable_data	: ORIGIN = 0x7fe00000, LENGTH = 2M
-	dtb_region	: ORIGIN = 0x80000000, LENGTH = 2M
 }
diff --git a/guest/pvmfw/src/dice.rs b/guest/pvmfw/src/dice.rs
index 8be73a4..470711f 100644
--- a/guest/pvmfw/src/dice.rs
+++ b/guest/pvmfw/src/dice.rs
@@ -36,8 +36,10 @@
 #[derive(Debug)]
 pub enum Error {
     /// Error in CBOR operations
+    #[allow(dead_code)]
     CborError(ciborium::value::Error),
     /// Error in DICE operations
+    #[allow(dead_code)]
     DiceError(diced_open_dice::DiceError),
 }
 
diff --git a/guest/rialto/idmap.S b/guest/rialto/idmap.S
index 7281d9b..9b5375a 100644
--- a/guest/rialto/idmap.S
+++ b/guest/rialto/idmap.S
@@ -28,7 +28,6 @@
 .set .PAGE_SIZE, .SZ_4K
 
 .set .ORIGIN_ADDR, 2 * .SZ_1G
-.set .DTB_ADDR, .ORIGIN_ADDR + (0 * .SZ_2M)
 .set .TEXT_ADDR, .ORIGIN_ADDR + (1 * .SZ_2M)
 .set .DATA_ADDR, .ORIGIN_ADDR + (2 * .SZ_2M)
 
@@ -60,7 +59,7 @@
 	.balign .PAGE_SIZE, 0				// unmapped
 
 	/* level 2 */
-0:	.quad		.L_BLOCK_RO  | .DTB_ADDR	// DT provided by VMM
+0:	.quad		0x0				// 2 MiB unmapped
 	.quad		.L_BLOCK_MEM_XIP | .TEXT_ADDR	// 2 MiB of DRAM containing image
 	.quad		.L_BLOCK_MEM | .DATA_ADDR	// 2 MiB of writable DRAM
 	.balign .PAGE_SIZE, 0				// unmapped
diff --git a/guest/rialto/image.ld b/guest/rialto/image.ld
index 368acbb..95ffdf8 100644
--- a/guest/rialto/image.ld
+++ b/guest/rialto/image.ld
@@ -16,7 +16,6 @@
 
 MEMORY
 {
-	dtb_region	: ORIGIN = 0x80000000, LENGTH = 2M
 	image		: ORIGIN = 0x80200000, LENGTH = 2M
 	writable_data	: ORIGIN = 0x80400000, LENGTH = 2M
 }
diff --git a/libs/libvmbase/example/image.ld b/libs/libvmbase/example/image.ld
index 368acbb..95ffdf8 100644
--- a/libs/libvmbase/example/image.ld
+++ b/libs/libvmbase/example/image.ld
@@ -16,7 +16,6 @@
 
 MEMORY
 {
-	dtb_region	: ORIGIN = 0x80000000, LENGTH = 2M
 	image		: ORIGIN = 0x80200000, LENGTH = 2M
 	writable_data	: ORIGIN = 0x80400000, LENGTH = 2M
 }
diff --git a/libs/libvmbase/example/src/layout.rs b/libs/libvmbase/example/src/layout.rs
index fc578bc..49e4aa7 100644
--- a/libs/libvmbase/example/src/layout.rs
+++ b/libs/libvmbase/example/src/layout.rs
@@ -29,8 +29,6 @@
 }
 
 pub fn print_addresses() {
-    let dtb = layout::dtb_range();
-    info!("dtb:        {}..{} ({} bytes)", dtb.start, dtb.end, dtb.end - dtb.start);
     let text = layout::text_range();
     info!("text:       {}..{} ({} bytes)", text.start, text.end, text.end - text.start);
     let rodata = layout::rodata_range();
diff --git a/libs/libvmbase/example/src/main.rs b/libs/libvmbase/example/src/main.rs
index da82b17..a01f619 100644
--- a/libs/libvmbase/example/src/main.rs
+++ b/libs/libvmbase/example/src/main.rs
@@ -26,6 +26,7 @@
 use crate::layout::{boot_stack_range, print_addresses, DEVICE_REGION};
 use crate::pci::{check_pci, get_bar_region};
 use aarch64_paging::paging::MemoryRegion;
+use aarch64_paging::paging::VirtualAddress;
 use aarch64_paging::MapError;
 use alloc::{vec, vec::Vec};
 use core::ptr::addr_of_mut;
@@ -35,7 +36,10 @@
 use log::{debug, error, info, trace, warn, LevelFilter};
 use vmbase::{
     bionic, configure_heap,
-    layout::{dtb_range, rodata_range, scratch_range, text_range},
+    layout::{
+        crosvm::{FDT_MAX_SIZE, MEM_START},
+        rodata_range, scratch_range, text_range,
+    },
     linker, logger, main,
     memory::{PageTable, SIZE_64KB},
 };
@@ -47,7 +51,7 @@
 main!(main);
 configure_heap!(SIZE_64KB);
 
-fn init_page_table(pci_bar_range: &MemoryRegion) -> Result<(), MapError> {
+fn init_page_table(dtb: &MemoryRegion, pci_bar_range: &MemoryRegion) -> Result<(), MapError> {
     let mut page_table = PageTable::default();
 
     page_table.map_device(&DEVICE_REGION)?;
@@ -55,7 +59,7 @@
     page_table.map_rodata(&rodata_range().into())?;
     page_table.map_data(&scratch_range().into())?;
     page_table.map_data(&boot_stack_range().into())?;
-    page_table.map_rodata(&dtb_range().into())?;
+    page_table.map_rodata(dtb)?;
     page_table.map_device(pci_bar_range)?;
 
     info!("Activating IdMap...");
@@ -76,15 +80,16 @@
     info!("Hello world");
     info!("x0={:#018x}, x1={:#018x}, x2={:#018x}, x3={:#018x}", arg0, arg1, arg2, arg3);
     print_addresses();
-    assert_eq!(arg0, dtb_range().start.0 as u64);
     check_data();
     check_stack_guard();
 
     info!("Checking FDT...");
-    let fdt = dtb_range();
-    let fdt_size = fdt.end.0 - fdt.start.0;
+    let fdt_addr = usize::try_from(arg0).unwrap();
+    // We are about to access the region so check that it matches our page tables in idmap.S.
+    assert_eq!(fdt_addr, MEM_START);
     // SAFETY: The DTB range is valid, writable memory, and we don't construct any aliases to it.
-    let fdt = unsafe { core::slice::from_raw_parts_mut(fdt.start.0 as *mut u8, fdt_size) };
+    let fdt = unsafe { core::slice::from_raw_parts_mut(fdt_addr as *mut u8, FDT_MAX_SIZE) };
+    let fdt_region = (VirtualAddress(fdt_addr)..VirtualAddress(fdt_addr + fdt.len())).into();
     let fdt = Fdt::from_mut_slice(fdt).unwrap();
     info!("FDT passed verification.");
     check_fdt(fdt);
@@ -96,7 +101,7 @@
 
     check_alloc();
 
-    init_page_table(&get_bar_region(&pci_info)).unwrap();
+    init_page_table(&fdt_region, &get_bar_region(&pci_info)).unwrap();
 
     check_data();
     check_dice();
diff --git a/libs/libvmbase/sections.ld b/libs/libvmbase/sections.ld
index c7ef0ec..01b7e39 100644
--- a/libs/libvmbase/sections.ld
+++ b/libs/libvmbase/sections.ld
@@ -29,12 +29,6 @@
 
 SECTIONS
 {
-	.dtb (NOLOAD) : {
-		dtb_begin = .;
-		. += LENGTH(dtb_region);
-		dtb_end = .;
-	} >dtb_region
-
 	/*
 	 * Collect together the code. This is page aligned so it can be mapped
 	 * as executable-only.
diff --git a/libs/libvmbase/src/layout.rs b/libs/libvmbase/src/layout.rs
index 5ac435f..adcb2fa 100644
--- a/libs/libvmbase/src/layout.rs
+++ b/libs/libvmbase/src/layout.rs
@@ -60,11 +60,6 @@
     }};
 }
 
-/// Memory reserved for the DTB.
-pub fn dtb_range() -> Range<VirtualAddress> {
-    linker_region!(dtb_begin, dtb_end)
-}
-
 /// Executable code.
 pub fn text_range() -> Range<VirtualAddress> {
     linker_region!(text_begin, text_end)
diff --git a/tests/ferrochrome/Android.bp b/tests/ferrochrome/Android.bp
index f165b8f..f1b7f27 100644
--- a/tests/ferrochrome/Android.bp
+++ b/tests/ferrochrome/Android.bp
@@ -19,8 +19,7 @@
     name: "ferrochrome-tests.sh",
     srcs: ["ferrochrome.sh"],
     out: ["ferrochrome-tests"],
-    // This breaks shebang, but test will execute the script with bash
-    cmd: "echo \"set -x\" > $(out); cat $(in) >> $(out)",
+    cmd: "sed '2 i set -x' $(in) > $(out)",
 }
 
 sh_binary_host {
diff --git a/tests/ferrochrome/AndroidTest.xml b/tests/ferrochrome/AndroidTest.xml
index 9eaaed3..6c975be 100644
--- a/tests/ferrochrome/AndroidTest.xml
+++ b/tests/ferrochrome/AndroidTest.xml
@@ -35,13 +35,19 @@
         <option name="run-command" value="mkdir /data/local/tmp" />
         <option name="teardown-command" value="pkill vmlauncher" />
         <option name="teardown-command" value="rm /data/local/tmp/chromiumos_base_image.bin" />
+        <option name="teardown-command" value="rm -rf /data/local/tmp/ferrochrome_screenshots" />
     </target_preparer>
 
-    <test class="com.android.tradefed.testtype.binary.ExecutableHostTest" >
+    <test class="com.android.tradefed.testtype.binary.ExecutableHostTest">
         <option name="binary" value="ferrochrome-tests" />
         <option name="relative-path-execution" value="true" />
         <option name="runtime-hint" value="10m" />
         <option name="per-binary-timeout" value="20m" />
     </test>
+
+    <metrics_collector class="com.android.tradefed.device.metric.FilePullerLogCollector">
+        <option name="directory-keys" value="/data/local/tmp/ferrochrome_screenshots" />
+        <option name="collect-on-run-ended-only" value="true" />
+    </metrics_collector>
 </configuration>
 
diff --git a/tests/ferrochrome/assets/vm_config.json b/tests/ferrochrome/assets/vm_config.json
index 3053626..53e3b72 100644
--- a/tests/ferrochrome/assets/vm_config.json
+++ b/tests/ferrochrome/assets/vm_config.json
@@ -7,14 +7,32 @@
             "writable": true
         }
     ],
+    "protected": false,
+    "cpu_topology": "match_host",
+    "platform_version": "~1.0",
+    "memory_mib": 8096,
+    "debuggable": true,
+    "console_out": true,
+    "connect_console": true,
+    "console_input_device": "hvc0",
+    "network": true,
+    "input": {
+        "touchscreen": true,
+        "keyboard": true,
+        "mouse": true,
+        "trackpad": true,
+        "switches": true
+    },
+    "audio": {
+        "speaker": true,
+        "microphone": true
+    },
     "gpu": {
         "backend": "virglrenderer",
         "context_types": ["virgl2"]
     },
-    "params": "root=/dev/vda3 rootwait noinitrd ro enforcing=0 cros_debug cros_secure",
-    "protected": false,
-    "cpu_topology": "match_host",
-    "platform_version": "~1.0",
-    "memory_mib" : 8096,
-    "console_input_device": "hvc0"
+    "display": {
+        "scale": "0.77",
+        "refresh_rate": "30"
+    }
 }
diff --git a/tests/ferrochrome/ferrochrome.sh b/tests/ferrochrome/ferrochrome.sh
index 683b82e..b9b9fbc 100755
--- a/tests/ferrochrome/ferrochrome.sh
+++ b/tests/ferrochrome/ferrochrome.sh
@@ -21,12 +21,13 @@
 
 FECR_GS_URL="https://storage.googleapis.com/chromiumos-image-archive/ferrochrome-public"
 FECR_DEFAULT_VERSION="R128-15958.0.0"
+FECR_DEFAULT_SCREENSHOT_DIR="/data/local/tmp/ferrochrome_screenshots"  # Hardcoded at AndroidTest.xml
 FECR_TEST_IMAGE="chromiumos_test_image"
 FECR_BASE_IMAGE="chromiumos_base_image"
 FECR_DEVICE_DIR="/data/local/tmp"
 FECR_IMAGE_VM_CONFIG_JSON="chromiumos_base_image.bin"  # hardcoded at vm_config.json
 FECR_CONFIG_PATH="/data/local/tmp/vm_config.json"  # hardcoded at VmLauncherApp
-FECR_CONSOLE_LOG_PATH="/data/data/\${pkg_name}/files/console.log"
+FECR_CONSOLE_LOG_PATH="files/cros.log" # log file name is ${vm_name}.log
 FECR_TEST_IMAGE_BOOT_COMPLETED_LOG="Have fun and send patches!"
 FECR_BASE_IMAGE_BOOT_COMPLETED_LOG="Chrome started, our work is done, exiting"
 FECR_BOOT_TIMEOUT="300" # 5 minutes (300 seconds)
@@ -74,6 +75,7 @@
 fecr_verbose=""
 fecr_image="${FECR_DEFAULT_IMAGE}"
 fecr_boot_completed_log="${FECR_DEFAULT_BOOT_COMPLETED_LOG}"
+fecr_screenshot_dir="${FECR_DEFAULT_SCREENSHOT_DIR}"
 
 # Parse parameters
 while (( "${#}" )); do
@@ -153,18 +155,20 @@
 adb shell svc power stayon true
 adb shell wm dismiss-keyguard
 
+echo "Granting runtime permissions to ensure VmLauncher is focused"
+adb shell pm grant ${pkg_name} android.permission.RECORD_AUDIO
+
 echo "Starting ferrochrome"
 adb shell am start-activity -a ${ACTION_NAME} > /dev/null
 
-if [[ $(adb shell getprop ro.fw.mu.headless_system_user) == "true" ]]; then
-  current_user=$(adb shell am get-current-user)
-  log_path="/data/user/${current_user}/${pkg_name}/files/console.log"
-else
-  log_path="/data/data/${pkg_name}/files/console.log"
-fi
+# HSUM aware log path
+current_user=$(adb shell am get-current-user)
+log_path="/data/user/${current_user}/${pkg_name}/${FECR_CONSOLE_LOG_PATH}"
 fecr_start_time=${EPOCHSECONDS}
 
+adb shell mkdir -p "${fecr_screenshot_dir}"
 while [[ $((EPOCHSECONDS - fecr_start_time)) -lt ${FECR_BOOT_TIMEOUT} ]]; do
+  adb shell screencap -a -p "${fecr_screenshot_dir}/screenshot-${EPOCHSECONDS}.png"
   adb shell grep -soF \""${fecr_boot_completed_log}"\" "${log_path}" && exit 0 || true
   sleep 10
 done