Merge "[Thread] flagging Android Thread APIs" into main
diff --git a/Tethering/Android.bp b/Tethering/Android.bp
index 4478b1e..e69b872 100644
--- a/Tethering/Android.bp
+++ b/Tethering/Android.bp
@@ -95,6 +95,7 @@
     ],
     static_libs: [
         "NetworkStackApiCurrentShims",
+        "net-utils-device-common-struct",
     ],
     apex_available: ["com.android.tethering"],
     lint: { strict_updatability_linting: true },
@@ -109,6 +110,7 @@
     ],
     static_libs: [
         "NetworkStackApiStableShims",
+        "net-utils-device-common-struct",
     ],
     apex_available: ["com.android.tethering"],
     lint: { strict_updatability_linting: true },
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/util/SyncStateMachineTest.kt b/Tethering/tests/unit/src/com/android/networkstack/tethering/util/SyncStateMachineTest.kt
new file mode 100644
index 0000000..3a57fdd
--- /dev/null
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/util/SyncStateMachineTest.kt
@@ -0,0 +1,294 @@
+/**
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.networkstack.tethering.util
+
+import android.os.Message
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.util.State
+import com.android.networkstack.tethering.util.SyncStateMachine.StateInfo
+import java.util.ArrayDeque
+import java.util.ArrayList
+import kotlin.test.assertFailsWith
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.any
+import org.mockito.Mockito.inOrder
+import org.mockito.Mockito.spy
+import org.mockito.Mockito.verifyNoMoreInteractions
+
+private const val MSG_INVALID = -1
+private const val MSG_1 = 1
+private const val MSG_2 = 2
+private const val MSG_3 = 3
+private const val MSG_4 = 4
+private const val MSG_5 = 5
+private const val MSG_6 = 6
+private const val MSG_7 = 7
+private const val ARG_1 = 100
+private const val ARG_2 = 200
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class SynStateMachineTest {
+    private val mState1 = spy(object : TestState(MSG_1) {})
+    private val mState2 = spy(object : TestState(MSG_2) {})
+    private val mState3 = spy(object : TestState(MSG_3) {})
+    private val mState4 = spy(object : TestState(MSG_4) {})
+    private val mState5 = spy(object : TestState(MSG_5) {})
+    private val mState6 = spy(object : TestState(MSG_6) {})
+    private val mState7 = spy(object : TestState(MSG_7) {})
+    private val mInOrder = inOrder(mState1, mState2, mState3, mState4, mState5, mState6, mState7)
+    // Lazy initialize to make sure running in test thread.
+    private val mSM by lazy {
+        SyncStateMachine("TestSyncStateMachine", Thread.currentThread(), true /* debug */)
+    }
+    private val mAllStates = ArrayList<StateInfo>()
+
+    private val mMsgProcessedResults = ArrayDeque<Pair<State, Int>>()
+
+    open inner class TestState(val expected: Int) : State() {
+        // Control destination state in obj field for testing.
+        override fun processMessage(msg: Message): Boolean {
+            mMsgProcessedResults.add(this to msg.what)
+            assertEquals(ARG_1, msg.arg1)
+            assertEquals(ARG_2, msg.arg2)
+
+            if (msg.what == expected) {
+                msg.obj?.let { mSM.transitionTo(it as State) }
+                return true
+            }
+
+            return false
+        }
+    }
+
+    private fun verifyNoMoreInteractions() {
+        verifyNoMoreInteractions(mState1, mState2, mState3, mState4, mState5, mState6)
+    }
+
+    private fun processMessage(what: Int, toState: State?) {
+        mSM.processMessage(what, ARG_1, ARG_2, toState)
+    }
+
+    private fun verifyMessageProcessedBy(what: Int, vararg processedStates: State) {
+        for (state in processedStates) {
+            // InOrder.verify can't check the Message content here because SyncSM will recycle the
+            // message after it's been processed. SyncSM reuses the same Message instance for all
+            // messages it processes. So, if using InOrder.verify to verify the content of a message
+            // after SyncSM has processed it, the content would be wrong.
+            mInOrder.verify(state).processMessage(any())
+            val (processedState, msgWhat) = mMsgProcessedResults.remove()
+            assertEquals(state, processedState)
+            assertEquals(what, msgWhat)
+        }
+        assertTrue(mMsgProcessedResults.isEmpty())
+    }
+
+    @Test
+    fun testInitialState() {
+        // mState1 -> initial
+        //    |
+        // mState2
+        mAllStates.add(StateInfo(mState1, null))
+        mAllStates.add(StateInfo(mState2, mState1))
+        mSM.addAllStates(mAllStates)
+
+        mSM.start(mState1)
+        mInOrder.verify(mState1).enter()
+        verifyNoMoreInteractions()
+    }
+
+    @Test
+    fun testStartFromLeafState() {
+        // mState1 -> initial
+        //    |
+        // mState2
+        //    |
+        // mState3
+        mAllStates.add(StateInfo(mState1, null))
+        mAllStates.add(StateInfo(mState2, mState1))
+        mAllStates.add(StateInfo(mState3, mState2))
+        mSM.addAllStates(mAllStates)
+
+        mSM.start(mState3)
+        mInOrder.verify(mState1).enter()
+        mInOrder.verify(mState2).enter()
+        mInOrder.verify(mState3).enter()
+        verifyNoMoreInteractions()
+    }
+
+    private fun verifyStart() {
+        mSM.addAllStates(mAllStates)
+        mSM.start(mState1)
+        mInOrder.verify(mState1).enter()
+        verifyNoMoreInteractions()
+    }
+
+    fun addState(state: State, parent: State? = null) {
+        mAllStates.add(StateInfo(state, parent))
+    }
+
+    @Test
+    fun testAddState() {
+        // Add duplicated states.
+        mAllStates.add(StateInfo(mState1, null))
+        mAllStates.add(StateInfo(mState1, null))
+        assertFailsWith(IllegalStateException::class) {
+            mSM.addAllStates(mAllStates)
+        }
+    }
+
+    @Test
+    fun testProcessMessage() {
+        // mState1
+        //    |
+        // mState2
+        addState(mState1)
+        addState(mState2, mState1)
+        verifyStart()
+
+        processMessage(MSG_1, null)
+        verifyMessageProcessedBy(MSG_1, mState1)
+        verifyNoMoreInteractions()
+    }
+
+    @Test
+    fun testTwoStates() {
+        // mState1 <-initial, mState2
+        addState(mState1)
+        addState(mState2)
+        verifyStart()
+
+        // Test transition to mState2
+        processMessage(MSG_1, mState2)
+        verifyMessageProcessedBy(MSG_1, mState1)
+        mInOrder.verify(mState1).exit()
+        mInOrder.verify(mState2).enter()
+        verifyNoMoreInteractions()
+
+        // If set destState to mState2 (current state), no state transition.
+        processMessage(MSG_2, mState2)
+        verifyMessageProcessedBy(MSG_2, mState2)
+        verifyNoMoreInteractions()
+    }
+
+    @Test
+    fun testTwoStateTrees() {
+        //    mState1 -> initial  mState4
+        //    /     \             /     \
+        // mState2 mState3     mState5 mState6
+        addState(mState1)
+        addState(mState2, mState1)
+        addState(mState3, mState1)
+        addState(mState4)
+        addState(mState5, mState4)
+        addState(mState6, mState4)
+        verifyStart()
+
+        //    mState1 -> current     mState4
+        //    /     \                /     \
+        // mState2 mState3 -> dest mState5 mState6
+        processMessage(MSG_1, mState3)
+        verifyMessageProcessedBy(MSG_1, mState1)
+        mInOrder.verify(mState3).enter()
+        verifyNoMoreInteractions()
+
+        //           mState1                     mState4
+        //           /     \                     /     \
+        // dest <- mState2 mState3 -> current mState5 mState6
+        processMessage(MSG_1, mState2)
+        verifyMessageProcessedBy(MSG_1, mState3, mState1)
+        mInOrder.verify(mState3).exit()
+        mInOrder.verify(mState2).enter()
+        verifyNoMoreInteractions()
+
+        //               mState1          mState4
+        //               /     \          /     \
+        // current <- mState2 mState3 mState5 mState6 -> dest
+        processMessage(MSG_2, mState6)
+        verifyMessageProcessedBy(MSG_2, mState2)
+        mInOrder.verify(mState2).exit()
+        mInOrder.verify(mState1).exit()
+        mInOrder.verify(mState4).enter()
+        mInOrder.verify(mState6).enter()
+        verifyNoMoreInteractions()
+    }
+
+    @Test
+    fun testMultiDepthTransition() {
+        //      mState1 -> current
+        //    |          \
+        //  mState2         mState6
+        //    |   \           |
+        //  mState3 mState5  mState7
+        //    |
+        //  mState4
+        addState(mState1)
+        addState(mState2, mState1)
+        addState(mState6, mState1)
+        addState(mState3, mState2)
+        addState(mState5, mState2)
+        addState(mState7, mState6)
+        addState(mState4, mState3)
+        verifyStart()
+
+        //      mState1 -> current
+        //    |          \
+        //  mState2         mState6
+        //    |   \           |
+        //  mState3 mState5  mState7
+        //    |
+        //  mState4 -> dest
+        processMessage(MSG_1, mState4)
+        verifyMessageProcessedBy(MSG_1, mState1)
+        mInOrder.verify(mState2).enter()
+        mInOrder.verify(mState3).enter()
+        mInOrder.verify(mState4).enter()
+        verifyNoMoreInteractions()
+
+        //            mState1
+        //        /            \
+        //  mState2             mState6
+        //    |   \                 \
+        //  mState3 mState5 -> dest  mState7
+        //    |
+        //  mState4 -> current
+        processMessage(MSG_1, mState5)
+        verifyMessageProcessedBy(MSG_1, mState4, mState3, mState2, mState1)
+        mInOrder.verify(mState4).exit()
+        mInOrder.verify(mState3).exit()
+        mInOrder.verify(mState5).enter()
+        verifyNoMoreInteractions()
+
+        //            mState1
+        //        /              \
+        //  mState2               mState6
+        //    |   \                    \
+        //  mState3 mState5 -> current  mState7 -> dest
+        //    |
+        //  mState4
+        processMessage(MSG_2, mState7)
+        verifyMessageProcessedBy(MSG_2, mState5, mState2)
+        mInOrder.verify(mState5).exit()
+        mInOrder.verify(mState2).exit()
+        mInOrder.verify(mState6).enter()
+        mInOrder.verify(mState7).enter()
+        verifyNoMoreInteractions()
+    }
+}
diff --git a/common/Android.bp b/common/Android.bp
index c982431..1d73a46 100644
--- a/common/Android.bp
+++ b/common/Android.bp
@@ -25,6 +25,7 @@
 // as the above target may not exist
 // depending on the branch
 
