Ensure allowlist token

Bug: 320659371
Test: NotificationManagerServiceTest
Test: NotificationManager
Change-Id: I9d2d848a9beb2258affe58c0d0089e3a493d6de2
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index 0a34d36..b734315 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -3540,15 +3540,12 @@
      * Sets the token used for background operations for the pending intents associated with this
      * notification.
      *
-     * This token is automatically set during deserialization for you, you usually won't need to
-     * call this unless you want to change the existing token, if any.
-     *
      * @hide
      */
-    public void clearAllowlistToken() {
-        mAllowlistToken = null;
+    public void overrideAllowlistToken(IBinder token) {
+        mAllowlistToken = token;
         if (publicVersion != null) {
-            publicVersion.clearAllowlistToken();
+            publicVersion.overrideAllowlistToken(token);
         }
     }
 
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 6fa737d..7c3d43f 100755
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -727,7 +727,7 @@
 
     private static final int MY_UID = Process.myUid();
     private static final int MY_PID = Process.myPid();
-    private static final IBinder ALLOWLIST_TOKEN = new Binder();
+    static final IBinder ALLOWLIST_TOKEN = new Binder();
     protected RankingHandler mRankingHandler;
     private long mLastOverRateLogTime;
     private float mMaxPackageEnqueueRate = DEFAULT_MAX_NOTIFICATION_ENQUEUE_RATE;
@@ -4759,7 +4759,7 @@
                     // Remove background token before returning notification to untrusted app, this
                     // ensures the app isn't able to perform background operations that are
                     // associated with notification interactions.
-                    notification.clearAllowlistToken();
+                    notification.overrideAllowlistToken(null);
                     return new StatusBarNotification(
                             sbn.getPackageName(),
                             sbn.getOpPkg(),
@@ -7634,6 +7634,8 @@
             }
         }
 
+        notification.overrideAllowlistToken(ALLOWLIST_TOKEN);
+
         // Remote views? Are they too big?
         checkRemoteViews(pkg, tag, id, notification);
     }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index 3426cbe..f7a4742 100755
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -323,6 +323,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -12862,6 +12863,35 @@
     }
 
     @Test
+    public void fixNotification_customAllowlistToken()
+            throws Exception {
+        Notification n = new Notification.Builder(mContext, "test")
+                .build();
+        try {
+            Field allowlistToken = Class.forName("android.app.Notification").
+                    getDeclaredField("mAllowlistToken");
+            allowlistToken.setAccessible(true);
+            allowlistToken.set(n, new Binder());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        mService.fixNotification(n, PKG, "tag", 9, 0, mUid, NOT_FOREGROUND_SERVICE, true);
+
+        IBinder actual = null;
+        try {
+            Field allowlistToken = Class.forName("android.app.Notification").
+                    getDeclaredField("mAllowlistToken");
+            allowlistToken.setAccessible(true);
+            actual = (IBinder) allowlistToken.get(n);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        assertTrue(mService.ALLOWLIST_TOKEN == actual);
+    }
+
+    @Test
     public void testCancelAllNotifications_IgnoreUserInitiatedJob() throws Exception {
         when(mJsi.isNotificationAssociatedWithAnyUserInitiatedJobs(anyInt(), anyInt(), anyString()))
                 .thenReturn(true);