Merge "Send onBlockedStatusChanged callback based on bpf map contents" into main
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 5320c81..cc8eef8 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -486,8 +486,30 @@
     private final boolean mBackgroundFirewallChainEnabled;
 
     /**
-     * Stale copy of uid blocked reasons provided by NPMS. As long as they are accessed only in
-     * internal handler thread, they don't need a lock.
+     * Uids ConnectivityService tracks blocked status of to send blocked status callbacks.
+     * Key is uid based on mAsUid of registered networkRequestInfo
+     * Value is count of registered networkRequestInfo
+     *
+     * This is necessary because when a firewall chain is enabled or disabled, that affects all UIDs
+     * on the system, not just UIDs on that firewall chain. For example, entering doze mode affects
+     * all UIDs that are not on the dozable chain. ConnectivityService doesn't know which UIDs are
+     * running. But it only needs to send onBlockedStatusChanged to UIDs that have at least one
+     * NetworkCallback registered.
+     *
+     * UIDs are added to this list on the binder thread when processing requestNetwork and similar
+     * IPCs. They are removed from this list on the handler thread, when the callback unregistration
+     * is fully processed. They cannot be unregistered when the unregister IPC is processed because
+     * sometimes requests are unregistered on the handler thread.
+     */
+    @GuardedBy("mBlockedStatusTrackingUids")
+    private final SparseIntArray mBlockedStatusTrackingUids = new SparseIntArray();
+
+    /**
+     * Stale copy of UID blocked reasons. This is used to send onBlockedStatusChanged
+     * callbacks. This is only used on the handler thread, so it does not require a lock.
+     * On U-, the blocked reasons come from NPMS.
+     * On V+, the blocked reasons come from the BPF map contents and only maintains blocked reasons
+     * of uids that register network callbacks.
      */
     private final SparseIntArray mUidBlockedReasons = new SparseIntArray();
 
@@ -791,11 +813,10 @@
     private static final int EVENT_SET_PROFILE_NETWORK_PREFERENCE = 50;
 
     /**
-     * Event to specify that reasons for why an uid is blocked changed.
-     * arg1 = uid
-     * arg2 = blockedReasons
+     * Event to update blocked reasons for uids.
+     * obj = List of Pair(uid, blockedReasons)
      */
