Add bytecode rewriter to inject SplitCompat.install() call

Bug: 930797
Change-Id: I85048ef330fce15ed098876ad9aff8d7b4cdd367
Reviewed-on: https://chromium-review.googlesource.com/c/1476179Reviewed-by: Andrew Grieve <agrieve@chromium.org>
Commit-Queue: Piotr Bialecki <bialpio@chromium.org>
Cr-Commit-Position: refs/heads/master@{#633837}

build/android/bytecode/BUILD.gn

@@ -14,8 +14,9 @@ java_binary("java_bytecode_rewriter") {
     "java/org/chromium/bytecode/ClassPathValidator.java",
     "java/org/chromium/bytecode/CustomClassLoaderClassWriter.java",
     "java/org/chromium/bytecode/CustomResourcesClassAdapter.java",
-    "java/org/chromium/bytecode/TypeUtils.java",
+    "java/org/chromium/bytecode/SplitCompatClassAdapter.java",
     "java/org/chromium/bytecode/ThreadAssertionClassAdapter.java",
+    "java/org/chromium/bytecode/TypeUtils.java",
   ]
   main_class = "org.chromium.bytecode.ByteCodeProcessor"
   deps = [

build/android/bytecode/java/org/chromium/bytecode/ByteCodeProcessor.java

@@ -61,6 +61,7 @@ class ByteCodeProcessor {
     private static ClassLoader sDirectClassPathClassLoader;
     private static ClassLoader sFullClassPathClassLoader;
     private static Set<String> sFullClassPathJarPaths;
+    private static Set<String> sSplitCompatClassNames;
     private static ClassPathValidator sValidator;
 
     private static class EntryDataPair {
@@ -129,6 +130,10 @@ class ByteCodeProcessor {
             chain = new CustomResourcesClassAdapter(
                     chain, reader.getClassName(), reader.getSuperName(), sFullClassPathClassLoader);
         }
+        if (!sSplitCompatClassNames.isEmpty()) {
+            chain = new SplitCompatClassAdapter(
+                    chain, sSplitCompatClassNames, sFullClassPathClassLoader);
+        }
         reader.accept(chain, 0);
         byte[] patchedByteCode = writer.toByteArray();
         return EntryDataPair.create(entry.getName(), patchedByteCode);
@@ -235,6 +240,13 @@ class ByteCodeProcessor {
         currIndex += directJarsLength;
         sDirectClassPathClassLoader = loadJars(directClassPathJarPaths);
 
+        // Load list of class names that need to be fixed.
+        int splitCompatClassNamesLength = Integer.parseInt(args[currIndex++]);
+        sSplitCompatClassNames = new HashSet<>();
+        sSplitCompatClassNames.addAll(Arrays.asList(
+                Arrays.copyOfRange(args, currIndex, currIndex + splitCompatClassNamesLength)));
+        currIndex += splitCompatClassNamesLength;
+
         // Load all jars that are on the classpath for the input jar for analyzing class hierarchy.
         sFullClassPathJarPaths = new HashSet<>();
         sFullClassPathJarPaths.clear();

build/android/bytecode/java/org/chromium/bytecode/SplitCompatClassAdapter.java

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package org.chromium.bytecode;
+
+import static org.objectweb.asm.Opcodes.ACC_PROTECTED;
+import static org.objectweb.asm.Opcodes.ALOAD;
+import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
+import static org.objectweb.asm.Opcodes.INVOKESTATIC;
+import static org.objectweb.asm.Opcodes.RETURN;
+
+import static org.chromium.bytecode.TypeUtils.CONTEXT;
+import static org.chromium.bytecode.TypeUtils.VOID;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+import java.util.Set;
+
+/**
+ * A ClassVisitor for injecting ModuleInstaller.initActivity(activity) method call
+ * into Activity's attachBaseContext() method. The goal is to eventually invoke
+ * SplitCompat.install() method if running with the binary that has bundle support
+ * enabled. This needs to happen for activities that were not built with SplitCompat
+ * support.
+ */
+class SplitCompatClassAdapter extends ClassVisitor {
+    private static final String ANDROID_APP_ACTIVITY_CLASS_NAME = "android/app/Activity";
+    private static final String ATTACH_BASE_CONTEXT_METHOD_NAME = "attachBaseContext";
+    private static final String ATTACH_BASE_CONTEXT_DESCRIPTOR =
+            TypeUtils.getMethodDescriptor(VOID, CONTEXT);
+
+    private static final String MODULE_INSTALLER_CLASS_NAME =
+            "org/chromium/components/module_installer/ModuleInstaller";
+    private static final String INIT_ACTIVITY_METHOD_NAME = "initActivity";
+    private static final String INIT_ACTIVITY_DESCRIPTOR =
+            TypeUtils.getMethodDescriptor(VOID, CONTEXT);
+
+    private boolean mShouldTransform;
+
+    private Set<String> mClassNames;
+
+    private ClassLoader mClassLoader;
+
+    /**
+     * Creates instance of SplitCompatClassAdapter.
+     *
+     * @param visitor
+     * @param classNames Names of classes into which the attachBaseContext method will be
+     *         injected. Currently, we'll only consider classes for bytecode rewriting only if
+     *         they inherit directly from android.app.Activity & not already contain
+     *         attachBaseContext method.
+     * @param classLoader
+     */
+    SplitCompatClassAdapter(ClassVisitor visitor, Set<String> classNames, ClassLoader classLoader) {
+        super(Opcodes.ASM5, visitor);
+
+        mShouldTransform = false;
+        mClassNames = classNames;
+        mClassLoader = classLoader;
+    }
+
+    @Override
+    public void visit(int version, int access, String name, String signature, String superName,
+            String[] interfaces) {
+        super.visit(version, access, name, signature, superName, interfaces);
+
+        if (mClassNames.contains(name)) {
+            if (!isSubclassOfActivity(name)) {
+                throw new RuntimeException(name
+                        + " should be transformed but does not inherit from android.app.Activity");
+            }
+
+            mShouldTransform = true;
+        }
+    }
+
+    @Override
+    public MethodVisitor visitMethod(
+            int access, String name, String descriptor, String signature, String[] exceptions) {
+        // Check if current method matches attachBaseContext & we're supposed to emit code - if so,
+        // fail.
+        if (mShouldTransform && name.equals(ATTACH_BASE_CONTEXT_METHOD_NAME)) {
+            throw new RuntimeException(ATTACH_BASE_CONTEXT_METHOD_NAME + " method already exists");
+        }
+
+        return super.visitMethod(access, name, descriptor, signature, exceptions);
+    }
+
+    @Override
+    public void visitEnd() {
+        if (mShouldTransform) {
+            // If we reached this place, it means we're rewriting a class that inherits from
+            // Activity and there was no exception thrown due to existence of attachBaseContext
+            // method - emit code.
+            emitAttachBaseContext();
+        }
+
+        super.visitEnd();
+    }
+
+    /**
+     * Generates:
+     *
+     * <pre>
+     * protected void attachBaseContext(Context base) {
+     *     super.attachBaseContext(base);
+     *     ModuleInstaller.initActivity(this);
+     * }
+     * </pre>
+     */
+    private void emitAttachBaseContext() {
+        MethodVisitor mv = super.visitMethod(ACC_PROTECTED, ATTACH_BASE_CONTEXT_METHOD_NAME,
+                ATTACH_BASE_CONTEXT_DESCRIPTOR, null, null);
+        mv.visitCode();
+        mv.visitVarInsn(ALOAD, 0); // load "this" on stack
+        mv.visitVarInsn(ALOAD, 1); // load first method parameter on stack (Context)
+        mv.visitMethodInsn(INVOKESPECIAL, ANDROID_APP_ACTIVITY_CLASS_NAME,
+                ATTACH_BASE_CONTEXT_METHOD_NAME,
+                ATTACH_BASE_CONTEXT_DESCRIPTOR); // invoke super's attach base context
+        mv.visitVarInsn(ALOAD, 0); // load "this" on stack
+        mv.visitMethodInsn(INVOKESTATIC, MODULE_INSTALLER_CLASS_NAME, INIT_ACTIVITY_METHOD_NAME,
+                INIT_ACTIVITY_DESCRIPTOR);
+        mv.visitInsn(RETURN);
+        mv.visitMaxs(2, 2); // max stack size - 2, max locals - 2
+        mv.visitEnd();
+    }
+
+    /**
+     * Checks whether passed in class inherits from android.app.Activity.
+     * @param name Name of the class to be checked.
+     * @return true if class inherits from android.app.Activity, false otherwise.
+     */
+    private boolean isSubclassOfActivity(String name) {
+        Class<?> activityClass = loadClass(ANDROID_APP_ACTIVITY_CLASS_NAME);
+        Class<?> candidateClass = loadClass(name);
+        return activityClass.isAssignableFrom(candidateClass);
+    }
+
+    private Class<?> loadClass(String className) {
+        try {
+            return mClassLoader.loadClass(className.replace('/', '.'));
+        } catch (ClassNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

build/android/gyp/bytecode_processor.py

@@ -36,6 +36,11 @@ def main(argv):
   _AddSwitch(parser, '--enable-assert')
   _AddSwitch(parser, '--enable-thread-annotations')
   _AddSwitch(parser, '--enable-check-class-path')
+  parser.add_argument(
+      '--split-compat-class-names',
+      action='append',
+      default=[],
+      help='Names of classes that need to be made SplitCompat-enabled.')
   args = parser.parse_args(argv)
 
   sdk_jars = build_utils.ParseGnList(args.sdk_classpath_jars)
@@ -48,6 +53,9 @@ def main(argv):
   for a in args.extra_jars:
     extra_classpath_jars.extend(build_utils.ParseGnList(a))
 
+  split_compat_class_names = build_utils.ParseGnList(
+      args.split_compat_class_names)
+
   if args.verbose:
     verbose = '--verbose'
   else:
@@ -58,7 +66,9 @@ def main(argv):
       args.enable_assert, args.enable_custom_resources,
       args.enable_thread_annotations, args.enable_check_class_path,
       str(len(sdk_jars))
-  ] + sdk_jars + [str(len(direct_jars))] + direct_jars + extra_classpath_jars)
+  ] + sdk_jars + [str(len(direct_jars))] + direct_jars + [
+      str(len(split_compat_class_names))
+  ] + split_compat_class_names + extra_classpath_jars)
   subprocess.check_call(cmd)
 
 

build/config/android/internal_rules.gni

@@ -1486,8 +1486,10 @@ if (enable_java_templates) {
 
     _desugar = defined(invoker.supports_android) && invoker.supports_android
     _emma_instrument = invoker.emma_instrument
+    _enable_split_compat = defined(invoker.split_compat_class_names)
     _enable_bytecode_rewriter =
-        _enable_assert || _enable_custom_resources || _enable_thread_annotations
+        _enable_assert || _enable_custom_resources ||
+        _enable_thread_annotations || _enable_split_compat
     _is_prebuilt = defined(invoker.is_prebuilt) && invoker.is_prebuilt
     _enable_bytecode_checks = !defined(invoker.enable_bytecode_checks) ||
                               invoker.enable_bytecode_checks
@@ -1570,6 +1572,10 @@ if (enable_java_templates) {
         if (_enable_bytecode_checks) {
           args += [ "--enable-check-class-path" ]
         }
+        if (_enable_split_compat) {
+          args += [ "--split-compat-class-names" ] +
+                  invoker.split_compat_class_names
+        }
         args += [
           "--direct-classpath-jars",
           "@FileArg($_rebased_build_config:javac:classpath)",
@@ -3493,6 +3499,7 @@ if (enable_java_templates) {
                                    "enable_bytecode_rewriter",
                                    "jar_excluded_patterns",
                                    "jar_included_patterns",
+                                   "split_compat_class_names",
                                  ])
           is_prebuilt = _is_prebuilt
           supports_android = _supports_android

build/config/android/rules.gni

@@ -3587,6 +3587,10 @@ if (enable_java_templates) {
   #   extract_native_libraries: Whether to extract .so files found in the .aar.
   #       If the file contains .so, either extract_native_libraries or
   #       ignore_native_libraries must be set.
+  #   split_compat_class_names: Names of the classes that will have their
+  #       bytecode rewritten to inject the call to SplitCompat.install().
+  #       Used to make dependencies compatible with SplitCompat to immediately
+  #       access resources brought in by the modules.
   #   create_srcjar: If false, does not create an R.java file.
   #   TODO(jbudorick@): remove this arguments after crbug.com/522043 is fixed.
   #   requires_android: Whether this target can only be used for compiling
@@ -3763,6 +3767,7 @@ if (enable_java_templates) {
                                  "jar_included_patterns",
                                  "proguard_configs",
                                  "requires_android",
+                                 "split_compat_class_names",
                                  "testonly",
                                ])
         if (!defined(deps)) {

components/module_installer/android/java/src-impl/org/chromium/components/module_installer/ModuleInstaller.java

@@ -49,6 +49,17 @@ public class ModuleInstaller {
         updateCrashKeys();
     }
 
+    /**
+     * Needs to be called in attachBaseContext of the activities that want to have access to
+     * splits prior to application restart.
+     *
+     * For details, see:
+     * https://developer.android.com/reference/com/google/android/play/core/splitcompat/SplitCompat.html#install(android.content.Context)
+     */
+    public static void initActivity(Context context) {
+        SplitCompat.install(context);
+    }
+
     /** Writes fully installed and emulated modules to crash keys. */
     public static void updateCrashKeys() {
         Context context = ContextUtils.getApplicationContext();

components/module_installer/android/java/src-stub/org/chromium/components/module_installer/ModuleInstaller.java

@@ -4,11 +4,16 @@
 
 package org.chromium.components.module_installer;
 
+import android.content.Context;
+
 import org.chromium.base.VisibleForTesting;
 
 /** Dummy fallback of ModuleInstaller for APK builds. */
 public class ModuleInstaller {
     public static void init() {}
+
+    public static void initActivity(Context context) {}
+
     public static void updateCrashKeys(){};
 
     public static void install(

components/module_installer/android/java/src-test/org/chromium/components/module_installer/ModuleInstaller.java

@@ -4,6 +4,8 @@
 
 package org.chromium.components.module_installer;
 
+import android.content.Context;
+
 import org.chromium.base.VisibleForTesting;
 
 import java.util.HashSet;
@@ -14,6 +16,9 @@ public class ModuleInstaller {
     private static Set<String> sModulesRequestedDeffered = new HashSet<>();
 
     public static void init() {}
+
+    public static void initActivity(Context context) {}
+
     public static void updateCrashKeys(){};
 
     public static void install(

third_party/android_deps/BUILD.gn

@@ -539,6 +539,7 @@ android_aar_prebuilt("com_google_ar_core_java") {
   aar_path = "libs/com_google_ar_core/core-1.6.0.aar"
   info_path = "libs/com_google_ar_core/com_google_ar_core.info"
   extract_native_libraries = true
+  split_compat_class_names = [ "com/google/ar/core/InstallActivity" ]
 }
 
 java_prebuilt("com_google_dagger_dagger_java") {

tools/android/roll/android_deps/buildSrc/src/main/groovy/BuildConfigGenerator.groovy

@@ -170,6 +170,9 @@ class BuildConfigGenerator extends DefaultTask {
                 // Target .aar file contains .so libraries that need to be extracted,
                 // and android_aar_prebuilt template will fail if it's not set explictly.
                 sb.append('  extract_native_libraries = true\n')
+                // InstallActivity class is downloaded as a part of DFM & we need to inject
+                // a call to SplitCompat.install() into it.
+                sb.append('  split_compat_class_names = [ "com/google/ar/core/InstallActivity" ]\n')
                 break
         }
     }