Uid Sandboxing of Pullers

Overall flow of implementation:
1. parsing the config in MetricsManager to store the uids per atom. It
follows the mAllowedLogSources logic very closely
2. MetricsManager register itself as a PullUidProvider with the
PullerManager.
3. Metrics pass the config key when pulling (for both registering
receivers and normal pulls) , and the puller manager gets
the allowed uids from the PullUidProvider for that config.
4. PullerManager keys receivers by <atomId, configKey> so that it can
look up the uids for that atom using the PullUidProvider as well.
5. Added shell subscriber support. Hardcode a default of AID_SYSTEM for
them and also allow packages per atom. This involved adding a second
interface to Pull that simply accepts the uids, since I didnt want to
make the ShellSubscriber a PullUidProvider as well.
6. Change adb shell cmd stats pull-source to allow users to specify a
package. Default to AID_SYSTEM as well.

Notes:
The feature is flagged off right now, since configs do not pass in the
desired package. Another approach could be to hardcode in the current
mapping, but that doesn't work for OEM pulled atoms.

Test: m statsd
Test: bit statsd_test:* with useUids = false
Test: bit statsd_test:* with useUids = true
Bug: 144099783
Bug: 151978258

Change-Id: I4a7481d7402a52b9beb4ea28b102803f9e50e79f
diff --git a/cmds/statsd/src/external/StatsPullerManager.cpp b/cmds/statsd/src/external/StatsPullerManager.cpp
index 8b6a5a1..4ffa3d8 100644
--- a/cmds/statsd/src/external/StatsPullerManager.cpp
+++ b/cmds/statsd/src/external/StatsPullerManager.cpp
@@ -47,27 +47,73 @@
 StatsPullerManager::StatsPullerManager()
     : kAllPullAtomInfo({
               // TrainInfo.
-              {{.atomTag = util::TRAIN_INFO}, new TrainInfoPuller()},
+              {{.atomTag = util::TRAIN_INFO, .uid = -1}, new TrainInfoPuller()},
       }),
       mNextPullTimeNs(NO_ALARM_UPDATE) {
 }
 