-    private static final int EVENT_UID_BLOCKED_REASON_CHANGED = 51;
+    private static final int EVENT_BLOCKED_REASONS_CHANGED = 51;
 
     /**
      * Event to register a new network offer
@@ -1828,7 +1849,12 @@
         // To ensure uid state is synchronized with Network Policy, register for
         // NetworkPolicyManagerService events must happen prior to NetworkPolicyManagerService
         // reading existing policy from disk.
-        mPolicyManager.registerNetworkPolicyCallback(null, mPolicyCallback);
+        // If shouldTrackUidsForBlockedStatusCallbacks() is true (On V+), ConnectivityService
+        // updates blocked reasons when firewall chain and data saver status is updated based on
+        // bpf map contents instead of receiving callbacks from NPMS
+        if (!shouldTrackUidsForBlockedStatusCallbacks()) {
+            mPolicyManager.registerNetworkPolicyCallback(null, mPolicyCallback);
+        }
 
         final PowerManager powerManager = (PowerManager) context.getSystemService(
                 Context.POWER_SERVICE);
@@ -3315,14 +3341,38 @@
     private final NetworkPolicyCallback mPolicyCallback = new NetworkPolicyCallback() {
         @Override
         public void onUidBlockedReasonChanged(int uid, @BlockedReason int blockedReasons) {
-            mHandler.sendMessage(mHandler.obtainMessage(EVENT_UID_BLOCKED_REASON_CHANGED,
-                    uid, blockedReasons));
+            if (shouldTrackUidsForBlockedStatusCallbacks()) {
+                Log.wtf(TAG, "Received unexpected NetworkPolicy callback");
+                return;
+            }
+            mHandler.sendMessage(mHandler.obtainMessage(
+                    EVENT_BLOCKED_REASONS_CHANGED,
+                    List.of(new Pair<>(uid, blockedReasons))));
         }
     };
 
-    private void handleUidBlockedReasonChanged(int uid, @BlockedReason int blockedReasons) {
-        maybeNotifyNetworkBlockedForNewState(uid, blockedReasons);
-        setUidBlockedReasons(uid, blockedReasons);
+    private boolean shouldTrackUidsForBlockedStatusCallbacks() {
+        return mDeps.isAtLeastV();
+    }
+
+    @VisibleForTesting
+    void handleBlockedReasonsChanged(List<Pair<Integer, Integer>> reasonsList) {
+        for (Pair<Integer, Integer> reasons: reasonsList) {
+            final int uid = reasons.first;
+            final int blockedReasons = reasons.second;
+            if (shouldTrackUidsForBlockedStatusCallbacks()) {
+                synchronized (mBlockedStatusTrackingUids) {
+                    if (mBlockedStatusTrackingUids.get(uid) == 0) {
+                        // This uid is not tracked anymore.
+                        // This can happen if the network request is unregistered while
+                        // EVENT_BLOCKED_REASONS_CHANGED is posted but not processed yet.
+                        continue;
+                    }
+                }
+            }
+            maybeNotifyNetworkBlockedForNewState(uid, blockedReasons);
+            setUidBlockedReasons(uid, blockedReasons);
+        }
     }
 
     static final class UidFrozenStateChangedArgs {
@@ -5498,6 +5548,12 @@
                 // If there is an active request, then for sure there is a satisfier.
                 nri.getSatisfier().addRequest(activeRequest);
             }
+
+            if (shouldTrackUidsForBlockedStatusCallbacks()
+                    && isAppRequest(nri)
+                    && !nri.mUidTrackedForBlockedStatus) {
+                Log.wtf(TAG, "Registered nri is not tracked for sending blocked status: " + nri);
+            }
         }
 
         if (mFlags.noRematchAllRequestsOnRegister()) {
@@ -5737,7 +5793,7 @@
         }
 
         if (untrackUids) {
-            maybeUntrackUid(nri);
+            maybeUntrackUidAndClearBlockedReasons(nri);
         }
         mNetworkRequestInfoLogs.log("RELEASE " + nri);
         checkNrisConsistency(nri);
@@ -6510,8 +6566,8 @@
                     handlePrivateDnsValidationUpdate(
                             (PrivateDnsValidationUpdate) msg.obj);
                     break;
-                case EVENT_UID_BLOCKED_REASON_CHANGED:
-                    handleUidBlockedReasonChanged(msg.arg1, msg.arg2);
+                case EVENT_BLOCKED_REASONS_CHANGED:
+                    handleBlockedReasonsChanged((List) msg.obj);
                     break;
                 case EVENT_SET_REQUIRE_VPN_FOR_UIDS:
                     handleSetRequireVpnForUids(toBool(msg.arg1), (UidRange[]) msg.obj);
@@ -7376,6 +7432,11 @@
         // maximum limit of registered callbacks per UID.
         final int mAsUid;
 
+        // Flag to indicate that uid of this nri is tracked for sending blocked status callbacks.
+        // It is always true on V+ if mMessenger != null. As such, it's not strictly necessary.
+        // it's used only as a safeguard to avoid double counting or leaking.
+        boolean mUidTrackedForBlockedStatus;
+
         // Preference order of this request.
         final int mPreferenceOrder;
 
@@ -7499,6 +7560,7 @@
             mPerUidCounter = nri.mPerUidCounter;
             mCallbackFlags = nri.mCallbackFlags;
             mCallingAttributionTag = nri.mCallingAttributionTag;
+            mUidTrackedForBlockedStatus = nri.mUidTrackedForBlockedStatus;
             mPreferenceOrder = PREFERENCE_ORDER_INVALID;
             linkDeathRecipient();
         }
@@ -7586,7 +7648,8 @@
                     + " " + mRequests
                     + (mPendingIntent == null ? "" : " to trigger " + mPendingIntent)
                     + " callback flags: " + mCallbackFlags
-                    + " order: " + mPreferenceOrder;
+                    + " order: " + mPreferenceOrder
+                    + " isUidTracked: " + mUidTrackedForBlockedStatus;
         }
     }
 
@@ -8066,29 +8129,75 @@
         return true;
     }
 
-    private void trackUid(final NetworkRequestInfo nri) {
-        if (nri.mMessenger == null && nri.mPendingIntent == null) {
-            Log.wtf(TAG, "trackUid is called for nri with "
-                    + "null mMessenger and mPendingIntent: " + nri);
+    private boolean isAppRequest(NetworkRequestInfo nri) {
+        return nri.mMessenger != null || nri.mPendingIntent != null;
+    }
+
+    private void trackUidAndMaybePostCurrentBlockedReason(final NetworkRequestInfo nri) {
+        if (!isAppRequest(nri)) {
+            Log.wtf(TAG, "trackUidAndMaybePostCurrentBlockedReason is called for non app"
+                    + "request: " + nri);
             return;
         }
         nri.mPerUidCounter.incrementCountOrThrow(nri.mUid);
+
+        // If nri.mMessenger is null, this nri does not have NetworkCallback so ConnectivityService
+        // does not need to send onBlockedStatusChanged callback for this uid and does not need to
+        // track the uid in mBlockedStatusTrackingUids
+        if (!shouldTrackUidsForBlockedStatusCallbacks() || nri.mMessenger == null) {
+            return;
+        }
+        if (nri.mUidTrackedForBlockedStatus) {
+            Log.wtf(TAG, "Nri is already tracked for sending blocked status: " + nri);
+            return;
+        }
+        nri.mUidTrackedForBlockedStatus = true;
+        synchronized (mBlockedStatusTrackingUids) {
+            final int uid = nri.mAsUid;
+            final int count = mBlockedStatusTrackingUids.get(uid, 0);
+            if (count == 0) {
+                mHandler.sendMessage(mHandler.obtainMessage(EVENT_BLOCKED_REASONS_CHANGED,
+                        List.of(new Pair<>(uid, mBpfNetMaps.getUidNetworkingBlockedReasons(uid)))));
+            }
+            mBlockedStatusTrackingUids.put(uid, count + 1);
+        }
     }
 
     private void trackUidAndRegisterNetworkRequest(final int event, NetworkRequestInfo nri) {
-        // Once trackUid is called, the register network
+        // Post the update of the UID's blocked reasons before posting the message that registers
+        // the callback. This is necessary because if the callback immediately matches a request,
+        // the onBlockedStatusChanged must be called with the correct blocked reasons.
+        // Also, once trackUidAndMaybePostCurrentBlockedReason is called, the register network
         // request event must be posted, because otherwise the counter for uid will never be
         // decremented.
-        trackUid(nri);
+        trackUidAndMaybePostCurrentBlockedReason(nri);
         mHandler.sendMessage(mHandler.obtainMessage(event, nri));
     }
 
-    private void maybeUntrackUid(final NetworkRequestInfo nri) {
-        if (nri.mMessenger == null && nri.mPendingIntent == null) {
+    private void maybeUntrackUidAndClearBlockedReasons(final NetworkRequestInfo nri) {
+        if (!isAppRequest(nri)) {
             // Not an app request.
             return;
         }
         nri.mPerUidCounter.decrementCount(nri.mUid);
+
+        if (!shouldTrackUidsForBlockedStatusCallbacks() || nri.mMessenger == null) {
+            return;
+        }
+        if (!nri.mUidTrackedForBlockedStatus) {
+            Log.wtf(TAG, "Nri is not tracked for sending blocked status: " + nri);
+            return;
+        }
+        nri.mUidTrackedForBlockedStatus = false;
+        synchronized (mBlockedStatusTrackingUids) {
+            final int count = mBlockedStatusTrackingUids.get(nri.mAsUid);
+            if (count > 1) {
+                mBlockedStatusTrackingUids.put(nri.mAsUid, count - 1);
+            } else {
+                mBlockedStatusTrackingUids.delete(nri.mAsUid);
+                mUidBlockedReasons.delete(nri.mAsUid);
+            }
+        }
     }
 
     @Override
@@ -13302,14 +13411,17 @@
         final ArraySet<NetworkRequestInfo> perAppCallbackRequestsToUpdate =
                 getPerAppCallbackRequestsToUpdate();
         final ArraySet<NetworkRequestInfo> nrisToRegister = new ArraySet<>(nris);
-        // This method does not need to modify perUidCounter because:
+        // This method does not need to modify perUidCounter and mBlockedStatusTrackingUids because:
         // - |nris| only contains per-app network requests created by ConnectivityService which
-        //   are internal requests and so do not need to be tracked in perUidCounter.
+        //    are internal requests and have no messenger and are not associated with any callbacks,
+        //    and so do not need to be tracked in perUidCounter and mBlockedStatusTrackingUids.
         // - The requests in perAppCallbackRequestsToUpdate are removed, modified, and re-added,
         //   but the same number of requests is removed and re-added, and none of the requests
-        //   changes mUid, so the perUidCounter before and after this method remains the same.
-        //   Re-adding the requests does not modify perUidCounter (that is done when the app
-        //   registers the request), so removing them must not modify perUidCounter.
+        //   changes mUid and mAsUid, so the perUidCounter and mBlockedStatusTrackingUids before
+        //   and after this method remains the same. Re-adding the requests does not modify
+        //   perUidCounter and mBlockedStatusTrackingUids (that is done when the app registers the
+        //   request), so removing them must not modify perUidCounter and mBlockedStatusTrackingUids
+        //   either.
         // TODO(b/341228979): Modify nris in place instead of removing them and re-adding them
         handleRemoveNetworkRequests(perAppCallbackRequestsToUpdate,
                 false /* untrackUids */);
