Merge changes Ife5e91de,I46c5dfeb

* changes:
  Change the sanitizeCaps signature in preparation for moving
  Clarify sanitization of caps from NetworkAgent
diff --git a/Tethering/common/TetheringLib/src/android/net/ITetheringEventCallback.aidl b/Tethering/common/TetheringLib/src/android/net/ITetheringEventCallback.aidl
index b4e3ba4..836761f 100644
--- a/Tethering/common/TetheringLib/src/android/net/ITetheringEventCallback.aidl
+++ b/Tethering/common/TetheringLib/src/android/net/ITetheringEventCallback.aidl
@@ -36,4 +36,5 @@
     void onTetherStatesChanged(in TetherStatesParcel states);
     void onTetherClientsChanged(in List<TetheredClient> clients);
     void onOffloadStatusChanged(int status);
+    void onSupportedTetheringTypes(long supportedBitmap);
 }
diff --git a/Tethering/common/TetheringLib/src/android/net/TetheringCallbackStartedParcel.aidl b/Tethering/common/TetheringLib/src/android/net/TetheringCallbackStartedParcel.aidl
index 253eacb..f33f846 100644
--- a/Tethering/common/TetheringLib/src/android/net/TetheringCallbackStartedParcel.aidl
+++ b/Tethering/common/TetheringLib/src/android/net/TetheringCallbackStartedParcel.aidl
@@ -26,7 +26,7 @@
  * @hide
  */
 parcelable TetheringCallbackStartedParcel {
-    boolean tetheringSupported;
+    long supportedTypes;
     Network upstreamNetwork;
     TetheringConfigurationParcel config;
     TetherStatesParcel states;
diff --git a/Tethering/common/TetheringLib/src/android/net/TetheringManager.java b/Tethering/common/TetheringLib/src/android/net/TetheringManager.java
index 6f9b33e..b3f0cf2 100644
--- a/Tethering/common/TetheringLib/src/android/net/TetheringManager.java
+++ b/Tethering/common/TetheringLib/src/android/net/TetheringManager.java
@@ -183,6 +183,12 @@
      */
     public static final int TETHERING_WIGIG = 6;
 
+    /**
+     * The int value of last tethering type.
+     * @hide
+     */
+    public static final int MAX_TETHERING_TYPE = TETHERING_WIGIG;
+
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
     @IntDef(value = {
@@ -520,6 +526,9 @@
         }
 
         @Override
+        public void onSupportedTetheringTypes(long supportedBitmap) { }
+
+        @Override
         public void onUpstreamChanged(Network network) { }
 
         @Override
@@ -1033,15 +1042,29 @@
         /**
          * Called when tethering supported status changed.
          *
+         * <p>This callback will be called immediately after the callback is
+         * registered, and never be called if there is changes afterward.
+         *
+         * <p>Tethering may be disabled via system properties, device configuration, or device
+         * policy restrictions.
+         *
+         * @param supported whether any tethering type is supported.
+         */
+        default void onTetheringSupported(boolean supported) {}
+
+        /**
+         * Called when tethering supported status changed.
+         *
          * <p>This will be called immediately after the callback is registered, and may be called
          * multiple times later upon changes.
          *
          * <p>Tethering may be disabled via system properties, device configuration, or device
          * policy restrictions.
          *
-         * @param supported The new supported status
+         * @param supportedTypes a set of @TetheringType which is supported.
+         * @hide
          */
-        default void onTetheringSupported(boolean supported) {}
+        default void onSupportedTetheringTypes(@NonNull Set<Integer> supportedTypes) {}
 
         /**
          * Called when tethering upstream changed.
@@ -1339,7 +1362,8 @@
                 @Override
                 public void onCallbackStarted(TetheringCallbackStartedParcel parcel) {
                     executor.execute(() -> {
-                        callback.onTetheringSupported(parcel.tetheringSupported);
+                        callback.onSupportedTetheringTypes(unpackBits(parcel.supportedTypes));
+                        callback.onTetheringSupported(parcel.supportedTypes != 0);
                         callback.onUpstreamChanged(parcel.upstreamNetwork);
                         sendErrorCallbacks(parcel.states);
                         sendRegexpsChanged(parcel.config);
@@ -1358,6 +1382,13 @@
                     });
                 }
 
+                @Override
+                public void onSupportedTetheringTypes(long supportedBitmap) {
+                    executor.execute(() -> {
+                        callback.onSupportedTetheringTypes(unpackBits(supportedBitmap));
+                    });
+                }
+
                 private void sendRegexpsChanged(TetheringConfigurationParcel parcel) {
                     callback.onTetherableInterfaceRegexpsChanged(new TetheringInterfaceRegexps(
                             parcel.tetherableBluetoothRegexs,
@@ -1396,6 +1427,23 @@
     }
 
     /**
+     * Unpack bitmap to a set of bit position intergers.
+     * @hide
+     */
+    public static ArraySet<Integer> unpackBits(long val) {
+        final ArraySet<Integer> result = new ArraySet<>(Long.bitCount(val));
+        int bitPos = 0;
+        while (val != 0) {
+            if ((val & 1) == 1) result.add(bitPos);
+
+            val = val >>> 1;
+            bitPos++;
+        }
+
+        return result;
+    }
+
+    /**
      * Remove tethering event callback previously registered with
      * {@link #registerTetheringEventCallback}.
      *
diff --git a/Tethering/src/com/android/networkstack/tethering/Tethering.java b/Tethering/src/com/android/networkstack/tethering/Tethering.java
index 89ed620..0c59b61 100644
--- a/Tethering/src/com/android/networkstack/tethering/Tethering.java
+++ b/Tethering/src/com/android/networkstack/tethering/Tethering.java
@@ -281,6 +281,11 @@
     private BluetoothPan mBluetoothPan;
     private PanServiceListener mBluetoothPanListener;
     private ArrayList<Pair<Boolean, IIntResultListener>> mPendingPanRequests;
+    // AIDL doesn't support Set<Integer>. Maintain a int bitmap here. When the bitmap is passed to
+    // TetheringManager, TetheringManager would convert it to a set of Integer types.
+    // mSupportedTypeBitmap should always be updated inside tethering internal thread but it may be
+    // read from binder thread which called TetheringService directly.
+    private volatile long mSupportedTypeBitmap;
 
     public Tethering(TetheringDependencies deps) {
         mLog.mark("Tethering.constructed");
@@ -512,6 +517,8 @@
         mUpstreamNetworkMonitor.setUpstreamConfig(mConfig.chooseUpstreamAutomatically,
                 mConfig.isDunRequired);
         reportConfigurationChanged(mConfig.toStableParcelable());
+
+        updateSupportedDownstreams(mConfig);
     }
 
     private void maybeDunSettingChanged() {
@@ -1550,26 +1557,6 @@
         return mConfig;
     }
 
-    boolean hasAnySupportedDownstream() {
-        if ((mConfig.tetherableUsbRegexs.length != 0)
-                || (mConfig.tetherableWifiRegexs.length != 0)
-                || (mConfig.tetherableBluetoothRegexs.length != 0)) {
-            return true;
-        }
-
-        // Before T, isTetheringSupported would return true if wifi, usb and bluetooth tethering are
-        // disabled (whole tethering settings would be hidden). This means tethering would also not
-        // support wifi p2p, ethernet tethering and mirrorlink. This is wrong but probably there are
-        // some devices in the field rely on this to disable tethering entirely.
-        if (!SdkLevel.isAtLeastT()) return false;
-
-        return (mConfig.tetherableWifiP2pRegexs.length != 0)
-                || (mConfig.tetherableNcmRegexs.length != 0)
-                || isEthernetSupported();
-    }
-
-    // TODO: using EtherentManager new API to check whether ethernet is supported when the API is
-    // ready to use.
     private boolean isEthernetSupported() {
         return mContext.getSystemService(Context.ETHERNET_SERVICE) != null;
     }
@@ -2359,7 +2346,7 @@
         mHandler.post(() -> {
             mTetheringEventCallbacks.register(callback, new CallbackCookie(hasListPermission));
             final TetheringCallbackStartedParcel parcel = new TetheringCallbackStartedParcel();
-            parcel.tetheringSupported = isTetheringSupported();
+            parcel.supportedTypes = mSupportedTypeBitmap;
             parcel.upstreamNetwork = mTetherUpstream;
             parcel.config = mConfig.toStableParcelable();
             parcel.states =
@@ -2398,6 +2385,22 @@
         });
     }
 
+    private void reportTetheringSupportedChange(final long supportedBitmap) {
+        final int length = mTetheringEventCallbacks.beginBroadcast();
+        try {
+            for (int i = 0; i < length; i++) {
+                try {
+                    mTetheringEventCallbacks.getBroadcastItem(i).onSupportedTetheringTypes(
+                            supportedBitmap);
+                } catch (RemoteException e) {
+                    // Not really very much to do here.
+                }
+            }
+        } finally {
+            mTetheringEventCallbacks.finishBroadcast();
+        }
+    }
+
     private void reportUpstreamChanged(UpstreamNetworkState ns) {
         final int length = mTetheringEventCallbacks.beginBroadcast();
         final Network network = (ns != null) ? ns.network : null;
@@ -2482,18 +2485,56 @@
         }
     }
 
+    private void updateSupportedDownstreams(final TetheringConfiguration config) {
+        final long preSupportedBitmap = mSupportedTypeBitmap;
+
+        if (!isTetheringAllowed() || mEntitlementMgr.isProvisioningNeededButUnavailable()) {
+            mSupportedTypeBitmap = 0;
+        } else {
+            mSupportedTypeBitmap = makeSupportedDownstreams(config);
+        }
+
+        if (preSupportedBitmap != mSupportedTypeBitmap) {
+            reportTetheringSupportedChange(mSupportedTypeBitmap);
+        }
+    }
+
+    private long makeSupportedDownstreams(final TetheringConfiguration config) {
+        long types = 0;
+        if (config.tetherableUsbRegexs.length != 0) types |= (1 << TETHERING_USB);
+
+        if (config.tetherableWifiRegexs.length != 0) types |= (1 << TETHERING_WIFI);
+
+        if (config.tetherableBluetoothRegexs.length != 0) types |= (1 << TETHERING_BLUETOOTH);
+
+        // Before T, isTetheringSupported would return true if wifi, usb and bluetooth tethering are
+        // disabled (whole tethering settings would be hidden). This means tethering would also not
+        // support wifi p2p, ethernet tethering and mirrorlink. This is wrong but probably there are
+        // some devices in the field rely on this to disable tethering entirely.
+        if (!SdkLevel.isAtLeastT() && types == 0) return types;
+
+        if (config.tetherableNcmRegexs.length != 0) types |= (1 << TETHERING_NCM);
+
+        if (config.tetherableWifiP2pRegexs.length != 0) types |= (1 << TETHERING_WIFI_P2P);
+
+        if (isEthernetSupported()) types |= (1 << TETHERING_ETHERNET);
+
+        return types;
+    }
+
     // if ro.tether.denied = true we default to no tethering
     // gservices could set the secure setting to 1 though to enable it on a build where it
     // had previously been turned off.
-    boolean isTetheringSupported() {
+    boolean isTetheringAllowed() {
         final int defaultVal = mDeps.isTetheringDenied() ? 0 : 1;
         final boolean tetherSupported = Settings.Global.getInt(mContext.getContentResolver(),
                 Settings.Global.TETHER_SUPPORTED, defaultVal) != 0;
-        final boolean tetherEnabledInSettings = tetherSupported
+        return tetherSupported
                 && !mUserManager.hasUserRestriction(UserManager.DISALLOW_CONFIG_TETHERING);
+    }
 
-        return tetherEnabledInSettings && hasAnySupportedDownstream()
-                && !mEntitlementMgr.isProvisioningNeededButUnavailable();
+    boolean isTetheringSupported() {
+        return mSupportedTypeBitmap > 0;
     }
 
     private void dumpBpf(IndentingPrintWriter pw) {
diff --git a/Tethering/src/com/android/networkstack/tethering/TetheringService.java b/Tethering/src/com/android/networkstack/tethering/TetheringService.java
index f147e10..96ddfa0 100644
--- a/Tethering/src/com/android/networkstack/tethering/TetheringService.java
+++ b/Tethering/src/com/android/networkstack/tethering/TetheringService.java
@@ -237,7 +237,7 @@
                     listener.onResult(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION);
                     return true;
                 }
-                if (!mTethering.isTetheringSupported()) {
+                if (!mTethering.isTetheringSupported() || !mTethering.isTetheringAllowed()) {
                     listener.onResult(TETHER_ERROR_UNSUPPORTED);
                     return true;
                 }
@@ -255,7 +255,7 @@
                 receiver.send(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION, null);
                 return true;
             }
-            if (!mTethering.isTetheringSupported()) {
+            if (!mTethering.isTetheringSupported() || !mTethering.isTetheringAllowed()) {
                 receiver.send(TETHER_ERROR_UNSUPPORTED, null);
                 return true;
             }
diff --git a/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java b/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
index 86dca1c..ef4f052 100644
--- a/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
+++ b/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
@@ -192,12 +192,13 @@
         mUiAutomation.adoptShellPermissionIdentity(
                 MANAGE_TEST_NETWORKS, NETWORK_SETTINGS, TETHER_PRIVILEGED, ACCESS_NETWORK_STATE,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS, DUMP);
-        mRunTests = mTm.isTetheringSupported() && mEm != null;
-        assumeTrue(mRunTests);
-
         mHandlerThread = new HandlerThread(getClass().getSimpleName());
         mHandlerThread.start();
         mHandler = new Handler(mHandlerThread.getLooper());
+
+        mRunTests = isEthernetTetheringSupported();
+        assumeTrue(mRunTests);
+
         mTetheredInterfaceRequester = new TetheredInterfaceRequester(mHandler, mEm);
     }
 
@@ -225,7 +226,6 @@
             mHandler.post(() -> reader.stop());
             mDownstreamReader = null;
         }
-        mHandlerThread.quitSafely();
         mTetheredInterfaceRequester.release();
         mEm.setIncludeTestInterfaces(false);
         maybeDeleteTestInterface();
@@ -236,6 +236,7 @@
         try {
             if (mRunTests) cleanUp();
         } finally {
+            mHandlerThread.quitSafely();
             mUiAutomation.dropShellPermissionIdentity();
         }
     }
@@ -410,6 +411,23 @@
         // client, which is not possible in this test.
     }
 
+    private boolean isEthernetTetheringSupported() throws Exception {
+        final CompletableFuture<Boolean> future = new CompletableFuture<>();
+        final TetheringEventCallback callback = new TetheringEventCallback() {
+            @Override
+            public void onSupportedTetheringTypes(Set<Integer> supportedTypes) {
+                future.complete(supportedTypes.contains(TETHERING_ETHERNET));
+            }
+        };
+
+        try {
+            mTm.registerTetheringEventCallback(mHandler::post, callback);
+            return future.get(TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        } finally {
+            mTm.unregisterTetheringEventCallback(callback);
+        }
+    }
+
     private static final class MyTetheringEventCallback implements TetheringEventCallback {
         private final TetheringManager mTm;
         private final CountDownLatch mTetheringStartedLatch = new CountDownLatch(1);
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java
index 9db8f16..e114cb5 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java
@@ -24,6 +24,7 @@
 import static android.net.TetheringManager.TETHER_ERROR_NO_ACCESS_TETHERING_PERMISSION;
 import static android.net.TetheringManager.TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION;
 import static android.net.TetheringManager.TETHER_ERROR_NO_ERROR;
+import static android.net.TetheringManager.TETHER_ERROR_UNSUPPORTED;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -139,23 +140,27 @@
     }
 
     private void runAsNoPermission(final TestTetheringCall test) throws Exception {
-        runTetheringCall(test, new String[0]);
+        runTetheringCall(test, true /* isTetheringAllowed */, new String[0]);
     }
 
     private void runAsTetherPrivileged(final TestTetheringCall test) throws Exception {
-        runTetheringCall(test, TETHER_PRIVILEGED);
+        runTetheringCall(test, true /* isTetheringAllowed */, TETHER_PRIVILEGED);
     }
 
     private void runAsAccessNetworkState(final TestTetheringCall test) throws Exception {
-        runTetheringCall(test, ACCESS_NETWORK_STATE);
+        runTetheringCall(test, true /* isTetheringAllowed */, ACCESS_NETWORK_STATE);
     }
 
     private void runAsWriteSettings(final TestTetheringCall test) throws Exception {
-        runTetheringCall(test, WRITE_SETTINGS);
+        runTetheringCall(test, true /* isTetheringAllowed */, WRITE_SETTINGS);
     }
 