-bool StatsPullerManager::Pull(int tagId, vector<shared_ptr<LogEvent>>* data) {
-    AutoMutex _l(mLock);
-    return PullLocked(tagId, data);
+bool StatsPullerManager::Pull(int tagId, const ConfigKey& configKey,
+                              vector<shared_ptr<LogEvent>>* data, bool useUids) {
+    std::lock_guard<std::mutex> _l(mLock);
+    return PullLocked(tagId, configKey, data, useUids);
 }
 
-bool StatsPullerManager::PullLocked(int tagId, vector<shared_ptr<LogEvent>>* data) {
-    VLOG("Initiating pulling %d", tagId);
+bool StatsPullerManager::Pull(int tagId, const vector<int32_t>& uids,
+                              vector<std::shared_ptr<LogEvent>>* data, bool useUids) {
+    std::lock_guard<std::mutex> _l(mLock);
+    return PullLocked(tagId, uids, data, useUids);
+}
 
-    if (kAllPullAtomInfo.find({.atomTag = tagId}) != kAllPullAtomInfo.end()) {
-        bool ret = kAllPullAtomInfo.find({.atomTag = tagId})->second->Pull(data);
-        VLOG("pulled %d items", (int)data->size());
-        if (!ret) {
-            StatsdStats::getInstance().notePullFailed(tagId);
+bool StatsPullerManager::PullLocked(int tagId, const ConfigKey& configKey,
+                                    vector<shared_ptr<LogEvent>>* data, bool useUids) {
+    vector<int32_t> uids;
+    if (useUids) {
+        auto uidProviderIt = mPullUidProviders.find(configKey);
+        if (uidProviderIt == mPullUidProviders.end()) {
+            ALOGE("Error pulling tag %d. No pull uid provider for config key %s", tagId,
+                  configKey.ToString().c_str());
+            return false;
         }
-        return ret;
+        sp<PullUidProvider> pullUidProvider = uidProviderIt->second.promote();
+        if (pullUidProvider == nullptr) {
+            ALOGE("Error pulling tag %d, pull uid provider for config %s is gone.", tagId,
+                  configKey.ToString().c_str());
+            return false;
+        }
+        uids = pullUidProvider->getPullAtomUids(tagId);
+    }
+    return PullLocked(tagId, uids, data, useUids);
+}
+
+bool StatsPullerManager::PullLocked(int tagId, const vector<int32_t>& uids,
+                                    vector<shared_ptr<LogEvent>>* data, bool useUids) {
+    VLOG("Initiating pulling %d", tagId);
+    if (useUids) {
+        for (int32_t uid : uids) {
+            PullerKey key = {.atomTag = tagId, .uid = uid};
+            auto pullerIt = kAllPullAtomInfo.find(key);
+            if (pullerIt != kAllPullAtomInfo.end()) {
+                bool ret = pullerIt->second->Pull(data);
+                VLOG("pulled %zu items", data->size());
+                if (!ret) {
+                    StatsdStats::getInstance().notePullFailed(tagId);
+                }
+                return ret;
+            }
+        }
+        ALOGW("StatsPullerManager: Unknown tagId %d", tagId);
+        return false;  // Return early since we don't know what to pull.
     } else {
+        PullerKey key = {.atomTag = tagId, .uid = -1};
+        auto pullerIt = kAllPullAtomInfo.find(key);
+        if (pullerIt != kAllPullAtomInfo.end()) {
+            bool ret = pullerIt->second->Pull(data);
+            VLOG("pulled %zu items", data->size());
+            if (!ret) {
+                StatsdStats::getInstance().notePullFailed(tagId);
+            }
+            return ret;
+        }
         ALOGW("StatsPullerManager: Unknown tagId %d", tagId);
         return false;  // Return early since we don't know what to pull.
     }
@@ -96,7 +142,7 @@
 
 void StatsPullerManager::SetStatsCompanionService(
         shared_ptr<IStatsCompanionService> statsCompanionService) {
-    AutoMutex _l(mLock);
+    std::lock_guard<std::mutex> _l(mLock);
     shared_ptr<IStatsCompanionService> tmpForLock = mStatsCompanionService;
     mStatsCompanionService = statsCompanionService;
     for (const auto& pulledAtom : kAllPullAtomInfo) {
@@ -107,10 +153,11 @@
     }
 }
 
-void StatsPullerManager::RegisterReceiver(int tagId, wp<PullDataReceiver> receiver,
-                                              int64_t nextPullTimeNs, int64_t intervalNs) {
-    AutoMutex _l(mLock);
-    auto& receivers = mReceivers[tagId];
+void StatsPullerManager::RegisterReceiver(int tagId, const ConfigKey& configKey,
+                                          wp<PullDataReceiver> receiver, int64_t nextPullTimeNs,
+                                          int64_t intervalNs) {
+    std::lock_guard<std::mutex> _l(mLock);
+    auto& receivers = mReceivers[{.atomTag = tagId, .configKey = configKey}];
     for (auto it = receivers.begin(); it != receivers.end(); it++) {
         if (it->receiver == receiver) {
             VLOG("Receiver already registered of %d", (int)receivers.size());
@@ -142,13 +189,15 @@
     VLOG("Puller for tagId %d registered of %d", tagId, (int)receivers.size());
 }
 
-void StatsPullerManager::UnRegisterReceiver(int tagId, wp<PullDataReceiver> receiver) {
-    AutoMutex _l(mLock);
-    if (mReceivers.find(tagId) == mReceivers.end()) {
+void StatsPullerManager::UnRegisterReceiver(int tagId, const ConfigKey& configKey,
+                                            wp<PullDataReceiver> receiver) {
+    std::lock_guard<std::mutex> _l(mLock);
+    auto receiversIt = mReceivers.find({.atomTag = tagId, .configKey = configKey});
+    if (receiversIt == mReceivers.end()) {
         VLOG("Unknown pull code or no receivers: %d", tagId);
         return;
     }
-    auto& receivers = mReceivers.find(tagId)->second;
+    std::list<ReceiverInfo>& receivers = receiversIt->second;
     for (auto it = receivers.begin(); it != receivers.end(); it++) {
         if (receiver == it->receiver) {
             receivers.erase(it);
@@ -158,16 +207,26 @@
     }
 }
 
+void StatsPullerManager::RegisterPullUidProvider(const ConfigKey& configKey,
+                                                 wp<PullUidProvider> provider) {
+    std::lock_guard<std::mutex> _l(mLock);
+    mPullUidProviders[configKey] = provider;
+}
+
+void StatsPullerManager::UnregisterPullUidProvider(const ConfigKey& configKey) {
+    std::lock_guard<std::mutex> _l(mLock);
+    mPullUidProviders.erase(configKey);
+}
+
 void StatsPullerManager::OnAlarmFired(int64_t elapsedTimeNs) {
-    AutoMutex _l(mLock);
+    std::lock_guard<std::mutex> _l(mLock);
     int64_t wallClockNs = getWallClockNs();
 
     int64_t minNextPullTimeNs = NO_ALARM_UPDATE;
 
-    vector<pair<int, vector<ReceiverInfo*>>> needToPull =
-            vector<pair<int, vector<ReceiverInfo*>>>();
+    vector<pair<const ReceiverKey*, vector<ReceiverInfo*>>> needToPull;
     for (auto& pair : mReceivers) {
-        vector<ReceiverInfo*> receivers = vector<ReceiverInfo*>();
+        vector<ReceiverInfo*> receivers;
         if (pair.second.size() != 0) {
             for (ReceiverInfo& receiverInfo : pair.second) {
                 if (receiverInfo.nextPullTimeNs <= elapsedTimeNs) {
@@ -179,17 +238,16 @@
                 }
             }
             if (receivers.size() > 0) {
-                needToPull.push_back(make_pair(pair.first, receivers));
+                needToPull.push_back(make_pair(&pair.first, receivers));
             }
         }
     }
-
     for (const auto& pullInfo : needToPull) {
         vector<shared_ptr<LogEvent>> data;
-        bool pullSuccess = PullLocked(pullInfo.first, &data);
+        bool pullSuccess = PullLocked(pullInfo.first->atomTag, pullInfo.first->configKey, &data);
         if (pullSuccess) {
-            StatsdStats::getInstance().notePullDelay(
-                    pullInfo.first, getElapsedRealtimeNs() - elapsedTimeNs);
+            StatsdStats::getInstance().notePullDelay(pullInfo.first->atomTag,
+                                                     getElapsedRealtimeNs() - elapsedTimeNs);
         } else {
             VLOG("pull failed at %lld, will try again later", (long long)elapsedTimeNs);
         }
@@ -248,18 +306,19 @@
 void StatsPullerManager::RegisterPullAtomCallback(const int uid, const int32_t atomTag,
                                                   const int64_t coolDownNs, const int64_t timeoutNs,
                                                   const vector<int32_t>& additiveFields,
-                                                  const shared_ptr<IPullAtomCallback>& callback) {
-    AutoMutex _l(mLock);
+                                                  const shared_ptr<IPullAtomCallback>& callback,
+                                                  bool useUid) {
+    std::lock_guard<std::mutex> _l(mLock);
     VLOG("RegisterPullerCallback: adding puller for tag %d", atomTag);
     // TODO(b/146439412): linkToDeath with the callback so that we can remove it
     // and delete the puller.
     StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/true);
-    kAllPullAtomInfo[{.atomTag = atomTag}] =
+    kAllPullAtomInfo[{.atomTag = atomTag, .uid = useUid ? uid : -1}] =
             new StatsCallbackPuller(atomTag, callback, coolDownNs, timeoutNs, additiveFields);
 }
 
 void StatsPullerManager::UnregisterPullAtomCallback(const int uid, const int32_t atomTag) {
-    AutoMutex _l(mLock);
+    std::lock_guard<std::mutex> _l(mLock);
     StatsdStats::getInstance().notePullerCallbackRegistrationChanged(atomTag, /*registered=*/false);
     kAllPullAtomInfo.erase({.atomTag = atomTag});
 }