Merge "Track active AttributionSources, and finish chains on death" into sc-dev

services/core/java/com/android/server/pm/permission/PermissionManagerService.java

@@ -172,6 +172,7 @@ import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Set;
 import java.util.WeakHashMap;
@@ -180,6 +181,7 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiFunction;
 
@@ -219,6 +221,10 @@ public class PermissionManagerService extends IPermissionManager.Stub {
     /** If the permission of the value is granted, so is the key */
     private static final Map<String, String> FULLER_PERMISSION_MAP = new HashMap<>();
 
+    /** Map of IBinder -> Running AttributionSource */
+    private static final ConcurrentHashMap<IBinder, RegisteredAttribution>
+            sRunningAttributionSources = new ConcurrentHashMap<>();
+
     static {
         FULLER_PERMISSION_MAP.put(Manifest.permission.ACCESS_COARSE_LOCATION,
                 Manifest.permission.ACCESS_FINE_LOCATION);
@@ -5555,8 +5561,15 @@ public class PermissionManagerService extends IPermissionManager.Stub {
 
         @Override
         public void finishDataDelivery(int op,
+                @NonNull AttributionSourceState attributionSourceState, boolean fromDataSource) {
+            finishDataDelivery(mContext, op, attributionSourceState,
+                    fromDataSource);
+        }
+
+        private static void finishDataDelivery(Context context, int op,
                 @NonNull AttributionSourceState attributionSourceState, boolean fromDatasource) {
             Objects.requireNonNull(attributionSourceState);
+            AppOpsManager appOpsManager = context.getSystemService(AppOpsManager.class);
 
             if (op == AppOpsManager.OP_NONE) {
                 return;
@@ -5573,7 +5586,7 @@ public class PermissionManagerService extends IPermissionManager.Stub {
                 // If the call is from a datasource we need to vet only the chain before it. This
                 // way we can avoid the datasource creating an attribution context for every call.
                 if (!(fromDatasource && current.asState() == attributionSourceState)
-                        && next != null && !current.isTrusted(mContext)) {
+                        && next != null && !current.isTrusted(context)) {
                     return;
                 }
 
@@ -5589,20 +5602,20 @@ public class PermissionManagerService extends IPermissionManager.Stub {
                         ? current : next;
 
                 if (selfAccess) {
-                    final String resolvedPackageName = resolvePackageName(mContext, accessorSource);
+                    final String resolvedPackageName = resolvePackageName(context, accessorSource);
                     if (resolvedPackageName == null) {
                         return;
                     }
-                    mAppOpsManager.finishOp(accessorSource.getToken(), op,
+                    appOpsManager.finishOp(accessorSource.getToken(), op,
                             accessorSource.getUid(), resolvedPackageName,
                             accessorSource.getAttributionTag());
                 } else {
                     final AttributionSource resolvedAttributionSource =
-                            resolveAttributionSource(mContext, accessorSource);
+                            resolveAttributionSource(context, accessorSource);
                     if (resolvedAttributionSource.getPackageName() == null) {
                         return;
                     }
-                    mAppOpsManager.finishProxyOp(AppOpsManager.opToPublicName(op),
+                    appOpsManager.finishProxyOp(AppOpsManager.opToPublicName(op),
                             resolvedAttributionSource, skipCurrentFinish);
                 }
 
@@ -5610,6 +5623,11 @@ public class PermissionManagerService extends IPermissionManager.Stub {
                     return;
                 }
 
+                RegisteredAttribution registered =
+                        sRunningAttributionSources.remove(current.getToken());
+                if (registered != null) {
+                    registered.unregister();
+                }
                 current = next;
             }
         }
@@ -5827,6 +5845,12 @@ public class PermissionManagerService extends IPermissionManager.Stub {
                     }
                 }
 
+                if (startDataDelivery) {
+                    RegisteredAttribution registered = new RegisteredAttribution(context, op,
+                            current, fromDatasource);
+                    sRunningAttributionSources.put(current.getToken(), registered);
+                }
+
                 if (next == null || next.getNext() == null) {
                     return PermissionChecker.PERMISSION_GRANTED;
                 }
@@ -6153,4 +6177,43 @@ public class PermissionManagerService extends IPermissionManager.Stub {
                     attributionSource));
         }
     }
+
+    private static final class RegisteredAttribution {
+        private final DeathRecipient mDeathRecipient;
+        private final IBinder mToken;
+        private final AtomicBoolean mFinished;
+
+        RegisteredAttribution(Context context, int op, AttributionSource source,
+                boolean fromDatasource) {
+            mFinished = new AtomicBoolean(false);
+            mDeathRecipient = () -> {
+                if (unregister()) {
+                    PermissionCheckerService
+                            .finishDataDelivery(context, op, source.asState(), fromDatasource);
+                }
+            };
+            mToken = source.getToken();
+            if (mToken != null) {
+                try {
+                    mToken.linkToDeath(mDeathRecipient, 0);
+                } catch (RemoteException e) {
+                    mDeathRecipient.binderDied();
+                }
+            }
+        }
+
+        public boolean unregister() {
+            if (mFinished.compareAndSet(false, true)) {
+                try {
+                    if (mToken != null) {
+                        mToken.unlinkToDeath(mDeathRecipient, 0);
+                    }
+                } catch (NoSuchElementException e) {
+                    // do nothing
+                }
+                return true;
+            }
+            return false;
+        }
+    }
 }