Merge "[avb] Define rust API for image verification and connect pvmfw"
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index ffcdc51..63b5628 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -49,7 +49,9 @@
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.SystemApi;
+import android.content.ComponentCallbacks2;
 import android.content.Context;
+import android.content.res.Configuration;
 import android.os.Binder;
 import android.os.IBinder;
 import android.os.ParcelFileDescriptor;
@@ -61,6 +63,7 @@
 import android.system.virtualizationservice.IVirtualMachine;
 import android.system.virtualizationservice.IVirtualMachineCallback;
 import android.system.virtualizationservice.IVirtualizationService;
+import android.system.virtualizationservice.MemoryTrimLevel;
 import android.system.virtualizationservice.PartitionType;
 import android.system.virtualizationservice.VirtualMachineAppConfig;
 import android.system.virtualizationservice.VirtualMachineState;
@@ -193,6 +196,56 @@
      */
     @NonNull private final List<ExtraApkSpec> mExtraApks;
 
+    private class MemoryManagementCallbacks implements ComponentCallbacks2 {
+        @Override
+        public void onConfigurationChanged(@NonNull Configuration newConfig) {}
+
+        @Override
+        public void onLowMemory() {}
+
+        @Override
+        public void onTrimMemory(int level) {
+            @MemoryTrimLevel int vmTrimLevel;
+
+            switch (level) {
+                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_CRITICAL:
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_CRITICAL;
+                    break;
+                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_LOW:
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_LOW;
+                    break;
+                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_MODERATE:
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_MODERATE;
+                    break;
+                case ComponentCallbacks2.TRIM_MEMORY_BACKGROUND:
+                case ComponentCallbacks2.TRIM_MEMORY_MODERATE:
+                case ComponentCallbacks2.TRIM_MEMORY_COMPLETE:
+                    /* Release as much memory as we can. The app is on the LMKD LRU kill list. */
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_CRITICAL;
+                    break;
+                default:
+                    /* Treat unrecognised messages as generic low-memory warnings. */
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_LOW;
+                    break;
+            }
+
+            synchronized (mLock) {
+                try {
+                    if (mVirtualMachine != null) {
+                        mVirtualMachine.onTrimMemory(vmTrimLevel);
+                    }
+                } catch (Exception e) {
+                    /* Caller doesn't want our exceptions. Log them instead. */
+                    Log.w(TAG, "TrimMemory failed: ", e);
+                }
+            }
+        }
+    }
+
+    @NonNull private final MemoryManagementCallbacks mMemoryManagementCallbacks;
+
+    @NonNull private final Context mContext;
+
     // A note on lock ordering:
     // You can take mLock while holding VirtualMachineManager.sCreateLock, but not vice versa.
     // We never take any other lock while holding mCallbackLock; therefore you can
@@ -268,6 +321,8 @@
         mInstanceFilePath = new File(thisVmDir, INSTANCE_IMAGE_FILE);
         mIdsigFilePath = new File(thisVmDir, IDSIG_FILE);
         mExtraApks = setupExtraApks(context, config, thisVmDir);
+        mMemoryManagementCallbacks = new MemoryManagementCallbacks();
+        mContext = context;
     }
 
     /**
@@ -696,6 +751,7 @@
                                 executeCallback((cb) -> cb.onRamdump(VirtualMachine.this, ramdump));
                             }
                         });
+                mContext.registerComponentCallbacks(mMemoryManagementCallbacks);
                 service.asBinder().linkToDeath(deathRecipient, 0);
                 mVirtualMachine.start();
             } catch (IOException | IllegalStateException | ServiceSpecificException e) {
@@ -772,6 +828,7 @@
             }
             try {
                 mVirtualMachine.stop();
+                mContext.unregisterComponentCallbacks(mMemoryManagementCallbacks);
                 mVirtualMachine = null;
             } catch (RemoteException e) {
                 throw e.rethrowAsRuntimeException();
@@ -797,6 +854,7 @@
             try {
                 if (stateToStatus(mVirtualMachine.getState()) == STATUS_RUNNING) {
                     mVirtualMachine.stop();
+                    mContext.unregisterComponentCallbacks(mMemoryManagementCallbacks);
                     mVirtualMachine = null;
                 }
             } catch (RemoteException e) {
diff --git a/microdroid/init.rc b/microdroid/init.rc
index 7d04557..310cf2b 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -45,7 +45,7 @@
 
     setprop ro.debuggable ${ro.boot.microdroid.debuggable:-0}
 
-on property:dev.bootcomplete=1
+on property:microdroid_manager.init_done=1
     # Stop ueventd to save memory
     stop ueventd
 
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index f8e7d34..3859785 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -27,6 +27,12 @@
     /** Path to the APK contents path. */
     const String VM_APK_CONTENTS_PATH = "/mnt/apk";
 
+    /**
+     * Path to the encrypted storage. Note the path will not exist if encrypted storage
+     * is not enabled.
+     */
+    const String ENCRYPTEDSTORE_MOUNTPOINT = "/mnt/encryptedstore";
+
     /** Notifies that the payload is ready to serve. */
     void notifyPayloadReady();
 
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 4f94bb4..0e45461 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -29,6 +29,7 @@
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
     VM_APK_CONTENTS_PATH,
     VM_PAYLOAD_SERVICE_SOCKET_NAME,
