Fix settings show data usage with the uid of a removed package

When an app is removed, the stats still be included in
the persist snapshot in NetworkStatsFactory which
causes settings show the data usage with the uid of the
removed app. Thus, remove the stats of the removed
package from the persist snapshot when NetworkStatsService
receives ACTION_UID_REMOVED intent.

Bug: 239899930
Bug: 209360825
Test: FrameworksNetTests:NetworkStatsServiceTest
      FrameworksNetTests:NetworkStatsFactoryTest
Change-Id: I73cca367ac6bf0d2d29ef0a7d94500f1e6917dcb
diff --git a/service-t/src/com/android/server/net/NetworkStatsFactory.java b/service-t/src/com/android/server/net/NetworkStatsFactory.java
index c9d1718..8161f50 100644
--- a/service-t/src/com/android/server/net/NetworkStatsFactory.java
+++ b/service-t/src/com/android/server/net/NetworkStatsFactory.java
@@ -296,6 +296,16 @@
         return mTunAnd464xlatAdjustedStats.clone();
     }
 
+    /**
+     * Remove stats from {@code mPersistSnapshot} and {@code mTunAnd464xlatAdjustedStats} for the
+     * given uids.
+     */
+    public void removeUidsLocked(int[] uids) {
+        synchronized (mPersistentDataLock) {
+            mPersistSnapshot.removeUids(uids);
+            mTunAnd464xlatAdjustedStats.removeUids(uids);
+        }
+    }
 
     public void assertEquals(NetworkStats expected, NetworkStats actual) {
         if (expected.size() != actual.size()) {
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index c4ffdec..309c9dd 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -2469,13 +2469,13 @@
         mUidRecorder.removeUidsLocked(uids);
         mUidTagRecorder.removeUidsLocked(uids);
 
+        mStatsFactory.removeUidsLocked(uids);
         // Clear kernel stats associated with UID
         for (int uid : uids) {
             deleteKernelTagData(uid);
         }
-
-       // TODO: Remove the UID's entries from mOpenSessionCallsPerUid and
-       // mOpenSessionCallsPerCaller
+        // TODO: Remove the UID's entries from mOpenSessionCallsPerUid and
+        // mOpenSessionCallsPerCaller
     }
 
     /**
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java b/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java
index 14455fa..04db6d3 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsFactoryTest.java
@@ -25,6 +25,7 @@
 import static android.net.NetworkStats.SET_ALL;
 import static android.net.NetworkStats.SET_DEFAULT;
 import static android.net.NetworkStats.SET_FOREGROUND;
+import static android.net.NetworkStats.TAG_ALL;
 import static android.net.NetworkStats.TAG_NONE;
 import static android.net.NetworkStats.UID_ALL;
 
@@ -89,6 +90,7 @@
         // related to networkStatsFactory is compiled to a minimal native library and loaded here.
         System.loadLibrary("networkstatsfactorytestjni");
         doReturn(mBpfNetMaps).when(mDeps).createBpfNetMaps(any());
+
         mFactory = new NetworkStatsFactory(mContext, mDeps);
         mFactory.updateUnderlyingNetworkInfos(new UnderlyingNetworkInfo[0]);
     }
@@ -462,6 +464,46 @@
         assertNoStatsEntry(stats, "wlan0", 1029, SET_DEFAULT, 0x0);
     }
 
+    @Test
+    public void testRemoveUidsStats() throws Exception {
+        final NetworkStats stats = new NetworkStats(SystemClock.elapsedRealtime(), 1)
+                .insertEntry(TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 16L, 1L, 16L, 1L, 0L)
+                .insertEntry(TEST_IFACE, UID_BLUE, SET_DEFAULT, TAG_NONE,
+                        256L, 16L, 512L, 32L, 0L)
+                .insertEntry(TEST_IFACE, UID_GREEN, SET_DEFAULT, TAG_NONE, 64L, 3L, 1024L, 8L, 0L);
+
+        doReturn(stats).when(mDeps).getNetworkStatsDetail(anyInt(), any(),
+                anyInt());
+
+        final String[] ifaces = new String[]{TEST_IFACE};
+        final NetworkStats res = mFactory.readNetworkStatsDetail(UID_ALL, ifaces, TAG_ALL);
+
+        // Verify that the result of the mocked stats are expected.
+        assertValues(res, TEST_IFACE, UID_RED, 16L, 1L, 16L, 1L);
+        assertValues(res, TEST_IFACE, UID_BLUE, 256L, 16L, 512L, 32L);
+        assertValues(res, TEST_IFACE, UID_GREEN, 64L, 3L, 1024L, 8L);
+
+        // Assume the apps were removed.
+        final int[] removedUids = new int[]{UID_RED, UID_BLUE};
+        mFactory.removeUidsLocked(removedUids);
+
+        // Return empty stats for reading the result of removing uids stats later.
+        doReturn(buildEmptyStats()).when(mDeps).getNetworkStatsDetail(anyInt(), any(),
+                anyInt());
+
+        final NetworkStats removedUidsStats =
+                mFactory.readNetworkStatsDetail(UID_ALL, ifaces, TAG_ALL);
+
+        // Verify that the stats of the removed uids were removed.
+        assertValues(removedUidsStats, TEST_IFACE, UID_RED, 0L, 0L, 0L, 0L);
+        assertValues(removedUidsStats, TEST_IFACE, UID_BLUE, 0L, 0L, 0L, 0L);
+        assertValues(removedUidsStats, TEST_IFACE, UID_GREEN, 64L, 3L, 1024L, 8L);
+    }
+
+    private NetworkStats buildEmptyStats() {
+        return new NetworkStats(SystemClock.elapsedRealtime(), 0);
+    }
+
     private NetworkStats parseNetworkStatsFromGoldenSample(int resourceId, int initialSize,
             boolean consumeHeader, boolean checkActive, boolean isUidData) throws IOException {
         final NetworkStats stats =
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index be44946..992c078 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -84,6 +84,7 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -172,6 +173,7 @@
 import java.time.ZonedDateTime;
 import java.time.temporal.ChronoUnit;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -2027,6 +2029,59 @@
         }
     }
 
+    @Test
+    public void testStatsFactoryRemoveUids() throws Exception {
+        // pretend that network comes online
+        mockDefaultSettings();
+        NetworkStateSnapshot[] states = new NetworkStateSnapshot[] {buildWifiState()};
+        mockNetworkStatsSummary(buildEmptyStats());
+        mockNetworkStatsUidDetail(buildEmptyStats());
+
+        mService.notifyNetworkStatus(NETWORKS_WIFI, states, getActiveIface(states),
+                new UnderlyingNetworkInfo[0]);
+
+        // Create some traffic
+        incrementCurrentTime(HOUR_IN_MILLIS);
+        mockDefaultSettings();
+        final NetworkStats stats = new NetworkStats(getElapsedRealtime(), 1)
+                .insertEntry(TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 16L, 1L, 16L, 1L, 0L)
+                .insertEntry(TEST_IFACE, UID_BLUE, SET_DEFAULT, TAG_NONE,
+                        4096L, 258L, 512L, 32L, 0L)
+                .insertEntry(TEST_IFACE, UID_GREEN, SET_DEFAULT, TAG_NONE, 64L, 3L, 1024L, 8L, 0L);
+        mockNetworkStatsUidDetail(stats);
+
+        forcePollAndWaitForIdle();
+
+        // Verify service recorded history
+        assertUidTotal(sTemplateWifi, UID_RED, 16L, 1L, 16L, 1L, 0);
+        assertUidTotal(sTemplateWifi, UID_BLUE, 4096L, 258L, 512L, 32L, 0);
+        assertUidTotal(sTemplateWifi, UID_GREEN, 64L, 3L, 1024L, 8L, 0);
+
+        // Simulate that the apps are removed.
+        final Intent intentBlue = new Intent(ACTION_UID_REMOVED);
+        intentBlue.putExtra(EXTRA_UID, UID_BLUE);
+        mServiceContext.sendBroadcast(intentBlue);
+
+        final Intent intentRed = new Intent(ACTION_UID_REMOVED);
+        intentRed.putExtra(EXTRA_UID, UID_RED);
+        mServiceContext.sendBroadcast(intentRed);
+
+        final int[] removedUids = {UID_BLUE, UID_RED};
+
+        final ArgumentCaptor<int[]> removedUidsCaptor = ArgumentCaptor.forClass(int[].class);
+        verify(mStatsFactory, times(2)).removeUidsLocked(removedUidsCaptor.capture());
+        final List<int[]> captureRemovedUids = removedUidsCaptor.getAllValues();
+        // Simulate that the stats are removed in NetworkStatsFactory.
+        if (captureRemovedUids.contains(removedUids)) {
+            stats.removeUids(removedUids);
+        }
+
+        // Verify the stats of the removed uid is removed.
+        assertUidTotal(sTemplateWifi, UID_RED, 0L, 0L, 0L, 0L, 0);
+        assertUidTotal(sTemplateWifi, UID_BLUE, 0L, 0L, 0L, 0L, 0);
+        assertUidTotal(sTemplateWifi, UID_GREEN, 64L, 3L, 1024L, 8L, 0);
+    }
+
     private void assertShouldRunComparison(boolean expected, boolean isDebuggable) {
         assertEquals("shouldRunComparison (debuggable=" + isDebuggable + "): ",
                 expected, mService.shouldRunComparison());