NetworkStatsService: Fix getDetailedUidStats to take VPNs into account.

(cherry picked from commit 720133f79d)

This API is similar to one provided by NetworkStatsFactory with the
difference that NSS also migrates traffic from VPN UID to other apps.

Since traffic can only be migrated over NetworkStats delta, NSS
therefore maintains NetworkStats snapshot across all UIDs/ifaces/tags.

This snapshot gets updated whenever NSS records a new snapshot
(based on various hooks such as VPN updating its underlying networks,
network getting lost, etc.), or getDetailedUidStats API is invoked by
one of its callers.

Bug: 113122541
Bug: 120145746
Test: atest FrameworksNetTests
Test: manually verified that battery stats are migrating traffic off of
TUN (after patching above CL where we point BatteryStats to use this
API).

Change-Id: I4b8d7c5b6905a4a12c1806dfd35c2c4c63610404

core/java/android/net/NetworkStats.java

@@ -34,6 +34,7 @@ import libcore.util.EmptyArray;
 import java.io.CharArrayWriter;
 import java.io.PrintWriter;
 import java.util.Arrays;
+import java.util.function.Predicate;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
@@ -994,23 +995,33 @@ public class NetworkStats implements Parcelable {
         if (limitUid == UID_ALL && limitTag == TAG_ALL && limitIfaces == INTERFACES_ALL) {
             return;
         }
+        filter(e -> (limitUid == UID_ALL || limitUid == e.uid)
+            && (limitTag == TAG_ALL || limitTag == e.tag)
+            && (limitIfaces == INTERFACES_ALL
+                    || ArrayUtils.contains(limitIfaces, e.iface)));
+    }
 
+    /**
+     * Only keep entries with {@link #set} value less than {@link #SET_DEBUG_START}.
+     *
+     * <p>This mutates the original structure in place.
+     */
+    public void filterDebugEntries() {
+        filter(e -> e.set < SET_DEBUG_START);
+    }
+
+    private void filter(Predicate<Entry> predicate) {
         Entry entry = new Entry();
         int nextOutputEntry = 0;
         for (int i = 0; i < size; i++) {
             entry = getValues(i, entry);
-            final boolean matches =
-                    (limitUid == UID_ALL || limitUid == entry.uid)
-                    && (limitTag == TAG_ALL || limitTag == entry.tag)
-                    && (limitIfaces == INTERFACES_ALL
-                            || ArrayUtils.contains(limitIfaces, entry.iface));
-
-            if (matches) {
-                setValues(nextOutputEntry, entry);
+            if (predicate.test(entry)) {
+                if (nextOutputEntry != i) {
+                    setValues(nextOutputEntry, entry);
+                }
                 nextOutputEntry++;
             }
         }
-
         size = nextOutputEntry;
     }
 

services/core/java/com/android/server/net/NetworkStatsFactory.java

@@ -263,6 +263,10 @@ public class NetworkStatsFactory {
         return stats;
     }
 
+    /**
+     * @deprecated Use NetworkStatsService#getDetailedUidStats which also accounts for
+     * VPN traffic
+     */
     public NetworkStats readNetworkStatsDetail() throws IOException {
         return readNetworkStatsDetail(UID_ALL, null, TAG_ALL, null);
     }

services/core/java/com/android/server/net/NetworkStatsService.java

@@ -293,6 +293,22 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
     /** Data layer operation counters for splicing into other structures. */
     private NetworkStats mUidOperations = new NetworkStats(0L, 10);
 
+    /**
+     * Snapshot containing most recent network stats for all UIDs across all interfaces and tags
+     * since boot.
+     *
+     * <p>Maintains migrated VPN stats which are result of performing TUN migration on {@link
+     * #mLastUidDetailSnapshot}.
+     */
+    @GuardedBy("mStatsLock")
+    private NetworkStats mTunAdjustedStats;
+    /**
+     * Used by {@link #mTunAdjustedStats} to migrate VPN traffic over delta between this snapshot
+     * and latest snapshot.
+     */
+    @GuardedBy("mStatsLock")
+    private NetworkStats mLastUidDetailSnapshot;
+
     /** Must be set in factory by calling #setHandler. */
     private Handler mHandler;
     private Handler.Callback mHandlerCallback;
@@ -812,15 +828,39 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
     @Override
     public NetworkStats getDetailedUidStats(String[] requiredIfaces) {
         try {
+            // Get the latest snapshot from NetworkStatsFactory.
+            // TODO: Querying for INTERFACES_ALL may incur performance penalty. Consider restricting
+            // this to limited set of ifaces.
+            NetworkStats uidDetailStats = getNetworkStatsUidDetail(INTERFACES_ALL);
+
+            // Migrate traffic from VPN UID over delta and update mTunAdjustedStats.
+            NetworkStats result;
+            synchronized (mStatsLock) {
+                migrateTunTraffic(uidDetailStats, mVpnInfos);
+                result = mTunAdjustedStats.clone();
+            }
+
+            // Apply filter based on ifacesToQuery.
             final String[] ifacesToQuery =
                     NetworkStatsFactory.augmentWithStackedInterfaces(requiredIfaces);
-            return getNetworkStatsUidDetail(ifacesToQuery);
+            result.filter(UID_ALL, ifacesToQuery, TAG_ALL);
+            return result;
         } catch (RemoteException e) {
             Log.wtf(TAG, "Error compiling UID stats", e);
             return new NetworkStats(0L, 0);
         }
     }
 
+    @VisibleForTesting
+    NetworkStats getTunAdjustedStats() {
+        synchronized (mStatsLock) {
+            if (mTunAdjustedStats == null) {
+                return null;
+            }
+            return mTunAdjustedStats.clone();
+        }
+    }
+
     @Override
     public String[] getMobileIfaces() {
         return mMobileIfaces;
@@ -1295,6 +1335,34 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         // a race condition between the service handler thread and the observer's
         mStatsObservers.updateStats(xtSnapshot, uidSnapshot, new ArrayMap<>(mActiveIfaces),
                 new ArrayMap<>(mActiveUidIfaces), vpnArray, currentTime);
+
+        migrateTunTraffic(uidSnapshot, vpnArray);
+    }
+
+    /**
+     * Updates {@link #mTunAdjustedStats} with the delta containing traffic migrated off of VPNs.
+     */
+    @GuardedBy("mStatsLock")
+    private void migrateTunTraffic(NetworkStats uidDetailStats, VpnInfo[] vpnInfoArray) {
+        if (mTunAdjustedStats == null) {
+            // Either device booted or system server restarted, hence traffic cannot be migrated
+            // correctly without knowing the past state of VPN's underlying networks.
+            mTunAdjustedStats = uidDetailStats;
+            mLastUidDetailSnapshot = uidDetailStats;
+            return;
+        }
+        // Migrate delta traffic from VPN to other apps.
+        NetworkStats delta = uidDetailStats.subtract(mLastUidDetailSnapshot);
+        for (VpnInfo info : vpnInfoArray) {
+            delta.migrateTun(info.ownerUid, info.vpnIface, info.underlyingIfaces);
+        }
+        // Filter out debug entries as that may lead to over counting.
+        delta.filterDebugEntries();
+        // Update #mTunAdjustedStats with migrated delta.
+        mTunAdjustedStats.combineAllValues(delta);
+        mTunAdjustedStats.setElapsedRealtime(uidDetailStats.getElapsedRealtime());
+        // Update last snapshot.
+        mLastUidDetailSnapshot = uidDetailStats;
     }
 
     /**