-    private void runTetheringCall(final TestTetheringCall test, String... permissions)
-            throws Exception {
+    private void runAsTetheringDisallowed(final TestTetheringCall test) throws Exception {
+        runTetheringCall(test, false /* isTetheringAllowed */, TETHER_PRIVILEGED);
+    }
+
+    private void runTetheringCall(final TestTetheringCall test, boolean isTetheringAllowed,
+            String... permissions) throws Exception {
         // Allow the test to run even if ACCESS_NETWORK_STATE was granted at the APK level
         if (!CollectionUtils.contains(permissions, ACCESS_NETWORK_STATE)) {
             mMockConnector.setPermission(ACCESS_NETWORK_STATE, PERMISSION_DENIED);
@@ -164,6 +169,7 @@
         if (permissions.length > 0) mUiAutomation.adoptShellPermissionIdentity(permissions);
         try {
             when(mTethering.isTetheringSupported()).thenReturn(true);
+            when(mTethering.isTetheringAllowed()).thenReturn(isTetheringAllowed);
             test.runTetheringCall(new TestTetheringResult());
         } finally {
             mUiAutomation.dropShellPermissionIdentity();
@@ -180,6 +186,7 @@
     private void runTether(final TestTetheringResult result) throws Exception {
         mTetheringConnector.tether(TEST_IFACE_NAME, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG, result);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         verify(mTethering).tether(TEST_IFACE_NAME, IpServer.STATE_TETHERED, result);
     }
 
@@ -203,12 +210,22 @@
             verify(mTethering).isTetherProvisioningRequired();
             verifyNoMoreInteractionsForTethering();
         });
+
+        runAsTetheringDisallowed((result) -> {
+            mTetheringConnector.tether(TEST_IFACE_NAME, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG,
+                    result);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            result.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runUnTether(final TestTetheringResult result) throws Exception {
         mTetheringConnector.untether(TEST_IFACE_NAME, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG,
                 result);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         verify(mTethering).untether(eq(TEST_IFACE_NAME), eq(result));
     }
 
@@ -232,6 +249,15 @@
             verify(mTethering).isTetherProvisioningRequired();
             verifyNoMoreInteractionsForTethering();
         });
+
+        runAsTetheringDisallowed((result) -> {
+            mTetheringConnector.untether(TEST_IFACE_NAME, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG,
+                    result);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            result.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runSetUsbTethering(final TestTetheringResult result) throws Exception {
@@ -243,6 +269,7 @@
         mTetheringConnector.setUsbTethering(true /* enable */, TEST_CALLER_PKG,
                 TEST_ATTRIBUTION_TAG, result);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         verify(mTethering).setUsbTethering(eq(true) /* enable */, any(IIntResultListener.class));
         result.assertResult(TETHER_ERROR_NO_ERROR);
     }
@@ -268,6 +295,14 @@
             verifyNoMoreInteractionsForTethering();
         });
 
+        runAsTetheringDisallowed((result) -> {
+            mTetheringConnector.setUsbTethering(true /* enable */, TEST_CALLER_PKG,
+                    TEST_ATTRIBUTION_TAG, result);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            result.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runStartTethering(final TestTetheringResult result,
@@ -275,6 +310,7 @@
         mTetheringConnector.startTethering(request, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG,
                 result);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         verify(mTethering).startTethering(eq(request), eq(TEST_CALLER_PKG), eq(result));
     }
 
@@ -301,6 +337,15 @@
             verify(mTethering).isTetherProvisioningRequired();
             verifyNoMoreInteractionsForTethering();
         });
+
+        runAsTetheringDisallowed((result) -> {
+            mTetheringConnector.startTethering(request, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG,
+                    result);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            result.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runStartTetheringAndVerifyNoPermission(final TestTetheringResult result)
@@ -337,6 +382,7 @@
         mTetheringConnector.stopTethering(TETHERING_WIFI, TEST_CALLER_PKG,
                 TEST_ATTRIBUTION_TAG, result);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         verify(mTethering).stopTethering(TETHERING_WIFI);
         result.assertResult(TETHER_ERROR_NO_ERROR);
     }
@@ -361,6 +407,15 @@
             verify(mTethering).isTetherProvisioningRequired();
             verifyNoMoreInteractionsForTethering();
         });
+
+        runAsTetheringDisallowed((result) -> {
+            mTetheringConnector.stopTethering(TETHERING_WIFI, TEST_CALLER_PKG,
+                    TEST_ATTRIBUTION_TAG, result);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            result.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runRequestLatestTetheringEntitlementResult() throws Exception {
@@ -368,6 +423,7 @@
         mTetheringConnector.requestLatestTetheringEntitlementResult(TETHERING_WIFI, result,
                 true /* showEntitlementUi */, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         verify(mTethering).requestLatestTetheringEntitlementResult(eq(TETHERING_WIFI),
                 eq(result), eq(true) /* showEntitlementUi */);
     }
@@ -392,6 +448,16 @@
             verify(mTethering).isTetherProvisioningRequired();
             verifyNoMoreInteractionsForTethering();
         });
+
+        runAsTetheringDisallowed((none) -> {
+            final MyResultReceiver receiver = new MyResultReceiver(null);
+            mTetheringConnector.requestLatestTetheringEntitlementResult(TETHERING_WIFI, receiver,
+                    true /* showEntitlementUi */, TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            receiver.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runRegisterTetheringEventCallback() throws Exception {
@@ -419,6 +485,12 @@
             runRegisterTetheringEventCallback();
             verifyNoMoreInteractionsForTethering();
         });
+
+        // should still be able to register callback even tethering is restricted.
+        runAsTetheringDisallowed((result) -> {
+            runRegisterTetheringEventCallback();
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runUnregisterTetheringEventCallback() throws Exception {
@@ -446,11 +518,19 @@
             runUnregisterTetheringEventCallback();
             verifyNoMoreInteractionsForTethering();
         });
+
+        // should still be able to unregister callback even tethering is restricted.
+        runAsTetheringDisallowed((result) -> {
+            runUnregisterTetheringEventCallback();
+            verifyNoMoreInteractionsForTethering();
+        });
+
     }
 
     private void runStopAllTethering(final TestTetheringResult result) throws Exception {
         mTetheringConnector.stopAllTethering(TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG, result);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         verify(mTethering).untetherAll();
         result.assertResult(TETHER_ERROR_NO_ERROR);
     }
@@ -474,11 +554,20 @@
             verify(mTethering).isTetherProvisioningRequired();
             verifyNoMoreInteractionsForTethering();
         });
+
+        runAsTetheringDisallowed((result) -> {
+            mTetheringConnector.stopAllTethering(TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG, result);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            result.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private void runIsTetheringSupported(final TestTetheringResult result) throws Exception {
         mTetheringConnector.isTetheringSupported(TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG, result);
         verify(mTethering).isTetheringSupported();
+        verify(mTethering).isTetheringAllowed();
         result.assertResult(TETHER_ERROR_NO_ERROR);
     }
 
@@ -502,6 +591,15 @@
             verify(mTethering).isTetherProvisioningRequired();
             verifyNoMoreInteractionsForTethering();
         });
+
+        runAsTetheringDisallowed((result) -> {
+            mTetheringConnector.isTetheringSupported(TEST_CALLER_PKG, TEST_ATTRIBUTION_TAG,
+                    result);
+            verify(mTethering).isTetheringSupported();
+            verify(mTethering).isTetheringAllowed();
+            result.assertResult(TETHER_ERROR_UNSUPPORTED);
+            verifyNoMoreInteractionsForTethering();
+        });
     }
 
     private class ConnectorSupplier<T> implements Supplier<T> {
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
index b402bc3..8de7836 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
@@ -144,6 +144,7 @@
 import android.net.TetheringCallbackStartedParcel;
 import android.net.TetheringConfigurationParcel;
 import android.net.TetheringInterface;
+import android.net.TetheringManager;
 import android.net.TetheringRequestParcel;
 import android.net.dhcp.DhcpLeaseParcelable;
 import android.net.dhcp.DhcpServerCallbacks;
@@ -175,6 +176,7 @@
 import android.telephony.PhoneStateListener;
 import android.telephony.TelephonyManager;
 import android.test.mock.MockContentResolver;
+import android.util.ArraySet;
 
 import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
@@ -217,6 +219,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.Vector;
 
 @RunWith(AndroidJUnit4.class)
@@ -1729,6 +1732,7 @@
         private final ArrayList<TetherStatesParcel> mTetherStates = new ArrayList<>();
         private final ArrayList<Integer> mOffloadStatus = new ArrayList<>();
         private final ArrayList<List<TetheredClient>> mTetheredClients = new ArrayList<>();
+        private final ArrayList<Long> mSupportedBitmaps = new ArrayList<>();
 
         // This function will remove the recorded callbacks, so it must be called once for
         // each callback. If this is called after multiple callback, the order matters.
@@ -1781,6 +1785,10 @@
             assertTrue(leases.containsAll(result));
         }
 
+        public void expectSupportedTetheringTypes(Set<Integer> expectedTypes) {
+            assertEquals(expectedTypes, TetheringManager.unpackBits(mSupportedBitmaps.remove(0)));
+        }
+
         @Override
         public void onUpstreamChanged(Network network) {
             mActualUpstreams.add(network);
@@ -1813,11 +1821,17 @@
             mTetherStates.add(parcel.states);
             mOffloadStatus.add(parcel.offloadStatus);
             mTetheredClients.add(parcel.tetheredClients);
+            mSupportedBitmaps.add(parcel.supportedTypes);
         }
 
         @Override
         public void onCallbackStopped(int errorCode) { }
 
+        @Override
+        public void onSupportedTetheringTypes(long supportedBitmap) {
+            mSupportedBitmaps.add(supportedBitmap);
+        }
+
         public void assertNoUpstreamChangeCallback() {
             assertTrue(mActualUpstreams.isEmpty());
         }
@@ -2945,53 +2959,81 @@
         runStopUSBTethering();
     }
 
