Add deny firewall chain for OEM

Bug: 207773349
Bug: 208371987
Test: atest
CtsNetTestCases:android.net.cts.ConnectivityManagerTest#testFirewallBlocking
--iterations 50 && atest ConnectivityServiceTest --iterations 10

Change-Id: I60d5540821abcced03356f366775f16ee369d7f9
(cherry picked from commit d980149817948d11de0631caee8aee3172e4e159)
Merged-In: I60d5540821abcced03356f366775f16ee369d7f9
diff --git a/bpf_progs/bpf_shared.h b/bpf_progs/bpf_shared.h
index 634fbf4..2afb789 100644
--- a/bpf_progs/bpf_shared.h
+++ b/bpf_progs/bpf_shared.h
@@ -133,6 +133,8 @@
     LOW_POWER_STANDBY_MATCH = (1 << 6),
     IIF_MATCH = (1 << 7),
     LOCKDOWN_VPN_MATCH = (1 << 8),
+    OEM_DENY_1_MATCH = (1 << 9),
+    OEM_DENY_2_MATCH = (1 << 10),
 };
 
 enum BpfPermissionMatch {
diff --git a/bpf_progs/netd.c b/bpf_progs/netd.c
index b4ef7eb..c1d0897 100644
--- a/bpf_progs/netd.c
+++ b/bpf_progs/netd.c
@@ -216,6 +216,12 @@
         if ((enabledRules & LOW_POWER_STANDBY_MATCH) && !(uidRules & LOW_POWER_STANDBY_MATCH)) {
             return BPF_DROP;
         }
+        if ((enabledRules & OEM_DENY_1_MATCH) && (uidRules & OEM_DENY_1_MATCH)) {
+            return BPF_DROP;
+        }
+        if ((enabledRules & OEM_DENY_2_MATCH) && (uidRules & OEM_DENY_2_MATCH)) {
+            return BPF_DROP;
+        }
     }
     if (direction == BPF_INGRESS && skb->ifindex != 1) {
         if (uidRules & IIF_MATCH) {
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 4ecc8a1..9c04e8c 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -992,6 +992,20 @@
      */
     public static final int FIREWALL_CHAIN_LOCKDOWN_VPN = 6;
 
+    /**
+     * Firewall chain used for OEM-specific application restrictions.
+     * Denylist of apps that will not have network access due to OEM-specific restrictions.
+     * @hide
+     */
+    public static final int FIREWALL_CHAIN_OEM_DENY_1 = 7;
+
+    /**
+     * Firewall chain used for OEM-specific application restrictions.
+     * Denylist of apps that will not have network access due to OEM-specific restrictions.
+     * @hide
+     */
+    public static final int FIREWALL_CHAIN_OEM_DENY_2 = 8;
+
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
     @IntDef(flag = false, prefix = "FIREWALL_CHAIN_", value = {
@@ -1000,7 +1014,9 @@
         FIREWALL_CHAIN_POWERSAVE,
         FIREWALL_CHAIN_RESTRICTED,
         FIREWALL_CHAIN_LOW_POWER_STANDBY,
-        FIREWALL_CHAIN_LOCKDOWN_VPN
+        FIREWALL_CHAIN_LOCKDOWN_VPN,
+        FIREWALL_CHAIN_OEM_DENY_1,
+        FIREWALL_CHAIN_OEM_DENY_2
     })
     public @interface FirewallChain {}
     // LINT.ThenChange(packages/modules/Connectivity/service/native/include/Common.h)
diff --git a/service/native/TrafficController.cpp b/service/native/TrafficController.cpp
index 70c7c34..548ecbe 100644
--- a/service/native/TrafficController.cpp
+++ b/service/native/TrafficController.cpp
@@ -74,6 +74,8 @@
 const char* TrafficController::LOCAL_POWERSAVE = "fw_powersave";
 const char* TrafficController::LOCAL_RESTRICTED = "fw_restricted";
 const char* TrafficController::LOCAL_LOW_POWER_STANDBY = "fw_low_power_standby";
+const char* TrafficController::LOCAL_OEM_DENY_1 = "fw_oem_deny_1";
+const char* TrafficController::LOCAL_OEM_DENY_2 = "fw_oem_deny_2";
 
 static_assert(BPF_PERMISSION_INTERNET == INetd::PERMISSION_INTERNET,
               "Mismatch between BPF and AIDL permissions: PERMISSION_INTERNET");
@@ -99,6 +101,8 @@
     FLAG_MSG_TRANS(matchType, LOW_POWER_STANDBY_MATCH, match);
     FLAG_MSG_TRANS(matchType, IIF_MATCH, match);
     FLAG_MSG_TRANS(matchType, LOCKDOWN_VPN_MATCH, match);
+    FLAG_MSG_TRANS(matchType, OEM_DENY_1_MATCH, match);
+    FLAG_MSG_TRANS(matchType, OEM_DENY_2_MATCH, match);
     if (match) {
         return StringPrintf("Unknown match: %u", match);
     }
@@ -335,6 +339,10 @@
             return ALLOWLIST;
         case LOCKDOWN:
             return DENYLIST;
+        case OEM_DENY_1:
+            return DENYLIST;
+        case OEM_DENY_2:
+            return DENYLIST;
         case NONE:
         default:
             return DENYLIST;
@@ -363,6 +371,12 @@
         case LOCKDOWN:
             res = updateOwnerMapEntry(LOCKDOWN_VPN_MATCH, uid, rule, type);
             break;
+        case OEM_DENY_1:
+            res = updateOwnerMapEntry(OEM_DENY_1_MATCH, uid, rule, type);
+            break;
+        case OEM_DENY_2:
+            res = updateOwnerMapEntry(OEM_DENY_2_MATCH, uid, rule, type);
+            break;
         case NONE:
         default:
             ALOGW("Unknown child chain: %d", chain);
@@ -440,6 +454,10 @@
         res = replaceRulesInMap(RESTRICTED_MATCH, uids);
     } else if (!name.compare(LOCAL_LOW_POWER_STANDBY)) {
         res = replaceRulesInMap(LOW_POWER_STANDBY_MATCH, uids);
+    } else if (!name.compare(LOCAL_OEM_DENY_1)) {
+        res = replaceRulesInMap(OEM_DENY_1_MATCH, uids);
+    } else if (!name.compare(LOCAL_OEM_DENY_2)) {
+        res = replaceRulesInMap(OEM_DENY_2_MATCH, uids);
     } else {
         ALOGE("unknown chain name: %s", name.c_str());
         return -EINVAL;
@@ -479,6 +497,12 @@
         case LOW_POWER_STANDBY:
             match = LOW_POWER_STANDBY_MATCH;
             break;
+        case OEM_DENY_1:
+            match = OEM_DENY_1_MATCH;
+            break;
+        case OEM_DENY_2:
+            match = OEM_DENY_2_MATCH;
+            break;
         default:
             return -EINVAL;
     }
diff --git a/service/native/TrafficControllerTest.cpp b/service/native/TrafficControllerTest.cpp
index c44b9d6..9ec50af 100644
--- a/service/native/TrafficControllerTest.cpp
+++ b/service/native/TrafficControllerTest.cpp
@@ -308,6 +308,8 @@
     checkUidOwnerRuleForChain(RESTRICTED, RESTRICTED_MATCH);
     checkUidOwnerRuleForChain(LOW_POWER_STANDBY, LOW_POWER_STANDBY_MATCH);
     checkUidOwnerRuleForChain(LOCKDOWN, LOCKDOWN_VPN_MATCH);
+    checkUidOwnerRuleForChain(OEM_DENY_1, OEM_DENY_1_MATCH);
+    checkUidOwnerRuleForChain(OEM_DENY_2, OEM_DENY_2_MATCH);
     ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(NONE, TEST_UID, ALLOW, ALLOWLIST));
     ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(INVALID_CHAIN, TEST_UID, ALLOW, ALLOWLIST));
 }
