Fetch context URL from the top-most app, send with share if available.

- Add contextUrl to ScreenshotData.
- Fork AssistContentRequester from launcher, adapt slightly for sysui
  dagger.
- Request the assist content (async) if the task ID is available.
- Set the URL as EXTRA_TEXT if available. Don't set the subject in the
  share if we have EXTRA_TEXT, as they seem to conflict with each other
  in some apps.
- Slight refactor in ActionIntentCreator to make the various ways to
  create share intents clearer.

Flag: All code changes (beyond variable existence) are protected by the
      SCREENSHOT_METADATA flag, currently disabled. The work profile
      screenshots flag is also required for this to work properly (in
      teamfood).
Bug: 242791070
Test: atest com.android.systemui.screenshot
Test: Verify that context URL from chrome passes all the way through
      when sharing with Gmail.
Change-Id: I39c3503147631f4ef79088f23c201bbad7405d16
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ActionIntentCreator.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ActionIntentCreator.kt
index 017e57f..310baaf 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ActionIntentCreator.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ActionIntentCreator.kt
@@ -25,8 +25,22 @@
 import com.android.systemui.R
 
 object ActionIntentCreator {
+    /** @return a chooser intent to share the given URI. */
+    fun createShareIntent(uri: Uri) = createShareIntent(uri, null, null)
+
     /** @return a chooser intent to share the given URI with the optional provided subject. */
-    fun createShareIntent(uri: Uri, subject: String?): Intent {
+    fun createShareIntentWithSubject(uri: Uri, subject: String?) =
+        createShareIntent(uri, subject = subject)
+
+    /** @return a chooser intent to share the given URI with the optional provided extra text. */
+    fun createShareIntentWithExtraText(uri: Uri, extraText: String?) =
+        createShareIntent(uri, extraText = extraText)
+
+    private fun createShareIntent(
+        uri: Uri,
+        subject: String? = null,
+        extraText: String? = null
+    ): Intent {
         // Create a share intent, this will always go through the chooser activity first
         // which should not trigger auto-enter PiP
         val sharingIntent =
@@ -43,6 +57,7 @@
                     )
 
                 putExtra(Intent.EXTRA_SUBJECT, subject)
+                putExtra(Intent.EXTRA_TEXT, extraText)
                 addFlags(Intent.FLAG_GRANT_READ_URI_PERMISSION)
                 addFlags(Intent.FLAG_GRANT_WRITE_URI_PERMISSION)
             }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/AssistContentRequester.java b/packages/SystemUI/src/com/android/systemui/screenshot/AssistContentRequester.java
new file mode 100644
index 0000000..ab8fc65
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/AssistContentRequester.java
@@ -0,0 +1,147 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot;
+import android.app.ActivityTaskManager;
+import android.app.IActivityTaskManager;
+import android.app.IAssistDataReceiver;
+import android.app.assist.AssistContent;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.os.Bundle;
+import android.os.RemoteException;
+import android.util.Log;
+
+import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.dagger.qualifiers.Background;
+import com.android.systemui.dagger.qualifiers.Main;
+
+import java.lang.ref.WeakReference;
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+import java.util.concurrent.Executor;
+
+import javax.inject.Inject;
+
+/**
+ * Can be used to request the AssistContent from a provided task id, useful for getting the web uri
+ * if provided from the task.
+ *
+ * Forked from
+ * packages/apps/Launcher3/quickstep/src/com/android/quickstep/util/AssistContentRequester.java
+ */
+@SysUISingleton
+public class AssistContentRequester {
+    private static final String TAG = "AssistContentRequester";
+    private static final String ASSIST_KEY_CONTENT = "content";
+
+    /** For receiving content, called on the main thread. */
+    public interface Callback {
+        /**
+         * Called when the {@link android.app.assist.AssistContent} of the requested task is
+         * available.
+         **/
+        void onAssistContentAvailable(AssistContent assistContent);
+    }
+
+    private final IActivityTaskManager mActivityTaskManager;
+    private final String mPackageName;
+    private final Executor mCallbackExecutor;
+    private final Executor mSystemInteractionExecutor;
+    private final String mAttributionTag;
+
+    // If system loses the callback, our internal cache of original callback will also get cleared.
+    private final Map<Object, Callback> mPendingCallbacks =
+            Collections.synchronizedMap(new WeakHashMap<>());
+
+    @Inject
+    public AssistContentRequester(Context context, @Main Executor mainExecutor,
+            @Background Executor bgExecutor) {
+        mActivityTaskManager = ActivityTaskManager.getService();
+        mPackageName = context.getApplicationContext().getPackageName();
+        mCallbackExecutor = mainExecutor;
+        mSystemInteractionExecutor = bgExecutor;
+        mAttributionTag = context.getAttributionTag();
+    }
+
+    /**
+     * Request the {@link AssistContent} from the task with the provided id.
+     *
+     * @param taskId to query for the content.
+     * @param callback to call when the content is available, called on the main thread.
+     */
+    public void requestAssistContent(final int taskId, final Callback callback) {
+        // ActivityTaskManager interaction here is synchronous, so call off the main thread.
+        mSystemInteractionExecutor.execute(() -> {
+            try {
+                mActivityTaskManager.requestAssistDataForTask(
+                        new AssistDataReceiver(callback, this), taskId, mPackageName,
+                        mAttributionTag);
+            } catch (RemoteException e) {
+                Log.e(TAG, "Requesting assist content failed for task: " + taskId, e);
+            }
+        });
+    }
+
+    private void executeOnMainExecutor(Runnable callback) {
+        mCallbackExecutor.execute(callback);
+    }
+
+    private static final class AssistDataReceiver extends IAssistDataReceiver.Stub {
+
+        // The AssistDataReceiver binder callback object is passed to a system server, that may
+        // keep hold of it for longer than the lifetime of the AssistContentRequester object,
+        // potentially causing a memory leak. In the callback passed to the system server, only
+        // keep a weak reference to the parent object and lookup its callback if it still exists.
+        private final WeakReference<AssistContentRequester> mParentRef;
+        private final Object mCallbackKey = new Object();
+
+        AssistDataReceiver(Callback callback, AssistContentRequester parent) {
+            parent.mPendingCallbacks.put(mCallbackKey, callback);
+            mParentRef = new WeakReference<>(parent);
+        }
+
+        @Override
+        public void onHandleAssistData(Bundle data) {
+            if (data == null) {
+                return;
+            }
+
+            final AssistContent content = data.getParcelable(ASSIST_KEY_CONTENT);
+            if (content == null) {
+                Log.e(TAG, "Received AssistData, but no AssistContent found");
+                return;
+            }
+
+            AssistContentRequester requester = mParentRef.get();
+            if (requester != null) {
+                Callback callback = requester.mPendingCallbacks.get(mCallbackKey);
+                if (callback != null) {
+                    requester.executeOnMainExecutor(
+                            () -> callback.onAssistContentAvailable(content));
+                } else {
+                    Log.d(TAG, "Callback received after calling UI was disposed of");
+                }
+            } else {
+                Log.d(TAG, "Callback received after Requester was collected");
+            }
+        }
+
+        @Override
+        public void onHandleAssistScreenshot(Bitmap screenshot) {}
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotActivity.java b/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotActivity.java
index d64b33b..ca8e101 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotActivity.java
@@ -366,7 +366,7 @@
 
     private void doShare(Uri uri) {
         if (mFeatureFlags.isEnabled(Flags.SCREENSHOT_WORK_PROFILE_POLICY)) {
-            Intent shareIntent = ActionIntentCreator.INSTANCE.createShareIntent(uri, null);
+            Intent shareIntent = ActionIntentCreator.INSTANCE.createShareIntent(uri);
             mActionExecutor.launchIntentAsync(shareIntent, null,
                     mScreenshotUserHandle.getIdentifier(), false);
         } else {
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index 02719e3fc..adff6e1 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -44,6 +44,7 @@
 import android.app.ExitTransitionCoordinator.ExitTransitionCallbacks;
 import android.app.ICompatCameraControlCallback;
 import android.app.Notification;
+import android.app.assist.AssistContent;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Context;
@@ -281,6 +282,7 @@
     private final ActionIntentExecutor mActionExecutor;
     private final UserManager mUserManager;
     private final WorkProfileMessageController mWorkProfileMessageController;
+    private final AssistContentRequester mAssistContentRequester;
 
     private final OnBackInvokedCallback mOnBackInvokedCallback = () -> {
         if (DEBUG_INPUT) {
@@ -328,7 +330,8 @@
             ScreenshotNotificationSmartActionsProvider screenshotNotificationSmartActionsProvider,
             ActionIntentExecutor actionExecutor,
             UserManager userManager,
-            WorkProfileMessageController workProfileMessageController
+            WorkProfileMessageController workProfileMessageController,
+            AssistContentRequester assistContentRequester
     ) {
         mScreenshotSmartActions = screenshotSmartActions;
         mNotificationsController = screenshotNotificationsController;
@@ -361,6 +364,7 @@
         mActionExecutor = actionExecutor;
         mUserManager = userManager;
         mWorkProfileMessageController = workProfileMessageController;
+        mAssistContentRequester = assistContentRequester;
 
         mAccessibilityManager = AccessibilityManager.getInstance(mContext);
 
@@ -466,7 +470,18 @@
                     mContext.getDrawable(R.drawable.overlay_badge_background),
                     screenshot.getUserHandle()));
         }
-        mScreenshotView.setScreenshot(mScreenBitmap, screenshot.getInsets());
+        mScreenshotView.setScreenshot(screenshot);
+
+        if (screenshot.getTaskId() >= 0) {
+            mAssistContentRequester.requestAssistContent(screenshot.getTaskId(),
+                    new AssistContentRequester.Callback() {
+                        @Override
+                        public void onAssistContentAvailable(AssistContent assistContent) {
+                            screenshot.setContextUrl(assistContent.getWebUri());
+                        }
+                    });
+        }
+
         if (DEBUG_WINDOW) {
             Log.d(TAG, "setContentView: " + mScreenshotView);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
index 1d1b968..c43e4b4 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
@@ -4,6 +4,7 @@
 import android.graphics.Bitmap
 import android.graphics.Insets
 import android.graphics.Rect
+import android.net.Uri
 import android.os.UserHandle
 import android.view.WindowManager.ScreenshotSource
 import android.view.WindowManager.ScreenshotType
@@ -21,6 +22,8 @@
     var taskId: Int,
     var insets: Insets,
     var bitmap: Bitmap?,
+    /** App-provided URL representing the content the user was looking at in the screenshot. */
+    var contextUrl: Uri? = null,
 ) {
     val packageNameString: String
         get() = if (topComponent == null) "" else topComponent!!.packageName
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
index 7c013a8..bd4ea11 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
@@ -39,6 +39,7 @@
 import android.app.Notification;
 import android.app.PendingIntent;
 import android.content.Context;
+import android.content.Intent;
 import android.content.res.ColorStateList;
 import android.content.res.Resources;
 import android.graphics.Bitmap;
@@ -168,6 +169,8 @@
 
     private final ArrayList<OverlayActionChip> mSmartChips = new ArrayList<>();
     private PendingInteraction mPendingInteraction;
+    // Should only be set/used if the SCREENSHOT_METADATA flag is set.
+    private ScreenshotData mScreenshotData;
 
     private final InteractionJankMonitor mInteractionJankMonitor;
     private long mDefaultTimeoutOfTimeoutHandler;
@@ -477,6 +480,13 @@
         mScreenshotPreview.setImageDrawable(createScreenDrawable(mResources, bitmap, screenInsets));
     }
 
+    void setScreenshot(ScreenshotData screenshot) {
+        mScreenshotData = screenshot;
+        setScreenshot(screenshot.getBitmap(), screenshot.getInsets());
+        mScreenshotPreview.setImageDrawable(createScreenDrawable(mResources, screenshot.getBitmap(),
+                screenshot.getInsets()));
+    }
+
     void setPackageName(String packageName) {
         mPackageName = packageName;
     }
@@ -815,9 +825,17 @@
             mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_SHARE_TAPPED, 0, mPackageName);
             if (mFlags.isEnabled(Flags.SCREENSHOT_WORK_PROFILE_POLICY)) {
                 prepareSharedTransition();
-                mActionExecutor.launchIntentAsync(
-                        ActionIntentCreator.INSTANCE.createShareIntent(
-                                imageData.uri, imageData.subject),
+
+                Intent shareIntent;
+                if (mFlags.isEnabled(Flags.SCREENSHOT_METADATA) && mScreenshotData != null
+                        && mScreenshotData.getContextUrl() != null) {
+                    shareIntent = ActionIntentCreator.INSTANCE.createShareIntentWithExtraText(
+                            imageData.uri, mScreenshotData.getContextUrl().toString());
+                } else {
+                    shareIntent = ActionIntentCreator.INSTANCE.createShareIntentWithSubject(
+                            imageData.uri, imageData.subject);
+                }
+                mActionExecutor.launchIntentAsync(shareIntent,
                         imageData.shareTransition.get().bundle,
                         imageData.owner.getIdentifier(), false);
             } else {
@@ -1119,6 +1137,7 @@
         mQuickShareChip = null;
         setAlpha(1);
         mScreenshotStatic.setAlpha(1);
+        mScreenshotData = null;
     }
 
     private void startSharedTransition(ActionTransition transition) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ActionIntentCreatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ActionIntentCreatorTest.kt
index b6a595b..7ba2cf7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ActionIntentCreatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ActionIntentCreatorTest.kt
@@ -35,9 +35,33 @@
     @Test
     fun testCreateShareIntent() {
         val uri = Uri.parse("content://fake")
+
+        val output = ActionIntentCreator.createShareIntent(uri)
+
+        assertThat(output.action).isEqualTo(Intent.ACTION_CHOOSER)
+        assertFlagsSet(
+            Intent.FLAG_ACTIVITY_NEW_TASK or
+                Intent.FLAG_ACTIVITY_CLEAR_TASK or
+                Intent.FLAG_GRANT_READ_URI_PERMISSION,
+            output.flags
+        )
+
+        val wrappedIntent = output.getParcelableExtra(Intent.EXTRA_INTENT, Intent::class.java)
+        assertThat(wrappedIntent?.action).isEqualTo(Intent.ACTION_SEND)
+        assertThat(wrappedIntent?.data).isEqualTo(uri)
+        assertThat(wrappedIntent?.type).isEqualTo("image/png")
+        assertThat(wrappedIntent?.getStringExtra(Intent.EXTRA_SUBJECT)).isNull()
+        assertThat(wrappedIntent?.getStringExtra(Intent.EXTRA_TEXT)).isNull()
+        assertThat(wrappedIntent?.getParcelableExtra(Intent.EXTRA_STREAM, Uri::class.java))
+            .isEqualTo(uri)
+    }
+
+    @Test
+    fun testCreateShareIntentWithSubject() {
+        val uri = Uri.parse("content://fake")
         val subject = "Example subject"
 
-        val output = ActionIntentCreator.createShareIntent(uri, subject)
+        val output = ActionIntentCreator.createShareIntentWithSubject(uri, subject)
 
         assertThat(output.action).isEqualTo(Intent.ACTION_CHOOSER)
         assertFlagsSet(
@@ -52,16 +76,34 @@
         assertThat(wrappedIntent?.data).isEqualTo(uri)
         assertThat(wrappedIntent?.type).isEqualTo("image/png")
         assertThat(wrappedIntent?.getStringExtra(Intent.EXTRA_SUBJECT)).isEqualTo(subject)
+        assertThat(wrappedIntent?.getStringExtra(Intent.EXTRA_TEXT)).isNull()
         assertThat(wrappedIntent?.getParcelableExtra(Intent.EXTRA_STREAM, Uri::class.java))
             .isEqualTo(uri)
     }
 
     @Test
-    fun testCreateShareIntent_noSubject() {
+    fun testCreateShareIntentWithExtraText() {
         val uri = Uri.parse("content://fake")
-        val output = ActionIntentCreator.createShareIntent(uri, null)
+        val extraText = "Extra text"
+
+        val output = ActionIntentCreator.createShareIntentWithExtraText(uri, extraText)
+
+        assertThat(output.action).isEqualTo(Intent.ACTION_CHOOSER)
+        assertFlagsSet(
+            Intent.FLAG_ACTIVITY_NEW_TASK or
+                Intent.FLAG_ACTIVITY_CLEAR_TASK or
+                Intent.FLAG_GRANT_READ_URI_PERMISSION,
+            output.flags
+        )
+
         val wrappedIntent = output.getParcelableExtra(Intent.EXTRA_INTENT, Intent::class.java)
+        assertThat(wrappedIntent?.action).isEqualTo(Intent.ACTION_SEND)
+        assertThat(wrappedIntent?.data).isEqualTo(uri)
+        assertThat(wrappedIntent?.type).isEqualTo("image/png")
         assertThat(wrappedIntent?.getStringExtra(Intent.EXTRA_SUBJECT)).isNull()
+        assertThat(wrappedIntent?.getStringExtra(Intent.EXTRA_TEXT)).isEqualTo(extraText)
+        assertThat(wrappedIntent?.getParcelableExtra(Intent.EXTRA_STREAM, Uri::class.java))
+            .isEqualTo(uri)
     }
 
     @Test