Merge "rialto: Remove early ID-mapping for DTB" into main
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..ccfff95
--- /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>();
 }