+    public static ArraySet<Integer> getAllSupportedTetheringTypes() {
+        return new ArraySet<>(new Integer[] { TETHERING_USB, TETHERING_NCM, TETHERING_WIFI,
+                TETHERING_WIFI_P2P, TETHERING_BLUETOOTH, TETHERING_ETHERNET });
+    }
+
     @Test
     public void testTetheringSupported() throws Exception {
+        final ArraySet<Integer> expectedTypes = getAllSupportedTetheringTypes();
+        // Check tethering is supported after initialization.
         setTetheringSupported(true /* supported */);
-        updateConfigAndVerifySupported(true /* supported */);
+        TestTetheringEventCallback callback = new TestTetheringEventCallback();
+        mTethering.registerTetheringEventCallback(callback);
+        mLooper.dispatchAll();
+        updateConfigAndVerifySupported(callback, expectedTypes);
 
         // Could disable tethering supported by settings.
         Settings.Global.putInt(mContentResolver, Settings.Global.TETHER_SUPPORTED, 0);
-        updateConfigAndVerifySupported(false /* supported */);
+        updateConfigAndVerifySupported(callback, new ArraySet<>());
 
         // Could disable tethering supported by user restriction.
         setTetheringSupported(true /* supported */);
+        updateConfigAndVerifySupported(callback, expectedTypes);
         when(mUserManager.hasUserRestriction(
                 UserManager.DISALLOW_CONFIG_TETHERING)).thenReturn(true);
-        updateConfigAndVerifySupported(false /* supported */);
+        updateConfigAndVerifySupported(callback, new ArraySet<>());
 
         // Tethering is supported if it has any supported downstream.
         setTetheringSupported(true /* supported */);
+        updateConfigAndVerifySupported(callback, expectedTypes);
+        // Usb tethering is not supported:
+        expectedTypes.remove(TETHERING_USB);
         when(mResources.getStringArray(R.array.config_tether_usb_regexs))
                 .thenReturn(new String[0]);
-        updateConfigAndVerifySupported(true /* supported */);
+        updateConfigAndVerifySupported(callback, expectedTypes);
+        // Wifi tethering is not supported:
+        expectedTypes.remove(TETHERING_WIFI);
         when(mResources.getStringArray(R.array.config_tether_wifi_regexs))
                 .thenReturn(new String[0]);
-        updateConfigAndVerifySupported(true /* supported */);
-
+        updateConfigAndVerifySupported(callback, expectedTypes);
+        // Bluetooth tethering is not supported:
+        expectedTypes.remove(TETHERING_BLUETOOTH);
+        when(mResources.getStringArray(R.array.config_tether_bluetooth_regexs))
+                .thenReturn(new String[0]);
 
         if (isAtLeastT()) {
-            when(mResources.getStringArray(R.array.config_tether_bluetooth_regexs))
-                    .thenReturn(new String[0]);
-            updateConfigAndVerifySupported(true /* supported */);
+            updateConfigAndVerifySupported(callback, expectedTypes);
+
+            // P2p tethering is not supported:
+            expectedTypes.remove(TETHERING_WIFI_P2P);
             when(mResources.getStringArray(R.array.config_tether_wifi_p2p_regexs))
                     .thenReturn(new String[0]);
-            updateConfigAndVerifySupported(true /* supported */);
+            updateConfigAndVerifySupported(callback, expectedTypes);
+            // Ncm tethering is not supported:
+            expectedTypes.remove(TETHERING_NCM);
             when(mResources.getStringArray(R.array.config_tether_ncm_regexs))
                     .thenReturn(new String[0]);
-            updateConfigAndVerifySupported(true /* supported */);
+            updateConfigAndVerifySupported(callback, expectedTypes);
+            // Ethernet tethering (last supported type) is not supported:
+            expectedTypes.remove(TETHERING_ETHERNET);
             mForceEthernetServiceUnavailable = true;
-            updateConfigAndVerifySupported(false /* supported */);
+            updateConfigAndVerifySupported(callback, new ArraySet<>());
+
         } else {
-            when(mResources.getStringArray(R.array.config_tether_bluetooth_regexs))
-                    .thenReturn(new String[0]);
-            updateConfigAndVerifySupported(false /* supported */);
+            // If wifi, usb and bluetooth are all not supported, all the types are not supported.
+            expectedTypes.clear();
+            updateConfigAndVerifySupported(callback, expectedTypes);
         }
     }
 
