Merge "Add Data Saver utils in Firewall test helper class" into main
diff --git a/framework/src/android/net/BpfNetMapsConstants.java b/framework/src/android/net/BpfNetMapsConstants.java
index 36848e7..8086809 100644
--- a/framework/src/android/net/BpfNetMapsConstants.java
+++ b/framework/src/android/net/BpfNetMapsConstants.java
@@ -16,6 +16,15 @@
 
 package android.net;
 
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_DOZABLE;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_LOW_POWER_STANDBY;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_1;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_2;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_3;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_POWERSAVE;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_RESTRICTED;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_STANDBY;
+
 import android.util.Pair;
 
 import com.android.net.module.util.Struct;
@@ -66,7 +75,6 @@
     public static final long OEM_DENY_1_MATCH = (1 << 9);
     public static final long OEM_DENY_2_MATCH = (1 << 10);
     public static final long OEM_DENY_3_MATCH = (1 << 11);
-    // LINT.ThenChange(../../../../bpf_progs/netd.h)
 
     public static final List<Pair<Long, String>> MATCH_LIST = Arrays.asList(
             Pair.create(HAPPY_BOX_MATCH, "HAPPY_BOX_MATCH"),
@@ -82,4 +90,29 @@
             Pair.create(OEM_DENY_2_MATCH, "OEM_DENY_2_MATCH"),
             Pair.create(OEM_DENY_3_MATCH, "OEM_DENY_3_MATCH")
     );
+
+    /**
+     * List of all firewall allow chains.
+     *
+     * Allow chains mean the firewall denies all uids by default, uids must be explicitly allowed.
+     */
+    public static final List<Integer> ALLOW_CHAINS = List.of(
+            FIREWALL_CHAIN_DOZABLE,
+            FIREWALL_CHAIN_POWERSAVE,
+            FIREWALL_CHAIN_RESTRICTED,
+            FIREWALL_CHAIN_LOW_POWER_STANDBY
+    );
+
+    /**
+     * List of all firewall deny chains.
+     *
+     * Deny chains mean the firewall allows all uids by default, uids must be explicitly denied.
+     */
+    public static final List<Integer> DENY_CHAINS = List.of(
+            FIREWALL_CHAIN_STANDBY,
+            FIREWALL_CHAIN_OEM_DENY_1,
+            FIREWALL_CHAIN_OEM_DENY_2,
+            FIREWALL_CHAIN_OEM_DENY_3
+    );
+    // LINT.ThenChange(../../../../bpf_progs/netd.h)
 }
diff --git a/framework/src/android/net/BpfNetMapsReader.java b/framework/src/android/net/BpfNetMapsReader.java
index 49e874a..9bce9cd 100644
--- a/framework/src/android/net/BpfNetMapsReader.java
+++ b/framework/src/android/net/BpfNetMapsReader.java
@@ -57,10 +57,42 @@
     private final IBpfMap<S32, UidOwnerValue> mUidOwnerMap;
     private final Dependencies mDeps;
 
-    public BpfNetMapsReader() {
+    // Bitmaps for calculating whether a given uid is blocked by firewall chains.
+    private static final long sMaskDropIfSet;
+    private static final long sMaskDropIfUnset;
+
+    static {
+        long maskDropIfSet = 0L;
+        long maskDropIfUnset = 0L;
+
+        for (int chain : BpfNetMapsConstants.ALLOW_CHAINS) {
+            final long match = getMatchByFirewallChain(chain);
+            maskDropIfUnset |= match;
+        }
+        for (int chain : BpfNetMapsConstants.DENY_CHAINS) {
+            final long match = getMatchByFirewallChain(chain);
+            maskDropIfSet |= match;
+        }
+        sMaskDropIfSet = maskDropIfSet;
+        sMaskDropIfUnset = maskDropIfUnset;
+    }
+
+    private static class SingletonHolder {
+        static final BpfNetMapsReader sInstance = new BpfNetMapsReader();
+    }
+
+    @NonNull
+    public static BpfNetMapsReader getInstance() {
+        return SingletonHolder.sInstance;
+    }
+
+    private BpfNetMapsReader() {
         this(new Dependencies());
     }
 
+    // While the production code uses the singleton to optimize for performance and deal with
+    // concurrent access, the test needs to use a non-static approach for dependency injection and
+    // mocking virtual bpf maps.
     @VisibleForTesting
     public BpfNetMapsReader(@NonNull Dependencies deps) {
         if (!SdkLevel.isAtLeastT()) {
@@ -176,4 +208,33 @@
                     "Unable to get uid rule status: " + Os.strerror(e.errno));
         }
     }
+
+    /**
+     * Return whether the network is blocked by firewall chains for the given uid.
+     *
+     * @param uid The target uid.
+     *
+     * @return True if the network is blocked. Otherwise, false.
+     * @throws ServiceSpecificException if the read fails.
+     *
+     * @hide
+     */
+    public boolean isUidBlockedByFirewallChains(final int uid) {
+        throwIfPreT("isUidBlockedByFirewallChains is not available on pre-T devices");
+
+        final long uidRuleConfig;
+        final long uidMatch;
+        try {
+            uidRuleConfig = mConfigurationMap.getValue(UID_RULES_CONFIGURATION_KEY).val;
+            final UidOwnerValue value = mUidOwnerMap.getValue(new S32(uid));
+            uidMatch = (value != null) ? value.rule : 0L;
+        } catch (ErrnoException e) {
+            throw new ServiceSpecificException(e.errno,
+                    "Unable to get firewall chain status: " + Os.strerror(e.errno));
+        }
+
+        final boolean blockedByAllowChains = 0 != (uidRuleConfig & ~uidMatch & sMaskDropIfUnset);
+        final boolean blockedByDenyChains = 0 != (uidRuleConfig & uidMatch & sMaskDropIfSet);
+        return blockedByAllowChains || blockedByDenyChains;
+    }
 }
diff --git a/framework/src/android/net/BpfNetMapsUtils.java b/framework/src/android/net/BpfNetMapsUtils.java
index 28d5891..e9c9137 100644
--- a/framework/src/android/net/BpfNetMapsUtils.java
+++ b/framework/src/android/net/BpfNetMapsUtils.java
@@ -16,6 +16,8 @@
 
 package android.net;
 
+import static android.net.BpfNetMapsConstants.ALLOW_CHAINS;
+import static android.net.BpfNetMapsConstants.DENY_CHAINS;
 import static android.net.BpfNetMapsConstants.DOZABLE_MATCH;
 import static android.net.BpfNetMapsConstants.LOW_POWER_STANDBY_MATCH;
 import static android.net.BpfNetMapsConstants.MATCH_LIST;
@@ -82,26 +84,18 @@
     }
 
     /**
-     * Get if the chain is allow list or not.
+     * Get whether the chain is an allow-list or a deny-list.
      *
      * ALLOWLIST means the firewall denies all by default, uids must be explicitly allowed
-     * DENYLIST means the firewall allows all by default, uids must be explicitly denyed
+     * DENYLIST means the firewall allows all by default, uids must be explicitly denied
      */
     public static boolean isFirewallAllowList(final int chain) {
-        switch (chain) {
-            case FIREWALL_CHAIN_DOZABLE:
-            case FIREWALL_CHAIN_POWERSAVE:
-            case FIREWALL_CHAIN_RESTRICTED:
-            case FIREWALL_CHAIN_LOW_POWER_STANDBY:
-                return true;
-            case FIREWALL_CHAIN_STANDBY:
-            case FIREWALL_CHAIN_OEM_DENY_1:
-            case FIREWALL_CHAIN_OEM_DENY_2:
-            case FIREWALL_CHAIN_OEM_DENY_3:
-                return false;
-            default:
-                throw new ServiceSpecificException(EINVAL, "Invalid firewall chain: " + chain);
+        if (ALLOW_CHAINS.contains(chain)) {
+            return true;
+        } else if (DENY_CHAINS.contains(chain)) {
+            return false;
         }
+        throw new ServiceSpecificException(EINVAL, "Invalid firewall chain: " + chain);
     }
 
     /**
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 32058a4..ad76012 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -6198,6 +6198,36 @@
         }
     }
 
+    /**
+     * Return whether the network is blocked for the given uid.
+     *
+     * Similar to {@link NetworkPolicyManager#isUidNetworkingBlocked}, but directly reads the BPF
+     * maps and therefore considerably faster. For use by the NetworkStack process only.
+     *
+     * @param uid The target uid.
+     * @return True if all networking is blocked. Otherwise, false.
+     * @throws IllegalStateException if the map cannot be opened.
+     * @throws ServiceSpecificException if the read fails.
+     * @hide
+     */
+    // This isn't protected by a standard Android permission since it can't
+    // afford to do IPC for performance reasons. Instead, the access control
+    // is provided by linux file group permission AID_NET_BW_ACCT and the
+    // selinux context fs_bpf_net*.
+    // Only the system server process and the network stack have access.
+    // TODO: Expose api when ready.
+    // @SystemApi(client = MODULE_LIBRARIES)
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)  // BPF maps were only mainlined in T
+    @RequiresPermission(NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK)
+    public boolean isUidNetworkingBlocked(int uid) {
+        final BpfNetMapsReader reader = BpfNetMapsReader.getInstance();
+
+        return reader.isUidBlockedByFirewallChains(uid);
+
+        // TODO: If isNetworkMetered is true, check the data saver switch, penalty box
+        //  and happy box rules.
+    }
+
     /** @hide */
     public IBinder getCompanionDeviceManagerProxyService() {
         try {
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index cc3f019..c74f229 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -519,9 +519,9 @@
         }
     }
 