@@ -13547,10 +13659,17 @@
             throw new IllegalStateException(e);
         }
 
-        try {
-            mBpfNetMaps.setDataSaverEnabled(enable);
-        } catch (ServiceSpecificException | UnsupportedOperationException e) {
-            Log.e(TAG, "Failed to set data saver " + enable + " : " + e);
+        synchronized (mBlockedStatusTrackingUids) {
+            try {
+                mBpfNetMaps.setDataSaverEnabled(enable);
+            } catch (ServiceSpecificException | UnsupportedOperationException e) {
+                Log.e(TAG, "Failed to set data saver " + enable + " : " + e);
+                return;
+            }
+
+            if (shouldTrackUidsForBlockedStatusCallbacks()) {
+                updateTrackingUidsBlockedReasons();
+            }
         }
     }
 
@@ -13586,10 +13705,17 @@
             throw new IllegalArgumentException("setUidFirewallRule with invalid rule: " + rule);
         }
 
-        try {
-            mBpfNetMaps.setUidRule(chain, uid, firewallRule);
-        } catch (ServiceSpecificException e) {
-            throw new IllegalStateException(e);
+        synchronized (mBlockedStatusTrackingUids) {
+            try {
+                mBpfNetMaps.setUidRule(chain, uid, firewallRule);
+            } catch (ServiceSpecificException e) {
+                throw new IllegalStateException(e);
+            }
+            if (shouldTrackUidsForBlockedStatusCallbacks()
+                    && mBlockedStatusTrackingUids.get(uid, 0) != 0) {
+                mHandler.sendMessage(mHandler.obtainMessage(EVENT_BLOCKED_REASONS_CHANGED,
+                        List.of(new Pair<>(uid, mBpfNetMaps.getUidNetworkingBlockedReasons(uid)))));
+            }
         }
     }
 
@@ -13667,10 +13793,15 @@
                     "Chain (" + chain + ") can not be controlled by setFirewallChainEnabled");
         }
 
-        try {
-            mBpfNetMaps.setChildChain(chain, enable);
-        } catch (ServiceSpecificException e) {
-            throw new IllegalStateException(e);
+        synchronized (mBlockedStatusTrackingUids) {
+            try {
+                mBpfNetMaps.setChildChain(chain, enable);
+            } catch (ServiceSpecificException e) {
+                throw new IllegalStateException(e);
+            }
+            if (shouldTrackUidsForBlockedStatusCallbacks()) {
+                updateTrackingUidsBlockedReasons();
+            }
         }
 
         if (mDeps.isAtLeastU() && enable) {
@@ -13682,6 +13813,22 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
+    @GuardedBy("mBlockedStatusTrackingUids")
+    private void updateTrackingUidsBlockedReasons() {
+        if (mBlockedStatusTrackingUids.size() == 0) {
+            return;
+        }
+        final ArrayList<Pair<Integer, Integer>> uidBlockedReasonsList = new ArrayList<>();
+        for (int i = 0; i < mBlockedStatusTrackingUids.size(); i++) {
+            final int uid = mBlockedStatusTrackingUids.keyAt(i);
+            uidBlockedReasonsList.add(
+                    new Pair<>(uid, mBpfNetMaps.getUidNetworkingBlockedReasons(uid)));
+        }
+        mHandler.sendMessage(mHandler.obtainMessage(EVENT_BLOCKED_REASONS_CHANGED,
+                uidBlockedReasonsList));
+    }
+
     @Override
     public boolean getFirewallChainEnabled(final int chain) {
         enforceNetworkStackOrSettingsPermission();
@@ -13706,7 +13853,13 @@
             return;
         }
 
-        mBpfNetMaps.replaceUidChain(chain, uids);
+        synchronized (mBlockedStatusTrackingUids) {
+            mBpfNetMaps.replaceUidChain(chain, uids);
+
+            if (shouldTrackUidsForBlockedStatusCallbacks()) {
+                updateTrackingUidsBlockedReasons();
+            }
+        }
     }
 
     @Override
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 5ed4696..bf40462 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -36,11 +36,22 @@
 import static android.content.pm.PackageManager.FEATURE_WIFI_DIRECT;
 import static android.content.pm.PackageManager.GET_PERMISSIONS;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
+import static android.net.ConnectivityManager.BLOCKED_METERED_REASON_ADMIN_DISABLED;
+import static android.net.ConnectivityManager.BLOCKED_METERED_REASON_USER_RESTRICTED;
+import static android.net.ConnectivityManager.BLOCKED_REASON_APP_BACKGROUND;
+import static android.net.ConnectivityManager.BLOCKED_REASON_APP_STANDBY;
+import static android.net.ConnectivityManager.BLOCKED_REASON_BATTERY_SAVER;
+import static android.net.ConnectivityManager.BLOCKED_REASON_DOZE;
+import static android.net.ConnectivityManager.BLOCKED_REASON_LOW_POWER_STANDBY;
+import static android.net.ConnectivityManager.BLOCKED_REASON_OEM_DENY;
+import static android.net.ConnectivityManager.BLOCKED_REASON_RESTRICTED_MODE;
 import static android.net.ConnectivityManager.EXTRA_NETWORK;
 import static android.net.ConnectivityManager.EXTRA_NETWORK_REQUEST;
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_BACKGROUND;
 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_METERED_DENY_ADMIN;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_METERED_DENY_USER;
 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;
@@ -196,6 +207,7 @@
 import com.android.testutils.CompatUtil;
 import com.android.testutils.ConnectivityModuleTest;
 import com.android.testutils.DevSdkIgnoreRule;
+import com.android.testutils.DevSdkIgnoreRule.IgnoreAfter;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import com.android.testutils.DeviceConfigRule;
 import com.android.testutils.DeviceInfoUtils;
@@ -2472,8 +2484,11 @@
         }
     }
 