-    private void updateConfigAndVerifySupported(boolean supported) {
+    private void updateConfigAndVerifySupported(final TestTetheringEventCallback callback,
+            final ArraySet<Integer> expectedTypes) {
         sendConfigurationChanged();
-        assertEquals(supported, mTethering.isTetheringSupported());
+
+        assertEquals(expectedTypes.size() > 0, mTethering.isTetheringSupported());
+        callback.expectSupportedTetheringTypes(expectedTypes);
     }
     // TODO: Test that a request for hotspot mode doesn't interfere with an
     // already operating tethering mode interface.
diff --git a/bpf_progs/dscpPolicy.c b/bpf_progs/dscpPolicy.c
index d3e7b38..f308931 100644
--- a/bpf_progs/dscpPolicy.c
+++ b/bpf_progs/dscpPolicy.c
@@ -34,21 +34,10 @@
 #include "dscpPolicy.h"
 
 #define ECN_MASK 3
-#define IP4_OFFSET(field, header) (header + offsetof(struct iphdr, field))
-#define UPDATE_TOS(dscp, tos) (dscp << 2) | (tos & ECN_MASK)
-#define UPDATE_PRIORITY(dscp) ((dscp >> 2) + 0x60)
-#define UPDATE_FLOW_LABEL(dscp, flow_lbl) ((dscp & 0xf) << 6) + (flow_lbl >> 6)
+#define IP4_OFFSET(field, header) ((header) + offsetof(struct iphdr, field))
+#define UPDATE_TOS(dscp, tos) ((dscp) << 2) | ((tos) & ECN_MASK)
 
-DEFINE_BPF_MAP_GRW(switch_comp_map, ARRAY, int, uint64_t, 1, AID_SYSTEM)
-
-DEFINE_BPF_MAP_GRW(ipv4_socket_to_policies_map_A, HASH, uint64_t, RuleEntry, MAX_POLICIES,
-                   AID_SYSTEM)
-DEFINE_BPF_MAP_GRW(ipv4_socket_to_policies_map_B, HASH, uint64_t, RuleEntry, MAX_POLICIES,
-                   AID_SYSTEM)
-DEFINE_BPF_MAP_GRW(ipv6_socket_to_policies_map_A, HASH, uint64_t, RuleEntry, MAX_POLICIES,
-                   AID_SYSTEM)
-DEFINE_BPF_MAP_GRW(ipv6_socket_to_policies_map_B, HASH, uint64_t, RuleEntry, MAX_POLICIES,
-                   AID_SYSTEM)
+DEFINE_BPF_MAP_GRW(socket_policy_cache_map, HASH, uint64_t, RuleEntry, CACHE_MAP_SIZE, AID_SYSTEM)
 
 DEFINE_BPF_MAP_GRW(ipv4_dscp_policies_map, ARRAY, uint32_t, DscpPolicy, MAX_POLICIES, AID_SYSTEM)
 DEFINE_BPF_MAP_GRW(ipv6_dscp_policies_map, ARRAY, uint32_t, DscpPolicy, MAX_POLICIES, AID_SYSTEM)
@@ -62,14 +51,7 @@
 
     if (data + l2_header_size > data_end) return;
 
-    int zero = 0;
     int hdr_size = 0;
-    uint64_t* selected_map = bpf_switch_comp_map_lookup_elem(&zero);
-
-    // use this with HASH map so map lookup only happens once policies have been added?
-    if (!selected_map) {
-        return;
-    }
 
     // used for map lookup
     uint64_t cookie = bpf_get_socket_cookie(skb);
@@ -80,9 +62,8 @@
     uint8_t protocol = 0;  // TODO: Use are reserved value? Or int (-1) and cast to uint below?
     struct in6_addr src_ip = {};
     struct in6_addr dst_ip = {};
-    uint8_t tos = 0;       // Only used for IPv4
-    uint8_t priority = 0;  // Only used for IPv6
-    uint8_t flow_lbl = 0;  // Only used for IPv6
+    uint8_t tos = 0;             // Only used for IPv4
+    uint32_t old_first_u32 = 0;  // Only used for IPv6
     if (ipv4) {
         const struct iphdr* const iph = (void*)(eth + 1);
         hdr_size = l2_header_size + sizeof(struct iphdr);
@@ -115,8 +96,7 @@
         src_ip = ip6h->saddr;
         dst_ip = ip6h->daddr;
         protocol = ip6h->nexthdr;
-        priority = ip6h->priority;
-        flow_lbl = ip6h->flow_lbl[0];
+        old_first_u32 = *(uint32_t*)ip6h;
     }
 
     switch (protocol) {
@@ -139,42 +119,30 @@
             return;
     }
 
-    RuleEntry* existing_rule;
-    if (ipv4) {
-        if (*selected_map == MAP_A) {
-            existing_rule = bpf_ipv4_socket_to_policies_map_A_lookup_elem(&cookie);
-        } else {
-            existing_rule = bpf_ipv4_socket_to_policies_map_B_lookup_elem(&cookie);
-        }
-    } else {
-        if (*selected_map == MAP_A) {
-            existing_rule = bpf_ipv6_socket_to_policies_map_A_lookup_elem(&cookie);
-        } else {
-            existing_rule = bpf_ipv6_socket_to_policies_map_B_lookup_elem(&cookie);
-        }
-    }
+    RuleEntry* existing_rule = bpf_socket_policy_cache_map_lookup_elem(&cookie);
 
     if (existing_rule && v6_equal(src_ip, existing_rule->src_ip) &&
             v6_equal(dst_ip, existing_rule->dst_ip) && skb->ifindex == existing_rule->ifindex &&
         ntohs(sport) == htons(existing_rule->src_port) &&
         ntohs(dport) == htons(existing_rule->dst_port) && protocol == existing_rule->proto) {
+        if (existing_rule->dscp_val < 0) return;
         if (ipv4) {
             uint8_t newTos = UPDATE_TOS(existing_rule->dscp_val, tos);
             bpf_l3_csum_replace(skb, IP4_OFFSET(check, l2_header_size), htons(tos), htons(newTos),
                                 sizeof(uint16_t));
             bpf_skb_store_bytes(skb, IP4_OFFSET(tos, l2_header_size), &newTos, sizeof(newTos), 0);
         } else {
-            uint8_t new_priority = UPDATE_PRIORITY(existing_rule->dscp_val);
-            uint8_t new_flow_label = UPDATE_FLOW_LABEL(existing_rule->dscp_val, flow_lbl);
-            bpf_skb_store_bytes(skb, 0 + l2_header_size, &new_priority, sizeof(uint8_t), 0);
-            bpf_skb_store_bytes(skb, 1 + l2_header_size, &new_flow_label, sizeof(uint8_t), 0);
+            uint32_t new_first_u32 =
+                htonl(ntohl(old_first_u32) & 0xF03FFFFF | (existing_rule->dscp_val << 22));
+            bpf_skb_store_bytes(skb, l2_header_size, &new_first_u32, sizeof(uint32_t),
+                BPF_F_RECOMPUTE_CSUM);
         }
         return;
     }
 
     // Linear scan ipv4_dscp_policies_map since no stored params match skb.
-    int best_score = -1;
-    uint32_t best_match = 0;
+    int best_score = 0;
+    int8_t new_dscp = -1;
 
     for (register uint64_t i = 0; i < MAX_POLICIES; i++) {
         int score = 0;
@@ -222,35 +190,11 @@
         }
 
         if (score > best_score && temp_mask == policy->present_fields) {
-            best_match = i;
             best_score = score;
+            new_dscp = policy->dscp_val;
         }
     }
 