@@ -319,6 +321,8 @@
     checkUidMapReplace("fw_powersave", uids, POWERSAVE_MATCH);
     checkUidMapReplace("fw_restricted", uids, RESTRICTED_MATCH);
     checkUidMapReplace("fw_low_power_standby", uids, LOW_POWER_STANDBY_MATCH);
+    checkUidMapReplace("fw_oem_deny_1", uids, OEM_DENY_1_MATCH);
+    checkUidMapReplace("fw_oem_deny_2", uids, OEM_DENY_2_MATCH);
     ASSERT_EQ(-EINVAL, mTc.replaceUidOwnerMap("unknow", true, uids));
 }
 
diff --git a/service/native/include/Common.h b/service/native/include/Common.h
index 847acec..48f68ea 100644
--- a/service/native/include/Common.h
+++ b/service/native/include/Common.h
@@ -36,6 +36,8 @@
     RESTRICTED = 4,
     LOW_POWER_STANDBY = 5,
     LOCKDOWN = 6,
+    OEM_DENY_1 = 7,
+    OEM_DENY_2 = 8,
     INVALID_CHAIN
 };
 // LINT.ThenChange(packages/modules/Connectivity/framework/src/android/net/ConnectivityManager.java)
diff --git a/service/native/include/TrafficController.h b/service/native/include/TrafficController.h
index d3d52e2..7a36e1e 100644
--- a/service/native/include/TrafficController.h
+++ b/service/native/include/TrafficController.h
@@ -88,6 +88,8 @@
     static const char* LOCAL_POWERSAVE;
     static const char* LOCAL_RESTRICTED;
     static const char* LOCAL_LOW_POWER_STANDBY;
