Fix PowerManagerService handling of work chains.

When a wakelock with specific set of flags is requested
PowerManagerService crashes with NPE because
applyWakeLockFlagsOnAcquireLocked does not handle work chains and
requires UID set directly on the work source instead of the work chain.

This fixes the behavior by handling the work chains as well as
introducing additional emptiness checks.

There may be multiple work chains in a work source, and only the first
one is used. I don't know the details of this method, but it looks like
it's a reasonable decision.

Bug: 113868145
Test: atest FrameworksCoreTests:PowerManagerTest

Change-Id: Idbfa1d0cd32ddc021caf6443138688d3f8b7c946

3 files changed, 47 insertions, 6 deletions

diff --git a/core/tests/coretests/AndroidManifest.xml b/core/tests/coretests/AndroidManifest.xml
index d704957a7666..f60d8d0adcdd 100644
--- a/core/tests/coretests/AndroidManifest.xml
+++ b/core/tests/coretests/AndroidManifest.xml
@@ -61,6 +61,7 @@
     <uses-permission android:name="android.permission.READ_LOGS"/>
     <uses-permission android:name="android.permission.READ_PHONE_STATE" />
     <uses-permission android:name="android.permission.READ_SMS"/>
+    <uses-permission android:name="android.permission.UPDATE_DEVICE_STATS" />
     <uses-permission android:name="android.permission.USE_CREDENTIALS" />
     <uses-permission android:name="android.permission.WAKE_LOCK" />
     <uses-permission android:name="android.permission.WRITE_CONTACTS" />
diff --git a/core/tests/coretests/src/android/os/PowerManagerTest.java b/core/tests/coretests/src/android/os/PowerManagerTest.java
index 0d250b825d97..da17b56bc0fa 100644
--- a/core/tests/coretests/src/android/os/PowerManagerTest.java
+++ b/core/tests/coretests/src/android/os/PowerManagerTest.java
@@ -64,7 +64,7 @@ public class PowerManagerTest extends AndroidTestCase {
 
         // TODO: Some sort of functional test (maybe not in the unit test here?)
         // that confirms that things are really happening e.g. screen power, keyboard power.
-}
+    }
 
     /**
      * Confirm that we can't create dysfunctional wakelocks.
@@ -85,6 +85,25 @@ public class PowerManagerTest extends AndroidTestCase {
     }
 
     /**
+     * Ensure that we can have work sources with work chains when uid is not set directly on work
+     * source, and that this doesn't crash system server.
+     *
+     * @throws Exception
+     */
+    @SmallTest
+    public void testWakeLockWithWorkChains() throws Exception {
+        PowerManager.WakeLock wakeLock = mPm.newWakeLock(
+                PowerManager.SCREEN_DIM_WAKE_LOCK | PowerManager.ACQUIRE_CAUSES_WAKEUP,
+                "TEST_LOCK");
+        WorkSource workSource = new WorkSource();
+        WorkSource.WorkChain workChain = workSource.createWorkChain();
+        workChain.addNode(1000, "test");
+        wakeLock.setWorkSource(workSource);
+
+        doTestWakeLock(wakeLock);
+    }
+
+    /**
      * Apply a few tests to a wakelock to make sure it's healthy.
      *
      * @param wl The wakelock to be tested.
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index b3f2a27cf99a..43a9c782b8d6 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -1079,18 +1079,39 @@ public final class PowerManagerService extends SystemService
         return false;
     }
 
+    private static WorkChain getFirstNonEmptyWorkChain(WorkSource workSource) {
+        if (workSource.getWorkChains() == null) {
+            return null;
+        }
+
+        for (WorkChain workChain: workSource.getWorkChains()) {
+            if (workChain.getSize() > 0) {
+                return workChain;
+            }
+        }
+
+        return null;
+    }
+
     private void applyWakeLockFlagsOnAcquireLocked(WakeLock wakeLock, int uid) {
         if ((wakeLock.mFlags & PowerManager.ACQUIRE_CAUSES_WAKEUP) != 0
                 && isScreenLock(wakeLock)) {
             String opPackageName;
             int opUid;
-            if (wakeLock.mWorkSource != null && wakeLock.mWorkSource.getName(0) != null) {
-                opPackageName = wakeLock.mWorkSource.getName(0);
-                opUid = wakeLock.mWorkSource.get(0);
+            if (wakeLock.mWorkSource != null && !wakeLock.mWorkSource.isEmpty()) {
+                WorkSource workSource = wakeLock.mWorkSource;
+                WorkChain workChain = getFirstNonEmptyWorkChain(workSource);
+                if (workChain != null) {
+                    opPackageName = workChain.getAttributionTag();
+                    opUid = workChain.getAttributionUid();
+                } else {
+                    opPackageName = workSource.getName(0) != null
+                            ? workSource.getName(0) : wakeLock.mPackageName;
+                    opUid = workSource.get(0);
+                }
             } else {
                 opPackageName = wakeLock.mPackageName;
-                opUid = wakeLock.mWorkSource != null ? wakeLock.mWorkSource.get(0)
-                        : wakeLock.mOwnerUid;
+                opUid = wakeLock.mOwnerUid;
             }
             wakeUpNoUpdateLocked(SystemClock.uptimeMillis(), wakeLock.mTag, opUid,
                     opPackageName, opUid);