Merge "[Partial Screenshare] Set WindowContainerToken for task capture" into tm-qpr-dev

media/java/android/media/projection/IMediaProjection.aidl

@@ -17,7 +17,7 @@
 package android.media.projection;
 
 import android.media.projection.IMediaProjectionCallback;
-import android.window.WindowContainerToken;
+import android.os.IBinder;
 
 /** {@hide} */
 interface IMediaProjection {
@@ -31,14 +31,14 @@ interface IMediaProjection {
     void unregisterCallback(IMediaProjectionCallback callback);
 
     /**
-     * Returns the {@link android.window.WindowContainerToken} identifying the task to record, or
+     * Returns the {@link android.os.IBinder} identifying the task to record, or {@code null} if
-     * {@code null} if there is none.
+     * there is none.
      */
-    WindowContainerToken getTaskRecordingWindowContainerToken();
+    IBinder getLaunchCookie();
 
     /**
-     * Updates the {@link android.window.WindowContainerToken} identifying the task to record, or
+     * Updates the {@link android.os.IBinder} identifying the task to record, or {@code null} if
-     * {@code null} if there is none.
+     * there is none.
      */
-    void setTaskRecordingWindowContainerToken(in WindowContainerToken token);
+    void setLaunchCookie(in IBinder launchCookie);
 }

media/java/android/media/projection/MediaProjection.java

@@ -25,13 +25,13 @@ import android.hardware.display.DisplayManager;
 import android.hardware.display.VirtualDisplay;
 import android.hardware.display.VirtualDisplayConfig;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.util.ArrayMap;
 import android.util.Log;
 import android.view.ContentRecordingSession;
 import android.view.Surface;
-import android.window.WindowContainerToken;
 
 import java.util.Map;
 
@@ -172,18 +172,16 @@ public final class MediaProjection {
             @NonNull VirtualDisplayConfig.Builder virtualDisplayConfig,
             @Nullable VirtualDisplay.Callback callback, @Nullable Handler handler) {
         try {
-            final WindowContainerToken taskWindowContainerToken =
+            final IBinder launchCookie = mImpl.getLaunchCookie();
-                    mImpl.getTaskRecordingWindowContainerToken();
             Context windowContext = null;
             ContentRecordingSession session;
-            if (taskWindowContainerToken == null) {
+            if (launchCookie == null) {
                 windowContext = mContext.createWindowContext(mContext.getDisplayNoVerify(),
                         TYPE_APPLICATION, null /* options */);
                 session = ContentRecordingSession.createDisplaySession(
                         windowContext.getWindowContextToken());
             } else {
-                session = ContentRecordingSession.createTaskSession(
+                session = ContentRecordingSession.createTaskSession(launchCookie);
-                        taskWindowContainerToken.asBinder());
             }
             virtualDisplayConfig.setWindowManagerMirroring(true);
             final DisplayManager dm = mContext.getSystemService(DisplayManager.class);

services/core/java/com/android/server/media/projection/MediaProjectionManagerService.java

@@ -46,7 +46,6 @@ import android.os.UserHandle;
 import android.util.ArrayMap;
 import android.util.Slog;
 import android.view.ContentRecordingSession;
-import android.window.WindowContainerToken;
 
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.DumpUtils;
@@ -433,7 +432,7 @@ public final class MediaProjectionManagerService extends SystemService
         private IBinder mToken;
         private IBinder.DeathRecipient mDeathEater;
         private boolean mRestoreSystemAlertWindow;
-        private WindowContainerToken mTaskRecordingWindowContainerToken = null;
+        private IBinder mLaunchCookie = null;
 
         MediaProjection(int type, int uid, String packageName, int targetSdkVersion,
                 boolean isPrivileged) {
@@ -609,14 +608,13 @@ public final class MediaProjectionManagerService extends SystemService
         }
 
         @Override // Binder call
-        public void setTaskRecordingWindowContainerToken(WindowContainerToken token) {
+        public void setLaunchCookie(IBinder launchCookie) {
-            // TODO(b/221417940) set the task id to record from sysui, for the package chosen.
+            mLaunchCookie = launchCookie;
-            mTaskRecordingWindowContainerToken = token;
         }
 
         @Override // Binder call
-        public WindowContainerToken getTaskRecordingWindowContainerToken() {
+        public IBinder getLaunchCookie() {
-            return mTaskRecordingWindowContainerToken;
+            return mLaunchCookie;
         }
 
         public MediaProjectionInfo getProjectionInfo() {

services/core/java/com/android/server/wm/ContentRecordingController.java

@@ -63,6 +63,7 @@ final class ContentRecordingController {
      */
     void setContentRecordingSessionLocked(@Nullable ContentRecordingSession incomingSession,
             @NonNull WindowManagerService wmService) {
+        // TODO(b/219761722) handle a null session arriving due to task setup failing
         if (incomingSession != null && (!ContentRecordingSession.isValid(incomingSession)
                 || ContentRecordingSession.isSameDisplay(mSession, incomingSession))) {
             // Ignore an invalid session, or a session for the same display as currently recording.

services/core/java/com/android/server/wm/WindowManagerService.java

@@ -48,6 +48,7 @@ import static android.provider.Settings.Global.DEVELOPMENT_FORCE_DESKTOP_MODE_ON
 import static android.provider.Settings.Global.DEVELOPMENT_FORCE_RESIZABLE_ACTIVITIES;
 import static android.provider.Settings.Global.DEVELOPMENT_RENDER_SHADOWS_IN_COMPOSITOR;
 import static android.provider.Settings.Global.DEVELOPMENT_WM_DISPLAY_SETTINGS_PATH;
+import static android.view.ContentRecordingSession.RECORD_CONTENT_TASK;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.DISPLAY_IME_POLICY_FALLBACK_DISPLAY;
@@ -288,6 +289,7 @@ import android.view.displayhash.VerifiedDisplayHash;
 import android.window.ClientWindowFrames;
 import android.window.ITaskFpsCallback;
 import android.window.TaskSnapshot;
+import android.window.WindowContainerToken;
 
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
@@ -8281,6 +8283,26 @@ public class WindowManagerService extends IWindowManager.Stub
         @Override
         public void setContentRecordingSession(@Nullable ContentRecordingSession incomingSession) {
             synchronized (mGlobalLock) {
+                // Allow the controller to handle teardown or a non-task session.
+                if (incomingSession == null
+                        || incomingSession.getContentToRecord() != RECORD_CONTENT_TASK) {
+                    mContentRecordingController.setContentRecordingSessionLocked(incomingSession,
+                            WindowManagerService.this);
+                    return;
+                }
+                // For a task session, find the activity identified by the launch cookie.
+                final WindowContainerToken wct = getTaskWindowContainerTokenForLaunchCookie(
+                        incomingSession.getTokenToRecord());
+                if (wct == null) {
+                    Slog.w(TAG, "Handling a new recording session; unable to find the "
+                            + "WindowContainerToken");
+                    mContentRecordingController.setContentRecordingSessionLocked(null,
+                            WindowManagerService.this);
+                    return;
+                }
+                // Replace the launch cookie in the session details with the task's
+                // WindowContainerToken.
+                incomingSession.setTokenToRecord(wct.asBinder());
                 mContentRecordingController.setContentRecordingSessionLocked(incomingSession,
                         WindowManagerService.this);
             }
@@ -8538,6 +8560,38 @@ public class WindowManagerService extends IWindowManager.Stub
         mAtmService.setFocusedTask(task.mTaskId, touchedActivity);
     }
 
+    /**
+     * Retrieve the {@link WindowContainerToken} of the task that contains the activity started
+     * with the given launch cookie.
+     *
+     * @param launchCookie the launch cookie set on the {@link ActivityOptions} when starting an
+     *                     activity
+     * @return a token representing the task containing the activity started with the given launch
+     * cookie, or {@code null} if the token couldn't be found.
+     */
+    @VisibleForTesting
+    @Nullable
+    WindowContainerToken getTaskWindowContainerTokenForLaunchCookie(@NonNull IBinder launchCookie) {
+        // Find the activity identified by the launch cookie.
+        final ActivityRecord targetActivity = mRoot.getActivity(
+                activity -> activity.mLaunchCookie == launchCookie);
+        if (targetActivity == null) {
+            Slog.w(TAG, "Unable to find the activity for this launch cookie");
+            return null;
+        }
+        if (targetActivity.getTask() == null) {
+            Slog.w(TAG, "Unable to find the task for this launch cookie");
+            return null;
+        }
+        WindowContainerToken taskWindowContainerToken =
+                targetActivity.getTask().mRemoteToken.toWindowContainerToken();
+        if (taskWindowContainerToken == null) {
+            Slog.w(TAG, "Unable to find the WindowContainerToken for " + targetActivity.getName());
+            return null;
+        }
+        return taskWindowContainerToken;
+    }
+
     /**
      * You need ALLOW_SLIPPERY_TOUCHES permission to be able to set FLAG_SLIPPERY.
      */

services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java

@@ -61,6 +61,7 @@ import android.view.InsetsState;
 import android.view.InsetsVisibilities;
 import android.view.View;
 import android.view.WindowManager;
+import android.window.WindowContainerToken;
 
 import androidx.test.filters.SmallTest;
 
@@ -317,4 +318,76 @@ public class WindowManagerServiceTests extends WindowTestsBase {
         verify(mWm.mInputManager).setInTouchMode(
                 !currentTouchMode, callingPid, callingUid, /* hasPermission= */ false);
     }
+
+    @Test
+    public void testGetTaskWindowContainerTokenForLaunchCookie_nullCookie() {
+        WindowContainerToken wct = mWm.getTaskWindowContainerTokenForLaunchCookie(null);
+        assertThat(wct).isNull();
+    }
+
+    @Test
+    public void testGetTaskWindowContainerTokenForLaunchCookie_invalidCookie() {
+        Binder cookie = new Binder("test cookie");
+        WindowContainerToken wct = mWm.getTaskWindowContainerTokenForLaunchCookie(cookie);
+        assertThat(wct).isNull();
+
+        final ActivityRecord testActivity = new ActivityBuilder(mAtm)
+                .setCreateTask(true)
+                .build();
+
+        wct = mWm.getTaskWindowContainerTokenForLaunchCookie(cookie);
+        assertThat(wct).isNull();
+    }
+
+    @Test
+    public void testGetTaskWindowContainerTokenForLaunchCookie_validCookie() {
+        final Binder cookie = new Binder("ginger cookie");
+        final WindowContainerToken launchRootTask = mock(WindowContainerToken.class);
+        setupActivityWithLaunchCookie(cookie, launchRootTask);
+
+        WindowContainerToken wct = mWm.getTaskWindowContainerTokenForLaunchCookie(cookie);
+        assertThat(wct).isEqualTo(launchRootTask);
+    }
+
+    @Test
+    public void testGetTaskWindowContainerTokenForLaunchCookie_multipleCookies() {
+        final Binder cookie1 = new Binder("ginger cookie");
+        final WindowContainerToken launchRootTask1 = mock(WindowContainerToken.class);
+        setupActivityWithLaunchCookie(cookie1, launchRootTask1);
+
+        setupActivityWithLaunchCookie(new Binder("choc chip cookie"),
+                mock(WindowContainerToken.class));
+
+        setupActivityWithLaunchCookie(new Binder("peanut butter cookie"),
+                mock(WindowContainerToken.class));
+
+        WindowContainerToken wct = mWm.getTaskWindowContainerTokenForLaunchCookie(cookie1);
+        assertThat(wct).isEqualTo(launchRootTask1);
+    }
+
+    @Test
+    public void testGetTaskWindowContainerTokenForLaunchCookie_multipleCookies_noneValid() {
+        setupActivityWithLaunchCookie(new Binder("ginger cookie"),
+                mock(WindowContainerToken.class));
+
+        setupActivityWithLaunchCookie(new Binder("choc chip cookie"),
+                mock(WindowContainerToken.class));
+
+        setupActivityWithLaunchCookie(new Binder("peanut butter cookie"),
+                mock(WindowContainerToken.class));
+
+        WindowContainerToken wct = mWm.getTaskWindowContainerTokenForLaunchCookie(
+                new Binder("some other cookie"));
+        assertThat(wct).isNull();
+    }
+
+    private void setupActivityWithLaunchCookie(IBinder launchCookie, WindowContainerToken wct) {
+        final WindowContainer.RemoteToken remoteToken = mock(WindowContainer.RemoteToken.class);
+        when(remoteToken.toWindowContainerToken()).thenReturn(wct);
+        final ActivityRecord testActivity = new ActivityBuilder(mAtm)
+                .setCreateTask(true)
+                .build();
+        testActivity.mLaunchCookie = launchCookie;
+        testActivity.getTask().mRemoteToken = remoteToken;
+    }
 }