+    static const char* LOCAL_OEM_DENY_1;
+    static const char* LOCAL_OEM_DENY_2;
 
   private:
     /*
@@ -149,7 +151,7 @@
      * the map right now:
      * - Entry with UID_RULES_CONFIGURATION_KEY:
      *    Store the configuration for the current uid rules. It indicates the device
-     *    is in doze/powersave/standby/restricted/low power standby mode.
+     *    is in doze/powersave/standby/restricted/low power standby/oem deny mode.
      * - Entry with CURRENT_STATS_MAP_CONFIGURATION_KEY:
      *    Stores the current live stats map that kernel program is writing to.
      *    Userspace can do scraping and cleaning job on the other one depending on the
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 0a6c2bd..edcdfa8 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -11363,6 +11363,8 @@
         final int defaultRule;
         switch (chain) {
             case ConnectivityManager.FIREWALL_CHAIN_STANDBY:
+            case ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_1:
+            case ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_2:
                 defaultRule = FIREWALL_RULE_ALLOW;
                 break;
             case ConnectivityManager.FIREWALL_CHAIN_DOZABLE:
@@ -11412,6 +11414,12 @@
                     mBpfNetMaps.replaceUidChain("fw_low_power_standby", true /* isAllowList */,
                             uids);
                     break;
+                case ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_1:
+                    mBpfNetMaps.replaceUidChain("fw_oem_deny_1", false /* isAllowList */, uids);
+                    break;
+                case ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_2:
+                    mBpfNetMaps.replaceUidChain("fw_oem_deny_2", false /* isAllowList */, uids);
+                    break;
                 default:
                     throw new IllegalArgumentException("replaceFirewallChain with invalid chain: "
                             + chain);
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index bdda82a..fd6c165 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -37,6 +37,10 @@
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.net.ConnectivityManager.EXTRA_NETWORK;
 import static android.net.ConnectivityManager.EXTRA_NETWORK_REQUEST;