+    // TODO: Use a Handler instead of a StateMachine since there are no state changes.
     private class NsdStateMachine extends StateMachine {
 
-        private final DefaultState mDefaultState = new DefaultState();
         private final EnabledState mEnabledState = new EnabledState();
 
         @Override
@@ -591,124 +591,12 @@
 
         NsdStateMachine(String name, Handler handler) {
             super(name, handler);
-            addState(mDefaultState);
-                addState(mEnabledState, mDefaultState);
+            addState(mEnabledState);
             State initialState = mEnabledState;
             setInitialState(initialState);
             setLogRecSize(25);
         }
 
-        class DefaultState extends State {
-            @Override
-            public boolean processMessage(Message msg) {
-                final ClientInfo cInfo;
-                final int clientRequestId = msg.arg2;
-                switch (msg.what) {
-                    case NsdManager.REGISTER_CLIENT:
-                        final ConnectorArgs arg = (ConnectorArgs) msg.obj;
-                        final INsdManagerCallback cb = arg.callback;
-                        try {
-                            cb.asBinder().linkToDeath(arg.connector, 0);
-                            final String tag = "Client" + arg.uid + "-" + mClientNumberId++;
-                            final NetworkNsdReportedMetrics metrics =
-                                    mDeps.makeNetworkNsdReportedMetrics(
-                                            (int) mClock.elapsedRealtime());
-                            cInfo = new ClientInfo(cb, arg.uid, arg.useJavaBackend,
-                                    mServiceLogs.forSubComponent(tag), metrics);
-                            mClients.put(arg.connector, cInfo);
-                        } catch (RemoteException e) {
-                            Log.w(TAG, "Client request id " + clientRequestId
-                                    + " has already died");
-                        }
-                        break;
-                    case NsdManager.UNREGISTER_CLIENT:
-                        final NsdServiceConnector connector = (NsdServiceConnector) msg.obj;
-                        cInfo = mClients.remove(connector);
-                        if (cInfo != null) {
-                            cInfo.expungeAllRequests();
-                            if (cInfo.isPreSClient()) {
-                                mLegacyClientCount -= 1;
-                            }
-                        }
-                        maybeStopMonitoringSocketsIfNoActiveRequest();
-                        maybeScheduleStop();
-                        break;
-                    case NsdManager.DISCOVER_SERVICES:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cInfo.onDiscoverServicesFailedImmediately(clientRequestId,
-                                    NsdManager.FAILURE_INTERNAL_ERROR, true /* isLegacy */);
-                        }
-                       break;
-                    case NsdManager.STOP_DISCOVERY:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cInfo.onStopDiscoveryFailed(
-                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
-                        }
-                        break;
-                    case NsdManager.REGISTER_SERVICE:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cInfo.onRegisterServiceFailedImmediately(clientRequestId,
-                                    NsdManager.FAILURE_INTERNAL_ERROR, true /* isLegacy */);
-                        }
-                        break;
-                    case NsdManager.UNREGISTER_SERVICE:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cInfo.onUnregisterServiceFailed(
-                                    clientRequestId, NsdManager.FAILURE_INTERNAL_ERROR);
-                        }
-                        break;
-                    case NsdManager.RESOLVE_SERVICE:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cInfo.onResolveServiceFailedImmediately(clientRequestId,
-                                    NsdManager.FAILURE_INTERNAL_ERROR, true /* isLegacy */);
-                        }
-                        break;
-                    case NsdManager.STOP_RESOLUTION:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cInfo.onStopResolutionFailed(
-                                    clientRequestId, NsdManager.FAILURE_OPERATION_NOT_RUNNING);
-                        }
-                        break;
-                    case NsdManager.REGISTER_SERVICE_CALLBACK:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cInfo.onServiceInfoCallbackRegistrationFailed(
-                                    clientRequestId, NsdManager.FAILURE_BAD_PARAMETERS);
-                        }
-                        break;
-                    case NsdManager.DAEMON_CLEANUP:
-                        maybeStopDaemon();
-                        break;
-                    // This event should be only sent by the legacy (target SDK < S) clients.
-                    // Mark the sending client as legacy.
-                    case NsdManager.DAEMON_STARTUP:
-                        cInfo = getClientInfoForReply(msg);
-                        if (cInfo != null) {
-                            cancelStop();
-                            cInfo.setPreSClient();
-                            mLegacyClientCount += 1;
-                            maybeStartDaemon();
-                        }
-                        break;
-                    default:
-                        Log.e(TAG, "Unhandled " + msg);
-                        return NOT_HANDLED;
-                }
-                return HANDLED;
-            }
-
-            private ClientInfo getClientInfoForReply(Message msg) {
-                final ListenerArgs args = (ListenerArgs) msg.obj;
-                return mClients.get(args.connector);
-            }
-        }
-
         class EnabledState extends State {
             @Override
             public void enter() {
@@ -793,6 +681,11 @@
                 removeRequestMap(clientRequestId, transactionId, clientInfo);
             }
 
+            private ClientInfo getClientInfoForReply(Message msg) {
+                final ListenerArgs args = (ListenerArgs) msg.obj;
+                return mClients.get(args.connector);
+            }
+
             @Override
             public boolean processMessage(Message msg) {
                 final ClientInfo clientInfo;
@@ -1214,7 +1107,51 @@
                     case NsdManager.UNREGISTER_OFFLOAD_ENGINE:
                         mOffloadEngines.unregister((IOffloadEngine) msg.obj);
                         break;
+                    case NsdManager.REGISTER_CLIENT:
+                        final ConnectorArgs arg = (ConnectorArgs) msg.obj;
+                        final INsdManagerCallback cb = arg.callback;
+                        try {
+                            cb.asBinder().linkToDeath(arg.connector, 0);
+                            final String tag = "Client" + arg.uid + "-" + mClientNumberId++;
+                            final NetworkNsdReportedMetrics metrics =
+                                    mDeps.makeNetworkNsdReportedMetrics(
+                                            (int) mClock.elapsedRealtime());
+                            clientInfo = new ClientInfo(cb, arg.uid, arg.useJavaBackend,
+                                    mServiceLogs.forSubComponent(tag), metrics);
+                            mClients.put(arg.connector, clientInfo);
+                        } catch (RemoteException e) {
+                            Log.w(TAG, "Client request id " + clientRequestId
+                                    + " has already died");
+                        }
+                        break;
+                    case NsdManager.UNREGISTER_CLIENT:
+                        final NsdServiceConnector connector = (NsdServiceConnector) msg.obj;
+                        clientInfo = mClients.remove(connector);
+                        if (clientInfo != null) {
+                            clientInfo.expungeAllRequests();
+                            if (clientInfo.isPreSClient()) {
+                                mLegacyClientCount -= 1;
+                            }
+                        }
+                        maybeStopMonitoringSocketsIfNoActiveRequest();
+                        maybeScheduleStop();
+                        break;
+                    case NsdManager.DAEMON_CLEANUP:
+                        maybeStopDaemon();
+                        break;
+                    // This event should be only sent by the legacy (target SDK < S) clients.
+                    // Mark the sending client as legacy.
+                    case NsdManager.DAEMON_STARTUP:
+                        clientInfo = getClientInfoForReply(msg);
+                        if (clientInfo != null) {
+                            cancelStop();
+                            clientInfo.setPreSClient();
+                            mLegacyClientCount += 1;
+                            maybeStartDaemon();
+                        }
+                        break;
                     default:
+                        Log.wtf(TAG, "Unhandled " + msg);
                         return NOT_HANDLED;
                 }
                 return HANDLED;
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
index e07d380..42a6b0d 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
@@ -68,6 +68,8 @@
 
     @NonNull
     private final SharedLog mSharedLog;
+    @NonNull
+    private final byte[] mPacketCreationBuffer;
 
     /**
      * Callbacks called by {@link MdnsInterfaceAdvertiser} to report status updates.
@@ -205,6 +207,7 @@
         mCbHandler = new Handler(looper);
         mReplySender = deps.makeReplySender(sharedLog.getTag(), looper, socket,
                 packetCreationBuffer, sharedLog);
+        mPacketCreationBuffer = packetCreationBuffer;
         mAnnouncer = deps.makeMdnsAnnouncer(sharedLog.getTag(), looper, mReplySender,
                 mAnnouncingCallback, sharedLog);
         mProber = deps.makeMdnsProber(sharedLog.getTag(), looper, mReplySender, mProbingCallback,
@@ -390,12 +393,13 @@
      * @param serviceId The serviceId.
      * @return the raw offload payload
      */
+    @NonNull
     public byte[] getRawOffloadPayload(int serviceId) {
         try {
-            return MdnsUtils.createRawDnsPacket(mReplySender.getPacketCreationBuffer(),
+            return MdnsUtils.createRawDnsPacket(mPacketCreationBuffer,
                     mRecordRepository.getOffloadPacket(serviceId));
         } catch (IOException | IllegalArgumentException e) {
-            mSharedLog.wtf("Cannot create rawOffloadPacket: " + e.getMessage());
+            mSharedLog.wtf("Cannot create rawOffloadPacket: ", e);
             return new byte[0];
         }
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java b/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java
index abf5d99..ea3af5e 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java
@@ -99,11 +99,6 @@
         return PACKET_SENT;
     }
 
-    /** Get the packetCreationBuffer */
-    public byte[] getPacketCreationBuffer() {
-        return mPacketCreationBuffer;
-    }
-
     /**
      * Cancel all pending sends.
      */
diff --git a/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java b/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
index d0f3d9a..4d79f9d 100644
--- a/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
+++ b/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
@@ -35,6 +35,7 @@
 import java.nio.charset.Charset;
 import java.nio.charset.CharsetEncoder;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
 
@@ -211,9 +212,7 @@
         }
 
         final int len = writer.getWritePosition();
-        final byte[] outBuffer = new byte[len];
-        System.arraycopy(packetCreationBuffer, 0, outBuffer, 0, len);
-        return outBuffer;
+        return Arrays.copyOfRange(packetCreationBuffer, 0, len);
     }
 
     /**
diff --git a/service-t/src/com/android/server/net/NetworkStatsObservers.java b/service-t/src/com/android/server/net/NetworkStatsObservers.java
index 1cd670a..21cf351 100644
--- a/service-t/src/com/android/server/net/NetworkStatsObservers.java
+++ b/service-t/src/com/android/server/net/NetworkStatsObservers.java
@@ -142,6 +142,11 @@
 
     @VisibleForTesting
     protected Looper getHandlerLooperLocked() {
+        // TODO: Currently, callbacks are dispatched on this thread if the caller register
+        //  callback without supplying a Handler. To ensure that the service handler thread
+        //  is not blocked by client code, the observers must create their own thread. Once
+        //  all callbacks are dispatched outside of the handler thread, the service handler
+        //  thread can be used here.
         HandlerThread handlerThread = new HandlerThread(TAG);
         handlerThread.start();
         return handlerThread.getLooper();
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index 6ade124..14ab2a1 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -32,6 +32,7 @@
 import static android.net.BpfNetMapsConstants.UID_RULES_CONFIGURATION_KEY;
 import static android.net.BpfNetMapsUtils.PRE_T;
 import static android.net.BpfNetMapsUtils.getMatchByFirewallChain;
+import static android.net.BpfNetMapsUtils.isFirewallAllowList;
 import static android.net.BpfNetMapsUtils.matchToString;
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_DOZABLE;
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_LOW_POWER_STANDBY;
@@ -352,29 +353,6 @@
         mDeps = deps;
     }
 
-    /**
-     * Get if the chain is allow list or not.
-     *
-     * ALLOWLIST means the firewall denies all by default, uids must be explicitly allowed
-     * DENYLIST means the firewall allows all by default, uids must be explicitly denyed
-     */
-    public boolean isFirewallAllowList(final int chain) {
-        switch (chain) {
-            case FIREWALL_CHAIN_DOZABLE:
-            case FIREWALL_CHAIN_POWERSAVE:
-            case FIREWALL_CHAIN_RESTRICTED:
-            case FIREWALL_CHAIN_LOW_POWER_STANDBY:
-                return true;
-            case FIREWALL_CHAIN_STANDBY:
-            case FIREWALL_CHAIN_OEM_DENY_1:
-            case FIREWALL_CHAIN_OEM_DENY_2:
-            case FIREWALL_CHAIN_OEM_DENY_3:
-                return false;
-            default:
-                throw new ServiceSpecificException(EINVAL, "Invalid firewall chain: " + chain);
-        }
-    }
-
     private void maybeThrow(final int err, final String msg) {
         if (err != 0) {
             throw new ServiceSpecificException(err, msg + ": " + Os.strerror(err));
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index b31b94c..f52a1a2 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -136,6 +136,7 @@
 import android.content.pm.PackageManager;
 import android.content.res.XmlResourceParser;
 import android.database.ContentObserver;
+import android.net.BpfNetMapsUtils;
 import android.net.CaptivePortal;
 import android.net.CaptivePortalData;
 import android.net.ConnectionInfo;
@@ -5273,7 +5274,14 @@
 
     private boolean isNetworkPotentialSatisfier(
             @NonNull final NetworkAgentInfo candidate, @NonNull final NetworkRequestInfo nri) {
-        // listen requests won't keep up a network satisfying it. If this is not a multilayer
+        // While destroyed network sometimes satisfy requests (including occasionally newly
+        // satisfying requests), *potential* satisfiers are networks that might beat a current
+        // champion if they validate. As such, a destroyed network is never a potential satisfier,
+        // because it's never a good idea to keep a destroyed network in case it validates.
+        // For example, declaring it a potential satisfier would keep an unvalidated destroyed
+        // candidate after it's been replaced by another unvalidated network.
+        if (candidate.isDestroyed()) return false;
+        // Listen requests won't keep up a network satisfying it. If this is not a multilayer
         // request, return immediately. For multilayer requests, check to see if any of the
         // multilayer requests may have a potential satisfier.
         if (!nri.isMultilayerRequest() && (nri.mRequests.get(0).isListen()
@@ -5291,8 +5299,12 @@
             if (req.isListen() || req.isListenForBest()) {
                 continue;
             }
-            // If this Network is already the best Network for a request, or if
-            // there is hope for it to become one if it validated, then it is needed.
+            // If there is hope for this network might validate and subsequently become the best
+            // network for that request, then it is needed. Note that this network can't already
+            // be the best for this request, or it would be the current satisfier, and therefore
+            // there would be no need to call this method to find out if it is a *potential*
+            // satisfier ("unneeded", the only caller, only calls this if this network currently
+            // satisfies no request).
             if (candidate.satisfies(req)) {
                 // As soon as a network is found that satisfies a request, return. Specifically for
                 // multilayer requests, returning as soon as a NetworkAgentInfo satisfies a request
@@ -12678,7 +12690,7 @@
 
     private void closeSocketsForFirewallChainLocked(final int chain)
             throws ErrnoException, SocketException, InterruptedIOException {
-        if (mBpfNetMaps.isFirewallAllowList(chain)) {
+        if (BpfNetMapsUtils.isFirewallAllowList(chain)) {
             // Allowlist means the firewall denies all by default, uids must be explicitly allowed
             // So, close all non-system socket owned by uids that are not explicitly allowed
             Set<Range<Integer>> ranges = new ArraySet<>();
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 11345d3..cda8d06 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -692,8 +692,10 @@
 
     /**
      * Dump AutomaticOnOffKeepaliveTracker state.
+     * This should be only be called in ConnectivityService handler thread.
      */
     public void dump(IndentingPrintWriter pw) {
+        ensureRunningOnHandlerThread();
         mKeepaliveTracker.dump(pw);
         // Reading DeviceConfig will check if the calling uid and calling package name are the same.
         // Clear calling identity to align the calling uid and package so that it won't fail if cts
@@ -712,6 +714,9 @@
         pw.increaseIndent();
         mEventLog.reverseDump(pw);
         pw.decreaseIndent();
+
+        pw.println();
+        mKeepaliveStatsTracker.dump(pw);
     }
 
     /**
diff --git a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
index 0c2ed18..7a8b41b 100644
--- a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
@@ -34,6 +34,7 @@
 import android.telephony.SubscriptionInfo;
 import android.telephony.SubscriptionManager;
 import android.telephony.TelephonyManager;
+import android.util.IndentingPrintWriter;
 import android.util.Log;
 import android.util.SparseArray;
 import android.util.SparseIntArray;
@@ -73,6 +74,9 @@
 public class KeepaliveStatsTracker {
     private static final String TAG = KeepaliveStatsTracker.class.getSimpleName();
     private static final int INVALID_KEEPALIVE_ID = -1;
+    // 1 hour acceptable deviation in metrics collection duration time.
+    private static final long MAX_EXPECTED_DURATION_MS =
+            AutomaticOnOffKeepaliveTracker.METRICS_COLLECTION_DURATION_MS + 1 * 60 * 60 * 1_000L;
 
     @NonNull private final Handler mConnectivityServiceHandler;
     @NonNull private final Dependencies mDependencies;
@@ -709,6 +713,36 @@
         return mEnabled.get();
     }
 
+    /**
+     * Checks the DailykeepaliveInfoReported for the following:
+     * 1. total active durations/lifetimes <= total registered durations/lifetimes.
+     * 2. Total time in Durations == total time in Carrier lifetime stats
+     * 3. The total elapsed real time spent is within expectations.
+     */
+    @VisibleForTesting
+    public boolean allMetricsExpected(DailykeepaliveInfoReported dailyKeepaliveInfoReported) {
+        int totalRegistered = 0;
+        int totalActiveDurations = 0;
+        int totalTimeSpent = 0;
+        for (DurationForNumOfKeepalive durationForNumOfKeepalive: dailyKeepaliveInfoReported
+                .getDurationPerNumOfKeepalive().getDurationForNumOfKeepaliveList()) {
+            final int n = durationForNumOfKeepalive.getNumOfKeepalive();
+            totalRegistered += durationForNumOfKeepalive.getKeepaliveRegisteredDurationsMsec() * n;
+            totalActiveDurations += durationForNumOfKeepalive.getKeepaliveActiveDurationsMsec() * n;
+            totalTimeSpent += durationForNumOfKeepalive.getKeepaliveRegisteredDurationsMsec();
+        }
+        int totalLifetimes = 0;
+        int totalActiveLifetimes = 0;
+        for (KeepaliveLifetimeForCarrier keepaliveLifetimeForCarrier: dailyKeepaliveInfoReported
+                .getKeepaliveLifetimePerCarrier().getKeepaliveLifetimeForCarrierList()) {
+            totalLifetimes += keepaliveLifetimeForCarrier.getLifetimeMsec();
+            totalActiveLifetimes += keepaliveLifetimeForCarrier.getActiveLifetimeMsec();
+        }
+        return totalActiveDurations <= totalRegistered && totalActiveLifetimes <= totalLifetimes
+                && totalLifetimes == totalRegistered && totalActiveLifetimes == totalActiveDurations
+                && totalTimeSpent <= MAX_EXPECTED_DURATION_MS;
+    }
+
     /** Writes the stored metrics to ConnectivityStatsLog and resets. */
     public void writeAndResetMetrics() {
         ensureRunningOnHandlerThread();
@@ -724,9 +758,21 @@
         }
 
         final DailykeepaliveInfoReported dailyKeepaliveInfoReported = buildAndResetMetrics();
+        if (!allMetricsExpected(dailyKeepaliveInfoReported)) {
+            Log.wtf(TAG, "Unexpected metrics values: " + dailyKeepaliveInfoReported.toString());
+        }
         mDependencies.writeStats(dailyKeepaliveInfoReported);
     }
 
+    /** Dump KeepaliveStatsTracker state. */
+    public void dump(IndentingPrintWriter pw) {
+        ensureRunningOnHandlerThread();
+        pw.println("KeepaliveStatsTracker enabled: " + isEnabled());
+        pw.increaseIndent();
+        pw.println(buildKeepaliveMetrics().toString());
+        pw.decreaseIndent();
+    }
+
     private void ensureRunningOnHandlerThread() {
         if (mConnectivityServiceHandler.getLooper().getThread() != Thread.currentThread()) {
             throw new IllegalStateException(
diff --git a/staticlibs/device/com/android/net/module/util/Struct.java b/staticlibs/device/com/android/net/module/util/Struct.java
index dc0d19b..ff7a711 100644
--- a/staticlibs/device/com/android/net/module/util/Struct.java
+++ b/staticlibs/device/com/android/net/module/util/Struct.java
@@ -422,7 +422,14 @@
                 final byte[] address = new byte[isIpv6 ? 16 : 4];
                 buf.get(address);
                 try {
-                    value = InetAddress.getByAddress(address);
+                    if (isIpv6) {
+                        // Using Inet6Address.getByAddress since InetAddress.getByAddress converts
+                        // v4-mapped v6 address to v4 address internally and returns Inet4Address.
+                        value = Inet6Address.getByAddress(
+                                null /* host */, address, -1 /* scope_id */);
+                    } else {
+                        value = InetAddress.getByAddress(address);
+                    }
                 } catch (UnknownHostException e) {
                     throw new IllegalArgumentException("illegal length of IP address", e);
                 }
diff --git a/staticlibs/framework/com/android/net/module/util/InetAddressUtils.java b/staticlibs/framework/com/android/net/module/util/InetAddressUtils.java
index 40fc59f..4b27a97 100644
--- a/staticlibs/framework/com/android/net/module/util/InetAddressUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/InetAddressUtils.java
@@ -21,6 +21,7 @@
 import android.util.Log;
 
 
+import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.net.UnknownHostException;
@@ -32,6 +33,7 @@
 public class InetAddressUtils {
 
     private static final String TAG = InetAddressUtils.class.getSimpleName();
+    private static final int INET4_ADDR_LENGTH = 4;
     private static final int INET6_ADDR_LENGTH = 16;
 
     /**
@@ -93,5 +95,29 @@
         }
     }
 
+    /**
+     * Create a v4-mapped v6 address from v4 address
+     *
+     * @param v4Addr Inet4Address which is converted to v4-mapped v6 address
+     * @return v4-mapped v6 address
+     */
+    public static Inet6Address v4MappedV6Address(@NonNull final Inet4Address v4Addr) {
+        final byte[] v6AddrBytes = new byte[INET6_ADDR_LENGTH];
+        v6AddrBytes[10] = (byte) 0xFF;
+        v6AddrBytes[11] = (byte) 0xFF;
+        System.arraycopy(v4Addr.getAddress(), 0 /* srcPos */, v6AddrBytes, 12 /* dstPos */,
+                INET4_ADDR_LENGTH);
+        try {
+            // Using Inet6Address.getByAddress since InetAddress.getByAddress converts v4-mapped v6
+            // address to v4 address internally and returns Inet4Address
+            return Inet6Address.getByAddress(null /* host */, v6AddrBytes, -1 /* scope_id */);
+        } catch (UnknownHostException impossible) {
+            // getByAddress throws UnknownHostException when the argument is the invalid length
+            // but INET6_ADDR_LENGTH(16) is the valid length.
+            Log.wtf(TAG, "Failed to generate v4-mapped v6 address from " + v4Addr, impossible);
+            return null;
+        }
+    }
+
     private InetAddressUtils() {}
 }
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/InetAddressUtilsTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/InetAddressUtilsTest.java
index bb2b933..66427fc 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/InetAddressUtilsTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/InetAddressUtilsTest.java
@@ -18,6 +18,7 @@
 
 import static junit.framework.Assert.assertEquals;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
@@ -30,6 +31,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetAddress;
 
@@ -92,4 +94,17 @@
         assertEquals(localAddrStr + "%" + scopeId, updatedLocalAddr.getHostAddress());
         assertEquals(scopeId, updatedLocalAddr.getScopeId());
     }
+
+    @Test
+    public void testV4MappedV6Address() throws Exception {
+        final Inet4Address v4Addr = (Inet4Address) InetAddress.getByName("192.0.2.1");
+        final Inet6Address v4MappedV6Address = InetAddressUtils.v4MappedV6Address(v4Addr);
+        final byte[] expectedAddrBytes = new byte[]{
+                (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00,
+                (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00,
+                (byte) 0x00, (byte) 0x00, (byte) 0xff, (byte) 0xff,
+                (byte) 0xc0, (byte) 0x00, (byte) 0x02, (byte) 0x01,
+        };
+        assertArrayEquals(expectedAddrBytes, v4MappedV6Address.getAddress());
+    }
 }
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/StructTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/StructTest.java
index b4da043..a39b7a3 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/StructTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/StructTest.java
@@ -765,6 +765,14 @@
                 msg.writeToBytes(ByteOrder.BIG_ENDIAN));
     }
 
+    @Test
+    public void testV4MappedV6Address() {
+        final IpAddressMessage msg = doParsingMessageTest("c0a86401"
+                + "00000000000000000000ffffc0a86401", IpAddressMessage.class, ByteOrder.BIG_ENDIAN);
+        assertEquals(TEST_IPV4_ADDRESS, msg.ipv4Address);
+        assertEquals(InetAddressUtils.v4MappedV6Address(TEST_IPV4_ADDRESS), msg.ipv6Address);
+    }
+
     public static class WrongIpAddressType extends Struct {
         @Field(order = 0, type = Type.Ipv4Address) public byte[] ipv4Address;
         @Field(order = 1, type = Type.Ipv6Address) public byte[] ipv6Address;
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
index d92fb01..5893de7 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
@@ -752,27 +752,24 @@
         assertDelayedShellCommand("dumpsys deviceidle get deep", enabled ? "IDLE" : "ACTIVE");
     }
 
-    protected void setAppIdle(boolean enabled) throws Exception {
+    protected void setAppIdle(boolean isIdle) throws Exception {
+        setAppIdleNoAssert(isIdle);
+        assertAppIdle(isIdle);
+    }
+
+    protected void setAppIdleNoAssert(boolean isIdle) throws Exception {
         if (!isAppStandbySupported()) {
             return;
         }
-        Log.i(TAG, "Setting app idle to " + enabled);
-        executeSilentShellCommand("am set-inactive " + TEST_APP2_PKG + " " + enabled );
-        assertAppIdle(enabled);
+        Log.i(TAG, "Setting app idle to " + isIdle);
+        final String bucketName = isIdle ? "rare" : "active";
+        executeSilentShellCommand("am set-standby-bucket " + TEST_APP2_PKG + " " + bucketName);
     }
 
-    protected void setAppIdleNoAssert(boolean enabled) throws Exception {
-        if (!isAppStandbySupported()) {
-            return;
-        }
-        Log.i(TAG, "Setting app idle to " + enabled);
-        executeSilentShellCommand("am set-inactive " + TEST_APP2_PKG + " " + enabled );
-    }
-
-    protected void assertAppIdle(boolean enabled) throws Exception {
+    protected void assertAppIdle(boolean isIdle) throws Exception {
         try {
             assertDelayedShellCommand("am get-inactive " + TEST_APP2_PKG,
-                    30 /* maxTries */, 1 /* napTimeSeconds */, "Idle=" + enabled);
+                    30 /* maxTries */, 1 /* napTimeSeconds */, "Idle=" + isIdle);
         } catch (Throwable e) {
             throw e;
         }
diff --git a/tests/unit/java/android/net/BpfNetMapsReaderTest.kt b/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
index facb932..86a0acb 100644
--- a/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
+++ b/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
@@ -16,23 +16,31 @@
 
 package android.net
 
+import android.net.BpfNetMapsConstants.DOZABLE_MATCH
+import android.net.BpfNetMapsConstants.STANDBY_MATCH
 import android.net.BpfNetMapsConstants.UID_RULES_CONFIGURATION_KEY
 import android.net.BpfNetMapsUtils.getMatchByFirewallChain
-import android.os.Build
+import android.os.Build.VERSION_CODES
 import com.android.net.module.util.IBpfMap
 import com.android.net.module.util.Struct.S32
 import com.android.net.module.util.Struct.U32
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
 import com.android.testutils.DevSdkIgnoreRunner
 import com.android.testutils.TestBpfMap
+import java.lang.reflect.Modifier
+import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 import org.junit.Test
 import org.junit.runner.RunWith
 
+private const val TEST_UID1 = 1234
+private const val TEST_UID2 = TEST_UID1 + 1
+private const val NO_IIF = 0
+
 // pre-T devices does not support Bpf.
 @RunWith(DevSdkIgnoreRunner::class)
-@IgnoreUpTo(Build.VERSION_CODES.S_V2)
+@IgnoreUpTo(VERSION_CODES.S_V2)
 class BpfNetMapsReaderTest {
     private val testConfigurationMap: IBpfMap<S32, U32> = TestBpfMap()
     private val testUidOwnerMap: IBpfMap<S32, UidOwnerValue> = TestBpfMap()
@@ -66,4 +74,75 @@
         doTestIsChainEnabled(ConnectivityManager.FIREWALL_CHAIN_RESTRICTED)
         doTestIsChainEnabled(ConnectivityManager.FIREWALL_CHAIN_LOW_POWER_STANDBY)
     }
+
+    @Test
+    fun testFirewallChainList() {
+        // Verify that when a firewall chain constant is added, it should also be included in
+        // firewall chain list.
+        val declaredChains = ConnectivityManager::class.java.declaredFields.filter {
+            Modifier.isStatic(it.modifiers) && it.name.startsWith("FIREWALL_CHAIN_")
+        }
+        // Verify the size matches, this also verifies no common item in allow and deny chains.
+        assertEquals(BpfNetMapsConstants.ALLOW_CHAINS.size +
+                BpfNetMapsConstants.DENY_CHAINS.size, declaredChains.size)
+        declaredChains.forEach {
+            assertTrue(BpfNetMapsConstants.ALLOW_CHAINS.contains(it.get(null)) ||
+                    BpfNetMapsConstants.DENY_CHAINS.contains(it.get(null)))
+        }
+    }
+
+    private fun mockChainEnabled(chain: Int, enabled: Boolean) {
+        val config = testConfigurationMap.getValue(UID_RULES_CONFIGURATION_KEY).`val`
+        val newConfig = if (enabled) {
+            config or getMatchByFirewallChain(chain)
+        } else {
+            config and getMatchByFirewallChain(chain).inv()
+        }
+        testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(newConfig))
+    }
+
+    @Test
+    fun testIsUidNetworkingBlockedByFirewallChains_allowChain() {
+        // With everything disabled by default, verify the return value is false.
+        testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
+        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+
+        // Enable dozable chain but does not provide allowed list. Verify the network is blocked
+        // for all uids.
+        mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_DOZABLE, true)
+        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+
+        // Add uid1 to dozable allowed list. Verify the network is not blocked for uid1, while
+        // uid2 is blocked.
+        testUidOwnerMap.updateEntry(S32(TEST_UID1), UidOwnerValue(NO_IIF, DOZABLE_MATCH))
+        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+    }
+
+    @Test
+    fun testIsUidNetworkingBlockedByFirewallChains_denyChain() {
+        // Enable standby chain but does not provide denied list. Verify the network is allowed
+        // for all uids.
+        testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
+        mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_STANDBY, true)
+        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+
+        // Add uid1 to standby allowed list. Verify the network is blocked for uid1, while
+        // uid2 is not blocked.
+        testUidOwnerMap.updateEntry(S32(TEST_UID1), UidOwnerValue(NO_IIF, STANDBY_MATCH))
+        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+        assertFalse(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID2))
+    }
+
+    @Test
+    fun testIsUidNetworkingBlockedByFirewallChains_blockedWithAllowed() {
+        // Uids blocked by powersave chain but allowed by standby chain, verify the blocking
+        // takes higher priority.
+        testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
+        mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_POWERSAVE, true)
+        mockChainEnabled(ConnectivityManager.FIREWALL_CHAIN_STANDBY, true)
+        assertTrue(bpfNetMapsReader.isUidBlockedByFirewallChains(TEST_UID1))
+    }
 }
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index ea2b261..1dfc8c0 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -16,10 +16,12 @@
 
 package com.android.server;
 
+import static android.net.BpfNetMapsConstants.ALLOW_CHAINS;
 import static android.net.BpfNetMapsConstants.CURRENT_STATS_MAP_CONFIGURATION_KEY;
 import static android.net.BpfNetMapsConstants.DATA_SAVER_ENABLED_KEY;
 import static android.net.BpfNetMapsConstants.DATA_SAVER_DISABLED;
 import static android.net.BpfNetMapsConstants.DATA_SAVER_ENABLED;
+import static android.net.BpfNetMapsConstants.DENY_CHAINS;
 import static android.net.BpfNetMapsConstants.DOZABLE_MATCH;
 import static android.net.BpfNetMapsConstants.HAPPY_BOX_MATCH;
 import static android.net.BpfNetMapsConstants.IIF_MATCH;
@@ -119,20 +121,16 @@
     private static final int NO_IIF = 0;
     private static final int NULL_IIF = 0;
     private static final String CHAINNAME = "fw_dozable";
-    private static final List<Integer> FIREWALL_CHAINS = List.of(
-            FIREWALL_CHAIN_DOZABLE,
-            FIREWALL_CHAIN_STANDBY,
-            FIREWALL_CHAIN_POWERSAVE,
-            FIREWALL_CHAIN_RESTRICTED,
-            FIREWALL_CHAIN_LOW_POWER_STANDBY,
-            FIREWALL_CHAIN_OEM_DENY_1,
-            FIREWALL_CHAIN_OEM_DENY_2,
-            FIREWALL_CHAIN_OEM_DENY_3
-    );
 
     private static final long STATS_SELECT_MAP_A = 0;
     private static final long STATS_SELECT_MAP_B = 1;
 
+    private static final List<Integer> FIREWALL_CHAINS = new ArrayList<>();
+    static {
+        FIREWALL_CHAINS.addAll(ALLOW_CHAINS);
+        FIREWALL_CHAINS.addAll(DENY_CHAINS);
+    }
+
     private BpfNetMaps mBpfNetMaps;
 
     @Mock INetd mNetd;
@@ -617,7 +615,7 @@
         mUidOwnerMap.updateEntry(new S32(TEST_UID), new UidOwnerValue(TEST_IF_INDEX, IIF_MATCH));
 
         for (final int chain: testChains) {
-            final int ruleToAddMatch = mBpfNetMaps.isFirewallAllowList(chain)
+            final int ruleToAddMatch = BpfNetMapsUtils.isFirewallAllowList(chain)
                     ? FIREWALL_RULE_ALLOW : FIREWALL_RULE_DENY;
             mBpfNetMaps.setUidRule(chain, TEST_UID, ruleToAddMatch);
         }
@@ -625,7 +623,7 @@
         checkUidOwnerValue(TEST_UID, TEST_IF_INDEX, IIF_MATCH | getMatch(testChains));
 
         for (final int chain: testChains) {
-            final int ruleToRemoveMatch = mBpfNetMaps.isFirewallAllowList(chain)
+            final int ruleToRemoveMatch = BpfNetMapsUtils.isFirewallAllowList(chain)
                     ? FIREWALL_RULE_DENY : FIREWALL_RULE_ALLOW;
             mBpfNetMaps.setUidRule(chain, TEST_UID, ruleToRemoveMatch);
         }
@@ -705,11 +703,11 @@
         for (final int chain: FIREWALL_CHAINS) {
             final String testCase = "EnabledChains: " + enableChains + " CheckedChain: " + chain;
             if (enableChains.contains(chain)) {
-                final int expectedRule = mBpfNetMaps.isFirewallAllowList(chain)
+                final int expectedRule = BpfNetMapsUtils.isFirewallAllowList(chain)
                         ? FIREWALL_RULE_ALLOW : FIREWALL_RULE_DENY;
                 assertEquals(testCase, expectedRule, mBpfNetMaps.getUidRule(chain, TEST_UID));
             } else {
-                final int expectedRule = mBpfNetMaps.isFirewallAllowList(chain)
+                final int expectedRule = BpfNetMapsUtils.isFirewallAllowList(chain)
                         ? FIREWALL_RULE_DENY : FIREWALL_RULE_ALLOW;
                 assertEquals(testCase, expectedRule, mBpfNetMaps.getUidRule(chain, TEST_UID));
             }
@@ -752,7 +750,7 @@
     public void testGetUidRuleNoEntry() throws Exception {
         mUidOwnerMap.clear();
         for (final int chain: FIREWALL_CHAINS) {
-            final int expectedRule = mBpfNetMaps.isFirewallAllowList(chain)
+            final int expectedRule = BpfNetMapsUtils.isFirewallAllowList(chain)
                     ? FIREWALL_RULE_DENY : FIREWALL_RULE_ALLOW;
             assertEquals(expectedRule, mBpfNetMaps.getUidRule(chain, TEST_UID));
         }
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 82fb651..aae37e5 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -417,7 +417,6 @@
 import com.android.server.connectivity.QosCallbackTracker;
 import com.android.server.connectivity.TcpKeepaliveController;
 import com.android.server.connectivity.UidRangeUtils;
-import com.android.server.connectivity.Vpn;
 import com.android.server.connectivity.VpnProfileStore;
 import com.android.server.net.NetworkPinner;
 import com.android.testutils.DevSdkIgnoreRule;
@@ -1497,13 +1496,8 @@
         return uidRangesForUids(CollectionUtils.toIntArray(uids));
     }
 
-    private static Looper startHandlerThreadAndReturnLooper() {
-        final HandlerThread handlerThread = new HandlerThread("MockVpnThread");
-        handlerThread.start();
-        return handlerThread.getLooper();
-    }
-
-    private class MockVpn extends Vpn implements TestableNetworkCallback.HasNetwork {
+    // Helper class to mock vpn interaction.
+    private class MockVpn implements TestableNetworkCallback.HasNetwork {
         // Note : Please do not add any new instrumentation here. If you need new instrumentation,
         // please add it in CSTest and use subclasses of CSTest instead of adding more
         // tools in ConnectivityServiceTest.
@@ -1511,45 +1505,23 @@
         // Careful ! This is different from mNetworkAgent, because MockNetworkAgent does
         // not inherit from NetworkAgent.
         private TestNetworkAgentWrapper mMockNetworkAgent;
+        // Initialize a stored NetworkCapabilities following the defaults of VPN. The TransportInfo
+        // should at least be updated to a valid VPN type before usage, see registerAgent(...).
+        private NetworkCapabilities mNetworkCapabilities = new NetworkCapabilities.Builder()
+                .addTransportType(NetworkCapabilities.TRANSPORT_VPN)
+                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN)
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED)
+                .setTransportInfo(new VpnTransportInfo(
+                        VpnManager.TYPE_VPN_NONE,
+                        null /* sessionId */,
+                        false /* bypassable */,
+                        false /* longLivedTcpConnectionsExpensive */))
+                .build();
         private boolean mAgentRegistered = false;
 
         private int mVpnType = VpnManager.TYPE_VPN_SERVICE;
-        private UnderlyingNetworkInfo mUnderlyingNetworkInfo;
         private String mSessionKey;
 
-        public MockVpn(int userId) {
-            super(startHandlerThreadAndReturnLooper(), mServiceContext,
-                    new Dependencies() {
-                        @Override
-                        public boolean isCallerSystem() {
-                            return true;
-                        }
-
-                        @Override
-                        public DeviceIdleInternal getDeviceIdleInternal() {
-                            return mDeviceIdleInternal;
-                        }
-                    },
-                    mNetworkManagementService, mMockNetd, userId, mVpnProfileStore,
-                    new SystemServices(mServiceContext) {
-                        @Override
-                        public String settingsSecureGetStringForUser(String key, int userId) {
-                            switch (key) {
-                                // Settings keys not marked as @Readable are not readable from
-                                // non-privileged apps, unless marked as testOnly=true
-                                // (atest refuses to install testOnly=true apps), even if mocked
-                                // in the content provider, because
-                                // Settings.Secure.NameValueCache#getStringForUser checks the key
-                                // before querying the mock settings provider.
-                                case Settings.Secure.ALWAYS_ON_VPN_APP:
-                                    return null;
-                                default:
-                                    return super.settingsSecureGetStringForUser(key, userId);
-                            }
-                        }
-                    }, new Ikev2SessionCreator());
-        }
-
         public void setUids(Set<UidRange> uids) {
             mNetworkCapabilities.setUids(UidRange.toIntRanges(uids));
             if (mAgentRegistered) {
@@ -1561,7 +1533,6 @@
             mVpnType = vpnType;
         }
 
-        @Override
         public Network getNetwork() {
             return (mMockNetworkAgent == null) ? null : mMockNetworkAgent.getNetwork();
         }
@@ -1570,7 +1541,6 @@
             return null == mMockNetworkAgent ? null : mMockNetworkAgent.getNetworkAgentConfig();
         }
 
-        @Override
         public int getActiveVpnType() {
             return mVpnType;
         }
@@ -1584,14 +1554,11 @@
         private void registerAgent(boolean isAlwaysMetered, Set<UidRange> uids, LinkProperties lp)
                 throws Exception {
             if (mAgentRegistered) throw new IllegalStateException("already registered");
-            updateState(NetworkInfo.DetailedState.CONNECTING, "registerAgent");
-            mConfig = new VpnConfig();
-            mConfig.session = "MySession12345";
+            final String session = "MySession12345";
             setUids(uids);
             if (!isAlwaysMetered) mNetworkCapabilities.addCapability(NET_CAPABILITY_NOT_METERED);
-            mInterface = VPN_IFNAME;
             mNetworkCapabilities.setTransportInfo(new VpnTransportInfo(getActiveVpnType(),
-                    mConfig.session));
+                    session));
             mMockNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_VPN, lp,
                     mNetworkCapabilities);
             mMockNetworkAgent.waitForIdle(TIMEOUT_MS);
@@ -1605,9 +1572,7 @@
             mAgentRegistered = true;
             verify(mMockNetd).networkCreate(nativeNetworkConfigVpn(getNetwork().netId,
                     !mMockNetworkAgent.isBypassableVpn(), mVpnType));
-            updateState(NetworkInfo.DetailedState.CONNECTED, "registerAgent");
             mNetworkCapabilities.set(mMockNetworkAgent.getNetworkCapabilities());
-            mNetworkAgent = mMockNetworkAgent.getNetworkAgent();
         }
 
         private void registerAgent(Set<UidRange> uids) throws Exception {
@@ -1667,23 +1632,20 @@
         public void disconnect() {
             if (mMockNetworkAgent != null) {
                 mMockNetworkAgent.disconnect();
-                updateState(NetworkInfo.DetailedState.DISCONNECTED, "disconnect");
             }
             mAgentRegistered = false;
             setUids(null);
             // Remove NET_CAPABILITY_INTERNET or MockNetworkAgent will refuse to connect later on.
             mNetworkCapabilities.removeCapability(NET_CAPABILITY_INTERNET);
-            mInterface = null;
         }
 
-        private synchronized void startLegacyVpn() {
-            updateState(DetailedState.CONNECTING, "startLegacyVpn");
+        private void startLegacyVpn() {
+            // Do nothing.
         }
 
         // Mock the interaction of IkeV2VpnRunner start. In the context of ConnectivityService,
         // setVpnDefaultForUids() is the main interaction and a sessionKey is stored.
-        private synchronized void startPlatformVpn() {
-            updateState(DetailedState.CONNECTING, "startPlatformVpn");
+        private void startPlatformVpn() {
             mSessionKey = UUID.randomUUID().toString();
             // Assuming no disallowed applications
             final Set<Range<Integer>> ranges = UidRange.toIntRanges(Set.of(PRIMARY_UIDRANGE));
@@ -1692,7 +1654,6 @@
             waitForIdle();
         }
 
-        @Override
         public void startLegacyVpnPrivileged(VpnProfile profile,
                 @Nullable Network underlying, @NonNull LinkProperties egress) {
             switch (profile.type) {
@@ -1714,8 +1675,7 @@
             }
         }
 
-        @Override
-        public synchronized void stopVpnRunnerPrivileged() {
+        public void stopVpnRunnerPrivileged() {
             if (mSessionKey != null) {
                 // Clear vpn network preference.
                 mCm.setVpnDefaultForUids(mSessionKey, Collections.EMPTY_LIST);
@@ -1724,20 +1684,7 @@
             disconnect();
         }
 
-        @Override
-        public synchronized UnderlyingNetworkInfo getUnderlyingNetworkInfo() {
-            if (mUnderlyingNetworkInfo != null) return mUnderlyingNetworkInfo;
-
-            return super.getUnderlyingNetworkInfo();
-        }
-
-        private synchronized void setUnderlyingNetworkInfo(
-                UnderlyingNetworkInfo underlyingNetworkInfo) {
-            mUnderlyingNetworkInfo = underlyingNetworkInfo;
-        }
-
-        @Override
-        public synchronized boolean setUnderlyingNetworks(@Nullable Network[] networks) {
+        public boolean setUnderlyingNetworks(@Nullable Network[] networks) {
             if (!mAgentRegistered) return false;
             mMockNetworkAgent.setUnderlyingNetworks(
                     (networks == null) ? null : Arrays.asList(networks));
@@ -1774,11 +1721,6 @@
         waitForIdle();
     }
 
-    private void mockVpn(int uid) {
-        int userId = UserHandle.getUserId(uid);
-        mMockVpn = new MockVpn(userId);
-    }
-
     private void mockUidNetworkingBlocked() {
         doAnswer(i -> isUidBlocked(mBlockedReasons, i.getArgument(1))
         ).when(mNetworkPolicyManager).isUidNetworkingBlocked(anyInt(), anyBoolean());
@@ -1990,7 +1932,7 @@
         mService.systemReadyInternal();
         verify(mMockDnsResolver).registerUnsolicitedEventListener(any());
 
-        mockVpn(Process.myUid());
+        mMockVpn = new MockVpn();
         mCm.bindProcessToNetwork(null);
         mQosCallbackTracker = mock(QosCallbackTracker.class);
 
@@ -4346,7 +4288,9 @@
         testFactory.terminate();
         testFactory.assertNoRequestChanged();
         if (networkCallback != null) mCm.unregisterNetworkCallback(networkCallback);
-        handlerThread.quit();
+
+        handlerThread.quitSafely();
+        handlerThread.join();
     }
 
     @Test
@@ -4407,6 +4351,8 @@
         expectNoRequestChanged(testFactoryAll); // still seeing the request
 
         mWiFiAgent.disconnect();
+        handlerThread.quitSafely();
+        handlerThread.join();
     }
 
     @Test
@@ -4440,7 +4386,8 @@
                 }
             }
         }
-        handlerThread.quit();
+        handlerThread.quitSafely();
+        handlerThread.join();
     }
 
     @Test
@@ -6061,7 +6008,8 @@
             testFactory.assertNoRequestChanged();
         } finally {
             mCm.unregisterNetworkCallback(cellNetworkCallback);
-            handlerThread.quit();
+            handlerThread.quitSafely();
+            handlerThread.join();
         }
     }
 
@@ -6651,7 +6599,8 @@
             }
         } finally {
             testFactory.terminate();
-            handlerThread.quit();
+            handlerThread.quitSafely();
+            handlerThread.join();
         }
     }
 
