3 files changed, 33 insertions, 5 deletions

diff --git a/core/api/current.txt b/core/api/current.txt
index bcb21e5bd3d1..d26cf2e60121 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -10560,6 +10560,7 @@ package android.content {
   public final class ContextParams {
     method @Nullable public String getAttributionTag();
     method @Nullable public android.content.AttributionSource getNextAttributionSource();
+    method @NonNull public boolean shouldRegisterAttributionSource();
   }
 
   public static final class ContextParams.Builder {
@@ -10568,6 +10569,7 @@ package android.content {
     method @NonNull public android.content.ContextParams build();
     method @NonNull public android.content.ContextParams.Builder setAttributionTag(@Nullable String);
     method @NonNull public android.content.ContextParams.Builder setNextAttributionSource(@Nullable android.content.AttributionSource);
+    method @NonNull public android.content.ContextParams.Builder setShouldRegisterAttributionSource(boolean);
   }
 
   public class ContextWrapper extends android.content.Context {
diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index 5feafbed148c..a538247998e6 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -3455,20 +3455,20 @@ class ContextImpl extends Context {
         mOpPackageName = overrideOpPackageName != null ? overrideOpPackageName : opPackageName;
         mParams = Objects.requireNonNull(params);
         mAttributionSource = createAttributionSource(attributionTag, nextAttributionSource,
-                params.getRenouncedPermissions());
+                params.getRenouncedPermissions(), params.shouldRegisterAttributionSource());
         mContentResolver = new ApplicationContentResolver(this, mainThread);
     }
 
     private @NonNull AttributionSource createAttributionSource(@Nullable String attributionTag,
             @Nullable AttributionSource nextAttributionSource,
-            @Nullable Set<String> renouncedPermissions) {
+            @Nullable Set<String> renouncedPermissions, boolean shouldRegister) {
         AttributionSource attributionSource = new AttributionSource(Process.myUid(),
                 Process.myPid(), mOpPackageName, attributionTag,
                 (renouncedPermissions != null) ? renouncedPermissions.toArray(new String[0]) : null,
                 getDeviceId(), 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) {
+        if (nextAttributionSource != null || shouldRegister) {
             attributionSource = getSystemService(PermissionManager.class)
                     .registerAttributionSource(attributionSource);
         }
diff --git a/core/java/android/content/ContextParams.java b/core/java/android/content/ContextParams.java
index 5cc3a2420811..988a9c0ab9b3 100644
--- a/core/java/android/content/ContextParams.java
+++ b/core/java/android/content/ContextParams.java
@@ -50,17 +50,20 @@ public final class ContextParams {
     private final @Nullable String mAttributionTag;
     private final @Nullable AttributionSource mNext;
     private final @NonNull Set<String> mRenouncedPermissions;
+    private final boolean mShouldRegisterAttributionSource;
 
     /** {@hide} */
     public static final ContextParams EMPTY = new ContextParams.Builder().build();
 
     private ContextParams(@Nullable String attributionTag,
             @Nullable AttributionSource next,
-            @Nullable Set<String> renouncedPermissions) {
+            @Nullable Set<String> renouncedPermissions,
+            boolean shouldRegister) {
         mAttributionTag = attributionTag;
         mNext = next;
         mRenouncedPermissions = (renouncedPermissions != null)
                 ? renouncedPermissions : Collections.emptySet();
+        mShouldRegisterAttributionSource = shouldRegister;
     }
 
     /**
@@ -95,12 +98,22 @@ public final class ContextParams {
     }
 
     /**
+     * @return Whether the attribution source associated with the Context being created should be
+     * registered.
+     */
+    @NonNull
+    public boolean shouldRegisterAttributionSource() {
+        return mShouldRegisterAttributionSource;
+    }
+
+    /**
      * Builder for creating a {@link ContextParams}.
      */
     public static final class Builder {
         private @Nullable String mAttributionTag;
         private @NonNull Set<String> mRenouncedPermissions = Collections.emptySet();
         private @Nullable AttributionSource mNext;
+        private boolean mShouldRegisterAttributionSource;
 
         /**
          * Create a new builder.
@@ -159,6 +172,19 @@ public final class ContextParams {
         }
 
         /**
+         * Sets whether the attribution source associated with the context created from these params
+         * should be registered.
+         *
+         * @param shouldRegister Whether the attribution source associated with the Context being
+         *                       created should be registered.
+         */
+        @NonNull
+        public Builder setShouldRegisterAttributionSource(boolean shouldRegister) {
+            mShouldRegisterAttributionSource = shouldRegister;
+            return this;
+        }
+
+        /**
          * Sets permissions which have been voluntarily "renounced" by the
          * calling app.
          * <p>
@@ -205,7 +231,7 @@ public final class ContextParams {
         @NonNull
         public ContextParams build() {
             return new ContextParams(mAttributionTag, mNext,
-                    mRenouncedPermissions);
+                    mRenouncedPermissions, mShouldRegisterAttributionSource);
         }
     }
 }