Remove TaskInfo#topActivityToken usage in size compat for security

Security report shows that this can cause leak token of different app.

Replace the functionality with a callback to the TaskOrganizerController
to restart activity when size compat restart button is clicked.

Bug: 186776724
Test: manually verify the restart button still works
Change-Id: I097b9f02e8435e6765695b9d5a531a4e165bac66
Merged-In: I097b9f02e8435e6765695b9d5a531a4e165bac66

19 files changed, 147 insertions, 130 deletions

diff --git a/core/java/android/app/ActivityClient.java b/core/java/android/app/ActivityClient.java
index 633b986c06c9..bd4386885dd6 100644
--- a/core/java/android/app/ActivityClient.java
+++ b/core/java/android/app/ActivityClient.java
@@ -431,19 +431,6 @@ public class ActivityClient {
         }
     }
 
-    /**
-     * Restart the process and activity to adopt the latest configuration for size compat mode.
-     * This only takes effect for visible activity because invisible background activity can be
-     * restarted naturally when it becomes visible.
-     */
-    public void restartActivityProcessIfVisible(IBinder token) {
-        try {
-            getActivityClientController().restartActivityProcessIfVisible(token);
-        } catch (RemoteException e) {
-            e.rethrowFromSystemServer();
-        }
-    }
-
     /** Removes the snapshot of home task. */
     public void invalidateHomeTaskSnapshot(IBinder homeToken) {
         try {
diff --git a/core/java/android/app/IActivityClientController.aidl b/core/java/android/app/IActivityClientController.aidl
index 09b0c2f96b13..c6649692d848 100644
--- a/core/java/android/app/IActivityClientController.aidl
+++ b/core/java/android/app/IActivityClientController.aidl
@@ -117,16 +117,6 @@ interface IActivityClientController {
     oneway void setDisablePreviewScreenshots(in IBinder token, boolean disable);
 
     /**
-     * Restarts the activity by killing its process if it is visible. If the activity is not
-     * visible, the activity will not be restarted immediately and just keep the activity record in
-     * the stack. It also resets the current override configuration so the activity will use the
-     * configuration according to the latest state.
-     *
-     * @param activityToken The token of the target activity to restart.
-     */
-    void restartActivityProcessIfVisible(in IBinder activityToken);
-
-    /**
      * It should only be called from home activity to remove its outdated snapshot. The home
      * snapshot is used to speed up entering home from screen off. If the content of home activity
      * is significantly different from before taking the snapshot, then the home activity can use
diff --git a/core/java/android/app/TaskInfo.java b/core/java/android/app/TaskInfo.java
index b95412f00453..c1871be80462 100644
--- a/core/java/android/app/TaskInfo.java
+++ b/core/java/android/app/TaskInfo.java
@@ -187,13 +187,6 @@ public class TaskInfo {
     public ActivityInfo topActivityInfo;
 
     /**
-     * The top activity in this task.
-     * @hide
-     */
-    @Nullable
-    public IBinder topActivityToken;
-
-    /**
      * Whether the direct top activity is in size compat mode on foreground.
      * @hide
      */
@@ -356,8 +349,7 @@ public class TaskInfo {
         return displayId == that.displayId
                 && taskId == that.taskId
                 && topActivityInSizeCompat == that.topActivityInSizeCompat
-                // TopActivityToken and bounds are important if top activity is in size compat
-                && (!topActivityInSizeCompat || topActivityToken.equals(that.topActivityToken))
+                // Bounds are important if top activity is in size compat
                 && (!topActivityInSizeCompat || configuration.windowConfiguration.getBounds()
                     .equals(that.configuration.windowConfiguration.getBounds()))
                 && (!topActivityInSizeCompat || configuration.getLayoutDirection()
@@ -396,7 +388,6 @@ public class TaskInfo {
         parentTaskId = source.readInt();
         isFocused = source.readBoolean();
         isVisible = source.readBoolean();
-        topActivityToken = source.readStrongBinder();
         topActivityInSizeCompat = source.readBoolean();
         mTopActivityLocusId = source.readTypedObject(LocusId.CREATOR);
     }
@@ -434,7 +425,6 @@ public class TaskInfo {
         dest.writeInt(parentTaskId);
         dest.writeBoolean(isFocused);
         dest.writeBoolean(isVisible);
-        dest.writeStrongBinder(topActivityToken);
         dest.writeBoolean(topActivityInSizeCompat);
         dest.writeTypedObject(mTopActivityLocusId, flags);
     }
@@ -462,7 +452,6 @@ public class TaskInfo {
                 + " parentTaskId=" + parentTaskId
                 + " isFocused=" + isFocused
                 + " isVisible=" + isVisible
-                + " topActivityToken=" + topActivityToken
                 + " topActivityInSizeCompat=" + topActivityInSizeCompat
                 + " locusId= " + mTopActivityLocusId
                 + "}";
diff --git a/core/java/android/window/ITaskOrganizerController.aidl b/core/java/android/window/ITaskOrganizerController.aidl
index 2d0211e129bf..a833600e1fbc 100644
--- a/core/java/android/window/ITaskOrganizerController.aidl
+++ b/core/java/android/window/ITaskOrganizerController.aidl
@@ -61,4 +61,9 @@ interface ITaskOrganizerController {
      */
     void setInterceptBackPressedOnTaskRoot(in WindowContainerToken task,
             boolean interceptBackPressed);
+
+    /**
+     * Restarts the top activity in the given task by killing its process if it is visible.
+     */
+    void restartTaskTopActivityProcessIfVisible(in WindowContainerToken task);
 }
diff --git a/core/java/android/window/TaskOrganizer.java b/core/java/android/window/TaskOrganizer.java
index 73995491668a..c7c91cdd0941 100644
--- a/core/java/android/window/TaskOrganizer.java
+++ b/core/java/android/window/TaskOrganizer.java
@@ -223,6 +223,19 @@ public class TaskOrganizer extends WindowOrganizer {
     }
 
     /**
+     * Restarts the top activity in the given task by killing its process if it is visible.
+     * @hide
+     */
+    @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS)
+    public void restartTaskTopActivityProcessIfVisible(@NonNull WindowContainerToken task) {
+        try {
+            mTaskOrganizerController.restartTaskTopActivityProcessIfVisible(task);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Gets the executor to run callbacks on.
      * @hide
      */
diff --git a/data/etc/services.core.protolog.json b/data/etc/services.core.protolog.json
index ac5e2d0fcacb..b67988ee9646 100644
--- a/data/etc/services.core.protolog.json
+++ b/data/etc/services.core.protolog.json
@@ -151,6 +151,12 @@
       "group": "WM_DEBUG_ADD_REMOVE",
       "at": "com\/android\/server\/wm\/WindowManagerService.java"
     },
+    "-1963363332": {
+      "message": "Restart top activity process of Task taskId=%d",
+      "level": "VERBOSE",
+      "group": "WM_DEBUG_WINDOW_ORGANIZER",
+      "at": "com\/android\/server\/wm\/TaskOrganizerController.java"
+    },
     "-1949279037": {
       "message": "Attempted to add input method window with bad token %s.  Aborting.",
       "level": "WARN",
@@ -163,12 +169,6 @@
       "group": "WM_DEBUG_WINDOW_ORGANIZER",
       "at": "com\/android\/server\/wm\/TaskOrganizerController.java"
     },
-    "-1939358269": {
-      "message": "mRecentScreenshotAnimator finish",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_RECENTS_ANIMATIONS",
-      "at": "com\/android\/server\/wm\/RecentsAnimationController.java"
-    },
     "-1938839202": {
       "message": "SURFACE LEAK DESTROY: %s",
       "level": "INFO",
@@ -3499,12 +3499,6 @@
       "group": "WM_DEBUG_APP_TRANSITIONS",
       "at": "com\/android\/server\/wm\/AppTransitionController.java"
     },
-    "1984470582": {
-      "message": "Creating TaskScreenshotAnimatable: task: %s width: %d height: %d",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_RECENTS_ANIMATIONS",
-      "at": "com\/android\/server\/wm\/TaskScreenshotAnimatable.java"
-    },
     "1984782949": {
       "message": ">>> OPEN TRANSACTION animate",
       "level": "INFO",
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index 4b1955e56a6c..8cdb434d4f74 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -62,7 +62,8 @@ import java.util.function.Consumer;
  * Unified task organizer for all components in the shell.
  * TODO(b/167582004): may consider consolidating this class and TaskOrganizer
  */
-public class ShellTaskOrganizer extends TaskOrganizer {
+public class ShellTaskOrganizer extends TaskOrganizer implements
+        SizeCompatUIController.SizeCompatUICallback {
 
     // Intentionally using negative numbers here so the positive numbers can be used
     // for task id specific listeners that will be added later.
@@ -158,6 +159,9 @@ public class ShellTaskOrganizer extends TaskOrganizer {
             Context context, @Nullable SizeCompatUIController sizeCompatUI) {
         super(taskOrganizerController, mainExecutor);
         mSizeCompatUI = sizeCompatUI;
+        if (sizeCompatUI != null) {
+            sizeCompatUI.setSizeCompatUICallback(this);
+        }
     }
 
     @Override
@@ -481,6 +485,17 @@ public class ShellTaskOrganizer extends TaskOrganizer {
         }
     }
 
+    @Override
+    public void onSizeCompatRestartButtonClicked(int taskId) {
+        final TaskAppearedInfo info;
+        synchronized (mLock) {
+            info = mTasks.get(taskId);
+        }
+        if (info != null) {
+            restartTaskTopActivityProcessIfVisible(info.getTaskInfo().token);
+        }
+    }
+
     /**
      * Notifies {@link SizeCompatUIController} about the size compat info changed on the give Task
      * to update the UI accordingly.
@@ -499,13 +514,12 @@ public class ShellTaskOrganizer extends TaskOrganizer {
         if (taskListener == null || !taskListener.supportSizeCompatUI()
                 || !taskInfo.topActivityInSizeCompat) {
             mSizeCompatUI.onSizeCompatInfoChanged(taskInfo.displayId, taskInfo.taskId,
-                    null /* taskConfig */, null /* sizeCompatActivity*/,
-                    null /* taskListener */);
+                    null /* taskConfig */, null /* taskListener */);
             return;
         }
 
         mSizeCompatUI.onSizeCompatInfoChanged(taskInfo.displayId, taskInfo.taskId,
-                taskInfo.configuration, taskInfo.topActivityToken, taskListener);
+                taskInfo.configuration, taskListener);
     }
 
     private TaskListener getTaskListener(RunningTaskInfo runningTaskInfo) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java
index c981adee9b5c..1fc4d12def1f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUIController.java
@@ -20,7 +20,6 @@ import android.annotation.Nullable;
 import android.content.Context;
 import android.content.res.Configuration;
 import android.hardware.display.DisplayManager;
-import android.os.IBinder;
 import android.util.ArraySet;
 import android.util.Log;
 import android.util.SparseArray;
@@ -45,6 +44,13 @@ import java.util.function.Consumer;
  */
 public class SizeCompatUIController implements DisplayController.OnDisplaysChangedListener,
         DisplayImeController.ImePositionProcessor {
+
+    /** Callback for size compat UI interaction. */
+    public interface SizeCompatUICallback {
+        /** Called when the size compat restart button is clicked. */
+        void onSizeCompatRestartButtonClicked(int taskId);
+    }
+
     private static final String TAG = "SizeCompatUIController";
 
     /** Whether the IME is shown on display id. */
@@ -61,6 +67,8 @@ public class SizeCompatUIController implements DisplayController.OnDisplaysChang
     private final DisplayImeController mImeController;
     private final SyncTransactionQueue mSyncQueue;
 
+    private SizeCompatUICallback mCallback;
+
     /** Only show once automatically in the process life. */
     private boolean mHasShownHint;
 
@@ -76,29 +84,31 @@ public class SizeCompatUIController implements DisplayController.OnDisplaysChang
         mImeController.addPositionProcessor(this);
     }
 
+    /** Sets the callback for UI interactions. */
+    public void setSizeCompatUICallback(SizeCompatUICallback callback) {
+        mCallback = callback;
+    }
+
     /**
      * Called when the Task info changed. Creates and updates the size compat UI if there is an
      * activity in size compat, or removes the UI if there is no size compat activity.
-     *
      * @param displayId display the task and activity are in.
      * @param taskId task the activity is in.
      * @param taskConfig task config to place the size compat UI with.
-     * @param sizeCompatActivity the size compat activity in the task. Can be {@code null} if the
-     *                           top activity in this Task is not in size compat.
      * @param taskListener listener to handle the Task Surface placement.
      */
     public void onSizeCompatInfoChanged(int displayId, int taskId,
-            @Nullable Configuration taskConfig, @Nullable IBinder sizeCompatActivity,
+            @Nullable Configuration taskConfig,
             @Nullable ShellTaskOrganizer.TaskListener taskListener) {
-        if (taskConfig == null || sizeCompatActivity == null || taskListener == null) {
+        if (taskConfig == null || taskListener == null) {
             // Null token means the current foreground activity is not in size compatibility mode.
             removeLayout(taskId);
         } else if (mActiveLayouts.contains(taskId)) {
             // UI already exists, update the UI layout.
-            updateLayout(taskId, taskConfig, sizeCompatActivity, taskListener);
+            updateLayout(taskId, taskConfig, taskListener);
         } else {
             // Create a new size compat UI.
-            createLayout(displayId, taskId, taskConfig, sizeCompatActivity, taskListener);
+            createLayout(displayId, taskId, taskConfig, taskListener);
         }
     }
 
@@ -137,7 +147,7 @@ public class SizeCompatUIController implements DisplayController.OnDisplaysChang
     }
 
     private void createLayout(int displayId, int taskId, Configuration taskConfig,
-            IBinder activityToken, ShellTaskOrganizer.TaskListener taskListener) {
+            ShellTaskOrganizer.TaskListener taskListener) {
         final Context context = getOrCreateDisplayContext(displayId);
         if (context == null) {
             Log.e(TAG, "Cannot get context for display " + displayId);
@@ -145,17 +155,16 @@ public class SizeCompatUIController implements DisplayController.OnDisplaysChang
         }
 
         final SizeCompatUILayout layout = createLayout(context, displayId, taskId, taskConfig,
-                activityToken, taskListener);
+                taskListener);
         mActiveLayouts.put(taskId, layout);
         layout.createSizeCompatButton(isImeShowingOnDisplay(displayId));
     }
 
     @VisibleForTesting
     SizeCompatUILayout createLayout(Context context, int displayId, int taskId,
-            Configuration taskConfig, IBinder activityToken,
-            ShellTaskOrganizer.TaskListener taskListener) {
-        final SizeCompatUILayout layout = new SizeCompatUILayout(mSyncQueue, context, taskConfig,
-                taskId, activityToken, taskListener, mDisplayController.getDisplayLayout(displayId),
+            Configuration taskConfig, ShellTaskOrganizer.TaskListener taskListener) {
+        final SizeCompatUILayout layout = new SizeCompatUILayout(mSyncQueue, mCallback, context,
+                taskConfig, taskId, taskListener, mDisplayController.getDisplayLayout(displayId),
                 mHasShownHint);
         // Only show hint for the first time.
         mHasShownHint = true;
@@ -163,13 +172,12 @@ public class SizeCompatUIController implements DisplayController.OnDisplaysChang
     }
 
     private void updateLayout(int taskId, Configuration taskConfig,
-            IBinder sizeCompatActivity,
             ShellTaskOrganizer.TaskListener taskListener) {
         final SizeCompatUILayout layout = mActiveLayouts.get(taskId);
         if (layout == null) {
             return;
         }
-        layout.updateSizeCompatInfo(taskConfig, sizeCompatActivity, taskListener,
+        layout.updateSizeCompatInfo(taskConfig, taskListener,
                 isImeShowingOnDisplay(layout.getDisplayId()));
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUILayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUILayout.java
index c6d994ecde8d..a5e96d14dde6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUILayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sizecompatui/SizeCompatUILayout.java
@@ -23,13 +23,11 @@ import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_TRUSTED_OVERL
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
 
 import android.annotation.Nullable;
-import android.app.ActivityClient;
 import android.content.Context;
 import android.content.res.Configuration;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.os.Binder;
-import android.os.IBinder;
 import android.view.SurfaceControl;
 import android.view.View;
 import android.view.WindowManager;
@@ -48,11 +46,11 @@ class SizeCompatUILayout {
     private static final String TAG = "SizeCompatUILayout";
 
     private final SyncTransactionQueue mSyncQueue;
+    private final SizeCompatUIController.SizeCompatUICallback mCallback;
     private Context mContext;
     private Configuration mTaskConfig;
     private final int mDisplayId;
     private final int mTaskId;
-    private IBinder mActivityToken;
     private ShellTaskOrganizer.TaskListener mTaskListener;
     private DisplayLayout mDisplayLayout;
 
@@ -72,15 +70,16 @@ class SizeCompatUILayout {
     final int mPopupOffsetY;
     boolean mShouldShowHint;
 
-    SizeCompatUILayout(SyncTransactionQueue syncQueue, Context context, Configuration taskConfig,
-            int taskId, IBinder activityToken, ShellTaskOrganizer.TaskListener taskListener,
+    SizeCompatUILayout(SyncTransactionQueue syncQueue,
+            SizeCompatUIController.SizeCompatUICallback callback, Context context,
+            Configuration taskConfig, int taskId, ShellTaskOrganizer.TaskListener taskListener,
             DisplayLayout displayLayout, boolean hasShownHint) {
         mSyncQueue = syncQueue;
+        mCallback = callback;
         mContext = context.createConfigurationContext(taskConfig);
         mTaskConfig = taskConfig;
         mDisplayId = mContext.getDisplayId();
         mTaskId = taskId;
-        mActivityToken = activityToken;
         mTaskListener = taskListener;
         mDisplayLayout = displayLayout;
         mShouldShowHint = !hasShownHint;
@@ -141,12 +140,11 @@ class SizeCompatUILayout {
     }
 
     /** Called when size compat info changed. */
-    void updateSizeCompatInfo(Configuration taskConfig, IBinder activityToken,
+    void updateSizeCompatInfo(Configuration taskConfig,
             ShellTaskOrganizer.TaskListener taskListener, boolean isImeShowing) {
         final Configuration prevTaskConfig = mTaskConfig;
         final ShellTaskOrganizer.TaskListener prevTaskListener = mTaskListener;
         mTaskConfig = taskConfig;
-        mActivityToken = activityToken;
         mTaskListener = taskListener;
 
         // Update configuration.
@@ -253,7 +251,7 @@ class SizeCompatUILayout {
 
     /** Called when the restart button is clicked. */
     void onRestartButtonClicked() {
-        ActivityClient.getInstance().restartActivityProcessIfVisible(mActivityToken);
+        mCallback.onSizeCompatRestartButtonClicked(mTaskId);
     }
 
     /** Called when the restart button is long clicked. */
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
index df0a856db73c..cf5bd3ace806 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
@@ -48,6 +48,7 @@ import android.view.SurfaceControl;
 import android.window.ITaskOrganizer;
 import android.window.ITaskOrganizerController;
 import android.window.TaskAppearedInfo;
+import android.window.WindowContainerToken;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -289,7 +290,6 @@ public class ShellTaskOrganizerTests {
     public void testOnSizeCompatActivityChanged() {
         final RunningTaskInfo taskInfo1 = createTaskInfo(12, WINDOWING_MODE_FULLSCREEN);
         taskInfo1.displayId = DEFAULT_DISPLAY;
-        taskInfo1.topActivityToken = mock(IBinder.class);
         taskInfo1.topActivityInSizeCompat = false;
         final TrackingTaskListener taskListener = new TrackingTaskListener();
         mOrganizer.addListenerForType(taskListener, TASK_LISTENER_TYPE_FULLSCREEN);
@@ -297,23 +297,22 @@ public class ShellTaskOrganizerTests {
 
         // sizeCompatActivity is null if top activity is not in size compat.
         verify(mSizeCompatUI).onSizeCompatInfoChanged(taskInfo1.displayId, taskInfo1.taskId,
-                null /* taskConfig */, null /* sizeCompatActivity*/, null /* taskListener */);
+                null /* taskConfig */, null /* taskListener */);
 
         // sizeCompatActivity is non-null if top activity is in size compat.
         clearInvocations(mSizeCompatUI);
         final RunningTaskInfo taskInfo2 =
                 createTaskInfo(taskInfo1.taskId, taskInfo1.getWindowingMode());
         taskInfo2.displayId = taskInfo1.displayId;
-        taskInfo2.topActivityToken = taskInfo1.topActivityToken;
         taskInfo2.topActivityInSizeCompat = true;
         mOrganizer.onTaskInfoChanged(taskInfo2);
         verify(mSizeCompatUI).onSizeCompatInfoChanged(taskInfo1.displayId, taskInfo1.taskId,
-                taskInfo1.configuration, taskInfo1.topActivityToken, taskListener);
+                taskInfo1.configuration, taskListener);
 
         clearInvocations(mSizeCompatUI);
         mOrganizer.onTaskVanished(taskInfo1);
         verify(mSizeCompatUI).onSizeCompatInfoChanged(taskInfo1.displayId, taskInfo1.taskId,
-                null /* taskConfig */, null /* sizeCompatActivity*/, null /* taskListener */);
+                null /* taskConfig */, null /* taskListener */);
     }
 
     @Test
@@ -433,6 +432,18 @@ public class ShellTaskOrganizerTests {
         assertEquals(listener.invisibleLocusTasks.size(), 0);
     }
 
+    @Test
+    public void testOnSizeCompatRestartButtonClicked() throws RemoteException {
+        RunningTaskInfo task1 = createTaskInfo(1, WINDOWING_MODE_MULTI_WINDOW);
+        task1.token = mock(WindowContainerToken.class);
+
+        mOrganizer.onTaskAppeared(task1, null);
+
+        mOrganizer.onSizeCompatRestartButtonClicked(task1.taskId);
+
+        verify(mTaskOrganizerController).restartTaskTopActivityProcessIfVisible(task1.token);
+    }
+
     private static RunningTaskInfo createTaskInfo(int taskId, int windowingMode) {
         RunningTaskInfo taskInfo = new RunningTaskInfo();
         taskInfo.taskId = taskId;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatHintPopupTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatHintPopupTest.java
index 9845d4650d20..10fd7d705967 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatHintPopupTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatHintPopupTest.java
@@ -22,7 +22,6 @@ import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.verify;
 
 import android.content.res.Configuration;
-import android.os.IBinder;
 import android.testing.AndroidTestingRunner;
 import android.view.LayoutInflater;
 import android.widget.Button;
@@ -52,7 +51,7 @@ import org.mockito.MockitoAnnotations;
 public class SizeCompatHintPopupTest extends ShellTestCase {
 
     @Mock private SyncTransactionQueue mSyncTransactionQueue;
-    @Mock private IBinder mActivityToken;
+    @Mock private SizeCompatUIController.SizeCompatUICallback mCallback;
     @Mock private ShellTaskOrganizer.TaskListener mTaskListener;
     @Mock private DisplayLayout mDisplayLayout;
 
@@ -64,8 +63,9 @@ public class SizeCompatHintPopupTest extends ShellTestCase {
         MockitoAnnotations.initMocks(this);
 
         final int taskId = 1;
-        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mContext, new Configuration(),
-                taskId, mActivityToken, mTaskListener, mDisplayLayout, false /* hasShownHint*/);
+        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mCallback, mContext,
+                new Configuration(), taskId, mTaskListener, mDisplayLayout,
+                false /* hasShownHint */);
         mHint = (SizeCompatHintPopup)
                 LayoutInflater.from(mContext).inflate(R.layout.size_compat_mode_hint, null);
         mHint.inject(mLayout);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButtonTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButtonTest.java
index 5a43925a5677..a20a5e9e8d91 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButtonTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatRestartButtonTest.java
@@ -22,7 +22,6 @@ import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.verify;
 
 import android.content.res.Configuration;
-import android.os.IBinder;
 import android.testing.AndroidTestingRunner;
 import android.view.LayoutInflater;
 import android.widget.ImageButton;
@@ -51,8 +50,10 @@ import org.mockito.MockitoAnnotations;
 @SmallTest
 public class SizeCompatRestartButtonTest extends ShellTestCase {
 
+    private static final int TASK_ID = 1;
+
     @Mock private SyncTransactionQueue mSyncTransactionQueue;
-    @Mock private IBinder mActivityToken;
+    @Mock private SizeCompatUIController.SizeCompatUICallback mCallback;
     @Mock private ShellTaskOrganizer.TaskListener mTaskListener;
     @Mock private DisplayLayout mDisplayLayout;
 
@@ -63,9 +64,9 @@ public class SizeCompatRestartButtonTest extends ShellTestCase {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
-        final int taskId = 1;
-        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mContext, new Configuration(),
-                taskId, mActivityToken, mTaskListener, mDisplayLayout, false /* hasShownHint*/);
+        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mCallback, mContext,
+                new Configuration(), TASK_ID, mTaskListener, mDisplayLayout,
+                false /* hasShownHint */);
         mButton = (SizeCompatRestartButton)
                 LayoutInflater.from(mContext).inflate(R.layout.size_compat_ui, null);
         mButton.inject(mLayout);
@@ -75,12 +76,11 @@ public class SizeCompatRestartButtonTest extends ShellTestCase {
 
     @Test
     public void testOnClick() {
-        doNothing().when(mLayout).onRestartButtonClicked();
-
         final ImageButton button = mButton.findViewById(R.id.size_compat_restart_button);
         button.performClick();
 
         verify(mLayout).onRestartButtonClicked();
+        verify(mCallback).onSizeCompatRestartButtonClicked(TASK_ID);
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java
index 806a90b7832a..8839f58ea889 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUIControllerTest.java
@@ -27,7 +27,6 @@ import static org.mockito.Mockito.verify;
 
 import android.content.Context;
 import android.content.res.Configuration;
-import android.os.IBinder;
 import android.testing.AndroidTestingRunner;
 
 import androidx.test.filters.SmallTest;
@@ -61,7 +60,6 @@ public class SizeCompatUIControllerTest extends ShellTestCase {
     private @Mock DisplayController mMockDisplayController;
     private @Mock DisplayLayout mMockDisplayLayout;
     private @Mock DisplayImeController mMockImeController;
-    private @Mock IBinder mMockActivityToken;
     private @Mock ShellTaskOrganizer.TaskListener mMockTaskListener;
     private @Mock SyncTransactionQueue mMockSyncQueue;
     private @Mock SizeCompatUILayout mMockLayout;
@@ -77,8 +75,7 @@ public class SizeCompatUIControllerTest extends ShellTestCase {
                 mMockImeController, mMockSyncQueue) {
             @Override
             SizeCompatUILayout createLayout(Context context, int displayId, int taskId,
-                    Configuration taskConfig, IBinder activityToken,
-                    ShellTaskOrganizer.TaskListener taskListener) {
+                    Configuration taskConfig, ShellTaskOrganizer.TaskListener taskListener) {
                 return mMockLayout;
             }
         };
@@ -97,21 +94,21 @@ public class SizeCompatUIControllerTest extends ShellTestCase {
 
         // Verify that the restart button is added with non-null size compat info.
         mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
-                mMockActivityToken, mMockTaskListener);
+                mMockTaskListener);
 
         verify(mController).createLayout(any(), eq(DISPLAY_ID), eq(TASK_ID), eq(taskConfig),
-                eq(mMockActivityToken), eq(mMockTaskListener));
+                eq(mMockTaskListener));
 
         // Verify that the restart button is updated with non-null new size compat info.
         final Configuration newTaskConfig = new Configuration();
         mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, newTaskConfig,
-                mMockActivityToken, mMockTaskListener);
+                mMockTaskListener);
 
-        verify(mMockLayout).updateSizeCompatInfo(taskConfig, mMockActivityToken, mMockTaskListener,
+        verify(mMockLayout).updateSizeCompatInfo(taskConfig, mMockTaskListener,
                 false /* isImeShowing */);
 
         // Verify that the restart button is removed with null size compat info.
-        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, null, null, mMockTaskListener);
+        mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, null, mMockTaskListener);
 
         verify(mMockLayout).release();
     }
@@ -120,7 +117,7 @@ public class SizeCompatUIControllerTest extends ShellTestCase {
     public void testOnDisplayRemoved() {
         final Configuration taskConfig = new Configuration();
         mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
-                mMockActivityToken, mMockTaskListener);
+                mMockTaskListener);
 
         mController.onDisplayRemoved(DISPLAY_ID + 1);
 
@@ -135,7 +132,7 @@ public class SizeCompatUIControllerTest extends ShellTestCase {
     public void testOnDisplayConfigurationChanged() {
         final Configuration taskConfig = new Configuration();
         mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
-                mMockActivityToken, mMockTaskListener);
+                mMockTaskListener);
 
         final Configuration newTaskConfig = new Configuration();
         mController.onDisplayConfigurationChanged(DISPLAY_ID + 1, newTaskConfig);
@@ -151,7 +148,7 @@ public class SizeCompatUIControllerTest extends ShellTestCase {
     public void testChangeButtonVisibilityOnImeShowHide() {
         final Configuration taskConfig = new Configuration();
         mController.onSizeCompatInfoChanged(DISPLAY_ID, TASK_ID, taskConfig,
-                mMockActivityToken, mMockTaskListener);
+                mMockTaskListener);
 
         mController.onImeVisibilityChanged(DISPLAY_ID, true /* isShowing */);
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUILayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUILayoutTest.java
index f33cfe86224f..ee4c81547bbd 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUILayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sizecompatui/SizeCompatUILayoutTest.java
@@ -21,20 +21,16 @@ import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.clearInvocations;
-import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
-import android.app.ActivityClient;
 import android.content.res.Configuration;
 import android.graphics.Rect;
-import android.os.IBinder;
 import android.testing.AndroidTestingRunner;
 import android.view.DisplayInfo;
 import android.view.SurfaceControl;
@@ -66,7 +62,7 @@ public class SizeCompatUILayoutTest extends ShellTestCase {
     private static final int TASK_ID = 1;
 
     @Mock private SyncTransactionQueue mSyncTransactionQueue;
-    @Mock private IBinder mActivityToken;
+    @Mock private SizeCompatUIController.SizeCompatUICallback mCallback;
     @Mock private ShellTaskOrganizer.TaskListener mTaskListener;
     @Mock private DisplayLayout mDisplayLayout;
     @Mock private SizeCompatRestartButton mButton;
@@ -80,8 +76,9 @@ public class SizeCompatUILayoutTest extends ShellTestCase {
         MockitoAnnotations.initMocks(this);
         mTaskConfig = new Configuration();
 
-        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mContext, new Configuration(),
-                TASK_ID, mActivityToken, mTaskListener, mDisplayLayout, false /* hasShownHint*/);
+        mLayout = new SizeCompatUILayout(mSyncTransactionQueue, mCallback, mContext,
+                new Configuration(), TASK_ID, mTaskListener, mDisplayLayout,
+                false /* hasShownHint */);
 
         spyOn(mLayout);
         spyOn(mLayout.mButtonWindowManager);
@@ -145,7 +142,7 @@ public class SizeCompatUILayoutTest extends ShellTestCase {
 
         // No diff
         clearInvocations(mLayout);
-        mLayout.updateSizeCompatInfo(mTaskConfig, mActivityToken, mTaskListener,
+        mLayout.updateSizeCompatInfo(mTaskConfig, mTaskListener,
                 false /* isImeShowing */);
 
         verify(mLayout, never()).updateButtonSurfacePosition();
@@ -156,7 +153,7 @@ public class SizeCompatUILayoutTest extends ShellTestCase {
         clearInvocations(mLayout);
         final ShellTaskOrganizer.TaskListener newTaskListener = mock(
                 ShellTaskOrganizer.TaskListener.class);
-        mLayout.updateSizeCompatInfo(mTaskConfig, mActivityToken, newTaskListener,
+        mLayout.updateSizeCompatInfo(mTaskConfig, newTaskListener,
                 false /* isImeShowing */);
 
         verify(mLayout).release();
@@ -166,7 +163,7 @@ public class SizeCompatUILayoutTest extends ShellTestCase {
         clearInvocations(mLayout);
         final Configuration newTaskConfiguration = new Configuration();
         newTaskConfiguration.windowConfiguration.setBounds(new Rect(0, 1000, 0, 2000));
-        mLayout.updateSizeCompatInfo(newTaskConfiguration, mActivityToken, newTaskListener,
+        mLayout.updateSizeCompatInfo(newTaskConfiguration, newTaskListener,
                 false /* isImeShowing */);
 
         verify(mLayout).updateButtonSurfacePosition();
@@ -228,12 +225,9 @@ public class SizeCompatUILayoutTest extends ShellTestCase {
 
     @Test
     public void testOnRestartButtonClicked() {
-        spyOn(ActivityClient.getInstance());
-        doNothing().when(ActivityClient.getInstance()).restartActivityProcessIfVisible(any());
-
         mLayout.onRestartButtonClicked();
 
-        verify(ActivityClient.getInstance()).restartActivityProcessIfVisible(mActivityToken);
+        verify(mCallback).onSizeCompatRestartButtonClicked(TASK_ID);
     }
 
     @Test
diff --git a/services/core/java/com/android/server/wm/ActivityClientController.java b/services/core/java/com/android/server/wm/ActivityClientController.java
index 26f475ec8b29..3a4faf73bfe1 100644
--- a/services/core/java/com/android/server/wm/ActivityClientController.java
+++ b/services/core/java/com/android/server/wm/ActivityClientController.java
@@ -1069,8 +1069,7 @@ class ActivityClientController extends IActivityClientController.Stub {
         }
     }
 
-    @Override
-    public void restartActivityProcessIfVisible(IBinder token) {
+    void restartActivityProcessIfVisible(IBinder token) {
         ActivityTaskManagerService.enforceTaskPermission("restartActivityProcess");
         final long callingId = Binder.clearCallingIdentity();
         try {
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 325f10f65af2..9608a5fec2b3 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -4114,9 +4114,6 @@ class Task extends WindowContainer<WindowContainer> {
         info.topActivityInfo = mReuseActivitiesReport.top != null
                 ? mReuseActivitiesReport.top.info
                 : null;
-        info.topActivityToken = mReuseActivitiesReport.top != null
-                ? mReuseActivitiesReport.top.appToken
-                : null;
         // Whether the direct top activity is in size compat mode on foreground.
         info.topActivityInSizeCompat = mReuseActivitiesReport.top != null
                 && mReuseActivitiesReport.top.getOrganizedTask() == this
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 3a2ca80f2e12..f0318db0326a 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -922,6 +922,34 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
         }
     }
 
+    @Override
+    public void restartTaskTopActivityProcessIfVisible(WindowContainerToken token) {
+        enforceTaskPermission("restartTopActivityProcessIfVisible()");
+        final long origId = Binder.clearCallingIdentity();
+        try {
+            synchronized (mGlobalLock) {
+                final WindowContainer wc = WindowContainer.fromBinder(token.asBinder());
+                if (wc == null) {
+                    Slog.w(TAG, "Could not resolve window from token");
+                    return;
+                }
+                final Task task = wc.asTask();
+                if (task == null) {
+                    Slog.w(TAG, "Could not resolve task from token");
+                    return;
+                }
+                ProtoLog.v(WM_DEBUG_WINDOW_ORGANIZER,
+                        "Restart top activity process of Task taskId=%d", task.mTaskId);
+                final ActivityRecord activity = task.getTopNonFinishingActivity();
+                if (activity != null) {
+                    activity.restartProcessIfVisible();
+                }
+            }
+        } finally {
+            Binder.restoreCallingIdentity(origId);
+        }
+    }
+
     public boolean handleInterceptBackPressedOnTaskRoot(Task task) {
         if (task == null || !task.isOrganized()
                 || !mInterceptBackPressedOnRootTasks.contains(task.mTaskId)) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java b/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
index 4872ec511ccc..e1c61a747ea3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
@@ -595,7 +595,6 @@ public class SizeCompatTests extends WindowTestsBase {
         verify(mTask).onSizeCompatActivityChanged();
         ActivityManager.RunningTaskInfo taskInfo = mTask.getTaskInfo();
 
-        assertEquals(mActivity.appToken, taskInfo.topActivityToken);
         assertTrue(taskInfo.topActivityInSizeCompat);
 
         // Make the activity resizable again by restarting it
@@ -611,7 +610,6 @@ public class SizeCompatTests extends WindowTestsBase {
         verify(mTask).onSizeCompatActivityChanged();
         taskInfo = mTask.getTaskInfo();
 
-        assertEquals(mActivity.appToken, taskInfo.topActivityToken);
         assertFalse(taskInfo.topActivityInSizeCompat);
     }
 
@@ -630,7 +628,6 @@ public class SizeCompatTests extends WindowTestsBase {
         verify(mTask).onSizeCompatActivityChanged();
         ActivityManager.RunningTaskInfo taskInfo = mTask.getTaskInfo();
 
-        assertEquals(mActivity.appToken, taskInfo.topActivityToken);
         assertTrue(taskInfo.topActivityInSizeCompat);
 
         // Create another Task to hold another size compat activity.
@@ -651,7 +648,6 @@ public class SizeCompatTests extends WindowTestsBase {
         verify(mTask, never()).onSizeCompatActivityChanged();
         taskInfo = secondTask.getTaskInfo();
 
-        assertEquals(secondActivity.appToken, taskInfo.topActivityToken);
         assertTrue(taskInfo.topActivityInSizeCompat);
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index a1b3159825fb..d6a8401f5b18 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -1229,7 +1229,6 @@ public class WindowOrganizerTests extends WindowTestsBase {
         verify(organizer).onTaskInfoChanged(infoCaptor.capture());
         RunningTaskInfo info = infoCaptor.getValue();
         assertEquals(rootTask.mTaskId, info.taskId);
-        assertEquals(activity.appToken, info.topActivityToken);
         assertTrue(info.topActivityInSizeCompat);
 
         // Ensure task info show top activity that is not in foreground as not in size compat.
@@ -1240,7 +1239,6 @@ public class WindowOrganizerTests extends WindowTestsBase {
         verify(organizer).onTaskInfoChanged(infoCaptor.capture());
         info = infoCaptor.getValue();
         assertEquals(rootTask.mTaskId, info.taskId);
-        assertEquals(activity.appToken, info.topActivityToken);
         assertFalse(info.topActivityInSizeCompat);
 
         // Ensure task info show non size compat top activity as not in size compat.
@@ -1252,7 +1250,6 @@ public class WindowOrganizerTests extends WindowTestsBase {
         verify(organizer).onTaskInfoChanged(infoCaptor.capture());
         info = infoCaptor.getValue();
         assertEquals(rootTask.mTaskId, info.taskId);
-        assertEquals(activity.appToken, info.topActivityToken);
         assertFalse(info.topActivityInSizeCompat);
     }