@@ -10268,7 +10217,6 @@
 
         // Init lockdown state to simulate LockdownVpnTracker behavior.
         mCm.setLegacyLockdownVpnEnabled(true);
-        mMockVpn.setEnableTeardown(false);
         final List<Range<Integer>> ranges =
                 intRangesPrimaryExcludingUids(Collections.EMPTY_LIST /* excludedeUids */);
         mCm.setRequireVpnForUids(true /* requireVpn */, ranges);
@@ -10585,13 +10533,11 @@
         final boolean allowlist = true;
         final boolean denylist = false;
 
-        doReturn(true).when(mBpfNetMaps).isFirewallAllowList(anyInt());
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_DOZABLE, allowlist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_POWERSAVE, allowlist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_RESTRICTED, allowlist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_LOW_POWER_STANDBY, allowlist);
 
-        doReturn(false).when(mBpfNetMaps).isFirewallAllowList(anyInt());
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_STANDBY, denylist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_OEM_DENY_1, denylist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_OEM_DENY_2, denylist);
@@ -12596,9 +12542,6 @@
         mMockVpn.establish(new LinkProperties(), vpnOwnerUid, vpnRange);
         assertVpnUidRangesUpdated(true, vpnRange, vpnOwnerUid);
 
-        final UnderlyingNetworkInfo underlyingNetworkInfo =
-                new UnderlyingNetworkInfo(vpnOwnerUid, VPN_IFNAME, new ArrayList<>());
-        mMockVpn.setUnderlyingNetworkInfo(underlyingNetworkInfo);
         mDeps.setConnectionOwnerUid(42);
     }
 