-    uint8_t new_tos = 0;  // Can 0 be used as default forwarding value?
-    uint8_t new_dscp = 0;
-    uint8_t new_priority = 0;
-    uint8_t new_flow_lbl = 0;
-    if (best_score > 0) {
-        DscpPolicy* policy;
-        if (ipv4) {
-            policy = bpf_ipv4_dscp_policies_map_lookup_elem(&best_match);
-        } else {
-            policy = bpf_ipv6_dscp_policies_map_lookup_elem(&best_match);
-        }
-
-        if (policy) {
-            new_dscp = policy->dscp_val;
-            if (ipv4) {
-                new_tos = UPDATE_TOS(new_dscp, tos);
-            } else {
-                new_priority = UPDATE_PRIORITY(new_dscp);
-                new_flow_lbl = UPDATE_FLOW_LABEL(new_dscp, flow_lbl);
-            }
-        }
-    } else
-        return;
-
     RuleEntry value = {
         .src_ip = src_ip,
         .dst_ip = dst_ip,
@@ -261,34 +205,26 @@
         .dscp_val = new_dscp,
     };
 
-    // Update map with new policy.
-    if (ipv4) {
-        if (*selected_map == MAP_A) {
-            bpf_ipv4_socket_to_policies_map_A_update_elem(&cookie, &value, BPF_ANY);
-        } else {
-            bpf_ipv4_socket_to_policies_map_B_update_elem(&cookie, &value, BPF_ANY);
-        }
-    } else {
-        if (*selected_map == MAP_A) {
-            bpf_ipv6_socket_to_policies_map_A_update_elem(&cookie, &value, BPF_ANY);
-        } else {
-            bpf_ipv6_socket_to_policies_map_B_update_elem(&cookie, &value, BPF_ANY);
-        }
-    }
+    // Update cache with found policy.
+    bpf_socket_policy_cache_map_update_elem(&cookie, &value, BPF_ANY);
+
+    if (new_dscp < 0) return;
 
     // Need to store bytes after updating map or program will not load.
-    if (ipv4 && new_tos != (tos & 252)) {
+    if (ipv4) {
+        uint8_t new_tos = UPDATE_TOS(new_dscp, tos);
         bpf_l3_csum_replace(skb, IP4_OFFSET(check, l2_header_size), htons(tos), htons(new_tos), 2);
         bpf_skb_store_bytes(skb, IP4_OFFSET(tos, l2_header_size), &new_tos, sizeof(new_tos), 0);
-    } else if (!ipv4 && (new_priority != priority || new_flow_lbl != flow_lbl)) {
-        bpf_skb_store_bytes(skb, l2_header_size, &new_priority, sizeof(new_priority), 0);
-        bpf_skb_store_bytes(skb, l2_header_size + 1, &new_flow_lbl, sizeof(new_flow_lbl), 0);
+    } else {
+        uint32_t new_first_u32 = htonl(ntohl(old_first_u32) & 0xF03FFFFF | (new_dscp << 22));
+        bpf_skb_store_bytes(skb, l2_header_size, &new_first_u32, sizeof(uint32_t),
+            BPF_F_RECOMPUTE_CSUM);
     }
     return;
 }
 
