diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index 648ada2..d7c5a06 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -24,6 +24,8 @@
 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_DENY;
 import static android.system.OsConstants.EINVAL;
 import static android.system.OsConstants.ENODEV;
 import static android.system.OsConstants.ENOENT;
@@ -35,7 +37,6 @@
 import android.system.ErrnoException;
 import android.system.Os;
 import android.util.Log;
-import android.util.SparseLongArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -95,22 +96,6 @@
     @VisibleForTesting public static final long OEM_DENY_3_MATCH = (1 << 11);
     // LINT.ThenChange(packages/modules/Connectivity/bpf_progs/bpf_shared.h)
 
-    // TODO: Use Java BpfMap instead of JNI code (TrafficController) for map update.
-    // Currently, BpfNetMaps uses TrafficController for map update and TrafficController
-    // (changeUidOwnerRule and toggleUidOwnerMap) also does conversion from "firewall chain" to
-    // "match". Migrating map update from JNI to Java BpfMap will solve this duplication.
-    private static final SparseLongArray FIREWALL_CHAIN_TO_MATCH = new SparseLongArray();
-    static {
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_DOZABLE, DOZABLE_MATCH);
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_STANDBY, STANDBY_MATCH);
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_POWERSAVE, POWERSAVE_MATCH);
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_RESTRICTED, RESTRICTED_MATCH);
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_LOW_POWER_STANDBY, LOW_POWER_STANDBY_MATCH);
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_OEM_DENY_1, OEM_DENY_1_MATCH);
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_OEM_DENY_2, OEM_DENY_2_MATCH);
-        FIREWALL_CHAIN_TO_MATCH.put(FIREWALL_CHAIN_OEM_DENY_3, OEM_DENY_3_MATCH);
-    }
-
     /**
      * Set configurationMap for test.
      */
