Save zygoteInit callbacks to invoke after process forked

Bridge/src/main/java/com/elderdrivers/riru/xposed/Main.java

@@ -55,6 +55,7 @@ public class Main implements KeepAll {
         sAppDataDir = appDataDir;
         sIsGlobalMode = isGlobalMode;
         sIsDynamicModules = isDynamicModules;
+        Router.prepare(false);
         if (isGlobalMode) {
             // do bootstrap hooking only once in zygote process
             Router.onProcessForked(false);
@@ -78,6 +79,7 @@ public class Main implements KeepAll {
                 // load modules for each app process on its forked
                 Router.loadModulesSafely();
             }
+            Router.callZygoteInits();
         } else {
             // in zygote process, res is child zygote pid
             // don't print log here, see https://github.com/RikkaApps/Riru/blob/77adfd6a4a6a81bfd20569c910bc4854f2f84f5e/riru-core/jni/main/jni_native_method.cpp#L55-L66
@@ -97,8 +99,10 @@ public class Main implements KeepAll {
 //        Utils.logD(sForkSystemServerPramsStr + " = " + pid);
         if (pid == 0) {
             // in system_server process
+            Router.prepare(true);
             Router.onProcessForked(true);
             Router.loadModulesSafely();
+            Router.callZygoteInits();
         } else {
             // in zygote process, res is child zygote pid
             // don't print log here, see https://github.com/RikkaApps/Riru/blob/77adfd6a4a6a81bfd20569c910bc4854f2f84f5e/riru-core/jni/main/jni_native_method.cpp#L55-L66

Bridge/src/main/java/com/elderdrivers/riru/xposed/entry/Router.java

@@ -7,11 +7,19 @@ import com.elderdrivers.riru.xposed.entry.bootstrap.SysInnerHookInfo;
 import com.elderdrivers.riru.xposed.entry.hooker.SystemMainHooker;
 import com.elderdrivers.riru.xposed.util.Utils;
 
+import java.util.Map;
+
+import de.robv.android.xposed.IXposedHookZygoteInit;
 import de.robv.android.xposed.XposedBridge;
 import de.robv.android.xposed.XposedInit;
 
 public class Router {
 
+    public static void prepare(boolean isSystem) {
+        // this flag is needed when loadModules
+        XposedInit.startsSystemServer = isSystem;
+    }
+
     public static void onProcessForked(boolean isSystem) {
         // Initialize the Xposed framework
         try {
@@ -31,6 +39,16 @@ public class Router {
         }
     }
 
+    public static void callZygoteInits() {
+        for (Map.Entry<IXposedHookZygoteInit, IXposedHookZygoteInit.StartupParam> entry : XposedBridge.sZygoteInitCallbacks.entrySet()) {
+            try {
+                entry.getKey().initZygote(entry.getValue());
+            } catch (Throwable t) {
+                Utils.logE("Failed to load class " + entry.getKey().getClass(), t);
+            }
+        }
+    }
+
     public static void startBootstrapHook(boolean isSystem) {
         Utils.logD("startBootstrapHook starts: isSystem = " + isSystem);
         ClassLoader classLoader = XposedBridge.BOOTCLASSLOADER;

Bridge/src/main/java/de/robv/android/xposed/XposedBridge.java

@@ -65,6 +65,7 @@ public final class XposedBridge {
 	// built-in handlers
 	private static final Map<Member, CopyOnWriteSortedSet<XC_MethodHook>> sHookedMethodCallbacks = new HashMap<>();
 	public static final CopyOnWriteSortedSet<XC_LoadPackage> sLoadedPackageCallbacks = new CopyOnWriteSortedSet<>();
+	public static final HashMap<IXposedHookZygoteInit, IXposedHookZygoteInit.StartupParam> sZygoteInitCallbacks = new HashMap<>();
 	/*package*/ static final CopyOnWriteSortedSet<XC_InitPackageResources> sInitPackageResourcesCallbacks = new CopyOnWriteSortedSet<>();
 
 	private XposedBridge() {}
@@ -373,6 +374,13 @@ public final class XposedBridge {
 		}
 	}
 
+	public static void hookZygoteInit(IXposedHookZygoteInit moduleInstance,
+									  IXposedHookZygoteInit.StartupParam param) {
+		synchronized (sZygoteInitCallbacks) {
+			sZygoteInitCallbacks.put(moduleInstance, param);
+		}
+	}
+
 	/**
 	 * Adds a callback to be executed when the resources for an app are initialized.
 	 *

Bridge/src/main/java/de/robv/android/xposed/XposedInit.java

@@ -34,7 +34,7 @@ import static de.robv.android.xposed.XposedHelpers.setStaticLongField;
 
 public final class XposedInit {
     private static final String TAG = XposedBridge.TAG;
-    private static boolean startsSystemServer = false;
+    public static boolean startsSystemServer = false;
     private static final String startClassName = ""; // ed: no support for tool process anymore
 
     public static final String INSTALLER_PACKAGE_NAME = "com.solohsu.android.edxp.manager";
@@ -61,7 +61,6 @@ public final class XposedInit {
         if (!bootstrapHooked.compareAndSet(false, true)) {
             return;
         }
-        startsSystemServer = isSystem;
         Router.startBootstrapHook(isSystem);
         // MIUI
         if (findFieldIfExists(ZygoteInit.class, "BOOT_START_TIME") != null) {
@@ -209,7 +208,8 @@ public final class XposedInit {
                             IXposedHookZygoteInit.StartupParam param = new IXposedHookZygoteInit.StartupParam();
                             param.modulePath = apk;
                             param.startsSystemServer = startsSystemServer;
-                            ((IXposedHookZygoteInit) moduleInstance).initZygote(param);
+//                            ((IXposedHookZygoteInit) moduleInstance).initZygote(param);
+                            XposedBridge.hookZygoteInit((IXposedHookZygoteInit) moduleInstance, param);
                         }
 
                         if (moduleInstance instanceof IXposedHookLoadPackage)