+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_RULE_ALLOW;
+import static android.net.ConnectivityManager.FIREWALL_RULE_DENY;
 import static android.net.ConnectivityManager.PROFILE_NETWORK_PREFERENCE_ENTERPRISE;
 import static android.net.ConnectivityManager.TYPE_BLUETOOTH;
 import static android.net.ConnectivityManager.TYPE_ETHERNET;
@@ -204,6 +208,8 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.OutputStream;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
 import java.net.HttpURLConnection;
 import java.net.Inet4Address;
 import java.net.Inet6Address;
@@ -218,6 +224,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
+import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
@@ -255,6 +262,7 @@
     private static final int NETWORK_CALLBACK_TIMEOUT_MS = 30_000;
     private static final int LISTEN_ACTIVITY_TIMEOUT_MS = 5_000;
     private static final int NO_CALLBACK_TIMEOUT_MS = 100;
+    private static final int SOCKET_TIMEOUT_MS = 100;
     private static final int NUM_TRIES_MULTIPATH_PREF_CHECK = 20;
     private static final long INTERVAL_MULTIPATH_PREF_CHECK_MS = 500;
     // device could have only one interface: data, wifi.
@@ -3316,6 +3324,111 @@
         assertTrue(dumpOutput, dumpOutput.contains("BPF map content"));
     }
 
+    private void checkFirewallBlocking(final DatagramSocket srcSock, final DatagramSocket dstSock,
+            final boolean expectBlock) throws Exception {
+        final Random random = new Random();
+        final byte[] sendData = new byte[100];
+        random.nextBytes(sendData);
+
+        final DatagramPacket pkt = new DatagramPacket(sendData, sendData.length,
+                InetAddresses.parseNumericAddress("::1"), dstSock.getLocalPort());
+        try {
+            srcSock.send(pkt);
+        } catch (IOException e) {
+            if (expectBlock) {
+                return;
+            }
+            fail("Expect not to be blocked by firewall but sending packet was blocked");
+        }
+
+        if (expectBlock) {
+            fail("Expect to be blocked by firewall but sending packet was not blocked");
+        }
+
+        dstSock.receive(pkt);
+        assertArrayEquals(sendData, pkt.getData());
+    }
+
+    private static final boolean EXPECT_PASS = false;
+    private static final boolean EXPECT_BLOCK = true;
+
+    private void doTestFirewallBlockingDenyRule(final int chain) {
+        runWithShellPermissionIdentity(() -> {
+            try (DatagramSocket srcSock = new DatagramSocket();
+                 DatagramSocket dstSock = new DatagramSocket()) {
+                dstSock.setSoTimeout(SOCKET_TIMEOUT_MS);
+
+                // No global config, No uid config
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+
+                // Has global config, No uid config
+                mCm.setFirewallChainEnabled(chain, true /* enable */);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+
+                // Has global config, Has uid config
+                mCm.setUidFirewallRule(chain, Process.myUid(), FIREWALL_RULE_DENY);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_BLOCK);
+
+                // No global config, Has uid config
+                mCm.setFirewallChainEnabled(chain, false /* enable */);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+
+                // No global config, No uid config
+                mCm.setUidFirewallRule(chain, Process.myUid(), FIREWALL_RULE_ALLOW);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+            } finally {
+                mCm.setFirewallChainEnabled(chain, false /* enable */);
+                mCm.setUidFirewallRule(chain, Process.myUid(), FIREWALL_RULE_ALLOW);
+            }
+        }, NETWORK_SETTINGS);
+    }
+
+    private void doTestFirewallBlockingAllowRule(final int chain) {
+        runWithShellPermissionIdentity(() -> {
+            try (DatagramSocket srcSock = new DatagramSocket();
+                 DatagramSocket dstSock = new DatagramSocket()) {
+                dstSock.setSoTimeout(SOCKET_TIMEOUT_MS);
+
+                // No global config, No uid config
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+
+                // Has global config, No uid config
+                mCm.setFirewallChainEnabled(chain, true /* enable */);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_BLOCK);
+
+                // Has global config, Has uid config
+                mCm.setUidFirewallRule(chain, Process.myUid(), FIREWALL_RULE_ALLOW);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+
+                // No global config, Has uid config
+                mCm.setFirewallChainEnabled(chain, false /* enable */);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+
+                // No global config, No uid config
+                mCm.setUidFirewallRule(chain, Process.myUid(), FIREWALL_RULE_DENY);
+                checkFirewallBlocking(srcSock, dstSock, EXPECT_PASS);
+            } finally {
+                mCm.setFirewallChainEnabled(chain, false /* enable */);
+                mCm.setUidFirewallRule(chain, Process.myUid(), FIREWALL_RULE_DENY);
+            }
+        }, NETWORK_SETTINGS);
+    }
+
+    @Test @IgnoreUpTo(SC_V2)
+    public void testFirewallBlocking() {
+        // Following tests affect the actual state of networking on the device after the test.
+        // This might cause unexpected behaviour of the device. So, we skip them for now.
+        // We will enable following tests after adding the logic of firewall state restoring.
+        // doTestFirewallBlockingAllowRule(FIREWALL_CHAIN_DOZABLE);
+        // doTestFirewallBlockingAllowRule(FIREWALL_CHAIN_POWERSAVE);
+        // doTestFirewallBlockingAllowRule(FIREWALL_CHAIN_RESTRICTED);
+        // doTestFirewallBlockingAllowRule(FIREWALL_CHAIN_LOW_POWER_STANDBY);
+
+        // doTestFirewallBlockingDenyRule(FIREWALL_CHAIN_STANDBY);
+        doTestFirewallBlockingDenyRule(FIREWALL_CHAIN_OEM_DENY_1);
+        doTestFirewallBlockingDenyRule(FIREWALL_CHAIN_OEM_DENY_2);
+    }
+
     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 44550e6..74731c3 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -52,8 +52,16 @@
 import static android.net.ConnectivityManager.CONNECTIVITY_ACTION;
 import static android.net.ConnectivityManager.EXTRA_NETWORK_INFO;
 import static android.net.ConnectivityManager.EXTRA_NETWORK_TYPE;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_DOZABLE;
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_LOCKDOWN_VPN;
+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_POWERSAVE;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_RESTRICTED;
+import static android.net.ConnectivityManager.FIREWALL_CHAIN_STANDBY;
 import static android.net.ConnectivityManager.FIREWALL_RULE_ALLOW;
