From 61582b63d00864950bb635e1bb8f57d546020d80 Mon Sep 17 00:00:00 2001
From: Oizaro <98434320+Oizaro@users.noreply.github.com>
Date: Tue, 8 Mar 2022 20:18:49 +0100
Subject: [PATCH] Use new method for lifting idle restrictions temporarily when
 receiving PUSH messages

---
 .../java/org/microg/gms/gcm/McsService.java   | 61 ++++++++++++++-----
 1 file changed, 45 insertions(+), 16 deletions(-)

diff --git a/play-services-core/src/main/java/org/microg/gms/gcm/McsService.java b/play-services-core/src/main/java/org/microg/gms/gcm/McsService.java
index cd6472a3..9533f353 100644
--- a/play-services-core/src/main/java/org/microg/gms/gcm/McsService.java
+++ b/play-services-core/src/main/java/org/microg/gms/gcm/McsService.java
@@ -54,6 +54,7 @@ import static org.microg.gms.gcm.McsConstants.MSG_OUTPUT_ERROR;
 import static org.microg.gms.gcm.McsConstants.MSG_OUTPUT_READY;
 import static org.microg.gms.gcm.McsConstants.MSG_TEARDOWN;
 
+import android.annotation.SuppressLint;
 import android.app.AlarmManager;
 import android.app.PendingIntent;
 import android.app.Service;
@@ -77,6 +78,8 @@ import android.os.SystemClock;
 import android.os.UserHandle;
 import android.util.Log;
 
+import androidx.annotation.Nullable;
+import androidx.annotation.RequiresApi;
 import androidx.legacy.content.WakefulBroadcastReceiver;
 
 import com.mgoogle.android.gms.R;
@@ -159,9 +162,18 @@ public class McsService extends Service implements Handler.Callback {
 
     private static final int maxTtl = 24 * 60 * 60;
 
-    private Object deviceIdleController;
+    @Nullable
     private Method getUserIdMethod;
+    @Nullable
+    private Object deviceIdleController;
+    @Nullable
     private Method addPowerSaveTempWhitelistAppMethod;
+    @Nullable
+    @RequiresApi(Build.VERSION_CODES.S)
+    private Object powerExemptionManager;
+    @Nullable
+    @RequiresApi(Build.VERSION_CODES.S)
+    private Method addToTemporaryAllowListMethod;
 
     private class HandlerThread extends Thread {
 
@@ -190,6 +202,7 @@ public class McsService extends Service implements Handler.Callback {
     }
 
     @Override
+    @SuppressLint("PrivateApi")
     public void onCreate() {
         super.onCreate();
         TriggerReceiver.register(this);
@@ -199,20 +212,27 @@ public class McsService extends Service implements Handler.Callback {
         powerManager = (PowerManager) getSystemService(POWER_SERVICE);
         if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M && checkSelfPermission("android.permission.CHANGE_DEVICE_IDLE_TEMP_WHITELIST") == PackageManager.PERMISSION_GRANTED) {
             try {
-                String deviceIdleControllerName = "deviceidle";
-                try {
-                    Field field = Context.class.getField("DEVICE_IDLE_CONTROLLER");
-                    deviceIdleControllerName = (String) field.get(null);
-                } catch (Exception ignored) {
-                }
-                IBinder binder = (IBinder) Class.forName("android.os.ServiceManager")
-                        .getMethod("getService", String.class).invoke(null, deviceIdleControllerName);
-                if (binder != null) {
-                    deviceIdleController = Class.forName("android.os.IDeviceIdleController$Stub")
-                            .getMethod("asInterface", IBinder.class).invoke(null, binder);
-                    getUserIdMethod = UserHandle.class.getMethod("getUserId", int.class);
-                    addPowerSaveTempWhitelistAppMethod = deviceIdleController.getClass()
-                            .getMethod("addPowerSaveTempWhitelistApp", String.class, long.class, int.class, String.class);
+                if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.S) {
+                    Class<?> powerExemptionManagerClass = Class.forName("android.os.PowerExemptionManager");
+                    powerExemptionManager = getSystemService(powerExemptionManagerClass);
+                    addToTemporaryAllowListMethod =
+                            powerExemptionManagerClass.getMethod("addToTemporaryAllowList", String.class, int.class, String.class, long.class);
+                } else {
+                    String deviceIdleControllerName = "deviceidle";
+                    try {
+                        Field field = Context.class.getField("DEVICE_IDLE_CONTROLLER");
+                        deviceIdleControllerName = (String) field.get(null);
+                    } catch (Exception ignored) {
+                    }
+                    IBinder binder = (IBinder) Class.forName("android.os.ServiceManager")
+                            .getMethod("getService", String.class).invoke(null, deviceIdleControllerName);
+                    if (binder != null) {
+                        deviceIdleController = Class.forName("android.os.IDeviceIdleController$Stub")
+                                .getMethod("asInterface", IBinder.class).invoke(null, binder);
+                        getUserIdMethod = UserHandle.class.getMethod("getUserId", int.class);
+                        addPowerSaveTempWhitelistAppMethod = deviceIdleController.getClass()
+                                .getMethod("addPowerSaveTempWhitelistApp", String.class, long.class, int.class, String.class);
+                    }
                 }
             } catch (Exception e) {
                 Log.w(TAG, e);
@@ -609,7 +629,16 @@ public class McsService extends Service implements Handler.Callback {
     }
 
     private void addPowerSaveTempWhitelistApp(String packageName) {
-        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.S) {
+            try {
+                if (addToTemporaryAllowListMethod != null && powerExemptionManager != null) {
+                    logd(this, "Adding app " + packageName + " to the temp allowlist");
+                    addToTemporaryAllowListMethod.invoke(powerExemptionManager, packageName, 0, "GCM Push", 10000);
+                }
+            } catch (Exception e) {
+                Log.e(TAG, "Error adding app" + packageName + " to the temp allowlist.", e);
+            }
+        } else if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
             try {
                 if (getUserIdMethod != null && addPowerSaveTempWhitelistAppMethod != null && deviceIdleController != null) {
                     int userId = (int) getUserIdMethod.invoke(null, getPackageManager().getApplicationInfo(packageName, 0).uid);