@@ -15378,6 +15321,8 @@
         expectNoRequestChanged(oemPaidFactory);
         internetFactory.expectRequestAdd();
         mCm.unregisterNetworkCallback(wifiCallback);
+        handlerThread.quitSafely();
+        handlerThread.join();
     }
 
     /**
@@ -15742,6 +15687,8 @@
             assertTrue(testFactory.getMyStartRequested());
         } finally {
             testFactory.terminate();
+            handlerThread.quitSafely();
+            handlerThread.join();
         }
     }
 
diff --git a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
index 986c389..8e19c01 100644
--- a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
@@ -77,6 +77,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
+import com.android.internal.util.IndentingPrintWriter;
 import com.android.server.connectivity.AutomaticOnOffKeepaliveTracker.AutomaticOnOffKeepalive;
 import com.android.server.connectivity.KeepaliveTracker.KeepaliveInfo;
 import com.android.testutils.DevSdkIgnoreRule;
@@ -94,6 +95,7 @@
 import org.mockito.MockitoAnnotations;
 
 import java.io.FileDescriptor;
+import java.io.StringWriter;
 import java.net.Inet4Address;
 import java.net.InetAddress;
 import java.net.Socket;
@@ -974,4 +976,19 @@
         // The keepalive should be removed in AutomaticOnOffKeepaliveTracker.
         assertNull(getAutoKiForBinder(testInfo.binder));
     }
+
+    @Test
+    public void testDumpDoesNotCrash() throws Exception {
+        final TestKeepaliveInfo testInfo1 = doStartNattKeepalive();
+        final TestKeepaliveInfo testInfo2 = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(TEST_SLOT, testInfo1.kpd);
+        checkAndProcessKeepaliveStart(TEST_SLOT + 1, testInfo2.kpd);
+        final AutomaticOnOffKeepalive autoKi1  = getAutoKiForBinder(testInfo1.binder);
+        doPauseKeepalive(autoKi1);
+
+        final StringWriter stringWriter = new StringWriter();
+        final IndentingPrintWriter pw = new IndentingPrintWriter(stringWriter, "   ");
+        visibleOnHandlerThread(mTestHandler, () -> mAOOKeepaliveTracker.dump(pw));
+        assertFalse(stringWriter.toString().isEmpty());
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java b/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
index 90a0edd..1b964e2 100644
--- a/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/KeepaliveStatsTrackerTest.java
@@ -37,6 +37,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.content.BroadcastReceiver;
@@ -1293,5 +1294,18 @@
                 expectRegisteredDurations,
                 expectActiveDurations,
                 new KeepaliveCarrierStats[0]);
+
+        assertTrue(mKeepaliveStatsTracker.allMetricsExpected(dailyKeepaliveInfoReported));
+
+        // Write time after 26 hours.
+        final int writeTime2 = 26 * 60 * 60 * 1000;
+        setElapsedRealtime(writeTime2);
+
+        visibleOnHandlerThread(mTestHandler, () -> mKeepaliveStatsTracker.writeAndResetMetrics());
+        verify(mDependencies, times(2)).writeStats(dailyKeepaliveInfoReportedCaptor.capture());
+        final DailykeepaliveInfoReported dailyKeepaliveInfoReported2 =
+                dailyKeepaliveInfoReportedCaptor.getValue();
+
+        assertFalse(mKeepaliveStatsTracker.allMetricsExpected(dailyKeepaliveInfoReported2));
     }
 }
diff --git a/tests/unit/java/com/android/server/connectivityservice/CSDestroyedNetworkTests.kt b/tests/unit/java/com/android/server/connectivityservice/CSDestroyedNetworkTests.kt
new file mode 100644
index 0000000..d40c035
--- /dev/null
+++ b/tests/unit/java/com/android/server/connectivityservice/CSDestroyedNetworkTests.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2023 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.server
+
+import android.net.NetworkCapabilities
+import android.net.NetworkCapabilities.TRANSPORT_WIFI
+import android.net.NetworkRequest
+import com.android.testutils.RecorderCallback.CallbackEntry.Lost
+import com.android.testutils.TestableNetworkCallback
+import org.junit.Test
+
+private const val LONG_TIMEOUT_MS = 5_000
+
+class CSDestroyedNetworkTests : CSTest() {
+    @Test
+    fun testDestroyNetworkNotKeptWhenUnvalidated() {
+        val nc = NetworkCapabilities.Builder()
+                .addTransportType(TRANSPORT_WIFI)
+                .build()
+
+        val nr = NetworkRequest.Builder()
+                .clearCapabilities()
+                .addTransportType(TRANSPORT_WIFI)
+                .build()
+        val cbRequest = TestableNetworkCallback()
+        val cbCallback = TestableNetworkCallback()
+        cm.requestNetwork(nr, cbRequest)
+        cm.registerNetworkCallback(nr, cbCallback)
+
+        val firstAgent = Agent(nc = nc)
+        firstAgent.connect()
+        cbCallback.expectAvailableCallbacks(firstAgent.network, validated = false)
+
+        firstAgent.unregisterAfterReplacement(LONG_TIMEOUT_MS)
+
+        val secondAgent = Agent(nc = nc)
+        secondAgent.connect()
+        cbCallback.expectAvailableCallbacks(secondAgent.network, validated = false)
+
+        cbCallback.expect<Lost>(timeoutMs = 500) { it.network == firstAgent.network }
+    }
+}
diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
index 094ded3..8860895 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSAgentWrapper.kt
@@ -34,7 +34,6 @@
 import android.net.NetworkTestResultParcelable
 import android.net.networkstack.NetworkStackClientBase
 import android.os.HandlerThread
-import com.android.modules.utils.build.SdkLevel
 import com.android.testutils.RecorderCallback.CallbackEntry.Available
 import com.android.testutils.RecorderCallback.CallbackEntry.Lost
 import com.android.testutils.TestableNetworkCallback
@@ -168,5 +167,6 @@
         cb.eventuallyExpect<Lost> { it.network == agent.network }
     }
 
+    fun unregisterAfterReplacement(timeoutMs: Int) = agent.unregisterAfterReplacement(timeoutMs)
     fun sendLocalNetworkConfig(lnc: LocalNetworkConfig) = agent.sendLocalNetworkConfig(lnc)
 }
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java b/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
index c477b2c..e62ac74 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
@@ -125,7 +125,7 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
-        mObserverHandlerThread = new HandlerThread("HandlerThread");
+        mObserverHandlerThread = new HandlerThread("NetworkStatsObserversTest");
         mObserverHandlerThread.start();
         final Looper observerLooper = mObserverHandlerThread.getLooper();
         mStatsObservers = new NetworkStatsObservers() {
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index e8d5c66..92a5b64 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -120,6 +120,7 @@
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.IBinder;
+import android.os.Looper;
 import android.os.PowerManager;
 import android.os.SimpleClock;
 import android.provider.Settings;
@@ -284,6 +285,7 @@
     private @Mock PersistentInt mImportLegacyFallbacksCounter;
     private @Mock Resources mResources;
     private Boolean mIsDebuggable;
+    private HandlerThread mObserverHandlerThread;
 
     private class MockContext extends BroadcastInterceptingContext {
         private final Context mBaseContext;
@@ -365,10 +367,23 @@
         PowerManager.WakeLock wakeLock =
                 powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
 
-        mHandlerThread = new HandlerThread("HandlerThread");
+        mHandlerThread = new HandlerThread("NetworkStatsServiceTest-HandlerThread");
         final NetworkStatsService.Dependencies deps = makeDependencies();
+        // Create a separate thread for observers to run on. This thread cannot be the same
+        // as the handler thread, because the observer callback is fired on this thread, and
+        // it should not be blocked by client code. Additionally, creating the observers
+        // object requires a looper, which can only be obtained after a thread has been started.
+        mObserverHandlerThread = new HandlerThread("NetworkStatsServiceTest-ObserversThread");
+        mObserverHandlerThread.start();
+        final Looper observerLooper = mObserverHandlerThread.getLooper();
+        final NetworkStatsObservers statsObservers = new NetworkStatsObservers() {
+            @Override
+            protected Looper getHandlerLooperLocked() {
+                return observerLooper;
+            }
+        };
         mService = new NetworkStatsService(mServiceContext, mNetd, mAlarmManager, wakeLock,
-                mClock, mSettings, mStatsFactory, new NetworkStatsObservers(), deps);
+                mClock, mSettings, mStatsFactory, statsObservers, deps);
 
         mElapsedRealtime = 0L;
 
@@ -545,8 +560,14 @@
         mSession.close();
         mService = null;
 
-        mHandlerThread.quitSafely();
-        mHandlerThread.join();
+        if (mHandlerThread != null) {
+            mHandlerThread.quitSafely();
+            mHandlerThread.join();
+        }
+        if (mObserverHandlerThread != null) {
+            mObserverHandlerThread.quitSafely();
+            mObserverHandlerThread.join();
+        }
     }
 
     private void initWifiStats(NetworkStateSnapshot snapshot) throws Exception {