+    ENCRYPTEDSTORE_MOUNTPOINT,
 };
 use anyhow::{anyhow, bail, ensure, Context, Error, Result};
 use apkverify::{get_public_key_der, verify, V4Signature};
@@ -86,7 +87,6 @@
 const ENCRYPTEDSTORE_BIN: &str = "/system/bin/encryptedstore";
 const ENCRYPTEDSTORE_KEY_IDENTIFIER: &str = "encryptedstore_key";
 const ENCRYPTEDSTORE_KEYSIZE: u32 = 32;
-const ENCRYPTEDSTORE_MOUNTPOINT: &str = "/mnt/encryptedstore";
 
 #[derive(thiserror::Error, Debug)]
 enum MicrodroidError {
@@ -433,12 +433,17 @@
 
     register_vm_payload_service(allow_restricted_apis, service.clone(), dice_context)?;
 
+    // Wait for encryptedstore to finish mounting the storage (if enabled) before setting
+    // microdroid_manager.init_done. Reason is init stops uneventd after that.
+    // Encryptedstore, however requires ueventd
     if let Some(mut child) = encryptedstore_child {
         let exitcode = child.wait().context("Wait for encryptedstore child")?;
         ensure!(exitcode.success(), "Unable to prepare encrypted storage. Exitcode={}", exitcode);
     }
 
     wait_for_property_true("dev.bootcomplete").context("failed waiting for dev.bootcomplete")?;
+    system_properties::write("microdroid_manager.init_done", "1")
+        .context("set microdroid_manager.init_done")?;
     info!("boot completed, time to run payload");
     exec_task(task, service).context("Failed to run payload")
 }
diff --git a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
index 16e4893..c8c8660 100644
--- a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
@@ -30,6 +30,9 @@
     /** Returns an entry from /proc/meminfo. */
     long getMemInfoEntry(String name);
 
+    /** Allocates anonymous memory and returns the raw pointer. */
+    long allocAnonMemory(long mb);
+
     /**
      * Initializes the vsock server on VM.
      * @return the server socket file descriptor.
diff --git a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
index eda4f75..077c74f 100644
--- a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
@@ -19,6 +19,8 @@
 interface ITestService {
     const int SERVICE_PORT = 5678;
 
+    const int ECHO_REVERSE_PORT = 6789;
+
     /* add two integers. */
     int addInteger(int a, int b);
 
@@ -39,4 +41,9 @@
 
     /* get the encrypted storage path. */
     String getEncryptedStoragePath();
+
+    /* start a simple vsock server on ECHO_REVERSE_PORT that reads a line at a time and echoes
+     * each line reverse.
+     */
+    void runEchoReverseServer();
 }
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index 14a0e39..3c3faf2 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -27,6 +27,7 @@
 import android.app.Instrumentation;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
+import android.os.Process;
 import android.os.RemoteException;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineConfig;
@@ -279,6 +280,58 @@
         return runInShell(TAG, mInstrumentation.getUiAutomation(), command);
     }
 
