Merge "reduce log level when no interface is present for rate limiting"
diff --git a/Tethering/tests/privileged/src/com/android/net/module/util/BpfBitmapTest.java b/Tethering/tests/privileged/src/com/android/net/module/util/BpfBitmapTest.java
new file mode 100644
index 0000000..2112396
--- /dev/null
+++ b/Tethering/tests/privileged/src/com/android/net/module/util/BpfBitmapTest.java
@@ -0,0 +1,103 @@
+/*
+ * Copyright (C) 2022 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 com.android.networkstack.tethering;
+
+import static com.android.networkstack.tethering.util.TetheringUtils.getTetheringJniLibraryName;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.net.module.util.BpfBitmap;
+import com.android.testutils.DevSdkIgnoreRunner;
+
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@SmallTest
+@RunWith(DevSdkIgnoreRunner.class)
+public final class BpfBitmapTest {
+    private static final String TEST_BITMAP_PATH =
+            "/sys/fs/bpf/tethering/map_test_bitmap";
+
+    private static final int mTestData[] = {0,1,2,6,63,64,72};
+    private BpfBitmap mTestBitmap;
+
+    @Before
+    public void setUp() throws Exception {
+        mTestBitmap = new BpfBitmap(TEST_BITMAP_PATH);
+        mTestBitmap.clear();
+        assertTrue(mTestBitmap.isEmpty());
+    }
+
+    @Test
+    public void testSet() throws Exception {
+        for (int i : mTestData) {
+            mTestBitmap.set(i);
+            assertFalse(mTestBitmap.isEmpty());
+            assertTrue(mTestBitmap.get(i));
+            // Check that the next item in the bitmap is unset since test data is in
+            // ascending order.
+            assertFalse(mTestBitmap.get(i + 1));
+        }
+    }
+
+    @Test
+    public void testSetThenUnset() throws Exception {
+        for (int i : mTestData) {
+            mTestBitmap.set(i);
+            assertFalse(mTestBitmap.isEmpty());
+            assertTrue(mTestBitmap.get(i));
+            // Since test unsets all test data during each iteration, ensure all other
+            // bit are unset.
+            for (int j = 0; j < 128; ++j) if (j != i) assertFalse(mTestBitmap.get(j));
+            mTestBitmap.unset(i);
+        }
+    }
+
+    @Test
+    public void testSetAllThenUnsetAll() throws Exception {
+        for (int i : mTestData) {
+            mTestBitmap.set(i);
+        }
+
+        for (int i : mTestData) {
+            mTestBitmap.unset(i);
+            if (i < mTestData.length)
+                assertFalse(mTestBitmap.isEmpty());
+            assertFalse(mTestBitmap.get(i));
+        }
+        assertTrue(mTestBitmap.isEmpty());
+    }
+
+    @Test
+    public void testClear() throws Exception {
+        for (int i = 0; i < 128; ++i) {
+            mTestBitmap.set(i);
+        }
+        assertFalse(mTestBitmap.isEmpty());
+        mTestBitmap.clear();
+        assertTrue(mTestBitmap.isEmpty());
+    }
+}
diff --git a/bpf_progs/test.c b/bpf_progs/test.c
index a76e346..c9c73f1 100644
--- a/bpf_progs/test.c
+++ b/bpf_progs/test.c
@@ -28,6 +28,9 @@
 // Used only by TetheringPrivilegedTests, not by production code.
 DEFINE_BPF_MAP_GRW(tether_downstream6_map, HASH, TetherDownstream6Key, Tether6Value, 16,
                    AID_NETWORK_STACK)
+// Used only by BpfBitmapTest, not by production code.
+DEFINE_BPF_MAP_GRW(bitmap, ARRAY, int, uint64_t, 2,
+                   AID_NETWORK_STACK)
 
 DEFINE_BPF_PROG_KVER("xdp/drop_ipv4_udp_ether", AID_ROOT, AID_NETWORK_STACK,
                       xdp_test, KVER(5, 9, 0))
diff --git a/service/jni/com_android_server_BpfNetMaps.cpp b/service/jni/com_android_server_BpfNetMaps.cpp
index 85cfc09..f13c68d 100644
--- a/service/jni/com_android_server_BpfNetMaps.cpp
+++ b/service/jni/com_android_server_BpfNetMaps.cpp
@@ -24,6 +24,7 @@
 #include <nativehelper/JNIHelp.h>
 #include <nativehelper/ScopedUtfChars.h>
 #include <nativehelper/ScopedPrimitiveArray.h>
+#include <netjniutils/netjniutils.h>
 #include <net/if.h>
 #include <vector>
 
@@ -108,6 +109,7 @@
     }
 
     size_t size = uids.size();
+    static_assert(sizeof(*(uids.get())) == sizeof(int32_t));
     std::vector<int32_t> data ((int32_t *)&uids[0], (int32_t*)&uids[size]);
     int res = mTc.replaceUidOwnerMap(chainName, isAllowlist, data);
     if (res) {
@@ -144,6 +146,7 @@
     }
 
     size_t size = uids.size();
+    static_assert(sizeof(*(uids.get())) == sizeof(int32_t));
     std::vector<int32_t> data ((int32_t *)&uids[0], (int32_t*)&uids[size]);
     Status status = mTc.addUidInterfaceRules(ifIndex, data);
     if (!isOk(status)) {
@@ -159,6 +162,7 @@
     }
 
     size_t size = uids.size();
+    static_assert(sizeof(*(uids.get())) == sizeof(int32_t));
     std::vector<int32_t> data ((int32_t *)&uids[0], (int32_t*)&uids[size]);
     Status status = mTc.removeUidInterfaceRules(data);
     if (!isOk(status)) {
@@ -186,6 +190,15 @@
     mTc.setPermissionForUids(permission, data);
 }
 
+static void native_dump(JNIEnv* env, jobject clazz, jobject javaFd, jboolean verbose) {
+    int fd = netjniutils::GetNativeFileDescriptor(env, javaFd);
+    if (fd < 0) {
+        jniThrowExceptionFmt(env, "java/io/IOException", "Invalid file descriptor");
+        return;
+    }
+    mTc.dump(fd, verbose);
+}
+
 /*
  * JNI registration.
  */