@@ -203,11 +188,50 @@
      */
     @VisibleForTesting
     public long getMatchByFirewallChain(final int chain) {
-        final long match = FIREWALL_CHAIN_TO_MATCH.get(chain, NO_MATCH);
-        if (match == NO_MATCH) {
-            throw new ServiceSpecificException(EINVAL, "Invalid firewall chain: " + chain);
+        switch (chain) {
+            case FIREWALL_CHAIN_DOZABLE:
+                return DOZABLE_MATCH;
+            case FIREWALL_CHAIN_STANDBY:
+                return STANDBY_MATCH;
+            case FIREWALL_CHAIN_POWERSAVE:
+                return POWERSAVE_MATCH;
+            case FIREWALL_CHAIN_RESTRICTED:
+                return RESTRICTED_MATCH;
+            case FIREWALL_CHAIN_LOW_POWER_STANDBY:
+                return LOW_POWER_STANDBY_MATCH;
+            case FIREWALL_CHAIN_OEM_DENY_1:
+                return OEM_DENY_1_MATCH;
+            case FIREWALL_CHAIN_OEM_DENY_2:
+                return OEM_DENY_2_MATCH;
+            case FIREWALL_CHAIN_OEM_DENY_3:
+                return OEM_DENY_3_MATCH;
+            default:
+                throw new ServiceSpecificException(EINVAL, "Invalid firewall chain: " + chain);
         }
-        return match;
+    }
+
+    /**
+     * 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
+     */
+    @VisibleForTesting
+    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) {
@@ -412,9 +436,17 @@
      *                                  cause of the failure.
      */
     public void setUidRule(final int childChain, final int uid, final int firewallRule) {
-        synchronized (sUidOwnerMap) {
-            final int err = native_setUidRule(childChain, uid, firewallRule);
-            maybeThrow(err, "Unable to set uid rule");
+        throwIfPreT("setUidRule is not available on pre-T devices");
+
+        final long match = getMatchByFirewallChain(childChain);
+        final boolean isAllowList = isFirewallAllowList(childChain);
+        final boolean add = (firewallRule == FIREWALL_RULE_ALLOW && isAllowList)
+                || (firewallRule == FIREWALL_RULE_DENY && !isAllowList);
+
+        if (add) {
+            addRule(uid, match, "setUidRule");
+        } else {
+            removeRule(uid, match, "setUidRule");
         }
     }
 
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index 98d0905..0718952 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -24,6 +24,8 @@
 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_DENY;
 import static android.net.INetd.PERMISSION_INTERNET;
 
 import static com.android.server.BpfNetMaps.DOZABLE_MATCH;
@@ -124,12 +126,16 @@
         verify(mNetd).trafficSetNetPermForUids(PERMISSION_INTERNET, TEST_UIDS);
     }
 
-    private void doTestIsChainEnabled(final List<Integer> enableChains) throws Exception {
+    private long getMatch(final List<Integer> chains) {
         long match = 0;
-        for (final int chain: enableChains) {
+        for (final int chain: chains) {
             match |= mBpfNetMaps.getMatchByFirewallChain(chain);
         }
-        mConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, new U32(match));
+        return match;
+    }
+
+    private void doTestIsChainEnabled(final List<Integer> enableChains) throws Exception {
+        mConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, new U32(getMatch(enableChains)));
 
         for (final int chain: FIREWALL_CHAINS) {
             final String testCase = "EnabledChains: " + enableChains + " CheckedChain: " + chain;
@@ -557,4 +563,90 @@
         doTestRemoveUidInterfaceRules(NO_IIF, DOZABLE_MATCH,
                 NO_IIF, DOZABLE_MATCH | POWERSAVE_MATCH | RESTRICTED_MATCH);
     }
+
+    private void doTestSetUidRule(final List<Integer> testChains) throws Exception {
+        mUidOwnerMap.updateEntry(new U32(TEST_UID), new UidOwnerValue(TEST_IF_INDEX, IIF_MATCH));
+
+        for (final int chain: testChains) {
+            final int ruleToAddMatch = mBpfNetMaps.isFirewallAllowList(chain)
+                    ? FIREWALL_RULE_ALLOW : FIREWALL_RULE_DENY;
+            mBpfNetMaps.setUidRule(chain, TEST_UID, ruleToAddMatch);
+        }
+
+        checkUidOwnerValue(TEST_UID, TEST_IF_INDEX, IIF_MATCH | getMatch(testChains));
+
+        for (final int chain: testChains) {
+            final int ruleToRemoveMatch = mBpfNetMaps.isFirewallAllowList(chain)
+                    ? FIREWALL_RULE_DENY : FIREWALL_RULE_ALLOW;
+            mBpfNetMaps.setUidRule(chain, TEST_UID, ruleToRemoveMatch);
+        }
+
+        checkUidOwnerValue(TEST_UID, TEST_IF_INDEX, IIF_MATCH);
+    }
+
+    private void doTestSetUidRule(final int testChain) throws Exception {
+        doTestSetUidRule(List.of(testChain));
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testSetUidRule() throws Exception {
+        doTestSetUidRule(FIREWALL_CHAIN_DOZABLE);
+        doTestSetUidRule(FIREWALL_CHAIN_STANDBY);
+        doTestSetUidRule(FIREWALL_CHAIN_POWERSAVE);
+        doTestSetUidRule(FIREWALL_CHAIN_RESTRICTED);
+        doTestSetUidRule(FIREWALL_CHAIN_LOW_POWER_STANDBY);
+        doTestSetUidRule(FIREWALL_CHAIN_OEM_DENY_1);
+        doTestSetUidRule(FIREWALL_CHAIN_OEM_DENY_2);
+        doTestSetUidRule(FIREWALL_CHAIN_OEM_DENY_3);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testSetUidRuleMultipleChain() throws Exception {
+        doTestSetUidRule(List.of(
+                FIREWALL_CHAIN_DOZABLE,
+                FIREWALL_CHAIN_STANDBY));
+        doTestSetUidRule(List.of(
+                FIREWALL_CHAIN_DOZABLE,
+                FIREWALL_CHAIN_STANDBY,
+                FIREWALL_CHAIN_POWERSAVE,
+                FIREWALL_CHAIN_RESTRICTED));
+        doTestSetUidRule(FIREWALL_CHAINS);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testSetUidRuleRemoveRuleFromUidWithNoRule() {
+        final Class<ServiceSpecificException> expected = ServiceSpecificException.class;
+        assertThrows(expected,
+                () -> mBpfNetMaps.setUidRule(FIREWALL_CHAIN_DOZABLE, TEST_UID, FIREWALL_RULE_DENY));
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testSetUidRuleInvalidChain() {
+        final Class<ServiceSpecificException> expected = ServiceSpecificException.class;
+        assertThrows(expected,
+                () -> mBpfNetMaps.setUidRule(-1 /* childChain */, TEST_UID, FIREWALL_RULE_ALLOW));
+        assertThrows(expected,
+                () -> mBpfNetMaps.setUidRule(1000 /* childChain */, TEST_UID, FIREWALL_RULE_ALLOW));
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testSetUidRuleInvalidRule() {
+        final Class<ServiceSpecificException> expected = ServiceSpecificException.class;
+        assertThrows(expected, () ->
+                mBpfNetMaps.setUidRule(FIREWALL_CHAIN_DOZABLE, TEST_UID, -1 /* firewallRule */));
+        assertThrows(expected, () ->
+                mBpfNetMaps.setUidRule(FIREWALL_CHAIN_DOZABLE, TEST_UID, 1000 /* firewallRule */));
+    }
+
+    @Test
+    @IgnoreAfter(Build.VERSION_CODES.S_V2)
+    public void testSetUidRuleBeforeT() {
+        assertThrows(UnsupportedOperationException.class, () ->
+                mBpfNetMaps.setUidRule(FIREWALL_CHAIN_DOZABLE, TEST_UID, FIREWALL_RULE_ALLOW));
+    }
 }