+    private static class CrosvmStats {
+        public final long mHostRss;
+        public final long mHostPss;
+        public final long mGuestRss;
+        public final long mGuestPss;
+
+        CrosvmStats(Function<String, String> shellExecutor) {
+            try {
+                List<Integer> crosvmPids =
+                        ProcessUtil.getProcessMap(shellExecutor).entrySet().stream()
+                                .filter(e -> e.getValue().contains("crosvm"))
+                                .map(e -> e.getKey())
+                                .collect(java.util.stream.Collectors.toList());
+                if (crosvmPids.size() != 1) {
+                    throw new IllegalStateException(
+                            "expected to find exactly one crosvm processes, found "
+                                    + crosvmPids.size());
+                }
+
+                long hostRss = 0;
+                long hostPss = 0;
+                long guestRss = 0;
+                long guestPss = 0;
+                boolean hasGuestMaps = false;
+                for (ProcessUtil.SMapEntry entry :
+                        ProcessUtil.getProcessSmaps(crosvmPids.get(0), shellExecutor)) {
+                    long rss = entry.metrics.get("Rss");
+                    long pss = entry.metrics.get("Pss");
+                    if (entry.name.contains("crosvm_guest")) {
+                        guestRss += rss;
+                        guestPss += pss;
+                        hasGuestMaps = true;
+                    } else {
+                        hostRss += rss;
+                        hostPss += pss;
+                    }
+                }
+                if (!hasGuestMaps) {
+                    throw new IllegalStateException(
+                            "found no crosvm_guest smap entry in crosvm process");
+                }
+                mHostRss = hostRss;
+                mHostPss = hostPss;
+                mGuestRss = guestRss;
+                mGuestPss = guestPss;
+            } catch (Exception e) {
+                Log.e(TAG, "Error inside onPayloadReady():" + e);
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
     @Test
     public void testMemoryUsage() throws Exception {
         final String vmName = "test_vm_mem_usage";
@@ -299,10 +352,10 @@
         double mem_buffers = (double) listener.mBuffers / 1024.0;
         double mem_cached = (double) listener.mCached / 1024.0;
         double mem_slab = (double) listener.mSlab / 1024.0;
-        double mem_crosvm_host_rss = (double) listener.mCrosvmHostRss / 1024.0;
-        double mem_crosvm_host_pss = (double) listener.mCrosvmHostPss / 1024.0;
-        double mem_crosvm_guest_rss = (double) listener.mCrosvmGuestRss / 1024.0;
-        double mem_crosvm_guest_pss = (double) listener.mCrosvmGuestPss / 1024.0;
+        double mem_crosvm_host_rss = (double) listener.mCrosvm.mHostRss / 1024.0;
+        double mem_crosvm_host_pss = (double) listener.mCrosvm.mHostPss / 1024.0;
+        double mem_crosvm_guest_rss = (double) listener.mCrosvm.mGuestRss / 1024.0;
+        double mem_crosvm_guest_pss = (double) listener.mCrosvm.mGuestPss / 1024.0;
 
         double mem_kernel = mem_overall - mem_total;
         double mem_used = mem_total - mem_free - mem_buffers - mem_cached - mem_slab;
@@ -327,7 +380,7 @@
             mShellExecutor = shellExecutor;
         }
 
-        public Function<String, String> mShellExecutor;
+        public final Function<String, String> mShellExecutor;
 
         public long mMemTotal;
         public long mMemFree;
@@ -336,10 +389,7 @@
         public long mCached;
         public long mSlab;
 
-        public long mCrosvmHostRss;
-        public long mCrosvmHostPss;
-        public long mCrosvmGuestRss;
-        public long mCrosvmGuestPss;
+        public CrosvmStats mCrosvm;
 
         @Override
         public void onPayloadReady(VirtualMachine vm, IBenchmarkService service)
@@ -350,39 +400,80 @@
             mBuffers = service.getMemInfoEntry("Buffers");
             mCached = service.getMemInfoEntry("Cached");
             mSlab = service.getMemInfoEntry("Slab");
+            mCrosvm = new CrosvmStats(mShellExecutor);
+        }
+    }
 
+    @Test
+    public void testMemoryReclaim() throws Exception {
+        final String vmName = "test_vm_mem_reclaim";
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_io.json")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .setMemoryMib(256)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
+        MemoryReclaimListener listener = new MemoryReclaimListener(this::executeCommand);
+        BenchmarkVmListener.create(listener).runToFinish(TAG, vm);
+
+        double mem_pre_crosvm_host_rss = (double) listener.mPreCrosvm.mHostRss / 1024.0;
+        double mem_pre_crosvm_host_pss = (double) listener.mPreCrosvm.mHostPss / 1024.0;
+        double mem_pre_crosvm_guest_rss = (double) listener.mPreCrosvm.mGuestRss / 1024.0;
+        double mem_pre_crosvm_guest_pss = (double) listener.mPreCrosvm.mGuestPss / 1024.0;
+        double mem_post_crosvm_host_rss = (double) listener.mPostCrosvm.mHostRss / 1024.0;
+        double mem_post_crosvm_host_pss = (double) listener.mPostCrosvm.mHostPss / 1024.0;
+        double mem_post_crosvm_guest_rss = (double) listener.mPostCrosvm.mGuestRss / 1024.0;
+        double mem_post_crosvm_guest_pss = (double) listener.mPostCrosvm.mGuestPss / 1024.0;
+
+        Bundle bundle = new Bundle();
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_host_rss_MB", mem_pre_crosvm_host_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_host_pss_MB", mem_pre_crosvm_host_pss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_guest_rss_MB", mem_pre_crosvm_guest_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_guest_pss_MB", mem_pre_crosvm_guest_pss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_host_rss_MB", mem_post_crosvm_host_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_host_pss_MB", mem_post_crosvm_host_pss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_guest_rss_MB", mem_post_crosvm_guest_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_guest_pss_MB", mem_post_crosvm_guest_pss);
+        mInstrumentation.sendStatus(0, bundle);
+    }
+
+    private static class MemoryReclaimListener implements BenchmarkVmListener.InnerListener {
+        MemoryReclaimListener(Function<String, String> shellExecutor) {
+            mShellExecutor = shellExecutor;
+        }
+
+        public final Function<String, String> mShellExecutor;
+
+        public CrosvmStats mPreCrosvm;
+        public CrosvmStats mPostCrosvm;
+
+        @Override
+        @SuppressWarnings("ReturnValueIgnored")
+        public void onPayloadReady(VirtualMachine vm, IBenchmarkService service)
+                throws RemoteException {
+            // Allocate 256MB of anonymous memory. This will fill all guest
+            // memory and cause swapping to start.
+            service.allocAnonMemory(256);
+            mPreCrosvm = new CrosvmStats(mShellExecutor);
+            // Send a memory trim hint to cause memory reclaim.
+            mShellExecutor.apply("am send-trim-memory " + Process.myPid() + " RUNNING_CRITICAL");
+            // Give time for the memory reclaim to do its work.
             try {
-                List<Integer> crosvmPids =
-                        ProcessUtil.getProcessMap(mShellExecutor).entrySet().stream()
-                                .filter(e -> e.getValue().contains("crosvm"))
-                                .map(e -> e.getKey())
-                                .collect(java.util.stream.Collectors.toList());
-                if (crosvmPids.size() != 1) {
-                    throw new RuntimeException(
-                            "expected to find exactly one crosvm processes, found "
-                                    + crosvmPids.size());
-                }
-
-                mCrosvmHostRss = 0;
-                mCrosvmHostPss = 0;
-                mCrosvmGuestRss = 0;
-                mCrosvmGuestPss = 0;
-                for (ProcessUtil.SMapEntry entry :
-                        ProcessUtil.getProcessSmaps(crosvmPids.get(0), mShellExecutor)) {
-                    long rss = entry.metrics.get("Rss");
-                    long pss = entry.metrics.get("Pss");
-                    if (entry.name.contains("crosvm_guest")) {
-                        mCrosvmGuestRss += rss;
-                        mCrosvmGuestPss += pss;
-                    } else {
-                        mCrosvmHostRss += rss;
-                        mCrosvmHostPss += pss;
-                    }
-                }
-            } catch (Exception e) {
-                Log.e(TAG, "Error inside onPayloadReady():" + e);
-                throw new RuntimeException(e);
+                Thread.sleep(isCuttlefish() ? 10000 : 5000);
+            } catch (InterruptedException e) {
+                Log.e(TAG, "Interrupted sleep:" + e);
+                Thread.currentThread().interrupt();
             }
+            mPostCrosvm = new CrosvmStats(mShellExecutor);
         }
     }
 
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 70ec7db..56963e6 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -77,6 +77,11 @@
         return ndk::ScopedAStatus::ok();
     }
 