@@ -216,6 +229,8 @@
     (void*)native_swapActiveStatsMap},
     {"native_setPermissionForUids", "(I[I)V",
     (void*)native_setPermissionForUids},
+    {"native_dump", "(Ljava/io/FileDescriptor;Z)V",
+    (void*)native_dump},
 };
 // clang-format on
 
diff --git a/service/native/TrafficController.cpp b/service/native/TrafficController.cpp
index 1cbfd94..393ee0a 100644
--- a/service/native/TrafficController.cpp
+++ b/service/native/TrafficController.cpp
@@ -601,8 +601,10 @@
     }
 }
 
-void TrafficController::dump(DumpWriter& dw, bool verbose) {
+void TrafficController::dump(int fd, bool verbose) {
     std::lock_guard guard(mMutex);
+    DumpWriter dw(fd);
+
     ScopedIndent indentTop(dw);
     dw.println("TrafficController");
 
diff --git a/service/native/include/TrafficController.h b/service/native/include/TrafficController.h
index 6fe117f..79e75ac 100644
--- a/service/native/include/TrafficController.h
+++ b/service/native/include/TrafficController.h
@@ -62,7 +62,7 @@
     netdutils::Status updateOwnerMapEntry(UidOwnerMatchType match, uid_t uid, FirewallRule rule,
                                           FirewallType type) EXCLUDES(mMutex);
 
-    void dump(netdutils::DumpWriter& dw, bool verbose) EXCLUDES(mMutex);
+    void dump(int fd, bool verbose) EXCLUDES(mMutex);
 
     netdutils::Status replaceRulesInMap(UidOwnerMatchType match, const std::vector<int32_t>& uids)
             EXCLUDES(mMutex);
diff --git a/service/native/libs/libclat/bpfhelper.cpp b/service/native/libs/libclat/bpfhelper.cpp
index 6e230d0..00785ad 100644
--- a/service/native/libs/libclat/bpfhelper.cpp
+++ b/service/native/libs/libclat/bpfhelper.cpp
@@ -28,16 +28,6 @@
 #define DEVICEPREFIX "v4-"
 
 using android::base::unique_fd;
-using android::net::RAWIP;
-using android::net::getClatEgress4MapFd;
-using android::net::getClatIngress6MapFd;
-using android::net::getClatEgress4ProgFd;
-using android::net::getClatIngress6ProgFd;
-using android::net::tcQdiscAddDevClsact;
-using android::net::tcFilterAddDevEgressClatIpv4;
-using android::net::tcFilterAddDevIngressClatIpv6;
-using android::net::tcFilterDelDevEgressClatIpv4;
-using android::net::tcFilterDelDevIngressClatIpv6;
 using android::bpf::BpfMap;
 
 BpfMap<ClatEgress4Key, ClatEgress4Value> mClatEgress4Map;
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index ddee275..f2ca18b 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -16,6 +16,8 @@
 
 package com.android.server;
 
+import static android.system.OsConstants.EOPNOTSUPP;
+
 import android.net.INetd;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
@@ -24,6 +26,9 @@
 
 import com.android.modules.utils.build.SdkLevel;
 
+import java.io.FileDescriptor;
+import java.io.IOException;
+
 /**
  * BpfNetMaps is responsible for providing traffic controller relevant functionality.
  *
@@ -273,6 +278,23 @@
         native_setPermissionForUids(permissions, uids);
     }
 
+    /**
+     * Dump BPF maps
+     *
+     * @param fd file descriptor to output
+     * @throws IOException when file descriptor is invalid.
+     * @throws ServiceSpecificException when the method is called on an unsupported device.
+     */
+    public void dump(final FileDescriptor fd, boolean verbose)
+            throws IOException, ServiceSpecificException {
+        if (USE_NETD) {
+            throw new ServiceSpecificException(
+                    EOPNOTSUPP, "dumpsys connectivity trafficcontroller dump not available on pre-T"
+                    + " devices, use dumpsys netd trafficcontroller instead.");
+        }
+        native_dump(fd, verbose);
+    }
+
     private static native void native_init();
     private native int native_addNaughtyApp(int uid);
     private native int native_removeNaughtyApp(int uid);
@@ -285,4 +307,5 @@
     private native int native_removeUidInterfaceRules(int[] uids);
     private native int native_swapActiveStatsMap();
     private native void native_setPermissionForUids(int permissions, int[] uids);
+    private native void native_dump(FileDescriptor fd, boolean verbose);
 }
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 271dd17..d30b3de 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -313,6 +313,7 @@
     public static final String SHORT_ARG = "--short";
     private static final String NETWORK_ARG = "networks";
     private static final String REQUEST_ARG = "requests";
+    private static final String TRAFFICCONTROLLER_ARG = "trafficcontroller";
 
     private static final boolean DBG = true;
     private static final boolean DDBG = Log.isLoggable(TAG, Log.DEBUG);
@@ -645,8 +646,9 @@
      * Event for NetworkMonitor to inform ConnectivityService that the probe status has changed.
      * Both of the arguments are bitmasks, and the value of bits come from
      * INetworkMonitor.NETWORK_VALIDATION_PROBE_*.
-     * arg1 = A bitmask to describe which probes are completed.
-     * arg2 = A bitmask to describe which probes are successful.
+     * arg1 = unused
+     * arg2 = netId
+     * obj = A Pair of integers: the bitmasks of, respectively, completed and successful probes.
      */
     public static final int EVENT_PROBE_STATUS_CHANGED = 45;
 
@@ -1409,6 +1411,8 @@
                 // converting rateInBytesPerSecond from long to int is safe here because the
                 // setting's range is limited to INT_MAX.
                 // TODO: add long/uint64 support to tcFilterAddDevIngressPolice.
+                Log.i(TAG,
+                        "enableIngressRateLimit on " + iface + ": " + rateInBytesPerSecond + "B/s");
                 TcUtils.tcFilterAddDevIngressPolice(params.index, TC_PRIO_POLICE, (short) ETH_P_ALL,
                         (int) rateInBytesPerSecond, TC_POLICE_BPF_PROG_PATH);
             } catch (IOException e) {
@@ -1430,6 +1434,8 @@
                 return;
             }
             try {
+                Log.i(TAG,
+                        "disableIngressRateLimit on " + iface);
                 TcUtils.tcFilterDelDev(params.index, true, TC_PRIO_POLICE, (short) ETH_P_ALL);
             } catch (IOException e) {
                 loge("TcUtils.tcFilterDelDev(ifaceIndex=" + params.index
@@ -1752,7 +1758,7 @@
 
         // Watch for ingress rate limit changes.
         mSettingsObserver.observe(
-                Settings.Secure.getUriFor(
+                Settings.Global.getUriFor(
                         ConnectivitySettingsManager.INGRESS_RATE_LIMIT_BYTES_PER_SECOND),
                 EVENT_INGRESS_RATE_LIMIT_CHANGED);
     }
@@ -3212,6 +3218,10 @@
         } else if (CollectionUtils.contains(args, REQUEST_ARG)) {
             dumpNetworkRequests(pw);
             return;
+        } else if (CollectionUtils.contains(args, TRAFFICCONTROLLER_ARG)) {
+            boolean verbose = !CollectionUtils.contains(args, SHORT_ARG);
+            dumpTrafficController(pw, fd, verbose);
+            return;
         }
 
         pw.print("NetworkProviders for:");
@@ -3429,6 +3439,17 @@
         }
     }
 
