Creating PendingIntentRef

Update #setDataFetchOperation to avoid using intentsender.

Bug: 146074295
Test: Ran GTS Tests
Change-Id: I7df5c6441725aa4e46fac18925664c9455f50fb9
diff --git a/apex/statsd/aidl/Android.bp b/apex/statsd/aidl/Android.bp
index aed6ad9..f8325d4 100644
--- a/apex/statsd/aidl/Android.bp
+++ b/apex/statsd/aidl/Android.bp
@@ -18,6 +18,7 @@
 filegroup {
     name: "statsd_aidl",
     srcs: [
+        "android/os/IPendingIntentRef.aidl",
         "android/os/IPullAtomCallback.aidl",
         "android/os/IPullAtomResultReceiver.aidl",
         "android/os/IStatsCompanionService.aidl",
diff --git a/apex/statsd/aidl/android/os/IPendingIntentRef.aidl b/apex/statsd/aidl/android/os/IPendingIntentRef.aidl
new file mode 100644
index 0000000..6b9e467
--- /dev/null
+++ b/apex/statsd/aidl/android/os/IPendingIntentRef.aidl
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2019 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 android.os;
+
+import android.os.StatsDimensionsValue;
+
+/**
+  * Binder interface to hold a PendingIntent for StatsCompanionService.
+  * {@hide}
+  */
+interface IPendingIntentRef {
+
+    /**
+     * Sends a broadcast to the specified PendingIntent that it should getData now.
+     * This should be only called from StatsCompanionService.
+     */
+     oneway void sendDataBroadcast(long lastReportTimeNs);
+
+    /**
+     * Send a broadcast to the specified PendingIntent notifying it that the list of active configs
+     * has changed. This should be only called from StatsCompanionService.
+     */
+     oneway void sendActiveConfigsChangedBroadcast(in long[] configIds);
+
+     /**
+      * Send a broadcast to the specified PendingIntent, along with the other information
+      * specified. This should only be called from StatsCompanionService.
+      */
+     oneway void sendSubscriberBroadcast(long configUid, long configId, long subscriptionId,
+                                         long subscriptionRuleId, in String[] cookies,
+                                         in StatsDimensionsValue dimensionsValue);
+}
\ No newline at end of file
diff --git a/apex/statsd/aidl/android/os/IStatsCompanionService.aidl b/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
index 5a6118e..ec3a226 100644
--- a/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
+++ b/apex/statsd/aidl/android/os/IStatsCompanionService.aidl
@@ -66,9 +66,6 @@
     /** Pull the specified data. Results will be sent to statsd when complete. */
     StatsLogEventWrapper[] pullData(int pullCode);
 
-    /** Send a broadcast to the specified PendingIntent's as IBinder that it should getData now. */
-    oneway void sendDataBroadcast(in IBinder intentSender, long lastReportTimeNs);
-
     /**
      * Send a broadcast to the specified PendingIntent's as IBinder notifying it that the list
      * of active configs has changed.
diff --git a/apex/statsd/aidl/android/os/IStatsManagerService.aidl b/apex/statsd/aidl/android/os/IStatsManagerService.aidl
index 45ba3a2..151cdbd 100644
--- a/apex/statsd/aidl/android/os/IStatsManagerService.aidl
+++ b/apex/statsd/aidl/android/os/IStatsManagerService.aidl
@@ -30,12 +30,19 @@
      * memory consumed by the metrics for this configuration approach the pre-defined limits. There
      * can be at most one listener per config key.
      *
-     * Requires Manifest.permission.DUMP.
+     * Requires Manifest.permission.DUMP and Manifest.permission.PACKAGE_USAGE_STATS.
      */
-    void setDataFetchOperation(long configKey, in PendingIntent pendingIntent,
+    void setDataFetchOperation(long configId, in PendingIntent pendingIntent,
         in String packageName);
 
     /**
+     * Removes the data fetch operation for the specified configuration.
+     *
+     * Requires Manifest.permission.DUMP and Manifest.permission.PACKAGE_USAGE_STATS.
+     */
+    void removeDataFetchOperation(long configId, in String packageName);
+
+    /**
      * Registers the given pending intent for this packagename. This intent is invoked when the
      * active status of any of the configs sent by this package changes and will contain a list of
      * config ids that are currently active. It also returns the list of configs that are currently
diff --git a/apex/statsd/aidl/android/os/IStatsd.aidl b/apex/statsd/aidl/android/os/IStatsd.aidl
index cce79fa..986d483 100644
--- a/apex/statsd/aidl/android/os/IStatsd.aidl
+++ b/apex/statsd/aidl/android/os/IStatsd.aidl
@@ -17,6 +17,7 @@
 package android.os;
 
 import android.os.IStatsPullerCallback;
+import android.os.IPendingIntentRef;
 import android.os.IPullAtomCallback;
 import android.os.ParcelFileDescriptor;
 
@@ -114,14 +115,15 @@
      *
      * Requires Manifest.permission.DUMP.
      */
-    void setDataFetchOperation(long configKey, in IBinder intentSender, in String packageName);
+    void setDataFetchOperation(long configId, in IPendingIntentRef pendingIntentRef,
+                               int callingUid);
 
     /**
      * Removes the data fetch operation for the specified configuration.
      *
      * Requires Manifest.permission.DUMP.
      */
-    void removeDataFetchOperation(long configKey, in String packageName);
+    void removeDataFetchOperation(long configId, int callingUid);
 
     /**
      * Registers the given pending intent for this packagename. This intent is invoked when the
diff --git a/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java b/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java
index 71b52e2..6a0bf62 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsCompanion.java
@@ -16,7 +16,13 @@
 
 package com.android.server.stats;
 
+import android.app.PendingIntent;
 import android.content.Context;
+import android.content.Intent;
+import android.os.Binder;
+import android.os.IPendingIntentRef;
+import android.os.Process;
+import android.os.StatsDimensionsValue;
 import android.util.Slog;
 
 import com.android.server.SystemService;
@@ -28,6 +34,13 @@
     private static final String TAG = "StatsCompanion";
     private static final boolean DEBUG = false;
 
+    static void enforceStatsCompanionPermission(Context context) {
+        if (Binder.getCallingPid() == Process.myPid()) {
+            return;
+        }
+        context.enforceCallingPermission(android.Manifest.permission.STATSCOMPANION, null);
+    }
+
     /**
      * Lifecycle class for both {@link StatsCompanionService} and {@link StatsManagerService}.
      */
@@ -63,8 +76,57 @@
             super.onBootPhase(phase);
             if (phase == PHASE_THIRD_PARTY_APPS_CAN_START) {
                 mStatsCompanionService.systemReady();
-                mStatsManagerService.systemReady();
             }
         }
     }
+
+    /**
+     * Wrapper for {@link PendingIntent}. Allows Statsd to send PendingIntents.
+     */
+    public static class PendingIntentRef extends IPendingIntentRef.Stub {
+
+        private static final String TAG = "PendingIntentRef";
+
+        /**
+         * The last report time is provided with each intent registered to
+         * StatsManager#setFetchReportsOperation. This allows easy de-duping in the receiver if
+         * statsd is requesting the client to retrieve the same statsd data. The last report time
+         * corresponds to the last_report_elapsed_nanos that will provided in the current
+         * ConfigMetricsReport, and this timestamp also corresponds to the
+         * current_report_elapsed_nanos of the most recently obtained ConfigMetricsReport.
+         */
+        private static final String EXTRA_LAST_REPORT_TIME = "android.app.extra.LAST_REPORT_TIME";
+        private static final int CODE_DATA_BROADCAST = 1;
+
+        private final PendingIntent mPendingIntent;
+        private final Context mContext;
+
+        public PendingIntentRef(PendingIntent pendingIntent, Context context) {
+            mPendingIntent = pendingIntent;
+            mContext = context;
+        }
+
+        @Override
+        public void sendDataBroadcast(long lastReportTimeNs) {
+            enforceStatsCompanionPermission(mContext);
+            Intent intent = new Intent();
+            intent.putExtra(EXTRA_LAST_REPORT_TIME, lastReportTimeNs);
+            try {
+                mPendingIntent.send(mContext, CODE_DATA_BROADCAST, intent, null, null);
+            } catch (PendingIntent.CanceledException e) {
+                Slog.w(TAG, "Unable to send PendingIntent");
+            }
+        }
+
+        @Override
+        public void sendActiveConfigsChangedBroadcast(long[] configIds) {
+            // no-op
+        }
+
+        @Override
+        public void sendSubscriberBroadcast(long configUid, long configId, long subscriptionId,
+                long subscriptionRuleId, String[] cookies, StatsDimensionsValue dimensionsValue) {
+            // no-op
+        }
+    }
 }