+    ndk::ScopedAStatus allocAnonMemory(int64_t mb, int64_t* out) override {
+        *out = (int64_t)(long)alloc_anon_memory((long)mb);
+        return ndk::ScopedAStatus::ok();
+    }
+
     ndk::ScopedAStatus initVsockServer(int32_t port, int32_t* out) override {
         auto res = io_vsock::init_vsock_server(port);
         if (res.ok()) {
@@ -131,6 +136,17 @@
         return {file_size_mb / elapsed_seconds};
     }
 
+    void* alloc_anon_memory(long mb) {
+        long bytes = mb << 20;
+        void* p = malloc(bytes);
+        /*
+         * Heap memory is demand allocated. Dirty all pages to ensure
+         * all are allocated.
+         */
+        memset(p, 0x55, bytes);
+        return p;
+    }
+
     Result<size_t> read_meminfo_entry(const std::string& stat) {
         std::ifstream fs("/proc/meminfo");
         if (!fs.is_open()) {
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 679f6e7..35b9e61 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -31,6 +31,9 @@
 
 import android.content.Context;
 import android.os.Build;
+import android.os.ParcelFileDescriptor;
+import android.os.ParcelFileDescriptor.AutoCloseInputStream;
+import android.os.ParcelFileDescriptor.AutoCloseOutputStream;
 import android.os.SystemProperties;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineCallback;
@@ -55,11 +58,17 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 
+import java.io.BufferedReader;
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
 import java.io.RandomAccessFile;
+import java.io.Writer;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -68,6 +77,7 @@
 import java.util.OptionalLong;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 
 import co.nstant.in.cbor.CborDecoder;
 import co.nstant.in.cbor.model.Array;
@@ -113,6 +123,7 @@
                 newVmConfigBuilder()
                         .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
                         .setMemoryMib(minMemoryRequired())
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
 
@@ -126,6 +137,25 @@
     }
 
     @Test
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
+    public void createAndRunNoDebugVm() throws Exception {
+        assumeSupportedKernel();
+
+        // For most of our tests we use a debug VM so failures can be diagnosed.
+        // But we do need non-debug VMs to work, so run one.
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setMemoryMib(minMemoryRequired())
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
+
+        TestResults testResults = runVmTestService(vm);
+        assertThat(testResults.mException).isNull();
+    }
+
+    @Test
     @CddTest(
             requirements = {
                 "9.17/C-1-1",
@@ -160,6 +190,7 @@
                 newVmConfigBuilder()
                         .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
                         .setMemoryMib(minMemoryRequired())
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
         try (VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config)) {
@@ -183,6 +214,56 @@
 
     @Test
     @CddTest(requirements = {"9.17/C-1-1"})
+    public void connectVsock() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setMemoryMib(minMemoryRequired())
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_vsock", config);
+
+        AtomicReference<Exception> exception = new AtomicReference<>();
+        AtomicReference<String> response = new AtomicReference<>();
+        String request = "Look not into the abyss";
+
+        VmEventListener listener =
+                new VmEventListener() {
+                    @Override
+                    public void onPayloadReady(VirtualMachine vm) {
+                        try (vm) {
+                            ITestService testService =
+                                    ITestService.Stub.asInterface(
+                                            vm.connectToVsockServer(ITestService.SERVICE_PORT));
+                            testService.runEchoReverseServer();
+
+                            ParcelFileDescriptor pfd =
+                                    vm.connectVsock(ITestService.ECHO_REVERSE_PORT);
+                            try (InputStream input = new AutoCloseInputStream(pfd);
+                                    OutputStream output = new AutoCloseOutputStream(pfd)) {
+                                BufferedReader reader =
+                                        new BufferedReader(new InputStreamReader(input));
+                                Writer writer = new OutputStreamWriter(output);
+                                writer.write(request + "\n");
+                                writer.flush();
+                                response.set(reader.readLine());
+                            }
+                        } catch (Exception e) {
+                            exception.set(e);
+                        }
+                    }
+                };
+        listener.runToFinish(TAG, vm);
+        if (exception.get() != null) {
+            throw new RuntimeException(exception.get());
+        }
+        assertThat(response.get()).isEqualTo(new StringBuilder(request).reverse().toString());
+    }
+
+    @Test
+    @CddTest(requirements = {"9.17/C-1-1"})
     public void vmConfigUnitTests() {
         VirtualMachineConfig minimal =
                 newVmConfigBuilder().setPayloadBinaryPath("binary/path").build();
@@ -305,6 +386,7 @@
                         .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
                         .setApkPath(getContext().getPackageCodePath())
                         .setMemoryMib(minMemoryRequired())
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_explicit_apk_path", config);
@@ -322,6 +404,7 @@
                 newVmConfigBuilder()
                         .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
                         .setApkPath("relative/path/to.apk")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
                         .setMemoryMib(minMemoryRequired());
         assertThrows(IllegalArgumentException.class, () -> builder.build());
     }
@@ -347,6 +430,7 @@
                 newVmConfigBuilder()
                         .setPayloadConfigPath("assets/vm_config_extra_apk.json")
                         .setMemoryMib(minMemoryRequired())
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_extra_apk", config);
 
@@ -713,7 +797,6 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                        .setDebugLevel(DEBUG_LEVEL_NONE)
                         .build();
 
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
@@ -763,7 +846,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         String vmNameOrig = "test_vm_orig";
         String vmNameImport = "test_vm_import";
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index c0a8c0e..8a0019d 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -13,33 +13,40 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 #include <aidl/com/android/microdroid/testservice/BnTestService.h>
 #include <android-base/file.h>
 #include <android-base/properties.h>
 #include <android-base/result.h>
-#include <android/binder_auto_utils.h>
-#include <android/binder_manager.h>
+#include <android/log.h>
 #include <fcntl.h>
 #include <fsverity_digests.pb.h>
 #include <linux/vm_sockets.h>
 #include <stdint.h>
 #include <stdio.h>
-#include <sys/ioctl.h>
 #include <sys/system_properties.h>
 #include <unistd.h>
 #include <vm_main.h>
 #include <vm_payload_restricted.h>
 
 #include <string>
+#include <thread>
 
+using android::base::borrowed_fd;
 using android::base::ErrnoError;
 using android::base::Error;
 using android::base::Result;
+using android::base::unique_fd;
+
+using aidl::com::android::microdroid::testservice::BnTestService;
+using ndk::ScopedAStatus;
 
 extern void testlib_sub();
 
 namespace {
 
+constexpr char TAG[] = "testbinary";
+
 template <typename T>
 Result<T> report_test(std::string name, Result<T> result) {
     auto property = "debug.microdroid.test." + name;
@@ -48,72 +55,156 @@
         outcome << "PASS";
     } else {
         outcome << "FAIL: " << result.error();
-        // Pollute stderr with the error in case the property is truncated.
-        std::cerr << "[" << name << "] test failed: " << result.error() << "\n";
+        // Log the error in case the property is truncated.
+        std::string message = name + ": " + outcome.str();
+        __android_log_write(ANDROID_LOG_WARN, TAG, message.c_str());
     }
     __system_property_set(property.c_str(), outcome.str().c_str());
     return result;
 }
 
+Result<void> run_echo_reverse_server(borrowed_fd listening_fd) {
+    struct sockaddr_vm client_sa = {};
+    socklen_t client_sa_len = sizeof(client_sa);
+    unique_fd connect_fd{accept4(listening_fd.get(), (struct sockaddr*)&client_sa, &client_sa_len,
+                                 SOCK_CLOEXEC)};
+    if (!connect_fd.ok()) {
+        return ErrnoError() << "Failed to accept vsock connection";
+    }
+
+    unique_fd input_fd{fcntl(connect_fd, F_DUPFD_CLOEXEC, 0)};
+    if (!input_fd.ok()) {
+        return ErrnoError() << "Failed to dup";
+    }
+    FILE* input = fdopen(input_fd.release(), "r");
+    if (!input) {
+        return ErrnoError() << "Failed to fdopen";
+    }
+
+    char* line = nullptr;
+    size_t size = 0;
+    if (getline(&line, &size, input) < 0) {
+        return ErrnoError() << "Failed to read";
+    }
+
+    if (fclose(input) != 0) {
+        return ErrnoError() << "Failed to fclose";
+    }
+
+    std::string_view original = line;
+    if (!original.empty() && original.back() == '\n') {
+        original = original.substr(0, original.size() - 1);
+    }
+
+    std::string reversed(original.rbegin(), original.rend());
+
+    if (write(connect_fd, reversed.data(), reversed.size()) < 0) {
+        return ErrnoError() << "Failed to write";
+    }
+
+    return {};
+}
+
+Result<void> start_echo_reverse_server() {
+    unique_fd server_fd{TEMP_FAILURE_RETRY(socket(AF_VSOCK, SOCK_STREAM | SOCK_CLOEXEC, 0))};
+    if (!server_fd.ok()) {
+        return ErrnoError() << "Failed to create vsock socket";
+    }
+    struct sockaddr_vm server_sa = (struct sockaddr_vm){
+            .svm_family = AF_VSOCK,
+            .svm_port = BnTestService::ECHO_REVERSE_PORT,
+            .svm_cid = VMADDR_CID_ANY,
+    };
+    int ret = TEMP_FAILURE_RETRY(bind(server_fd, (struct sockaddr*)&server_sa, sizeof(server_sa)));
+    if (ret < 0) {
+        return ErrnoError() << "Failed to bind vsock socket";
+    }
+    ret = TEMP_FAILURE_RETRY(listen(server_fd, /*backlog=*/1));
+    if (ret < 0) {
+        return ErrnoError() << "Failed to listen";
+    }
+
+    std::thread accept_thread{[listening_fd = std::move(server_fd)] {
+        auto result = run_echo_reverse_server(listening_fd);
+        if (!result.ok()) {
+            __android_log_write(ANDROID_LOG_ERROR, TAG, result.error().message().c_str());
+            // Make sure the VM exits so the test will fail solidly
+            exit(1);
+        }
+    }};
+    accept_thread.detach();
+
+    return {};
+}
+
 Result<void> start_test_service() {
-    class TestService : public aidl::com::android::microdroid::testservice::BnTestService {
-        ndk::ScopedAStatus addInteger(int32_t a, int32_t b, int32_t* out) override {
+    class TestService : public BnTestService {
+        ScopedAStatus addInteger(int32_t a, int32_t b, int32_t* out) override {
             *out = a + b;
-            return ndk::ScopedAStatus::ok();
+            return ScopedAStatus::ok();
         }
 
-        ndk::ScopedAStatus readProperty(const std::string& prop, std::string* out) override {
+        ScopedAStatus readProperty(const std::string& prop, std::string* out) override {
             *out = android::base::GetProperty(prop, "");
             if (out->empty()) {
                 std::string msg = "cannot find property " + prop;
-                return ndk::ScopedAStatus::fromExceptionCodeWithMessage(EX_SERVICE_SPECIFIC,
-                                                                        msg.c_str());
+                return ScopedAStatus::fromExceptionCodeWithMessage(EX_SERVICE_SPECIFIC,
+                                                                   msg.c_str());
             }
 
-            return ndk::ScopedAStatus::ok();
+            return ScopedAStatus::ok();
         }
 
-        ndk::ScopedAStatus insecurelyExposeVmInstanceSecret(std::vector<uint8_t>* out) override {
+        ScopedAStatus insecurelyExposeVmInstanceSecret(std::vector<uint8_t>* out) override {
             const uint8_t identifier[] = {1, 2, 3, 4};
             out->resize(32);
             AVmPayload_getVmInstanceSecret(identifier, sizeof(identifier), out->data(),
                                            out->size());
-            return ndk::ScopedAStatus::ok();
+            return ScopedAStatus::ok();
         }
 
-        ndk::ScopedAStatus insecurelyExposeAttestationCdi(std::vector<uint8_t>* out) override {
+        ScopedAStatus insecurelyExposeAttestationCdi(std::vector<uint8_t>* out) override {
             size_t cdi_size = AVmPayload_getDiceAttestationCdi(nullptr, 0);
             out->resize(cdi_size);
             AVmPayload_getDiceAttestationCdi(out->data(), out->size());
-            return ndk::ScopedAStatus::ok();
+            return ScopedAStatus::ok();
         }
 
-        ndk::ScopedAStatus getBcc(std::vector<uint8_t>* out) override {
+        ScopedAStatus getBcc(std::vector<uint8_t>* out) override {
             size_t bcc_size = AVmPayload_getDiceAttestationChain(nullptr, 0);
             out->resize(bcc_size);
             AVmPayload_getDiceAttestationChain(out->data(), out->size());
-            return ndk::ScopedAStatus::ok();
+            return ScopedAStatus::ok();
         }
 
-        ndk::ScopedAStatus getApkContentsPath(std::string* out) override {
+        ScopedAStatus getApkContentsPath(std::string* out) override {
             const char* path_c = AVmPayload_getApkContentsPath();
             if (path_c == nullptr) {
-                return ndk::ScopedAStatus::
+                return ScopedAStatus::
                         fromServiceSpecificErrorWithMessage(0, "Failed to get APK contents path");
             }
-            std::string path(path_c);
-            *out = path;
-            return ndk::ScopedAStatus::ok();
+            *out = path_c;
+            return ScopedAStatus::ok();
         }
 
-        ndk::ScopedAStatus getEncryptedStoragePath(std::string* out) override {
+        ScopedAStatus getEncryptedStoragePath(std::string* out) override {
             const char* path_c = AVmPayload_getEncryptedStoragePath();
             if (path_c == nullptr) {
                 out->clear();
             } else {
                 *out = path_c;
             }
-            return ndk::ScopedAStatus::ok();
+            return ScopedAStatus::ok();
+        }
+
+        virtual ::ScopedAStatus runEchoReverseServer() override {
+            auto result = start_echo_reverse_server();
+            if (result.ok()) {
+                return ScopedAStatus::ok();
+            } else {
+                std::string message = result.error().message();
+                return ScopedAStatus::fromServiceSpecificErrorWithMessage(-1, message.c_str());
+            }
         }
     };
     auto testService = ndk::SharedRefBase::make<TestService>();
@@ -143,14 +234,10 @@
 } // Anonymous namespace
 
 extern "C" int AVmPayload_main() {
-    // disable buffering to communicate seamlessly
-    setvbuf(stdin, nullptr, _IONBF, 0);
-    setvbuf(stdout, nullptr, _IONBF, 0);
-    setvbuf(stderr, nullptr, _IONBF, 0);
+    __android_log_write(ANDROID_LOG_INFO, TAG, "Hello Microdroid");
 
-    printf("Hello Microdroid");
+    // Make sure we can call into other shared libraries.
     testlib_sub();
-    printf("\n");
 
     // Extra apks may be missing; this is not a fatal error
     report_test("extra_apk", verify_apk());
@@ -160,7 +247,7 @@
     if (auto res = start_test_service(); res.ok()) {
         return 0;
     } else {
-        std::cerr << "starting service failed: " << res.error() << "\n";
+        __android_log_write(ANDROID_LOG_ERROR, TAG, res.error().message().c_str());
         return 1;
     }
 }
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index b767013..da56f76 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -51,6 +51,7 @@
         "libshared_child",
         "libstatslog_virtualization_rust",
         "libtombstoned_client_rust",
+        "libvm_control",
         "libvmconfig",
         "libzip",
         "libvsock",
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
index d9d9a61..d76b586 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
@@ -16,6 +16,7 @@
 package android.system.virtualizationservice;
 
 import android.system.virtualizationservice.IVirtualMachineCallback;
+import android.system.virtualizationservice.MemoryTrimLevel;
 import android.system.virtualizationservice.VirtualMachineState;
 
 interface IVirtualMachine {
@@ -41,6 +42,9 @@
      */
     void stop();
 
+    /** Communicate app low-memory notifications to the VM. */
+    void onTrimMemory(MemoryTrimLevel level);
+
     /** Open a vsock connection to the CID of the VM on the given port. */
     ParcelFileDescriptor connectVsock(int port);
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/MemoryTrimLevel.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/MemoryTrimLevel.aidl
new file mode 100644
index 0000000..9ed9e99
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/MemoryTrimLevel.aidl
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.system.virtualizationservice;
+
+/**
+ * Memory trim levels propagated from the app to the VM.
+ */
+@Backing(type="int")
+enum MemoryTrimLevel {
+    /* Same meaning as in ComponentCallbacks2 */
+    TRIM_MEMORY_RUNNING_CRITICAL = 0,
+    TRIM_MEMORY_RUNNING_LOW = 1,
+    TRIM_MEMORY_RUNNING_MODERATE = 2,
+}
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 040c0d8..7d24a32 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -27,6 +27,7 @@
     IVirtualMachine::{BnVirtualMachine, IVirtualMachine},
     IVirtualMachineCallback::IVirtualMachineCallback,
     IVirtualizationService::IVirtualizationService,
+    MemoryTrimLevel::MemoryTrimLevel,
     Partition::Partition,
     PartitionType::PartitionType,
     VirtualMachineAppConfig::{Payload::Payload, VirtualMachineAppConfig},
@@ -961,6 +962,13 @@
         })
     }
 
+    fn onTrimMemory(&self, level: MemoryTrimLevel) -> binder::Result<()> {
+        self.instance.trim_memory(level).map_err(|e| {
+            error!("Error trimming VM with CID {}: {:?}", self.instance.cid, e);
+            Status::new_service_specific_error_str(-1, Some(e.to_string()))
+        })
+    }
+
     fn connectVsock(&self, port: i32) -> binder::Result<ParcelFileDescriptor> {
         if !matches!(&*self.instance.vm_state.lock().unwrap(), VmState::Running { .. }) {
             return Err(Status::new_service_specific_error_str(-1, Some("VM is not running")));
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 00c33ef..0fdc293 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -39,7 +39,10 @@
 use std::sync::{Arc, Condvar, Mutex};
 use std::time::{Duration, SystemTime};
 use std::thread;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::DeathReason::DeathReason;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+    DeathReason::DeathReason,
+    MemoryTrimLevel::MemoryTrimLevel,
+};
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IGlobalVmContext::IGlobalVmContext;
 use binder::Strong;
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
@@ -48,6 +51,7 @@
 
 /// external/crosvm
 use base::UnixSeqpacketListener;
+use vm_control::{BalloonControlCommand, VmRequest, VmResponse};
 
 const CROSVM_PATH: &str = "/apex/com.android.virt/bin/crosvm";
 
@@ -172,7 +176,7 @@
 
             // If this fails and returns an error, `self` will be left in the `Failed` state.
             let child =
-                Arc::new(run_vm(config, &instance.temporary_directory, failure_pipe_write)?);
+                Arc::new(run_vm(config, &instance.crosvm_control_socket_path, failure_pipe_write)?);
 
             let instance_monitor_status = instance.clone();
             let child_monitor_status = child.clone();
@@ -229,6 +233,8 @@
     vm_context: VmContext,
     /// The CID assigned to the VM for vsock communication.
     pub cid: Cid,
+    /// Path to crosvm control socket
+    crosvm_control_socket_path: PathBuf,
     /// The name of the VM.
     pub name: String,
     /// Whether the VM is a protected VM.
@@ -286,6 +292,7 @@
             vm_state: Mutex::new(VmState::NotStarted { config }),
             vm_context,
             cid,
+            crosvm_control_socket_path: temporary_directory.join("crosvm.sock"),
             name,
             protected,
             temporary_directory,
@@ -465,6 +472,46 @@
         }
     }
 
+    /// Responds to memory-trimming notifications by inflating the virtio
+    /// balloon to reclaim guest memory.
+    pub fn trim_memory(&self, level: MemoryTrimLevel) -> Result<(), Error> {
+        let request = VmRequest::BalloonCommand(BalloonControlCommand::Stats {});
+        match vm_control::client::handle_request(&request, &self.crosvm_control_socket_path) {
+            Ok(VmResponse::BalloonStats { stats, balloon_actual: _ }) => {
+                if let Some(total_memory) = stats.total_memory {
+                    // Reclaim up to 50% of total memory assuming worst case
+                    // most memory is anonymous and must be swapped to zram
+                    // with an approximate 2:1 compression ratio.
+                    let pct = match level {
+                        MemoryTrimLevel::TRIM_MEMORY_RUNNING_CRITICAL => 50,
+                        MemoryTrimLevel::TRIM_MEMORY_RUNNING_LOW => 30,
+                        MemoryTrimLevel::TRIM_MEMORY_RUNNING_MODERATE => 10,
+                        _ => bail!("Invalid memory trim level {:?}", level),
+                    };
+                    let command =
+                        BalloonControlCommand::Adjust { num_bytes: total_memory * pct / 100 };
+                    if let Err(e) = vm_control::client::handle_request(
+                        &VmRequest::BalloonCommand(command),
+                        &self.crosvm_control_socket_path,
+                    ) {
+                        bail!("Error sending balloon adjustment: {:?}", e);
+                    }
+                }
+            }
+            Ok(VmResponse::Err(e)) => {
+                // ENOTSUP is returned when the balloon protocol is not initialised. This
+                // can occur for numerous reasons: Guest is still booting, guest doesn't
+                // support ballooning, host doesn't support ballooning. We don't log or
+                // raise an error in this case: trim is just a hint and we can ignore it.
+                if e.errno() != libc::ENOTSUP {
+                    bail!("Errno return when requesting balloon stats: {}", e.errno())
+                }
+            }
+            e => bail!("Error requesting balloon stats: {:?}", e),
+        }
+        Ok(())
+    }
+
     /// Checks if ramdump has been created. If so, send a notification to the user with the handle
     /// to read the ramdump.
     fn handle_ramdump(&self) -> Result<(), Error> {
@@ -604,7 +651,7 @@
 /// Starts an instance of `crosvm` to manage a new VM.
 fn run_vm(
     config: CrosvmConfig,
-    temporary_directory: &Path,
+    crosvm_control_socket_path: &Path,
     failure_pipe_write: File,
 ) -> Result<SharedChild, Error> {
     validate_config(&config)?;
@@ -710,9 +757,8 @@
         command.arg(add_preserved_fd(&mut preserved_fds, kernel));
     }
 
-    let control_server_socket =
-        UnixSeqpacketListener::bind(temporary_directory.join("crosvm.sock"))
-            .context("failed to create control server")?;
+    let control_server_socket = UnixSeqpacketListener::bind(crosvm_control_socket_path)
+        .context("failed to create control server")?;
     command.arg("--socket").arg(add_preserved_fd(&mut preserved_fds, &control_server_socket));
 
     debug!("Preserving FDs {:?}", preserved_fds);
diff --git a/vm_payload/src/api.rs b/vm_payload/src/api.rs
index a79c0bb..28b440e 100644
--- a/vm_payload/src/api.rs
+++ b/vm_payload/src/api.rs
@@ -18,7 +18,7 @@
 #![warn(unsafe_op_in_unsafe_fn)]
 
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
+    ENCRYPTEDSTORE_MOUNTPOINT, IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
 use anyhow::{ensure, bail, Context, Result};
 use binder::{Strong, unstable_api::{AIBinder, new_spibinder}};
 use lazy_static::lazy_static;
@@ -28,6 +28,7 @@
 use std::ffi::CString;
 use std::fmt::Debug;
 use std::os::raw::{c_char, c_void};
+use std::path::Path;
 use std::ptr;
 use std::sync::{Mutex, atomic::{AtomicBool, Ordering}};
 
@@ -35,6 +36,8 @@
     static ref VM_APK_CONTENTS_PATH_C: CString =
         CString::new(VM_APK_CONTENTS_PATH).expect("CString::new failed");
     static ref PAYLOAD_CONNECTION: Mutex<Option<Strong<dyn IVmPayloadService>>> = Mutex::default();
+    static ref VM_ENCRYPTED_STORAGE_PATH_C: CString =
+        CString::new(ENCRYPTEDSTORE_MOUNTPOINT).expect("CString::new failed");
 }
 
 static ALREADY_NOTIFIED: AtomicBool = AtomicBool::new(false);
@@ -249,12 +252,15 @@
 /// Gets the path to the APK contents.
 #[no_mangle]
 pub extern "C" fn AVmPayload_getApkContentsPath() -> *const c_char {
-    (*VM_APK_CONTENTS_PATH_C).as_ptr()
+    VM_APK_CONTENTS_PATH_C.as_ptr()
 }
 
 /// Gets the path to the VM's encrypted storage.
 #[no_mangle]
 pub extern "C" fn AVmPayload_getEncryptedStoragePath() -> *const c_char {
-    // TODO(b/254454578): Return a real path if storage is present
-    ptr::null()
+    if Path::new(ENCRYPTEDSTORE_MOUNTPOINT).exists() {
+        VM_ENCRYPTED_STORAGE_PATH_C.as_ptr()
+    } else {
+        ptr::null()
+    }
 }