Add VirtualizationModuleFrameworkInitializer

VirtualizationModuleFrameworkInitializer.registerServiceWrappers() is
called in SystemServiceRegistry during zygote initialization to register
Manager classes inside the com.android.virt APEX with the
Context.getSystemService() lookup API.

For the devices without AVF support the
Context.getSystemService(VirtualMachineManager.class) will return null.
A CTS test is added to verify this behaviour.

Bug: 249093790
Test: presubmit
Change-Id: Ib8a94643e5f7ccbfa55dd659d93d02d146af6443
Merged-In: Ib8a94643e5f7ccbfa55dd659d93d02d146af6443
(cherry picked from commit 63594352ae1e5fed298b7d4b35281a393bcfe5e2)
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 8376c1f..dd81738 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -30,6 +30,11 @@
       "name": "AVFHostTestCases"
     }
   ],
+  "postsubmit": [
+    {
+      "name": "CtsMicrodroidDisabledTestCases"
+    }
+  ],
   "imports": [
     {
       "path": "packages/modules/Virtualization/apkdmverity"
diff --git a/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index 8e870ea..77f2ee7 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -256,7 +256,8 @@
                     builder.setDebugLevel(VirtualMachineConfig.DEBUG_LEVEL_FULL);
                 }
                 VirtualMachineConfig config = builder.build();
-                VirtualMachineManager vmm = VirtualMachineManager.getInstance(getApplication());
+                VirtualMachineManager vmm =
+                        getApplication().getSystemService(VirtualMachineManager.class);
                 mVirtualMachine = vmm.getOrCreate(VM_NAME, config);
                 try {
                     mVirtualMachine.setConfig(config);
diff --git a/javalib/api/module-lib-current.txt b/javalib/api/module-lib-current.txt
index d802177..4d59764 100644
--- a/javalib/api/module-lib-current.txt
+++ b/javalib/api/module-lib-current.txt
@@ -1 +1,9 @@
 // Signature format: 2.0
+package android.system.virtualmachine {
+
+  public class VirtualizationFrameworkInitializer {
+    method public static void registerServiceWrappers();
+  }
+
+}
+
diff --git a/javalib/api/system-current.txt b/javalib/api/system-current.txt
index ea2d23e..850bffc 100644
--- a/javalib/api/system-current.txt
+++ b/javalib/api/system-current.txt
@@ -98,7 +98,6 @@
     method public void delete(@NonNull String) throws android.system.virtualmachine.VirtualMachineException;
     method @Nullable public android.system.virtualmachine.VirtualMachine get(@NonNull String) throws android.system.virtualmachine.VirtualMachineException;
     method public int getCapabilities();
-    method @NonNull public static android.system.virtualmachine.VirtualMachineManager getInstance(@NonNull android.content.Context);
     method @NonNull public android.system.virtualmachine.VirtualMachine getOrCreate(@NonNull String, @NonNull android.system.virtualmachine.VirtualMachineConfig) throws android.system.virtualmachine.VirtualMachineException;
     method @NonNull public android.system.virtualmachine.VirtualMachine importFromDescriptor(@NonNull String, @NonNull android.system.virtualmachine.VirtualMachineDescriptor) throws android.system.virtualmachine.VirtualMachineException;
     field public static final int CAPABILITY_NON_PROTECTED_VM = 2; // 0x2
diff --git a/javalib/jarjar-rules.txt b/javalib/jarjar-rules.txt
index dd8ad2d..726f9aa 100644
--- a/javalib/jarjar-rules.txt
+++ b/javalib/jarjar-rules.txt
@@ -1,9 +1,10 @@
 # Rules for the android.system.virtualmachine java_sdk_library.
 
-# This is the root of the API, everything we care about should be
-# reachable from here.
-# (This gets rid of all the android.sysprop classes we don't use.)
+# Keep the API surface, most of it is accessible from VirtualMachineManager
 keep android.system.virtualmachine.VirtualMachineManager
+# VirtualizationModuleFrameworkInitializer is not accessible from
+# VirtualMachineManager, we need to explicitly keep it.
+keep android.system.virtualmachine.VirtualizationFrameworkInitializer
 
 # We statically link PlatformProperties, rename to avoid clashes.
 rule android.sysprop.** com.android.system.virtualmachine.sysprop.@1
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index a520ab4..c179498 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -21,10 +21,11 @@
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.RequiresFeature;
 import android.annotation.RequiresPermission;
-import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
 import android.content.Context;
+import android.content.pm.PackageManager;
 import android.sysprop.HypervisorProperties;
 import android.util.ArrayMap;
 
@@ -34,7 +35,6 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.ref.WeakReference;
 import java.util.Map;
-import java.util.WeakHashMap;
 
 /**
  * Manages {@link VirtualMachine virtual machine} instances created by an app. Each instance is
@@ -46,9 +46,13 @@
  *
  * <p>The app can then start, stop and otherwise interact with the VM.
  *
+ * <p>An instance of {@link VirtualMachineManager} can be obtained by calling {@link
+ * Context#getSystemService(Class)}.
+ *
  * @hide
  */
 @SystemApi
+@RequiresFeature(PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK)
 public class VirtualMachineManager {
     /**
      * A lock used to synchronize the creation of virtual machines. It protects {@link #mVmsByName},
@@ -59,14 +63,11 @@
 
     @NonNull private final Context mContext;
 
-    private VirtualMachineManager(@NonNull Context context) {
-        mContext = context;
+    /** @hide */
+    public VirtualMachineManager(@NonNull Context context) {
+        mContext = requireNonNull(context);
     }
 
-    @GuardedBy("sInstances")
-    private static final Map<Context, WeakReference<VirtualMachineManager>> sInstances =
-            new WeakHashMap<>();
-
     @GuardedBy("sCreateLock")
     private final Map<String, WeakReference<VirtualMachine>> mVmsByName = new ArrayMap<>();
 
@@ -93,27 +94,6 @@
     public static final int CAPABILITY_NON_PROTECTED_VM = 2;
 
     /**
-     * Returns the per-context instance.
-     *
-     * @hide
-     */
-    @SystemApi
-    @NonNull
-    @SuppressLint("ManagerLookup") // TODO(b/249093790): remove
-    public static VirtualMachineManager getInstance(@NonNull Context context) {
-        requireNonNull(context, "context must not be null");
-        synchronized (sInstances) {
-            VirtualMachineManager vmm =
-                    sInstances.containsKey(context) ? sInstances.get(context).get() : null;
-            if (vmm == null) {
-                vmm = new VirtualMachineManager(context);
-                sInstances.put(context, new WeakReference<>(vmm));
-            }
-            return vmm;
-        }
-    }
-
-    /**
      * Returns a set of flags indicating what this implementation of virtualization is capable of.
      *
      * @see #CAPABILITY_PROTECTED_VM
diff --git a/javalib/src/android/system/virtualmachine/VirtualizationFrameworkInitializer.java b/javalib/src/android/system/virtualmachine/VirtualizationFrameworkInitializer.java
new file mode 100644
index 0000000..30ac425
--- /dev/null
+++ b/javalib/src/android/system/virtualmachine/VirtualizationFrameworkInitializer.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 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.virtualmachine;
+
+import static android.content.pm.PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK;
+
+import android.annotation.SystemApi;
+import android.app.SystemServiceRegistry;
+import android.content.Context;
+
+/**
+ * Holds initialization code for virtualization module
+ *
+ * @hide
+ */
+@SystemApi(client = SystemApi.Client.MODULE_LIBRARIES)
+public class VirtualizationFrameworkInitializer {
+
+    private VirtualizationFrameworkInitializer() {}
+
+    /**
+     * Called by the static initializer in the {@link SystemServiceRegistry}, and registers {@link
+     * VirtualMachineManager} to the {@link Context}. so that it's accessible from {@link
+     * Context#getSystemService(String)}.
+     */
+    public static void registerServiceWrappers() {
+        // Note: it's important that the getPackageManager().hasSystemFeature() check is executed
+        // in the lambda, and not directly in the registerServiceWrappers method. The
+        // registerServiceWrappers is called during Zygote static initialization, and at that
+        // point the PackageManager is not available yet.
+        //
+        // On the other hand, the lambda is executed after the app calls Context.getSystemService
+        // (VirtualMachineManager.class), at which point the PackageManager is available. The
+        // result of the lambda is cached on per-context basis.
+        SystemServiceRegistry.registerContextAwareService(
+                Context.VIRTUALIZATION_SERVICE,
+                VirtualMachineManager.class,
+                ctx ->
+                        ctx.getPackageManager().hasSystemFeature(FEATURE_VIRTUALIZATION_FRAMEWORK)
+                                ? new VirtualMachineManager(ctx)
+                                : null);
+    }
+}
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 24e2049..b043c72 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -123,7 +123,7 @@
                 .that(ctx.getPackageManager().hasSystemFeature(FEATURE_VIRTUALIZATION_FRAMEWORK))
                 .isTrue();
 
-        mInner = new Inner(ctx, protectedVm, VirtualMachineManager.getInstance(ctx));
+        mInner = new Inner(ctx, protectedVm, ctx.getSystemService(VirtualMachineManager.class));
 
         int capabilities = mInner.getVirtualMachineManager().getCapabilities();
         if (protectedVm) {
diff --git a/tests/no_avf/Android.bp b/tests/no_avf/Android.bp
new file mode 100644
index 0000000..fd0d5e2
--- /dev/null
+++ b/tests/no_avf/Android.bp
@@ -0,0 +1,21 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_test {
+    name: "CtsMicrodroidDisabledTestCases",
+    test_suites: [
+        "cts",
+        "general-tests",
+    ],
+    srcs: ["src/**/*.java"],
+    static_libs: [
+        "androidx.test.runner",
+        "androidx.test.ext.junit",
+        "compatibility-common-util-devicesidelib",
+        "truth-prebuilt",
+    ],
+    sdk_version: "test_current",
+    compile_multilib: "both",
+    min_sdk_version: "UpsideDownCake",
+}
diff --git a/tests/no_avf/AndroidManifest.xml b/tests/no_avf/AndroidManifest.xml
new file mode 100644
index 0000000..4a1304e
--- /dev/null
+++ b/tests/no_avf/AndroidManifest.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 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.
+-->
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+      package="com.android.nomicrodroid.test">
+    <application />
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.nomicrodroid.test"
+        android:label="CTS test for devices without Microdroid support" />
+</manifest>
diff --git a/tests/no_avf/AndroidTest.xml b/tests/no_avf/AndroidTest.xml
new file mode 100644
index 0000000..1e93887
--- /dev/null
+++ b/tests/no_avf/AndroidTest.xml
@@ -0,0 +1,29 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 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.
+-->
+<configuration description="Runs Microdroid device-side tests.">
+    <option name="test-suite-tag" value="cts" />
+    <option name="config-descriptor:metadata" key="component" value="security" />
+    <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
+    <option name="config-descriptor:metadata" key="parameter" value="not_multi_abi" />
+    <option name="config-descriptor:metadata" key="parameter" value="secondary_user" />
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="test-file-name" value="CtsMicrodroidDisabledTestCases.apk" />
+    </target_preparer>
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.nomicrodroid.test" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+    </test>
+</configuration>
diff --git a/tests/no_avf/README.md b/tests/no_avf/README.md
new file mode 100644
index 0000000..b96dc97
--- /dev/null
+++ b/tests/no_avf/README.md
@@ -0,0 +1 @@
+CTS tests for devices that don't support AVF.
\ No newline at end of file
diff --git a/tests/no_avf/src/com/android/nomicrodroid/test/NoMicrodroidTest.java b/tests/no_avf/src/com/android/nomicrodroid/test/NoMicrodroidTest.java
new file mode 100644
index 0000000..0982e35
--- /dev/null
+++ b/tests/no_avf/src/com/android/nomicrodroid/test/NoMicrodroidTest.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 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 com.android.nomicrodroid.test;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.TruthJUnit.assume;
+
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.system.virtualmachine.VirtualMachineManager;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.compatibility.common.util.CddTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests to validate that devices without support for AVF (Android Virtualization Framework) are set
+ * up correctly.
+ */
+@RunWith(JUnit4.class)
+public class NoMicrodroidTest {
+
+    @Before
+    public void setUp() {
+        final PackageManager pm =
+                InstrumentationRegistry.getInstrumentation().getTargetContext().getPackageManager();
+        assume().withMessage("Device supports AVF")
+                .that(pm.hasSystemFeature(PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK))
+                .isFalse();
+    }
+
+    @CddTest(requirements = {"9.17/C-1-1"})
+    @Test
+    public void testVirtualMachineManagerLookupReturnsNull() {
+        final Context ctx = InstrumentationRegistry.getInstrumentation().getTargetContext();
+        assertThat(ctx.getSystemService(VirtualMachineManager.class)).isNull();
+    }
+}