+// The library requires the final artifact to contain net-utils-device-common-struct.
 java_library {
     name: "connectivity-net-module-utils-bpf",
     srcs: [
@@ -40,8 +41,9 @@
     libs: [
         "androidx.annotation_annotation",
         "framework-connectivity.stubs.module_lib",
-    ],
-    static_libs: [
+        // For libraries which are statically linked in framework-connectivity, do not
+        // statically link here because callers of this library might already have a static
+        // version linked.
         "net-utils-device-common-struct",
     ],
     apex_available: [
diff --git a/common/flags.aconfig b/common/flags.aconfig
index 7235202..2e552a1 100644
--- a/common/flags.aconfig
+++ b/common/flags.aconfig
@@ -20,10 +20,3 @@
   description: "Remove expired services from MdnsServiceCache"
   bug: "304649384"
 }
-
-flag {
-  name: "set_data_saver_via_cm"
-  namespace: "android_core_networking"
-  description: "Set data saver through ConnectivityManager API"
-  bug: "297836825"
-}
diff --git a/framework/Android.bp b/framework/Android.bp
index 182c558..449e652 100644
--- a/framework/Android.bp
+++ b/framework/Android.bp
@@ -101,7 +101,7 @@
         "framework-connectivity-javastream-protos",
     ],
     impl_only_static_libs: [
-        "net-utils-device-common-struct",
+        "net-utils-device-common-bpf",
     ],
     libs: [
         "androidx.annotation_annotation",
@@ -130,7 +130,7 @@
         // to generate the SDK stubs.
         // Even if the library is included in "impl_only_static_libs" of defaults. This is still
         // needed because java_library which doesn't understand "impl_only_static_libs".
-        "net-utils-device-common-struct",
+        "net-utils-device-common-bpf",
     ],
     libs: [
         // This cannot be in the defaults clause above because if it were, it would be used
diff --git a/framework/api/module-lib-current.txt b/framework/api/module-lib-current.txt
index 782e20a..193bd92 100644
--- a/framework/api/module-lib-current.txt
+++ b/framework/api/module-lib-current.txt
@@ -24,7 +24,6 @@
     method @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_SETTINGS, android.Manifest.permission.NETWORK_SETUP_WIZARD, android.Manifest.permission.NETWORK_STACK, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public void setAcceptPartialConnectivity(@NonNull android.net.Network, boolean, boolean);
     method @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_SETTINGS, android.Manifest.permission.NETWORK_SETUP_WIZARD, android.Manifest.permission.NETWORK_STACK, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public void setAcceptUnvalidated(@NonNull android.net.Network, boolean, boolean);
     method @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_SETTINGS, android.Manifest.permission.NETWORK_SETUP_WIZARD, android.Manifest.permission.NETWORK_STACK, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public void setAvoidUnvalidated(@NonNull android.net.Network);
-    method @FlaggedApi("com.android.net.flags.set_data_saver_via_cm") @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_SETTINGS, android.Manifest.permission.NETWORK_STACK, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public void setDataSaverEnabled(boolean);
     method @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_SETTINGS, android.Manifest.permission.NETWORK_STACK, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public void setFirewallChainEnabled(int, boolean);
     method @RequiresPermission(android.Manifest.permission.NETWORK_STACK) public void setGlobalProxy(@Nullable android.net.ProxyInfo);
     method @RequiresPermission(anyOf={android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK, android.Manifest.permission.NETWORK_STACK, android.Manifest.permission.NETWORK_SETTINGS}) public void setLegacyLockdownVpnEnabled(boolean);
diff --git a/framework/src/android/net/BpfNetMapsReader.java b/framework/src/android/net/BpfNetMapsReader.java
new file mode 100644
index 0000000..49e874a
--- /dev/null
+++ b/framework/src/android/net/BpfNetMapsReader.java
@@ -0,0 +1,179 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net;
+
+import static android.net.BpfNetMapsConstants.CONFIGURATION_MAP_PATH;
+import static android.net.BpfNetMapsConstants.UID_OWNER_MAP_PATH;
+import static android.net.BpfNetMapsConstants.UID_RULES_CONFIGURATION_KEY;
+import static android.net.BpfNetMapsUtils.getMatchByFirewallChain;
+import static android.net.BpfNetMapsUtils.isFirewallAllowList;
+import static android.net.BpfNetMapsUtils.throwIfPreT;
+import static android.net.ConnectivityManager.FIREWALL_RULE_ALLOW;
+import static android.net.ConnectivityManager.FIREWALL_RULE_DENY;
+
+import android.annotation.NonNull;
+import android.annotation.RequiresApi;
+import android.os.Build;
+import android.os.ServiceSpecificException;
+import android.system.ErrnoException;
+import android.system.Os;
+
+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.IBpfMap;
+import com.android.net.module.util.Struct;
+import com.android.net.module.util.Struct.S32;
+import com.android.net.module.util.Struct.U32;
+
+/**
+ * A helper class to *read* java BpfMaps.
+ * @hide
+ */
+@RequiresApi(Build.VERSION_CODES.TIRAMISU)  // BPF maps were only mainlined in T
+public class BpfNetMapsReader {
+    // Locally store the handle of bpf maps. The FileDescriptors are statically cached inside the
+    // BpfMap implementation.
+
+    // Bpf map to store various networking configurations, the format of the value is different
+    // for different keys. See BpfNetMapsConstants#*_CONFIGURATION_KEY for keys.
+    private final IBpfMap<S32, U32> mConfigurationMap;
+    // Bpf map to store per uid traffic control configurations.
+    // See {@link UidOwnerValue} for more detail.
+    private final IBpfMap<S32, UidOwnerValue> mUidOwnerMap;
+    private final Dependencies mDeps;
+
+    public BpfNetMapsReader() {
+        this(new Dependencies());
+    }
+
+    @VisibleForTesting
+    public BpfNetMapsReader(@NonNull Dependencies deps) {
+        if (!SdkLevel.isAtLeastT()) {
+            throw new UnsupportedOperationException(
+                    BpfNetMapsReader.class.getSimpleName() + " is not supported below Android T");
+        }
+        mDeps = deps;
+        mConfigurationMap = mDeps.getConfigurationMap();
+        mUidOwnerMap = mDeps.getUidOwnerMap();
+    }
+
+    /**
+     * Dependencies of BpfNetMapReader, for injection in tests.
+     */
+    @VisibleForTesting
+    public static class Dependencies {
+        /** Get the configuration map. */
+        public IBpfMap<S32, U32> getConfigurationMap() {
+            try {
+                return new BpfMap<>(CONFIGURATION_MAP_PATH, BpfMap.BPF_F_RDONLY,
+                        S32.class, U32.class);
+            } catch (ErrnoException e) {
+                throw new IllegalStateException("Cannot open configuration map", e);
+            }
+        }
+
+        /** Get the uid owner map. */
+        public IBpfMap<S32, UidOwnerValue> getUidOwnerMap() {
+            try {
+                return new BpfMap<>(UID_OWNER_MAP_PATH, BpfMap.BPF_F_RDONLY,
+                        S32.class, UidOwnerValue.class);
+            } catch (ErrnoException e) {
+                throw new IllegalStateException("Cannot open uid owner map", e);
+            }
+        }
+    }
+
+    /**
+     * Get the specified firewall chain's status.
+     *
+     * @param chain target chain
+     * @return {@code true} if chain is enabled, {@code false} if chain is not enabled.
+     * @throws UnsupportedOperationException if called on pre-T devices.
+     * @throws ServiceSpecificException in case of failure, with an error code indicating the
+     *                                  cause of the failure.
+     */
+    public boolean isChainEnabled(final int chain) {
+        return isChainEnabled(mConfigurationMap, chain);
+    }
+
+    /**
+     * Get firewall rule of specified firewall chain on specified uid.
+     *
+     * @param chain target chain
+     * @param uid        target uid
+     * @return either {@link ConnectivityManager#FIREWALL_RULE_ALLOW} or
+     *         {@link ConnectivityManager#FIREWALL_RULE_DENY}.
+     * @throws UnsupportedOperationException if called on pre-T devices.
+     * @throws ServiceSpecificException in case of failure, with an error code indicating the
+     *                                  cause of the failure.
+     */
+    public int getUidRule(final int chain, final int uid) {
+        return getUidRule(mUidOwnerMap, chain, uid);
+    }
+
+    /**
+     * Get the specified firewall chain's status.
+     *
+     * @param configurationMap target configurationMap
+     * @param chain target chain
+     * @return {@code true} if chain is enabled, {@code false} if chain is not enabled.
+     * @throws UnsupportedOperationException if called on pre-T devices.
+     * @throws ServiceSpecificException in case of failure, with an error code indicating the
+     *                                  cause of the failure.
+     */
+    public static boolean isChainEnabled(
+            final IBpfMap<Struct.S32, Struct.U32> configurationMap, final int chain) {
+        throwIfPreT("isChainEnabled is not available on pre-T devices");
+
+        final long match = getMatchByFirewallChain(chain);
+        try {
+            final Struct.U32 config = configurationMap.getValue(UID_RULES_CONFIGURATION_KEY);
+            return (config.val & match) != 0;
+        } catch (ErrnoException e) {
+            throw new ServiceSpecificException(e.errno,
+                    "Unable to get firewall chain status: " + Os.strerror(e.errno));
+        }
+    }
+
+    /**
+     * Get firewall rule of specified firewall chain on specified uid.
+     *
+     * @param uidOwnerMap target uidOwnerMap.
+     * @param chain target chain.
+     * @param uid target uid.
+     * @return either FIREWALL_RULE_ALLOW or FIREWALL_RULE_DENY
+     * @throws UnsupportedOperationException if called on pre-T devices.
+     * @throws ServiceSpecificException      in case of failure, with an error code indicating the
+     *                                       cause of the failure.
+     */
+    public static int getUidRule(final IBpfMap<Struct.S32, UidOwnerValue> uidOwnerMap,
+            final int chain, final int uid) {
+        throwIfPreT("getUidRule is not available on pre-T devices");
+
+        final long match = getMatchByFirewallChain(chain);
+        final boolean isAllowList = isFirewallAllowList(chain);
+        try {
+            final UidOwnerValue uidMatch = uidOwnerMap.getValue(new Struct.S32(uid));
+            final boolean isMatchEnabled = uidMatch != null && (uidMatch.rule & match) != 0;
+            return isMatchEnabled == isAllowList ? FIREWALL_RULE_ALLOW : FIREWALL_RULE_DENY;
+        } catch (ErrnoException e) {
+            throw new ServiceSpecificException(e.errno,
+                    "Unable to get uid rule status: " + Os.strerror(e.errno));
+        }
+    }
+}
diff --git a/framework/src/android/net/BpfNetMapsUtils.java b/framework/src/android/net/BpfNetMapsUtils.java
index d464e3d..28d5891 100644
--- a/framework/src/android/net/BpfNetMapsUtils.java
+++ b/framework/src/android/net/BpfNetMapsUtils.java
@@ -39,6 +39,8 @@
 import android.os.ServiceSpecificException;
 import android.util.Pair;
 
+import com.android.modules.utils.build.SdkLevel;
+
 import java.util.StringJoiner;
 
 /**
@@ -124,4 +126,15 @@
         }
         return sj.toString();
     }
+
+    public static final boolean PRE_T = !SdkLevel.isAtLeastT();
+
+    /**
+     * Throw UnsupportedOperationException if SdkLevel is before T.
+     */
+    public static void throwIfPreT(final String msg) {
+        if (PRE_T) {
+            throw new UnsupportedOperationException(msg);
+        }
+    }
 }
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 9e879c2..915c20d 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -26,7 +26,6 @@
 import static android.net.QosCallback.QosCallbackRegistrationException;
 
 import android.annotation.CallbackExecutor;
-import android.annotation.FlaggedApi;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -116,14 +115,6 @@
     private static final String TAG = "ConnectivityManager";
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
 
-    // TODO : remove this class when udc-mainline-prod is abandoned and android.net.flags.Flags is
-    // available here
-    /** @hide */
-    public static class Flags {
-        static final String SET_DATA_SAVER_VIA_CM =
-                "com.android.net.flags.set_data_saver_via_cm";
-    }
-
     /**
      * A change in network connectivity has occurred. A default connection has either
      * been established or lost. The NetworkInfo for the affected network is
@@ -5967,28 +5958,6 @@
     }
 
     /**
-     * Sets data saver switch.
-     *
-     * @param enable True if enable.
-     * @throws IllegalStateException if failed.
-     * @hide
-     */
-    @FlaggedApi(Flags.SET_DATA_SAVER_VIA_CM)
-    @SystemApi(client = MODULE_LIBRARIES)
-    @RequiresPermission(anyOf = {
-            android.Manifest.permission.NETWORK_SETTINGS,
-            android.Manifest.permission.NETWORK_STACK,
-            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK
-    })
-    public void setDataSaverEnabled(final boolean enable) {
-        try {
-            mService.setDataSaverEnabled(enable);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-    }
-
-    /**
      * Adds the specified UID to the list of UIds that are allowed to use data on metered networks
      * even when background data is restricted. The deny list takes precedence over the allow list.
      *
diff --git a/framework/src/android/net/IConnectivityManager.aidl b/framework/src/android/net/IConnectivityManager.aidl
index 92e1ea1..fe27773 100644
--- a/framework/src/android/net/IConnectivityManager.aidl
+++ b/framework/src/android/net/IConnectivityManager.aidl
@@ -240,8 +240,6 @@
 
     void setTestAllowBadWifiUntil(long timeMs);
 
-    void setDataSaverEnabled(boolean enable);
-
     void updateMeteredNetworkAllowList(int uid, boolean add);
 
     void updateMeteredNetworkDenyList(int uid, boolean add);
diff --git a/service/src/com/android/server/UidOwnerValue.java b/framework/src/android/net/UidOwnerValue.java
similarity index 86%
rename from service/src/com/android/server/UidOwnerValue.java
rename to framework/src/android/net/UidOwnerValue.java
index d6c0e0d..e8ae604 100644
--- a/service/src/com/android/server/UidOwnerValue.java
+++ b/framework/src/android/net/UidOwnerValue.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022 The Android Open Source Project
+ * Copyright (C) 2023 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,11 +14,15 @@
  * limitations under the License.
  */
 
-package com.android.server;
+package android.net;
 
 import com.android.net.module.util.Struct;
 
-/** Value type for per uid traffic control configuration map  */
+/**
+ * Value type for per uid traffic control configuration map.
+ *
+ * @hide
+ */
 public class UidOwnerValue extends Struct {
     // Allowed interface index. Only applicable if IIF_MATCH is set in the rule bitmask below.
     @Field(order = 0, type = Type.S32)
diff --git a/service/Android.bp b/service/Android.bp
index 8e59e86..250693f 100644
--- a/service/Android.bp
+++ b/service/Android.bp
@@ -188,7 +188,6 @@
         "dnsresolver_aidl_interface-V11-java",
         "modules-utils-shell-command-handler",
         "net-utils-device-common",
-        "net-utils-device-common-bpf",
         "net-utils-device-common-ip",
         "net-utils-device-common-netlink",
         "net-utils-services-common",
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index 6a34a24..671c4ac 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -26,6 +26,7 @@
 import static android.net.BpfNetMapsConstants.UID_OWNER_MAP_PATH;
 import static android.net.BpfNetMapsConstants.UID_PERMISSION_MAP_PATH;
 import static android.net.BpfNetMapsConstants.UID_RULES_CONFIGURATION_KEY;
+import static android.net.BpfNetMapsUtils.PRE_T;
 import static android.net.BpfNetMapsUtils.getMatchByFirewallChain;
 import static android.net.BpfNetMapsUtils.matchToString;
 import static android.net.ConnectivityManager.FIREWALL_CHAIN_DOZABLE;
@@ -51,7 +52,9 @@
 
 import android.app.StatsManager;
 import android.content.Context;
+import android.net.BpfNetMapsReader;
 import android.net.INetd;
+import android.net.UidOwnerValue;
 import android.os.Build;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
@@ -92,7 +95,6 @@
  * {@hide}
  */
 public class BpfNetMaps {
-    private static final boolean PRE_T = !SdkLevel.isAtLeastT();
     static {
         if (!PRE_T) {
             System.loadLibrary("service-connectivity");
@@ -298,6 +300,7 @@
     }
 
     /** Constructor used after T that doesn't need to use netd anymore. */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public BpfNetMaps(final Context context) {
         this(context, null);
 
@@ -420,6 +423,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void addNaughtyApp(final int uid) {
         throwIfPreT("addNaughtyApp is not available on pre-T devices");
 
@@ -438,6 +442,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void removeNaughtyApp(final int uid) {
         throwIfPreT("removeNaughtyApp is not available on pre-T devices");
 
@@ -456,6 +461,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void addNiceApp(final int uid) {
         throwIfPreT("addNiceApp is not available on pre-T devices");
 
@@ -474,6 +480,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void removeNiceApp(final int uid) {
         throwIfPreT("removeNiceApp is not available on pre-T devices");
 
@@ -494,6 +501,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void setChildChain(final int childChain, final boolean enable) {
         throwIfPreT("setChildChain is not available on pre-T devices");
 
@@ -523,18 +531,14 @@
      * @throws UnsupportedOperationException if called on pre-T devices.
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
+     *
+     * @deprecated Use {@link BpfNetMapsReader#isChainEnabled} instead.
      */
+    // TODO: Migrate the callers to use {@link BpfNetMapsReader#isChainEnabled} instead.
+    @Deprecated
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public boolean isChainEnabled(final int childChain) {
-        throwIfPreT("isChainEnabled is not available on pre-T devices");
-
-        final long match = getMatchByFirewallChain(childChain);
-        try {
-            final U32 config = sConfigurationMap.getValue(UID_RULES_CONFIGURATION_KEY);
-            return (config.val & match) != 0;
-        } catch (ErrnoException e) {
-            throw new ServiceSpecificException(e.errno,
-                    "Unable to get firewall chain status: " + Os.strerror(e.errno));
-        }
+        return BpfNetMapsReader.isChainEnabled(sConfigurationMap, childChain);
     }
 
     private Set<Integer> asSet(final int[] uids) {
@@ -554,6 +558,7 @@
      * @throws UnsupportedOperationException if called on pre-T devices.
      * @throws IllegalArgumentException if {@code chain} is not a valid chain.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void replaceUidChain(final int chain, final int[] uids) {
         throwIfPreT("replaceUidChain is not available on pre-T devices");
 
@@ -638,6 +643,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void setUidRule(final int childChain, final int uid, final int firewallRule) {
         throwIfPreT("setUidRule is not available on pre-T devices");
 
@@ -667,20 +673,12 @@
      * @throws UnsupportedOperationException if called on pre-T devices.
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
+     *
+     * @deprecated use {@link BpfNetMapsReader#getUidRule} instead.
      */
+    // TODO: Migrate the callers to use {@link BpfNetMapsReader#getUidRule} instead.
     public int getUidRule(final int childChain, final int uid) {
-        throwIfPreT("isUidChainEnabled is not available on pre-T devices");
-
-        final long match = getMatchByFirewallChain(childChain);
-        final boolean isAllowList = isFirewallAllowList(childChain);
-        try {
-            final UidOwnerValue uidMatch = sUidOwnerMap.getValue(new S32(uid));
-            final boolean isMatchEnabled = uidMatch != null && (uidMatch.rule & match) != 0;
-            return isMatchEnabled == isAllowList ? FIREWALL_RULE_ALLOW : FIREWALL_RULE_DENY;
-        } catch (ErrnoException e) {
-            throw new ServiceSpecificException(e.errno,
-                    "Unable to get uid rule status: " + Os.strerror(e.errno));
-        }
+        return BpfNetMapsReader.getUidRule(sUidOwnerMap, childChain, uid);
     }
 
     private Set<Integer> getUidsMatchEnabled(final int childChain) throws ErrnoException {
@@ -830,6 +828,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void updateUidLockdownRule(final int uid, final boolean add) {
         throwIfPreT("updateUidLockdownRule is not available on pre-T devices");
 
@@ -852,6 +851,7 @@
      * @throws ServiceSpecificException in case of failure, with an error code indicating the
      *                                  cause of the failure.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void swapActiveStatsMap() {
         throwIfPreT("swapActiveStatsMap is not available on pre-T devices");
 
@@ -927,6 +927,7 @@
     }
 
     /** Register callback for statsd to pull atom. */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void setPullAtomCallback(final Context context) {
         throwIfPreT("setPullAtomCallback is not available on pre-T devices");
 
@@ -1016,6 +1017,7 @@
      * @throws IOException when file descriptor is invalid.
      * @throws ServiceSpecificException when the method is called on an unsupported device.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void dump(final IndentingPrintWriter pw, final FileDescriptor fd, boolean verbose)
             throws IOException, ServiceSpecificException {
         if (PRE_T) {
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index ada5860..2797b47 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -12587,20 +12587,6 @@
     }
 
     @Override
-    public void setDataSaverEnabled(final boolean enable) {
-        enforceNetworkStackOrSettingsPermission();
-        try {
-            final boolean ret = mNetd.bandwidthEnableDataSaver(enable);
-            if (!ret) {
-                throw new IllegalStateException("Error when changing iptables: " + enable);
-            }
-        } catch (RemoteException e) {
-            // Lack of permission or binder errors.
-            throw new IllegalStateException(e);
-        }
-    }
-
-    @Override
     public void updateMeteredNetworkAllowList(final int uid, final boolean add) {
         enforceNetworkStackOrSettingsPermission();
 
diff --git a/staticlibs/Android.bp b/staticlibs/Android.bp
index 621759e..0bcb757 100644
--- a/staticlibs/Android.bp
+++ b/staticlibs/Android.bp
@@ -181,6 +181,8 @@
     },
 }
 
+// The net-utils-device-common-netlink library requires the callers to contain
+// net-utils-device-common-struct.
 java_library {
     name: "net-utils-device-common-netlink",
     srcs: [
@@ -192,12 +194,13 @@
         "//packages/modules/Connectivity:__subpackages__",
         "//packages/modules/NetworkStack:__subpackages__",
     ],
-    static_libs: [
-        "net-utils-device-common-struct",
-    ],
     libs: [
         "androidx.annotation_annotation",
         "framework-connectivity.stubs.module_lib",
+        // For libraries which are statically linked in framework-connectivity, do not
+        // statically link here because callers of this library might already have a static
+        // version linked.
+        "net-utils-device-common-struct",
     ],
     apex_available: [
         "com.android.tethering",
@@ -209,6 +212,8 @@
     },
 }
 
+// The net-utils-device-common-ip library requires the callers to contain
+// net-utils-device-common-struct.
 java_library {
     // TODO : this target should probably be folded into net-utils-device-common
     name: "net-utils-device-common-ip",
diff --git a/staticlibs/testutils/Android.bp b/staticlibs/testutils/Android.bp
index 5fe7ac3..a5e5afb 100644
--- a/staticlibs/testutils/Android.bp
+++ b/staticlibs/testutils/Android.bp
@@ -38,6 +38,7 @@
         "net-utils-device-common",
         "net-utils-device-common-async",
         "net-utils-device-common-netlink",
+        "net-utils-device-common-struct",
         "net-utils-device-common-wear",
         "modules-utils-build_system",
     ],
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/TestBpfMap.java b/staticlibs/testutils/devicetests/com/android/testutils/TestBpfMap.java
index 733bd98..70f20d6 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/TestBpfMap.java
+++ b/staticlibs/testutils/devicetests/com/android/testutils/TestBpfMap.java
@@ -43,6 +43,8 @@
 public class TestBpfMap<K extends Struct, V extends Struct> implements IBpfMap<K, V> {
     private final ConcurrentHashMap<K, V> mMap = new ConcurrentHashMap<>();
 
+    public TestBpfMap() {}
+
     // TODO: Remove this constructor
     public TestBpfMap(final Class<K> key, final Class<V> value) {
     }
diff --git a/staticlibs/testutils/host/com/android/testutils/ConnectivityTestTargetPreparer.kt b/staticlibs/testutils/host/com/android/testutils/ConnectivityTestTargetPreparer.kt
index eb94781..600a623 100644
--- a/staticlibs/testutils/host/com/android/testutils/ConnectivityTestTargetPreparer.kt
+++ b/staticlibs/testutils/host/com/android/testutils/ConnectivityTestTargetPreparer.kt
@@ -128,7 +128,7 @@
         if (testInfo.device.getApiLevel() < 31) return
         testInfo.exec("cmd connectivity set-chain3-enabled $enableChain")
         enablePkgs.forEach { (pkg, allow) ->
-            testInfo.exec("cmd connectivity set-package-networking-enabled $pkg $allow")
+            testInfo.exec("cmd connectivity set-package-networking-enabled $allow $pkg")
         }
     }
 
diff --git a/tests/integration/Android.bp b/tests/integration/Android.bp
index 12919ae..f705e34 100644
--- a/tests/integration/Android.bp
+++ b/tests/integration/Android.bp
@@ -45,6 +45,7 @@
         // order-dependent setup.
         "NetworkStackApiStableLib",
         "androidx.test.ext.junit",
+        "compatibility-device-util-axt",
         "frameworks-net-integration-testutils",
         "kotlin-reflect",
         "mockito-target-extended-minus-junit4",
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 e264b55..9b082a4 100644
--- a/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
+++ b/tests/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
@@ -40,11 +40,14 @@
 import android.os.IBinder
 import android.os.SystemConfigManager
 import android.os.UserHandle
+import android.os.VintfRuntimeInfo
 import android.testing.TestableContext
 import android.util.Log
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.compatibility.common.util.SystemUtil
 import com.android.connectivity.resources.R
+import com.android.net.module.util.BpfUtils
 import com.android.server.BpfNetMaps
 import com.android.server.ConnectivityService
 import com.android.server.NetworkAgentWrapper
@@ -53,6 +56,7 @@
 import com.android.server.connectivity.MockableSystemProperties
 import com.android.server.connectivity.MultinetworkPolicyTracker
 import com.android.server.connectivity.ProxyTracker
+import com.android.testutils.DeviceInfoUtils
 import com.android.testutils.RecorderCallback.CallbackEntry.LinkPropertiesChanged
 import com.android.testutils.TestableNetworkCallback
 import kotlin.test.assertEquals
@@ -60,6 +64,7 @@
 import kotlin.test.assertTrue
 import kotlin.test.fail
 import org.junit.After
+import org.junit.Assume
 import org.junit.Before
 import org.junit.BeforeClass
 import org.junit.Test
@@ -302,4 +307,25 @@
                     !it.hasCapability(NET_CAPABILITY_VALIDATED)
         }
     }
+
+    private fun isBpfGetCgroupProgramIdSupportedByKernel(): Boolean {
+        val kVersionString = VintfRuntimeInfo.getKernelVersion()
+        return DeviceInfoUtils.compareMajorMinorVersion(kVersionString, "4.19") >= 0
+    }
+
+    @Test
+    fun testBpfProgramAttachStatus() {
+        Assume.assumeTrue(isBpfGetCgroupProgramIdSupportedByKernel())
+
+        listOf(
+                BpfUtils.BPF_CGROUP_INET_INGRESS,
+                BpfUtils.BPF_CGROUP_INET_EGRESS,
+                BpfUtils.BPF_CGROUP_INET_SOCK_CREATE
+        ).forEach {
+            val ret = SystemUtil.runShellCommand(InstrumentationRegistry.getInstrumentation(),
+                    "cmd connectivity bpf-get-cgroup-program-id $it").trim()
+
+            assertTrue(Integer.parseInt(ret) > 0, "Unexpected output $ret for type $it")
+        }
+    }
 }
diff --git a/tests/unit/java/android/net/BpfNetMapsReaderTest.kt b/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
new file mode 100644
index 0000000..facb932
--- /dev/null
+++ b/tests/unit/java/android/net/BpfNetMapsReaderTest.kt
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net
+
+import android.net.BpfNetMapsConstants.UID_RULES_CONFIGURATION_KEY
+import android.net.BpfNetMapsUtils.getMatchByFirewallChain
+import android.os.Build
+import com.android.net.module.util.IBpfMap
+import com.android.net.module.util.Struct.S32
+import com.android.net.module.util.Struct.U32
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import com.android.testutils.DevSdkIgnoreRunner
+import com.android.testutils.TestBpfMap
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+
+// pre-T devices does not support Bpf.
+@RunWith(DevSdkIgnoreRunner::class)
+@IgnoreUpTo(Build.VERSION_CODES.S_V2)
+class BpfNetMapsReaderTest {
+    private val testConfigurationMap: IBpfMap<S32, U32> = TestBpfMap()
+    private val testUidOwnerMap: IBpfMap<S32, UidOwnerValue> = TestBpfMap()
+    private val bpfNetMapsReader = BpfNetMapsReader(
+        TestDependencies(testConfigurationMap, testUidOwnerMap))
+
+    class TestDependencies(
+        private val configMap: IBpfMap<S32, U32>,
+        private val uidOwnerMap: IBpfMap<S32, UidOwnerValue>
+    ) : BpfNetMapsReader.Dependencies() {
+        override fun getConfigurationMap() = configMap
+        override fun getUidOwnerMap() = uidOwnerMap
+    }
+
+    private fun doTestIsChainEnabled(chain: Int) {
+        testConfigurationMap.updateEntry(
+            UID_RULES_CONFIGURATION_KEY,
+            U32(getMatchByFirewallChain(chain))
+        )
+        assertTrue(bpfNetMapsReader.isChainEnabled(chain))
+        testConfigurationMap.updateEntry(UID_RULES_CONFIGURATION_KEY, U32(0))
+        assertFalse(bpfNetMapsReader.isChainEnabled(chain))
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testIsChainEnabled() {
+        doTestIsChainEnabled(ConnectivityManager.FIREWALL_CHAIN_DOZABLE)
+        doTestIsChainEnabled(ConnectivityManager.FIREWALL_CHAIN_STANDBY)
+        doTestIsChainEnabled(ConnectivityManager.FIREWALL_CHAIN_POWERSAVE)
+        doTestIsChainEnabled(ConnectivityManager.FIREWALL_CHAIN_RESTRICTED)
+        doTestIsChainEnabled(ConnectivityManager.FIREWALL_CHAIN_LOW_POWER_STANDBY)
+    }
+}
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index 5f280c6..da5f7e1 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -66,6 +66,7 @@
 import android.content.Context;
 import android.net.BpfNetMapsUtils;
 import android.net.INetd;
+import android.net.UidOwnerValue;
 import android.os.Build;
 import android.os.ServiceSpecificException;
 import android.system.ErrnoException;
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 8bca4dd..194cec3 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -935,12 +935,6 @@
         return appUid + (firstSdkSandboxUid - Process.FIRST_APPLICATION_UID);
     }
 
-    // This function assumes the UID range for user 0 ([1, 99999])
-    private static UidRangeParcel[] uidRangeParcelsExcludingUids(Integer... excludedUids) {
-        final List<Integer> uids = Arrays.asList(excludedUids);
-        return intToUidRangeStableParcels(intRangesPrimaryExcludingUids(uids));
-    }
-
     // Create the list of ranges for the primary user (User 0), excluding excludedUids.
     private static List<Range<Integer>> intRangesPrimaryExcludingUids(List<Integer> excludedUids) {
         final List<Integer> excludedUidsList = new ArrayList<>(excludedUids);
@@ -9527,8 +9521,16 @@
         assertNotNull(mCm.getActiveNetworkForUid(restrictedUid));
 
         // Enable always-on VPN lockdown. The main user loses network access because no VPN is up.
-        final ArrayList<String> allowList = new ArrayList<>();
-        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
+        // Coverage in VpnTest.
+        final List<Integer> excludedUids = new ArrayList<>();
+        excludedUids.add(VPN_UID);
+        if (mDeps.isAtLeastT()) {
+            // On T onwards, the corresponding SDK sandbox UID should also be excluded
+            excludedUids.add(toSdkSandboxUid(VPN_UID));
+        }
+        final List<Range<Integer>> primaryRanges = intRangesPrimaryExcludingUids(excludedUids);
+        mCm.setRequireVpnForUids(true, primaryRanges);
+
         waitForIdle();
         assertNull(mCm.getActiveNetworkForUid(uid));
         // This is arguably overspecified: a UID that is not running doesn't have an active network.
@@ -9540,12 +9542,6 @@
         // TODO: check that VPN app within restricted profile still has access, etc.
         // Add a restricted user.
         // This is equivalent to `mMockVpn.onUserAdded(RESTRICTED_USER);`, coverage in VpnTest.
-        final List<Integer> excludedUids = new ArrayList<Integer>();
-        excludedUids.add(VPN_UID);
-        if (mDeps.isAtLeastT()) {
-            // On T onwards, the corresponding SDK sandbox UID should also be excluded
-            excludedUids.add(toSdkSandboxUid(VPN_UID));
-        }
         final List<Range<Integer>> restrictedRanges =
                 intRangesExcludingUids(RESTRICTED_USER, excludedUids);
         mCm.setRequireVpnForUids(true, restrictedRanges);
@@ -9563,7 +9559,8 @@
         assertNull(mCm.getActiveNetworkForUid(uid));
         assertNotNull(mCm.getActiveNetworkForUid(restrictedUid));
 
-        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
+        mCm.setRequireVpnForUids(false, primaryRanges);
+
         waitForIdle();
     }
 
@@ -10016,18 +10013,20 @@
                 new Handler(ConnectivityThread.getInstanceLooper()));
 
         final int uid = Process.myUid();
-        final ArrayList<String> allowList = new ArrayList<>();
-        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
-        waitForIdle();
 
-        final Set<Integer> excludedUids = new ArraySet<Integer>();
+        // Enable always-on VPN lockdown, coverage in VpnTest.
+        final List<Integer> excludedUids = new ArrayList<Integer>();
         excludedUids.add(VPN_UID);
         if (mDeps.isAtLeastT()) {
             // On T onwards, the corresponding SDK sandbox UID should also be excluded
             excludedUids.add(toSdkSandboxUid(VPN_UID));
         }
-        final UidRangeParcel[] uidRangeParcels = uidRangeParcelsExcludingUids(
-                excludedUids.toArray(new Integer[0]));
+
+        final List<Range<Integer>> primaryRanges = intRangesPrimaryExcludingUids(excludedUids);
+        mCm.setRequireVpnForUids(true, primaryRanges);
+        waitForIdle();
+
+        final UidRangeParcel[] uidRangeParcels = intToUidRangeStableParcels(primaryRanges);
         InOrder inOrder = inOrder(mMockNetd);
         expectNetworkRejectNonSecureVpn(inOrder, true, uidRangeParcels);
 
@@ -10047,7 +10046,8 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
 
         // Disable lockdown, expect to see the network unblocked.
-        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
+        mCm.setRequireVpnForUids(false, primaryRanges);
+        waitForIdle();
         callback.expect(BLOCKED_STATUS, mWiFiAgent, cb -> !cb.getBlocked());
         defaultCallback.expect(BLOCKED_STATUS, mWiFiAgent, cb -> !cb.getBlocked());
         vpnUidCallback.assertNoCallback();
@@ -10060,22 +10060,25 @@
         assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
         assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
 
-        // Add our UID to the allowlist and re-enable lockdown, expect network is not blocked.
-        allowList.add(TEST_PACKAGE_NAME);
-        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
+        // Add our UID to the allowlist, expect network is not blocked. Coverage in VpnTest.
+        excludedUids.add(uid);
+        if (mDeps.isAtLeastT()) {
+            // On T onwards, the corresponding SDK sandbox UID should also be excluded
+            excludedUids.add(toSdkSandboxUid(uid));
+        }
+        final List<Range<Integer>> primaryRangesExcludingUid =
+                intRangesPrimaryExcludingUids(excludedUids);
+        mCm.setRequireVpnForUids(true, primaryRangesExcludingUid);
+        waitForIdle();
+
         callback.assertNoCallback();
         defaultCallback.assertNoCallback();
         vpnUidCallback.assertNoCallback();
         vpnUidDefaultCallback.assertNoCallback();
         vpnDefaultCallbackAsUid.assertNoCallback();
 
-        excludedUids.add(uid);
-        if (mDeps.isAtLeastT()) {
-            // On T onwards, the corresponding SDK sandbox UID should also be excluded
-            excludedUids.add(toSdkSandboxUid(uid));
-        }
-        final UidRangeParcel[] uidRangeParcelsAlsoExcludingUs = uidRangeParcelsExcludingUids(
-                excludedUids.toArray(new Integer[0]));
+        final UidRangeParcel[] uidRangeParcelsAlsoExcludingUs =
+                intToUidRangeStableParcels(primaryRangesExcludingUid);
         expectNetworkRejectNonSecureVpn(inOrder, true, uidRangeParcelsAlsoExcludingUs);
         assertEquals(mWiFiAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
         assertEquals(mWiFiAgent.getNetwork(), mCm.getActiveNetwork());
@@ -10098,15 +10101,15 @@
         assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
         assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
 
-        // Disable lockdown, remove our UID from the allowlist, and re-enable lockdown.
-        // Everything should now be blocked.
-        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
+        // Disable lockdown
+        mCm.setRequireVpnForUids(false, primaryRangesExcludingUid);
         waitForIdle();
         expectNetworkRejectNonSecureVpn(inOrder, false, uidRangeParcelsAlsoExcludingUs);
-        allowList.clear();
-        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
+        // Remove our UID from the allowlist, and re-enable lockdown.
+        mCm.setRequireVpnForUids(true, primaryRanges);
         waitForIdle();
         expectNetworkRejectNonSecureVpn(inOrder, true, uidRangeParcels);
+        // Everything should now be blocked.
         defaultCallback.expect(BLOCKED_STATUS, mWiFiAgent, cb -> cb.getBlocked());
         assertBlockedCallbackInAnyOrder(callback, true, mWiFiAgent, mCellAgent);
         vpnUidCallback.assertNoCallback();
@@ -10119,7 +10122,7 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
 
         // Disable lockdown. Everything is unblocked.
-        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
+        mCm.setRequireVpnForUids(false, primaryRanges);
         defaultCallback.expect(BLOCKED_STATUS, mWiFiAgent, cb -> !cb.getBlocked());
         assertBlockedCallbackInAnyOrder(callback, false, mWiFiAgent, mCellAgent);
         vpnUidCallback.assertNoCallback();
@@ -10132,7 +10135,7 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
 
         // Enable lockdown and connect a VPN. The VPN is not blocked.
-        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
+        mCm.setRequireVpnForUids(true, primaryRanges);
         defaultCallback.expect(BLOCKED_STATUS, mWiFiAgent, cb -> cb.getBlocked());
         assertBlockedCallbackInAnyOrder(callback, true, mWiFiAgent, mCellAgent);
         vpnUidCallback.assertNoCallback();