diff --git a/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java b/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
index e4f7300..d160f22 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
@@ -85,7 +85,6 @@
 import android.os.Looper;
 import android.os.ParcelFileDescriptor;
 import android.os.Parcelable;
-import android.os.Process;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.StatFs;
@@ -203,18 +202,8 @@
     private static final int PACKAGE_NAME_FIELD_ID = 4;
     private static final int INSTALLER_FIELD_ID = 5;
 
-    public static final int CODE_DATA_BROADCAST = 1;
     public static final int CODE_SUBSCRIBER_BROADCAST = 1;
     public static final int CODE_ACTIVE_CONFIGS_BROADCAST = 1;
-    /**
-     * The last report time is provided with each intent registered to
-     * StatsManager#setFetchReportsOperation. This allows easy de-duping in the receiver if
-     * statsd is requesting the client to retrieve the same statsd data. The last report time
-     * corresponds to the last_report_elapsed_nanos that will provided in the current
-     * ConfigMetricsReport, and this timestamp also corresponds to the
-     * current_report_elapsed_nanos of the most recently obtained ConfigMetricsReport.
-     */
-    public static final String EXTRA_LAST_REPORT_TIME = "android.app.extra.LAST_REPORT_TIME";
     public static final int DEATH_THRESHOLD = 10;
     /**
      * Which native processes to snapshot memory for.
@@ -455,21 +444,8 @@
     }
 
     @Override
-    public void sendDataBroadcast(IBinder intentSenderBinder, long lastReportTimeNs) {
-        enforceCallingPermission();
-        IntentSender intentSender = new IntentSender(intentSenderBinder);
-        Intent intent = new Intent();
-        intent.putExtra(EXTRA_LAST_REPORT_TIME, lastReportTimeNs);
-        try {
-            intentSender.sendIntent(mContext, CODE_DATA_BROADCAST, intent, null, null);
-        } catch (IntentSender.SendIntentException e) {
-            Slog.w(TAG, "Unable to send using IntentSender");
-        }
-    }
-
-    @Override
     public void sendActiveConfigsChangedBroadcast(IBinder intentSenderBinder, long[] configIds) {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         IntentSender intentSender = new IntentSender(intentSenderBinder);
         Intent intent = new Intent();
         intent.putExtra(StatsManager.EXTRA_STATS_ACTIVE_CONFIG_KEYS, configIds);
@@ -487,7 +463,7 @@
     public void sendSubscriberBroadcast(IBinder intentSenderBinder, long configUid, long configKey,
             long subscriptionId, long subscriptionRuleId, String[] cookies,
             StatsDimensionsValue dimensionsValue) {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         IntentSender intentSender = new IntentSender(intentSenderBinder);
         Intent intent =
                 new Intent()
@@ -770,7 +746,7 @@
 
     @Override // Binder call
     public void setAnomalyAlarm(long timestampMs) {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) Slog.d(TAG, "Setting anomaly alarm for " + timestampMs);
         final long callingToken = Binder.clearCallingIdentity();
         try {
@@ -786,7 +762,7 @@
 
     @Override // Binder call
     public void cancelAnomalyAlarm() {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) Slog.d(TAG, "Cancelling anomaly alarm");
         final long callingToken = Binder.clearCallingIdentity();
         try {
@@ -798,7 +774,7 @@
 
     @Override // Binder call
     public void setAlarmForSubscriberTriggering(long timestampMs) {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) {
             Slog.d(TAG,
                     "Setting periodic alarm in about " + (timestampMs
@@ -817,7 +793,7 @@
 
     @Override // Binder call
     public void cancelAlarmForSubscriberTriggering() {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) {
             Slog.d(TAG, "Cancelling periodic alarm");
         }
@@ -831,7 +807,7 @@
 
     @Override // Binder call
     public void setPullingAlarm(long nextPullTimeMs) {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) {
             Slog.d(TAG, "Setting pulling alarm in about "
                     + (nextPullTimeMs - SystemClock.elapsedRealtime()));
@@ -849,7 +825,7 @@
 
     @Override // Binder call
     public void cancelPullingAlarm() {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) {
             Slog.d(TAG, "Cancelling pulling alarm");
         }
@@ -2455,7 +2431,7 @@
      */
     @Override // Binder call
     public StatsLogEventWrapper[] pullData(int tagId) {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) {
             Slog.d(TAG, "Pulling " + tagId);
         }