-DEFINE_BPF_PROG_KVER("schedcls/set_dscp_ether", AID_ROOT, AID_SYSTEM,
-                     schedcls_set_dscp_ether, KVER(5, 15, 0))
+DEFINE_BPF_PROG_KVER("schedcls/set_dscp_ether", AID_ROOT, AID_SYSTEM, schedcls_set_dscp_ether,
+                     KVER(5, 15, 0))
 (struct __sk_buff* skb) {
     if (skb->pkt_type != PACKET_HOST) return TC_ACT_PIPE;
 
diff --git a/bpf_progs/dscpPolicy.h b/bpf_progs/dscpPolicy.h
index 455a121..c1db6ab 100644
--- a/bpf_progs/dscpPolicy.h
+++ b/bpf_progs/dscpPolicy.h
@@ -14,9 +14,8 @@
  * limitations under the License.
  */
 
+#define CACHE_MAP_SIZE 1024
 #define MAX_POLICIES 16
-#define MAP_A 1
-#define MAP_B 2
 
 #define SRC_IP_MASK_FLAG     1
 #define DST_IP_MASK_FLAG     2
@@ -51,7 +50,7 @@
     __be16 dst_port_start;
     __be16 dst_port_end;
     uint8_t proto;
-    uint8_t dscp_val;
+    int8_t dscp_val;  // -1 none, or 0..63 DSCP value
     uint8_t present_fields;
     uint8_t pad[3];
 } DscpPolicy;
@@ -64,7 +63,7 @@
     __be16 src_port;
     __be16 dst_port;
     __u8 proto;
-    __u8 dscp_val;
+    __s8 dscp_val;  // -1 none, or 0..63 DSCP value
     __u8 pad[2];
 } RuleEntry;
-STRUCT_SIZE(RuleEntry, 2 * 16 + 1 * 4 + 2 * 2 + 2 * 1 + 2);  // 44
\ No newline at end of file
+STRUCT_SIZE(RuleEntry, 2 * 16 + 1 * 4 + 2 * 2 + 2 * 1 + 2);  // 44
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 6ccd77e..28f0699 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -5949,7 +5949,7 @@
      *
      * @param chain target chain to replace.
      * @param uids The list of UIDs to be placed into chain.
-     * @throws IllegalStateException if replacing the firewall chain failed.
+     * @throws UnsupportedOperationException if called on pre-T devices.
      * @throws IllegalArgumentException if {@code chain} is not a valid chain.
      * @hide
      */
diff --git a/service-t/src/com/android/server/net/NetworkStatsFactory.java b/service-t/src/com/android/server/net/NetworkStatsFactory.java
index b628251..21d1a00 100644
--- a/service-t/src/com/android/server/net/NetworkStatsFactory.java
+++ b/service-t/src/com/android/server/net/NetworkStatsFactory.java
@@ -164,7 +164,7 @@
     }
 
     public NetworkStatsFactory(@NonNull Context ctx) {
-        this(ctx, new File("/proc/"), true, new BpfNetMaps());
+        this(ctx, new File("/proc/"), true, new BpfNetMaps(ctx));
     }
 
     @VisibleForTesting
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index d7c5a06..3befd86 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -31,9 +31,11 @@
 import static android.system.OsConstants.ENOENT;
 import static android.system.OsConstants.EOPNOTSUPP;
 
+import android.content.Context;
 import android.net.INetd;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
+import android.provider.DeviceConfig;
 import android.system.ErrnoException;
 import android.system.Os;
 import android.util.Log;
@@ -42,10 +44,15 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.BpfMap;
+import com.android.net.module.util.DeviceConfigUtils;
 import com.android.net.module.util.Struct.U32;
 
 import java.io.FileDescriptor;
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * BpfNetMaps is responsible for providing traffic controller relevant functionality.
@@ -66,6 +73,10 @@
     // Use legacy netd for releases before T.
     private static boolean sInitialized = false;
 
+    private static Boolean sEnableJavaBpfMap = null;
+    private static final String BPF_NET_MAPS_ENABLE_JAVA_BPF_MAP =
+            "bpf_net_maps_enable_java_bpf_map";
+
     // Lock for sConfigurationMap entry for UID_RULES_CONFIGURATION_KEY.
     // This entry is not accessed by others.
     // BpfNetMaps acquires this lock while sequence of read, modify, and write.
@@ -97,6 +108,14 @@
     // LINT.ThenChange(packages/modules/Connectivity/bpf_progs/bpf_shared.h)
 
     /**
+     * Set sEnableJavaBpfMap for test.
+     */
+    @VisibleForTesting
+    public static void setEnableJavaBpfMapForTest(boolean enable) {
+        sEnableJavaBpfMap = enable;
+    }
+
+    /**
      * Set configurationMap for test.
      */
     @VisibleForTesting
@@ -143,8 +162,15 @@
      * Initializes the class if it is not already initialized. This method will open maps but not
      * cause any other effects. This method may be called multiple times on any thread.
      */