+    private void dumpTrafficController(IndentingPrintWriter pw, final FileDescriptor fd,
+            boolean verbose) {
+        try {
+            mBpfNetMaps.dump(fd, verbose);
+        } catch (ServiceSpecificException e) {
+            pw.println(e.getMessage());
+        } catch (IOException e) {
+            loge("Dump BPF maps failed, " + e);
+        }
+    }
+
     private void dumpAllRequestInfoLogsToLogcat() {
         try (PrintWriter logPw = new PrintWriter(new Writer() {
             @Override
@@ -3602,19 +3623,21 @@
         }
 
         private boolean maybeHandleNetworkMonitorMessage(Message msg) {
+            final int netId = msg.arg2;
+            final NetworkAgentInfo nai = getNetworkAgentInfoForNetId(netId);
             switch (msg.what) {
                 default:
                     return false;
                 case EVENT_PROBE_STATUS_CHANGED: {
-                    final Integer netId = (Integer) msg.obj;
-                    final NetworkAgentInfo nai = getNetworkAgentInfoForNetId(netId);
                     if (nai == null) {
                         break;
                     }
+                    final int probesCompleted = ((Pair<Integer, Integer>) msg.obj).first;
+                    final int probesSucceeded = ((Pair<Integer, Integer>) msg.obj).second;
                     final boolean probePrivateDnsCompleted =
-                            ((msg.arg1 & NETWORK_VALIDATION_PROBE_PRIVDNS) != 0);
+                            ((probesCompleted & NETWORK_VALIDATION_PROBE_PRIVDNS) != 0);
                     final boolean privateDnsBroken =
-                            ((msg.arg2 & NETWORK_VALIDATION_PROBE_PRIVDNS) == 0);
+                            ((probesSucceeded & NETWORK_VALIDATION_PROBE_PRIVDNS) == 0);
                     if (probePrivateDnsCompleted) {
                         if (nai.networkCapabilities.isPrivateDnsBroken() != privateDnsBroken) {
                             nai.networkCapabilities.setPrivateDnsBroken(privateDnsBroken);
@@ -3641,7 +3664,6 @@
                 case EVENT_NETWORK_TESTED: {
                     final NetworkTestedResults results = (NetworkTestedResults) msg.obj;
 
-                    final NetworkAgentInfo nai = getNetworkAgentInfoForNetId(results.mNetId);
                     if (nai == null) break;
 
                     handleNetworkTested(nai, results.mTestResult,
@@ -3649,9 +3671,7 @@
                     break;
                 }
                 case EVENT_PROVISIONING_NOTIFICATION: {
-                    final int netId = msg.arg2;
                     final boolean visible = toBool(msg.arg1);
-                    final NetworkAgentInfo nai = getNetworkAgentInfoForNetId(netId);
                     // If captive portal status has changed, update capabilities or disconnect.
                     if (nai != null && (visible != nai.lastCaptivePortalDetected)) {
                         nai.lastCaptivePortalDetected = visible;
@@ -3685,14 +3705,12 @@
                     break;
                 }
                 case EVENT_PRIVATE_DNS_CONFIG_RESOLVED: {
-                    final NetworkAgentInfo nai = getNetworkAgentInfoForNetId(msg.arg2);
                     if (nai == null) break;
 
                     updatePrivateDns(nai, (PrivateDnsConfig) msg.obj);
                     break;
                 }
                 case EVENT_CAPPORT_DATA_CHANGED: {
-                    final NetworkAgentInfo nai = getNetworkAgentInfoForNetId(msg.arg2);
                     if (nai == null) break;
                     handleCapportApiDataUpdate(nai, (CaptivePortalData) msg.obj);
                     break;
@@ -3832,6 +3850,7 @@
             // the same looper so messages will be processed in sequence.
             final Message msg = mTrackerHandler.obtainMessage(
                     EVENT_NETWORK_TESTED,
+                    0, mNetId,
                     new NetworkTestedResults(
                             mNetId, p.result, p.timestampMillis, p.redirectUrl));
             mTrackerHandler.sendMessage(msg);
@@ -3869,7 +3888,7 @@
         public void notifyProbeStatusChanged(int probesCompleted, int probesSucceeded) {
             mTrackerHandler.sendMessage(mTrackerHandler.obtainMessage(
                     EVENT_PROBE_STATUS_CHANGED,
-                    probesCompleted, probesSucceeded, new Integer(mNetId)));
+                    0, mNetId, new Pair<>(probesCompleted, probesSucceeded)));
         }
 
         @Override
@@ -10684,8 +10703,11 @@
     }
 
     private boolean canNetworkBeRateLimited(@NonNull final NetworkAgentInfo networkAgent) {
-        if (!networkAgent.networkCapabilities.hasCapability(NET_CAPABILITY_INTERNET)) {
-            // rate limits only apply to networks that provide internet connectivity.
+        final NetworkCapabilities agentCaps = networkAgent.networkCapabilities;
+        // Only test networks (they cannot hold NET_CAPABILITY_INTERNET) and networks that provide
+        // internet connectivity can be rate limited.
+        if (!agentCaps.hasCapability(NET_CAPABILITY_INTERNET) && !agentCaps.hasTransport(
+                TRANSPORT_TEST)) {
             return false;
         }
 
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
index 5778b0d..7150918 100755
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
@@ -334,9 +334,8 @@
             @Nullable ProxyInfo proxyInfo,
             @Nullable ArrayList<Network> underlyingNetworks, boolean isAlwaysMetered)
             throws Exception {
-        startVpn(addresses, routes, new String[0] /* excludedRoutes */, allowedApplications,
-                disallowedApplications, proxyInfo, underlyingNetworks, isAlwaysMetered,
-                false /* addRoutesByIpPrefix */);
+        startVpn(addresses, routes, excludedRoutes, allowedApplications, disallowedApplications,
+                proxyInfo, underlyingNetworks, isAlwaysMetered, false /* addRoutesByIpPrefix */);
     }
 
     private void startVpn(
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index ea64252..8f6ff19 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -2174,7 +2174,7 @@
     private void waitForAvailable(
             @NonNull final TestableNetworkCallback cb, @NonNull final Network expectedNetwork) {
         cb.expectAvailableCallbacks(expectedNetwork, false /* suspended */,
-                true /* validated */,
+                null /* validated */,
                 false /* blocked */, NETWORK_CALLBACK_TIMEOUT_MS);
     }
 
@@ -3259,6 +3259,19 @@
         assertTrue(dumpOutput, dumpOutput.contains("Active default network"));
     }
 
+    @Test @IgnoreUpTo(SC_V2)
+    public void testDumpBpfNetMaps() throws Exception {
+        final String[] args = new String[] {"--short", "trafficcontroller"};
+        String dumpOutput = DumpTestUtils.dumpServiceWithShellPermission(
+                Context.CONNECTIVITY_SERVICE, args);
+        assertTrue(dumpOutput, dumpOutput.contains("TrafficController"));
+        assertFalse(dumpOutput, dumpOutput.contains("BPF map content"));
+
+        dumpOutput = DumpTestUtils.dumpServiceWithShellPermission(
+                Context.CONNECTIVITY_SERVICE, args[1]);
+        assertTrue(dumpOutput, dumpOutput.contains("BPF map content"));
+    }
+
     private void unregisterRegisteredCallbacks() {
         for (NetworkCallback callback: mRegisteredCallbacks) {
             mCm.unregisterNetworkCallback(callback);
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 16d71d6..aa4e4bb 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -66,6 +66,7 @@
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
@@ -78,6 +79,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -108,6 +110,7 @@
 import android.os.PowerManager;
 import android.os.SimpleClock;
 import android.provider.Settings;
+import android.system.ErrnoException;
 import android.telephony.TelephonyManager;
 
 import androidx.annotation.Nullable;
@@ -125,6 +128,7 @@
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
 import com.android.testutils.HandlerUtils;
+import com.android.testutils.TestBpfMap;
 import com.android.testutils.TestableNetworkStatsProviderBinder;
 
 import libcore.testing.io.TestIoUtils;
@@ -143,6 +147,7 @@
 import java.time.ZoneOffset;
 import java.util.Objects;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Tests for {@link NetworkStatsService}.
@@ -198,11 +203,16 @@
     private HandlerThread mHandlerThread;
     @Mock
     private LocationPermissionChecker mLocationPermissionChecker;
-    private @Mock IBpfMap<U32, U8> mUidCounterSetMap;
-    private @Mock IBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap;
-    private @Mock IBpfMap<StatsMapKey, StatsMapValue> mStatsMapA;
-    private @Mock IBpfMap<StatsMapKey, StatsMapValue> mStatsMapB;
-    private @Mock IBpfMap<UidStatsMapKey, StatsMapValue> mAppUidStatsMap;
+    private TestBpfMap<U32, U8> mUidCounterSetMap = spy(new TestBpfMap<>(U32.class, U8.class));
+
+    private TestBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap = new TestBpfMap<>(
+            CookieTagMapKey.class, CookieTagMapValue.class);
+    private TestBpfMap<StatsMapKey, StatsMapValue> mStatsMapA = new TestBpfMap<>(StatsMapKey.class,
+            StatsMapValue.class);
+    private TestBpfMap<StatsMapKey, StatsMapValue> mStatsMapB = new TestBpfMap<>(StatsMapKey.class,
+            StatsMapValue.class);
+    private TestBpfMap<UidStatsMapKey, StatsMapValue> mAppUidStatsMap = new TestBpfMap<>(
+            UidStatsMapKey.class, StatsMapValue.class);
 
     private NetworkStatsService mService;
     private INetworkStatsSession mSession;
@@ -1922,4 +1932,70 @@
     private void waitForIdle() {
         HandlerUtils.waitForIdle(mHandlerThread, WAIT_TIMEOUT);
     }
+
+    private boolean cookieTagMapContainsUid(int uid) throws ErrnoException {
+        final AtomicBoolean found = new AtomicBoolean();
+        mCookieTagMap.forEach((k, v) -> {
+            if (v.uid == uid) {
+                found.set(true);
+            }
+        });
+        return found.get();
+    }
+
+    private static <K extends StatsMapKey, V extends StatsMapValue> boolean statsMapContainsUid(
+            TestBpfMap<K, V> map, int uid) throws ErrnoException {
+        final AtomicBoolean found = new AtomicBoolean();
+        map.forEach((k, v) -> {
+            if (k.uid == uid) {
+                found.set(true);
+            }
+        });
+        return found.get();
+    }
+
+    private void initBpfMapsWithTagData(int uid) throws ErrnoException {
+        // key needs to be unique, use some offset from uid.
+        mCookieTagMap.insertEntry(new CookieTagMapKey(1000 + uid), new CookieTagMapValue(uid, 1));
+        mCookieTagMap.insertEntry(new CookieTagMapKey(2000 + uid), new CookieTagMapValue(uid, 2));
+
+        mStatsMapA.insertEntry(new StatsMapKey(uid, 1, 0, 10), new StatsMapValue(5, 5000, 3, 3000));
+        mStatsMapA.insertEntry(new StatsMapKey(uid, 2, 0, 10), new StatsMapValue(5, 5000, 3, 3000));
+
+        mStatsMapB.insertEntry(new StatsMapKey(uid, 1, 0, 10), new StatsMapValue(0, 0, 0, 0));
+
+        mAppUidStatsMap.insertEntry(new UidStatsMapKey(uid), new StatsMapValue(10, 10000, 6, 6000));
+
+        mUidCounterSetMap.insertEntry(new U32(uid), new U8((short) 1));
+
+        assertTrue(cookieTagMapContainsUid(uid));
+        assertTrue(statsMapContainsUid(mStatsMapA, uid));
+        assertTrue(statsMapContainsUid(mStatsMapB, uid));
+        assertTrue(mAppUidStatsMap.containsKey(new UidStatsMapKey(uid)));
+        assertTrue(mUidCounterSetMap.containsKey(new U32(uid)));
+    }
+
+    @Test
+    public void testRemovingUidRemovesTagDataForUid() throws ErrnoException {
+        initBpfMapsWithTagData(UID_BLUE);
+        initBpfMapsWithTagData(UID_RED);
+
+        final Intent intent = new Intent(ACTION_UID_REMOVED);
+        intent.putExtra(EXTRA_UID, UID_BLUE);
+        mServiceContext.sendBroadcast(intent);
+
+        // assert that all UID_BLUE related tag data has been removed from the maps.
+        assertFalse(cookieTagMapContainsUid(UID_BLUE));
+        assertFalse(statsMapContainsUid(mStatsMapA, UID_BLUE));
+        assertFalse(statsMapContainsUid(mStatsMapB, UID_BLUE));
+        assertFalse(mAppUidStatsMap.containsKey(new UidStatsMapKey(UID_BLUE)));
+        assertFalse(mUidCounterSetMap.containsKey(new U32(UID_BLUE)));
+
+        // assert that UID_RED related tag data is still in the maps.
+        assertTrue(cookieTagMapContainsUid(UID_RED));
+        assertTrue(statsMapContainsUid(mStatsMapA, UID_RED));
+        assertTrue(statsMapContainsUid(mStatsMapB, UID_RED));
+        assertTrue(mAppUidStatsMap.containsKey(new UidStatsMapKey(UID_RED)));
+        assertTrue(mUidCounterSetMap.containsKey(new U32(UID_RED)));
+    }
 }