Keep the deviceId in AttributionSource in sync with Context

Also re-registers the AttributionSource if needed

Change-Id: I81bbba6a041cc87aef77cfa1f21c1016b35a10a2
Fix: 333272850
Test: presubmit
(cherry picked from commit 02f4520f2933e2d54e99a178e61619fa1b3d7439)

1 files changed, 20 insertions, 3 deletions

diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index af56cb4d55b2..c15105a15c7a 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -19,6 +19,7 @@ package android.app;
 import static android.content.pm.PackageManager.PERMISSION_DENIED;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.os.StrictMode.vmIncorrectContextUseEnabled;
+import static android.permission.flags.Flags.shouldRegisterAttributionSource;
 import static android.view.WindowManager.LayoutParams.WindowType;
 
 import android.annotation.CallbackExecutor;
@@ -3094,7 +3095,8 @@ class ContextImpl extends Context {
             int deviceId = vdm.getDeviceIdForDisplayId(displayId);
             if (deviceId != mDeviceId) {
                 mDeviceId = deviceId;
-                mAttributionSource = mAttributionSource.withDeviceId(mDeviceId);
+                mAttributionSource =
+                        createAttributionSourceWithDeviceId(mAttributionSource, mDeviceId);
                 notifyOnDeviceChangedListeners(mDeviceId);
             }
         }
@@ -3117,6 +3119,7 @@ class ContextImpl extends Context {
 
         if (mDeviceId != updatedDeviceId) {
             mDeviceId = updatedDeviceId;
+            mAttributionSource = createAttributionSourceWithDeviceId(mAttributionSource, mDeviceId);
             notifyOnDeviceChangedListeners(updatedDeviceId);
         }
     }
@@ -3485,8 +3488,22 @@ class ContextImpl extends Context {
                 deviceId, nextAttributionSource);
         // If we want to access protected data on behalf of another app we need to
         // tell the OS that we opt in to participate in the attribution chain.
-        if (nextAttributionSource != null || shouldRegister) {
-            attributionSource = getSystemService(PermissionManager.class)
+        return registerAttributionSourceIfNeeded(attributionSource, shouldRegister);
+    }
+
+    private @NonNull AttributionSource createAttributionSourceWithDeviceId(
+            @NonNull AttributionSource oldSource, int deviceId) {
+        boolean shouldRegister = false;
+        if (shouldRegisterAttributionSource()) {
+            shouldRegister = mParams.shouldRegisterAttributionSource();
+        }
+        return registerAttributionSourceIfNeeded(oldSource.withDeviceId(deviceId), shouldRegister);
+    }
+
+    private @NonNull AttributionSource registerAttributionSourceIfNeeded(
+            @NonNull AttributionSource attributionSource, boolean shouldRegister) {
+        if (shouldRegister || attributionSource.getNext() != null) {
+            return getSystemService(PermissionManager.class)
                     .registerAttributionSource(attributionSource);
         }
         return attributionSource;