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/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);
+    }
+}