Add a runtime check to ensure that system server jars are prefetched.

We prefetch standalone system server jars in ZygoteInit based on the
STANDALONE_SYSTEMSERVER_JARS environment variable, so that they can take
the advantage of AOT compilation. This CL adds a check to disallow jars
that are not prefetched, which reminds developers to make appropriate
changes so that their jars will be in the environment variable.

Bug: 203198541
Test: 1. Build a system image.
  2. The device boots.
Test: 1. Remove an entry from PRODUCT_APEX_STANDALONE_SYSTEM_SERVER_JARS
  2. Build a system image.
  3. The device does not boot and encounters the following error:
     java.lang.RuntimeException: Creating a ClassLoader from /apex/com.android.wifi/javalib/service-wifi.jar is not allowed. Please make sure that the jar is listed in `PRODUCT_APEX_STANDALONE_SYSTEM_SERVER_JARS` in the Makefile and added as a `standalone_contents` of a `systemserverclasspath_fragment` in `Android.bp`.
Change-Id: I275d75ac37194a4d8fd491529b7cdb697dc04e37

core/java/com/android/internal/os/SystemServerClassLoaderFactory.java

@@ -29,22 +29,66 @@ public final class SystemServerClassLoaderFactory {
     private static final ArrayMap<String, PathClassLoader> sLoadedPaths = new ArrayMap<>();
 
     /**
-     * Creates and caches a ClassLoader for the jar at the given path, or returns a cached
+     * Creates and caches a ClassLoader for the jar at the given path.
-     * ClassLoader if it exists.
+     *
+     * This method should only be called by ZygoteInit to prefetch jars. For other users, use
+     * {@link getOrCreateClassLoader} instead.
      *
      * The parent class loader should always be the system server class loader. Changing it has
      * implications that require discussion with the mainline team.
      *
      * @hide for internal use only
      */
-    public static PathClassLoader getOrCreateClassLoader(String path, ClassLoader parent) {
+    /* package */ static PathClassLoader createClassLoader(String path, ClassLoader parent) {
-        PathClassLoader pathClassLoader = sLoadedPaths.get(path);
+        if (sLoadedPaths.containsKey(path)) {
-        if (pathClassLoader == null) {
+            throw new IllegalStateException("A ClassLoader for " + path + " already exists");
-            pathClassLoader = (PathClassLoader) ClassLoaderFactory.createClassLoader(
-                    path, /*librarySearchPath=*/null, /*libraryPermittedPath=*/null, parent,
-                    Build.VERSION.SDK_INT, /*isNamespaceShared=*/true , /*classLoaderName=*/null);
-            sLoadedPaths.put(path, pathClassLoader);
         }
+        PathClassLoader pathClassLoader = (PathClassLoader) ClassLoaderFactory.createClassLoader(
+                path, /*librarySearchPath=*/null, /*libraryPermittedPath=*/null, parent,
+                Build.VERSION.SDK_INT, /*isNamespaceShared=*/true , /*classLoaderName=*/null);
+        sLoadedPaths.put(path, pathClassLoader);
         return pathClassLoader;
     }
+
+    /**
+     * Returns a cached ClassLoader to be used at runtime for the jar at the given path. Or, creates
+     * one if it is not prefetched and is allowed to be created at runtime.
+     *
+     * The parent class loader should always be the system server class loader. Changing it has
+     * implications that require discussion with the mainline team.
+     *
+     * @hide for internal use only
+     */
+    public static PathClassLoader getOrCreateClassLoader(
+            String path, ClassLoader parent, boolean isTestOnly) {
+        PathClassLoader pathClassLoader = sLoadedPaths.get(path);
+        if (pathClassLoader != null) {
+            return pathClassLoader;
+        }
+        if (!allowClassLoaderCreation(path, isTestOnly)) {
+            throw new RuntimeException("Creating a ClassLoader from " + path + " is not allowed. "
+                    + "Please make sure that the jar is listed in "
+                    + "`PRODUCT_APEX_STANDALONE_SYSTEM_SERVER_JARS` in the Makefile and added as a "
+                    + "`standalone_contents` of a `systemserverclasspath_fragment` in "
+                    + "`Android.bp`.");
+        }
+        return createClassLoader(path, parent);
+    }
+
+    /**
+     * Returns whether a class loader for the jar is allowed to be created at runtime.
+     */
+    private static boolean allowClassLoaderCreation(String path, boolean isTestOnly) {
+        // Currently, we only enforce prefetching for APEX jars.
+        if (!path.startsWith("/apex/")) {
+            return true;
+        }
+        // APEXes for testing only are okay to ignore.
+        if (isTestOnly) {
+            return true;
+        }
+        return false;
+    }
+
+
 }

core/java/com/android/internal/os/ZygoteInit.java

@@ -586,7 +586,7 @@ public class ZygoteInit {
         }
         for (String jar : envStr.split(":")) {
             try {
-                SystemServerClassLoaderFactory.getOrCreateClassLoader(
+                SystemServerClassLoaderFactory.createClassLoader(
                         jar, getOrCreateSystemServerClassLoader());
             } catch (Error e) {
                 // We don't want the process to crash for this error because prefetching is just an

services/core/java/com/android/server/SystemServiceManager.java

@@ -21,6 +21,8 @@ import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.app.ActivityManager;
 import android.content.Context;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageInfo;
 import android.content.pm.UserInfo;
 import android.os.Environment;
 import android.os.SystemClock;
@@ -38,6 +40,7 @@ import com.android.internal.util.Preconditions;
 import com.android.server.SystemService.TargetUser;
 import com.android.server.SystemService.UserCompletedEventType;
 import com.android.server.am.EventLogTags;
+import com.android.server.pm.ApexManager;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.utils.TimingsTraceAndSlog;
 
@@ -47,6 +50,8 @@ import java.io.File;
 import java.io.PrintWriter;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -147,12 +152,31 @@ public final class SystemServiceManager implements Dumpable {
      * @return The service instance.
      */
     public SystemService startServiceFromJar(String className, String path) {
-        PathClassLoader pathClassLoader = SystemServerClassLoaderFactory.getOrCreateClassLoader(
+        PathClassLoader pathClassLoader =
-                path, this.getClass().getClassLoader());
+                SystemServerClassLoaderFactory.getOrCreateClassLoader(
+                        path, this.getClass().getClassLoader(), isJarInTestApex(path));
         final Class<SystemService> serviceClass = loadClassFromLoader(className, pathClassLoader);
         return startService(serviceClass);
     }
 
+    /**
+     * Returns true if the jar is in a test APEX.
+     */
+    private static boolean isJarInTestApex(String pathStr) {
+        Path path = Paths.get(pathStr);
+        if (path.getNameCount() >= 2 && path.getName(0).toString().equals("apex")) {
+            String apexModuleName = path.getName(1).toString();
+            ApexManager apexManager = ApexManager.getInstance();
+            String packageName = apexManager.getActivePackageNameForApexModuleName(apexModuleName);
+            PackageInfo packageInfo = apexManager.getPackageInfo(
+                    packageName, ApexManager.MATCH_ACTIVE_PACKAGE);
+            if (packageInfo != null) {
+                return (packageInfo.applicationInfo.flags & ApplicationInfo.FLAG_TEST_ONLY) != 0;
+            }
+        }
+        return false;
+    }
+
     /*
      * Loads and initializes a class from the given classLoader. Returns the class.
      */

services/core/java/com/android/server/pm/ApexManager.java

@@ -347,6 +347,13 @@ public abstract class ApexManager {
     @Nullable
     public abstract String getApexModuleNameForPackageName(String apexPackageName);
 
+    /**
+     * Returns the package name of the active APEX whose name is {@code apexModuleName}. If not
+     * found, returns {@code null}.
+     */
+    @Nullable
+    public abstract String getActivePackageNameForApexModuleName(String apexModuleName);
+
     /**
      * Copies the CE apex data directory for the given {@code userId} to a backup location, for use
      * in case of rollback.
@@ -484,6 +491,12 @@ public abstract class ApexManager {
         @GuardedBy("mLock")
         private ArrayMap<String, String> mPackageNameToApexModuleName;
 
+        /**
+         * Reverse mapping of {@link #mPackageNameToApexModuleName}, for active packages only.
+         */
+        @GuardedBy("mLock")
+        private ArrayMap<String, String> mApexModuleNameToActivePackageName;
+
         /**
          * Whether an APEX package is active or not.
          *
@@ -552,6 +565,7 @@ public abstract class ApexManager {
             try {
                 mAllPackagesCache = new ArrayList<>();
                 mPackageNameToApexModuleName = new ArrayMap<>();
+                mApexModuleNameToActivePackageName = new ArrayMap<>();
                 allPkgs = waitForApexService().getAllPackages();
             } catch (RemoteException re) {
                 Slog.e(TAG, "Unable to retrieve packages from apexservice: " + re.toString());
@@ -634,6 +648,13 @@ public abstract class ApexManager {
                                             + packageInfo.packageName);
                         }
                         activePackagesSet.add(packageInfo.packageName);
+                        if (mApexModuleNameToActivePackageName.containsKey(ai.moduleName)) {
+                            throw new IllegalStateException(
+                                    "Two active packages have the same APEX module name: "
+                                            + ai.moduleName);
+                        }
+                        mApexModuleNameToActivePackageName.put(
+                                ai.moduleName, packageInfo.packageName);
                     }
                     if (ai.isFactory) {
                         // Don't throw when the duplicating APEX is VNDK APEX
@@ -966,6 +987,16 @@ public abstract class ApexManager {
             }
         }
 
+        @Override
+        @Nullable
+        public String getActivePackageNameForApexModuleName(String apexModuleName) {
+            synchronized (mLock) {
+                Preconditions.checkState(mApexModuleNameToActivePackageName != null,
+                        "APEX packages have not been scanned");
+                return mApexModuleNameToActivePackageName.get(apexModuleName);
+            }
+        }
+
         @Override
         public boolean snapshotCeData(int userId, int rollbackId, String apexPackageName) {
             String apexModuleName;
@@ -1390,6 +1421,12 @@ public abstract class ApexManager {
             return null;
         }
 
+        @Override
+        @Nullable
+        public String getActivePackageNameForApexModuleName(String apexModuleName) {
+            return null;
+        }
+
         @Override
         public boolean snapshotCeData(int userId, int rollbackId, String apexPackageName) {
             throw new UnsupportedOperationException();

services/tests/servicestests/src/com/android/server/pm/ApexManagerTest.java

@@ -489,6 +489,20 @@ public class ApexManagerTest {
         assertThat(e).hasMessageThat().contains("Failed to collect certificates from ");
     }
 
+    @Test
+    public void testGetActivePackageNameForApexModuleName() throws Exception {
+        final String moduleName = "com.android.module_name";
+
+        ApexInfo[] apexInfo = createApexInfoForTestPkg(true, false);
+        apexInfo[0].moduleName = moduleName;
+        when(mApexService.getAllPackages()).thenReturn(apexInfo);
+        mApexManager.scanApexPackagesTraced(mPackageParser2,
+                ParallelPackageParser.makeExecutorService());
+
+        assertThat(mApexManager.getActivePackageNameForApexModuleName(moduleName))
+                .isEqualTo(TEST_APEX_PKG);
+    }
+
     private ApexInfo createApexInfoForTestPkg(boolean isActive, boolean isFactory, int version) {
         File apexFile = extractResource(TEST_APEX_PKG,  TEST_APEX_FILE_NAME);
         ApexInfo apexInfo = new ApexInfo();