-    private static synchronized void ensureInitialized() {
+    private static synchronized void ensureInitialized(final Context context) {
         if (sInitialized) return;
+        if (sEnableJavaBpfMap == null) {
+            sEnableJavaBpfMap = DeviceConfigUtils.isFeatureEnabled(context,
+                    DeviceConfig.NAMESPACE_TETHERING, BPF_NET_MAPS_ENABLE_JAVA_BPF_MAP,
+                    SdkLevel.isAtLeastU() /* defaultValue */);
+        }
+        Log.d(TAG, "BpfNetMaps is initialized with sEnableJavaBpfMap=" + sEnableJavaBpfMap);
+
         setBpfMaps();
         native_init();
         sInitialized = true;
@@ -164,20 +190,20 @@
     }
 
     /** Constructor used after T that doesn't need to use netd anymore. */
-    public BpfNetMaps() {
-        this(null);
+    public BpfNetMaps(final Context context) {
+        this(context, null);
 
         if (PRE_T) throw new IllegalArgumentException("BpfNetMaps need to use netd before T");
     }
 
-    public BpfNetMaps(final INetd netd) {
-        this(netd, new Dependencies());
+    public BpfNetMaps(final Context context, final INetd netd) {
+        this(context, netd, new Dependencies());
     }
 
     @VisibleForTesting
-    public BpfNetMaps(final INetd netd, final Dependencies deps) {
+    public BpfNetMaps(final Context context, final INetd netd, final Dependencies deps) {
         if (!PRE_T) {
-            ensureInitialized();
+            ensureInitialized(context);
         }
         mNetd = netd;
         mDeps = deps;
@@ -404,25 +430,46 @@
 
     /**
      * Replaces the contents of the specified UID-based firewall chain.
+     * Enables the chain for specified uids and disables the chain for non-specified uids.
      *
-     * The chain may be an allowlist chain or a denylist chain. A denylist chain contains DROP
-     * rules for the specified UIDs and a RETURN rule at the end. An allowlist chain contains RETURN
-     * rules for the system UID range (0 to {@code UID_APP} - 1), RETURN rules for the specified
-     * UIDs, and a DROP rule at the end. The chain will be created if it does not exist.
-     *
-     * @param chainName   The name of the chain to replace.
-     * @param isAllowlist Whether this is an allowlist or denylist chain.
+     * @param chain       Target chain.
      * @param uids        The list of UIDs to allow/deny.
-     * @return 0 if the chain was successfully replaced, errno otherwise.
+     * @throws UnsupportedOperationException if called on pre-T devices.
+     * @throws IllegalArgumentException if {@code chain} is not a valid chain.
      */
-    public int replaceUidChain(final String chainName, final boolean isAllowlist,
-            final int[] uids) {
-        synchronized (sUidOwnerMap) {
-            final int err = native_replaceUidChain(chainName, isAllowlist, uids);
-            if (err != 0) {
-                Log.e(TAG, "replaceUidChain failed: " + Os.strerror(-err));
+    public void replaceUidChain(final int chain, final int[] uids) {
+        throwIfPreT("replaceUidChain is not available on pre-T devices");
+
+        final long match;
+        try {
+            match = getMatchByFirewallChain(chain);
+        } catch (ServiceSpecificException e) {
+            // Throws IllegalArgumentException to keep the behavior of
+            // ConnectivityManager#replaceFirewallChain API
+            throw new IllegalArgumentException("Invalid firewall chain: " + chain);
+        }
+        final Set<Integer> uidSet = Arrays.stream(uids).boxed().collect(Collectors.toSet());
+        final Set<Integer> uidSetToRemoveRule = new HashSet<>();
+        try {
+            synchronized (sUidOwnerMap) {
+                sUidOwnerMap.forEach((uid, config) -> {
+                    // config could be null if there is a concurrent entry deletion.
+                    // http://b/220084230.
+                    if (config != null
+                            && !uidSet.contains((int) uid.val) && (config.rule & match) != 0) {
+                        uidSetToRemoveRule.add((int) uid.val);
+                    }
+                });
+
+                for (final int uid : uidSetToRemoveRule) {
+                    removeRule(uid, match, "replaceUidChain");
+                }
+                for (final int uid : uids) {
+                    addRule(uid, match, "replaceUidChain");
+                }
             }
-            return -err;
+        } catch (ErrnoException | ServiceSpecificException e) {
+            Log.e(TAG, "replaceUidChain failed: " + e);
         }
     }
 
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 5602f26..7b5026d 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -1399,8 +1399,8 @@
          * @param netd
          * @return BpfNetMaps implementation.
          */
-        public BpfNetMaps getBpfNetMaps(INetd netd) {
-            return new BpfNetMaps(netd);
+        public BpfNetMaps getBpfNetMaps(Context context, INetd netd) {
+            return new BpfNetMaps(context, netd);
         }
 
         /**
@@ -1529,7 +1529,7 @@
         mProxyTracker = mDeps.makeProxyTracker(mContext, mHandler);
 
         mNetd = netd;
-        mBpfNetMaps = mDeps.getBpfNetMaps(netd);
+        mBpfNetMaps = mDeps.getBpfNetMaps(mContext, netd);
         mTelephonyManager = (TelephonyManager) mContext.getSystemService(Context.TELEPHONY_SERVICE);
         mAppOpsManager = (AppOpsManager) mContext.getSystemService(Context.APP_OPS_SERVICE);
         mLocationPermissionChecker = mDeps.makeLocationPermissionChecker(mContext);
@@ -11399,39 +11399,6 @@
     public void replaceFirewallChain(final int chain, final int[] uids) {
         enforceNetworkStackOrSettingsPermission();
 
-        try {
-            switch (chain) {
-                case ConnectivityManager.FIREWALL_CHAIN_DOZABLE:
-                    mBpfNetMaps.replaceUidChain("fw_dozable", true /* isAllowList */, uids);
-                    break;
-                case ConnectivityManager.FIREWALL_CHAIN_STANDBY:
-                    mBpfNetMaps.replaceUidChain("fw_standby", false /* isAllowList */, uids);
-                    break;
-                case ConnectivityManager.FIREWALL_CHAIN_POWERSAVE:
-                    mBpfNetMaps.replaceUidChain("fw_powersave", true /* isAllowList */, uids);
-                    break;
-                case ConnectivityManager.FIREWALL_CHAIN_RESTRICTED:
-                    mBpfNetMaps.replaceUidChain("fw_restricted", true /* isAllowList */, uids);
-                    break;
-                case ConnectivityManager.FIREWALL_CHAIN_LOW_POWER_STANDBY:
-                    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;
-                case ConnectivityManager.FIREWALL_CHAIN_OEM_DENY_3:
-                    mBpfNetMaps.replaceUidChain("fw_oem_deny_3", false /* isAllowList */, uids);
-                    break;
-                default:
-                    throw new IllegalArgumentException("replaceFirewallChain with invalid chain: "
-                            + chain);
-            }
-        } catch (ServiceSpecificException e) {
-            throw new IllegalStateException(e);
-        }
+        mBpfNetMaps.replaceUidChain(chain, uids);
     }
 }
diff --git a/service/src/com/android/server/connectivity/DscpPolicyTracker.java b/service/src/com/android/server/connectivity/DscpPolicyTracker.java
index 8cb3213..2bfad10 100644
--- a/service/src/com/android/server/connectivity/DscpPolicyTracker.java
+++ b/service/src/com/android/server/connectivity/DscpPolicyTracker.java
@@ -185,7 +185,7 @@
                         new DscpPolicyValue(policy.getSourceAddress(),
                             policy.getDestinationAddress(), ifIndex,
                             policy.getSourcePort(), policy.getDestinationPortRange(),
-                            (short) policy.getProtocol(), (short) policy.getDscpValue()));
+                            (short) policy.getProtocol(), (byte) policy.getDscpValue()));
             }
 
             // Add v6 policy to mBpfDscpIpv6Policies if source and destination address
@@ -196,7 +196,7 @@
                         new DscpPolicyValue(policy.getSourceAddress(),
                                 policy.getDestinationAddress(), ifIndex,
                                 policy.getSourcePort(), policy.getDestinationPortRange(),
-                                (short) policy.getProtocol(), (short) policy.getDscpValue()));
+                                (short) policy.getProtocol(), (byte) policy.getDscpValue()));
             }
 
             ifacePolicies.put(policy.getPolicyId(), addIndex);
diff --git a/service/src/com/android/server/connectivity/DscpPolicyValue.java b/service/src/com/android/server/connectivity/DscpPolicyValue.java
index 6e4e7eb..4bb41da 100644
--- a/service/src/com/android/server/connectivity/DscpPolicyValue.java
+++ b/service/src/com/android/server/connectivity/DscpPolicyValue.java
@@ -52,8 +52,8 @@
     @Field(order = 6, type = Type.U8)
     public final short proto;
 
-    @Field(order = 7, type = Type.U8)
-    public final short dscp;
+    @Field(order = 7, type = Type.S8)
+    public final byte dscp;
 
     @Field(order = 8, type = Type.U8, padding = 3)
     public final short mask;