+    // On V+, ConnectivityService generates blockedReasons based on bpf map contents even if the
+    // otherUid does not exist on device. So if allowlist chain (e.g. background chain) is enabled,
+    // blockedReasons for otherUid will not be BLOCKED_REASON_NONE.
     @AppModeFull(reason = "Cannot get WifiManager in instant app mode")
-    @Test
+    @Test @IgnoreAfter(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
     public void testBlockedStatusCallback() throws Exception {
         // Cannot use @IgnoreUpTo(Build.VERSION_CODES.R) because this test also requires API 31
         // shims, and @IgnoreUpTo does not check that.
@@ -3713,6 +3728,260 @@
                 Process.myUid() + 1, EXPECT_OPEN);
     }
 
+    private int getBlockedReason(final int chain) {
+        switch(chain) {
+            case FIREWALL_CHAIN_DOZABLE:
+                return BLOCKED_REASON_DOZE;
+            case  FIREWALL_CHAIN_POWERSAVE:
+                return BLOCKED_REASON_BATTERY_SAVER;
+            case  FIREWALL_CHAIN_RESTRICTED:
+                return BLOCKED_REASON_RESTRICTED_MODE;
+            case  FIREWALL_CHAIN_LOW_POWER_STANDBY:
+                return BLOCKED_REASON_LOW_POWER_STANDBY;
+            case  FIREWALL_CHAIN_BACKGROUND:
+                return BLOCKED_REASON_APP_BACKGROUND;
+            case  FIREWALL_CHAIN_STANDBY:
+                return BLOCKED_REASON_APP_STANDBY;
+            case FIREWALL_CHAIN_METERED_DENY_USER:
+                return BLOCKED_METERED_REASON_USER_RESTRICTED;
+            case FIREWALL_CHAIN_METERED_DENY_ADMIN:
+                return BLOCKED_METERED_REASON_ADMIN_DISABLED;
+            case FIREWALL_CHAIN_OEM_DENY_1:
+            case FIREWALL_CHAIN_OEM_DENY_2:
+            case FIREWALL_CHAIN_OEM_DENY_3:
+                return BLOCKED_REASON_OEM_DENY;
+            default:
+                throw new IllegalArgumentException(
+                        "Failed to find blockedReasons for chain: " + chain);
+        }
+    }
+
+    private void doTestBlockedReasons_setUidFirewallRule(final int chain, final boolean metered)
+            throws Exception {
+        assumeTrue(mPackageManager.hasSystemFeature(FEATURE_WIFI));
+
+        // Store current Wi-Fi metered value and update metered value
+        final Network currentWifiNetwork = mCtsNetUtils.ensureWifiConnected();
+        final NetworkCapabilities wifiNetworkCapabilities = callWithShellPermissionIdentity(
+                () -> mCm.getNetworkCapabilities(currentWifiNetwork));
+        final String ssid = unquoteSSID(wifiNetworkCapabilities.getSsid());
+        final boolean oldMeteredValue = wifiNetworkCapabilities.isMetered();
+        final Network wifiNetwork =
+                setWifiMeteredStatusAndWait(ssid, metered, true /* waitForValidation */);
+
+        // Store current firewall chains status. This test operates on the chain that is passed in,
+        // but also always operates on FIREWALL_CHAIN_METERED_DENY_USER to ensure that metered
+        // chains are tested as well.
+        final int myUid = Process.myUid();
+        final boolean wasChainEnabled = runWithShellPermissionIdentity(
+                () -> mCm.getFirewallChainEnabled(chain), NETWORK_SETTINGS);
+        final int previousFirewallRule = runWithShellPermissionIdentity(
+                () -> mCm.getUidFirewallRule(chain, myUid));
+        final int previousMeteredDenyFirewallRule = runWithShellPermissionIdentity(
+                () -> mCm.getUidFirewallRule(FIREWALL_CHAIN_METERED_DENY_USER, myUid));
+
+        final DetailedBlockedStatusCallback cb = new DetailedBlockedStatusCallback();
+        networkCallbackRule.requestNetwork(makeWifiNetworkRequest(), cb);
+        testAndCleanup(() -> {
+            int blockedReasonsWithoutChain = BLOCKED_REASON_NONE;
+            int blockedReasonsWithChain = getBlockedReason(chain);
+            int blockedReasonsWithChainAndLockDown =
+                    getBlockedReason(chain) | BLOCKED_REASON_LOCKDOWN_VPN;
+            if (metered) {
+                blockedReasonsWithoutChain |= BLOCKED_METERED_REASON_USER_RESTRICTED;
+                blockedReasonsWithChain |= BLOCKED_METERED_REASON_USER_RESTRICTED;
+                blockedReasonsWithChainAndLockDown |= BLOCKED_METERED_REASON_USER_RESTRICTED;
+            }
+
+            // Set RULE_DENY on target chain and metered deny chain
+            runWithShellPermissionIdentity(() -> {
+                mCm.setFirewallChainEnabled(chain, true /* enable */);
+                mCm.setUidFirewallRule(chain, myUid, FIREWALL_RULE_DENY);
+                mCm.setUidFirewallRule(FIREWALL_CHAIN_METERED_DENY_USER, myUid,
+                        FIREWALL_RULE_DENY);
+            }, NETWORK_SETTINGS);
+            cb.eventuallyExpectBlockedStatusCallback(wifiNetwork, blockedReasonsWithChain);
+
+            // Set VPN lockdown
+            final Range<Integer> myUidRange = new Range<>(myUid, myUid);
+            runWithShellPermissionIdentity(() -> setRequireVpnForUids(
+                    true /* requireVpn */, List.of(myUidRange)), NETWORK_SETTINGS);
+            cb.eventuallyExpectBlockedStatusCallback(wifiNetwork,
+                    blockedReasonsWithChainAndLockDown);
+
+            // Unset VPN lockdown
+            runWithShellPermissionIdentity(() -> setRequireVpnForUids(
+                    false /* requireVpn */, List.of(myUidRange)), NETWORK_SETTINGS);
+            cb.eventuallyExpectBlockedStatusCallback(wifiNetwork, blockedReasonsWithChain);
+
+            // Set RULE_ALLOW on target chain
+            runWithShellPermissionIdentity(
+                    () -> mCm.setUidFirewallRule(chain, myUid, FIREWALL_RULE_ALLOW),
+                    NETWORK_SETTINGS);
+            cb.eventuallyExpectBlockedStatusCallback(wifiNetwork, blockedReasonsWithoutChain);
+
+            // Set RULE_ALLOW on metered deny chain
+            runWithShellPermissionIdentity(() -> mCm.setUidFirewallRule(
+                            FIREWALL_CHAIN_METERED_DENY_USER, myUid, FIREWALL_RULE_ALLOW),
+                    NETWORK_SETTINGS);
+            if (metered) {
+                cb.eventuallyExpectBlockedStatusCallback(wifiNetwork, BLOCKED_REASON_NONE);
+            }
+        }, /* cleanup */ () -> {
+            setWifiMeteredStatusAndWait(ssid, oldMeteredValue, false /* waitForValidation */);
+        }, /* cleanup */ () -> {
+            mCm.unregisterNetworkCallback(cb);
+        }, /* cleanup */ () -> {
+            runWithShellPermissionIdentity(() -> {
+                mCm.setFirewallChainEnabled(chain, wasChainEnabled);
+                try {
+                    mCm.setUidFirewallRule(chain, myUid, previousFirewallRule);
+                } catch (IllegalStateException ignored) {
+                    // Removing match causes an exception when the rule entry for the uid does
+                    // not exist. But this is fine and can be ignored.
+                }
+                try {
+                    mCm.setUidFirewallRule(FIREWALL_CHAIN_METERED_DENY_USER, myUid,
+                            previousMeteredDenyFirewallRule);
+                } catch (IllegalStateException ignored) {
+                    // Removing match causes an exception when the rule entry for the uid does
+                    // not exist. But this is fine and can be ignored.
+                }
+            }, NETWORK_SETTINGS);
+        });
+    }
+
+    @AppModeFull(reason = "Cannot get WifiManager in instant app mode")
+    @Test @IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE) @ConnectivityModuleTest
+    public void testBlockedReasons_setUidFirewallRule() throws Exception {
+        doTestBlockedReasons_setUidFirewallRule(FIREWALL_CHAIN_DOZABLE, true /* metered */);
+        doTestBlockedReasons_setUidFirewallRule(FIREWALL_CHAIN_STANDBY, false /* metered */);
+    }
+
+    private void doTestBlockedReasons_setFirewallChainEnabled(final int chain) {
+        // Store current firewall chains status.
+        final int myUid = Process.myUid();
+        // TODO(b/342508466): Use runAsShell
+        final boolean wasChainEnabled = runWithShellPermissionIdentity(
+                () -> mCm.getFirewallChainEnabled(chain), NETWORK_SETTINGS);
+        final int previousFirewallRule = runWithShellPermissionIdentity(
+                () -> mCm.getUidFirewallRule(chain, myUid), NETWORK_SETTINGS);
+
+        final DetailedBlockedStatusCallback cb = new DetailedBlockedStatusCallback();
+        networkCallbackRule.registerDefaultNetworkCallback(cb);
+        final Network network = cb.expect(CallbackEntry.AVAILABLE).getNetwork();
+        testAndCleanup(() -> {
+            // Disable chain and set RULE_DENY on target chain
+            runWithShellPermissionIdentity(() -> {
+                mCm.setFirewallChainEnabled(chain, false /* enable */);
+                mCm.setUidFirewallRule(chain, myUid, FIREWALL_RULE_DENY);
+            }, NETWORK_SETTINGS);
+            cb.eventuallyExpectBlockedStatusCallback(network, BLOCKED_REASON_NONE);
+
+            // Enable chain
+            runWithShellPermissionIdentity(() -> {
+                mCm.setFirewallChainEnabled(chain, true /* enable */);
+            }, NETWORK_SETTINGS);
+            cb.eventuallyExpectBlockedStatusCallback(network, getBlockedReason(chain));
+
+            // Disable chain
+            runWithShellPermissionIdentity(() -> {
+                mCm.setFirewallChainEnabled(chain, false /* enable */);
+            }, NETWORK_SETTINGS);
+            cb.eventuallyExpectBlockedStatusCallback(network, BLOCKED_REASON_NONE);
+        }, /* cleanup */ () -> {
+            runWithShellPermissionIdentity(() -> {
+                mCm.setFirewallChainEnabled(chain, wasChainEnabled);
+                try {
+                    mCm.setUidFirewallRule(chain, myUid, previousFirewallRule);
+                } catch (IllegalStateException ignored) {
+                    // Removing match causes an exception when the rule entry for the uid does
+                    // not exist. But this is fine and can be ignored.
+                }
+            }, NETWORK_SETTINGS);
+        });
+    }
+
+    @Test @IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE) @ConnectivityModuleTest
+    public void testBlockedReasons_setFirewallChainEnabled() {
+        doTestBlockedReasons_setFirewallChainEnabled(FIREWALL_CHAIN_POWERSAVE);
+        doTestBlockedReasons_setFirewallChainEnabled(FIREWALL_CHAIN_OEM_DENY_1);
+    }
+
+    private void doTestBlockedReasons_replaceFirewallChain(
+            final int chain, final boolean isAllowList) {
+        // Store current firewall chains status.
+        final int myUid = Process.myUid();
+        final boolean wasChainEnabled = runWithShellPermissionIdentity(
+                () -> mCm.getFirewallChainEnabled(chain), NETWORK_SETTINGS);
+        final int previousFirewallRule = runWithShellPermissionIdentity(
+                () -> mCm.getUidFirewallRule(chain, myUid), NETWORK_SETTINGS);
+
+        final DetailedBlockedStatusCallback cb = new DetailedBlockedStatusCallback();
+        networkCallbackRule.registerDefaultNetworkCallback(cb);
+        final Network network = cb.expect(CallbackEntry.AVAILABLE).getNetwork();
+        testAndCleanup(() -> {
+            cb.eventuallyExpectBlockedStatusCallback(network, BLOCKED_REASON_NONE);
+
+            // Remove uid from the target chain and enable chain
+            runWithShellPermissionIdentity(() -> {
+                // Note that this removes *all* UIDs from the chain, not just the UID that is
+                // being tested. This is probably OK since FIREWALL_CHAIN_OEM_DENY_2 is unused
+                // in AOSP and FIREWALL_CHAIN_BACKGROUND is probably empty when running this
+                // test (since nothing is in the foreground).
+                //
+                // TODO(b/342508466): add a getFirewallUidChainContents or similar method to fetch
+                // chain contents, and update this test to use it.
+                mCm.replaceFirewallChain(chain, new int[0]);
+                mCm.setFirewallChainEnabled(chain, true /* enable */);
+            }, NETWORK_SETTINGS);
+
+            if (isAllowList) {
+                cb.eventuallyExpectBlockedStatusCallback(network, getBlockedReason(chain));
+            } else {
+                cb.assertNoBlockedStatusCallback();
+            }
+
+            // Put uid on the target chain
+            runWithShellPermissionIdentity(
+                    () -> mCm.replaceFirewallChain(chain, new int[]{myUid}), NETWORK_SETTINGS);
+
+            if (isAllowList) {
+                cb.eventuallyExpectBlockedStatusCallback(network, BLOCKED_REASON_NONE);
+            } else {
+                cb.eventuallyExpectBlockedStatusCallback(network, getBlockedReason(chain));
+            }
+
+            // Remove uid from the target chain
+            runWithShellPermissionIdentity(
+                    () -> mCm.replaceFirewallChain(chain, new int[0]), NETWORK_SETTINGS);
+
+            if (isAllowList) {
+                cb.eventuallyExpectBlockedStatusCallback(network, getBlockedReason(chain));
+            } else {
+                cb.eventuallyExpectBlockedStatusCallback(network, BLOCKED_REASON_NONE);
+            }
+        }, /* cleanup */ () -> {
+            runWithShellPermissionIdentity(() -> {
+                mCm.setFirewallChainEnabled(chain, wasChainEnabled);
+                try {
+                    mCm.setUidFirewallRule(chain, myUid, previousFirewallRule);
+                } catch (IllegalStateException ignored) {
+                    // Removing match causes an exception when the rule entry for the uid does
+                    // not exist. But this is fine and can be ignored.
+                }
+            }, NETWORK_SETTINGS);
+        });
+    }
+
+    @Test @IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE) @ConnectivityModuleTest
+    public void testBlockedReasons_replaceFirewallChain() {
+        doTestBlockedReasons_replaceFirewallChain(
+                FIREWALL_CHAIN_BACKGROUND, true /* isAllowChain */);
+        doTestBlockedReasons_replaceFirewallChain(
+                FIREWALL_CHAIN_OEM_DENY_2, false /* isAllowChain */);
+    }
+
     private void assumeTestSApis() {
         // Cannot use @IgnoreUpTo(Build.VERSION_CODES.R) because this test also requires API 31
         // shims, and @IgnoreUpTo does not check that.
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index f70cbb1..6e63807 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -1746,7 +1746,15 @@
 
     private void setBlockedReasonChanged(int blockedReasons) {
         mBlockedReasons = blockedReasons;
-        mPolicyCallback.onUidBlockedReasonChanged(Process.myUid(), blockedReasons);
+        if (mDeps.isAtLeastV()) {
+            visibleOnHandlerThread(mCsHandlerThread.getThreadHandler(),
+                    () -> mService.handleBlockedReasonsChanged(
+                            List.of(new Pair<>(Process.myUid(), blockedReasons))
+
+                    ));
+        } else {
+            mPolicyCallback.onUidBlockedReasonChanged(Process.myUid(), blockedReasons);
+        }
     }
 
     private Nat464Xlat getNat464Xlat(NetworkAgentWrapper mna) {
@@ -1927,11 +1935,16 @@
         mService.mLingerDelayMs = TEST_LINGER_DELAY_MS;
         mService.mNascentDelayMs = TEST_NASCENT_DELAY_MS;
 
-        final ArgumentCaptor<NetworkPolicyCallback> policyCallbackCaptor =
-                ArgumentCaptor.forClass(NetworkPolicyCallback.class);
-        verify(mNetworkPolicyManager).registerNetworkPolicyCallback(any(),
-                policyCallbackCaptor.capture());
-        mPolicyCallback = policyCallbackCaptor.getValue();
+        if (mDeps.isAtLeastV()) {
+            verify(mNetworkPolicyManager, never()).registerNetworkPolicyCallback(any(), any());
+            mPolicyCallback = null;
+        } else {
+            final ArgumentCaptor<NetworkPolicyCallback> policyCallbackCaptor =
+                    ArgumentCaptor.forClass(NetworkPolicyCallback.class);
+            verify(mNetworkPolicyManager).registerNetworkPolicyCallback(any(),
+                    policyCallbackCaptor.capture());
+            mPolicyCallback = policyCallbackCaptor.getValue();
+        }
 
         // Create local CM before sending system ready so that we can answer
         // getSystemService() correctly.
diff --git a/tests/unit/java/com/android/server/connectivityservice/CSBlockedReasonsTest.kt b/tests/unit/java/com/android/server/connectivityservice/CSBlockedReasonsTest.kt
new file mode 100644
index 0000000..2694916
--- /dev/null
+++ b/tests/unit/java/com/android/server/connectivityservice/CSBlockedReasonsTest.kt
@@ -0,0 +1,346 @@
+/*
+ * Copyright (C) 2024 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.ConnectivityManager.BLOCKED_METERED_REASON_DATA_SAVER
+import android.net.ConnectivityManager.BLOCKED_METERED_REASON_USER_RESTRICTED
+import android.net.ConnectivityManager.BLOCKED_REASON_APP_BACKGROUND
+import android.net.ConnectivityManager.BLOCKED_REASON_DOZE
+import android.net.ConnectivityManager.BLOCKED_REASON_NONE
+import android.net.ConnectivityManager.FIREWALL_CHAIN_BACKGROUND
+import android.net.ConnectivityManager.FIREWALL_CHAIN_DOZABLE
+import android.net.ConnectivityManager.FIREWALL_CHAIN_METERED_DENY_USER
+import android.net.ConnectivityManager.FIREWALL_RULE_ALLOW
+import android.net.ConnectivityManager.FIREWALL_RULE_DENY
+import android.net.ConnectivitySettingsManager
+import android.net.Network
+import android.net.NetworkCapabilities
+import android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET
+import android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED
+import android.net.NetworkCapabilities.NET_CAPABILITY_NOT_SUSPENDED
+import android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED
+import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
+import android.net.NetworkCapabilities.TRANSPORT_WIFI
+import android.net.NetworkRequest
+import android.os.Build
+import android.os.Process
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import com.android.testutils.DevSdkIgnoreRunner
+import com.android.testutils.RecorderCallback.CallbackEntry.BlockedStatusInt
+import com.android.testutils.TestableNetworkCallback
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyBoolean
+import org.mockito.Mockito.doReturn
+
+private fun cellNc() = NetworkCapabilities.Builder()
+        .addTransportType(TRANSPORT_CELLULAR)
+        .addCapability(NET_CAPABILITY_INTERNET)
+        .addCapability(NET_CAPABILITY_NOT_SUSPENDED)
+        .addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
+        .build()
+private fun cellRequest() = NetworkRequest.Builder()
+        .addTransportType(TRANSPORT_CELLULAR)
+        .build()
+private fun wifiNc() = NetworkCapabilities.Builder()
+        .addTransportType(TRANSPORT_WIFI)
+        .addCapability(NET_CAPABILITY_INTERNET)
+        .addCapability(NET_CAPABILITY_NOT_VCN_MANAGED)
+        .addCapability(NET_CAPABILITY_NOT_METERED)
+        .build()
+private fun wifiRequest() = NetworkRequest.Builder()
+        .addTransportType(TRANSPORT_WIFI)
+        .build()
+
+@RunWith(DevSdkIgnoreRunner::class)
+@IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+class CSBlockedReasonsTest : CSTest() {
+
+    inner class DetailedBlockedStatusCallback : TestableNetworkCallback() {
+        override fun onBlockedStatusChanged(network: Network, blockedReasons: Int) {
+            history.add(BlockedStatusInt(network, blockedReasons))
+        }
+
+        fun expectBlockedStatusChanged(network: Network, blockedReasons: Int) {
+            expect<BlockedStatusInt>(network) { it.reason == blockedReasons }
+        }
+    }
+
+    @Test
+    fun testBlockedReasons_onAvailable() {
+        doReturn(BLOCKED_REASON_DOZE or BLOCKED_METERED_REASON_DATA_SAVER)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+
+        val cellAgent = Agent(nc = cellNc())
+        cellAgent.connect()
+        val wifiAgent = Agent(nc = wifiNc())
+        wifiAgent.connect()
+
+        val cellCb = DetailedBlockedStatusCallback()
+        val wifiCb = DetailedBlockedStatusCallback()
+        cm.requestNetwork(cellRequest(), cellCb)
+        cm.requestNetwork(wifiRequest(), wifiCb)
+
+        cellCb.expectAvailableCallbacks(
+                cellAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_DOZE or BLOCKED_METERED_REASON_DATA_SAVER
+        )
+        wifiCb.expectAvailableCallbacks(
+                wifiAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_DOZE
+        )
+
+        cellAgent.disconnect()
+        wifiAgent.disconnect()
+        cm.unregisterNetworkCallback(cellCb)
+        cm.unregisterNetworkCallback(wifiCb)
+    }
+
+    @Test
+    fun testBlockedReasons_dataSaverChanged() {
+        doReturn(BLOCKED_REASON_APP_BACKGROUND or BLOCKED_METERED_REASON_DATA_SAVER)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        doReturn(true).`when`(netd).bandwidthEnableDataSaver(anyBoolean())
+
+        val cellCb = DetailedBlockedStatusCallback()
+        val wifiCb = DetailedBlockedStatusCallback()
+        cm.requestNetwork(cellRequest(), cellCb)
+        cm.requestNetwork(wifiRequest(), wifiCb)
+
+        val cellAgent = Agent(nc = cellNc())
+        cellAgent.connect()
+        val wifiAgent = Agent(nc = wifiNc())
+        wifiAgent.connect()
+        cellCb.expectAvailableCallbacks(
+                cellAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_APP_BACKGROUND or BLOCKED_METERED_REASON_DATA_SAVER
+        )
+        wifiCb.expectAvailableCallbacks(
+                wifiAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_APP_BACKGROUND
+        )
+
+        // Disable data saver
+        doReturn(BLOCKED_REASON_APP_BACKGROUND)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.setDataSaverEnabled(false)
+        cellCb.expectBlockedStatusChanged(cellAgent.network, BLOCKED_REASON_APP_BACKGROUND)
+
+        // Enable data saver
+        doReturn(BLOCKED_REASON_APP_BACKGROUND or BLOCKED_METERED_REASON_DATA_SAVER)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.setDataSaverEnabled(true)
+        cellCb.expectBlockedStatusChanged(
+                cellAgent.network,
+                BLOCKED_REASON_APP_BACKGROUND or BLOCKED_METERED_REASON_DATA_SAVER
+        )
+        // BlockedStatus does not change for the non-metered network
+        wifiCb.assertNoCallback()
+
+        cellAgent.disconnect()
+        wifiAgent.disconnect()
+        cm.unregisterNetworkCallback(cellCb)
+        cm.unregisterNetworkCallback(wifiCb)
+    }
+
+    @Test
+    fun testBlockedReasons_setUidFirewallRule() {
+        doReturn(BLOCKED_REASON_DOZE or BLOCKED_METERED_REASON_USER_RESTRICTED)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+
+        val cellCb = DetailedBlockedStatusCallback()
+        val wifiCb = DetailedBlockedStatusCallback()
+        cm.requestNetwork(cellRequest(), cellCb)
+        cm.requestNetwork(wifiRequest(), wifiCb)
+
+        val cellAgent = Agent(nc = cellNc())
+        cellAgent.connect()
+        val wifiAgent = Agent(nc = wifiNc())
+        wifiAgent.connect()
+        cellCb.expectAvailableCallbacks(
+                cellAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_DOZE or BLOCKED_METERED_REASON_USER_RESTRICTED
+        )
+        wifiCb.expectAvailableCallbacks(
+                wifiAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_DOZE
+        )
+
+        // Set RULE_ALLOW on metered deny chain
+        doReturn(BLOCKED_REASON_DOZE)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.setUidFirewallRule(
+                FIREWALL_CHAIN_METERED_DENY_USER,
+                Process.myUid(),
+                FIREWALL_RULE_ALLOW
+        )
+        cellCb.expectBlockedStatusChanged(
+                cellAgent.network,
+                BLOCKED_REASON_DOZE
+        )
+        // BlockedStatus does not change for the non-metered network
+        wifiCb.assertNoCallback()
+
+        // Set RULE_DENY on metered deny chain
+        doReturn(BLOCKED_REASON_DOZE or BLOCKED_METERED_REASON_USER_RESTRICTED)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.setUidFirewallRule(
+                FIREWALL_CHAIN_METERED_DENY_USER,
+                Process.myUid(),
+                FIREWALL_RULE_DENY
+        )
+        cellCb.expectBlockedStatusChanged(
+                cellAgent.network,
+                BLOCKED_REASON_DOZE or BLOCKED_METERED_REASON_USER_RESTRICTED
+        )
+        // BlockedStatus does not change for the non-metered network
+        wifiCb.assertNoCallback()
+
+        cellAgent.disconnect()
+        wifiAgent.disconnect()
+        cm.unregisterNetworkCallback(cellCb)
+        cm.unregisterNetworkCallback(wifiCb)
+    }
+
+    @Test
+    fun testBlockedReasons_setFirewallChainEnabled() {
+        doReturn(BLOCKED_REASON_NONE)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+
+        val wifiCb = DetailedBlockedStatusCallback()
+        cm.requestNetwork(wifiRequest(), wifiCb)
+        val wifiAgent = Agent(nc = wifiNc())
+        wifiAgent.connect()
+        wifiCb.expectAvailableCallbacks(
+                wifiAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_NONE
+        )
+
+        // Enable dozable firewall chain
+        doReturn(BLOCKED_REASON_DOZE)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.setFirewallChainEnabled(FIREWALL_CHAIN_DOZABLE, true)
+        wifiCb.expectBlockedStatusChanged(
+                wifiAgent.network,
+                BLOCKED_REASON_DOZE
+        )
+
+        // Disable dozable firewall chain
+        doReturn(BLOCKED_REASON_NONE)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.setFirewallChainEnabled(FIREWALL_CHAIN_DOZABLE, false)
+        wifiCb.expectBlockedStatusChanged(
+                wifiAgent.network,
+                BLOCKED_REASON_NONE
+        )
+
+        wifiAgent.disconnect()
+        cm.unregisterNetworkCallback(wifiCb)
+    }
+
+    @Test
+    fun testBlockedReasons_replaceFirewallChain() {
+        doReturn(BLOCKED_REASON_APP_BACKGROUND)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+
+        val wifiCb = DetailedBlockedStatusCallback()
+        cm.requestNetwork(wifiRequest(), wifiCb)
+        val wifiAgent = Agent(nc = wifiNc())
+        wifiAgent.connect()
+        wifiCb.expectAvailableCallbacks(
+                wifiAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_APP_BACKGROUND
+        )
+
+        // Put uid on background firewall chain
+        doReturn(BLOCKED_REASON_NONE)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.replaceFirewallChain(FIREWALL_CHAIN_BACKGROUND, intArrayOf(Process.myUid()))
+        wifiCb.expectBlockedStatusChanged(
+                wifiAgent.network,
+                BLOCKED_REASON_NONE
+        )
+
+        // Remove uid from background firewall chain
+        doReturn(BLOCKED_REASON_APP_BACKGROUND)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+        cm.replaceFirewallChain(FIREWALL_CHAIN_BACKGROUND, intArrayOf())
+        wifiCb.expectBlockedStatusChanged(
+                wifiAgent.network,
+                BLOCKED_REASON_APP_BACKGROUND
+        )
+
+        wifiAgent.disconnect()
+        cm.unregisterNetworkCallback(wifiCb)
+    }
+
+    @Test
+    fun testBlockedReasons_perAppDefaultNetwork() {
+        doReturn(BLOCKED_METERED_REASON_USER_RESTRICTED)
+                .`when`(bpfNetMaps).getUidNetworkingBlockedReasons(Process.myUid())
+
+        val cellCb = DetailedBlockedStatusCallback()
+        val wifiCb = DetailedBlockedStatusCallback()
+        cm.requestNetwork(cellRequest(), cellCb)
+        cm.requestNetwork(wifiRequest(), wifiCb)
+
+        val cellAgent = Agent(nc = cellNc())
+        cellAgent.connect()
+        val wifiAgent = Agent(nc = wifiNc())
+        wifiAgent.connect()
+
+        val cb = DetailedBlockedStatusCallback()
+        cm.registerDefaultNetworkCallback(cb)
+        cb.expectAvailableCallbacks(
+                wifiAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_NONE
+        )
+
+        // CS must send correct blocked reasons after per app default network change
+        ConnectivitySettingsManager.setMobileDataPreferredUids(context, setOf(Process.myUid()))
+        service.updateMobileDataPreferredUids()
+        cb.expectAvailableCallbacks(
+                cellAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_METERED_REASON_USER_RESTRICTED
+        )
+
+        // Remove per app default network request
+        ConnectivitySettingsManager.setMobileDataPreferredUids(context, setOf())
+        service.updateMobileDataPreferredUids()
+        cb.expectAvailableCallbacks(
+                wifiAgent.network,
+                validated = false,
+                blockedReason = BLOCKED_REASON_NONE
+        )
+
+        cellAgent.disconnect()
+        wifiAgent.disconnect()
+        cm.unregisterNetworkCallback(cellCb)
+        cm.unregisterNetworkCallback(wifiCb)
+        cm.unregisterNetworkCallback(cb)
+    }
+}