@@ -2690,12 +2666,11 @@
 
     @Override // Binder call
     public void statsdReady() {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         if (DEBUG) {
             Slog.d(TAG, "learned that statsdReady");
         }
         sayHiToStatsd(); // tell statsd that we're ready too and link to it
-        mStatsManagerService.systemReady();
         mContext.sendBroadcastAsUser(new Intent(StatsManager.ACTION_STATSD_STARTED)
                         .addFlags(Intent.FLAG_RECEIVER_INCLUDE_BACKGROUND),
                 UserHandle.SYSTEM, android.Manifest.permission.DUMP);
@@ -2703,7 +2678,7 @@
 
     @Override
     public void triggerUidSnapshot() {
-        enforceCallingPermission();
+        StatsCompanion.enforceStatsCompanionPermission(mContext);
         synchronized (sStatsdLock) {
             final long token = Binder.clearCallingIdentity();
             try {
@@ -2716,13 +2691,6 @@
         }
     }
 
-    private void enforceCallingPermission() {
-        if (Binder.getCallingPid() == Process.myPid()) {
-            return;
-        }
-        mContext.enforceCallingPermission(android.Manifest.permission.STATSCOMPANION, null);
-    }
-
     @Override
     public void registerPullAtomCallback(int atomTag, long coolDownNs, long timeoutNs,
             int[] additiveFields, IPullAtomCallback pullerCallback) {
@@ -2817,6 +2785,7 @@
                                 + "alive.");
                 return;
             }
+            mStatsManagerService.statsdReady(sStatsd);
             if (DEBUG) Slog.d(TAG, "Saying hi to statsd");
             try {
                 sStatsd.statsCompanionReady();
@@ -2928,6 +2897,7 @@
         if (looperStats != null) {
             looperStats.reset();
         }
+        mStatsManagerService.statsdNotReady();
     }
 
     @Override
diff --git a/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java b/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java
index 24b7978..46fe597 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsManagerService.java
@@ -16,18 +16,27 @@
 
 package com.android.server.stats;
 
+import static com.android.server.stats.StatsCompanion.PendingIntentRef;
+
+import android.Manifest;
+import android.app.AppOpsManager;
 import android.app.PendingIntent;
 import android.content.Context;
-import android.os.IBinder;
+import android.os.Binder;
 import android.os.IStatsManagerService;
 import android.os.IStatsd;
+import android.os.Process;
 import android.os.RemoteException;
-import android.os.ServiceManager;
+import android.util.ArrayMap;
 import android.util.Slog;
 
 import com.android.internal.annotations.GuardedBy;
 
+import java.util.Map;
+import java.util.Objects;
+
 /**
+ * Service for {@link android.app.StatsManager}.
  * @hide
  */
 public class StatsManagerService extends IStatsManagerService.Stub {
@@ -35,22 +44,99 @@
     private static final String TAG = "StatsManagerService";
     private static final boolean DEBUG = false;
 
-    @GuardedBy("sStatsdLock")
-    private static IStatsd sStatsd;
-    private static final Object sStatsdLock = new Object();
+    private static final int STATSD_TIMEOUT_MILLIS = 5000;
+
+    private static final String USAGE_STATS_PERMISSION_OPS = "android:get_usage_stats";
+
+    @GuardedBy("mLock")
+    private IStatsd mStatsd;
+    private final Object mLock = new Object();
 
     private StatsCompanionService mStatsCompanionService;
+    private Context mContext;
+
+    @GuardedBy("mLock")
+    private ArrayMap<ConfigKey, PendingIntentRef> mDataFetchPirMap = new ArrayMap<>();
 
     public StatsManagerService(Context context) {
         super();
+        mContext = context;
+    }
+
+    private static class ConfigKey {
+        private int mUid;
+        private long mConfigId;
+
+        ConfigKey(int uid, long configId) {
+            mUid = uid;
+            mConfigId = configId;
+        }
+
+        public int getUid() {
+            return mUid;
+        }
+
+        public long getConfigId() {
+            return mConfigId;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(mUid, mConfigId);
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj instanceof ConfigKey) {
+                ConfigKey other = (ConfigKey) obj;
+                return this.mUid == other.getUid() && this.mConfigId == other.getConfigId();
+            }
+            return false;
+        }
     }
 
     @Override
-    public void setDataFetchOperation(long configKey, PendingIntent pendingIntent,
+    public void setDataFetchOperation(long configId, PendingIntent pendingIntent,
             String packageName) {
-        // no-op
-        if (DEBUG) {
-            Slog.d(TAG, "setDataFetchOperation");
+        enforceDumpAndUsageStatsPermission(packageName);
+        int callingUid = Binder.getCallingUid();
+        final long token = Binder.clearCallingIdentity();
+        PendingIntentRef pir = new PendingIntentRef(pendingIntent, mContext);
+        ConfigKey key = new ConfigKey(callingUid, configId);
+        // We add the PIR to a map so we can reregister if statsd is unavailable.
+        synchronized (mLock) {
+            mDataFetchPirMap.put(key, pir);
+        }
+        try {
+            IStatsd statsd = getStatsdNonblocking();
+            if (statsd != null) {
+                statsd.setDataFetchOperation(configId, pir, callingUid);
+            }
+        } catch (RemoteException e) {
+            Slog.e(TAG, "Failed to setDataFetchOperation with statsd");
+        } finally {
+            Binder.restoreCallingIdentity(token);
+        }
+    }
+
+    @Override
+    public void removeDataFetchOperation(long configId, String packageName) {
+        enforceDumpAndUsageStatsPermission(packageName);
+        int callingUid = Binder.getCallingUid();
+        final long token = Binder.clearCallingIdentity();
+        ConfigKey key = new ConfigKey(callingUid, configId);
+        synchronized (mLock) {
+            mDataFetchPirMap.remove(key);
+        }
+        try {
+            IStatsd statsd = getStatsdNonblocking();
+            if (statsd != null) {
+                statsd.removeDataFetchOperation(configId, callingUid);
+            }
+        } catch (RemoteException e) {
+            Slog.e(TAG, "Failed to removeDataFetchOperation with statsd");
+        } finally {
+            Binder.restoreCallingIdentity(token);
         }
     }
 
@@ -77,37 +163,100 @@
         mStatsCompanionService = statsCompanionService;
     }
 
-    void systemReady() {
-        if (DEBUG) {
-            Slog.d(TAG, "statsdReady");
+    private void enforceDumpAndUsageStatsPermission(String packageName) {
+        int callingUid = Binder.getCallingUid();
+        int callingPid = Binder.getCallingPid();
+
+        if (callingPid == Process.myPid()) {
+            return;
         }
-        setupStatsManagerService();
+        mContext.enforceCallingPermission(Manifest.permission.DUMP, null);
+        mContext.enforceCallingPermission(Manifest.permission.PACKAGE_USAGE_STATS, null);
+
+        AppOpsManager appOpsManager = (AppOpsManager) mContext
+                .getSystemService(Context.APP_OPS_SERVICE);
+        switch (appOpsManager.noteOp(USAGE_STATS_PERMISSION_OPS,
+                Binder.getCallingUid(), packageName, null, null)) {
+            case AppOpsManager.MODE_ALLOWED:
+            case AppOpsManager.MODE_DEFAULT:
+                break;
+            default:
+                throw new SecurityException(
+                        String.format("UID %d / PID %d lacks app-op %s",
+                                callingUid, callingPid, USAGE_STATS_PERMISSION_OPS)
+                );
+        }
     }
 
-    private void setupStatsManagerService() {
-        synchronized (sStatsdLock) {
-            if (sStatsd != null) {
-                if (DEBUG) {
-                    Slog.e(TAG, "Trying to fetch statsd, but it was already fetched",
-                            new IllegalStateException(
-                                    "sStatsd is not null when being fetched"));
+    /**
+     * Clients should call this if blocking until statsd to be ready is desired
+     *
+     * @return IStatsd object if statsd becomes ready within the timeout, null otherwise.
+     */
+    private IStatsd waitForStatsd() {
+        synchronized (mLock) {
+            if (mStatsd == null) {
+                try {
+                    mLock.wait(STATSD_TIMEOUT_MILLIS);
+                } catch (InterruptedException e) {
+                    Slog.e(TAG, "wait for statsd interrupted");
                 }
-                return;
             }
-            sStatsd = IStatsd.Stub.asInterface(ServiceManager.getService("stats"));
-            if (sStatsd == null) {
-                if (DEBUG) {
-                    Slog.d(TAG, "Failed to get stats service.");
-                }
-                return;
-            }
-            // Assume statsd is ready since this is called form statscompanion, link to statsd.
+            return mStatsd;
+        }
+    }
+
+    /**
+     * Clients should call this to receive a reference to statsd.
+     *
+     * @return IStatsd object if statsd is ready, null otherwise.
+     */
+    private IStatsd getStatsdNonblocking() {
+        synchronized (mLock) {
+            return mStatsd;
+        }
+    }
+
+    /**
+     * Called from {@link StatsCompanionService}.
+     *
+     * Tells StatsManagerService that Statsd is ready and updates
+     * Statsd with the contents of our local cache.
+     */
+    void statsdReady(IStatsd statsd) {
+        synchronized (mLock) {
+            mStatsd = statsd;
+            mLock.notify();
+        }
+        sayHiToStatsd(statsd);
+    }
+
+    /**
+     * Called from {@link StatsCompanionService}.
+     *
+     * Tells StatsManagerService that Statsd is no longer ready
+     * and we should no longer make binder calls with statsd.
+     */
+    void statsdNotReady() {
+        synchronized (mLock) {
+            mStatsd = null;
+        }
+    }
+
+    private void sayHiToStatsd(IStatsd statsd) {
+        if (statsd == null) {
+            return;
+        }
+        ArrayMap<ConfigKey, PendingIntentRef> dataFetchCopy;
+        synchronized (mLock) {
+            dataFetchCopy = new ArrayMap<>(mDataFetchPirMap);
+        }
+        for (Map.Entry<ConfigKey, PendingIntentRef> entry : dataFetchCopy.entrySet()) {
+            ConfigKey key = entry.getKey();
             try {
-                sStatsd.asBinder().linkToDeath((IBinder.DeathRecipient) () -> {
-                    sStatsd = null;
-                }, 0);
+                statsd.setDataFetchOperation(key.getConfigId(), entry.getValue(), key.getUid());
             } catch (RemoteException e) {
-                Slog.e(TAG, "linkToDeath(StatsdDeathRecipient) failed", e);
+                Slog.e(TAG, "Failed to setDataFetchOperation from pirMap");
             }
         }
     }
diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index bb3a094..7fecf46 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -170,18 +170,18 @@
             mUidMap, mPullerManager, mAnomalyAlarmMonitor, mPeriodicAlarmMonitor,
             getElapsedRealtimeNs(),
             [this](const ConfigKey& key) {
-                sp<IStatsCompanionService> sc = getStatsCompanionService();
-                auto receiver = mConfigManager->GetConfigReceiver(key);
-                if (sc == nullptr) {
-                    VLOG("Could not find StatsCompanionService");
+                sp<IPendingIntentRef> receiver = mConfigManager->GetConfigReceiver(key);
+                if (receiver == nullptr) {
+                    VLOG("Could not find a broadcast receiver for %s",
+                        key.ToString().c_str());
                     return false;
-                } else if (receiver == nullptr) {
-                    VLOG("Statscompanion could not find a broadcast receiver for %s",
-                         key.ToString().c_str());
-                    return false;
-                } else {
-                    sc->sendDataBroadcast(receiver, mProcessor->getLastReportTimeNs(key));
+                } else if (receiver->sendDataBroadcast(
+                           mProcessor->getLastReportTimeNs(key)).isOk()) {
                     return true;
+                } else {
+                    VLOG("Failed to send a broadcast for receiver %s",
+                        key.ToString().c_str());
+                    return false;
                 }
             },
             [this](const int& uid, const vector<int64_t>& activeConfigs) {
@@ -574,18 +574,19 @@
         return UNKNOWN_ERROR;
     }
     ConfigKey key(uid, StrToInt64(name));
-    auto receiver = mConfigManager->GetConfigReceiver(key);
-    sp<IStatsCompanionService> sc = getStatsCompanionService();
-    if (sc == nullptr) {
-        VLOG("Could not access statsCompanion");
-    } else if (receiver == nullptr) {
-        VLOG("Could not find receiver for %s, %s", args[1].c_str(), args[2].c_str())
-    } else {
-        sc->sendDataBroadcast(receiver, mProcessor->getLastReportTimeNs(key));
+    sp<IPendingIntentRef> receiver = mConfigManager->GetConfigReceiver(key);
+    if (receiver == nullptr) {
+        VLOG("Could not find receiver for %s, %s", args[1].c_str(), args[2].c_str());
+        return UNKNOWN_ERROR;
+    } else if (receiver->sendDataBroadcast(
+               mProcessor->getLastReportTimeNs(key)).isOk()) {
         VLOG("StatsService::trigger broadcast succeeded to %s, %s", args[1].c_str(),
              args[2].c_str());
+    } else {
+        VLOG("StatsService::trigger broadcast failed to %s, %s", args[1].c_str(),
+             args[2].c_str());
+        return UNKNOWN_ERROR;
     }
-
     return NO_ERROR;
 }
 
@@ -1185,23 +1186,21 @@
     return true;
 }
 
-Status StatsService::removeDataFetchOperation(int64_t key, const String16& packageName) {
-    ENFORCE_DUMP_AND_USAGE_STATS(packageName);
-
-    IPCThreadState* ipc = IPCThreadState::self();
-    ConfigKey configKey(ipc->getCallingUid(), key);
+Status StatsService::removeDataFetchOperation(int64_t key,
+                                              const int32_t callingUid) {
+    ENFORCE_UID(AID_SYSTEM);
+    ConfigKey configKey(callingUid, key);
     mConfigManager->RemoveConfigReceiver(configKey);
     return Status::ok();
 }
 
 Status StatsService::setDataFetchOperation(int64_t key,
-                                           const sp<android::IBinder>& intentSender,
-                                           const String16& packageName) {
-    ENFORCE_DUMP_AND_USAGE_STATS(packageName);
+                                           const sp<IPendingIntentRef>& pir,
+                                           const int32_t callingUid) {
+    ENFORCE_UID(AID_SYSTEM);
 
-    IPCThreadState* ipc = IPCThreadState::self();
-    ConfigKey configKey(ipc->getCallingUid(), key);
-    mConfigManager->SetConfigReceiver(configKey, intentSender);
+    ConfigKey configKey(callingUid, key);
+    mConfigManager->SetConfigReceiver(configKey, pir);
     if (StorageManager::hasConfigMetricsReport(configKey)) {
         VLOG("StatsService::setDataFetchOperation marking configKey %s to dump reports on disk",
              configKey.ToString().c_str());
diff --git a/cmds/statsd/src/StatsService.h b/cmds/statsd/src/StatsService.h
index de55ca9..9912d2f 100644
--- a/cmds/statsd/src/StatsService.h
+++ b/cmds/statsd/src/StatsService.h
@@ -30,6 +30,7 @@
 #include <android/frameworks/stats/1.0/IStats.h>
 #include <android/frameworks/stats/1.0/types.h>
 #include <android/os/BnStatsd.h>
+#include <android/os/IPendingIntentRef.h>
 #include <android/os/IStatsCompanionService.h>
 #include <android/os/IStatsd.h>
 #include <binder/IResultReceiver.h>
@@ -121,14 +122,14 @@
      * Binder call to let clients register the data fetch operation for a configuration.
      */
     virtual Status setDataFetchOperation(int64_t key,
-                                         const sp<android::IBinder>& intentSender,
-                                         const String16& packageName) override;
+                                         const sp<IPendingIntentRef>& pir,
+                                         const int32_t callingUid) override;
 
     /**
      * Binder call to remove the data fetch operation for the specified config key.
      */
     virtual Status removeDataFetchOperation(int64_t key,
-                                            const String16& packageName) override;
+                                            const int32_t callingUid) override;
 
     /**
      * Binder call to let clients register the active configs changed operation.
diff --git a/cmds/statsd/src/config/ConfigManager.cpp b/cmds/statsd/src/config/ConfigManager.cpp
index fc949b4..7bfb991 100644
--- a/cmds/statsd/src/config/ConfigManager.cpp
+++ b/cmds/statsd/src/config/ConfigManager.cpp
@@ -46,6 +46,23 @@
 using android::base::StringPrintf;
 using std::unique_ptr;
 
+class ConfigReceiverDeathRecipient : public android::IBinder::DeathRecipient {
+    public:
+        ConfigReceiverDeathRecipient(sp<ConfigManager> configManager, const ConfigKey& configKey):
+            mConfigManager(configManager),
+            mConfigKey(configKey) {}
+        ~ConfigReceiverDeathRecipient() override = default;
+    private:
+        sp<ConfigManager> mConfigManager;
+        ConfigKey mConfigKey;
+
+    void binderDied(const android::wp<android::IBinder>& who) override {
+        if (IInterface::asBinder(mConfigManager->GetConfigReceiver(mConfigKey)) == who.promote()) {
+             mConfigManager->RemoveConfigReceiver(mConfigKey);
+        }
+    }
+};
+
 ConfigManager::ConfigManager() {
 }
 
@@ -118,9 +135,11 @@
     }
 }
 
-void ConfigManager::SetConfigReceiver(const ConfigKey& key, const sp<IBinder>& intentSender) {
+void ConfigManager::SetConfigReceiver(const ConfigKey& key,
+                                      const sp<IPendingIntentRef>& pir) {
     lock_guard<mutex> lock(mMutex);
-    mConfigReceivers[key] = intentSender;
+    mConfigReceivers[key] = pir;
+    IInterface::asBinder(pir)->linkToDeath(new ConfigReceiverDeathRecipient(this, key));
 }
 
 void ConfigManager::RemoveConfigReceiver(const ConfigKey& key) {
@@ -266,7 +285,7 @@
     return ret;
 }
 
-const sp<android::IBinder> ConfigManager::GetConfigReceiver(const ConfigKey& key) const {
+const sp<IPendingIntentRef> ConfigManager::GetConfigReceiver(const ConfigKey& key) const {
     lock_guard<mutex> lock(mMutex);
 
     auto it = mConfigReceivers.find(key);
diff --git a/cmds/statsd/src/config/ConfigManager.h b/cmds/statsd/src/config/ConfigManager.h
index c064a51..1aeb355 100644
--- a/cmds/statsd/src/config/ConfigManager.h
+++ b/cmds/statsd/src/config/ConfigManager.h
@@ -20,6 +20,7 @@
 #include "config/ConfigKey.h"
 #include "config/ConfigListener.h"
 
+#include <android/os/IPendingIntentRef.h>
 #include <map>
 #include <mutex>
 #include <set>
@@ -64,12 +65,12 @@
     /**
      * Sets the broadcast receiver for a configuration key.
      */
-    void SetConfigReceiver(const ConfigKey& key, const sp<IBinder>& intentSender);
+    void SetConfigReceiver(const ConfigKey& key, const sp<IPendingIntentRef>& pir);
 
     /**
      * Returns the package name and class name representing the broadcast receiver for this config.
      */
-    const sp<android::IBinder> GetConfigReceiver(const ConfigKey& key) const;
+    const sp<IPendingIntentRef> GetConfigReceiver(const ConfigKey& key) const;
 
     /**
      * Returns all config keys registered.
@@ -141,10 +142,9 @@
     std::map<int, std::set<ConfigKey>> mConfigs;
 
     /**
-     * Each config key can be subscribed by up to one receiver, specified as IBinder from
-     * PendingIntent.
+     * Each config key can be subscribed by up to one receiver, specified as IPendingIntentRef.
      */
-    std::map<ConfigKey, sp<android::IBinder>> mConfigReceivers;
+    std::map<ConfigKey, sp<IPendingIntentRef>> mConfigReceivers;
 
     /**
      * Each uid can be subscribed by up to one receiver to notify that the list of active configs
diff --git a/core/java/android/app/StatsManager.java b/core/java/android/app/StatsManager.java
index b9893aa..a458a55 100644
--- a/core/java/android/app/StatsManager.java
+++ b/core/java/android/app/StatsManager.java
@@ -309,18 +309,16 @@
             throws StatsUnavailableException {
         synchronized (sLock) {
             try {
-                IStatsd service = getIStatsdLocked();
+                IStatsManagerService service = getIStatsManagerServiceLocked();
                 if (pendingIntent == null) {
                     service.removeDataFetchOperation(configKey, mContext.getOpPackageName());
                 } else {
-                    // Extracts IIntentSender from the PendingIntent and turns it into an IBinder.
-                    IBinder intentSender = pendingIntent.getTarget().asBinder();
-                    service.setDataFetchOperation(configKey, intentSender,
+                    service.setDataFetchOperation(configKey, pendingIntent,
                             mContext.getOpPackageName());
                 }
 
             } catch (RemoteException e) {
-                Slog.e(TAG, "Failed to connect to statsd when registering data listener.");
+                Slog.e(TAG, "Failed to connect to statsmanager when registering data listener.");
                 throw new StatsUnavailableException("could not connect", e);
             } catch (SecurityException e) {
                 throw new StatsUnavailableException(e.getMessage(), e);