@@ -100,7 +100,7 @@
             InetAddress.parseNumericAddress("::").getAddress();
 
     private short makeMask(final byte[] src46, final byte[] dst46, final int srcPort,
-            final int dstPortStart, final short proto, final short dscp) {
+            final int dstPortStart, final short proto, final byte dscp) {
         short mask = 0;
         if (src46 != EMPTY_ADDRESS_FIELD) {
             mask |= SRC_IP_MASK;
@@ -122,7 +122,7 @@
 
     private DscpPolicyValue(final InetAddress src46, final InetAddress dst46, final long ifIndex,
             final int srcPort, final int dstPortStart, final int dstPortEnd, final short proto,
-            final short dscp) {
+            final byte dscp) {
         this.src46 = toAddressField(src46);
         this.dst46 = toAddressField(dst46);
         this.ifIndex = ifIndex;
@@ -142,7 +142,7 @@
 
     public DscpPolicyValue(final InetAddress src46, final InetAddress dst46, final long ifIndex,
             final int srcPort, final Range<Integer> dstPort, final short proto,
-            final short dscp) {
+            final byte dscp) {
         this(src46, dst46, ifIndex, srcPort, dstPort != null ? dstPort.getLower() : -1,
                 dstPort != null ? dstPort.getUpper() : -1, proto, dscp);
     }
@@ -150,7 +150,7 @@
     public static final DscpPolicyValue NONE = new DscpPolicyValue(
             null /* src46 */, null /* dst46 */, 0 /* ifIndex */, -1 /* srcPort */,
             -1 /* dstPortStart */, -1 /* dstPortEnd */, (short) -1 /* proto */,
-            (short) 0 /* dscp */);
+            (byte) -1 /* dscp */);
 
     @Override
     public String toString() {
diff --git a/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt b/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
index efc24d3..73e4c0e 100644
--- a/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
+++ b/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
@@ -209,7 +209,7 @@
         doReturn(mock(ProxyTracker::class.java)).`when`(deps).makeProxyTracker(any(), any())
         doReturn(mock(MockableSystemProperties::class.java)).`when`(deps).systemProperties
         doReturn(TestNetIdManager()).`when`(deps).makeNetIdManager()
-        doReturn(mock(BpfNetMaps::class.java)).`when`(deps).getBpfNetMaps(any())
+        doReturn(mock(BpfNetMaps::class.java)).`when`(deps).getBpfNetMaps(any(), any())
         doAnswer { inv ->
             object : MultinetworkPolicyTracker(inv.getArgument(0), inv.getArgument(1),
                     inv.getArgument(2)) {
diff --git a/tests/mts/bpf_existence_test.cpp b/tests/mts/bpf_existence_test.cpp
index 0911ecf..c7e8b97 100644
--- a/tests/mts/bpf_existence_test.cpp
+++ b/tests/mts/bpf_existence_test.cpp
@@ -88,12 +88,8 @@
     SHARED "map_clatd_clat_egress4_map",
     SHARED "map_clatd_clat_ingress6_map",
     SHARED "map_dscpPolicy_ipv4_dscp_policies_map",
-    SHARED "map_dscpPolicy_ipv4_socket_to_policies_map_A",
-    SHARED "map_dscpPolicy_ipv4_socket_to_policies_map_B",
     SHARED "map_dscpPolicy_ipv6_dscp_policies_map",
-    SHARED "map_dscpPolicy_ipv6_socket_to_policies_map_A",
-    SHARED "map_dscpPolicy_ipv6_socket_to_policies_map_B",
-    SHARED "map_dscpPolicy_switch_comp_map",
+    SHARED "map_dscpPolicy_socket_policy_cache_map",
     NETD "map_netd_app_uid_stats_map",
     NETD "map_netd_configuration_map",
     NETD "map_netd_cookie_tag_map",
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index 0908ad2..ee72f85 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -70,7 +70,6 @@
         "java/com/android/server/connectivity/NetdEventListenerServiceTest.java",
         "java/com/android/server/connectivity/VpnTest.java",
         "java/com/android/server/net/ipmemorystore/*.java",
-        "java/com/android/server/connectivity/mdns/**/*.java",
     ]
 }
 
@@ -112,7 +111,8 @@
         "service-connectivity-pre-jarjar",
         "service-connectivity-tiramisu-pre-jarjar",
         "services.core-vpn",
-        "cts-net-utils"
+        "cts-net-utils",
+        "service-mdns",
     ],
     libs: [
         "android.net.ipsec.ike.stubs.module_lib",
@@ -145,7 +145,6 @@
     static_libs: [
         "services.core",
         "services.net",
-        "service-mdns",
     ],
     jni_libs: [
         "libandroid_net_connectivity_com_android_net_module_util_jni",
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index 0718952..2d09bf2 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -46,6 +46,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.verify;
 
+import android.content.Context;
 import android.net.INetd;
 import android.os.Build;
 import android.os.ServiceSpecificException;
@@ -102,6 +103,7 @@
 
     @Mock INetd mNetd;
     @Mock BpfNetMaps.Dependencies mDeps;
+    @Mock Context mContext;
     private final BpfMap<U32, U32> mConfigurationMap = new TestBpfMap<>(U32.class, U32.class);
     private final BpfMap<U32, UidOwnerValue> mUidOwnerMap =
             new TestBpfMap<>(U32.class, UidOwnerValue.class);
@@ -110,9 +112,10 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         doReturn(TEST_IF_INDEX).when(mDeps).getIfIndex(TEST_IF_NAME);
+        BpfNetMaps.setEnableJavaBpfMapForTest(true /* enable */);
         BpfNetMaps.setConfigurationMapForTest(mConfigurationMap);
         BpfNetMaps.setUidOwnerMapForTest(mUidOwnerMap);
-        mBpfNetMaps = new BpfNetMaps(mNetd, mDeps);
+        mBpfNetMaps = new BpfNetMaps(mContext, mNetd, mDeps);
     }
 
     @Test
@@ -649,4 +652,80 @@
         assertThrows(UnsupportedOperationException.class, () ->
                 mBpfNetMaps.setUidRule(FIREWALL_CHAIN_DOZABLE, TEST_UID, FIREWALL_RULE_ALLOW));
     }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testReplaceUidChain() throws Exception {
+        final int uid0 = TEST_UIDS[0];
+        final int uid1 = TEST_UIDS[1];
+
+        mBpfNetMaps.replaceUidChain(FIREWALL_CHAIN_DOZABLE, TEST_UIDS);
+
+        checkUidOwnerValue(uid0, NO_IIF, DOZABLE_MATCH);
+        checkUidOwnerValue(uid1, NO_IIF, DOZABLE_MATCH);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testReplaceUidChainWithOtherMatch() throws Exception {
+        final int uid0 = TEST_UIDS[0];
+        final int uid1 = TEST_UIDS[1];
+        final long match0 = POWERSAVE_MATCH;
+        final long match1 = POWERSAVE_MATCH | RESTRICTED_MATCH;
+        mUidOwnerMap.updateEntry(new U32(uid0), new UidOwnerValue(NO_IIF, match0));
+        mUidOwnerMap.updateEntry(new U32(uid1), new UidOwnerValue(NO_IIF, match1));
+
+        mBpfNetMaps.replaceUidChain(FIREWALL_CHAIN_DOZABLE, new int[]{uid1});
+
+        checkUidOwnerValue(uid0, NO_IIF, match0);
+        checkUidOwnerValue(uid1, NO_IIF, match1 | DOZABLE_MATCH);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testReplaceUidChainWithExistingIifMatch() throws Exception {
+        final int uid0 = TEST_UIDS[0];
+        final int uid1 = TEST_UIDS[1];
+        final long match0 = IIF_MATCH;
+        final long match1 = IIF_MATCH | POWERSAVE_MATCH | RESTRICTED_MATCH;
+        mUidOwnerMap.updateEntry(new U32(uid0), new UidOwnerValue(TEST_IF_INDEX, match0));
+        mUidOwnerMap.updateEntry(new U32(uid1), new UidOwnerValue(NULL_IIF, match1));
+
+        mBpfNetMaps.replaceUidChain(FIREWALL_CHAIN_DOZABLE, TEST_UIDS);
+
+        checkUidOwnerValue(uid0, TEST_IF_INDEX, match0 | DOZABLE_MATCH);
+        checkUidOwnerValue(uid1, NULL_IIF, match1 | DOZABLE_MATCH);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testReplaceUidChainRemoveExistingMatch() throws Exception {
+        final int uid0 = TEST_UIDS[0];
+        final int uid1 = TEST_UIDS[1];
+        final long match0 = IIF_MATCH | DOZABLE_MATCH;
+        final long match1 = IIF_MATCH | POWERSAVE_MATCH | RESTRICTED_MATCH;
+        mUidOwnerMap.updateEntry(new U32(uid0), new UidOwnerValue(TEST_IF_INDEX, match0));
+        mUidOwnerMap.updateEntry(new U32(uid1), new UidOwnerValue(NULL_IIF, match1));
+
+        mBpfNetMaps.replaceUidChain(FIREWALL_CHAIN_DOZABLE, new int[]{uid1});
+
+        checkUidOwnerValue(uid0, TEST_IF_INDEX, match0 & ~DOZABLE_MATCH);
+        checkUidOwnerValue(uid1, NULL_IIF, match1 | DOZABLE_MATCH);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testReplaceUidChainInvalidChain() {
+        final Class<IllegalArgumentException> expected = IllegalArgumentException.class;
+        assertThrows(expected, () -> mBpfNetMaps.replaceUidChain(-1 /* chain */, TEST_UIDS));
+        assertThrows(expected, () -> mBpfNetMaps.replaceUidChain(1000 /* chain */, TEST_UIDS));
+    }
+
+    @Test
+    @IgnoreAfter(Build.VERSION_CODES.S_V2)
+    public void testReplaceUidChainBeforeT() {
+        assertThrows(UnsupportedOperationException.class,
+                () -> mBpfNetMaps.replaceUidChain(FIREWALL_CHAIN_DOZABLE, TEST_UIDS));
+    }
+
 }
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 0919dfc..62073cd 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -2061,7 +2061,7 @@
         }
 
         @Override
-        public BpfNetMaps getBpfNetMaps(INetd netd) {
+        public BpfNetMaps getBpfNetMaps(Context context, INetd netd) {
             return mBpfNetMaps;
         }
 
@@ -9592,24 +9592,23 @@
         }
     }
 
-    private void doTestReplaceFirewallChain(final int chain, final String chainName,
-            final boolean allowList) {
+    private void doTestReplaceFirewallChain(final int chain) {
         final int[] uids = new int[] {1001, 1002};
         mCm.replaceFirewallChain(chain, uids);
-        verify(mBpfNetMaps).replaceUidChain(chainName, allowList, uids);
+        verify(mBpfNetMaps).replaceUidChain(chain, 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);
-        doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_3, "fw_oem_deny_3", false);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_DOZABLE);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_STANDBY);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_POWERSAVE);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_RESTRICTED);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_LOW_POWER_STANDBY);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_1);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_2);
+        doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_3);
     }
 
     @Test @IgnoreUpTo(SC_V2)
@@ -9620,8 +9619,6 @@
                 () -> 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)