+import static android.net.ConnectivityManager.FIREWALL_RULE_DEFAULT;
 import static android.net.ConnectivityManager.FIREWALL_RULE_DENY;
 import static android.net.ConnectivityManager.PROFILE_NETWORK_PREFERENCE_DEFAULT;
 import static android.net.ConnectivityManager.PROFILE_NETWORK_PREFERENCE_ENTERPRISE;
@@ -9547,6 +9555,95 @@
         verify(mBpfNetMaps, never()).removeUidInterfaceRules(any());
     }
 
+    private void doTestSetUidFirewallRule(final int chain, final int defaultRule) {
+        final int uid = 1001;
+        mCm.setUidFirewallRule(chain, uid, FIREWALL_RULE_ALLOW);
+        verify(mBpfNetMaps).setUidRule(chain, uid, FIREWALL_RULE_ALLOW);
+        reset(mBpfNetMaps);
+
+        mCm.setUidFirewallRule(chain, uid, FIREWALL_RULE_DENY);
+        verify(mBpfNetMaps).setUidRule(chain, uid, FIREWALL_RULE_DENY);
+        reset(mBpfNetMaps);
+
+        mCm.setUidFirewallRule(chain, uid, FIREWALL_RULE_DEFAULT);
+        verify(mBpfNetMaps).setUidRule(chain, uid, defaultRule);
+        reset(mBpfNetMaps);
+    }
+
+    @Test @IgnoreUpTo(SC_V2)
+    public void testSetUidFirewallRule() throws Exception {
+        doTestSetUidFirewallRule(FIREWALL_CHAIN_DOZABLE, FIREWALL_RULE_DENY);
+        doTestSetUidFirewallRule(FIREWALL_CHAIN_STANDBY, FIREWALL_RULE_ALLOW);
+        doTestSetUidFirewallRule(FIREWALL_CHAIN_POWERSAVE, FIREWALL_RULE_DENY);
+        doTestSetUidFirewallRule(FIREWALL_CHAIN_RESTRICTED, FIREWALL_RULE_DENY);
+        doTestSetUidFirewallRule(FIREWALL_CHAIN_LOW_POWER_STANDBY, FIREWALL_RULE_DENY);
+        doTestSetUidFirewallRule(FIREWALL_CHAIN_OEM_DENY_1, FIREWALL_RULE_ALLOW);
+        doTestSetUidFirewallRule(FIREWALL_CHAIN_OEM_DENY_2, FIREWALL_RULE_ALLOW);
+    }
+
+    @Test @IgnoreUpTo(SC_V2)
+    public void testSetFirewallChainEnabled() throws Exception {
+        final List<Integer> firewallChains = Arrays.asList(
+                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);
+        for (final int chain: firewallChains) {
+            mCm.setFirewallChainEnabled(chain, true /* enabled */);
+            verify(mBpfNetMaps).setChildChain(chain, true /* enable */);
+            reset(mBpfNetMaps);
+
+            mCm.setFirewallChainEnabled(chain, false /* enabled */);
+            verify(mBpfNetMaps).setChildChain(chain, false /* enable */);
+            reset(mBpfNetMaps);
+        }
+    }
+
+    private void doTestReplaceFirewallChain(final int chain, final String chainName,
+            final boolean allowList) {
+        final int[] uids = new int[] {1001, 1002};
+        mCm.replaceFirewallChain(chain, uids);
+        verify(mBpfNetMaps).replaceUidChain(chainName, allowList, uids);
+        reset(mBpfNetMaps);
+    }
+
+    @Test @IgnoreUpTo(SC_V2)
+    public void testReplaceFirewallChain() {
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_DOZABLE, "fw_dozable", true);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_STANDBY, "fw_standby", false);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_POWERSAVE, "fw_powersave",  true);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_RESTRICTED, "fw_restricted", true);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_LOW_POWER_STANDBY, "fw_low_power_standby", true);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_1, "fw_oem_deny_1", false);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_2, "fw_oem_deny_2", false);
+    }
+
+    @Test @IgnoreUpTo(SC_V2)
+    public void testInvalidFirewallChain() throws Exception {
+        final int uid = 1001;
+        final Class<IllegalArgumentException> expected = IllegalArgumentException.class;
+        assertThrows(expected,
+                () -> mCm.setUidFirewallRule(-1 /* chain */, uid, FIREWALL_RULE_ALLOW));
+        assertThrows(expected,
+                () -> mCm.setUidFirewallRule(100 /* chain */, uid, FIREWALL_RULE_ALLOW));
+        assertThrows(expected, () -> mCm.replaceFirewallChain(-1 /* chain */, new int[]{uid}));
+        assertThrows(expected, () -> mCm.replaceFirewallChain(100 /* chain */, new int[]{uid}));
+    }
+
+    @Test @IgnoreUpTo(SC_V2)
+    public void testInvalidFirewallRule() throws Exception {
+        final Class<IllegalArgumentException> expected = IllegalArgumentException.class;
+        assertThrows(expected,
+                () -> mCm.setUidFirewallRule(FIREWALL_CHAIN_DOZABLE,
+                        1001 /* uid */, -1 /* rule */));
+        assertThrows(expected,
+                () -> mCm.setUidFirewallRule(FIREWALL_CHAIN_DOZABLE,
+                        1001 /* uid */, 100 /* rule */));
+    }
+
     /**
      * Test mutable and requestable network capabilities such as
      * {@link NetworkCapabilities#NET_CAPABILITY_TRUSTED} and