diff --git a/OWNERS_core_networking b/OWNERS_core_networking
index 83f798a..6d8ed4a 100644
--- a/OWNERS_core_networking
+++ b/OWNERS_core_networking
@@ -1,10 +1,6 @@
-chiachangwang@google.com
-cken@google.com
 jchalard@google.com
 junyulai@google.com
-lifr@google.com
 lorenzo@google.com
-markchien@google.com
 martinwu@google.com
 maze@google.com
 motomuman@google.com
@@ -12,7 +8,5 @@
 prohr@google.com
 reminv@google.com
 satk@google.com
-waynema@google.com
 xiaom@google.com
-yumike@google.com
 yuyanghuang@google.com
diff --git a/TEST_MAPPING b/TEST_MAPPING
index ab3ed66..d8d4c21 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -246,6 +246,9 @@
         },
         {
           "exclude-annotation": "com.android.testutils.DnsResolverModuleTest"
+        },
+        {
+          "exclude-annotation": "com.android.testutils.NetworkStackModuleTest"
         }
       ]
     },
diff --git a/Tethering/Android.bp b/Tethering/Android.bp
index e4e6c70..19bcff9 100644
--- a/Tethering/Android.bp
+++ b/Tethering/Android.bp
@@ -74,6 +74,8 @@
         "net-utils-device-common-bpf",
         "net-utils-device-common-ip",
         "net-utils-device-common-netlink",
+        "net-utils-device-common-struct",
+        "net-utils-device-common-struct-base",
         "netd-client",
         "tetheringstatsprotos",
     ],
@@ -98,7 +100,6 @@
     ],
     static_libs: [
         "NetworkStackApiCurrentShims",
-        "net-utils-device-common-struct",
     ],
     apex_available: ["com.android.tethering"],
     lint: {
@@ -115,7 +116,6 @@
     ],
     static_libs: [
         "NetworkStackApiStableShims",
-        "net-utils-device-common-struct",
     ],
     apex_available: ["com.android.tethering"],
     lint: {
diff --git a/Tethering/src/android/net/ip/IpServer.java b/Tethering/src/android/net/ip/IpServer.java
index 544ba01..9e0c970 100644
--- a/Tethering/src/android/net/ip/IpServer.java
+++ b/Tethering/src/android/net/ip/IpServer.java
@@ -137,8 +137,8 @@
     private static final MacAddress NULL_MAC_ADDRESS = MacAddress.fromString("00:00:00:00:00:00");
 
     private static final String TAG = "IpServer";
-    private static final boolean DBG = false;
-    private static final boolean VDBG = false;
+    private static final boolean DBG = Log.isLoggable(TAG, Log.DEBUG);
+    private static final boolean VDBG = Log.isLoggable(TAG, Log.VERBOSE);
     private static final Class[] sMessageClasses = {
             IpServer.class
     };
diff --git a/Tethering/src/com/android/networkstack/tethering/Tethering.java b/Tethering/src/com/android/networkstack/tethering/Tethering.java
index 873961a..d85d92f 100644
--- a/Tethering/src/com/android/networkstack/tethering/Tethering.java
+++ b/Tethering/src/com/android/networkstack/tethering/Tethering.java
@@ -173,8 +173,8 @@
 public class Tethering {
 
     private static final String TAG = Tethering.class.getSimpleName();
-    private static final boolean DBG = false;
-    private static final boolean VDBG = false;
+    private static final boolean DBG = Log.isLoggable(TAG, Log.DEBUG);
+    private static final boolean VDBG = Log.isLoggable(TAG, Log.VERBOSE);
 
     private static final Class[] sMessageClasses = {
             Tethering.class, TetherMainSM.class, IpServer.class
@@ -241,9 +241,6 @@
     private final TetherMainSM mTetherMainSM;
     private final OffloadController mOffloadController;
     private final UpstreamNetworkMonitor mUpstreamNetworkMonitor;
-    // TODO: Figure out how to merge this and other downstream-tracking objects
-    // into a single coherent structure.
-    private final HashSet<IpServer> mForwardedDownstreams;
     private final VersionedBroadcastListener mCarrierConfigChange;
     private final TetheringDependencies mDeps;
     private final EntitlementManager mEntitlementMgr;
@@ -271,8 +268,6 @@
 
     private boolean mRndisEnabled;       // track the RNDIS function enabled state
     private boolean mNcmEnabled;         // track the NCM function enabled state
-    // True iff. WiFi tethering should be started when soft AP is ready.
-    private boolean mWifiTetherRequested;
     private Network mTetherUpstream;
     private TetherStatesParcel mTetherStatesParcel;
     private boolean mDataSaverEnabled = false;
@@ -329,7 +324,6 @@
                 (what, obj) -> {
                     mTetherMainSM.sendMessage(TetherMainSM.EVENT_UPSTREAM_CALLBACK, what, 0, obj);
                 });
-        mForwardedDownstreams = new HashSet<>();
 
         IntentFilter filter = new IntentFilter();
         filter.addAction(ACTION_CARRIER_CONFIG_CHANGED);
@@ -763,7 +757,6 @@
             }
             if ((enable && mgr.startTetheredHotspot(null /* use existing softap config */))
                     || (!enable && mgr.stopSoftAp())) {
-                mWifiTetherRequested = enable;
                 return TETHER_ERROR_NO_ERROR;
             }
         } finally {
@@ -1470,10 +1463,6 @@
     }
 
     private void disableWifiIpServing(String ifname, int apState) {
-        // Regardless of whether we requested this transition, the AP has gone
-        // down.  Don't try to tether again unless we're requested to do so.
-        mWifiTetherRequested = false;
-
         mLog.log("Canceling WiFi tethering request - interface=" + ifname + " state=" + apState);
 
         disableWifiIpServingCommon(TETHERING_WIFI, ifname);
@@ -1505,8 +1494,7 @@
     private void enableWifiIpServing(String ifname, int wifiIpMode) {
         mLog.log("request WiFi tethering - interface=" + ifname + " state=" + wifiIpMode);
 
-        // Map wifiIpMode values to IpServer.Callback serving states, inferring
-        // from mWifiTetherRequested as a final "best guess".
+        // Map wifiIpMode values to IpServer.Callback serving states.
         final int ipServingMode;
         switch (wifiIpMode) {
             case IFACE_IP_MODE_TETHERED:
@@ -1653,11 +1641,6 @@
         mLog.log(state.getName() + " got " + sMagicDecoderRing.get(what, Integer.toString(what)));
     }
 
-    private boolean upstreamWanted() {
-        if (!mForwardedDownstreams.isEmpty()) return true;
-        return mWifiTetherRequested;
-    }
-
     // Needed because the canonical source of upstream truth is just the
     // upstream interface set, |mCurrentUpstreamIfaceSet|.
     private boolean pertainsToCurrentUpstream(UpstreamNetworkState ns) {
@@ -1715,12 +1698,16 @@
         private final ArrayList<IpServer> mNotifyList;
         private final IPv6TetheringCoordinator mIPv6TetheringCoordinator;
         private final OffloadWrapper mOffload;
+        // TODO: Figure out how to merge this and other downstream-tracking objects
+        // into a single coherent structure.
+        private final HashSet<IpServer> mForwardedDownstreams;
 
         private static final int UPSTREAM_SETTLE_TIME_MS     = 10000;
 
         TetherMainSM(String name, Looper looper, TetheringDependencies deps) {
             super(name, looper);
 
+            mForwardedDownstreams = new HashSet<>();
             mInitialState = new InitialState();
             mTetherModeAliveState = new TetherModeAliveState();
             mSetIpForwardingEnabledErrorState = new SetIpForwardingEnabledErrorState();
@@ -2056,6 +2043,10 @@
             }
         }
 
+        private boolean upstreamWanted() {
+            return !mForwardedDownstreams.isEmpty();
+        }
+
         class TetherModeAliveState extends State {
             boolean mUpstreamWanted = false;
             boolean mTryCell = true;
@@ -2393,6 +2384,9 @@
                 hasCallingPermission(NETWORK_SETTINGS)
                         || hasCallingPermission(PERMISSION_MAINLINE_NETWORK_STACK)
                         || hasCallingPermission(NETWORK_STACK);
+        if (callback == null) {
+            throw new NullPointerException();
+        }
         mHandler.post(() -> {
             mTetheringEventCallbacks.register(callback, new CallbackCookie(hasListPermission));
             final TetheringCallbackStartedParcel parcel = new TetheringCallbackStartedParcel();
@@ -2651,7 +2645,7 @@
             }
             pw.println(" - lastError = " + tetherState.lastError);
         }
-        pw.println("Upstream wanted: " + upstreamWanted());
+        pw.println("Upstream wanted: " + mTetherMainSM.upstreamWanted());
         pw.println("Current upstream interface(s): " + mCurrentUpstreamIfaceSet);
         pw.decreaseIndent();
 
diff --git a/Tethering/src/com/android/networkstack/tethering/TetheringDependencies.java b/Tethering/src/com/android/networkstack/tethering/TetheringDependencies.java
index 9dfd225..3f86056 100644
--- a/Tethering/src/com/android/networkstack/tethering/TetheringDependencies.java
+++ b/Tethering/src/com/android/networkstack/tethering/TetheringDependencies.java
@@ -156,6 +156,7 @@
     /**
      * Get a reference to BluetoothAdapter to be used by tethering.
      */
+    @Nullable
     public abstract BluetoothAdapter getBluetoothAdapter();
 
     /**
diff --git a/Tethering/src/com/android/networkstack/tethering/TetheringService.java b/Tethering/src/com/android/networkstack/tethering/TetheringService.java
index aa73819..623f502 100644
--- a/Tethering/src/com/android/networkstack/tethering/TetheringService.java
+++ b/Tethering/src/com/android/networkstack/tethering/TetheringService.java
@@ -30,6 +30,7 @@
 
 import android.app.Service;
 import android.bluetooth.BluetoothAdapter;
+import android.bluetooth.BluetoothManager;
 import android.content.Context;
 import android.content.Intent;
 import android.net.IIntResultListener;
@@ -377,7 +378,11 @@
 
             @Override
             public BluetoothAdapter getBluetoothAdapter() {
-                return BluetoothAdapter.getDefaultAdapter();
+                final BluetoothManager btManager = getSystemService(BluetoothManager.class);
+                if (btManager == null) {
+                    return null;
+                }
+                return btManager.getAdapter();
             }
         };
     }
diff --git a/Tethering/tests/integration/Android.bp b/Tethering/tests/integration/Android.bp
index 07fa733..337d408 100644
--- a/Tethering/tests/integration/Android.bp
+++ b/Tethering/tests/integration/Android.bp
@@ -33,6 +33,7 @@
         "net-tests-utils",
         "net-utils-device-common",
         "net-utils-device-common-bpf",
+        "net-utils-device-common-struct-base",
         "testables",
         "connectivity-net-module-utils-bpf",
     ],
diff --git a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
index f696885..120b871 100644
--- a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
+++ b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
@@ -74,7 +74,6 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.BeforeClass;
 
 import java.io.FileDescriptor;
 import java.net.Inet4Address;
@@ -144,7 +143,7 @@
 
     private static final Context sContext =
             InstrumentationRegistry.getInstrumentation().getContext();
-    private static final EthernetManager sEm = sContext.getSystemService(EthernetManager.class);
+    protected static final EthernetManager sEm = sContext.getSystemService(EthernetManager.class);
     private static final TetheringManager sTm = sContext.getSystemService(TetheringManager.class);
     private static final PackageManager sPackageManager = sContext.getPackageManager();
     private static final CtsNetUtils sCtsNetUtils = new CtsNetUtils(sContext);
@@ -168,34 +167,6 @@
         return sContext;
     }
 
-    @BeforeClass
-    public static void setUpOnce() throws Exception {
-        // The first test case may experience tethering restart with IP conflict handling.
-        // Tethering would cache the last upstreams so that the next enabled tethering avoids
-        // picking up the address that is in conflict with the upstreams. To protect subsequent
-        // tests, turn tethering on and off before running them.
-        MyTetheringEventCallback callback = null;
-        TestNetworkInterface testIface = null;
-        assumeTrue(sEm != null);
-        try {
-            // If the physical ethernet interface is available, do nothing.
-            if (isInterfaceForTetheringAvailable()) return;
-
-            testIface = createTestInterface();
-            setIncludeTestInterfaces(true);
-
-            callback = enableEthernetTethering(testIface.getInterfaceName(), null);
-            callback.awaitUpstreamChanged(true /* throwTimeoutException */);
-        } catch (TimeoutException e) {
-            Log.d(TAG, "WARNNING " + e);
-        } finally {
-            maybeCloseTestInterface(testIface);
-            maybeUnregisterTetheringEventCallback(callback);
-
-            setIncludeTestInterfaces(false);
-        }
-    }
-
     @Before
     public void setUp() throws Exception {
         mHandlerThread = new HandlerThread(getClass().getSimpleName());
diff --git a/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java b/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
index 4949eaa..c54d1b4 100644
--- a/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
+++ b/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
@@ -59,6 +59,7 @@
 import com.android.testutils.NetworkStackModuleTest;
 import com.android.testutils.TapPacketReader;
 
+import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -150,6 +151,35 @@
             (byte) 0x01, (byte) 0x02, (byte) 0x03, (byte) 0x04  /* Address: 1.2.3.4 */
     };
 
+    /** Enable/disable tethering once before running the tests. */
+    @BeforeClass
+    public static void setUpOnce() throws Exception {
+        // The first test case may experience tethering restart with IP conflict handling.
+        // Tethering would cache the last upstreams so that the next enabled tethering avoids
+        // picking up the address that is in conflict with the upstreams. To protect subsequent
+        // tests, turn tethering on and off before running them.
+        MyTetheringEventCallback callback = null;
+        TestNetworkInterface testIface = null;
+        assumeTrue(sEm != null);
+        try {
+            // If the physical ethernet interface is available, do nothing.
+            if (isInterfaceForTetheringAvailable()) return;
+
+            testIface = createTestInterface();
+            setIncludeTestInterfaces(true);
+
+            callback = enableEthernetTethering(testIface.getInterfaceName(), null);
+            callback.awaitUpstreamChanged(true /* throwTimeoutException */);
+        } catch (TimeoutException e) {
+            Log.d(TAG, "WARNNING " + e);
+        } finally {
+            maybeCloseTestInterface(testIface);
+            maybeUnregisterTetheringEventCallback(callback);
+
+            setIncludeTestInterfaces(false);
+        }
+    }
+
     @Test
     public void testVirtualEthernetAlreadyExists() throws Exception {
         // This test requires manipulating packets. Skip if there is a physical Ethernet connected.
diff --git a/Tethering/tests/mts/Android.bp b/Tethering/tests/mts/Android.bp
index a80e49e..c4d5636 100644
--- a/Tethering/tests/mts/Android.bp
+++ b/Tethering/tests/mts/Android.bp
@@ -45,6 +45,7 @@
         "junit-params",
         "connectivity-net-module-utils-bpf",
         "net-utils-device-common-bpf",
+        "net-utils-device-common-struct-base",
     ],
 
     jni_libs: [
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 750bfce..9f430af 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringTest.java
@@ -3622,6 +3622,42 @@
                 InetAddresses.parseNumericAddress(ifaceConfig.ipv4Addr), ifaceConfig.prefixLength);
         assertFalse(sapPrefix.equals(lohsPrefix));
     }
+
+    @Test
+    public void testWifiTetheringWhenP2pActive() throws Exception {
+        initTetheringOnTestThread();
+        // Enable wifi P2P.
+        sendWifiP2pConnectionChanged(true, true, TEST_P2P_IFNAME);
+        verifyInterfaceServingModeStarted(TEST_P2P_IFNAME);
+        verifyTetheringBroadcast(TEST_P2P_IFNAME, EXTRA_AVAILABLE_TETHER);
+        verifyTetheringBroadcast(TEST_P2P_IFNAME, EXTRA_ACTIVE_LOCAL_ONLY);
+        verify(mUpstreamNetworkMonitor).startObserveAllNetworks();
+        // Verify never enable upstream if only P2P active.
+        verify(mUpstreamNetworkMonitor, never()).setTryCell(true);
+        assertEquals(TETHER_ERROR_NO_ERROR, mTethering.getLastErrorForTest(TEST_P2P_IFNAME));
+
+        when(mWifiManager.startTetheredHotspot(any())).thenReturn(true);
+        // Emulate pressing the WiFi tethering button.
+        mTethering.startTethering(createTetheringRequestParcel(TETHERING_WIFI), TEST_CALLER_PKG,
+                null);
+        mLooper.dispatchAll();
+        verify(mWifiManager).startTetheredHotspot(null);
+        verifyNoMoreInteractions(mWifiManager);
+
+        mTethering.interfaceStatusChanged(TEST_WLAN_IFNAME, true);
+        sendWifiApStateChanged(WIFI_AP_STATE_ENABLED, TEST_WLAN_IFNAME, IFACE_IP_MODE_TETHERED);
+
+        verifyTetheringBroadcast(TEST_WLAN_IFNAME, EXTRA_AVAILABLE_TETHER);
+        verify(mWifiManager).updateInterfaceIpState(
+                TEST_WLAN_IFNAME, WifiManager.IFACE_IP_MODE_UNSPECIFIED);
+
+        verify(mWifiManager).updateInterfaceIpState(TEST_WLAN_IFNAME, IFACE_IP_MODE_TETHERED);
+        verifyNoMoreInteractions(mWifiManager);
+        verifyTetheringBroadcast(TEST_WLAN_IFNAME, EXTRA_ACTIVE_TETHER);
+
+        verify(mUpstreamNetworkMonitor).setTryCell(true);
+    }
+
     // TODO: Test that a request for hotspot mode doesn't interfere with an
     // already operating tethering mode interface.
 }
diff --git a/common/Android.bp b/common/Android.bp
index 0048a0a..5fabf41 100644
--- a/common/Android.bp
+++ b/common/Android.bp
@@ -26,7 +26,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.
+// The library requires the final artifact to contain net-utils-device-common-struct-base.
 java_library {
     name: "connectivity-net-module-utils-bpf",
     srcs: [
@@ -45,7 +45,7 @@
         // 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",
+        "net-utils-device-common-struct-base",
     ],
     apex_available: [
         "com.android.tethering",
diff --git a/framework-t/src/android/app/usage/NetworkStatsManager.java b/framework-t/src/android/app/usage/NetworkStatsManager.java
index 7fa0661..18c839f 100644
--- a/framework-t/src/android/app/usage/NetworkStatsManager.java
+++ b/framework-t/src/android/app/usage/NetworkStatsManager.java
@@ -752,8 +752,8 @@
     /**
      * Query realtime mobile network usage statistics.
      *
-     * Return a snapshot of current UID network statistics, as it applies
-     * to the mobile radios of the device. The snapshot will include any
+     * Return a snapshot of current UID network statistics for both cellular and satellite (which
+     * also uses same mobile radio as cellular) when called. The snapshot will include any
      * tethering traffic, video calling data usage and count of
      * network operations set by {@link TrafficStats#incrementOperationCount}
      * made over a mobile radio.
diff --git a/framework-t/src/android/net/INetworkStatsService.aidl b/framework-t/src/android/net/INetworkStatsService.aidl
index c86f7fd..7f0c1fe 100644
--- a/framework-t/src/android/net/INetworkStatsService.aidl
+++ b/framework-t/src/android/net/INetworkStatsService.aidl
@@ -101,4 +101,7 @@
      * Note that invocation of any interface will be sent to all providers.
      */
      void setStatsProviderWarningAndLimitAsync(String iface, long warning, long limit);
+
+     /** Clear TrafficStats rate-limit caches. */
+     void clearTrafficStatsRateLimitCaches();
 }
diff --git a/framework-t/src/android/net/TrafficStats.java b/framework-t/src/android/net/TrafficStats.java
index a69b38d..77c8001 100644
--- a/framework-t/src/android/net/TrafficStats.java
+++ b/framework-t/src/android/net/TrafficStats.java
@@ -19,6 +19,7 @@
 import static android.annotation.SystemApi.Client.MODULE_LIBRARIES;
 
 import android.annotation.NonNull;
+import android.annotation.RequiresPermission;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
@@ -692,6 +693,27 @@
         return UNSUPPORTED;
     }
 
+    /** Clear TrafficStats rate-limit caches.
+     *
+     * This is mainly for {@link com.android.server.net.NetworkStatsService} to
+     * clear rate-limit cache to avoid caching for TrafficStats API results.
+     * Tests might get stale values after generating network traffic, which
+     * generally need to wait for cache expiry to get updated values.
+     *
+     * @hide
+     */
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_STACK,
+            android.Manifest.permission.NETWORK_SETTINGS})
+    public static void clearRateLimitCaches() {
+        try {
+            getStatsService().clearTrafficStatsRateLimitCaches();
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
     /**
      * Return the number of packets transmitted on the specified interface since the interface
      * was created. Statistics are measured at the network layer, so both TCP and
diff --git a/framework-t/src/android/net/nsd/NsdServiceInfo.java b/framework-t/src/android/net/nsd/NsdServiceInfo.java
index dba08a1..9491a9c 100644
--- a/framework-t/src/android/net/nsd/NsdServiceInfo.java
+++ b/framework-t/src/android/net/nsd/NsdServiceInfo.java
@@ -16,6 +16,8 @@
 
 package android.net.nsd;
 
+import static com.android.net.module.util.HexDump.toHexString;
+
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -39,6 +41,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.StringJoiner;
 
 /**
  * A class representing service information for network service discovery
@@ -541,11 +544,50 @@
                 .append(", network: ").append(mNetwork)
                 .append(", expirationTime: ").append(mExpirationTime);
 
-        byte[] txtRecord = getTxtRecord();
-        sb.append(", txtRecord: ").append(new String(txtRecord, StandardCharsets.UTF_8));
+        final StringJoiner txtJoiner =
+                new StringJoiner(", " /* delimiter */, "{" /* prefix */, "}" /* suffix */);
+
+        sb.append(", txtRecord: ");
+        for (int i = 0; i < mTxtRecord.size(); i++) {
+            txtJoiner.add(mTxtRecord.keyAt(i) + "=" + getPrintableTxtValue(mTxtRecord.valueAt(i)));
+        }
+        sb.append(txtJoiner.toString());
         return sb.toString();
     }
 
+    /**
+     * Returns printable string for {@code txtValue}.
+     *
+     * If {@code txtValue} contains non-printable ASCII characters, a HEX string with prefix "0x"
+     * will be returned. Otherwise, the ASCII string of {@code txtValue} is returned.
+     *
+     */
+    private static String getPrintableTxtValue(@Nullable byte[] txtValue) {
+        if (txtValue == null) {
+            return "(null)";
+        }
+
+        if (containsNonPrintableChars(txtValue)) {
+            return "0x" + toHexString(txtValue);
+        }
+
+        return new String(txtValue, StandardCharsets.US_ASCII);
+    }
+
+    /**
+     * Returns {@code true} if {@code txtValue} contains non-printable ASCII characters.
+     *
+     * The printable characters are in range of [32, 126].
+     */
+    private static boolean containsNonPrintableChars(byte[] txtValue) {
+        for (int i = 0; i < txtValue.length; i++) {
+            if (txtValue[i] < 32 || txtValue[i] > 126) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     /** Implement the Parcelable interface */
     public int describeContents() {
         return 0;
diff --git a/framework/Android.bp b/framework/Android.bp
index 52f2c7c..8787167 100644
--- a/framework/Android.bp
+++ b/framework/Android.bp
@@ -96,6 +96,7 @@
     ],
     impl_only_static_libs: [
         "net-utils-device-common-bpf",
+        "net-utils-device-common-struct-base",
     ],
     libs: [
         "androidx.annotation_annotation",
@@ -124,6 +125,7 @@
         // 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-bpf",
+        "net-utils-device-common-struct-base",
     ],
     libs: [
         // This cannot be in the defaults clause above because if it were, it would be used
diff --git a/framework/src/android/net/NetworkCapabilities.java b/framework/src/android/net/NetworkCapabilities.java
index 84a0d29..45efbfe 100644
--- a/framework/src/android/net/NetworkCapabilities.java
+++ b/framework/src/android/net/NetworkCapabilities.java
@@ -749,22 +749,22 @@
      * Network capabilities that are expected to be mutable, i.e., can change while a particular
      * network is connected.
      */
-    private static final long MUTABLE_CAPABILITIES = BitUtils.packBitList(
+    private static final long MUTABLE_CAPABILITIES =
             // TRUSTED can change when user explicitly connects to an untrusted network in Settings.
             // http://b/18206275
-            NET_CAPABILITY_TRUSTED,
-            NET_CAPABILITY_VALIDATED,
-            NET_CAPABILITY_CAPTIVE_PORTAL,
-            NET_CAPABILITY_NOT_ROAMING,
-            NET_CAPABILITY_FOREGROUND,
-            NET_CAPABILITY_NOT_CONGESTED,
-            NET_CAPABILITY_NOT_SUSPENDED,
-            NET_CAPABILITY_PARTIAL_CONNECTIVITY,
-            NET_CAPABILITY_TEMPORARILY_NOT_METERED,
-            NET_CAPABILITY_NOT_VCN_MANAGED,
+            (1L << NET_CAPABILITY_TRUSTED) |
+            (1L << NET_CAPABILITY_VALIDATED) |
+            (1L << NET_CAPABILITY_CAPTIVE_PORTAL) |
+            (1L << NET_CAPABILITY_NOT_ROAMING) |
+            (1L << NET_CAPABILITY_FOREGROUND) |
+            (1L << NET_CAPABILITY_NOT_CONGESTED) |
+            (1L << NET_CAPABILITY_NOT_SUSPENDED) |
+            (1L << NET_CAPABILITY_PARTIAL_CONNECTIVITY) |
+            (1L << NET_CAPABILITY_TEMPORARILY_NOT_METERED) |
+            (1L << NET_CAPABILITY_NOT_VCN_MANAGED) |
             // The value of NET_CAPABILITY_HEAD_UNIT is 32, which cannot use int to do bit shift,
             // otherwise there will be an overflow. Use long to do bit shift instead.
-            NET_CAPABILITY_HEAD_UNIT);
+            (1L << NET_CAPABILITY_HEAD_UNIT);
 
     /**
      * Network capabilities that are not allowed in NetworkRequests. This exists because the
@@ -784,10 +784,10 @@
     /**
      * Capabilities that are set by default when the object is constructed.
      */
-    private static final long DEFAULT_CAPABILITIES = BitUtils.packBitList(
-            NET_CAPABILITY_NOT_RESTRICTED,
-            NET_CAPABILITY_TRUSTED,
-            NET_CAPABILITY_NOT_VPN);
+    private static final long DEFAULT_CAPABILITIES =
+            (1L << NET_CAPABILITY_NOT_RESTRICTED) |
+            (1L << NET_CAPABILITY_TRUSTED) |
+            (1L << NET_CAPABILITY_NOT_VPN);
 
     /**
      * Capabilities that are managed by ConnectivityService.
@@ -795,11 +795,10 @@
      */
     @VisibleForTesting
     public static final long CONNECTIVITY_MANAGED_CAPABILITIES =
-            BitUtils.packBitList(
-                    NET_CAPABILITY_VALIDATED,
-                    NET_CAPABILITY_CAPTIVE_PORTAL,
-                    NET_CAPABILITY_FOREGROUND,
-                    NET_CAPABILITY_PARTIAL_CONNECTIVITY);
+            (1L << NET_CAPABILITY_VALIDATED) |
+            (1L << NET_CAPABILITY_CAPTIVE_PORTAL) |
+            (1L << NET_CAPABILITY_FOREGROUND) |
+            (1L << NET_CAPABILITY_PARTIAL_CONNECTIVITY);
 
     /**
      * Capabilities that are allowed for all test networks. This list must be set so that it is safe
@@ -808,15 +807,14 @@
      * IMS, SUPL, etc.
      */
     private static final long TEST_NETWORKS_ALLOWED_CAPABILITIES =
-            BitUtils.packBitList(
-            NET_CAPABILITY_NOT_METERED,
-            NET_CAPABILITY_TEMPORARILY_NOT_METERED,
-            NET_CAPABILITY_NOT_RESTRICTED,
-            NET_CAPABILITY_NOT_VPN,
-            NET_CAPABILITY_NOT_ROAMING,
-            NET_CAPABILITY_NOT_CONGESTED,
-            NET_CAPABILITY_NOT_SUSPENDED,
-            NET_CAPABILITY_NOT_VCN_MANAGED);
+            (1L << NET_CAPABILITY_NOT_METERED) |
+            (1L << NET_CAPABILITY_TEMPORARILY_NOT_METERED) |
+            (1L << NET_CAPABILITY_NOT_RESTRICTED) |
+            (1L << NET_CAPABILITY_NOT_VPN) |
+            (1L << NET_CAPABILITY_NOT_ROAMING) |
+            (1L << NET_CAPABILITY_NOT_CONGESTED) |
+            (1L << NET_CAPABILITY_NOT_SUSPENDED) |
+            (1L << NET_CAPABILITY_NOT_VCN_MANAGED);
 
     /**
      * Extra allowed capabilities for test networks that do not have TRANSPORT_CELLULAR. Test
@@ -824,7 +822,9 @@
      * the risk of being used by running apps.
      */
     private static final long TEST_NETWORKS_EXTRA_ALLOWED_CAPABILITIES_ON_NON_CELL =
-            BitUtils.packBitList(NET_CAPABILITY_CBS, NET_CAPABILITY_DUN, NET_CAPABILITY_RCS);
+            (1L << NET_CAPABILITY_CBS) |
+            (1L << NET_CAPABILITY_DUN) |
+            (1L << NET_CAPABILITY_RCS);
 
     /**
      * Adds the given capability to this {@code NetworkCapability} instance.
@@ -1174,7 +1174,7 @@
      * @hide
      */
     public void maybeMarkCapabilitiesRestricted() {
-        if (NetworkCapabilitiesUtils.inferRestrictedCapability(this)) {
+        if (NetworkCapabilitiesUtils.inferRestrictedCapability(mNetworkCapabilities)) {
             removeCapability(NET_CAPABILITY_NOT_RESTRICTED);
         }
     }
@@ -1368,12 +1368,11 @@
      * Allowed transports on an unrestricted test network (in addition to TRANSPORT_TEST).
      */
     private static final long UNRESTRICTED_TEST_NETWORKS_ALLOWED_TRANSPORTS =
-            BitUtils.packBitList(
-                    TRANSPORT_TEST,
-                    // Test eth networks are created with EthernetManager#setIncludeTestInterfaces
-                    TRANSPORT_ETHERNET,
-                    // Test VPN networks can be created but their UID ranges must be empty.
-                    TRANSPORT_VPN);
+            (1L << TRANSPORT_TEST) |
+            // Test eth networks are created with EthernetManager#setIncludeTestInterfaces
+            (1L << TRANSPORT_ETHERNET) |
+            // Test VPN networks can be created but their UID ranges must be empty.
+            (1L << TRANSPORT_VPN);
 
     /**
      * Adds the given transport type to this {@code NetworkCapability} instance.
@@ -1775,8 +1774,7 @@
                 // use the same specifier, TelephonyNetworkSpecifier.
                 && mTransportTypes != (1L << TRANSPORT_TEST)
                 && Long.bitCount(mTransportTypes & ~(1L << TRANSPORT_TEST)) != 1
-                && (mTransportTypes & ~(1L << TRANSPORT_TEST))
-                != (1 << TRANSPORT_CELLULAR | 1 << TRANSPORT_SATELLITE)) {
+                && !specifierAcceptableForMultipleTransports(mTransportTypes)) {
             throw new IllegalStateException("Must have a single non-test transport specified to "
                     + "use setNetworkSpecifier");
         }
@@ -1786,6 +1784,12 @@
         return this;
     }
 
+    private boolean specifierAcceptableForMultipleTransports(long transportTypes) {
+        return (transportTypes & ~(1L << TRANSPORT_TEST))
+                // Cellular and satellite use the same NetworkSpecifier.
+                == (1 << TRANSPORT_CELLULAR | 1 << TRANSPORT_SATELLITE);
+    }
+
     /**
      * Sets the optional transport specific information.
      *
diff --git a/netbpfload/netbpfload.mainline.rc b/netbpfload/netbpfload.mainline.rc
index d7202f7..d38a503 100644
--- a/netbpfload/netbpfload.mainline.rc
+++ b/netbpfload/netbpfload.mainline.rc
@@ -10,6 +10,7 @@
     capabilities CHOWN SYS_ADMIN NET_ADMIN
     group system root graphics network_stack net_admin net_bw_acct net_bw_stats net_raw
     user system
+    file /dev/kmsg w
     rlimit memlock 1073741824 1073741824
     oneshot
     reboot_on_failure reboot,bpfloader-failed
diff --git a/netd/BpfHandler.cpp b/netd/BpfHandler.cpp
index e6fc825..0d75c05 100644
--- a/netd/BpfHandler.cpp
+++ b/netd/BpfHandler.cpp
@@ -179,22 +179,24 @@
 }
 
 Status BpfHandler::init(const char* cg2_path) {
-    // Make sure BPF programs are loaded before doing anything
-    ALOGI("Waiting for BPF programs");
+    if (base::GetProperty("bpf.progs_loaded", "") != "1") {
+        // Make sure BPF programs are loaded before doing anything
+        ALOGI("Waiting for BPF programs");
 
-    if (true || !modules::sdklevel::IsAtLeastV()) {
-        waitForNetProgsLoaded();
-        ALOGI("Networking BPF programs are loaded");
+        if (true || !modules::sdklevel::IsAtLeastV()) {
+            waitForNetProgsLoaded();
+            ALOGI("Networking BPF programs are loaded");
 
-        if (!base::SetProperty("ctl.start", "mdnsd_loadbpf")) {
-            ALOGE("Failed to set property ctl.start=mdnsd_loadbpf, see dmesg for reason.");
-            abort();
+            if (!base::SetProperty("ctl.start", "mdnsd_loadbpf")) {
+                ALOGE("Failed to set property ctl.start=mdnsd_loadbpf, see dmesg for reason.");
+                abort();
+            }
+
+            ALOGI("Waiting for remaining BPF programs");
         }
 
-        ALOGI("Waiting for remaining BPF programs");
+        android::bpf::waitForProgsLoaded();
     }
-
-    android::bpf::waitForProgsLoaded();
     ALOGI("BPF programs are loaded");
 
     RETURN_IF_NOT_OK(initPrograms(cg2_path));
diff --git a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
index 42a922d..385adc6 100644
--- a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
+++ b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
@@ -247,12 +247,15 @@
      * @param isLegacy Whether this call is using legacy backend.
      * @param transactionId The transaction id of service resolution.
      * @param durationMs The duration before stop resolving the service.
+     * @param sentQueryCount The count of sent queries during resolving.
      */
-    public void reportServiceResolutionStop(boolean isLegacy, int transactionId, long durationMs) {
+    public void reportServiceResolutionStop(boolean isLegacy, int transactionId, long durationMs,
+            int sentQueryCount) {
         final Builder builder = makeReportedBuilder(isLegacy, transactionId);
         builder.setType(NsdEventType.NET_RESOLVE);
         builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_RESOLUTION_STOP);
         builder.setEventDurationMillisec(durationMs);
+        builder.setSentQueryCount(sentQueryCount);
         mDependencies.statsWrite(builder.build());
     }
 
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index aca386f..2e258ab 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -275,7 +275,7 @@
     }
 
     @VisibleForTesting
-    static class MdnsListener implements MdnsServiceBrowserListener {
+    abstract static class MdnsListener implements MdnsServiceBrowserListener {
         protected final int mClientRequestId;
         protected final int mTransactionId;
         @NonNull
@@ -321,6 +321,10 @@
 
         @Override
         public void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode) { }
+
+        // Ensure toString gets overridden
+        @NonNull
+        public abstract String toString();
     }
 
     private class DiscoveryListener extends MdnsListener {
@@ -351,13 +355,21 @@
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     DISCOVERY_QUERY_SENT_CALLBACK, new MdnsEvent(mClientRequestId));
         }
+
+        @NonNull
+        @Override
+        public String toString() {
+            return String.format("DiscoveryListener: serviceType=%s", getListenedServiceType());
+        }
     }
 
     private class ResolutionListener extends MdnsListener {
+        private final String mServiceName;
 
         ResolutionListener(int clientRequestId, int transactionId,
-                @NonNull String listenServiceType) {
+                @NonNull String listenServiceType, @NonNull String serviceName) {
             super(clientRequestId, transactionId, listenServiceType);
+            mServiceName = serviceName;
         }
 
         @Override
@@ -373,13 +385,22 @@
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     DISCOVERY_QUERY_SENT_CALLBACK, new MdnsEvent(mClientRequestId));
         }
+
+        @NonNull
+        @Override
+        public String toString() {
+            return String.format("ResolutionListener serviceName=%s, serviceType=%s",
+                    mServiceName, getListenedServiceType());
+        }
     }
 
     private class ServiceInfoListener extends MdnsListener {
+        private final String mServiceName;
 
         ServiceInfoListener(int clientRequestId, int transactionId,
-                @NonNull String listenServiceType) {
+                @NonNull String listenServiceType, @NonNull String serviceName) {
             super(clientRequestId, transactionId, listenServiceType);
+            this.mServiceName = serviceName;
         }
 
         @Override
@@ -410,6 +431,13 @@
             mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
                     DISCOVERY_QUERY_SENT_CALLBACK, new MdnsEvent(mClientRequestId));
         }
+
+        @NonNull
+        @Override
+        public String toString() {
+            return String.format("ServiceInfoListener serviceName=%s, serviceType=%s",
+                    mServiceName, getListenedServiceType());
+        }
     }
 
     private class SocketRequestMonitor implements MdnsSocketProvider.SocketRequestMonitor {
@@ -656,9 +684,12 @@
             }
 
             private void storeAdvertiserRequestMap(int clientRequestId, int transactionId,
-                    ClientInfo clientInfo, @Nullable Network requestedNetwork) {
+                    ClientInfo clientInfo, @NonNull NsdServiceInfo serviceInfo) {
+                final String serviceFullName =
+                        serviceInfo.getServiceName() + "." + serviceInfo.getServiceType();
                 clientInfo.mClientRequests.put(clientRequestId, new AdvertiserClientRequest(
-                        transactionId, requestedNetwork, mClock.elapsedRealtime()));
+                        transactionId, serviceInfo.getNetwork(), serviceFullName,
+                        mClock.elapsedRealtime()));
                 mTransactionIdToClientInfoMap.put(transactionId, clientInfo);
                 updateMulticastLock();
             }
@@ -1010,7 +1041,7 @@
                             mAdvertiser.addOrUpdateService(transactionId, serviceInfo,
                                     mdnsAdvertisingOptions, clientInfo.mUid);
                             storeAdvertiserRequestMap(clientRequestId, transactionId, clientInfo,
-                                    serviceInfo.getNetwork());
+                                    serviceInfo);
                         } else {
                             maybeStartDaemon();
                             transactionId = getUniqueId();
@@ -1104,9 +1135,12 @@
 
                             maybeStartMonitoringSockets();
                             final MdnsListener listener = new ResolutionListener(clientRequestId,
-                                    transactionId, resolveServiceType);
+                                    transactionId, resolveServiceType, info.getServiceName());
+                            final int ifaceIdx = info.getNetwork() != null
+                                    ? 0 : info.getInterfaceIndex();
                             final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
                                     .setNetwork(info.getNetwork())
+                                    .setInterfaceIndex(ifaceIdx)
                                     .setQueryMode(mMdnsFeatureFlags.isAggressiveQueryModeEnabled()
                                             ? AGGRESSIVE_QUERY_MODE
                                             : PASSIVE_QUERY_MODE)
@@ -1204,9 +1238,12 @@
 
                         maybeStartMonitoringSockets();
                         final MdnsListener listener = new ServiceInfoListener(clientRequestId,
-                                transactionId, resolveServiceType);
+                                transactionId, resolveServiceType, info.getServiceName());
+                        final int ifIndex = info.getNetwork() != null
+                                ? 0 : info.getInterfaceIndex();
                         final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
                                 .setNetwork(info.getNetwork())
+                                .setInterfaceIndex(ifIndex)
                                 .setQueryMode(mMdnsFeatureFlags.isAggressiveQueryModeEnabled()
                                         ? AGGRESSIVE_QUERY_MODE
                                         : PASSIVE_QUERY_MODE)
@@ -2545,6 +2582,17 @@
         // Dump state machine logs
         mNsdStateMachine.dump(fd, pw, args);
 
+        // Dump clients
+        pw.println();
+        pw.println("Active clients:");
+        pw.increaseIndent();
+        HandlerUtils.runWithScissorsForDump(mNsdStateMachine.getHandler(), () -> {
+            for (ClientInfo clientInfo : mClients.values()) {
+                pw.println(clientInfo.toString());
+            }
+        }, 10_000);
+        pw.decreaseIndent();
+
         // Dump service and clients logs
         pw.println();
         pw.println("Logs:");
@@ -2617,6 +2665,21 @@
         public int getSentQueryCount() {
             return mSentQueryCount;
         }
+
+        @NonNull
+        @Override
+        public String toString() {
+            return getRequestDescriptor() + " {" + mTransactionId
+                    + ", startTime " + mStartTimeMs
+                    + ", foundServices " + mFoundServiceCount
+                    + ", lostServices " + mLostServiceCount
+                    + ", fromCache " + mIsServiceFromCache
+                    + ", sentQueries " + mSentQueryCount
+                    + "}";
+        }
+
+        @NonNull
+        protected abstract String getRequestDescriptor();
     }
 
     private static class LegacyClientRequest extends ClientRequest {
@@ -2626,6 +2689,12 @@
             super(transactionId, startTimeMs);
             mRequestCode = requestCode;
         }
+
+        @NonNull
+        @Override
+        protected String getRequestDescriptor() {
+            return "Legacy (" + mRequestCode + ")";
+        }
     }
 
     private abstract static class JavaBackendClientRequest extends ClientRequest {
@@ -2645,9 +2714,20 @@
     }
 
     private static class AdvertiserClientRequest extends JavaBackendClientRequest {
+        @NonNull
+        private final String mServiceFullName;
+
         private AdvertiserClientRequest(int transactionId, @Nullable Network requestedNetwork,
-                long startTimeMs) {
+                @NonNull String serviceFullName, long startTimeMs) {
             super(transactionId, requestedNetwork, startTimeMs);
+            mServiceFullName = serviceFullName;
+        }
+
+        @NonNull
+        @Override
+        public String getRequestDescriptor() {
+            return String.format("Advertiser: serviceFullName=%s, net=%s",
+                    mServiceFullName, getRequestedNetwork());
         }
     }
 
@@ -2660,6 +2740,12 @@
             super(transactionId, requestedNetwork, startTimeMs);
             mListener = listener;
         }
+
+        @NonNull
+        @Override
+        public String getRequestDescriptor() {
+            return String.format("Discovery/%s, net=%s", mListener, getRequestedNetwork());
+        }
     }
 
     /* Information tracked per client */
@@ -2704,17 +2790,15 @@
         @Override
         public String toString() {
             StringBuilder sb = new StringBuilder();
-            sb.append("mResolvedService ").append(mResolvedService).append("\n");
-            sb.append("mIsLegacy ").append(mIsPreSClient).append("\n");
-            sb.append("mUseJavaBackend ").append(mUseJavaBackend).append("\n");
-            sb.append("mUid ").append(mUid).append("\n");
+            sb.append("mUid ").append(mUid).append(", ");
+            sb.append("mResolvedService ").append(mResolvedService).append(", ");
+            sb.append("mIsLegacy ").append(mIsPreSClient).append(", ");
+            sb.append("mUseJavaBackend ").append(mUseJavaBackend).append(", ");
+            sb.append("mClientRequests:\n");
             for (int i = 0; i < mClientRequests.size(); i++) {
                 int clientRequestId = mClientRequests.keyAt(i);
-                sb.append("clientRequestId ")
-                        .append(clientRequestId)
-                        .append(" transactionId ").append(mClientRequests.valueAt(i).mTransactionId)
-                        .append(" type ").append(
-                                mClientRequests.valueAt(i).getClass().getSimpleName())
+                sb.append("  ").append(clientRequestId)
+                        .append(": ").append(mClientRequests.valueAt(i).toString())
                         .append("\n");
             }
             return sb.toString();
@@ -2767,7 +2851,8 @@
                                 request.getSentQueryCount());
                     } else if (listener instanceof ResolutionListener) {
                         mMetrics.reportServiceResolutionStop(false /* isLegacy */, transactionId,
-                                request.calculateRequestDurationMs(mClock.elapsedRealtime()));
+                                request.calculateRequestDurationMs(mClock.elapsedRealtime()),
+                                request.getSentQueryCount());
                     } else if (listener instanceof ServiceInfoListener) {
                         mMetrics.reportServiceInfoCallbackUnregistered(transactionId,
                                 request.calculateRequestDurationMs(mClock.elapsedRealtime()),
@@ -2808,7 +2893,8 @@
                     case NsdManager.RESOLVE_SERVICE:
                         stopResolveService(transactionId);
                         mMetrics.reportServiceResolutionStop(true /* isLegacy */, transactionId,
-                                request.calculateRequestDurationMs(mClock.elapsedRealtime()));
+                                request.calculateRequestDurationMs(mClock.elapsedRealtime()),
+                                NO_SENT_QUERY_COUNT);
                         break;
                     case NsdManager.REGISTER_SERVICE:
                         unregisterService(transactionId);
@@ -3024,7 +3110,8 @@
             mMetrics.reportServiceResolutionStop(
                     isLegacyClientRequest(request),
                     request.mTransactionId,
-                    request.calculateRequestDurationMs(mClock.elapsedRealtime()));
+                    request.calculateRequestDurationMs(mClock.elapsedRealtime()),
+                    request.getSentQueryCount());
             try {
                 mCb.onStopResolutionSucceeded(listenerKey);
             } catch (RemoteException e) {
diff --git a/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java b/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
index e61555a..54943c7 100644
--- a/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
+++ b/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
@@ -84,6 +84,7 @@
     private final byte[] packetCreationBuffer = new byte[1500]; // TODO: use interface MTU
     @NonNull
     private final List<MdnsResponse> existingServices;
+    private final boolean isQueryWithKnownAnswer;
 
     EnqueueMdnsQueryCallable(
             @NonNull MdnsSocketClientBase requestSender,
@@ -98,7 +99,8 @@
             @NonNull MdnsUtils.Clock clock,
             @NonNull SharedLog sharedLog,
             @NonNull MdnsServiceTypeClient.Dependencies dependencies,
-            @NonNull Collection<MdnsResponse> existingServices) {
+            @NonNull Collection<MdnsResponse> existingServices,
+            boolean isQueryWithKnownAnswer) {
         weakRequestSender = new WeakReference<>(requestSender);
         serviceTypeLabels = TextUtils.split(serviceType, "\\.");
         this.subtypes = new ArrayList<>(subtypes);
@@ -112,6 +114,7 @@
         this.sharedLog = sharedLog;
         this.dependencies = dependencies;
         this.existingServices = new ArrayList<>(existingServices);
+        this.isQueryWithKnownAnswer = isQueryWithKnownAnswer;
     }
 
     /**
@@ -226,27 +229,27 @@
 
     private void sendPacket(MdnsSocketClientBase requestSender, InetSocketAddress address,
             MdnsPacket mdnsPacket) throws IOException {
-        final DatagramPacket packet = dependencies.getDatagramPacketFromMdnsPacket(
-                packetCreationBuffer, mdnsPacket, address);
+        final List<DatagramPacket> packets = dependencies.getDatagramPacketsFromMdnsPacket(
+                packetCreationBuffer, mdnsPacket, address, isQueryWithKnownAnswer);
         if (expectUnicastResponse) {
             // MdnsMultinetworkSocketClient is only available on T+
             if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU
                     && requestSender instanceof MdnsMultinetworkSocketClient) {
                 ((MdnsMultinetworkSocketClient) requestSender).sendPacketRequestingUnicastResponse(
-                        packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks);
+                        packets, socketKey, onlyUseIpv6OnIpv6OnlyNetworks);
             } else {
                 requestSender.sendPacketRequestingUnicastResponse(
-                        packet, onlyUseIpv6OnIpv6OnlyNetworks);
+                        packets, onlyUseIpv6OnIpv6OnlyNetworks);
             }
         } else {
             if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU
                     && requestSender instanceof MdnsMultinetworkSocketClient) {
                 ((MdnsMultinetworkSocketClient) requestSender)
                         .sendPacketRequestingMulticastResponse(
-                                packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks);
+                                packets, socketKey, onlyUseIpv6OnIpv6OnlyNetworks);
             } else {
                 requestSender.sendPacketRequestingMulticastResponse(
-                        packet, onlyUseIpv6OnIpv6OnlyNetworks);
+                        packets, onlyUseIpv6OnIpv6OnlyNetworks);
             }
         }
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
index 7b0c738..0ab7a76 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -241,11 +241,30 @@
             }
         }
         // Request the network for discovery.
+        // This requests sockets on all networks even if the searchOptions have a given interface
+        // index (with getNetwork==null, for local interfaces), and only uses matching interfaces
+        // in that case. While this is a simple solution to only use matching sockets, a better
+        // practice would be to only request the correct socket for discovery.
+        // TODO: avoid requesting extra sockets after migrating P2P and tethering networks to local
+        // NetworkAgents.
         socketClient.notifyNetworkRequested(listener, searchOptions.getNetwork(),
                 new MdnsSocketClientBase.SocketCreationCallback() {
                     @Override
                     public void onSocketCreated(@NonNull SocketKey socketKey) {
                         discoveryExecutor.ensureRunningOnHandlerThread();
+                        final int searchInterfaceIndex = searchOptions.getInterfaceIndex();
+                        if (searchOptions.getNetwork() == null
+                                && searchInterfaceIndex > 0
+                                // The interface index in options should only match interfaces that
+                                // do not have any Network; a matching Network should be provided
+                                // otherwise.
+                                && (socketKey.getNetwork() != null
+                                    || socketKey.getInterfaceIndex() != searchInterfaceIndex)) {
+                            sharedLog.i("Skipping " + socketKey + " as ifIndex "
+                                    + searchInterfaceIndex + " was requested.");
+                            return;
+                        }
+
                         // All listeners of the same service types shares the same
                         // MdnsServiceTypeClient.
                         MdnsServiceTypeClient serviceTypeClient =
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
index c1c7d5f..0b2003f 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
@@ -22,12 +22,10 @@
 import android.annotation.Nullable;
 import android.annotation.RequiresApi;
 import android.net.LinkAddress;
-import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
 import android.os.Build;
 import android.os.Handler;
 import android.os.Looper;
-import android.util.ArraySet;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.net.module.util.HexDump;
@@ -38,6 +36,7 @@
 
 import java.io.IOException;
 import java.net.InetSocketAddress;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -136,6 +135,15 @@
 
             mAnnouncer.startSending(info.getServiceId(), announcementInfo,
                     0L /* initialDelayMs */);
+
+            // Re-announce the services which have the same custom hostname.
+            final String hostname = mRecordRepository.getHostnameForServiceId(info.getServiceId());
+            if (hostname != null) {
+                final List<MdnsAnnouncer.AnnouncementInfo> announcementInfos =
+                        new ArrayList<>(mRecordRepository.restartAnnouncingForHostname(hostname));
+                announcementInfos.removeIf((i) -> i.getServiceId() == info.getServiceId());
+                reannounceServices(announcementInfos);
+            }
         }
     }
 
@@ -284,6 +292,7 @@
         if (!mRecordRepository.hasActiveService(id)) return;
         mProber.stop(id);
         mAnnouncer.stop(id);
+        final String hostname = mRecordRepository.getHostnameForServiceId(id);
         final MdnsAnnouncer.ExitAnnouncementInfo exitInfo = mRecordRepository.exitService(id);
         if (exitInfo != null) {
             // This effectively schedules onAllServicesRemoved(), as it is to be called when the
@@ -303,6 +312,17 @@
                 }
             });
         }
+        // Re-probe/re-announce the services which have the same custom hostname. These services
+        // were probed/announced using host addresses which were just removed so they should be
+        // re-probed/re-announced without those addresses.
+        if (hostname != null) {
+            final List<MdnsProber.ProbingInfo> probingInfos =
+                    mRecordRepository.restartProbingForHostname(hostname);
+            reprobeServices(probingInfos);
+            final List<MdnsAnnouncer.AnnouncementInfo> announcementInfos =
+                    mRecordRepository.restartAnnouncingForHostname(hostname);
+            reannounceServices(announcementInfos);
+        }
     }
 
     /**
@@ -447,4 +467,19 @@
             return new byte[0];
         }
     }
+
+    private void reprobeServices(List<MdnsProber.ProbingInfo> probingInfos) {
+        for (MdnsProber.ProbingInfo probingInfo : probingInfos) {
+            mProber.stop(probingInfo.getServiceId());
+            mProber.startProbing(probingInfo);
+        }
+    }
+
+    private void reannounceServices(List<MdnsAnnouncer.AnnouncementInfo> announcementInfos) {
+        for (MdnsAnnouncer.AnnouncementInfo announcementInfo : announcementInfos) {
+            mAnnouncer.stop(announcementInfo.getServiceId());
+            mAnnouncer.startSending(
+                    announcementInfo.getServiceId(), announcementInfo, 0 /* initialDelayMs */);
+        }
+    }
 }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
index 869ac9b..fcfb15f 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
@@ -27,6 +27,7 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.util.ArrayMap;
+import android.util.Log;
 
 import com.android.net.module.util.SharedLog;
 
@@ -213,24 +214,30 @@
         return true;
     }
 
-    private void sendMdnsPacket(@NonNull DatagramPacket packet, @NonNull SocketKey targetSocketKey,
-            boolean onlyUseIpv6OnIpv6OnlyNetworks) {
+    private void sendMdnsPackets(@NonNull List<DatagramPacket> packets,
+            @NonNull SocketKey targetSocketKey, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
         final MdnsInterfaceSocket socket = getTargetSocket(targetSocketKey);
         if (socket == null) {
             mSharedLog.e("No socket matches targetSocketKey=" + targetSocketKey);
             return;
         }
+        if (packets.isEmpty()) {
+            Log.wtf(TAG, "No mDns packets to send");
+            return;
+        }
 
-        final boolean isIpv6 = ((InetSocketAddress) packet.getSocketAddress()).getAddress()
-                instanceof Inet6Address;
-        final boolean isIpv4 = ((InetSocketAddress) packet.getSocketAddress()).getAddress()
-                instanceof Inet4Address;
+        final boolean isIpv6 = ((InetSocketAddress) packets.get(0).getSocketAddress())
+                .getAddress() instanceof Inet6Address;
+        final boolean isIpv4 = ((InetSocketAddress) packets.get(0).getSocketAddress())
+                .getAddress() instanceof Inet4Address;
         final boolean shouldQueryIpv6 = !onlyUseIpv6OnIpv6OnlyNetworks || !socket.hasJoinedIpv4();
         // Check ip capability and network before sending packet
         if ((isIpv6 && socket.hasJoinedIpv6() && shouldQueryIpv6)
                 || (isIpv4 && socket.hasJoinedIpv4())) {
             try {
-                socket.send(packet);
+                for (DatagramPacket packet : packets) {
+                    socket.send(packet);
+                }
             } catch (IOException e) {
                 mSharedLog.e("Failed to send a mDNS packet.", e);
             }
@@ -259,34 +266,34 @@
     }
 
     /**
-     * Send a mDNS request packet via given socket key that asks for multicast response.
+     * Send mDNS request packets via given socket key that asks for multicast response.
      */
-    public void sendPacketRequestingMulticastResponse(@NonNull DatagramPacket packet,
+    public void sendPacketRequestingMulticastResponse(@NonNull List<DatagramPacket> packets,
             @NonNull SocketKey socketKey, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
-        mHandler.post(() -> sendMdnsPacket(packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks));
+        mHandler.post(() -> sendMdnsPackets(packets, socketKey, onlyUseIpv6OnIpv6OnlyNetworks));
     }
 
     @Override
     public void sendPacketRequestingMulticastResponse(
-            @NonNull DatagramPacket packet, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
+            @NonNull List<DatagramPacket> packets, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
         throw new UnsupportedOperationException("This socket client need to specify the socket to"
                 + "send packet");
     }
 
     /**
-     * Send a mDNS request packet via given socket key that asks for unicast response.
+     * Send mDNS request packets via given socket key that asks for unicast response.
      *
      * <p>The socket client may use a null network to identify some or all interfaces, in which case
      * passing null sends the packet to these.
      */
-    public void sendPacketRequestingUnicastResponse(@NonNull DatagramPacket packet,
+    public void sendPacketRequestingUnicastResponse(@NonNull List<DatagramPacket> packets,
             @NonNull SocketKey socketKey, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
-        mHandler.post(() -> sendMdnsPacket(packet, socketKey, onlyUseIpv6OnIpv6OnlyNetworks));
+        mHandler.post(() -> sendMdnsPackets(packets, socketKey, onlyUseIpv6OnIpv6OnlyNetworks));
     }
 
     @Override
     public void sendPacketRequestingUnicastResponse(
-            @NonNull DatagramPacket packet, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
+            @NonNull List<DatagramPacket> packets, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
         throw new UnsupportedOperationException("This socket client need to specify the socket to"
                 + "send packet");
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsPacket.java b/service-t/src/com/android/server/connectivity/mdns/MdnsPacket.java
index 1fabd49..83ecabc 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsPacket.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsPacket.java
@@ -42,7 +42,7 @@
     @NonNull
     public final List<MdnsRecord> additionalRecords;
 
-    MdnsPacket(int flags,
+    public MdnsPacket(int flags,
             @NonNull List<MdnsRecord> questions,
             @NonNull List<MdnsRecord> answers,
             @NonNull List<MdnsRecord> authorityRecords,
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsRecord.java b/service-t/src/com/android/server/connectivity/mdns/MdnsRecord.java
index 4b43989..1f9f42b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsRecord.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsRecord.java
@@ -23,7 +23,8 @@
 import android.os.SystemClock;
 import android.text.TextUtils;
 
-import com.android.internal.annotations.VisibleForTesting;
+import androidx.annotation.VisibleForTesting;
+
 import com.android.server.connectivity.mdns.util.MdnsUtils;
 
 import java.io.IOException;
@@ -231,7 +232,7 @@
      * @param writer The writer to use.
      * @param now    The current system time. This is used when writing the updated TTL.
      */
-    @VisibleForTesting
+    @VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
     public final void write(MdnsPacketWriter writer, long now) throws IOException {
         writeHeaderFields(writer);
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsRecordRepository.java b/service-t/src/com/android/server/connectivity/mdns/MdnsRecordRepository.java
index ac64c3a..073e465 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsRecordRepository.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsRecordRepository.java
@@ -925,22 +925,79 @@
         }
     }
 
+    @Nullable
+    public String getHostnameForServiceId(int id) {
+        ServiceRegistration registration = mServices.get(id);
+        if (registration == null) {
+            return null;
+        }
+        return registration.serviceInfo.getHostname();
+    }
+
+    /**
+     * Restart probing the services which are being probed and using the given custom hostname.
+     *
+     * @return The list of {@link MdnsProber.ProbingInfo} to be used by advertiser.
+     */
+    public List<MdnsProber.ProbingInfo> restartProbingForHostname(@NonNull String hostname) {
+        final ArrayList<MdnsProber.ProbingInfo> probingInfos = new ArrayList<>();
+        forEachActiveServiceRegistrationWithHostname(
+                hostname,
+                (id, registration) -> {
+                    if (!registration.isProbing) {
+                        return;
+                    }
+                    probingInfos.add(makeProbingInfo(id, registration));
+                });
+        return probingInfos;
+    }
+
+    /**
+     * Restart announcing the services which are using the given custom hostname.
+     *
+     * @return The list of {@link MdnsAnnouncer.AnnouncementInfo} to be used by advertiser.
+     */
+    public List<MdnsAnnouncer.AnnouncementInfo> restartAnnouncingForHostname(
+            @NonNull String hostname) {
+        final ArrayList<MdnsAnnouncer.AnnouncementInfo> announcementInfos = new ArrayList<>();
+        forEachActiveServiceRegistrationWithHostname(
+                hostname,
+                (id, registration) -> {
+                    if (registration.isProbing) {
+                        return;
+                    }
+                    announcementInfos.add(makeAnnouncementInfo(id, registration));
+                });
+        return announcementInfos;
+    }
+
     /**
      * Called to indicate that probing succeeded for a service.
+     *
      * @param probeSuccessInfo The successful probing info.
      * @return The {@link MdnsAnnouncer.AnnouncementInfo} to send, now that probing has succeeded.
      */
     public MdnsAnnouncer.AnnouncementInfo onProbingSucceeded(
-            MdnsProber.ProbingInfo probeSuccessInfo)
-            throws IOException {
-
-        int serviceId = probeSuccessInfo.getServiceId();
+            MdnsProber.ProbingInfo probeSuccessInfo) throws IOException {
+        final int serviceId = probeSuccessInfo.getServiceId();
         final ServiceRegistration registration = mServices.get(serviceId);
         if (registration == null) {
             throw new IOException("Service is not registered: " + serviceId);
         }
         registration.setProbing(false);
 
+        return makeAnnouncementInfo(serviceId, registration);
+    }
+
+    /**
+     * Make the announcement info of the given service ID.
+     *
+     * @param serviceId The service ID.
+     * @param registration The service registration.
+     * @return The {@link MdnsAnnouncer.AnnouncementInfo} of the given service ID.
+     */
+    private MdnsAnnouncer.AnnouncementInfo makeAnnouncementInfo(
+            int serviceId, ServiceRegistration registration) {
         final Set<MdnsRecord> answersSet = new LinkedHashSet<>();
         final ArrayList<MdnsRecord> additionalAnswers = new ArrayList<>();
 
@@ -972,8 +1029,8 @@
         addNsecRecordsForUniqueNames(additionalAnswers,
                 mGeneralRecords.iterator(), registration.allRecords.iterator());
 
-        return new MdnsAnnouncer.AnnouncementInfo(
-                probeSuccessInfo.getServiceId(), new ArrayList<>(answersSet), additionalAnswers);
+        return new MdnsAnnouncer.AnnouncementInfo(serviceId,
+                new ArrayList<>(answersSet), additionalAnswers);
     }
 
     /**
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
index 086094b..73405ab 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
@@ -59,6 +59,7 @@
                             source.readInt(),
                             source.readInt() == 1,
                             source.readParcelable(null),
+                            source.readInt(),
                             source.readString(),
                             source.readInt() == 1,
                             source.readInt());
@@ -79,6 +80,8 @@
     private final boolean removeExpiredService;
     // The target network for searching. Null network means search on all possible interfaces.
     @Nullable private final Network mNetwork;
+    // If the target interface does not have a Network, set to the interface index, otherwise unset.
+    private final int mInterfaceIndex;
 
     /** Parcelable constructs for a {@link MdnsSearchOptions}. */
     MdnsSearchOptions(
@@ -86,6 +89,7 @@
             int queryMode,
             boolean removeExpiredService,
             @Nullable Network network,
+            int interfaceIndex,
             @Nullable String resolveInstanceName,
             boolean onlyUseIpv6OnIpv6OnlyNetworks,
             int numOfQueriesBeforeBackoff) {
@@ -98,6 +102,7 @@
         this.numOfQueriesBeforeBackoff = numOfQueriesBeforeBackoff;
         this.removeExpiredService = removeExpiredService;
         mNetwork = network;
+        mInterfaceIndex = interfaceIndex;
         this.resolveInstanceName = resolveInstanceName;
     }
 
@@ -148,15 +153,27 @@
     }
 
     /**
-     * Returns the network which the mdns query should target on.
+     * Returns the network which the mdns query should target.
      *
-     * @return the target network or null if search on all possible interfaces.
+     * @return the target network or null to search on all possible interfaces.
      */
     @Nullable
     public Network getNetwork() {
         return mNetwork;
     }
 
+
+    /**
+     * Returns the interface index which the mdns query should target.
+     *
+     * This is only set when the service is to be searched on an interface that does not have a
+     * Network, in which case {@link #getNetwork()} returns null.
+     * The interface index as per {@link java.net.NetworkInterface#getIndex}, or 0 if unset.
+     */
+    public int getInterfaceIndex() {
+        return mInterfaceIndex;
+    }
+
     /**
      * If non-null, queries should try to resolve all records of this specific service, rather than
      * discovering all services.
@@ -177,6 +194,7 @@
         out.writeInt(queryMode);
         out.writeInt(removeExpiredService ? 1 : 0);
         out.writeParcelable(mNetwork, 0);
+        out.writeInt(mInterfaceIndex);
         out.writeString(resolveInstanceName);
         out.writeInt(onlyUseIpv6OnIpv6OnlyNetworks ? 1 : 0);
         out.writeInt(numOfQueriesBeforeBackoff);
@@ -190,6 +208,7 @@
         private int numOfQueriesBeforeBackoff = 3;
         private boolean removeExpiredService;
         private Network mNetwork;
+        private int mInterfaceIndex;
         private String resolveInstanceName;
 
         private Builder() {
@@ -278,6 +297,16 @@
             return this;
         }
 
+        /**
+         * Set the interface index to use for the query, if not querying on a {@link Network}.
+         *
+         * @see MdnsSearchOptions#getInterfaceIndex()
+         */
+        public Builder setInterfaceIndex(int index) {
+            mInterfaceIndex = index;
+            return this;
+        }
+
         /** Builds a {@link MdnsSearchOptions} with the arguments supplied to this builder. */
         public MdnsSearchOptions build() {
             return new MdnsSearchOptions(
@@ -285,6 +314,7 @@
                     queryMode,
                     removeExpiredService,
                     mNetwork,
+                    mInterfaceIndex,
                     resolveInstanceName,
                     onlyUseIpv6OnIpv6OnlyNetworks,
                     numOfQueriesBeforeBackoff);
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
index bfcd0b4..b3bdbe0 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -30,7 +30,8 @@
 import android.util.ArrayMap;
 import android.util.Pair;
 
-import com.android.internal.annotations.VisibleForTesting;
+import androidx.annotation.VisibleForTesting;
+
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.SharedLog;
 import com.android.server.connectivity.mdns.util.MdnsUtils;
@@ -195,7 +196,7 @@
     /**
      * Dependencies of MdnsServiceTypeClient, for injection in tests.
      */
-    @VisibleForTesting
+    @VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
     public static class Dependencies {
         /**
          * @see Handler#sendMessageDelayed(Message, long)
@@ -227,13 +228,22 @@
         }
 
         /**
-         * Generate a DatagramPacket from given MdnsPacket and InetSocketAddress.
+         * Generate the DatagramPackets from given MdnsPacket and InetSocketAddress.
+         *
+         * <p> If the query with known answer feature is enabled and the MdnsPacket is too large for
+         *     a single DatagramPacket, it will be split into multiple DatagramPackets.
          */
-        public DatagramPacket getDatagramPacketFromMdnsPacket(@NonNull byte[] packetCreationBuffer,
-                @NonNull MdnsPacket packet, @NonNull InetSocketAddress address) throws IOException {
-            final byte[] queryBuffer =
-                    MdnsUtils.createRawDnsPacket(packetCreationBuffer, packet);
-            return new DatagramPacket(queryBuffer, 0, queryBuffer.length, address);
+        public List<DatagramPacket> getDatagramPacketsFromMdnsPacket(
+                @NonNull byte[] packetCreationBuffer, @NonNull MdnsPacket packet,
+                @NonNull InetSocketAddress address, boolean isQueryWithKnownAnswer)
+                throws IOException {
+            if (isQueryWithKnownAnswer) {
+                return MdnsUtils.createQueryDatagramPackets(packetCreationBuffer, packet, address);
+            } else {
+                final byte[] queryBuffer =
+                        MdnsUtils.createRawDnsPacket(packetCreationBuffer, packet);
+                return List.of(new DatagramPacket(queryBuffer, 0, queryBuffer.length, address));
+            }
         }
     }
 
@@ -742,7 +752,8 @@
                                 clock,
                                 sharedLog,
                                 dependencies,
-                                existingServices)
+                                existingServices,
+                                featureFlags.isQueryWithKnownAnswerEnabled())
                                 .call();
             } catch (RuntimeException e) {
                 sharedLog.e(String.format("Failed to run EnqueueMdnsQueryCallable for subtype: %s",
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
index 7b71e43..9cfcba1 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
@@ -25,6 +25,7 @@
 import android.net.wifi.WifiManager.MulticastLock;
 import android.os.SystemClock;
 import android.text.format.DateUtils;
+import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.net.module.util.SharedLog;
@@ -206,18 +207,18 @@
     }
 
     @Override
-    public void sendPacketRequestingMulticastResponse(@NonNull DatagramPacket packet,
+    public void sendPacketRequestingMulticastResponse(@NonNull List<DatagramPacket> packets,
             boolean onlyUseIpv6OnIpv6OnlyNetworks) {
-        sendMdnsPacket(packet, multicastPacketQueue, onlyUseIpv6OnIpv6OnlyNetworks);
+        sendMdnsPackets(packets, multicastPacketQueue, onlyUseIpv6OnIpv6OnlyNetworks);
     }
 
     @Override
-    public void sendPacketRequestingUnicastResponse(@NonNull DatagramPacket packet,
+    public void sendPacketRequestingUnicastResponse(@NonNull List<DatagramPacket> packets,
             boolean onlyUseIpv6OnIpv6OnlyNetworks) {
         if (useSeparateSocketForUnicast) {
-            sendMdnsPacket(packet, unicastPacketQueue, onlyUseIpv6OnIpv6OnlyNetworks);
+            sendMdnsPackets(packets, unicastPacketQueue, onlyUseIpv6OnIpv6OnlyNetworks);
         } else {
-            sendMdnsPacket(packet, multicastPacketQueue, onlyUseIpv6OnIpv6OnlyNetworks);
+            sendMdnsPackets(packets, multicastPacketQueue, onlyUseIpv6OnIpv6OnlyNetworks);
         }
     }
 
@@ -238,17 +239,21 @@
         return false;
     }
 
-    private void sendMdnsPacket(DatagramPacket packet, Queue<DatagramPacket> packetQueueToUse,
-            boolean onlyUseIpv6OnIpv6OnlyNetworks) {
+    private void sendMdnsPackets(List<DatagramPacket> packets,
+            Queue<DatagramPacket> packetQueueToUse, boolean onlyUseIpv6OnIpv6OnlyNetworks) {
         if (shouldStopSocketLoop && !MdnsConfigs.allowAddMdnsPacketAfterDiscoveryStops()) {
             sharedLog.w("sendMdnsPacket() is called after discovery already stopped");
             return;
         }
+        if (packets.isEmpty()) {
+            Log.wtf(TAG, "No mDns packets to send");
+            return;
+        }
 
-        final boolean isIpv4 = ((InetSocketAddress) packet.getSocketAddress()).getAddress()
-                instanceof Inet4Address;
-        final boolean isIpv6 = ((InetSocketAddress) packet.getSocketAddress()).getAddress()
-                instanceof Inet6Address;
+        final boolean isIpv4 = ((InetSocketAddress) packets.get(0).getSocketAddress())
+                .getAddress() instanceof Inet4Address;
+        final boolean isIpv6 = ((InetSocketAddress) packets.get(0).getSocketAddress())
+                .getAddress() instanceof Inet6Address;
         final boolean ipv6Only = multicastSocket != null && multicastSocket.isOnIPv6OnlyNetwork();
         if (isIpv4 && ipv6Only) {
             return;
@@ -258,10 +263,11 @@
         }
 
         synchronized (packetQueueToUse) {
-            while (packetQueueToUse.size() >= MdnsConfigs.mdnsPacketQueueMaxSize()) {
+            while ((packetQueueToUse.size() + packets.size())
+                    > MdnsConfigs.mdnsPacketQueueMaxSize()) {
                 packetQueueToUse.remove();
             }
-            packetQueueToUse.add(packet);
+            packetQueueToUse.addAll(packets);
         }
         triggerSendThread();
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
index b6000f0..b1a543a 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
@@ -23,6 +23,7 @@
 
 import java.io.IOException;
 import java.net.DatagramPacket;
+import java.util.List;
 
 /**
  * Base class for multicast socket client.
@@ -40,15 +41,15 @@
     void setCallback(@Nullable Callback callback);
 
     /**
-     * Send a mDNS request packet via given network that asks for multicast response.
+     * Send mDNS request packets via given network that asks for multicast response.
      */
-    void sendPacketRequestingMulticastResponse(@NonNull DatagramPacket packet,
+    void sendPacketRequestingMulticastResponse(@NonNull List<DatagramPacket> packets,
             boolean onlyUseIpv6OnIpv6OnlyNetworks);
 
     /**
-     * Send a mDNS request packet via given network that asks for unicast response.
+     * Send mDNS request packets via given network that asks for unicast response.
      */
-    void sendPacketRequestingUnicastResponse(@NonNull DatagramPacket packet,
+    void sendPacketRequestingUnicastResponse(@NonNull List<DatagramPacket> packets,
             boolean onlyUseIpv6OnIpv6OnlyNetworks);
 
     /*** Notify that the given network is requested for mdns discovery / resolution */
diff --git a/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java b/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
index d553210..3c11a24 100644
--- a/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
+++ b/service-t/src/com/android/server/connectivity/mdns/util/MdnsUtils.java
@@ -16,6 +16,8 @@
 
 package com.android.server.connectivity.mdns.util;
 
+import static com.android.server.connectivity.mdns.MdnsConstants.FLAG_TRUNCATED;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.net.Network;
@@ -23,6 +25,7 @@
 import android.os.Handler;
 import android.os.SystemClock;
 import android.util.ArraySet;
+import android.util.Pair;
 
 import com.android.server.connectivity.mdns.MdnsConstants;
 import com.android.server.connectivity.mdns.MdnsPacket;
@@ -30,13 +33,18 @@
 import com.android.server.connectivity.mdns.MdnsRecord;
 
 import java.io.IOException;
+import java.net.DatagramPacket;
+import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.nio.CharBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.CharsetEncoder;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
 /**
@@ -226,6 +234,100 @@
     }
 
     /**
+     * Writes the possible query content of an MdnsPacket into the data buffer.
+     *
+     * <p>This method is specifically for query packets. It writes the question and answer sections
+     *    into the data buffer only.
+     *
+     * @param packetCreationBuffer The data buffer for the query content.
+     * @param packet The MdnsPacket to be written into the data buffer.
+     * @return A Pair containing:
+     *         1. The remaining MdnsPacket data that could not fit in the buffer.
+     *         2. The length of the data written to the buffer.
+     */
+    @Nullable
+    private static Pair<MdnsPacket, Integer> writePossibleMdnsPacket(
+            @NonNull byte[] packetCreationBuffer, @NonNull MdnsPacket packet) throws IOException {
+        MdnsPacket remainingPacket;
+        final MdnsPacketWriter writer = new MdnsPacketWriter(packetCreationBuffer);
+        writer.writeUInt16(packet.transactionId); // Transaction ID
+
+        final int flagsPos = writer.getWritePosition();
+        writer.writeUInt16(0); // Flags, written later
+        writer.writeUInt16(0); // questions count, written later
+        writer.writeUInt16(0); // answers count, written later
+        writer.writeUInt16(0); // authority entries count, empty session for query
+        writer.writeUInt16(0); // additional records count, empty session for query
+
+        int writtenQuestions = 0;
+        int writtenAnswers = 0;
+        int lastValidPos = writer.getWritePosition();
+        try {
+            for (MdnsRecord record : packet.questions) {
+                // Questions do not have TTL or data
+                record.writeHeaderFields(writer);
+                writtenQuestions++;
+                lastValidPos = writer.getWritePosition();
+            }
+            for (MdnsRecord record : packet.answers) {
+                record.write(writer, 0L);
+                writtenAnswers++;
+                lastValidPos = writer.getWritePosition();
+            }
+            remainingPacket = null;
+        } catch (IOException e) {
+            // Went over the packet limit; truncate
+            if (writtenQuestions == 0 && writtenAnswers == 0) {
+                // No space to write even one record: just throw (as subclass of IOException)
+                throw e;
+            }
+
+            // Set the last valid position as the final position (not as a rewind)
+            writer.rewind(lastValidPos);
+            writer.clearRewind();
+
+            remainingPacket = new MdnsPacket(packet.flags,
+                    packet.questions.subList(
+                            writtenQuestions, packet.questions.size()),
+                    packet.answers.subList(
+                            writtenAnswers, packet.answers.size()),
+                    Collections.emptyList(), /* authorityRecords */
+                    Collections.emptyList() /* additionalRecords */);
+        }
+
+        final int len = writer.getWritePosition();
+        writer.rewind(flagsPos);
+        writer.writeUInt16(packet.flags | (remainingPacket == null ? 0 : FLAG_TRUNCATED));
+        writer.writeUInt16(writtenQuestions);
+        writer.writeUInt16(writtenAnswers);
+        writer.unrewind();
+
+        return Pair.create(remainingPacket, len);
+    }
+
+    /**
+     * Create Datagram packets from given MdnsPacket and InetSocketAddress.
+     *
+     * <p> If the MdnsPacket is too large for a single DatagramPacket, it will be split into
+     *     multiple DatagramPackets.
+     */
+    public static List<DatagramPacket> createQueryDatagramPackets(
+            @NonNull byte[] packetCreationBuffer, @NonNull MdnsPacket packet,
+            @NonNull InetSocketAddress destination) throws IOException {
+        final List<DatagramPacket> datagramPackets = new ArrayList<>();
+        MdnsPacket remainingPacket = packet;
+        while (remainingPacket != null) {
+            final Pair<MdnsPacket, Integer> result =
+                    writePossibleMdnsPacket(packetCreationBuffer, remainingPacket);
+            remainingPacket = result.first;
+            final int len = result.second;
+            final byte[] outBuffer = Arrays.copyOfRange(packetCreationBuffer, 0, len);
+            datagramPackets.add(new DatagramPacket(outBuffer, 0, outBuffer.length, destination));
+        }
+        return datagramPackets;
+    }
+
+    /**
      * Checks if the MdnsRecord needs to be renewed or not.
      *
      * <p>As per RFC6762 7.1 no need to query if remaining TTL is more than half the original one,
diff --git a/service-t/src/com/android/server/ethernet/EthernetNetworkFactory.java b/service-t/src/com/android/server/ethernet/EthernetNetworkFactory.java
index 0b54fdd..cadc04d 100644
--- a/service-t/src/com/android/server/ethernet/EthernetNetworkFactory.java
+++ b/service-t/src/com/android/server/ethernet/EthernetNetworkFactory.java
@@ -250,6 +250,17 @@
         return mTrackingInterfaces.containsKey(ifaceName);
     }
 
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+    @Nullable
+    protected String getHwAddress(@NonNull final String ifaceName) {
+        if (!hasInterface(ifaceName)) {
+            return null;
+        }
+
+        NetworkInterfaceState iface = mTrackingInterfaces.get(ifaceName);
+        return iface.mHwAddress;
+    }
+
     @VisibleForTesting
     static class NetworkInterfaceState {
         final String name;
diff --git a/service-t/src/com/android/server/ethernet/EthernetTracker.java b/service-t/src/com/android/server/ethernet/EthernetTracker.java
index bfcc171..71f289e 100644
--- a/service-t/src/com/android/server/ethernet/EthernetTracker.java
+++ b/service-t/src/com/android/server/ethernet/EthernetTracker.java
@@ -131,6 +131,10 @@
     // returned when a tethered interface is requested; until then, it remains in client mode. Its
     // current mode is reflected in mTetheringInterfaceMode.
     private String mTetheringInterface;
+    // If the tethering interface is in server mode, it is not tracked by factory. The HW address
+    // must be maintained by the EthernetTracker. Its current mode is reflected in
+    // mTetheringInterfaceMode.
+    private String mTetheringInterfaceHwAddr;
     private int mTetheringInterfaceMode = INTERFACE_MODE_CLIENT;
     // Tracks whether clients were notified that the tethered interface is available
     private boolean mTetheredInterfaceWasAvailable = false;
@@ -460,7 +464,7 @@
             if (!include) {
                 removeTestData();
             }
-            mHandler.post(() -> trackAvailableInterfaces());
+            trackAvailableInterfaces();
         });
     }
 
@@ -582,6 +586,7 @@
         removeInterface(iface);
         if (iface.equals(mTetheringInterface)) {
             mTetheringInterface = null;
+            mTetheringInterfaceHwAddr = null;
         }
         broadcastInterfaceStateChange(iface);
     }
@@ -610,13 +615,14 @@
             return;
         }
 
+        final String hwAddress = config.hwAddr;
+
         if (getInterfaceMode(iface) == INTERFACE_MODE_SERVER) {
             maybeUpdateServerModeInterfaceState(iface, true);
+            mTetheringInterfaceHwAddr = hwAddress;
             return;
         }
 
-        final String hwAddress = config.hwAddr;
-
         NetworkCapabilities nc = mNetworkCapabilities.get(iface);
         if (nc == null) {
             // Try to resolve using mac address
diff --git a/service-t/src/com/android/server/net/NetworkStatsObservers.java b/service-t/src/com/android/server/net/NetworkStatsObservers.java
index cab29e3..21cf351 100644
--- a/service-t/src/com/android/server/net/NetworkStatsObservers.java
+++ b/service-t/src/com/android/server/net/NetworkStatsObservers.java
@@ -32,6 +32,7 @@
 import android.net.NetworkTemplate;
 import android.net.netstats.IUsageCallback;
 import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
@@ -45,7 +46,6 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.net.module.util.PerUidCounter;
 
-import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -78,11 +78,8 @@
     // Sequence number of DataUsageRequests
     private final AtomicInteger mNextDataUsageRequestId = new AtomicInteger();
 
-    private final Handler mHandler;
-
-    NetworkStatsObservers(@NonNull Looper looper) {
-        mHandler = new Handler(Objects.requireNonNull(looper), mHandlerCallback);
-    }
+    // Lazily instantiated when an observer is registered.
+    private volatile Handler mHandler;
 
     /**
      * Creates a wrapper that contains the caller context and a normalized request.
@@ -103,7 +100,7 @@
         if (LOG) Log.d(TAG, "Registering observer for " + requestInfo);
         mDataUsageRequestsPerUid.incrementCountOrThrow(callingUid);
 
-        mHandler.sendMessage(mHandler.obtainMessage(MSG_REGISTER, requestInfo));
+        getHandler().sendMessage(mHandler.obtainMessage(MSG_REGISTER, requestInfo));
         return request;
     }
 
@@ -113,7 +110,7 @@
      * <p>It will unregister the observer asynchronously, so it is safe to call from any thread.
      */
     public void unregister(DataUsageRequest request, int callingUid) {
-        mHandler.sendMessage(mHandler.obtainMessage(MSG_UNREGISTER, callingUid, 0 /* ignore */,
+        getHandler().sendMessage(mHandler.obtainMessage(MSG_UNREGISTER, callingUid, 0 /* ignore */,
                 request));
     }
 
@@ -128,10 +125,34 @@
                 long currentTime) {
         StatsContext statsContext = new StatsContext(xtSnapshot, uidSnapshot, activeIfaces,
                 activeUidIfaces, currentTime);
-        mHandler.sendMessage(mHandler.obtainMessage(MSG_UPDATE_STATS, statsContext));
+        getHandler().sendMessage(mHandler.obtainMessage(MSG_UPDATE_STATS, statsContext));
     }
 
-    private final Handler.Callback mHandlerCallback = new Handler.Callback() {
+    private Handler getHandler() {
+        if (mHandler == null) {
+            synchronized (this) {
+                if (mHandler == null) {
+                    if (LOGV) Log.v(TAG, "Creating handler");
+                    mHandler = new Handler(getHandlerLooperLocked(), mHandlerCallback);
+                }
+            }
+        }
+        return mHandler;
+    }
+
+    @VisibleForTesting
+    protected Looper getHandlerLooperLocked() {
+        // TODO: Currently, callbacks are dispatched on this thread if the caller register
+        //  callback without supplying a Handler. To ensure that the service handler thread
+        //  is not blocked by client code, the observers must create their own thread. Once
+        //  all callbacks are dispatched outside of the handler thread, the service handler
+        //  thread can be used here.
+        HandlerThread handlerThread = new HandlerThread(TAG);
+        handlerThread.start();
+        return handlerThread.getLooper();
+    }
+
+    private Handler.Callback mHandlerCallback = new Handler.Callback() {
         @Override
         public boolean handleMessage(Message msg) {
             switch (msg.what) {
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 64b17eb..d43c3da 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -16,6 +16,7 @@
 
 package com.android.server.net;
 
+import static android.Manifest.permission.NETWORK_SETTINGS;
 import static android.Manifest.permission.NETWORK_STATS_PROVIDER;
 import static android.Manifest.permission.READ_NETWORK_USAGE_HISTORY;
 import static android.Manifest.permission.UPDATE_DEVICE_STATS;
@@ -50,12 +51,17 @@
 import static android.net.NetworkTemplate.MATCH_WIFI;
 import static android.net.TrafficStats.KB_IN_BYTES;
 import static android.net.TrafficStats.MB_IN_BYTES;
+import static android.net.TrafficStats.TYPE_RX_BYTES;
+import static android.net.TrafficStats.TYPE_RX_PACKETS;
+import static android.net.TrafficStats.TYPE_TX_BYTES;
+import static android.net.TrafficStats.TYPE_TX_PACKETS;
 import static android.net.TrafficStats.UID_TETHERING;
 import static android.net.TrafficStats.UNSUPPORTED;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID_TAG;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_XT;
 import static android.os.Trace.TRACE_TAG_NETWORK;
+import static android.provider.DeviceConfig.NAMESPACE_TETHERING;
 import static android.system.OsConstants.ENOENT;
 import static android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID;
 import static android.text.format.DateUtils.DAY_IN_MILLIS;
@@ -64,6 +70,7 @@
 import static android.text.format.DateUtils.SECOND_IN_MILLIS;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PRIVATE;
+import static com.android.net.module.util.DeviceConfigUtils.getDeviceConfigPropertyInt;
 import static com.android.net.module.util.NetworkCapabilitiesUtils.getDisplayTransport;
 import static com.android.net.module.util.NetworkStatsUtils.LIMIT_GLOBAL_ALERT;
 import static com.android.server.net.NetworkStatsEventLogger.POLL_REASON_PERIODIC;
@@ -299,6 +306,12 @@
     static final String NETSTATS_FASTDATAINPUT_SUCCESSES_COUNTER_NAME = "fastdatainput.successes";
     static final String NETSTATS_FASTDATAINPUT_FALLBACKS_COUNTER_NAME = "fastdatainput.fallbacks";
 
+    static final String TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME =
+            "trafficstats_cache_expiry_duration_ms";
+    static final String TRAFFIC_STATS_CACHE_MAX_ENTRIES_NAME = "trafficstats_cache_max_entries";
+    static final int DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS = 1000;
+    static final int DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES = 400;
+
     private final Context mContext;
     private final NetworkStatsFactory mStatsFactory;
     private final AlarmManager mAlarmManager;
@@ -454,6 +467,13 @@
 
     private long mLastStatsSessionPoll;
 
+    private final TrafficStatsRateLimitCache mTrafficStatsTotalCache;
+    private final TrafficStatsRateLimitCache mTrafficStatsIfaceCache;
+    private final TrafficStatsRateLimitCache mTrafficStatsUidCache;
+    static final String TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG =
+            "trafficstats_rate_limit_cache_enabled_flag";
+    private final boolean mSupportTrafficStatsRateLimitCache;
+
     private final Object mOpenSessionCallsLock = new Object();
 
     /**
@@ -593,7 +613,7 @@
                 INetd.Stub.asInterface((IBinder) context.getSystemService(Context.NETD_SERVICE)),
                 alarmManager, wakeLock, getDefaultClock(),
                 new DefaultNetworkStatsSettings(), new NetworkStatsFactory(context),
-                new Dependencies());
+                new NetworkStatsObservers(), new Dependencies());
 
         return service;
     }
@@ -603,7 +623,8 @@
     @VisibleForTesting
     NetworkStatsService(Context context, INetd netd, AlarmManager alarmManager,
             PowerManager.WakeLock wakeLock, Clock clock, NetworkStatsSettings settings,
-            NetworkStatsFactory factory, @NonNull Dependencies deps) {
+            NetworkStatsFactory factory, NetworkStatsObservers statsObservers,
+            @NonNull Dependencies deps) {
         mContext = Objects.requireNonNull(context, "missing Context");
         mNetd = Objects.requireNonNull(netd, "missing Netd");
         mAlarmManager = Objects.requireNonNull(alarmManager, "missing AlarmManager");
@@ -611,6 +632,7 @@
         mSettings = Objects.requireNonNull(settings, "missing NetworkStatsSettings");
         mWakeLock = Objects.requireNonNull(wakeLock, "missing WakeLock");
         mStatsFactory = Objects.requireNonNull(factory, "missing factory");
+        mStatsObservers = Objects.requireNonNull(statsObservers, "missing NetworkStatsObservers");
         mDeps = Objects.requireNonNull(deps, "missing Dependencies");
         mStatsDir = mDeps.getOrCreateStatsDir();
         if (!mStatsDir.exists()) {
@@ -620,7 +642,6 @@
         final HandlerThread handlerThread = mDeps.makeHandlerThread();
         handlerThread.start();
         mHandler = new NetworkStatsHandler(handlerThread.getLooper());
-        mStatsObservers = new NetworkStatsObservers(handlerThread.getLooper());
         mNetworkStatsSubscriptionsMonitor = deps.makeSubscriptionsMonitor(mContext,
                 (command) -> mHandler.post(command) , this);
         mContentResolver = mContext.getContentResolver();
@@ -642,6 +663,16 @@
             mEventLogger = null;
         }
 
+        final long cacheExpiryDurationMs = mDeps.getTrafficStatsRateLimitCacheExpiryDuration();
+        final int cacheMaxEntries = mDeps.getTrafficStatsRateLimitCacheMaxEntries();
+        mSupportTrafficStatsRateLimitCache = mDeps.supportTrafficStatsRateLimitCache(mContext);
+        mTrafficStatsTotalCache = new TrafficStatsRateLimitCache(mClock,
+                cacheExpiryDurationMs, cacheMaxEntries);
+        mTrafficStatsIfaceCache = new TrafficStatsRateLimitCache(mClock,
+                cacheExpiryDurationMs, cacheMaxEntries);
+        mTrafficStatsUidCache = new TrafficStatsRateLimitCache(mClock,
+                cacheExpiryDurationMs, cacheMaxEntries);
+
         // TODO: Remove bpfNetMaps creation and always start SkDestroyListener
         // Following code is for the experiment to verify the SkDestroyListener refactoring. Based
         // on the experiment flag, BpfNetMaps starts C SkDestroyListener (existing code) or
@@ -695,7 +726,7 @@
          * Get the count of import legacy target attempts.
          */
         public int getImportLegacyTargetAttempts() {
-            return DeviceConfigUtils.getDeviceConfigPropertyInt(
+            return getDeviceConfigPropertyInt(
                     DeviceConfig.NAMESPACE_TETHERING,
                     NETSTATS_IMPORT_LEGACY_TARGET_ATTEMPTS,
                     DEFAULT_NETSTATS_IMPORT_LEGACY_TARGET_ATTEMPTS);
@@ -705,7 +736,7 @@
          * Get the count of using FastDataInput target attempts.
          */
         public int getUseFastDataInputTargetAttempts() {
-            return DeviceConfigUtils.getDeviceConfigPropertyInt(
+            return getDeviceConfigPropertyInt(
                     DeviceConfig.NAMESPACE_TETHERING,
                     NETSTATS_FASTDATAINPUT_TARGET_ATTEMPTS, 0);
         }
@@ -887,6 +918,76 @@
             return DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
                     ctx, CONFIG_ENABLE_NETWORK_STATS_EVENT_LOGGER);
         }
+
+        /**
+         * Get whether TrafficStats rate-limit cache is supported.
+         *
+         * This method should only be called once in the constructor,
+         * to ensure that the code does not need to deal with flag values changing at runtime.
+         */
+        public boolean supportTrafficStatsRateLimitCache(@NonNull Context ctx) {
+            return SdkLevel.isAtLeastV() && DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
+                    ctx, TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG);
+        }
+
+        /**
+         * Get TrafficStats rate-limit cache expiry.
+         *
+         * This method should only be called once in the constructor,
+         * to ensure that the code does not need to deal with flag values changing at runtime.
+         */
+        public int getTrafficStatsRateLimitCacheExpiryDuration() {
+            return getDeviceConfigPropertyInt(
+                    NAMESPACE_TETHERING, TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME,
+                    DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS);
+        }
+
+        /**
+         * Get TrafficStats rate-limit cache max entries.
+         *
+         * This method should only be called once in the constructor,
+         * to ensure that the code does not need to deal with flag values changing at runtime.
+         */
+        public int getTrafficStatsRateLimitCacheMaxEntries() {
+            return getDeviceConfigPropertyInt(
+                    NAMESPACE_TETHERING, TRAFFIC_STATS_CACHE_MAX_ENTRIES_NAME,
+                    DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES);
+        }
+
+        /**
+         * Retrieves native network total statistics.
+         *
+         * @return A NetworkStats.Entry containing the native statistics, or
+         *         null if an error occurs.
+         */
+        @Nullable
+        public NetworkStats.Entry nativeGetTotalStat() {
+            return NetworkStatsService.nativeGetTotalStat();
+        }
+
+        /**
+         * Retrieves native network interface statistics for the specified interface.
+         *
+         * @param iface The name of the network interface to query.
+         * @return A NetworkStats.Entry containing the native statistics for the interface, or
+         *         null if an error occurs.
+         */
+        @Nullable
+        public NetworkStats.Entry nativeGetIfaceStat(String iface) {
+            return NetworkStatsService.nativeGetIfaceStat(iface);
+        }
+
+        /**
+         * Retrieves native network uid statistics for the specified uid.
+         *
+         * @param uid The uid of the application to query.
+         * @return A NetworkStats.Entry containing the native statistics for the uid, or
+         *         null if an error occurs.
+         */
+        @Nullable
+        public NetworkStats.Entry nativeGetUidStat(int uid) {
+            return NetworkStatsService.nativeGetUidStat(uid);
+        }
     }
 
     /**
@@ -1778,6 +1879,8 @@
             if (transport == TRANSPORT_WIFI) {
                 ifaceSet = mAllWifiIfacesSinceBoot;
             } else if (transport == TRANSPORT_CELLULAR) {
+                // Since satellite networks appear under type mobile, this includes both cellular
+                // and satellite active interfaces
                 ifaceSet = mAllMobileIfacesSinceBoot;
             } else {
                 throw new IllegalArgumentException("Invalid transport " + transport);
@@ -1980,53 +2083,106 @@
         if (callingUid != android.os.Process.SYSTEM_UID && callingUid != uid) {
             return UNSUPPORTED;
         }
-        return getEntryValueForType(nativeGetUidStat(uid), type);
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
+
+        if (!mSupportTrafficStatsRateLimitCache) {
+            return getEntryValueForType(mDeps.nativeGetUidStat(uid), type);
+        }
+
+        final NetworkStats.Entry entry = mTrafficStatsUidCache.getOrCompute(IFACE_ALL, uid,
+                () -> mDeps.nativeGetUidStat(uid));
+
+        return getEntryValueForType(entry, type);
+    }
+
+    @Nullable
+    private NetworkStats.Entry getIfaceStatsInternal(@NonNull String iface) {
+        final NetworkStats.Entry entry = mDeps.nativeGetIfaceStat(iface);
+        if (entry == null) {
+            return null;
+        }
+        // When tethering offload is in use, nativeIfaceStats does not contain usage from
+        // offload, add it back here. Note that the included statistics might be stale
+        // since polling newest stats from hardware might impact system health and not
+        // suitable for TrafficStats API use cases.
+        entry.add(getProviderIfaceStats(iface));
+        return entry;
     }
 
     @Override
     public long getIfaceStats(@NonNull String iface, int type) {
         Objects.requireNonNull(iface);
-        final NetworkStats.Entry entry = nativeGetIfaceStat(iface);
-        final long value = getEntryValueForType(entry, type);
-        if (value == UNSUPPORTED) {
-            return UNSUPPORTED;
-        } else {
-            // When tethering offload is in use, nativeIfaceStats does not contain usage from
-            // offload, add it back here. Note that the included statistics might be stale
-            // since polling newest stats from hardware might impact system health and not
-            // suitable for TrafficStats API use cases.
-            entry.add(getProviderIfaceStats(iface));
-            return getEntryValueForType(entry, type);
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
+
+        if (!mSupportTrafficStatsRateLimitCache) {
+            return getEntryValueForType(getIfaceStatsInternal(iface), type);
         }
+
+        final NetworkStats.Entry entry = mTrafficStatsIfaceCache.getOrCompute(iface, UID_ALL,
+                () -> getIfaceStatsInternal(iface));
+
+        return getEntryValueForType(entry, type);
     }
 
     private long getEntryValueForType(@Nullable NetworkStats.Entry entry, int type) {
         if (entry == null) return UNSUPPORTED;
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
         switch (type) {
-            case TrafficStats.TYPE_RX_BYTES:
+            case TYPE_RX_BYTES:
                 return entry.rxBytes;
-            case TrafficStats.TYPE_TX_BYTES:
-                return entry.txBytes;
-            case TrafficStats.TYPE_RX_PACKETS:
+            case TYPE_RX_PACKETS:
                 return entry.rxPackets;
-            case TrafficStats.TYPE_TX_PACKETS:
+            case TYPE_TX_BYTES:
+                return entry.txBytes;
+            case TYPE_TX_PACKETS:
                 return entry.txPackets;
             default:
-                return UNSUPPORTED;
+                throw new IllegalStateException("Bug: Invalid type: "
+                        + type + " should not reach here.");
         }
     }
 
+    private boolean isEntryValueTypeValid(int type) {
+        switch (type) {
+            case TYPE_RX_BYTES:
+            case TYPE_RX_PACKETS:
+            case TYPE_TX_BYTES:
+            case TYPE_TX_PACKETS:
+                return true;
+            default :
+                return false;
+        }
+    }
+
+    @Nullable
+    private NetworkStats.Entry getTotalStatsInternal() {
+        final NetworkStats.Entry entry = mDeps.nativeGetTotalStat();
+        if (entry == null) {
+            return null;
+        }
+        entry.add(getProviderIfaceStats(IFACE_ALL));
+        return entry;
+    }
+
     @Override
     public long getTotalStats(int type) {
-        final NetworkStats.Entry entry = nativeGetTotalStat();
-        final long value = getEntryValueForType(entry, type);
-        if (value == UNSUPPORTED) {
-            return UNSUPPORTED;
-        } else {
-            // Refer to comment in getIfaceStats
-            entry.add(getProviderIfaceStats(IFACE_ALL));
-            return getEntryValueForType(entry, type);
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
+        if (!mSupportTrafficStatsRateLimitCache) {
+            return getEntryValueForType(getTotalStatsInternal(), type);
         }
+
+        final NetworkStats.Entry entry = mTrafficStatsTotalCache.getOrCompute(IFACE_ALL, UID_ALL,
+                () -> getTotalStatsInternal());
+
+        return getEntryValueForType(entry, type);
+    }
+
+    @Override
+    public void clearTrafficStatsRateLimitCaches() {
+        PermissionUtils.enforceNetworkStackPermissionOr(mContext, NETWORK_SETTINGS);
+        mTrafficStatsUidCache.clear();
+        mTrafficStatsIfaceCache.clear();
+        mTrafficStatsTotalCache.clear();
     }
 
     private NetworkStats.Entry getProviderIfaceStats(@Nullable String iface) {
@@ -2187,7 +2343,9 @@
         for (NetworkStateSnapshot snapshot : snapshots) {
             final int displayTransport =
                     getDisplayTransport(snapshot.getNetworkCapabilities().getTransportTypes());
-            final boolean isMobile = (NetworkCapabilities.TRANSPORT_CELLULAR == displayTransport);
+            // Consider satellite transport to support satellite stats appear as type_mobile
+            final boolean isMobile = NetworkCapabilities.TRANSPORT_CELLULAR == displayTransport
+                    || NetworkCapabilities.TRANSPORT_SATELLITE == displayTransport;
             final boolean isWifi = (NetworkCapabilities.TRANSPORT_WIFI == displayTransport);
             final boolean isDefault = CollectionUtils.contains(
                     mDefaultNetworks, snapshot.getNetwork());
@@ -2320,12 +2478,14 @@
     }
 
     /**
-     * For networks with {@code TRANSPORT_CELLULAR}, get ratType that was obtained through
-     * {@link PhoneStateListener}. Otherwise, return 0 given that other networks with different
-     * transport types do not actually fill this value.
+     * For networks with {@code TRANSPORT_CELLULAR} Or {@code TRANSPORT_SATELLITE}, get ratType
+     * that was obtained through {@link PhoneStateListener}. Otherwise, return 0 given that other
+     * networks with different transport types do not actually fill this value.
      */
     private int getRatTypeForStateSnapshot(@NonNull NetworkStateSnapshot state) {
-        if (!state.getNetworkCapabilities().hasTransport(NetworkCapabilities.TRANSPORT_CELLULAR)) {
+        if (!state.getNetworkCapabilities().hasTransport(NetworkCapabilities.TRANSPORT_CELLULAR)
+                && !state.getNetworkCapabilities()
+                .hasTransport(NetworkCapabilities.TRANSPORT_SATELLITE)) {
             return 0;
         }
 
@@ -2778,6 +2938,14 @@
             } catch (IOException e) {
                 pw.println("(failed to dump FastDataInput counters)");
             }
+            pw.print("trafficstats.cache.supported", mSupportTrafficStatsRateLimitCache);
+            pw.println();
+            pw.print(TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME,
+                    mDeps.getTrafficStatsRateLimitCacheExpiryDuration());
+            pw.println();
+            pw.print(TRAFFIC_STATS_CACHE_MAX_ENTRIES_NAME,
+                    mDeps.getTrafficStatsRateLimitCacheMaxEntries());
+            pw.println();
 
             pw.decreaseIndent();
 
diff --git a/service/Android.bp b/service/Android.bp
index c35c4f8..1d74efc 100644
--- a/service/Android.bp
+++ b/service/Android.bp
@@ -199,7 +199,9 @@
         "PlatformProperties",
         "service-connectivity-protos",
         "service-connectivity-stats-protos",
-        "net-utils-multicast-forwarding-structs",
+        // The required dependency net-utils-device-common-struct-base is in the classpath via
+        // framework-connectivity
+        "net-utils-device-common-struct",
     ],
     apex_available: [
         "com.android.tethering",
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index 123ad8f..a15a2bf 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -478,6 +478,7 @@
     private volatile boolean mLockdownEnabled;
 
     private final boolean mRequestRestrictedWifiEnabled;
+    private final boolean mBackgroundFirewallChainEnabled;
 
     /**
      * Stale copy of uid blocked reasons provided by NPMS. As long as they are accessed only in
@@ -1798,6 +1799,8 @@
         mLocationPermissionChecker = mDeps.makeLocationPermissionChecker(mContext);
         mRequestRestrictedWifiEnabled = mDeps.isAtLeastU()
                 && mDeps.isFeatureEnabled(context, REQUEST_RESTRICTED_WIFI);
+        mBackgroundFirewallChainEnabled = mDeps.isAtLeastV() && mDeps.isFeatureNotChickenedOut(
+                context, ConnectivityFlags.BACKGROUND_FIREWALL_CHAIN);
         mCarrierPrivilegeAuthenticator = mDeps.makeCarrierPrivilegeAuthenticator(
                 mContext, mTelephonyManager, mRequestRestrictedWifiEnabled,
                 this::handleUidCarrierPrivilegesLost, mHandler);
@@ -4148,6 +4151,9 @@
         pw.println();
         pw.println("Multicast routing supported: " +
                 (mMulticastRoutingCoordinatorService != null));
+
+        pw.println();
+        pw.println("Background firewall chain enabled: " + mBackgroundFirewallChainEnabled);
     }
 
     private void dumpNetworks(IndentingPrintWriter pw) {
@@ -13517,6 +13523,12 @@
     public void setUidFirewallRule(final int chain, final int uid, final int rule) {
         enforceNetworkStackOrSettingsPermission();
 
+        if (chain == FIREWALL_CHAIN_BACKGROUND && !mBackgroundFirewallChainEnabled) {
+            Log.i(TAG, "Ignoring operation setUidFirewallRule on the background chain because the"
+                    + " feature is disabled.");
+            return;
+        }
+
         // There are only two type of firewall rule: FIREWALL_RULE_ALLOW or FIREWALL_RULE_DENY
         int firewallRule = getFirewallRuleType(chain, rule);
 
@@ -13589,6 +13601,12 @@
     public void setFirewallChainEnabled(final int chain, final boolean enable) {
         enforceNetworkStackOrSettingsPermission();
 
+        if (chain == FIREWALL_CHAIN_BACKGROUND && !mBackgroundFirewallChainEnabled) {
+            Log.i(TAG, "Ignoring operation setFirewallChainEnabled on the background chain because"
+                    + " the feature is disabled.");
+            return;
+        }
+
         try {
             mBpfNetMaps.setChildChain(chain, enable);
         } catch (ServiceSpecificException e) {
@@ -13615,6 +13633,12 @@
     public void replaceFirewallChain(final int chain, final int[] uids) {
         enforceNetworkStackOrSettingsPermission();
 
+        if (chain == FIREWALL_CHAIN_BACKGROUND && !mBackgroundFirewallChainEnabled) {
+            Log.i(TAG, "Ignoring operation replaceFirewallChain on the background chain because"
+                    + " the feature is disabled.");
+            return;
+        }
+
         mBpfNetMaps.replaceUidChain(chain, uids);
     }
 
diff --git a/service/src/com/android/server/connectivity/ConnectivityFlags.java b/service/src/com/android/server/connectivity/ConnectivityFlags.java
index a55c683..176307d 100644
--- a/service/src/com/android/server/connectivity/ConnectivityFlags.java
+++ b/service/src/com/android/server/connectivity/ConnectivityFlags.java
@@ -42,6 +42,8 @@
     public static final String INGRESS_TO_VPN_ADDRESS_FILTERING =
             "ingress_to_vpn_address_filtering";
 
+    public static final String BACKGROUND_FIREWALL_CHAIN = "background_firewall_chain";
+
     private boolean mNoRematchAllRequestsOnRegister;
 
     /**
diff --git a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
index 48af9fa..21dbb45 100644
--- a/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveStatsTracker.java
@@ -29,6 +29,7 @@
 import android.net.TelephonyNetworkSpecifier;
 import android.net.TransportInfo;
 import android.net.wifi.WifiInfo;
+import android.os.Build;
 import android.os.Handler;
 import android.os.SystemClock;
 import android.telephony.SubscriptionInfo;
@@ -39,6 +40,8 @@
 import android.util.SparseArray;
 import android.util.SparseIntArray;
 
+import androidx.annotation.RequiresApi;
+
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.metrics.DailykeepaliveInfoReported;
 import com.android.metrics.DurationForNumOfKeepalive;
@@ -279,6 +282,7 @@
          *
          * @param dailyKeepaliveInfoReported the proto to write to statsD.
          */
+        @RequiresApi(Build.VERSION_CODES.TIRAMISU)
         public void writeStats(DailykeepaliveInfoReported dailyKeepaliveInfoReported) {
             ConnectivityStatsLog.write(
                     ConnectivityStatsLog.DAILY_KEEPALIVE_INFO_REPORTED,
diff --git a/service/src/com/android/server/connectivity/SatelliteAccessController.java b/service/src/com/android/server/connectivity/SatelliteAccessController.java
index b53abce..2cdc932 100644
--- a/service/src/com/android/server/connectivity/SatelliteAccessController.java
+++ b/service/src/com/android/server/connectivity/SatelliteAccessController.java
@@ -20,7 +20,10 @@
 import android.annotation.NonNull;
 import android.app.role.OnRoleHoldersChangedListener;
 import android.app.role.RoleManager;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.os.Handler;
@@ -49,7 +52,6 @@
     private final Context mContext;
     private final Dependencies mDeps;
     private final DefaultMessageRoleListener mDefaultMessageRoleListener;
-    private final UserManager mUserManager;
     private final Consumer<Set<Integer>> mCallback;
     private final Handler mConnectivityServiceHandler;
 
@@ -114,7 +116,6 @@
             @NonNull final Handler connectivityServiceInternalHandler) {
         mContext = c;
         mDeps = deps;
-        mUserManager = mContext.getSystemService(UserManager.class);
         mDefaultMessageRoleListener = new DefaultMessageRoleListener();
         mCallback = callback;
         mConnectivityServiceHandler = connectivityServiceInternalHandler;
@@ -165,9 +166,6 @@
     }
 
     // on Role sms change triggered by OnRoleHoldersChangedListener()
-    // TODO(b/326373613): using UserLifecycleListener, callback to be received when user removed for
-    // user delete scenario. This to be used to update uid list and ML Layer request can also be
-    // updated.
     private void onRoleSmsChanged(@NonNull UserHandle userHandle) {
         int userId = userHandle.getIdentifier();
         if (userId == Process.INVALID_UID) {
@@ -184,9 +182,8 @@
                 mAllUsersSatelliteNetworkFallbackUidCache.get(userId, new ArraySet<>());
 
         Log.i(TAG, "currentUser : role_sms_packages: " + userId + " : " + packageNames);
-        final Set<Integer> newUidsForUser = !packageNames.isEmpty()
-                ? updateSatelliteNetworkFallbackUidListCache(packageNames, userHandle)
-                : new ArraySet<>();
+        final Set<Integer> newUidsForUser =
+                updateSatelliteNetworkFallbackUidListCache(packageNames, userHandle);
         Log.i(TAG, "satellite_fallback_uid: " + newUidsForUser);
 
         // on Role change, update the multilayer request at ConnectivityService with updated
@@ -197,6 +194,11 @@
 
         mAllUsersSatelliteNetworkFallbackUidCache.put(userId, newUidsForUser);
 
+        // Update all users fallback cache for user, send cs fallback to update ML request
+        reportSatelliteNetworkFallbackUids();
+    }
+
+    private void reportSatelliteNetworkFallbackUids() {
         // Merge all uids of multiple users available
         Set<Integer> mergedSatelliteNetworkFallbackUidCache = new ArraySet<>();
         for (int i = 0; i < mAllUsersSatelliteNetworkFallbackUidCache.size(); i++) {
@@ -210,27 +212,48 @@
         mCallback.accept(mergedSatelliteNetworkFallbackUidCache);
     }
 
-    private List<String> getRoleSmsChangedPackageName(UserHandle userHandle) {
-        try {
-            return mDeps.getRoleHoldersAsUser(RoleManager.ROLE_SMS, userHandle);
-        } catch (RuntimeException e) {
-            Log.wtf(TAG, "Could not get package name at role sms change update due to: " + e);
-            return null;
-        }
-    }
-
-    /** Register OnRoleHoldersChangedListener */
     public void start() {
         mConnectivityServiceHandler.post(this::updateAllUserRoleSmsUids);
+
+        // register sms OnRoleHoldersChangedListener
         mDefaultMessageRoleListener.register();
+
+        // Monitor for User removal intent, to update satellite fallback uids.
+        IntentFilter userRemovedFilter = new IntentFilter(Intent.ACTION_USER_REMOVED);
+        mContext.registerReceiver(new BroadcastReceiver() {
+            @Override
+            public void onReceive(Context context, Intent intent) {
+                final String action = intent.getAction();
+                if (Intent.ACTION_USER_REMOVED.equals(action)) {
+                    final UserHandle userHandle = intent.getParcelableExtra(Intent.EXTRA_USER);
+                    if (userHandle == null) return;
+                    updateSatelliteFallbackUidListOnUserRemoval(userHandle.getIdentifier());
+                } else {
+                    Log.wtf(TAG, "received unexpected intent: " + action);
+                }
+            }
+        }, userRemovedFilter, null, mConnectivityServiceHandler);
+
     }
 
     private void updateAllUserRoleSmsUids() {
-        List<UserHandle> existingUsers = mUserManager.getUserHandles(true /* excludeDying */);
+        UserManager userManager = mContext.getSystemService(UserManager.class);
+        // get existing user handles of available users
+        List<UserHandle> existingUsers = userManager.getUserHandles(true /*excludeDying*/);
+
         // Iterate through the user handles and obtain their uids with role sms and satellite
         // communication permission
+        Log.i(TAG, "existing users: " + existingUsers);
         for (UserHandle userHandle : existingUsers) {
             onRoleSmsChanged(userHandle);
         }
     }
+
+    private void updateSatelliteFallbackUidListOnUserRemoval(int userIdRemoved) {
+        Log.i(TAG, "user id removed:" + userIdRemoved);
+        if (mAllUsersSatelliteNetworkFallbackUidCache.contains(userIdRemoved)) {
+            mAllUsersSatelliteNetworkFallbackUidCache.remove(userIdRemoved);
+            reportSatelliteNetworkFallbackUids();
+        }
+    }
 }
diff --git a/staticlibs/Android.bp b/staticlibs/Android.bp
index f7b42a6..ede6d3f 100644
--- a/staticlibs/Android.bp
+++ b/staticlibs/Android.bp
@@ -124,6 +124,8 @@
     ],
 }
 
+// The net-utils-device-common-bpf library requires the callers to contain
+// net-utils-device-common-struct-base.
 java_library {
     name: "net-utils-device-common-bpf",
     srcs: [
@@ -133,9 +135,7 @@
         "device/com/android/net/module/util/BpfUtils.java",
         "device/com/android/net/module/util/IBpfMap.java",
         "device/com/android/net/module/util/JniUtil.java",
-        "device/com/android/net/module/util/Struct.java",
         "device/com/android/net/module/util/TcUtils.java",
-        "framework/com/android/net/module/util/HexDump.java",
     ],
     sdk_version: "module_current",
     min_sdk_version: "30",
@@ -146,6 +146,7 @@
     libs: [
         "androidx.annotation_annotation",
         "framework-connectivity.stubs.module_lib",
+        "net-utils-device-common-struct-base",
     ],
     apex_available: [
         "com.android.tethering",
@@ -158,12 +159,9 @@
 }
 
 java_library {
-    name: "net-utils-device-common-struct",
+    name: "net-utils-device-common-struct-base",
     srcs: [
-        "device/com/android/net/module/util/Ipv6Utils.java",
-        "device/com/android/net/module/util/PacketBuilder.java",
         "device/com/android/net/module/util/Struct.java",
-        "device/com/android/net/module/util/structs/*.java",
     ],
     sdk_version: "module_current",
     min_sdk_version: "30",
@@ -176,6 +174,7 @@
     ],
     libs: [
         "androidx.annotation_annotation",
+        "framework-annotations-lib", // Required by InetAddressUtils.java
         "framework-connectivity.stubs.module_lib",
     ],
     apex_available: [
@@ -188,26 +187,30 @@
     },
 }
 
-// The net-utils-multicast-forwarding-structs library requires the callers to
-// contain net-utils-device-common-bpf.
+// The net-utils-device-common-struct library requires the callers to contain
+// net-utils-device-common-struct-base.
 java_library {
-    name: "net-utils-multicast-forwarding-structs",
+    name: "net-utils-device-common-struct",
     srcs: [
-        "device/com/android/net/module/util/structs/StructMf6cctl.java",
-        "device/com/android/net/module/util/structs/StructMif6ctl.java",
-        "device/com/android/net/module/util/structs/StructMrt6Msg.java",
+        "device/com/android/net/module/util/Ipv6Utils.java",
+        "device/com/android/net/module/util/PacketBuilder.java",
+        "device/com/android/net/module/util/structs/*.java",
     ],
     sdk_version: "module_current",
     min_sdk_version: "30",
     visibility: [
         "//packages/modules/Connectivity:__subpackages__",
+        "//packages/modules/NetworkStack:__subpackages__",
     ],
     libs: [
-        // Only Struct.java is needed from "net-utils-device-common-bpf"
-        "net-utils-device-common-bpf",
+        "androidx.annotation_annotation",
+        "framework-annotations-lib", // Required by IpUtils.java
+        "framework-connectivity.stubs.module_lib",
+        "net-utils-device-common-struct-base",
     ],
     apex_available: [
         "com.android.tethering",
+        "//apex_available:platform",
     ],
     lint: {
         strict_updatability_linting: true,
@@ -216,7 +219,7 @@
 }
 
 // The net-utils-device-common-netlink library requires the callers to contain
-// net-utils-device-common-struct.
+// net-utils-device-common-struct and net-utils-device-common-struct-base.
 java_library {
     name: "net-utils-device-common-netlink",
     srcs: [
@@ -235,6 +238,7 @@
         // statically link here because callers of this library might already have a static
         // version linked.
         "net-utils-device-common-struct",
+        "net-utils-device-common-struct-base",
     ],
     apex_available: [
         "com.android.tethering",
@@ -247,7 +251,7 @@
 }
 
 // The net-utils-device-common-ip library requires the callers to contain
-// net-utils-device-common-struct.
+// net-utils-device-common-struct and net-utils-device-common-struct-base.
 java_library {
     // TODO : this target should probably be folded into net-utils-device-common
     name: "net-utils-device-common-ip",
diff --git a/staticlibs/device/com/android/net/module/util/DeviceConfigUtils.java b/staticlibs/device/com/android/net/module/util/DeviceConfigUtils.java
index 5b7cbb8..0426ace 100644
--- a/staticlibs/device/com/android/net/module/util/DeviceConfigUtils.java
+++ b/staticlibs/device/com/android/net/module/util/DeviceConfigUtils.java
@@ -17,6 +17,7 @@
 package com.android.net.module.util;
 
 import static android.content.pm.PackageManager.MATCH_SYSTEM_ONLY;
+import static android.provider.DeviceConfig.NAMESPACE_CAPTIVEPORTALLOGIN;
 import static android.provider.DeviceConfig.NAMESPACE_CONNECTIVITY;
 import static android.provider.DeviceConfig.NAMESPACE_TETHERING;
 
@@ -203,6 +204,29 @@
                 () -> getTetheringModuleVersion(context));
     }
 
+    /**
+     * Check whether or not one specific experimental feature for a particular namespace from
+     * {@link DeviceConfig} is enabled by comparing module package version
+     * with current version of property. If this property version is valid, the corresponding
+     * experimental feature would be enabled, otherwise disabled.
+     *
+     * This is useful to ensure that if a module install is rolled back, flags are not left fully
+     * rolled out on a version where they have not been well tested.
+     *
+     * If the feature is disabled by default and enabled by flag push, this method should be used.
+     * If the feature is enabled by default and disabled by flag push (kill switch),
+     * {@link #isCaptivePortalLoginFeatureNotChickenedOut(Context, String)} should be used.
+     *
+     * @param context The global context information about an app environment.
+     * @param name The name of the property to look up.
+     * @return true if this feature is enabled, or false if disabled.
+     */
+    public static boolean isCaptivePortalLoginFeatureEnabled(@NonNull Context context,
+            @NonNull String name) {
+        return isFeatureEnabled(NAMESPACE_CAPTIVEPORTALLOGIN, name, false /* defaultEnabled */,
+                () -> getPackageVersion(context));
+    }
+
     private static boolean isFeatureEnabled(@NonNull String namespace,
             String name, boolean defaultEnabled, Supplier<Long> packageVersionSupplier) {
         final int flagValue = getDeviceConfigPropertyInt(namespace, name, 0 /* default value */);
diff --git a/staticlibs/device/com/android/net/module/util/netlink/NdOption.java b/staticlibs/device/com/android/net/module/util/netlink/NdOption.java
index defc88a..4f58380 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/NdOption.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/NdOption.java
@@ -67,6 +67,9 @@
             case StructNdOptRdnss.TYPE:
                 return StructNdOptRdnss.parse(buf);
 
+            case StructNdOptPio.TYPE:
+                return StructNdOptPio.parse(buf);
+
             default:
                 int newPosition = Math.min(buf.limit(), buf.position() + length * 8);
                 buf.position(newPosition);
diff --git a/staticlibs/device/com/android/net/module/util/netlink/NduseroptMessage.java b/staticlibs/device/com/android/net/module/util/netlink/NduseroptMessage.java
index 2e9a99b..586f19d 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/NduseroptMessage.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/NduseroptMessage.java
@@ -156,8 +156,9 @@
 
     @Override
     public String toString() {
-        return String.format("Nduseroptmsg(%d, %d, %d, %d, %d, %s)",
+        return String.format("Nduseroptmsg(family:%d, opts_len:%d, ifindex:%d, icmp_type:%d, "
+                + "icmp_code:%d, srcaddr: %s, %s)",
                 family, opts_len, ifindex, Byte.toUnsignedInt(icmp_type),
-                Byte.toUnsignedInt(icmp_code), srcaddr.getHostAddress());
+                Byte.toUnsignedInt(icmp_code), srcaddr.getHostAddress(), option);
     }
 }
diff --git a/staticlibs/device/com/android/net/module/util/netlink/StructNdOptPio.java b/staticlibs/device/com/android/net/module/util/netlink/StructNdOptPio.java
new file mode 100644
index 0000000..65541eb
--- /dev/null
+++ b/staticlibs/device/com/android/net/module/util/netlink/StructNdOptPio.java
@@ -0,0 +1,120 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.net.module.util.netlink;
+
+import android.net.IpPrefix;
+import android.util.Log;
+
+import androidx.annotation.NonNull;
+
+import com.android.net.module.util.HexDump;
+import com.android.net.module.util.Struct;
+import com.android.net.module.util.structs.PrefixInformationOption;
+
+import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
+import java.util.Objects;
+
+/**
+ * The Prefix Information Option. RFC 4861.
+ *
+ * 0                   1                   2                   3
+ * 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * |     Type      |    Length     | Prefix Length |L|A|R|P| Rsvd1 |
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * |                         Valid Lifetime                        |
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * |                       Preferred Lifetime                      |
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * |                           Reserved2                           |
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * |                                                               |
+ * +                                                               +
+ * |                                                               |
+ * +                            Prefix                             +
+ * |                                                               |
+ * +                                                               +
+ * |                                                               |
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ */
+public class StructNdOptPio extends NdOption {
+    private static final String TAG = StructNdOptPio.class.getSimpleName();
+    public static final int TYPE = 3;
+    public static final byte LENGTH = 4; // Length in 8-byte units
+
+    public final byte flags;
+    public final long preferred;
+    public final long valid;
+    @NonNull
+    public final IpPrefix prefix;
+
+    public StructNdOptPio(byte flags, long preferred, long valid, @NonNull final IpPrefix prefix) {
+        super((byte) TYPE, LENGTH);
+        this.prefix = Objects.requireNonNull(prefix, "prefix must not be null");
+        this.flags = flags;
+        this.preferred = preferred;
+        this.valid = valid;
+    }
+
+    /**
+     * Parses a PrefixInformation option from a {@link ByteBuffer}.
+     *
+     * @param buf The buffer from which to parse the option. The buffer's byte order must be
+     *            {@link java.nio.ByteOrder#BIG_ENDIAN}.
+     * @return the parsed option, or {@code null} if the option could not be parsed successfully.
+     */
+    public static StructNdOptPio parse(@NonNull ByteBuffer buf) {
+        if (buf == null || buf.remaining() < LENGTH * 8) return null;
+        try {
+            final PrefixInformationOption pio = Struct.parse(PrefixInformationOption.class, buf);
+            if (pio.type != TYPE) {
+                throw new IllegalArgumentException("Invalid type " + pio.type);
+            }
+            if (pio.length != LENGTH) {
+                throw new IllegalArgumentException("Invalid length " + pio.length);
+            }
+            return new StructNdOptPio(pio.flags, pio.preferredLifetime, pio.validLifetime,
+                    pio.getIpPrefix());
+        } catch (IllegalArgumentException | BufferUnderflowException e) {
+            // Not great, but better than throwing an exception that might crash the caller.
+            // Convention in this package is that null indicates that the option was truncated
+            // or malformed, so callers must already handle it.
+            Log.d(TAG, "Invalid PIO option: " + e);
+            return null;
+        }
+    }
+
+    protected void writeToByteBuffer(ByteBuffer buf) {
+        buf.put(PrefixInformationOption.build(prefix, flags, valid, preferred));
+    }
+
+    /** Outputs the wire format of the option to a new big-endian ByteBuffer. */
+    public ByteBuffer toByteBuffer() {
+        final ByteBuffer buf = ByteBuffer.allocate(Struct.getSize(PrefixInformationOption.class));
+        writeToByteBuffer(buf);
+        buf.flip();
+        return buf;
+    }
+
+    @Override
+    @NonNull
+    public String toString() {
+        return String.format("NdOptPio(flags:%s, preferred lft:%s, valid lft:%s, prefix:%s)",
+                HexDump.toHexString(flags), preferred, valid, prefix);
+    }
+}
diff --git a/staticlibs/device/com/android/net/module/util/structs/PrefixInformationOption.java b/staticlibs/device/com/android/net/module/util/structs/PrefixInformationOption.java
index 49d7654..bbbe571 100644
--- a/staticlibs/device/com/android/net/module/util/structs/PrefixInformationOption.java
+++ b/staticlibs/device/com/android/net/module/util/structs/PrefixInformationOption.java
@@ -21,11 +21,16 @@
 import android.net.IpPrefix;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.net.module.util.Struct;
+import com.android.net.module.util.Struct.Computed;
 import com.android.net.module.util.Struct.Field;
 import com.android.net.module.util.Struct.Type;
 
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 
@@ -71,7 +76,11 @@
     @Field(order = 7, type = Type.ByteArray, arraysize = 16)
     public final byte[] prefix;
 
-    PrefixInformationOption(final byte type, final byte length, final byte prefixLen,
+    @Computed
+    private final IpPrefix mIpPrefix;
+
+    @VisibleForTesting
+    public PrefixInformationOption(final byte type, final byte length, final byte prefixLen,
             final byte flags, final long validLifetime, final long preferredLifetime,
             final int reserved, @NonNull final byte[] prefix) {
         this.type = type;
@@ -82,6 +91,23 @@
         this.preferredLifetime = preferredLifetime;
         this.reserved = reserved;
         this.prefix = prefix;
+
+        try {
+            final Inet6Address addr = (Inet6Address) InetAddress.getByAddress(prefix);
+            mIpPrefix = new IpPrefix(addr, prefixLen);
+        } catch (UnknownHostException | ClassCastException e) {
+            // UnknownHostException should never happen unless prefix is null.
+            // ClassCastException can occur when prefix is an IPv6 mapped IPv4 address.
+            // Both scenarios should throw an exception in the context of Struct#parse().
+            throw new IllegalArgumentException(e);
+        }
+    }
+
+    /**
+     * Return the prefix {@link IpPrefix} included in the PIO.
+     */
+    public IpPrefix getIpPrefix() {
+        return mIpPrefix;
     }
 
     /**
diff --git a/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java b/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
index 54ce01e..7066131 100644
--- a/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
+++ b/staticlibs/framework/com/android/net/module/util/NetworkCapabilitiesUtils.java
@@ -39,12 +39,13 @@
 import static android.net.NetworkCapabilities.TRANSPORT_BLUETOOTH;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_ETHERNET;
+import static android.net.NetworkCapabilities.TRANSPORT_SATELLITE;
 import static android.net.NetworkCapabilities.TRANSPORT_USB;
 import static android.net.NetworkCapabilities.TRANSPORT_VPN;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI_AWARE;
 
-import static com.android.net.module.util.BitUtils.packBitList;
+import static com.android.net.module.util.BitUtils.packBits;
 import static com.android.net.module.util.BitUtils.unpackBits;
 
 import android.annotation.NonNull;
@@ -75,8 +76,8 @@
         TRANSPORT_BLUETOOTH,
         TRANSPORT_WIFI,
         TRANSPORT_ETHERNET,
-        TRANSPORT_USB
-
+        TRANSPORT_USB,
+        TRANSPORT_SATELLITE
         // Notably, TRANSPORT_TEST is not in this list as any network that has TRANSPORT_TEST and
         // one of the above transports should be counted as that transport, to keep tests as
         // realistic as possible.
@@ -88,41 +89,41 @@
       * and {@code FORCE_RESTRICTED_CAPABILITIES}.
      */
     @VisibleForTesting
-    public static final long RESTRICTED_CAPABILITIES = packBitList(
-            NET_CAPABILITY_BIP,
-            NET_CAPABILITY_CBS,
-            NET_CAPABILITY_DUN,
-            NET_CAPABILITY_EIMS,
-            NET_CAPABILITY_ENTERPRISE,
-            NET_CAPABILITY_FOTA,
-            NET_CAPABILITY_IA,
-            NET_CAPABILITY_IMS,
-            NET_CAPABILITY_MCX,
-            NET_CAPABILITY_RCS,
-            NET_CAPABILITY_VEHICLE_INTERNAL,
-            NET_CAPABILITY_VSIM,
-            NET_CAPABILITY_XCAP,
-            NET_CAPABILITY_MMTEL);
+    public static final long RESTRICTED_CAPABILITIES =
+            (1L << NET_CAPABILITY_BIP) |
+            (1L << NET_CAPABILITY_CBS) |
+            (1L << NET_CAPABILITY_DUN) |
+            (1L << NET_CAPABILITY_EIMS) |
+            (1L << NET_CAPABILITY_ENTERPRISE) |
+            (1L << NET_CAPABILITY_FOTA) |
+            (1L << NET_CAPABILITY_IA) |
+            (1L << NET_CAPABILITY_IMS) |
+            (1L << NET_CAPABILITY_MCX) |
+            (1L << NET_CAPABILITY_RCS) |
+            (1L << NET_CAPABILITY_VEHICLE_INTERNAL) |
+            (1L << NET_CAPABILITY_VSIM) |
+            (1L << NET_CAPABILITY_XCAP) |
+            (1L << NET_CAPABILITY_MMTEL);
 
     /**
      * Capabilities that force network to be restricted.
      * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted}.
      */
-    private static final long FORCE_RESTRICTED_CAPABILITIES = packBitList(
-            NET_CAPABILITY_ENTERPRISE,
-            NET_CAPABILITY_OEM_PAID,
-            NET_CAPABILITY_OEM_PRIVATE);
+    private static final long FORCE_RESTRICTED_CAPABILITIES =
+            (1L << NET_CAPABILITY_ENTERPRISE) |
+            (1L << NET_CAPABILITY_OEM_PAID) |
+            (1L << NET_CAPABILITY_OEM_PRIVATE);
 
     /**
      * Capabilities that suggest that a network is unrestricted.
      * See {@code NetworkCapabilities#maybeMarkCapabilitiesRestricted}.
      */
     @VisibleForTesting
-    public static final long UNRESTRICTED_CAPABILITIES = packBitList(
-            NET_CAPABILITY_INTERNET,
-            NET_CAPABILITY_MMS,
-            NET_CAPABILITY_SUPL,
-            NET_CAPABILITY_WIFI_P2P);
+    public static final long UNRESTRICTED_CAPABILITIES =
+            (1L << NET_CAPABILITY_INTERNET) |
+            (1L << NET_CAPABILITY_MMS) |
+            (1L << NET_CAPABILITY_SUPL) |
+            (1L << NET_CAPABILITY_WIFI_P2P);
 
     /**
      * Get a transport that can be used to classify a network when displaying its info to users.
@@ -158,28 +159,33 @@
      *
      * @return {@code true} if the network should be restricted.
      */
-    // TODO: Use packBits(nc.getCapabilities()) to check more easily using bit masks.
     public static boolean inferRestrictedCapability(NetworkCapabilities nc) {
+        return inferRestrictedCapability(packBits(nc.getCapabilities()));
+    }
+
+    /**
+     * Infers that all the capabilities it provides are typically provided by restricted networks
+     * or not.
+     *
+     * @param capabilities see {@link NetworkCapabilities#getCapabilities()}
+     *
+     * @return {@code true} if the network should be restricted.
+     */
+    public static boolean inferRestrictedCapability(long capabilities) {
         // Check if we have any capability that forces the network to be restricted.
-        for (int capability : unpackBits(FORCE_RESTRICTED_CAPABILITIES)) {
-            if (nc.hasCapability(capability)) {
-                return true;
-            }
+        if ((capabilities & FORCE_RESTRICTED_CAPABILITIES) != 0) {
+            return true;
         }
 
         // Verify there aren't any unrestricted capabilities.  If there are we say
         // the whole thing is unrestricted unless it is forced to be restricted.
-        for (int capability : unpackBits(UNRESTRICTED_CAPABILITIES)) {
-            if (nc.hasCapability(capability)) {
-                return false;
-            }
+        if ((capabilities & UNRESTRICTED_CAPABILITIES) != 0) {
+            return false;
         }
 
         // Must have at least some restricted capabilities.
-        for (int capability : unpackBits(RESTRICTED_CAPABILITIES)) {
-            if (nc.hasCapability(capability)) {
-                return true;
-            }
+        if ((capabilities & RESTRICTED_CAPABILITIES) != 0) {
+            return true;
         }
         return false;
     }
diff --git a/staticlibs/tests/unit/Android.bp b/staticlibs/tests/unit/Android.bp
index 4c226cc..fa466f8 100644
--- a/staticlibs/tests/unit/Android.bp
+++ b/staticlibs/tests/unit/Android.bp
@@ -25,6 +25,7 @@
         "net-utils-device-common-async",
         "net-utils-device-common-bpf",
         "net-utils-device-common-ip",
+        "net-utils-device-common-struct-base",
         "net-utils-device-common-wear",
     ],
     libs: [
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/DeviceConfigUtilsTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/DeviceConfigUtilsTest.java
index f32337d..9fb61d9 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/DeviceConfigUtilsTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/DeviceConfigUtilsTest.java
@@ -17,6 +17,7 @@
 package com.android.net.module.util;
 
 import static android.content.pm.PackageManager.MATCH_SYSTEM_ONLY;
+import static android.provider.DeviceConfig.NAMESPACE_CAPTIVEPORTALLOGIN;
 import static android.provider.DeviceConfig.NAMESPACE_CONNECTIVITY;
 import static android.provider.DeviceConfig.NAMESPACE_TETHERING;
 
@@ -233,8 +234,12 @@
                 TEST_EXPERIMENT_FLAG));
         doReturn(TEST_FLAG_VALUE_STRING).when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
+        doReturn(TEST_FLAG_VALUE_STRING).when(() -> DeviceConfig.getProperty(
+                NAMESPACE_CAPTIVEPORTALLOGIN, TEST_EXPERIMENT_FLAG));
         assertTrue(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertTrue(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertTrue(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
     }
     @Test
     public void testIsFeatureEnabledFeatureDefaultDisabled() throws Exception {
@@ -242,8 +247,12 @@
                 TEST_EXPERIMENT_FLAG));
         doReturn(null).when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
+        doReturn(null).when(() -> DeviceConfig.getProperty(NAMESPACE_CAPTIVEPORTALLOGIN,
+                TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertFalse(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
 
         // If the flag is unset, package info is not queried
         verify(mContext, never()).getPackageManager();
@@ -257,8 +266,12 @@
                 TEST_EXPERIMENT_FLAG));
         doReturn("1").when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
+        doReturn("1").when(() -> DeviceConfig.getProperty(NAMESPACE_CAPTIVEPORTALLOGIN,
+                TEST_EXPERIMENT_FLAG));
         assertTrue(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertTrue(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertTrue(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
 
         // If the feature is force enabled, package info is not queried
         verify(mContext, never()).getPackageManager();
@@ -272,8 +285,12 @@
                 TEST_EXPERIMENT_FLAG));
         doReturn("-1").when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
+        doReturn("-1").when(() -> DeviceConfig.getProperty(NAMESPACE_CAPTIVEPORTALLOGIN,
+                TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertFalse(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
 
         // If the feature is force disabled, package info is not queried
         verify(mContext, never()).getPackageManager();
@@ -290,24 +307,36 @@
                 TEST_EXPERIMENT_FLAG));
         doReturn("1").when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
+        doReturn("1").when(() -> DeviceConfig.getProperty(NAMESPACE_CAPTIVEPORTALLOGIN,
+                TEST_EXPERIMENT_FLAG));
         assertTrue(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertTrue(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertTrue(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
 
         // Feature should be disabled by flag value "999999999".
         doReturn("999999999").when(() -> DeviceConfig.getProperty(NAMESPACE_CONNECTIVITY,
                 TEST_EXPERIMENT_FLAG));
         doReturn("999999999").when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
+        doReturn("999999999").when(() -> DeviceConfig.getProperty(NAMESPACE_CAPTIVEPORTALLOGIN,
+                TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertFalse(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
 
         // If the flag is not set feature is disabled
         doReturn(null).when(() -> DeviceConfig.getProperty(NAMESPACE_CONNECTIVITY,
                 TEST_EXPERIMENT_FLAG));
         doReturn(null).when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
+        doReturn(null).when(() -> DeviceConfig.getProperty(NAMESPACE_CAPTIVEPORTALLOGIN,
+                TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertFalse(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
     }
 
     @Test
@@ -320,9 +349,13 @@
                 NAMESPACE_CONNECTIVITY, TEST_EXPERIMENT_FLAG));
         doReturn("0").when(() -> DeviceConfig.getProperty(
                 NAMESPACE_TETHERING, TEST_EXPERIMENT_FLAG));
+        doReturn("0").when(() -> DeviceConfig.getProperty(
+                NAMESPACE_CAPTIVEPORTALLOGIN, TEST_EXPERIMENT_FLAG));
 
         assertFalse(DeviceConfigUtils.isNetworkStackFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
         assertFalse(DeviceConfigUtils.isTetheringFeatureEnabled(mContext, TEST_EXPERIMENT_FLAG));
+        assertFalse(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
 
         doReturn(TEST_FLAG_VALUE_STRING).when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
@@ -343,6 +376,21 @@
     }
 
     @Test
+    public void testIsCaptivePortalLoginFeatureEnabledCaching() throws Exception {
+        doReturn(TEST_FLAG_VALUE_STRING).when(() -> DeviceConfig.getProperty(
+                NAMESPACE_CAPTIVEPORTALLOGIN, TEST_EXPERIMENT_FLAG));
+        assertTrue(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
+        assertTrue(DeviceConfigUtils.isCaptivePortalLoginFeatureEnabled(mContext,
+                TEST_EXPERIMENT_FLAG));
+
+        // Package info is only queried once
+        verify(mContext, times(1)).getPackageManager();
+        verify(mContext, times(1)).getPackageName();
+        verify(mPm, times(1)).getPackageInfo(anyString(), anyInt());
+    }
+
+    @Test
     public void testIsTetheringFeatureEnabledCaching() throws Exception {
         doReturn(TEST_FLAG_VALUE_STRING).when(() -> DeviceConfig.getProperty(NAMESPACE_TETHERING,
                 TEST_EXPERIMENT_FLAG));
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NduseroptMessageTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NduseroptMessageTest.java
index 4fc5ec2..d1df3a6 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NduseroptMessageTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NduseroptMessageTest.java
@@ -252,8 +252,10 @@
     public void testToString() {
         NduseroptMessage msg = parseNduseroptMessage(toBuffer(MSG_PREF64));
         assertNotNull(msg);
-        assertEquals("Nduseroptmsg(10, 16, 1431655765, 134, 0, fe80:2:3:4:5:6:7:8%1431655765)",
-                msg.toString());
+        final String expected = "Nduseroptmsg(family:10, opts_len:16, ifindex:1431655765, "
+                + "icmp_type:134, icmp_code:0, srcaddr: fe80:2:3:4:5:6:7:8%1431655765, "
+                + "NdOptPref64(2001:db8:3:4:5:6::/96, 10064))";
+        assertEquals(expected, msg.toString());
     }
 
     // Convenience method to parse a NduseroptMessage that's not part of a netlink message.
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/StructNdOptPioTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/StructNdOptPioTest.java
new file mode 100644
index 0000000..0d88829
--- /dev/null
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/StructNdOptPioTest.java
@@ -0,0 +1,205 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.net.module.util.netlink;
+
+import static com.android.net.module.util.NetworkStackConstants.INFINITE_LEASE;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+import android.net.IpPrefix;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.net.module.util.structs.PrefixInformationOption;
+
+import libcore.util.HexEncoding;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class StructNdOptPioTest {
+    private static final IpPrefix TEST_PREFIX = new IpPrefix("2a00:79e1:abc:f605::/64");
+    private static final byte TEST_PIO_FLAGS_P_UNSET = (byte) 0xC0; // L=1,A=1
+    private static final byte TEST_PIO_FLAGS_P_SET   = (byte) 0xD0; // L=1,A=1,P=1
+    private static final String PIO_BYTES =
+            "0304"                                // type=3, length=4
+            + "40"                                // prefix length=64
+            + "C0"                                // L=1,A=1
+            + "00278D00"                          // valid=259200
+            + "00093A80"                          // preferred=604800
+            + "00000000"                          // Reserved2
+            + "2A0079E10ABCF6050000000000000000"; // prefix=2a00:79e1:abc:f605::
+
+    private static final String PIO_WITH_P_FLAG_BYTES =
+            "0304"                                // type=3, length=4
+            + "40"                                // prefix length=64
+            + "D0"                                // L=1,A=1,P=1
+            + "00278D00"                          // valid=2592000
+            + "00093A80"                          // preferred=604800
+            + "00000000"                          // Reserved2
+            + "2A0079E10ABCF6050000000000000000"; // prefix=2a00:79e1:abc:f605::
+
+    private static final String PIO_WITH_P_FLAG_INFINITY_LIFETIME_BYTES =
+            "0304"                                // type=3, length=4
+            + "40"                                // prefix length=64
+            + "D0"                                // L=1,A=1,P=1
+            + "FFFFFFFF"                          // valid=infinity
+            + "FFFFFFFF"                          // preferred=infintiy
+            + "00000000"                          // Reserved2
+            + "2A0079E10ABCF6050000000000000000"; // prefix=2a00:79e1:abc:f605::
+
+    private static void assertPioOptMatches(final StructNdOptPio opt, int length, byte flags,
+            long preferred, long valid, final IpPrefix prefix) {
+        assertEquals(StructNdOptPio.TYPE, opt.type);
+        assertEquals(length, opt.length);
+        assertEquals(flags, opt.flags);
+        assertEquals(preferred, opt.preferred);
+        assertEquals(valid, opt.valid);
+        assertEquals(prefix, opt.prefix);
+    }
+
+    private static void assertToByteBufferMatches(final StructNdOptPio opt, final String expected) {
+        String actual = HexEncoding.encodeToString(opt.toByteBuffer().array());
+        assertEquals(expected, actual);
+    }
+
+    private static void doPioParsingTest(final String optionHexString, int length, byte flags,
+            long preferred, long valid, final IpPrefix prefix) {
+        final byte[] rawBytes = HexEncoding.decode(optionHexString);
+        final StructNdOptPio opt = StructNdOptPio.parse(ByteBuffer.wrap(rawBytes));
+        assertPioOptMatches(opt, length, flags, preferred, valid, prefix);
+        assertToByteBufferMatches(opt, optionHexString);
+    }
+
+    @Test
+    public void testParsingPioWithoutPFlag() {
+        doPioParsingTest(PIO_BYTES, 4 /* length */, TEST_PIO_FLAGS_P_UNSET,
+                604800 /* preferred */, 2592000 /* valid */, TEST_PREFIX);
+    }
+
+    @Test
+    public void testParsingPioWithPFlag() {
+        doPioParsingTest(PIO_WITH_P_FLAG_BYTES, 4 /* length */, TEST_PIO_FLAGS_P_SET,
+                604800 /* preferred */, 2592000 /* valid */, TEST_PREFIX);
+    }
+
+    @Test
+    public void testParsingPioWithPFlag_infinityLifetime() {
+        doPioParsingTest(PIO_WITH_P_FLAG_INFINITY_LIFETIME_BYTES, 4 /* length */,
+                TEST_PIO_FLAGS_P_SET,
+                Integer.toUnsignedLong(INFINITE_LEASE) /* preferred */,
+                Integer.toUnsignedLong(INFINITE_LEASE) /* valid */,
+                TEST_PREFIX);
+    }
+
+    @Test
+    public void testToByteBuffer() {
+        final StructNdOptPio pio =
+                new StructNdOptPio(TEST_PIO_FLAGS_P_UNSET, 604800 /* preferred */,
+                        2592000 /* valid */, TEST_PREFIX);
+        assertToByteBufferMatches(pio, PIO_BYTES);
+    }
+
+    @Test
+    public void testToByteBuffer_withPFlag() {
+        final StructNdOptPio pio =
+                new StructNdOptPio(TEST_PIO_FLAGS_P_SET, 604800 /* preferred */,
+                        2592000 /* valid */, TEST_PREFIX);
+        assertToByteBufferMatches(pio, PIO_WITH_P_FLAG_BYTES);
+    }
+
+    @Test
+    public void testToByteBuffer_infinityLifetime() {
+        final StructNdOptPio pio =
+                new StructNdOptPio(TEST_PIO_FLAGS_P_SET,
+                        Integer.toUnsignedLong(INFINITE_LEASE) /* preferred */,
+                        Integer.toUnsignedLong(INFINITE_LEASE) /* valid */, TEST_PREFIX);
+        assertToByteBufferMatches(pio, PIO_WITH_P_FLAG_INFINITY_LIFETIME_BYTES);
+    }
+
+    private static ByteBuffer makePioOption(byte type, byte length, byte prefixLen, byte flags,
+            long valid, long preferred, final byte[] prefix) {
+        final PrefixInformationOption pio = new PrefixInformationOption(type, length, prefixLen,
+                flags, valid, preferred, 0 /* reserved */, prefix);
+        return ByteBuffer.wrap(pio.writeToBytes(ByteOrder.BIG_ENDIAN));
+    }
+
+    @Test
+    public void testParsing_invalidOptionType() {
+        final ByteBuffer buf = makePioOption((byte) 24 /* wrong type:RIO */,
+                (byte) 4 /* length */, (byte) 64 /* prefixLen */, TEST_PIO_FLAGS_P_SET,
+                2592000 /* valid */, 604800 /* preferred */, TEST_PREFIX.getRawAddress());
+        assertNull(StructNdOptPio.parse(buf));
+    }
+
+    @Test
+    public void testParsing_invalidOptionLength() {
+        final ByteBuffer buf = makePioOption((byte) 24 /* wrong type:RIO */,
+                (byte) 3 /* wrong length */, (byte) 64 /* prefixLen */,
+                TEST_PIO_FLAGS_P_SET, 2592000 /* valid */, 604800 /* preferred */,
+                TEST_PREFIX.getRawAddress());
+        assertNull(StructNdOptPio.parse(buf));
+    }
+
+    @Test
+    public void testParsing_truncatedByteBuffer() {
+        final ByteBuffer buf = makePioOption((byte) 3 /* type */, (byte) 4 /* length */,
+                (byte) 64 /* prefixLen */, TEST_PIO_FLAGS_P_SET,
+                2592000 /* valid */, 604800 /* preferred */, TEST_PREFIX.getRawAddress());
+        final int len = buf.limit();
+        for (int i = 0; i < buf.limit() - 1; i++) {
+            buf.flip();
+            buf.limit(i);
+            assertNull("Option truncated to " + i + " bytes, should have returned null",
+                    StructNdOptPio.parse(buf));
+        }
+        buf.flip();
+        buf.limit(len);
+
+        final StructNdOptPio opt = StructNdOptPio.parse(buf);
+        assertPioOptMatches(opt, (byte) 4 /* length */, TEST_PIO_FLAGS_P_SET,
+                604800 /* preferred */, 2592000 /* valid */, TEST_PREFIX);
+    }
+
+    @Test
+    public void testParsing_invalidByteBufferLength() {
+        final ByteBuffer buf = makePioOption((byte) 3 /* type */, (byte) 4 /* length */,
+                (byte) 64 /* prefixLen */, TEST_PIO_FLAGS_P_SET,
+                2592000 /* valid */, 604800 /* preferred */, TEST_PREFIX.getRawAddress());
+        buf.limit(31); // less than 4 * 8
+        assertNull(StructNdOptPio.parse(buf));
+    }
+
+    @Test
+    public void testToString() {
+        final ByteBuffer buf = makePioOption((byte) 3 /* type */, (byte) 4 /* length */,
+                (byte) 64 /* prefixLen */, TEST_PIO_FLAGS_P_SET,
+                2592000 /* valid */, 604800 /* preferred */, TEST_PREFIX.getRawAddress());
+        final StructNdOptPio opt = StructNdOptPio.parse(buf);
+        final String expected = "NdOptPio"
+                + "(flags:D0, preferred lft:604800, valid lft:2592000,"
+                + " prefix:2a00:79e1:abc:f605::/64)";
+        assertEquals(expected, opt.toString());
+    }
+}
diff --git a/staticlibs/testutils/Android.bp b/staticlibs/testutils/Android.bp
index a8e5a69..9124ac0 100644
--- a/staticlibs/testutils/Android.bp
+++ b/staticlibs/testutils/Android.bp
@@ -40,6 +40,7 @@
         "net-utils-device-common-async",
         "net-utils-device-common-netlink",
         "net-utils-device-common-struct",
+        "net-utils-device-common-struct-base",
         "net-utils-device-common-wear",
         "modules-utils-build_system",
     ],
diff --git a/staticlibs/testutils/app/connectivitychecker/Android.bp b/staticlibs/testutils/app/connectivitychecker/Android.bp
index 5af8c14..394c6be 100644
--- a/staticlibs/testutils/app/connectivitychecker/Android.bp
+++ b/staticlibs/testutils/app/connectivitychecker/Android.bp
@@ -30,7 +30,6 @@
         "modules-utils-build_system",
         "net-tests-utils",
     ],
-    host_required: ["net-tests-utils-host-common"],
     lint: {
         strict_updatability_linting: true,
     },
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
new file mode 100644
index 0000000..93422ad
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
@@ -0,0 +1,227 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.testutils
+
+import android.net.ConnectivityManager
+import android.net.ConnectivityManager.NetworkCallback
+import android.net.Network
+import android.net.NetworkCapabilities
+import android.net.NetworkRequest
+import android.os.Handler
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.testutils.RecorderCallback.CallbackEntry
+import java.util.Collections
+import kotlin.test.fail
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * A rule to file [NetworkCallback]s to request or watch networks.
+ *
+ * The callbacks filed in test methods are automatically unregistered when the method completes.
+ */
+class AutoReleaseNetworkCallbackRule : NetworkCallbackHelper(), TestRule {
+    override fun apply(base: Statement, description: Description): Statement {
+        return RequestCellNetworkStatement(base, description)
+    }
+
+    private inner class RequestCellNetworkStatement(
+        private val base: Statement,
+        private val description: Description
+    ) : Statement() {
+        override fun evaluate() {
+            tryTest {
+                base.evaluate()
+            } cleanup {
+                unregisterAll()
+            }
+        }
+    }
+}
+
+/**
+ * Helps file [NetworkCallback]s to request or watch networks, keeping track of them for cleanup.
+ */
+open class NetworkCallbackHelper {
+    private val cm by lazy {
+        InstrumentationRegistry.getInstrumentation().context
+            .getSystemService(ConnectivityManager::class.java)
+            ?: fail("ConnectivityManager not found")
+    }
+    private val cbToCleanup = Collections.synchronizedSet(mutableSetOf<NetworkCallback>())
+    private var cellRequestCb: TestableNetworkCallback? = null
+
+    /**
+     * Convenience method to request a cell network, similarly to [requestNetwork].
+     *
+     * The rule will keep tract of a single cell network request, which can be unrequested manually
+     * using [unrequestCell].
+     */
+    fun requestCell(): Network {
+        if (cellRequestCb != null) {
+            fail("Cell network was already requested")
+        }
+        val cb = requestNetwork(
+            NetworkRequest.Builder()
+                .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR)
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                .build()
+        )
+        cellRequestCb = cb
+        return cb.expect<CallbackEntry.Available>(
+            errorMsg = "Cell network not available. " +
+                    "Please ensure the device has working mobile data."
+        ).network
+    }
+
+    /**
+     * Unrequest a cell network requested through [requestCell].
+     */
+    fun unrequestCell() {
+        val cb = cellRequestCb ?: fail("Cell network was not requested")
+        unregisterNetworkCallback(cb)
+        cellRequestCb = null
+    }
+
+    private fun addCallback(
+        cb: TestableNetworkCallback,
+        registrar: (TestableNetworkCallback) -> Unit
+    ): TestableNetworkCallback {
+        registrar(cb)
+        cbToCleanup.add(cb)
+        return cb
+    }
+
+    /**
+     * File a request for a Network.
+     *
+     * This will fail tests (throw) if the cell network cannot be obtained, or if it was already
+     * requested.
+     *
+     * Tests may call [unregisterNetworkCallback] once they are done using the returned [Network],
+     * otherwise it will be automatically unrequested after the test.
+     */
+    @JvmOverloads
+    fun requestNetwork(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler? = null
+    ) = addCallback(cb) {
+        if (handler == null) {
+            cm.requestNetwork(request, it)
+        } else {
+            cm.requestNetwork(request, it, handler)
+        }
+    }
+
+    /**
+     * Overload of [requestNetwork] that allows specifying a timeout.
+     */
+    @JvmOverloads
+    fun requestNetwork(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        timeoutMs: Int,
+    ) = addCallback(cb) { cm.requestNetwork(request, it, timeoutMs) }
+
+    /**
+     * File a callback for a NetworkRequest.
+     *
+     * This will fail tests (throw) if the cell network cannot be obtained, or if it was already
+     * requested.
+     *
+     * Tests may call [unregisterNetworkCallback] once they are done using the returned [Network],
+     * otherwise it will be automatically unrequested after the test.
+     */
+    @JvmOverloads
+    fun registerNetworkCallback(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback()
+    ) = addCallback(cb) { cm.registerNetworkCallback(request, it) }
+
+    /**
+     * @see ConnectivityManager.registerDefaultNetworkCallback
+     */
+    @JvmOverloads
+    fun registerDefaultNetworkCallback(
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler? = null
+    ) = addCallback(cb) {
+        if (handler == null) {
+            cm.registerDefaultNetworkCallback(it)
+        } else {
+            cm.registerDefaultNetworkCallback(it, handler)
+        }
+    }
+
+    /**
+     * @see ConnectivityManager.registerSystemDefaultNetworkCallback
+     */
+    @JvmOverloads
+    fun registerSystemDefaultNetworkCallback(
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.registerSystemDefaultNetworkCallback(it, handler) }
+
+    /**
+     * @see ConnectivityManager.registerDefaultNetworkCallbackForUid
+     */
+    @JvmOverloads
+    fun registerDefaultNetworkCallbackForUid(
+        uid: Int,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.registerDefaultNetworkCallbackForUid(uid, it, handler) }
+
+    /**
+     * @see ConnectivityManager.registerBestMatchingNetworkCallback
+     */
+    @JvmOverloads
+    fun registerBestMatchingNetworkCallback(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.registerBestMatchingNetworkCallback(request, it, handler) }
+
+    /**
+     * @see ConnectivityManager.requestBackgroundNetwork
+     */
+    @JvmOverloads
+    fun requestBackgroundNetwork(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.requestBackgroundNetwork(request, it, handler) }
+
+    /**
+     * Unregister a callback filed using registration methods in this class.
+     */
+    fun unregisterNetworkCallback(cb: NetworkCallback) {
+        cm.unregisterNetworkCallback(cb)
+        cbToCleanup.remove(cb)
+    }
+
+    /**
+     * Unregister all callbacks that were filed using registration methods in this class.
+     */
+    fun unregisterAll() {
+        cbToCleanup.forEach { cm.unregisterNetworkCallback(it) }
+        cbToCleanup.clear()
+        cellRequestCb = null
+    }
+}
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
index 8090d5b..3857810 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
@@ -86,6 +86,7 @@
         val callback = TestableNetworkCallback(timeoutMs = WIFI_CONNECT_TIMEOUT_MS)
         cm.registerNetworkCallback(NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_WIFI)
+                .addCapability(NET_CAPABILITY_INTERNET)
                 .build(), callback)
 
         return tryTest {
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
new file mode 100644
index 0000000..4185b05
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.testutils.com.android.testutils
+
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * A JUnit Rule that sets feature flags based on `@FeatureFlag` annotations.
+ *
+ * This rule enables dynamic control of feature flag states during testing.
+ *
+ * **Usage:**
+ * ```kotlin
+ * class MyTestClass {
+ *   @get:Rule
+ *   val setFeatureFlagsRule = SetFeatureFlagsRule(setFlagsMethod = (name, enabled) -> {
+ *     // Custom handling code.
+ *   })
+ *
+ *   // ... test methods with @FeatureFlag annotations
+ *   @FeatureFlag("FooBar1", true)
+ *   @FeatureFlag("FooBar2", false)
+ *   @Test
+ *   fun testFooBar() {}
+ * }
+ * ```
+ */
+class SetFeatureFlagsRule(val setFlagsMethod: (name: String, enabled: Boolean) -> Unit) : TestRule {
+    /**
+     * This annotation marks a test method as requiring a specific feature flag to be configured.
+     *
+     * Use this on test methods to dynamically control feature flag states during testing.
+     *
+     * @param name The name of the feature flag.
+     * @param enabled The desired state (true for enabled, false for disabled) of the feature flag.
+     */
+    @Target(AnnotationTarget.FUNCTION)
+    @Retention(AnnotationRetention.RUNTIME)
+    annotation class FeatureFlag(val name: String, val enabled: Boolean = true)
+
+    /**
+     * This method is the core of the rule, executed by the JUnit framework before each test method.
+     *
+     * It retrieves the test method's metadata.
+     * If any `@FeatureFlag` annotation is found, it passes every feature flag's name
+     * and enabled state into the user-specified lambda to apply custom actions.
+     */
+    override fun apply(base: Statement, description: Description): Statement {
+        return object : Statement() {
+            override fun evaluate() {
+                val testMethod = description.testClass.getMethod(description.methodName)
+                val featureFlagAnnotations = testMethod.getAnnotationsByType(
+                    FeatureFlag::class.java
+                )
+
+                for (featureFlagAnnotation in featureFlagAnnotations) {
+                    setFlagsMethod(featureFlagAnnotation.name, featureFlagAnnotation.enabled)
+                }
+
+                // Execute the test method, which includes methods annotated with
+                // @Before, @Test and @After.
+                base.evaluate()
+            }
+        }
+    }
+}
diff --git a/tests/cts/OWNERS b/tests/cts/OWNERS
index cb4ca59..acf506d 100644
--- a/tests/cts/OWNERS
+++ b/tests/cts/OWNERS
@@ -9,4 +9,6 @@
 # For incremental changes on EthernetManagerTest to increase coverage for existing behavior and for
 # testing bug fixes.
 per-file net/src/android/net/cts/EthernetManagerTest.kt = prohr@google.com #{LAST_RESORT_SUGGESTION}
+# Temporary ownership to develop APF CTS tests.
+per-file net/src/android/net/cts/ApfIntegrationTest.kt = prohr@google.com #{LAST_RESORT_SUGGESTION}
 
diff --git a/tests/cts/hostside/aidl/Android.bp b/tests/cts/hostside/aidl/Android.bp
index 18a5897..33761dc 100644
--- a/tests/cts/hostside/aidl/Android.bp
+++ b/tests/cts/hostside/aidl/Android.bp
@@ -21,9 +21,7 @@
     name: "CtsHostsideNetworkTestsAidl",
     sdk_version: "current",
     srcs: [
-        "com/android/cts/net/hostside/IMyService.aidl",
-        "com/android/cts/net/hostside/INetworkCallback.aidl",
-        "com/android/cts/net/hostside/INetworkStateObserver.aidl",
-        "com/android/cts/net/hostside/IRemoteSocketFactory.aidl",
+        "com/android/cts/net/hostside/*.aidl",
+        "com/android/cts/net/hostside/*.java",
     ],
 }
diff --git a/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl b/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl
index e7b2815..906024b 100644
--- a/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl
+++ b/tests/cts/hostside/aidl/com/android/cts/net/hostside/IMyService.aidl
@@ -19,11 +19,12 @@
 import android.app.job.JobInfo;
 
 import com.android.cts.net.hostside.INetworkCallback;
+import com.android.cts.net.hostside.NetworkCheckResult;
 
 interface IMyService {
     void registerBroadcastReceiver();
     int getCounters(String receiverName, String action);
-    String checkNetworkStatus();
+    NetworkCheckResult checkNetworkStatus(String customUrl);
     String getRestrictBackgroundStatus();
     void sendNotification(int notificationId, String notificationType);
     void registerNetworkCallback(in NetworkRequest request, in INetworkCallback cb);
diff --git a/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl b/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl
index 19198c5..8ef4659 100644
--- a/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl
+++ b/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl
@@ -16,8 +16,12 @@
 
 package com.android.cts.net.hostside;
 
+import android.net.NetworkInfo;
+
+import com.android.cts.net.hostside.NetworkCheckResult;
+
 interface INetworkStateObserver {
-    void onNetworkStateChecked(int resultCode, String resultData);
+    void onNetworkStateChecked(int resultCode, in NetworkCheckResult networkCheckResult);
 
     const int RESULT_SUCCESS_NETWORK_STATE_CHECKED = 0;
     const int RESULT_ERROR_UNEXPECTED_PROC_STATE = 1;
diff --git a/tests/cts/hostside/aidl/com/android/cts/net/hostside/NetworkCheckResult.aidl b/tests/cts/hostside/aidl/com/android/cts/net/hostside/NetworkCheckResult.aidl
new file mode 100644
index 0000000..cdd6b70
--- /dev/null
+++ b/tests/cts/hostside/aidl/com/android/cts/net/hostside/NetworkCheckResult.aidl
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.cts.net.hostside;
+
+import android.net.NetworkInfo;
+
+@JavaDerive(toString=true)
+parcelable NetworkCheckResult {
+   boolean connected;
+   String details;
+   NetworkInfo networkInfo;
+}
\ No newline at end of file
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
index 2ca8832..4437986 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
@@ -23,6 +23,7 @@
 import static android.net.ConnectivityManager.ACTION_RESTRICT_BACKGROUND_CHANGED;
 import static android.os.BatteryManager.BATTERY_PLUGGED_ANY;
 
+import static com.android.cts.net.arguments.InstrumentationArguments.ARG_CONNECTION_CHECK_CUSTOM_URL;
 import static com.android.cts.net.arguments.InstrumentationArguments.ARG_WAIVE_BIND_PRIORITY;
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.executeShellCommand;
 import static com.android.cts.net.hostside.NetworkPolicyTestUtils.forceRunJob;
@@ -51,6 +52,7 @@
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.net.ConnectivityManager;
+import android.net.NetworkInfo;
 import android.net.NetworkInfo.DetailedState;
 import android.net.NetworkInfo.State;
 import android.net.NetworkRequest;
@@ -78,10 +80,10 @@
 import org.junit.rules.RuleChain;
 import org.junit.runner.RunWith;
 
-import java.util.ArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
 
 /**
@@ -136,6 +138,7 @@
 
     private static final String KEY_NETWORK_STATE_OBSERVER = TEST_PKG + ".observer";
     private static final String KEY_SKIP_VALIDATION_CHECKS = TEST_PKG + ".skip_validation_checks";
+    private static final String KEY_CUSTOM_URL = TEST_PKG + ".custom_url";
 
     private static final String EMPTY_STRING = "";
 
@@ -165,6 +168,7 @@
     protected ConnectivityManager mCm;
     protected int mUid;
     private int mMyUid;
+    private @Nullable String mCustomUrl;
     private MyServiceClient mServiceClient;
     private DeviceConfigStateHelper mDeviceIdleDeviceConfigStateHelper;
     private PowerManager mPowerManager;
@@ -185,6 +189,11 @@
         mServiceClient = new MyServiceClient(mContext);
 
         final Bundle args = InstrumentationRegistry.getArguments();
+        mCustomUrl = args.getString(ARG_CONNECTION_CHECK_CUSTOM_URL);
+        if (mCustomUrl != null) {
+            Log.d(TAG, "Using custom URL " + mCustomUrl + " for network checks");
+        }
+
         final int bindPriorityFlags;
         if (Boolean.valueOf(args.getString(ARG_WAIVE_BIND_PRIORITY, "false"))) {
             bindPriorityFlags = Context.BIND_WAIVE_PRIORITY;
@@ -502,25 +511,23 @@
      */
     private String checkNetworkAccess(boolean expectAvailable,
             @Nullable final String expectedUnavailableError) throws Exception {
-        final String resultData = mServiceClient.checkNetworkStatus();
-        return checkForAvailabilityInResultData(resultData, expectAvailable,
+        final NetworkCheckResult checkResult = mServiceClient.checkNetworkStatus(mCustomUrl);
+        return checkForAvailabilityInNetworkCheckResult(checkResult, expectAvailable,
                 expectedUnavailableError);
     }
 
-    private String checkForAvailabilityInResultData(String resultData, boolean expectAvailable,
-            @Nullable final String expectedUnavailableError) {
-        if (resultData == null) {
-            assertNotNull("Network status from app2 is null", resultData);
-        }
-        // Network status format is described on MyBroadcastReceiver.checkNetworkStatus()
-        final String[] parts = resultData.split(NETWORK_STATUS_SEPARATOR);
-        assertEquals("Wrong network status: " + resultData, 5, parts.length);
-        final State state = parts[0].equals("null") ? null : State.valueOf(parts[0]);
-        final DetailedState detailedState = parts[1].equals("null")
-                ? null : DetailedState.valueOf(parts[1]);
-        final boolean connected = Boolean.valueOf(parts[2]);
-        final String connectionCheckDetails = parts[3];
-        final String networkInfo = parts[4];
+    private String checkForAvailabilityInNetworkCheckResult(NetworkCheckResult networkCheckResult,
+            boolean expectAvailable, @Nullable final String expectedUnavailableError) {
+        assertNotNull("NetworkCheckResult from app2 is null", networkCheckResult);
+
+        final NetworkInfo networkInfo = networkCheckResult.networkInfo;
+        assertNotNull("NetworkInfo from app2 is null", networkInfo);
+
+        final State state = networkInfo.getState();
+        final DetailedState detailedState = networkInfo.getDetailedState();
+
+        final boolean connected = networkCheckResult.connected;
+        final String connectionCheckDetails = networkCheckResult.details;
 
         final StringBuilder errors = new StringBuilder();
         final State expectedState;
@@ -926,33 +933,36 @@
         if (type == TYPE_COMPONENT_FOREGROUND_SERVICE) {
             startForegroundService();
             assertForegroundServiceNetworkAccess();
-            return;
         } else if (type == TYPE_COMPONENT_ACTIVTIY) {
             turnScreenOn();
             final CountDownLatch latch = new CountDownLatch(1);
             final Intent launchIntent = getIntentForComponent(type);
             final Bundle extras = new Bundle();
-            final ArrayList<Pair<Integer, String>> result = new ArrayList<>(1);
+            final AtomicReference<Pair<Integer, NetworkCheckResult>> result =
+                    new AtomicReference<>();
             extras.putBinder(KEY_NETWORK_STATE_OBSERVER, getNewNetworkStateObserver(latch, result));
             extras.putBoolean(KEY_SKIP_VALIDATION_CHECKS, !expectAvailable);
+            extras.putString(KEY_CUSTOM_URL, mCustomUrl);
             launchIntent.putExtras(extras);
             mContext.startActivity(launchIntent);
             if (latch.await(ACTIVITY_NETWORK_STATE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
-                final int resultCode = result.get(0).first;
-                final String resultData = result.get(0).second;
+                final int resultCode = result.get().first;
+                final NetworkCheckResult networkCheckResult = result.get().second;
                 if (resultCode == INetworkStateObserver.RESULT_SUCCESS_NETWORK_STATE_CHECKED) {
-                    final String error = checkForAvailabilityInResultData(
-                            resultData, expectAvailable, null /* expectedUnavailableError */);
+                    final String error = checkForAvailabilityInNetworkCheckResult(
+                            networkCheckResult, expectAvailable,
+                            null /* expectedUnavailableError */);
                     if (error != null) {
                         fail("Network is not available for activity in app2 (" + mUid + "): "
                                 + error);
                     }
                 } else if (resultCode == INetworkStateObserver.RESULT_ERROR_UNEXPECTED_PROC_STATE) {
-                    Log.d(TAG, resultData);
+                    Log.d(TAG, networkCheckResult.details);
                     // App didn't come to foreground when the activity is started, so try again.
                     assertTopNetworkAccess(true);
                 } else {
-                    fail("Unexpected resultCode=" + resultCode + "; received=[" + resultData + "]");
+                    fail("Unexpected resultCode=" + resultCode
+                            + "; networkCheckResult=[" + networkCheckResult + "]");
                 }
             } else {
                 fail("Timed out waiting for network availability status from app2's activity ("
@@ -960,10 +970,12 @@
             }
         } else if (type == TYPE_EXPEDITED_JOB) {
             final Bundle extras = new Bundle();
-            final ArrayList<Pair<Integer, String>> result = new ArrayList<>(1);
+            final AtomicReference<Pair<Integer, NetworkCheckResult>> result =
+                    new AtomicReference<>();
             final CountDownLatch latch = new CountDownLatch(1);
             extras.putBinder(KEY_NETWORK_STATE_OBSERVER, getNewNetworkStateObserver(latch, result));
             extras.putBoolean(KEY_SKIP_VALIDATION_CHECKS, !expectAvailable);
+            extras.putString(KEY_CUSTOM_URL, mCustomUrl);
             final JobInfo jobInfo = new JobInfo.Builder(TEST_JOB_ID, TEST_JOB_COMPONENT)
                     .setExpedited(true)
                     .setRequiredNetworkType(JobInfo.NETWORK_TYPE_ANY)
@@ -973,11 +985,12 @@
                     RESULT_SUCCESS, mServiceClient.scheduleJob(jobInfo));
             forceRunJob(TEST_APP2_PKG, TEST_JOB_ID);
             if (latch.await(JOB_NETWORK_STATE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
-                final int resultCode = result.get(0).first;
-                final String resultData = result.get(0).second;
+                final int resultCode = result.get().first;
+                final NetworkCheckResult networkCheckResult = result.get().second;
                 if (resultCode == INetworkStateObserver.RESULT_SUCCESS_NETWORK_STATE_CHECKED) {
-                    final String error = checkForAvailabilityInResultData(
-                            resultData, expectAvailable, null /* expectedUnavailableError */);
+                    final String error = checkForAvailabilityInNetworkCheckResult(
+                            networkCheckResult, expectAvailable,
+                            null /* expectedUnavailableError */);
                     if (error != null) {
                         Log.d(TAG, "Network state is unexpected, checking again. " + error);
                         // Right now we could end up in an unexpected state if expedited job
@@ -985,7 +998,8 @@
                         assertNetworkAccess(expectAvailable, false /* needScreenOn */);
                     }
                 } else {
-                    fail("Unexpected resultCode=" + resultCode + "; received=[" + resultData + "]");
+                    fail("Unexpected resultCode=" + resultCode
+                            + "; networkCheckResult=[" + networkCheckResult + "]");
                 }
             } else {
                 fail("Timed out waiting for network availability status from app2's expedited job ("
@@ -1028,11 +1042,12 @@
     }
 
     private Binder getNewNetworkStateObserver(final CountDownLatch latch,
-            final ArrayList<Pair<Integer, String>> result) {
+            final AtomicReference<Pair<Integer, NetworkCheckResult>> result) {
         return new INetworkStateObserver.Stub() {
             @Override
-            public void onNetworkStateChecked(int resultCode, String resultData) {
-                result.add(Pair.create(resultCode, resultData));
+            public void onNetworkStateChecked(int resultCode,
+                    NetworkCheckResult networkCheckResult) {
+                result.set(Pair.create(resultCode, networkCheckResult));
                 latch.countDown();
             }
         };
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/ConnOnActivityStartTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/ConnOnActivityStartTest.java
index c1d576d..3e22a23 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/ConnOnActivityStartTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/ConnOnActivityStartTest.java
@@ -69,6 +69,7 @@
     @RequiredProperties({BATTERY_SAVER_MODE})
     public void testStartActivity_batterySaver() throws Exception {
         setBatterySaverMode(true);
+        assertNetworkAccess(false, null);
         assertLaunchedActivityHasNetworkAccess("testStartActivity_batterySaver", null);
     }
 
@@ -76,6 +77,7 @@
     @RequiredProperties({DATA_SAVER_MODE, METERED_NETWORK})
     public void testStartActivity_dataSaver() throws Exception {
         setRestrictBackground(true);
+        assertNetworkAccess(false, null);
         assertLaunchedActivityHasNetworkAccess("testStartActivity_dataSaver", null);
     }
 
@@ -83,6 +85,7 @@
     @RequiredProperties({DOZE_MODE})
     public void testStartActivity_doze() throws Exception {
         setDozeMode(true);
+        assertNetworkAccess(false, null);
         // TODO (235284115): We need to turn on Doze every time before starting
         // the activity.
         assertLaunchedActivityHasNetworkAccess("testStartActivity_doze", null);
@@ -93,6 +96,7 @@
     public void testStartActivity_appStandby() throws Exception {
         turnBatteryOn();
         setAppIdle(true);
+        assertNetworkAccess(false, null);
         // TODO (235284115): We need to put the app into app standby mode every
         // time before starting the activity.
         assertLaunchedActivityHasNetworkAccess("testStartActivity_appStandby", null);
@@ -104,6 +108,7 @@
         assertLaunchedActivityHasNetworkAccess("testStartActivity_default", () -> {
             assertProcessStateBelow(PROCESS_STATE_TOP_SLEEPING);
             SystemClock.sleep(PROCESS_STATE_TRANSITION_DELAY_MS);
+            assertNetworkAccess(false, null);
         });
     }
 
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java
index 980ecd5..494192f 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/MyServiceClient.java
@@ -98,9 +98,10 @@
         return mService.getCounters(receiverName, action);
     }
 
-    public String checkNetworkStatus() throws RemoteException {
+    /** Retrieves the network state as observed from the bound test app */
+    public NetworkCheckResult checkNetworkStatus(String address) throws RemoteException {
         ensureServiceConnection();
-        return mService.checkNetworkStatus();
+        return mService.checkNetworkStatus(address);
     }
 
     public String getRestrictBackgroundStatus() throws RemoteException {
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
index f81a03d..0f86d78 100755
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
@@ -23,8 +23,8 @@
 import static android.content.pm.PackageManager.FEATURE_TELEPHONY;
 import static android.content.pm.PackageManager.FEATURE_WIFI;
 import static android.net.ConnectivityManager.TYPE_VPN;
-import static android.net.NetworkCapabilities.TRANSPORT_VPN;
 import static android.net.NetworkCapabilities.TRANSPORT_TEST;
+import static android.net.NetworkCapabilities.TRANSPORT_VPN;
 import static android.os.Process.INVALID_UID;
 import static android.system.OsConstants.AF_INET;
 import static android.system.OsConstants.AF_INET6;
@@ -123,6 +123,7 @@
 import com.android.net.module.util.ArrayTrackRecord;
 import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.PacketBuilder;
+import com.android.testutils.AutoReleaseNetworkCallbackRule;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import com.android.testutils.RecorderCallback;
@@ -233,9 +234,13 @@
     // The registered callbacks.
     private List<NetworkCallback> mRegisteredCallbacks = new ArrayList<>();
 
-    @Rule
+    @Rule(order = 1)
     public final DevSdkIgnoreRule mDevSdkIgnoreRule = new DevSdkIgnoreRule();
 
+    @Rule(order = 2)
+    public final AutoReleaseNetworkCallbackRule
+            mNetworkCallbackRule = new AutoReleaseNetworkCallbackRule();
+
     private boolean supportedHardware() {
         final PackageManager pm = getInstrumentation().getContext().getPackageManager();
         return !pm.hasSystemFeature("android.hardware.type.watch");
@@ -273,7 +278,6 @@
     public void tearDown() throws Exception {
         restorePrivateDnsSetting();
         mRemoteSocketFactoryClient.unbind();
-        mCtsNetUtils.tearDown();
         Log.i(TAG, "Stopping VPN");
         stopVpn();
         unregisterRegisteredCallbacks();
@@ -890,9 +894,7 @@
         testAndCleanup(() -> {
             // Ensure both of wifi and mobile data are connected.
             final Network wifiNetwork = mCtsNetUtils.ensureWifiConnected();
-            assertTrue("Wifi is not connected", (wifiNetwork != null));
-            final Network cellNetwork = mCtsNetUtils.connectToCell();
-            assertTrue("Mobile data is not connected", (cellNetwork != null));
+            final Network cellNetwork = mNetworkCallbackRule.requestCell();
             // Store current default network.
             final Network defaultNetwork = mCM.getActiveNetwork();
             // Start VPN and set empty array as its underlying networks.
diff --git a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java
index 37dc7a0..1c45579 100644
--- a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java
+++ b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java
@@ -15,10 +15,16 @@
  */
 package com.android.cts.net.hostside.app2;
 
+import static com.android.cts.net.hostside.INetworkStateObserver.RESULT_ERROR_OTHER;
+import static com.android.cts.net.hostside.INetworkStateObserver.RESULT_ERROR_UNEXPECTED_CAPABILITIES;
+import static com.android.cts.net.hostside.INetworkStateObserver.RESULT_ERROR_UNEXPECTED_PROC_STATE;
+
 import android.app.ActivityManager;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager.NameNotFoundException;
+import android.net.ConnectivityManager;
+import android.net.NetworkInfo;
 import android.os.AsyncTask;
 import android.os.Bundle;
 import android.os.Process;
@@ -26,6 +32,12 @@
 import android.util.Log;
 
 import com.android.cts.net.hostside.INetworkStateObserver;
+import com.android.cts.net.hostside.NetworkCheckResult;
+
+import java.net.HttpURLConnection;
+import java.net.InetAddress;
+import java.net.URL;
+import java.util.concurrent.TimeUnit;
 
 public final class Common {
 
@@ -48,6 +60,9 @@
     static final String ACTION_SNOOZE_WARNING =
             "com.android.server.net.action.SNOOZE_WARNING";
 
+    private static final String DEFAULT_TEST_URL =
+            "https://connectivitycheck.android.com/generate_204";
+
     static final String NOTIFICATION_TYPE_CONTENT = "CONTENT";
     static final String NOTIFICATION_TYPE_DELETE = "DELETE";
     static final String NOTIFICATION_TYPE_FULL_SCREEN = "FULL_SCREEN";
@@ -59,10 +74,12 @@
     static final String TEST_PKG = "com.android.cts.net.hostside";
     static final String KEY_NETWORK_STATE_OBSERVER = TEST_PKG + ".observer";
     static final String KEY_SKIP_VALIDATION_CHECKS = TEST_PKG + ".skip_validation_checks";
+    static final String KEY_CUSTOM_URL =  TEST_PKG + ".custom_url";
 
     static final int TYPE_COMPONENT_ACTIVTY = 0;
     static final int TYPE_COMPONENT_FOREGROUND_SERVICE = 1;
     static final int TYPE_COMPONENT_EXPEDITED_JOB = 2;
+    private static final int NETWORK_TIMEOUT_MS = (int) TimeUnit.SECONDS.toMillis(10);
 
     static int getUid(Context context) {
         final String packageName = context.getPackageName();
@@ -73,6 +90,15 @@
         }
     }
 
+    private static NetworkCheckResult createNetworkCheckResult(boolean connected, String details,
+            NetworkInfo networkInfo) {
+        final NetworkCheckResult checkResult = new NetworkCheckResult();
+        checkResult.connected = connected;
+        checkResult.details = details;
+        checkResult.networkInfo = networkInfo;
+        return checkResult;
+    }
+
     private static boolean validateComponentState(Context context, int componentType,
             INetworkStateObserver observer) throws RemoteException {
         final ActivityManager activityManager = context.getSystemService(ActivityManager.class);
@@ -80,9 +106,9 @@
             case TYPE_COMPONENT_ACTIVTY: {
                 final int procState = activityManager.getUidProcessState(Process.myUid());
                 if (procState != ActivityManager.PROCESS_STATE_TOP) {
-                    observer.onNetworkStateChecked(
-                            INetworkStateObserver.RESULT_ERROR_UNEXPECTED_PROC_STATE,
-                            "Unexpected procstate: " + procState);
+                    observer.onNetworkStateChecked(RESULT_ERROR_UNEXPECTED_PROC_STATE,
+                            createNetworkCheckResult(false, "Unexpected procstate: " + procState,
+                                    null));
                     return false;
                 }
                 return true;
@@ -90,9 +116,9 @@
             case TYPE_COMPONENT_FOREGROUND_SERVICE: {
                 final int procState = activityManager.getUidProcessState(Process.myUid());
                 if (procState != ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE) {
-                    observer.onNetworkStateChecked(
-                            INetworkStateObserver.RESULT_ERROR_UNEXPECTED_PROC_STATE,
-                            "Unexpected procstate: " + procState);
+                    observer.onNetworkStateChecked(RESULT_ERROR_UNEXPECTED_PROC_STATE,
+                            createNetworkCheckResult(false, "Unexpected procstate: " + procState,
+                                    null));
                     return false;
                 }
                 return true;
@@ -101,16 +127,17 @@
                 final int capabilities = activityManager.getUidProcessCapabilities(Process.myUid());
                 if ((capabilities
                         & ActivityManager.PROCESS_CAPABILITY_POWER_RESTRICTED_NETWORK) == 0) {
-                    observer.onNetworkStateChecked(
-                            INetworkStateObserver.RESULT_ERROR_UNEXPECTED_CAPABILITIES,
-                            "Unexpected capabilities: " + capabilities);
+                    observer.onNetworkStateChecked(RESULT_ERROR_UNEXPECTED_CAPABILITIES,
+                            createNetworkCheckResult(false,
+                                    "Unexpected capabilities: " + capabilities, null));
                     return false;
                 }
                 return true;
             }
             default: {
-                observer.onNetworkStateChecked(INetworkStateObserver.RESULT_ERROR_OTHER,
-                        "Unknown component type: " + componentType);
+                observer.onNetworkStateChecked(RESULT_ERROR_OTHER,
+                        createNetworkCheckResult(false, "Unknown component type: " + componentType,
+                                null));
                 return false;
             }
         }
@@ -131,6 +158,7 @@
         final INetworkStateObserver observer = INetworkStateObserver.Stub.asInterface(
                 extras.getBinder(KEY_NETWORK_STATE_OBSERVER));
         if (observer != null) {
+            final String customUrl = extras.getString(KEY_CUSTOM_URL);
             try {
                 final boolean skipValidation = extras.getBoolean(KEY_SKIP_VALIDATION_CHECKS);
                 if (!skipValidation && !validateComponentState(context, componentType, observer)) {
@@ -143,11 +171,64 @@
                 try {
                     observer.onNetworkStateChecked(
                             INetworkStateObserver.RESULT_SUCCESS_NETWORK_STATE_CHECKED,
-                            MyBroadcastReceiver.checkNetworkStatus(context));
+                            checkNetworkStatus(context, customUrl));
                 } catch (RemoteException e) {
                     Log.e(TAG, "Error occurred while notifying the observer: " + e);
                 }
             });
         }
     }
+
+    /**
+     * Checks whether the network is available by attempting a connection to the given address
+     * and returns a {@link NetworkCheckResult} object containing all the relevant details for
+     * debugging. Uses a default address if the given address is {@code null}.
+     *
+     * <p>
+     * The returned object has the following fields:
+     *
+     * <ul>
+     * <li>{@code connected}: whether or not the connection was successful.
+     * <li>{@code networkInfo}: the {@link NetworkInfo} describing the current active network as
+     * visible to this app.
+     * <li>{@code details}: A human readable string giving useful information about the success or
+     * failure.
+     * </ul>
+     */
+    static NetworkCheckResult checkNetworkStatus(Context context, String customUrl) {
+        final String address = (customUrl == null) ? DEFAULT_TEST_URL : customUrl;
+
+        // The current Android DNS resolver returns an UnknownHostException whenever network access
+        // is blocked. This can get cached in the current process-local InetAddress cache. Clearing
+        // the cache before attempting a connection ensures we never report a failure due to a
+        // negative cache entry.
+        InetAddress.clearDnsCache();
+
+        final ConnectivityManager cm = context.getSystemService(ConnectivityManager.class);
+
+        final NetworkInfo networkInfo = cm.getActiveNetworkInfo();
+        Log.d(TAG, "Running checkNetworkStatus() on thread "
+                + Thread.currentThread().getName() + " for UID " + getUid(context)
+                + "\n\tactiveNetworkInfo: " + networkInfo + "\n\tURL: " + address);
+        boolean checkStatus = false;
+        String checkDetails = "N/A";
+        try {
+            final URL url = new URL(address);
+            final HttpURLConnection conn = (HttpURLConnection) url.openConnection();
+            conn.setReadTimeout(NETWORK_TIMEOUT_MS);
+            conn.setConnectTimeout(NETWORK_TIMEOUT_MS / 2);
+            conn.setRequestMethod("GET");
+            conn.connect();
+            final int response = conn.getResponseCode();
+            checkStatus = true;
+            checkDetails = "HTTP response for " + address + ": " + response;
+        } catch (Exception e) {
+            checkStatus = false;
+            checkDetails = "Exception getting " + address + ": " + e;
+        }
+        final NetworkCheckResult result = createNetworkCheckResult(checkStatus, checkDetails,
+                networkInfo);
+        Log.d(TAG, "Offering: " + result);
+        return result;
+    }
 }
diff --git a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyBroadcastReceiver.java b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyBroadcastReceiver.java
index 825f2c9..1fd3745 100644
--- a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyBroadcastReceiver.java
+++ b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyBroadcastReceiver.java
@@ -30,7 +30,6 @@
 import static com.android.cts.net.hostside.app2.Common.NOTIFICATION_TYPE_DELETE;
 import static com.android.cts.net.hostside.app2.Common.NOTIFICATION_TYPE_FULL_SCREEN;
 import static com.android.cts.net.hostside.app2.Common.TAG;
-import static com.android.cts.net.hostside.app2.Common.getUid;
 
 import android.app.Notification;
 import android.app.Notification.Action;
@@ -42,15 +41,10 @@
 import android.content.Intent;
 import android.content.SharedPreferences;
 import android.net.ConnectivityManager;
-import android.net.NetworkInfo;
 import android.os.Bundle;
 import android.util.Log;
 import android.widget.Toast;
 
-import java.net.HttpURLConnection;
-import java.net.InetAddress;
-import java.net.URL;
-
 /**
  * Receiver used to:
  * <ol>
@@ -60,8 +54,6 @@
  */
 public class MyBroadcastReceiver extends BroadcastReceiver {
 
-    private static final int NETWORK_TIMEOUT_MS = 5 * 1000;
-
     private final String mName;
 
     public MyBroadcastReceiver() {
@@ -126,82 +118,6 @@
         return String.valueOf(apiStatus);
     }
 
-    private static final String NETWORK_STATUS_TEMPLATE = "%s|%s|%s|%s|%s";
-    /**
-     * Checks whether the network is available and return a string which can then be send as a
-     * result data for the ordered broadcast.
-     *
-     * <p>
-     * The string has the following format:
-     *
-     * <p><pre><code>
-     * NetinfoState|NetinfoDetailedState|RealConnectionCheck|RealConnectionCheckDetails|Netinfo
-     * </code></pre>
-     *
-     * <p>Where:
-     *
-     * <ul>
-     * <li>{@code NetinfoState}: enum value of {@link NetworkInfo.State}.
-     * <li>{@code NetinfoDetailedState}: enum value of {@link NetworkInfo.DetailedState}.
-     * <li>{@code RealConnectionCheck}: boolean value of a real connection check (i.e., an attempt
-     *     to access an external website.
-     * <li>{@code RealConnectionCheckDetails}: if HTTP output core or exception string of the real
-     *     connection attempt
-     * <li>{@code Netinfo}: string representation of the {@link NetworkInfo}.
-     * </ul>
-     *
-     * For example, if the connection was established fine, the result would be something like:
-     * <p><pre><code>
-     * CONNECTED|CONNECTED|true|200|[type: WIFI[], state: CONNECTED/CONNECTED, reason: ...]
-     * </code></pre>
-     *
-     */
-    // TODO: now that it uses Binder, it counl return a Bundle with the data parts instead...
-    static String checkNetworkStatus(Context context) {
-        final ConnectivityManager cm =
-                (ConnectivityManager) context.getSystemService(Context.CONNECTIVITY_SERVICE);
-        // TODO: connect to a hostside server instead
-        final String address = "http://example.com";
-        final NetworkInfo networkInfo = cm.getActiveNetworkInfo();
-        Log.d(TAG, "Running checkNetworkStatus() on thread "
-                + Thread.currentThread().getName() + " for UID " + getUid(context)
-                + "\n\tactiveNetworkInfo: " + networkInfo + "\n\tURL: " + address);
-        boolean checkStatus = false;
-        String checkDetails = "N/A";
-        try {
-            final URL url = new URL(address);
-            final HttpURLConnection conn = (HttpURLConnection) url.openConnection();
-            conn.setReadTimeout(NETWORK_TIMEOUT_MS);
-            conn.setConnectTimeout(NETWORK_TIMEOUT_MS / 2);
-            conn.setRequestMethod("GET");
-            conn.setDoInput(true);
-            conn.connect();
-            final int response = conn.getResponseCode();
-            checkStatus = true;
-            checkDetails = "HTTP response for " + address + ": " + response;
-        } catch (Exception e) {
-            checkStatus = false;
-            checkDetails = "Exception getting " + address + ": " + e;
-        }
-        // If the app tries to make a network connection in the foreground immediately after
-        // trying to do the same when it's network access was blocked, it could receive a
-        // UnknownHostException due to the cached DNS entry. So, clear the dns cache after
-        // every network access for now until we have a fix on the platform side.
-        InetAddress.clearDnsCache();
-        Log.d(TAG, checkDetails);
-        final String state, detailedState;
-        if (networkInfo != null) {
-            state = networkInfo.getState().name();
-            detailedState = networkInfo.getDetailedState().name();
-        } else {
-            state = detailedState = "null";
-        }
-        final String status = String.format(NETWORK_STATUS_TEMPLATE, state, detailedState,
-                Boolean.valueOf(checkStatus), checkDetails, networkInfo);
-        Log.d(TAG, "Offering " + status);
-        return status;
-    }
-
     /**
      * Sends a system notification containing actions with pending intents to launch the app's
      * main activitiy or service.
diff --git a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java
index 3ed5391..5010234 100644
--- a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java
+++ b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyService.java
@@ -21,7 +21,6 @@
 import static com.android.cts.net.hostside.app2.Common.ACTION_SNOOZE_WARNING;
 import static com.android.cts.net.hostside.app2.Common.DYNAMIC_RECEIVER;
 import static com.android.cts.net.hostside.app2.Common.TAG;
-import static com.android.networkstack.apishim.ConstantsShim.RECEIVER_EXPORTED;
 
 import android.app.NotificationChannel;
 import android.app.NotificationManager;
@@ -41,6 +40,7 @@
 
 import com.android.cts.net.hostside.IMyService;
 import com.android.cts.net.hostside.INetworkCallback;
+import com.android.cts.net.hostside.NetworkCheckResult;
 import com.android.modules.utils.build.SdkLevel;
 
 /**
@@ -56,9 +56,7 @@
 
     // TODO: move MyBroadcast static functions here - they were kept there to make git diff easier.
 
-    private IMyService.Stub mBinder =
-        new IMyService.Stub() {
-
+    private IMyService.Stub mBinder = new IMyService.Stub() {
         @Override
         public void registerBroadcastReceiver() {
             if (mReceiver != null) {
@@ -83,8 +81,8 @@
         }
 
         @Override
-        public String checkNetworkStatus() {
-            return MyBroadcastReceiver.checkNetworkStatus(getApplicationContext());
+        public NetworkCheckResult checkNetworkStatus(String customUrl) {
+            return Common.checkNetworkStatus(getApplicationContext(), customUrl);
         }
 
         @Override
@@ -94,7 +92,7 @@
 
         @Override
         public void sendNotification(int notificationId, String notificationType) {
-            MyBroadcastReceiver .sendNotification(getApplicationContext(), NOTIFICATION_CHANNEL_ID,
+            MyBroadcastReceiver.sendNotification(getApplicationContext(), NOTIFICATION_CHANNEL_ID,
                     notificationId, notificationType);
         }
 
@@ -170,7 +168,7 @@
                     .getSystemService(JobScheduler.class);
             return jobScheduler.schedule(jobInfo);
         }
-      };
+    };
 
     @Override
     public IBinder onBind(Intent intent) {
diff --git a/tests/cts/hostside/instrumentation_arguments/src/com/android/cts/net/arguments/InstrumentationArguments.java b/tests/cts/hostside/instrumentation_arguments/src/com/android/cts/net/arguments/InstrumentationArguments.java
index 472e347..911b129 100644
--- a/tests/cts/hostside/instrumentation_arguments/src/com/android/cts/net/arguments/InstrumentationArguments.java
+++ b/tests/cts/hostside/instrumentation_arguments/src/com/android/cts/net/arguments/InstrumentationArguments.java
@@ -18,4 +18,5 @@
 
 public interface InstrumentationArguments {
     String ARG_WAIVE_BIND_PRIORITY = "waive_bind_priority";
+    String ARG_CONNECTION_CHECK_CUSTOM_URL = "connection_check_custom_url";
 }
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java b/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
index 880e826..fff716d 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideConnOnActivityStartTest.java
@@ -47,29 +47,29 @@
 
     @Test
     public void testStartActivity_batterySaver() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_batterySaver");
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_CLASS, "testStartActivity_batterySaver");
     }
 
     @Test
     public void testStartActivity_dataSaver() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_dataSaver");
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_CLASS, "testStartActivity_dataSaver");
     }
 
     @FlakyTest(bugId = 231440256)
     @Test
     public void testStartActivity_doze() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_doze");
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_CLASS, "testStartActivity_doze");
     }
 
     @Test
     public void testStartActivity_appStandby() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_CLASS, "testStartActivity_appStandby");
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_CLASS, "testStartActivity_appStandby");
     }
 
     // TODO(b/321848487): Annotate with @RequiresFlagsEnabled to mirror the device-side test.
     @Test
     public void testStartActivity_default() throws Exception {
-        runDeviceTestsWithArgs(TEST_PKG, TEST_CLASS, "testStartActivity_default",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_CLASS, "testStartActivity_default",
                 Map.of(ARG_WAIVE_BIND_PRIORITY, "true"));
     }
 }
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideDefaultNetworkRestrictionsTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideDefaultNetworkRestrictionsTests.java
index 0d01fc1..faabbef 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideDefaultNetworkRestrictionsTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideDefaultNetworkRestrictionsTests.java
@@ -46,12 +46,12 @@
     }
 
     private void runMeteredTest(String methodName) throws DeviceNotAvailableException {
-        runDeviceTestsWithArgs(TEST_PKG, METERED_TEST_CLASS, methodName,
+        runDeviceTestsWithCustomOptions(TEST_PKG, METERED_TEST_CLASS, methodName,
                 Map.of(ARG_WAIVE_BIND_PRIORITY, "true"));
     }
 
     private void runNonMeteredTest(String methodName) throws DeviceNotAvailableException {
-        runDeviceTestsWithArgs(TEST_PKG, NON_METERED_TEST_CLASS, methodName,
+        runDeviceTestsWithCustomOptions(TEST_PKG, NON_METERED_TEST_CLASS, methodName,
                 Map.of(ARG_WAIVE_BIND_PRIORITY, "true"));
     }
 
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
index 361f7c7..c4bcdfd 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkCallbackTests.java
@@ -41,20 +41,20 @@
 
     @Test
     public void testOnBlockedStatusChanged_dataSaver() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkCallbackTest", "testOnBlockedStatusChanged_dataSaver");
     }
 
     @Test
     public void testOnBlockedStatusChanged_powerSaver() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkCallbackTest", "testOnBlockedStatusChanged_powerSaver");
     }
 
     // TODO(b/321848487): Annotate with @RequiresFlagsEnabled to mirror the device-side test.
     @Test
     public void testOnBlockedStatusChanged_default() throws Exception {
-        runDeviceTestsWithArgs(TEST_PKG, TEST_PKG + ".NetworkCallbackTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".NetworkCallbackTest",
                 "testOnBlockedStatusChanged_default", Map.of(ARG_WAIVE_BIND_PRIORITY, "true"));
     }
 }
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
index e97db58..4730b14 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkPolicyManagerTests.java
@@ -38,48 +38,48 @@
 
     @Test
     public void testIsUidNetworkingBlocked_withUidNotBlocked() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withUidNotBlocked");
     }
 
     @Test
     public void testIsUidNetworkingBlocked_withSystemUid() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest", "testIsUidNetworkingBlocked_withSystemUid");
     }
 
     @Test
     public void testIsUidNetworkingBlocked_withDataSaverMode() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withDataSaverMode");
     }
 
     @Test
     public void testIsUidNetworkingBlocked_withRestrictedNetworkingMode() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withRestrictedNetworkingMode");
     }
 
     @Test
     public void testIsUidNetworkingBlocked_withPowerSaverMode() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_withPowerSaverMode");
     }
 
     @Test
     public void testIsUidRestrictedOnMeteredNetworks() throws Exception {
-        runDeviceTests(TEST_PKG,
+        runDeviceTestsWithCustomOptions(TEST_PKG,
                 TEST_PKG + ".NetworkPolicyManagerTest", "testIsUidRestrictedOnMeteredNetworks");
     }
 
     // TODO(b/321848487): Annotate with @RequiresFlagsEnabled to mirror the device-side test.
     @Test
     public void testIsUidNetworkingBlocked_whenInBackground() throws Exception {
-        runDeviceTestsWithArgs(TEST_PKG, TEST_PKG + ".NetworkPolicyManagerTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".NetworkPolicyManagerTest",
                 "testIsUidNetworkingBlocked_whenInBackground",
                 Map.of(ARG_WAIVE_BIND_PRIORITY, "true"));
     }
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
index ca95ed6..d7dfa80 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideNetworkTestCase.java
@@ -16,12 +16,15 @@
 
 package com.android.cts.net;
 
+import static com.android.cts.net.arguments.InstrumentationArguments.ARG_CONNECTION_CHECK_CUSTOM_URL;
+
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
 import com.android.ddmlib.Log;
 import com.android.modules.utils.build.testing.DeviceSdkLevel;
+import com.android.tradefed.config.Option;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.invoker.TestInformation;
 import com.android.tradefed.targetprep.BuildError;
@@ -48,6 +51,10 @@
     protected static final String TEST_APP2_PKG = "com.android.cts.net.hostside.app2";
     protected static final String TEST_APP2_APK = "CtsHostsideNetworkTestsApp2.apk";
 
+    @Option(name = "custom-url", importance = Option.Importance.IF_UNSET,
+            description = "A custom url to use for testing network connections")
+    protected String mCustomUrl;
+
     @BeforeClassWithInfo
     public static void setUpOnceBase(TestInformation testInfo) throws Exception {
         DeviceSdkLevel deviceSdkLevel = new DeviceSdkLevel(testInfo.getDevice());
@@ -149,13 +156,31 @@
                 + packageName + ", u=" + currentUser);
     }
 
-    protected boolean runDeviceTestsWithArgs(String packageName, String className,
-            String methodName, Map<String, String> args) throws DeviceNotAvailableException {
+    protected boolean runDeviceTestsWithCustomOptions(String packageName, String className)
+            throws DeviceNotAvailableException {
+        return runDeviceTestsWithCustomOptions(packageName, className, null);
+    }
+
+    protected boolean runDeviceTestsWithCustomOptions(String packageName, String className,
+            String methodName) throws DeviceNotAvailableException {
+        return runDeviceTestsWithCustomOptions(packageName, className, methodName, null);
+    }
+
+    protected boolean runDeviceTestsWithCustomOptions(String packageName, String className,
+            String methodName, Map<String, String> testArgs) throws DeviceNotAvailableException {
         final DeviceTestRunOptions deviceTestRunOptions = new DeviceTestRunOptions(packageName)
                 .setTestClassName(className)
                 .setTestMethodName(methodName);
-        for (Map.Entry<String, String> arg : args.entrySet()) {
-            deviceTestRunOptions.addInstrumentationArg(arg.getKey(), arg.getValue());
+
+        // Currently there is only one custom option that the test exposes.
+        if (mCustomUrl != null) {
+            deviceTestRunOptions.addInstrumentationArg(ARG_CONNECTION_CHECK_CUSTOM_URL, mCustomUrl);
+        }
+        // Pass over any test specific arguments.
+        if (testArgs != null) {
+            for (Map.Entry<String, String> arg : testArgs.entrySet()) {
+                deviceTestRunOptions.addInstrumentationArg(arg.getKey(), arg.getValue());
+            }
         }
         return runDeviceTests(deviceTestRunOptions);
     }
diff --git a/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java b/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
index 9c3751d..7b9d3b5 100644
--- a/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
+++ b/tests/cts/hostside/src/com/android/cts/net/HostsideRestrictBackgroundNetworkTests.java
@@ -46,7 +46,7 @@
     @SecurityTest
     @Test
     public void testDataWarningReceiver() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DataWarningReceiverTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DataWarningReceiverTest",
                 "testSnoozeWarningNotReceived");
     }
 
@@ -56,25 +56,25 @@
 
     @Test
     public void testDataSaverMode_disabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_disabled");
     }
 
     @Test
     public void testDataSaverMode_whitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_whitelisted");
     }
 
     @Test
     public void testDataSaverMode_enabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_enabled");
     }
 
     @Test
     public void testDataSaverMode_blacklisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_blacklisted");
     }
 
@@ -97,13 +97,13 @@
 
     @Test
     public void testDataSaverMode_requiredWhitelistedPackages() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testGetRestrictBackgroundStatus_requiredWhitelistedPackages");
     }
 
     @Test
     public void testDataSaverMode_broadcastNotSentOnUnsupportedDevices() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DataSaverModeTest",
                 "testBroadcastNotSentOnUnsupportedDevices");
     }
 
@@ -113,19 +113,19 @@
 
     @Test
     public void testBatterySaverModeMetered_disabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
     @Test
     public void testBatterySaverModeMetered_whitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
     @Test
     public void testBatterySaverModeMetered_enabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".BatterySaverModeMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
@@ -149,19 +149,19 @@
 
     @Test
     public void testBatterySaverModeNonMetered_disabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
     @Test
     public void testBatterySaverModeNonMetered_whitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
     @Test
     public void testBatterySaverModeNonMetered_enabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".BatterySaverModeNonMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
@@ -171,31 +171,31 @@
 
     @Test
     public void testAppIdleMetered_disabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
     @Test
     public void testAppIdleMetered_whitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
     @Test
     public void testAppIdleMetered_tempWhitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_tempWhitelisted");
     }
 
     @Test
     public void testAppIdleMetered_enabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
     @Test
     public void testAppIdleMetered_idleWhitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testAppIdleNetworkAccess_idleWhitelisted");
     }
 
@@ -206,51 +206,51 @@
 
     @Test
     public void testAppIdleNonMetered_disabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
 
     @Test
     public void testAppIdleNonMetered_whitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
     @Test
     public void testAppIdleNonMetered_tempWhitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_tempWhitelisted");
     }
 
     @Test
     public void testAppIdleNonMetered_enabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
     @Test
     public void testAppIdleNonMetered_idleWhitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testAppIdleNetworkAccess_idleWhitelisted");
     }
 
     @Test
     public void testAppIdleNonMetered_whenCharging() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testAppIdleNetworkAccess_whenCharging");
     }
 
     @Test
     public void testAppIdleMetered_whenCharging() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleMeteredTest",
                 "testAppIdleNetworkAccess_whenCharging");
     }
 
     @Test
     public void testAppIdle_toast() throws Exception {
         // Check that showing a toast doesn't bring an app out of standby
-        runDeviceTests(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".AppIdleNonMeteredTest",
                 "testAppIdle_toast");
     }
 
@@ -260,25 +260,25 @@
 
     @Test
     public void testDozeModeMetered_disabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
     @Test
     public void testDozeModeMetered_whitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
     @Test
     public void testDozeModeMetered_enabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
     @Test
     public void testDozeModeMetered_enabledButWhitelistedOnNotificationAction() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeMeteredTest",
                 "testBackgroundNetworkAccess_enabledButWhitelistedOnNotificationAction");
     }
 
@@ -289,26 +289,26 @@
 
     @Test
     public void testDozeModeNonMetered_disabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
                 "testBackgroundNetworkAccess_disabled");
     }
 
     @Test
     public void testDozeModeNonMetered_whitelisted() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
                 "testBackgroundNetworkAccess_whitelisted");
     }
 
     @Test
     public void testDozeModeNonMetered_enabled() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
                 "testBackgroundNetworkAccess_enabled");
     }
 
     @Test
     public void testDozeModeNonMetered_enabledButWhitelistedOnNotificationAction()
             throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".DozeModeNonMeteredTest",
                 "testBackgroundNetworkAccess_enabledButWhitelistedOnNotificationAction");
     }
 
@@ -318,55 +318,55 @@
 
     @Test
     public void testDataAndBatterySaverModes_meteredNetwork() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDataAndBatterySaverModes_meteredNetwork");
     }
 
     @Test
     public void testDataAndBatterySaverModes_nonMeteredNetwork() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDataAndBatterySaverModes_nonMeteredNetwork");
     }
 
     @Test
     public void testDozeAndBatterySaverMode_powerSaveWhitelists() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDozeAndBatterySaverMode_powerSaveWhitelists");
     }
 
     @Test
     public void testDozeAndAppIdle_powerSaveWhitelists() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDozeAndAppIdle_powerSaveWhitelists");
     }
 
     @Test
     public void testAppIdleAndDoze_tempPowerSaveWhitelists() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndDoze_tempPowerSaveWhitelists");
     }
 
     @Test
     public void testAppIdleAndBatterySaver_tempPowerSaveWhitelists() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndBatterySaver_tempPowerSaveWhitelists");
     }
 
     @Test
     public void testDozeAndAppIdle_appIdleWhitelist() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testDozeAndAppIdle_appIdleWhitelist");
     }
 
     @Test
     public void testAppIdleAndDoze_tempPowerSaveAndAppIdleWhitelists() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndDoze_tempPowerSaveAndAppIdleWhitelists");
     }
 
     @Test
     public void testAppIdleAndBatterySaver_tempPowerSaveAndAppIdleWhitelists() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".MixedModesTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".MixedModesTest",
                 "testAppIdleAndBatterySaver_tempPowerSaveAndAppIdleWhitelists");
     }
 
@@ -376,13 +376,13 @@
 
     @Test
     public void testNetworkAccess_restrictedMode() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".RestrictedModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".RestrictedModeTest",
                 "testNetworkAccess");
     }
 
     @Test
     public void testNetworkAccess_restrictedMode_withBatterySaver() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".RestrictedModeTest",
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".RestrictedModeTest",
                 "testNetworkAccess_withBatterySaver");
     }
 
@@ -392,12 +392,12 @@
 
     @Test
     public void testMeteredNetworkAccess_expeditedJob() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".ExpeditedJobMeteredTest");
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".ExpeditedJobMeteredTest");
     }
 
     @Test
     public void testNonMeteredNetworkAccess_expeditedJob() throws Exception {
-        runDeviceTests(TEST_PKG, TEST_PKG + ".ExpeditedJobNonMeteredTest");
+        runDeviceTestsWithCustomOptions(TEST_PKG, TEST_PKG + ".ExpeditedJobNonMeteredTest");
     }
 
     /*******************
diff --git a/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt b/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
index 115210b..c883b78 100644
--- a/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
+++ b/tests/cts/multidevices/snippet/ConnectivityMultiDevicesSnippet.kt
@@ -35,6 +35,7 @@
 import android.net.wifi.WifiSsid
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.testutils.ConnectUtil
+import com.android.testutils.NetworkCallbackHelper
 import com.android.testutils.RecorderCallback.CallbackEntry.Available
 import com.android.testutils.RecorderCallback.CallbackEntry.CapabilitiesChanged
 import com.android.testutils.TestableNetworkCallback
@@ -48,9 +49,14 @@
     private val cm = context.getSystemService(ConnectivityManager::class.java)!!
     private val pm = context.packageManager
     private val ctsNetUtils = CtsNetUtils(context)
+    private val cbHelper = NetworkCallbackHelper()
     private val ctsTetheringUtils = CtsTetheringUtils(context)
     private var oldSoftApConfig: SoftApConfiguration? = null
 
+    override fun shutdown() {
+        cbHelper.unregisterAll()
+    }
+
     @Rpc(description = "Check whether the device has wifi feature.")
     fun hasWifiFeature() = pm.hasSystemFeature(FEATURE_WIFI)
 
@@ -65,13 +71,13 @@
     @Rpc(description = "Request cellular connection and ensure it is the default network.")
     fun requestCellularAndEnsureDefault() {
         ctsNetUtils.disableWifi()
-        val network = ctsNetUtils.connectToCell()
+        val network = cbHelper.requestCell()
         ctsNetUtils.expectNetworkIsSystemDefault(network)
     }
 
     @Rpc(description = "Unrequest cellular connection.")
     fun unrequestCellular() {
-        ctsNetUtils.disconnectFromCell()
+        cbHelper.unrequestCell()
     }
 
     @Rpc(description = "Ensure any wifi is connected and is the default network.")
diff --git a/tests/cts/net/src/android/net/cts/ApfIntegrationTest.kt b/tests/cts/net/src/android/net/cts/ApfIntegrationTest.kt
new file mode 100644
index 0000000..3be44f7
--- /dev/null
+++ b/tests/cts/net/src/android/net/cts/ApfIntegrationTest.kt
@@ -0,0 +1,239 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// ktlint does not allow annotating function argument literals inline. Disable the specific rule
+// since this negatively affects readability.
+@file:Suppress("ktlint:standard:comment-wrapping")
+
+package android.net.cts
+
+import android.Manifest.permission.WRITE_DEVICE_CONFIG
+import android.content.pm.PackageManager.FEATURE_WIFI
+import android.net.ConnectivityManager
+import android.net.NetworkCapabilities
+import android.net.NetworkRequest
+import android.net.apf.ApfCapabilities
+import android.os.Build
+import android.os.PowerManager
+import android.platform.test.annotations.AppModeFull
+import android.provider.DeviceConfig
+import android.provider.DeviceConfig.NAMESPACE_CONNECTIVITY
+import android.system.OsConstants
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.compatibility.common.util.PropertyUtil.getVsrApiLevel
+import com.android.compatibility.common.util.SystemUtil.runShellCommand
+import com.android.compatibility.common.util.SystemUtil.runShellCommandOrThrow
+import com.android.internal.util.HexDump
+import com.android.testutils.DevSdkIgnoreRule
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import com.android.testutils.DevSdkIgnoreRunner
+import com.android.testutils.NetworkStackModuleTest
+import com.android.testutils.RecorderCallback.CallbackEntry.LinkPropertiesChanged
+import com.android.testutils.SkipPresubmit
+import com.android.testutils.TestableNetworkCallback
+import com.android.testutils.runAsShell
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
+import com.google.common.truth.TruthJUnit.assume
+import java.lang.Thread
+import kotlin.random.Random
+import kotlin.test.assertNotNull
+import org.junit.After
+import org.junit.Before
+import org.junit.BeforeClass
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+private const val TAG = "ApfIntegrationTest"
+private const val TIMEOUT_MS = 2000L
+private const val APF_NEW_RA_FILTER_VERSION = "apf_new_ra_filter_version"
+private const val POLLING_INTERVAL_MS: Int = 100
+
+@AppModeFull(reason = "CHANGE_NETWORK_STATE permission can't be granted to instant apps")
+@RunWith(DevSdkIgnoreRunner::class)
+@NetworkStackModuleTest
+class ApfIntegrationTest {
+    companion object {
+        @BeforeClass
+        @JvmStatic
+        @Suppress("ktlint:standard:no-multi-spaces")
+        fun setupOnce() {
+            // TODO: check that there is no active wifi network. Otherwise, ApfFilter has already been
+            // created.
+            // APF adb cmds are only implemented in ApfFilter.java. Enable experiment to prevent
+            // LegacyApfFilter.java from being used.
+            runAsShell(WRITE_DEVICE_CONFIG) {
+                DeviceConfig.setProperty(
+                        NAMESPACE_CONNECTIVITY,
+                        APF_NEW_RA_FILTER_VERSION,
+                        "1",  // value => force enabled
+                        false // makeDefault
+                )
+            }
+        }
+    }
+
+    @get:Rule
+    val ignoreRule = DevSdkIgnoreRule()
+
+    private val context by lazy { InstrumentationRegistry.getInstrumentation().context }
+    private val cm by lazy { context.getSystemService(ConnectivityManager::class.java)!! }
+    private val pm by lazy { context.packageManager }
+    private val powerManager by lazy { context.getSystemService(PowerManager::class.java)!! }
+    private val wakeLock by lazy { powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG) }
+    private lateinit var ifname: String
+    private lateinit var networkCallback: TestableNetworkCallback
+    private lateinit var caps: ApfCapabilities
+
+    fun getApfCapabilities(): ApfCapabilities {
+        val caps = runShellCommand("cmd network_stack apf $ifname capabilities").trim()
+        if (caps.isEmpty()) {
+            return ApfCapabilities(0, 0, 0)
+        }
+        val (version, maxLen, packetFormat) = caps.split(",").map { it.toInt() }
+        return ApfCapabilities(version, maxLen, packetFormat)
+    }
+
+    fun pollingCheck(condition: () -> Boolean, timeout_ms: Int): Boolean {
+        var polling_time = 0
+        do {
+            Thread.sleep(POLLING_INTERVAL_MS.toLong())
+            polling_time += POLLING_INTERVAL_MS
+            if (condition()) return true
+        } while (polling_time < timeout_ms)
+        return false
+    }
+
+    fun turnScreenOff() {
+        if (!wakeLock.isHeld()) wakeLock.acquire()
+        runShellCommandOrThrow("input keyevent KEYCODE_SLEEP")
+        val result = pollingCheck({ !powerManager.isInteractive() }, timeout_ms = 2000)
+        assertThat(result).isTrue()
+    }
+
+    fun turnScreenOn() {
+        if (wakeLock.isHeld()) wakeLock.release()
+        runShellCommandOrThrow("input keyevent KEYCODE_WAKEUP")
+        val result = pollingCheck({ powerManager.isInteractive() }, timeout_ms = 2000)
+        assertThat(result).isTrue()
+    }
+
+    @Before
+    fun setUp() {
+        assume().that(pm.hasSystemFeature(FEATURE_WIFI)).isTrue()
+        // APF must run when the screen is off and the device is not interactive.
+        // TODO: consider running some of the tests with screen on (capabilities, read / write).
+        turnScreenOff()
+
+        networkCallback = TestableNetworkCallback()
+        cm.requestNetwork(
+                NetworkRequest.Builder()
+                        .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+                        .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                        .build(),
+                networkCallback
+        )
+        networkCallback.eventuallyExpect<LinkPropertiesChanged>(TIMEOUT_MS) {
+            ifname = assertNotNull(it.lp.interfaceName)
+            true
+        }
+        // It's possible the device does not support APF, in which case this command will not be
+        // successful. Ignore the error as testApfCapabilities() already asserts APF support on the
+        // respective VSR releases and all other tests are based on the capabilities indicated.
+        runShellCommand("cmd network_stack apf $ifname pause")
+        caps = getApfCapabilities()
+    }
+
+    @After
+    fun tearDown() {
+        if (::ifname.isInitialized) {
+            runShellCommand("cmd network_stack apf $ifname resume")
+        }
+        if (::networkCallback.isInitialized) {
+            cm.unregisterNetworkCallback(networkCallback)
+        }
+        turnScreenOn()
+    }
+
+    @Test
+    fun testApfCapabilities() {
+        // APF became mandatory in Android 14 VSR.
+        assume().that(getVsrApiLevel()).isAtLeast(34)
+
+        // ApfFilter does not support anything but ARPHRD_ETHER.
+        assertThat(caps.apfPacketFormat).isEqualTo(OsConstants.ARPHRD_ETHER)
+
+        // DEVICEs launching with Android 14 with CHIPSETs that set ro.board.first_api_level to 34:
+        // - [GMS-VSR-5.3.12-003] MUST return 4 or higher as the APF version number from calls to
+        //   the getApfPacketFilterCapabilities HAL method.
+        // - [GMS-VSR-5.3.12-004] MUST indicate at least 1024 bytes of usable memory from calls to
+        //   the getApfPacketFilterCapabilities HAL method.
+        // TODO: check whether above text should be changed "34 or higher"
+        // This should assert apfVersionSupported >= 4 as per the VSR requirements, but there are
+        // currently no tests for APFv6 and there cannot be a valid implementation as the
+        // interpreter has yet to be finalized.
+        assertThat(caps.apfVersionSupported).isEqualTo(4)
+        assertThat(caps.maximumApfProgramSize).isAtLeast(1024)
+
+        // DEVICEs launching with Android 15 (AOSP experimental) or higher with CHIPSETs that set
+        // ro.board.first_api_level or ro.board.api_level to 202404 or higher:
+        // - [GMS-VSR-5.3.12-009] MUST indicate at least 2000 bytes of usable memory from calls to
+        //   the getApfPacketFilterCapabilities HAL method.
+        if (getVsrApiLevel() >= 202404) {
+            assertThat(caps.maximumApfProgramSize).isAtLeast(2000)
+        }
+    }
+
+    // APF is backwards compatible, i.e. a v6 interpreter supports both v2 and v4 functionality.
+    fun assumeApfVersionSupportAtLeast(version: Int) {
+        assume().that(caps.apfVersionSupported).isAtLeast(version)
+    }
+
+    fun installProgram(bytes: ByteArray) {
+        val prog = HexDump.toHexString(bytes, 0 /* offset */, bytes.size, false /* upperCase */)
+        val result = runShellCommandOrThrow("cmd network_stack apf $ifname install $prog").trim()
+        // runShellCommandOrThrow only throws on S+.
+        assertThat(result).isEqualTo("success")
+    }
+
+    fun readProgram(): ByteArray {
+        val progHexString = runShellCommandOrThrow("cmd network_stack apf $ifname read").trim()
+        // runShellCommandOrThrow only throws on S+.
+        assertThat(progHexString).isNotEmpty()
+        return HexDump.hexStringToByteArray(progHexString)
+    }
+
+    @SkipPresubmit(reason = "This test takes longer than 1 minute, do not run it on presubmit.")
+    // APF integration is mostly broken before V, only run the full read / write test on V+.
+    @IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    @Test
+    fun testReadWriteProgram() {
+        assumeApfVersionSupportAtLeast(4)
+
+        // Only test down to 2 bytes. The first byte always stays PASS.
+        val program = ByteArray(caps.maximumApfProgramSize)
+        for (i in caps.maximumApfProgramSize downTo 2) {
+            // Randomize bytes in range [1, i). And install first [0, i) bytes of program.
+            // Note that only the very first instruction (PASS) is valid APF bytecode.
+            Random.nextBytes(program, 1 /* fromIndex */, i /* toIndex */)
+            installProgram(program.sliceArray(0..<i))
+
+            // Compare entire memory region.
+            val readResult = readProgram()
+            assertWithMessage("read/write $i byte prog failed").that(readResult).isEqualTo(program)
+        }
+    }
+}
diff --git a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
index 3e5d0ba..16a7b73 100644
--- a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
@@ -48,6 +48,7 @@
 import androidx.test.filters.SdkSuppress;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.testutils.AutoReleaseNetworkCallbackRule;
 import com.android.testutils.DevSdkIgnoreRule;
 
 import org.junit.Before;
@@ -67,7 +68,10 @@
 @RunWith(AndroidJUnit4.class)
 @SdkSuppress(minSdkVersion = Build.VERSION_CODES.R) // BatteryStatsManager did not exist on Q
 public class BatteryStatsManagerTest{
-    @Rule
+    @Rule(order = 1)
+    public final AutoReleaseNetworkCallbackRule
+            networkCallbackRule = new AutoReleaseNetworkCallbackRule();
+    @Rule(order = 2)
     public final DevSdkIgnoreRule ignoreRule = new DevSdkIgnoreRule();
     private static final String TAG = BatteryStatsManagerTest.class.getSimpleName();
     private static final String TEST_URL = "https://connectivitycheck.gstatic.com/generate_204";
@@ -145,7 +149,7 @@
             return;
         }
 
-        final Network cellNetwork = mCtsNetUtils.connectToCell();
+        final Network cellNetwork = networkCallbackRule.requestCell();
         final URL url = new URL(TEST_URL);
 
         // Get cellular battery stats
diff --git a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
index 99222dd..07e2024 100644
--- a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
+++ b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
@@ -47,6 +47,7 @@
 import com.android.modules.utils.build.SdkLevel.isAtLeastR
 import com.android.net.module.util.NetworkStackConstants.TEST_CAPTIVE_PORTAL_HTTPS_URL
 import com.android.net.module.util.NetworkStackConstants.TEST_CAPTIVE_PORTAL_HTTP_URL
+import com.android.testutils.AutoReleaseNetworkCallbackRule
 import com.android.testutils.DeviceConfigRule
 import com.android.testutils.RecorderCallback.CallbackEntry.CapabilitiesChanged
 import com.android.testutils.SkipMainlinePresubmit
@@ -101,9 +102,12 @@
 
     private val server = TestHttpServer("localhost")
 
-    @get:Rule
+    @get:Rule(order = 1)
     val deviceConfigRule = DeviceConfigRule(retryCountBeforeSIfConfigChanged = 5)
 
+    @get:Rule(order = 2)
+    val networkCallbackRule = AutoReleaseNetworkCallbackRule()
+
     companion object {
         @JvmStatic @BeforeClass
         fun setUpClass() {
@@ -144,15 +148,15 @@
         assumeTrue(pm.hasSystemFeature(FEATURE_WIFI))
         assumeFalse(pm.hasSystemFeature(FEATURE_WATCH))
         utils.ensureWifiConnected()
-        val cellNetwork = utils.connectToCell()
+        val cellNetwork = networkCallbackRule.requestCell()
 
         // Verify cell network is validated
         val cellReq = NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_CELLULAR)
                 .addCapability(NET_CAPABILITY_INTERNET)
                 .build()
-        val cellCb = TestableNetworkCallback(timeoutMs = TEST_TIMEOUT_MS)
-        cm.registerNetworkCallback(cellReq, cellCb)
+        val cellCb = networkCallbackRule.registerNetworkCallback(cellReq,
+            TestableNetworkCallback(timeoutMs = TEST_TIMEOUT_MS))
         val cb = cellCb.poll { it.network == cellNetwork &&
                 it is CapabilitiesChanged && it.caps.hasCapability(NET_CAPABILITY_VALIDATED)
         }
@@ -213,8 +217,6 @@
         } finally {
             cm.unregisterNetworkCallback(wifiCb)
             server.stop()
-            // disconnectFromCell should be called after connectToCell
-            utils.disconnectFromCell()
         }
     }
 
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index cdf8340..c0f1080 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -71,7 +71,9 @@
 import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VPN;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_PARTIAL_CONNECTIVITY;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED;
 import static android.net.NetworkCapabilities.TRANSPORT_TEST;
 import static android.net.NetworkCapabilities.TRANSPORT_VPN;
@@ -81,7 +83,6 @@
 import static android.net.cts.util.CtsNetUtils.HTTP_PORT;
 import static android.net.cts.util.CtsNetUtils.NETWORK_CALLBACK_ACTION;
 import static android.net.cts.util.CtsNetUtils.TEST_HOST;
-import static android.net.cts.util.CtsNetUtils.TestNetworkCallback;
 import static android.net.cts.util.CtsTetheringUtils.TestTetheringEventCallback;
 import static android.os.MessageQueue.OnFileDescriptorEventListener.EVENT_INPUT;
 import static android.os.Process.INVALID_UID;
@@ -191,6 +192,7 @@
 import com.android.networkstack.apishim.ConstantsShim;
 import com.android.networkstack.apishim.NetworkInformationShimImpl;
 import com.android.networkstack.apishim.common.ConnectivityManagerShim;
+import com.android.testutils.AutoReleaseNetworkCallbackRule;
 import com.android.testutils.CompatUtil;
 import com.android.testutils.ConnectivityModuleTest;
 import com.android.testutils.DevSdkIgnoreRule;
@@ -259,10 +261,14 @@
 
 @RunWith(AndroidJUnit4.class)
 public class ConnectivityManagerTest {
-    @Rule
+    @Rule(order = 1)
     public final DevSdkIgnoreRule ignoreRule = new DevSdkIgnoreRule();
 
-    @Rule
+    @Rule(order = 2)
+    public final AutoReleaseNetworkCallbackRule
+            networkCallbackRule = new AutoReleaseNetworkCallbackRule();
+
+    @Rule(order = 3)
     public final DeviceConfigRule mTestValidationConfigRule = new DeviceConfigRule(
             5 /* retryCountBeforeSIfConfigChanged */);
 
@@ -328,8 +334,6 @@
     private final ArraySet<Integer> mNetworkTypes = new ArraySet<>();
     private UiAutomation mUiAutomation;
     private CtsNetUtils mCtsNetUtils;
-    // The registered callbacks.
-    private List<NetworkCallback> mRegisteredCallbacks = new ArrayList<>();
     // Used for cleanup purposes.
     private final List<Range<Integer>> mVpnRequiredUidRanges = new ArrayList<>();
 
@@ -411,11 +415,6 @@
 
     @After
     public void tearDown() throws Exception {
-        // Release any NetworkRequests filed to connect mobile data.
-        if (mCtsNetUtils.cellConnectAttempted()) {
-            mCtsNetUtils.disconnectFromCell();
-        }
-
         if (TestUtils.shouldTestSApis()) {
             runWithShellPermissionIdentity(
                     () -> mCmShim.setRequireVpnForUids(false, mVpnRequiredUidRanges),
@@ -425,15 +424,12 @@
         // All tests in this class require a working Internet connection as they start. Make
         // sure there is still one as they end that's ready to use for the next test to use.
         mTestValidationConfigRule.runAfterNextCleanup(() -> {
-            final TestNetworkCallback callback = new TestNetworkCallback();
-            registerDefaultNetworkCallback(callback);
-            try {
-                assertNotNull("Couldn't restore Internet connectivity",
-                        callback.waitForAvailable());
-            } finally {
-                // Unregister all registered callbacks.
-                unregisterRegisteredCallbacks();
-            }
+            // mTestValidationConfigRule has higher order than networkCallbackRule, so
+            // networkCallbackRule is the outer rule and will be cleaned up after this method.
+            final TestableNetworkCallback callback =
+                    networkCallbackRule.registerDefaultNetworkCallback();
+            assertNotNull("Couldn't restore Internet connectivity",
+                    callback.eventuallyExpect(CallbackEntry.AVAILABLE));
         });
     }
 
@@ -555,7 +551,7 @@
             throws InterruptedException {
         assumeTrue(mPackageManager.hasSystemFeature(FEATURE_TELEPHONY));
         // Make sure cell is active to retrieve IMSI for verification in later step.
-        final Network cellNetwork = mCtsNetUtils.connectToCell();
+        final Network cellNetwork = networkCallbackRule.requestCell();
         final String subscriberId = getSubscriberIdForCellNetwork(cellNetwork);
         assertFalse(TextUtils.isEmpty(subscriberId));
 
@@ -853,7 +849,7 @@
         assumeTrue(mPackageManager.hasSystemFeature(FEATURE_TELEPHONY));
 
         Network wifiNetwork = mCtsNetUtils.ensureWifiConnected();
-        Network cellNetwork = mCtsNetUtils.connectToCell();
+        Network cellNetwork = networkCallbackRule.requestCell();
         // This server returns the requestor's IP address as the response body.
         URL url = new URL("http://google-ipv6test.appspot.com/ip.js?fmt=text");
         String wifiAddressString = httpGet(wifiNetwork, url);
@@ -993,10 +989,10 @@
         // default network.
         return new NetworkRequest.Builder()
                 .clearCapabilities()
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN)
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                .addCapability(NET_CAPABILITY_NOT_RESTRICTED)
+                .addCapability(NET_CAPABILITY_TRUSTED)
+                .addCapability(NET_CAPABILITY_NOT_VPN)
+                .addCapability(NET_CAPABILITY_INTERNET)
                 .build();
     }
 
@@ -1027,10 +1023,10 @@
         final String invalidPrivateDnsServer = "invalidhostname.example.com";
         final String goodPrivateDnsServer = "dns.google";
         mCtsNetUtils.storePrivateDnsSetting();
-        final TestableNetworkCallback cb = new TestableNetworkCallback();
         final NetworkRequest networkRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_INTERNET).build();
-        registerNetworkCallback(networkRequest, cb);
+        final TestableNetworkCallback cb =
+                networkCallbackRule.registerNetworkCallback(networkRequest);
         final Network networkForPrivateDns = mCm.getActiveNetwork();
         try {
             // Verifying the good private DNS sever
@@ -1068,24 +1064,27 @@
         assumeTrue(mPackageManager.hasSystemFeature(FEATURE_WIFI));
 
         // We will register for a WIFI network being available or lost.
-        final TestNetworkCallback callback = new TestNetworkCallback();
-        registerNetworkCallback(makeWifiNetworkRequest(), callback);
+        final TestableNetworkCallback callback = networkCallbackRule.registerNetworkCallback(
+                makeWifiNetworkRequest());
 
-        final TestNetworkCallback defaultTrackingCallback = new TestNetworkCallback();
-        registerDefaultNetworkCallback(defaultTrackingCallback);
+        final TestableNetworkCallback defaultTrackingCallback =
+                networkCallbackRule.registerDefaultNetworkCallback();
 
-        final TestNetworkCallback systemDefaultCallback = new TestNetworkCallback();
-        final TestNetworkCallback perUidCallback = new TestNetworkCallback();
-        final TestNetworkCallback bestMatchingCallback = new TestNetworkCallback();
+        final TestableNetworkCallback systemDefaultCallback = new TestableNetworkCallback();
+        final TestableNetworkCallback perUidCallback = new TestableNetworkCallback();
+        final TestableNetworkCallback bestMatchingCallback = new TestableNetworkCallback();
         final Handler h = new Handler(Looper.getMainLooper());
         if (TestUtils.shouldTestSApis()) {
             assertThrows(SecurityException.class, () ->
-                    registerSystemDefaultNetworkCallback(systemDefaultCallback, h));
+                    networkCallbackRule.registerSystemDefaultNetworkCallback(
+                            systemDefaultCallback, h));
             runWithShellPermissionIdentity(() -> {
-                registerSystemDefaultNetworkCallback(systemDefaultCallback, h);
-                registerDefaultNetworkCallbackForUid(Process.myUid(), perUidCallback, h);
+                networkCallbackRule.registerSystemDefaultNetworkCallback(systemDefaultCallback, h);
+                networkCallbackRule.registerDefaultNetworkCallbackForUid(Process.myUid(),
+                        perUidCallback, h);
             }, NETWORK_SETTINGS);
-            registerBestMatchingNetworkCallback(makeDefaultRequest(), bestMatchingCallback, h);
+            networkCallbackRule.registerBestMatchingNetworkCallback(
+                    makeDefaultRequest(), bestMatchingCallback, h);
         }
 
         Network wifiNetwork = null;
@@ -1094,24 +1093,22 @@
         // Now we should expect to get a network callback about availability of the wifi
         // network even if it was already connected as a state-based action when the callback
         // is registered.
-        wifiNetwork = callback.waitForAvailable();
+        wifiNetwork = callback.eventuallyExpect(CallbackEntry.AVAILABLE).getNetwork();
         assertNotNull("Did not receive onAvailable for TRANSPORT_WIFI request",
                 wifiNetwork);
 
-        final Network defaultNetwork = defaultTrackingCallback.waitForAvailable();
+        final Network defaultNetwork = defaultTrackingCallback.eventuallyExpect(
+                CallbackEntry.AVAILABLE).getNetwork();
         assertNotNull("Did not receive onAvailable on default network callback",
                 defaultNetwork);
 
         if (TestUtils.shouldTestSApis()) {
-            assertNotNull("Did not receive onAvailable on system default network callback",
-                    systemDefaultCallback.waitForAvailable());
-            final Network perUidNetwork = perUidCallback.waitForAvailable();
-            assertNotNull("Did not receive onAvailable on per-UID default network callback",
-                    perUidNetwork);
+            systemDefaultCallback.eventuallyExpect(CallbackEntry.AVAILABLE);
+            final Network perUidNetwork = perUidCallback.eventuallyExpect(CallbackEntry.AVAILABLE)
+                    .getNetwork();
             assertEquals(defaultNetwork, perUidNetwork);
-            final Network bestMatchingNetwork = bestMatchingCallback.waitForAvailable();
-            assertNotNull("Did not receive onAvailable on best matching network callback",
-                    bestMatchingNetwork);
+            final Network bestMatchingNetwork = bestMatchingCallback.eventuallyExpect(
+                    CallbackEntry.AVAILABLE).getNetwork();
             assertEquals(defaultNetwork, bestMatchingNetwork);
         }
     }
@@ -1123,8 +1120,8 @@
         final Handler h = new Handler(Looper.getMainLooper());
         // Verify registerSystemDefaultNetworkCallback can be accessed via
         // CONNECTIVITY_USE_RESTRICTED_NETWORKS permission.
-        runWithShellPermissionIdentity(() ->
-                        registerSystemDefaultNetworkCallback(new TestNetworkCallback(), h),
+        runWithShellPermissionIdentity(
+                () -> networkCallbackRule.registerSystemDefaultNetworkCallback(h),
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS);
     }
 
@@ -1294,15 +1291,14 @@
      */
     @AppModeFull(reason = "CHANGE_NETWORK_STATE permission can't be granted to instant apps")
     @Test
-    public void testRequestNetworkCallback() throws Exception {
-        final TestNetworkCallback callback = new TestNetworkCallback();
-        requestNetwork(new NetworkRequest.Builder()
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
-                .build(), callback);
+    public void testRequestNetworkCallback() {
+        final TestableNetworkCallback callback = networkCallbackRule.requestNetwork(
+                new NetworkRequest.Builder().addCapability(
+                                NET_CAPABILITY_INTERNET)
+                        .build());
 
         // Wait to get callback for availability of internet
-        Network internetNetwork = callback.waitForAvailable();
-        assertNotNull("Did not receive NetworkCallback#onAvailable for INTERNET", internetNetwork);
+        callback.eventuallyExpect(CallbackEntry.AVAILABLE).getNetwork();
     }
 
     /**
@@ -1320,16 +1316,13 @@
             }
         }
 
-        final TestNetworkCallback callback = new TestNetworkCallback();
-        requestNetwork(new NetworkRequest.Builder().addTransportType(TRANSPORT_WIFI).build(),
-                callback, 100);
 
+        final TestableNetworkCallback callback = networkCallbackRule.requestNetwork(
+                new NetworkRequest.Builder().addTransportType(TRANSPORT_WIFI).build(),
+                100 /* timeoutMs */);
         try {
             // Wait to get callback for unavailability of requested network
-            assertTrue("Did not receive NetworkCallback#onUnavailable",
-                    callback.waitForUnavailable());
-        } catch (InterruptedException e) {
-            fail("NetworkCallback wait was interrupted.");
+            callback.eventuallyExpect(CallbackEntry.UNAVAILABLE, 2_000 /* timeoutMs */);
         } finally {
             if (previousWifiEnabledState) {
                 mCtsNetUtils.connectToWifi();
@@ -1416,40 +1409,48 @@
             final boolean useSystemDefault)
             throws Exception {
         final CompletableFuture<Network> networkFuture = new CompletableFuture<>();
-        final NetworkCallback networkCallback = new NetworkCallback() {
-            @Override
-            public void onCapabilitiesChanged(Network network, NetworkCapabilities nc) {
-                if (!nc.hasTransport(targetTransportType)) return;
 
-                final boolean metered = !nc.hasCapability(NET_CAPABILITY_NOT_METERED);
-                final boolean validated = nc.hasCapability(NET_CAPABILITY_VALIDATED);
-                if (metered == requestedMeteredness && (!waitForValidation || validated)) {
-                    networkFuture.complete(network);
+        // Registering a callback here guarantees onCapabilitiesChanged is called immediately
+        // with the current setting. Therefore, if the setting has already been changed,
+        // this method will return right away, and if not, it'll wait for the setting to change.
+        final TestableNetworkCallback networkCallback;
+        if (useSystemDefault) {
+            networkCallback = runWithShellPermissionIdentity(() -> {
+                if (isAtLeastS()) {
+                    return networkCallbackRule.registerSystemDefaultNetworkCallback(
+                            new Handler(Looper.getMainLooper()));
+                } else {
+                    // registerSystemDefaultNetworkCallback is only supported on S+.
+                    return networkCallbackRule.requestNetwork(
+                            new NetworkRequest.Builder()
+                                    .clearCapabilities()
+                                    .addCapability(NET_CAPABILITY_NOT_RESTRICTED)
+                                    .addCapability(NET_CAPABILITY_TRUSTED)
+                                    .addCapability(NET_CAPABILITY_NOT_VPN)
+                                    .addCapability(NET_CAPABILITY_INTERNET)
+                                    .build(),
+                            new TestableNetworkCallback(),
+                            new Handler(Looper.getMainLooper()));
                 }
-            }
-        };
-
-        try {
-            // Registering a callback here guarantees onCapabilitiesChanged is called immediately
-            // with the current setting. Therefore, if the setting has already been changed,
-            // this method will return right away, and if not, it'll wait for the setting to change.
-            if (useSystemDefault) {
-                runWithShellPermissionIdentity(() ->
-                                registerSystemDefaultNetworkCallback(networkCallback,
-                                        new Handler(Looper.getMainLooper())),
-                        NETWORK_SETTINGS);
-            } else {
-                registerDefaultNetworkCallback(networkCallback);
-            }
-
-            // Changing meteredness on wifi involves reconnecting, which can take several seconds
-            // (involves re-associating, DHCP...).
-            return networkFuture.get(NETWORK_CALLBACK_TIMEOUT_MS, TimeUnit.MILLISECONDS);
-        } catch (TimeoutException e) {
-            throw new AssertionError("Timed out waiting for active network metered status to "
-                    + "change to " + requestedMeteredness + " ; network = "
-                    + mCm.getActiveNetwork(), e);
+            },
+            NETWORK_SETTINGS);
+        } else {
+            networkCallback = networkCallbackRule.registerDefaultNetworkCallback();
         }
+
+        return networkCallback.eventuallyExpect(
+                CallbackEntry.NETWORK_CAPS_UPDATED,
+                // Changing meteredness on wifi involves reconnecting, which can take several
+                // seconds (involves re-associating, DHCP...).
+                NETWORK_CALLBACK_TIMEOUT_MS,
+                cb -> {
+                    final NetworkCapabilities nc = cb.getCaps();
+                    if (!nc.hasTransport(targetTransportType)) return false;
+
+                    final boolean metered = !nc.hasCapability(NET_CAPABILITY_NOT_METERED);
+                    final boolean validated = nc.hasCapability(NET_CAPABILITY_VALIDATED);
+                    return metered == requestedMeteredness && (!waitForValidation || validated);
+                }).getNetwork();
     }
 
     private Network setWifiMeteredStatusAndWait(String ssid, boolean isMetered,
@@ -2025,7 +2026,7 @@
             return;
         }
 
-        final Network network = mCtsNetUtils.connectToCell();
+        final Network network = networkCallbackRule.requestCell();
         final int supported = getSupportedKeepalivesForNet(network);
         final InetAddress srcAddr = getFirstV4Address(network);
         assumeTrue("This test requires native IPv4", srcAddr != null);
@@ -2091,15 +2092,15 @@
     }
 
     private void verifyBindSocketToRestrictedNetworkDisallowed() throws Exception {
-        final TestableNetworkCallback testNetworkCb = new TestableNetworkCallback();
         final NetworkRequest testRequest = new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
-        runWithShellPermissionIdentity(() -> requestNetwork(testRequest, testNetworkCb),
+        final TestableNetworkCallback testNetworkCb = runWithShellPermissionIdentity(
+                () -> networkCallbackRule.requestNetwork(testRequest),
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS,
                 // CONNECTIVITY_INTERNAL is for requesting restricted network because shell does not
                 // have CONNECTIVITY_USE_RESTRICTED_NETWORKS on R.
@@ -2115,7 +2116,7 @@
                     NETWORK_CALLBACK_TIMEOUT_MS,
                     entry -> network.equals(entry.getNetwork())
                             && (!((CallbackEntry.CapabilitiesChanged) entry).getCaps()
-                            .hasCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)));
+                            .hasCapability(NET_CAPABILITY_NOT_RESTRICTED)));
             // CtsNetTestCases package doesn't hold CONNECTIVITY_USE_RESTRICTED_NETWORKS, so it
             // does not allow to bind socket to restricted network.
             assertThrows(IOException.class, () -> network.bindSocket(socket));
@@ -2200,8 +2201,7 @@
             registerCallbackAndWaitForAvailable(makeWifiNetworkRequest(), wifiCb);
         }
         if (supportTelephony) {
-            // connectToCell needs to be followed by disconnectFromCell, which is called in tearDown
-            mCtsNetUtils.connectToCell();
+            networkCallbackRule.requestCell();
             registerCallbackAndWaitForAvailable(makeCellNetworkRequest(), telephonyCb);
         }
 
@@ -2239,7 +2239,7 @@
 
     private void registerCallbackAndWaitForAvailable(@NonNull final NetworkRequest request,
             @NonNull final TestableNetworkCallback cb) {
-        registerNetworkCallback(request, cb);
+        networkCallbackRule.registerNetworkCallback(request, cb);
         waitForAvailable(cb);
     }
 
@@ -2347,21 +2347,13 @@
 
     private void verifySsidFromCallbackNetworkCapabilities(@NonNull String ssid, boolean hasSsid)
             throws Exception {
-        final CompletableFuture<NetworkCapabilities> foundNc = new CompletableFuture();
-        final NetworkCallback callback = new NetworkCallback() {
-            @Override
-            public void onCapabilitiesChanged(Network network, NetworkCapabilities nc) {
-                foundNc.complete(nc);
-            }
-        };
-
-        registerNetworkCallback(makeWifiNetworkRequest(), callback);
+        final TestableNetworkCallback callback =
+                networkCallbackRule.registerNetworkCallback(makeWifiNetworkRequest());
         // Registering a callback here guarantees onCapabilitiesChanged is called immediately
         // because WiFi network should be connected.
-        final NetworkCapabilities nc =
-                foundNc.get(NETWORK_CALLBACK_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        final NetworkCapabilities nc = callback.eventuallyExpect(
+                CallbackEntry.NETWORK_CAPS_UPDATED, NETWORK_CALLBACK_TIMEOUT_MS).getCaps();
         // Verify if ssid is contained in the NetworkCapabilities received from callback.
-        assertNotNull("NetworkCapabilities of the network is null", nc);
         assertEquals(hasSsid, Pattern.compile(ssid).matcher(nc.toString()).find());
     }
 
@@ -2390,8 +2382,8 @@
         final NetworkRequest testRequest = new NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_TEST)
                 // Test networks do not have NOT_VPN or TRUSTED capabilities by default
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_VPN)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         testNetworkInterface.getInterfaceName()))
                 .build();
@@ -2400,14 +2392,15 @@
         final TestableNetworkCallback callback = new TestableNetworkCallback();
         final Handler handler = new Handler(Looper.getMainLooper());
         assertThrows(SecurityException.class,
-                () -> requestBackgroundNetwork(testRequest, callback, handler));
+                () -> networkCallbackRule.requestBackgroundNetwork(testRequest, callback, handler));
 
         Network testNetwork = null;
         try {
             // Request background test network via Shell identity which has NETWORK_SETTINGS
             // permission granted.
             runWithShellPermissionIdentity(
-                    () -> requestBackgroundNetwork(testRequest, callback, handler),
+                    () -> networkCallbackRule.requestBackgroundNetwork(
+                            testRequest, callback, handler),
                     new String[] { android.Manifest.permission.NETWORK_SETTINGS });
 
             // Register the test network agent which has no foreground request associated to it.
@@ -2500,9 +2493,10 @@
         final int otherUid = UserHandle.getUid(5, Process.FIRST_APPLICATION_UID);
         final Handler handler = new Handler(Looper.getMainLooper());
 
-        registerDefaultNetworkCallback(myUidCallback, handler);
-        runWithShellPermissionIdentity(() -> registerDefaultNetworkCallbackForUid(
-                otherUid, otherUidCallback, handler), NETWORK_SETTINGS);
+        networkCallbackRule.registerDefaultNetworkCallback(myUidCallback, handler);
+        runWithShellPermissionIdentity(
+                () -> networkCallbackRule.registerDefaultNetworkCallbackForUid(
+                        otherUid, otherUidCallback, handler), NETWORK_SETTINGS);
 
         final Network defaultNetwork = myUidCallback.expect(CallbackEntry.AVAILABLE).getNetwork();
         final List<DetailedBlockedStatusCallback> allCallbacks =
@@ -2558,14 +2552,14 @@
         assertNotNull(info);
         assertEquals(DetailedState.CONNECTED, info.getDetailedState());
 
-        final TestableNetworkCallback callback = new TestableNetworkCallback();
+        final TestableNetworkCallback callback;
         try {
             mCmShim.setLegacyLockdownVpnEnabled(true);
 
             // setLegacyLockdownVpnEnabled is asynchronous and only takes effect when the
             // ConnectivityService handler thread processes it. Ensure it has taken effect by doing
             // something that blocks until the handler thread is idle.
-            registerDefaultNetworkCallback(callback);
+            callback = networkCallbackRule.registerDefaultNetworkCallback();
             waitForAvailable(callback);
 
             // Test one of the effects of setLegacyLockdownVpnEnabled: the fact that any NetworkInfo
@@ -2832,9 +2826,9 @@
     private void registerTestOemNetworkPreferenceCallbacks(
             @NonNull final TestableNetworkCallback defaultCallback,
             @NonNull final TestableNetworkCallback systemDefaultCallback) {
-        registerDefaultNetworkCallback(defaultCallback);
+        networkCallbackRule.registerDefaultNetworkCallback(defaultCallback);
         runWithShellPermissionIdentity(() ->
-                registerSystemDefaultNetworkCallback(systemDefaultCallback,
+                networkCallbackRule.registerSystemDefaultNetworkCallback(systemDefaultCallback,
                         new Handler(Looper.getMainLooper())), NETWORK_SETTINGS);
     }
 
@@ -2950,18 +2944,18 @@
                         + " unless device supports WiFi",
                 mPackageManager.hasSystemFeature(FEATURE_WIFI));
 
-        final TestNetworkCallback cb = new TestNetworkCallback();
         try {
             // Wait for partial connectivity to be detected on the network
             final Network network = preparePartialConnectivity();
 
-            requestNetwork(makeWifiNetworkRequest(), cb);
+            final TestableNetworkCallback cb = networkCallbackRule.requestNetwork(
+                    makeWifiNetworkRequest());
             runAsShell(NETWORK_SETTINGS, () -> {
                 // The always bit is verified in NetworkAgentTest
                 mCm.setAcceptPartialConnectivity(network, false /* accept */, false /* always */);
             });
             // Reject partial connectivity network should cause the network being torn down
-            assertEquals(network, cb.waitForLost());
+            assertEquals(network, cb.eventuallyExpect(CallbackEntry.LOST).getNetwork());
         } finally {
             mHttpServer.stop();
             // Wifi will not automatically reconnect to the network. ensureWifiDisconnected cannot
@@ -2989,17 +2983,17 @@
         assumeTrue("testAcceptPartialConnectivity_validatedNetwork cannot execute"
                         + " unless device supports WiFi and telephony", canRunTest);
 
-        final TestableNetworkCallback wifiCb = new TestableNetworkCallback();
         try {
             // Ensure at least one default network candidate connected.
-            mCtsNetUtils.connectToCell();
+            networkCallbackRule.requestCell();
 
             final Network wifiNetwork = prepareUnvalidatedNetwork();
             // Default network should not be wifi ,but checking that wifi is not the default doesn't
             // guarantee that it won't become the default in the future.
             assertNotEquals(wifiNetwork, mCm.getActiveNetwork());
 
-            registerNetworkCallback(makeWifiNetworkRequest(), wifiCb);
+            final TestableNetworkCallback wifiCb = networkCallbackRule.registerNetworkCallback(
+                    makeWifiNetworkRequest());
             runAsShell(NETWORK_SETTINGS, () -> {
                 mCm.setAcceptUnvalidated(wifiNetwork, false /* accept */, false /* always */);
             });
@@ -3026,19 +3020,19 @@
         assumeTrue("testSetAvoidUnvalidated cannot execute"
                 + " unless device supports WiFi and telephony", canRunTest);
 
-        final TestableNetworkCallback wifiCb = new TestableNetworkCallback();
-        final TestableNetworkCallback defaultCb = new TestableNetworkCallback();
         final int previousAvoidBadWifi =
                 ConnectivitySettingsManager.getNetworkAvoidBadWifi(mContext);
 
         allowBadWifi();
 
         try {
-            final Network cellNetwork = mCtsNetUtils.connectToCell();
+            final Network cellNetwork = networkCallbackRule.requestCell();
             final Network wifiNetwork = prepareValidatedNetwork();
 
-            registerDefaultNetworkCallback(defaultCb);
-            registerNetworkCallback(makeWifiNetworkRequest(), wifiCb);
+            final TestableNetworkCallback defaultCb =
+                    networkCallbackRule.registerDefaultNetworkCallback();
+            final TestableNetworkCallback wifiCb = networkCallbackRule.registerNetworkCallback(
+                    makeWifiNetworkRequest());
 
             // Verify wifi is the default network.
             defaultCb.eventuallyExpect(CallbackEntry.AVAILABLE, NETWORK_CALLBACK_TIMEOUT_MS,
@@ -3101,20 +3095,11 @@
         });
     }
 
-    private Network expectNetworkHasCapability(Network network, int expectedNetCap, long timeout)
-            throws Exception {
-        final CompletableFuture<Network> future = new CompletableFuture();
-        final NetworkCallback cb = new NetworkCallback() {
-            @Override
-            public void onCapabilitiesChanged(Network n, NetworkCapabilities nc) {
-                if (n.equals(network) && nc.hasCapability(expectedNetCap)) {
-                    future.complete(network);
-                }
-            }
-        };
-
-        registerNetworkCallback(new NetworkRequest.Builder().build(), cb);
-        return future.get(timeout, TimeUnit.MILLISECONDS);
+    private Network expectNetworkHasCapability(Network network, int expectedNetCap, long timeout) {
+        return networkCallbackRule.registerNetworkCallback(new NetworkRequest.Builder().build())
+                .eventuallyExpect(CallbackEntry.NETWORK_CAPS_UPDATED, timeout,
+                        cb -> cb.getNetwork().equals(network)
+                                && cb.getCaps().hasCapability(expectedNetCap)).getNetwork();
     }
 
     private void prepareHttpServer() throws Exception {
@@ -3214,8 +3199,6 @@
 
         if (supportWifi) {
             mCtsNetUtils.ensureWifiDisconnected(null /* wifiNetworkToCheck */);
-        } else {
-            mCtsNetUtils.disconnectFromCell();
         }
 
         final CompletableFuture<Boolean> future = new CompletableFuture<>();
@@ -3226,7 +3209,7 @@
             if (supportWifi) {
                 mCtsNetUtils.ensureWifiConnected();
             } else {
-                mCtsNetUtils.connectToCell();
+                networkCallbackRule.requestCell();
             }
             assertTrue(future.get(LISTEN_ACTIVITY_TIMEOUT_MS, TimeUnit.MILLISECONDS));
         }, () -> {
@@ -3267,13 +3250,14 @@
 
         // For testing mobile data preferred uids feature, it needs both wifi and cell network.
         final Network wifiNetwork = mCtsNetUtils.ensureWifiConnected();
-        final Network cellNetwork = mCtsNetUtils.connectToCell();
-        final TestableNetworkCallback defaultTrackingCb = new TestableNetworkCallback();
-        final TestableNetworkCallback systemDefaultCb = new TestableNetworkCallback();
+        final Network cellNetwork = networkCallbackRule.requestCell();
         final Handler h = new Handler(Looper.getMainLooper());
-        runWithShellPermissionIdentity(() -> registerSystemDefaultNetworkCallback(
-                systemDefaultCb, h), NETWORK_SETTINGS);
-        registerDefaultNetworkCallback(defaultTrackingCb);
+        final TestableNetworkCallback systemDefaultCb = runWithShellPermissionIdentity(
+                () -> networkCallbackRule.registerSystemDefaultNetworkCallback(h),
+                NETWORK_SETTINGS);
+
+        final TestableNetworkCallback defaultTrackingCb =
+                networkCallbackRule.registerDefaultNetworkCallback();
 
         try {
             // CtsNetTestCases uid is not listed in MOBILE_DATA_PREFERRED_UIDS setting, so the
@@ -3342,7 +3326,7 @@
         // Create test network agent with restricted network.
         final NetworkCapabilities nc = new NetworkCapabilities.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
@@ -3376,23 +3360,23 @@
                 mContext, originalUidsAllowedOnRestrictedNetworks), NETWORK_SETTINGS);
 
         // File a restricted network request with permission first to hold the connection.
-        final TestableNetworkCallback testNetworkCb = new TestableNetworkCallback();
         final NetworkRequest testRequest = new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
-        runWithShellPermissionIdentity(() -> requestNetwork(testRequest, testNetworkCb),
+        final TestableNetworkCallback testNetworkCb = runWithShellPermissionIdentity(
+                () -> networkCallbackRule.requestNetwork(testRequest),
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS);
 
         // File another restricted network request without permission.
         final TestableNetworkCallback restrictedNetworkCb = new TestableNetworkCallback();
         final NetworkRequest restrictedRequest = new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
@@ -3409,7 +3393,7 @@
                     NETWORK_CALLBACK_TIMEOUT_MS,
                     entry -> network.equals(entry.getNetwork())
                             && (!((CallbackEntry.CapabilitiesChanged) entry).getCaps()
-                            .hasCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)));
+                            .hasCapability(NET_CAPABILITY_NOT_RESTRICTED)));
             // CtsNetTestCases package doesn't hold CONNECTIVITY_USE_RESTRICTED_NETWORKS, so it
             // does not allow to bind socket to restricted network.
             assertThrows(IOException.class, () -> network.bindSocket(socket));
@@ -3427,13 +3411,13 @@
 
             if (TestUtils.shouldTestTApis()) {
                 // Uid is in allowed list. Try file network request again.
-                requestNetwork(restrictedRequest, restrictedNetworkCb);
+                networkCallbackRule.requestNetwork(restrictedRequest, restrictedNetworkCb);
                 // Verify that the network is restricted.
                 restrictedNetworkCb.eventuallyExpect(CallbackEntry.NETWORK_CAPS_UPDATED,
                         NETWORK_CALLBACK_TIMEOUT_MS,
                         entry -> network.equals(entry.getNetwork())
                                 && (!((CallbackEntry.CapabilitiesChanged) entry).getCaps()
-                                .hasCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)));
+                                .hasCapability(NET_CAPABILITY_NOT_RESTRICTED)));
             }
         } finally {
             agent.unregister();
@@ -3731,58 +3715,4 @@
         // shims, and @IgnoreUpTo does not check that.
         assumeTrue(TestUtils.shouldTestSApis());
     }
-
-    private void unregisterRegisteredCallbacks() {
-        for (NetworkCallback callback: mRegisteredCallbacks) {
-            mCm.unregisterNetworkCallback(callback);
-        }
-    }
-
-    private void registerDefaultNetworkCallback(NetworkCallback callback) {
-        mCm.registerDefaultNetworkCallback(callback);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerDefaultNetworkCallback(NetworkCallback callback, Handler handler) {
-        mCm.registerDefaultNetworkCallback(callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerNetworkCallback(NetworkRequest request, NetworkCallback callback) {
-        mCm.registerNetworkCallback(request, callback);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerSystemDefaultNetworkCallback(NetworkCallback callback, Handler handler) {
-        mCmShim.registerSystemDefaultNetworkCallback(callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerDefaultNetworkCallbackForUid(int uid, NetworkCallback callback,
-            Handler handler) throws Exception {
-        mCmShim.registerDefaultNetworkCallbackForUid(uid, callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void requestNetwork(NetworkRequest request, NetworkCallback callback) {
-        mCm.requestNetwork(request, callback);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void requestNetwork(NetworkRequest request, NetworkCallback callback, int timeoutSec) {
-        mCm.requestNetwork(request, callback, timeoutSec);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerBestMatchingNetworkCallback(NetworkRequest request,
-            NetworkCallback callback, Handler handler) {
-        mCm.registerBestMatchingNetworkCallback(request, callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void requestBackgroundNetwork(NetworkRequest request, NetworkCallback callback,
-            Handler handler) throws Exception {
-        mCmShim.requestBackgroundNetwork(request, callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
 }
diff --git a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
index a40ed0f..b703f77 100644
--- a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
@@ -81,6 +81,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.modules.utils.build.SdkLevel;
+import com.android.testutils.ConnectivityModuleTest;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import com.android.testutils.SkipMainlinePresubmit;
@@ -101,6 +102,7 @@
 import java.util.Map.Entry;
 import java.util.Set;
 
+@ConnectivityModuleTest
 @RunWith(AndroidJUnit4.class)
 @AppModeFull(reason = "Socket cannot bind in instant app mode")
 public class IpSecManagerTest extends IpSecBaseTest {
@@ -444,6 +446,11 @@
             long uidTxDelta = 0;
             long uidRxDelta = 0;
             for (int i = 0; i < 100; i++) {
+                // Clear TrafficStats cache is needed to avoid rate-limit caching for
+                // TrafficStats API results on V+ devices.
+                if (SdkLevel.isAtLeastV()) {
+                    runAsShell(NETWORK_SETTINGS, () -> TrafficStats.clearRateLimitCaches());
+                }
                 uidTxDelta = TrafficStats.getUidTxPackets(Os.getuid()) - uidTxPackets;
                 uidRxDelta = TrafficStats.getUidRxPackets(Os.getuid()) - uidRxPackets;
 
@@ -518,6 +525,11 @@
         }
 
         private static void initStatsChecker() throws Exception {
+            // Clear TrafficStats cache is needed to avoid rate-limit caching for
+            // TrafficStats API results on V+ devices.
+            if (SdkLevel.isAtLeastV()) {
+                runAsShell(NETWORK_SETTINGS, () -> TrafficStats.clearRateLimitCaches());
+            }
             uidTxBytes = TrafficStats.getUidTxBytes(Os.getuid());
             uidRxBytes = TrafficStats.getUidRxBytes(Os.getuid());
             uidTxPackets = TrafficStats.getUidTxPackets(Os.getuid());
diff --git a/tests/cts/net/src/android/net/cts/MultinetworkApiTest.java b/tests/cts/net/src/android/net/cts/MultinetworkApiTest.java
index bca18f5..73f65e0 100644
--- a/tests/cts/net/src/android/net/cts/MultinetworkApiTest.java
+++ b/tests/cts/net/src/android/net/cts/MultinetworkApiTest.java
@@ -16,10 +16,14 @@
 
 package android.net.cts;
 
+import static android.content.pm.PackageManager.FEATURE_TELEPHONY;
+import static android.content.pm.PackageManager.FEATURE_WIFI;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
+import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.provider.DeviceConfig.NAMESPACE_CONNECTIVITY;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
@@ -34,10 +38,12 @@
 import android.platform.test.annotations.AppModeFull;
 import android.system.ErrnoException;
 import android.system.OsConstants;
+import android.util.ArraySet;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import com.android.testutils.AutoReleaseNetworkCallbackRule;
 import com.android.testutils.DeviceConfigRule;
 
 import org.junit.Before;
@@ -45,11 +51,17 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.Set;
+
 @RunWith(AndroidJUnit4.class)
 public class MultinetworkApiTest {
-    @Rule
+    @Rule(order = 1)
     public final DeviceConfigRule mDeviceConfigRule = new DeviceConfigRule();
 
+    @Rule(order = 2)
+    public final AutoReleaseNetworkCallbackRule
+            mNetworkCallbackRule = new AutoReleaseNetworkCallbackRule();
+
     static {
         System.loadLibrary("nativemultinetwork_jni");
     }
@@ -74,9 +86,8 @@
     private ContentResolver mCR;
     private ConnectivityManager mCM;
     private CtsNetUtils mCtsNetUtils;
-    private String mOldMode;
-    private String mOldDnsSpecifier;
     private Context mContext;
+    private Network mRequestedCellNetwork;
 
     @Before
     public void setUp() throws Exception {
@@ -87,8 +98,8 @@
     }
 
     @Test
-    public void testGetaddrinfo() throws ErrnoException {
-        for (Network network : mCtsNetUtils.getTestableNetworks()) {
+    public void testGetaddrinfo() throws Exception {
+        for (Network network : getTestableNetworks()) {
             int errno = runGetaddrinfoCheck(network.getNetworkHandle());
             if (errno != 0) {
                 throw new ErrnoException(
@@ -99,12 +110,12 @@
 
     @Test
     @AppModeFull(reason = "CHANGE_NETWORK_STATE permission can't be granted to instant apps")
-    public void testSetprocnetwork() throws ErrnoException {
+    public void testSetprocnetwork() throws Exception {
         // Hopefully no prior test in this process space has set a default network.
         assertNull(mCM.getProcessDefaultNetwork());
         assertEquals(0, NetworkUtils.getBoundNetworkForProcess());
 
-        for (Network network : mCtsNetUtils.getTestableNetworks()) {
+        for (Network network : getTestableNetworks()) {
             mCM.setProcessDefaultNetwork(null);
             assertNull(mCM.getProcessDefaultNetwork());
 
@@ -123,7 +134,7 @@
             mCM.setProcessDefaultNetwork(null);
         }
 
-        for (Network network : mCtsNetUtils.getTestableNetworks()) {
+        for (Network network : getTestableNetworks()) {
             NetworkUtils.bindProcessToNetwork(0);
             assertNull(mCM.getBoundNetworkForProcess());
 
@@ -143,8 +154,8 @@
 
     @Test
     @AppModeFull(reason = "CHANGE_NETWORK_STATE permission can't be granted to instant apps")
-    public void testSetsocknetwork() throws ErrnoException {
-        for (Network network : mCtsNetUtils.getTestableNetworks()) {
+    public void testSetsocknetwork() throws Exception {
+        for (Network network : getTestableNetworks()) {
             int errno = runSetsocknetwork(network.getNetworkHandle());
             if (errno != 0) {
                 throw new ErrnoException(
@@ -154,8 +165,8 @@
     }
 
     @Test
-    public void testNativeDatagramTransmission() throws ErrnoException {
-        for (Network network : mCtsNetUtils.getTestableNetworks()) {
+    public void testNativeDatagramTransmission() throws Exception {
+        for (Network network : getTestableNetworks()) {
             int errno = runDatagramCheck(network.getNetworkHandle());
             if (errno != 0) {
                 throw new ErrnoException(
@@ -165,7 +176,7 @@
     }
 
     @Test
-    public void testNoSuchNetwork() {
+    public void testNoSuchNetwork() throws Exception {
         final Network eNoNet = new Network(54321);
         assertNull(mCM.getNetworkInfo(eNoNet));
 
@@ -178,9 +189,9 @@
     }
 
     @Test
-    public void testNetworkHandle() {
+    public void testNetworkHandle() throws Exception {
         // Test Network -> NetworkHandle -> Network results in the same Network.
-        for (Network network : mCtsNetUtils.getTestableNetworks()) {
+        for (Network network : getTestableNetworks()) {
             long networkHandle = network.getNetworkHandle();
             Network newNetwork = Network.fromNetworkHandle(networkHandle);
             assertEquals(newNetwork, network);
@@ -203,9 +214,7 @@
 
     @Test
     public void testResNApi() throws Exception {
-        final Network[] testNetworks = mCtsNetUtils.getTestableNetworks();
-
-        for (Network network : testNetworks) {
+        for (Network network : getTestableNetworks()) {
             // Throws AssertionError directly in jni function if test fail.
             runResNqueryCheck(network.getNetworkHandle());
             runResNsendCheck(network.getNetworkHandle());
@@ -241,7 +250,7 @@
         // b/144521720
         try {
             mCtsNetUtils.setPrivateDnsStrictMode(GOOGLE_PRIVATE_DNS_SERVER);
-            for (Network network : mCtsNetUtils.getTestableNetworks()) {
+            for (Network network : getTestableNetworks()) {
               // Wait for private DNS setting to propagate.
               mCtsNetUtils.awaitPrivateDnsSetting("NxDomain test wait private DNS setting timeout",
                         network, GOOGLE_PRIVATE_DNS_SERVER, true);
@@ -251,4 +260,44 @@
             mCtsNetUtils.restorePrivateDnsSetting();
         }
     }
+
+    /**
+     * Get all testable Networks with internet capability.
+     */
+    private Set<Network> getTestableNetworks() throws InterruptedException {
+        // Obtain cell and Wi-Fi through CtsNetUtils (which uses NetworkCallbacks), as they may have
+        // just been reconnected by the test using NetworkCallbacks, so synchronous calls may not
+        // yet return them (synchronous calls and callbacks should not be mixed for a given
+        // Network).
+        final Set<Network> testableNetworks = new ArraySet<>();
+        if (mContext.getPackageManager().hasSystemFeature(FEATURE_TELEPHONY)) {
+            if (mRequestedCellNetwork == null) {
+                mRequestedCellNetwork = mNetworkCallbackRule.requestCell();
+            }
+            assertNotNull("Cell network requested but not obtained", mRequestedCellNetwork);
+            testableNetworks.add(mRequestedCellNetwork);
+        }
+
+        if (mContext.getPackageManager().hasSystemFeature(FEATURE_WIFI)) {
+            testableNetworks.add(mCtsNetUtils.ensureWifiConnected());
+        }
+
+        // Obtain other networks through the synchronous API, if any.
+        for (Network network : mCtsNetUtils.getTestableNetworks()) {
+            final NetworkCapabilities nc = mCM.getNetworkCapabilities(network);
+            if (nc != null
+                    && !nc.hasTransport(TRANSPORT_WIFI)
+                    && !nc.hasTransport(TRANSPORT_CELLULAR)) {
+                testableNetworks.add(network);
+            }
+        }
+
+        // In practice this should not happen as getTestableNetworks throws if there is no network
+        // at all.
+        assertFalse("This device does not support WiFi nor cell data, and does not have any other "
+                        + "network connected. This test requires at least one internet-providing "
+                        + "network.",
+                testableNetworks.isEmpty());
+        return testableNetworks;
+    }
 }
diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index 61117df..6dd4857 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -2206,6 +2206,66 @@
     }
 
     @Test
+    fun testAdvertisingAndDiscovery_reregisterCustomHostWithDifferentAddresses_newAddressesFound() {
+        val si1 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.hostname = customHostname
+            it.hostAddresses = listOf(
+                    parseNumericAddress("192.0.2.23"),
+                    parseNumericAddress("2001:db8::1"))
+        }
+        val si2 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.serviceName = serviceName
+            it.serviceType = serviceType
+            it.hostname = customHostname
+            it.port = TEST_PORT
+        }
+        val si3 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.hostname = customHostname
+            it.hostAddresses = listOf(
+                    parseNumericAddress("192.0.2.24"),
+                    parseNumericAddress("2001:db8::2"))
+        }
+
+        val registrationRecord1 = NsdRegistrationRecord()
+        val registrationRecord2 = NsdRegistrationRecord()
+        val registrationRecord3 = NsdRegistrationRecord()
+
+        val discoveryRecord = NsdDiscoveryRecord()
+
+        tryTest {
+            registerService(registrationRecord1, si1)
+            registerService(registrationRecord2, si2)
+
+            nsdManager.unregisterService(registrationRecord1)
+            registrationRecord1.expectCallback<ServiceUnregistered>()
+
+            registerService(registrationRecord3, si3)
+
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network, Executor { it.run() }, discoveryRecord)
+            val discoveredInfo = discoveryRecord.waitForServiceDiscovered(
+                    serviceName, serviceType, testNetwork1.network)
+            val resolvedInfo = resolveService(discoveredInfo)
+
+            assertEquals(serviceName, discoveredInfo.serviceName)
+            assertEquals(TEST_PORT, resolvedInfo.port)
+            assertEquals(customHostname, resolvedInfo.hostname)
+            assertAddressEquals(
+                    listOf(parseNumericAddress("192.0.2.24"), parseNumericAddress("2001:db8::2")),
+                    resolvedInfo.hostAddresses)
+        } cleanupStep {
+            nsdManager.stopServiceDiscovery(discoveryRecord)
+            discoveryRecord.expectCallbackEventually<DiscoveryStopped>()
+        } cleanup {
+            nsdManager.unregisterService(registrationRecord2)
+            nsdManager.unregisterService(registrationRecord3)
+        }
+    }
+
+    @Test
     fun testServiceTypeClientRemovedAfterSocketDestroyed() {
         val si = makeTestServiceInfo(testNetwork1.network)
         // Register service on testNetwork1
diff --git a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
index 3d828a4..670889f 100644
--- a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
+++ b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
@@ -408,40 +408,10 @@
         return network;
     }
 
-    public Network connectToCell() throws InterruptedException {
-        if (cellConnectAttempted()) {
-            mCm.unregisterNetworkCallback(mCellNetworkCallback);
-        }
-        NetworkRequest cellRequest = new NetworkRequest.Builder()
-                .addTransportType(TRANSPORT_CELLULAR)
-                .addCapability(NET_CAPABILITY_INTERNET)
-                .build();
-        mCellNetworkCallback = new TestNetworkCallback();
-        mCm.requestNetwork(cellRequest, mCellNetworkCallback);
-        final Network cellNetwork = mCellNetworkCallback.waitForAvailable();
-        assertNotNull("Cell network not available. " +
-                "Please ensure the device has working mobile data.", cellNetwork);
-        return cellNetwork;
-    }
-
-    public void disconnectFromCell() {
-        if (!cellConnectAttempted()) {
-            throw new IllegalStateException("Cell connection not attempted");
-        }
-        mCm.unregisterNetworkCallback(mCellNetworkCallback);
-        mCellNetworkCallback = null;
-    }
-
     public boolean cellConnectAttempted() {
         return mCellNetworkCallback != null;
     }
 
-    public void tearDown() {
-        if (cellConnectAttempted()) {
-            disconnectFromCell();
-        }
-    }
-
     private NetworkRequest makeWifiNetworkRequest() {
         return new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
diff --git a/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt b/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt
index 765e56e..4780c5d 100644
--- a/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt
+++ b/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt
@@ -38,6 +38,7 @@
 import android.os.Build
 import android.os.Process
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.modules.utils.build.SdkLevel
 import com.android.server.net.integrationtests.NetworkStatsIntegrationTest.Direction.DOWNLOAD
 import com.android.server.net.integrationtests.NetworkStatsIntegrationTest.Direction.UPLOAD
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
@@ -214,6 +215,11 @@
         // In practice, for one way 10k download payload, the download usage is about
         // 11222~12880 bytes, with 14~17 packets. And the upload usage is about 1279~1626 bytes
         // with 14~17 packets, which is majorly contributed by TCP ACK packets.
+        // Clear TrafficStats cache is needed to avoid rate-limit caching for
+        // TrafficStats API results on V+ devices.
+        if (SdkLevel.isAtLeastV()) {
+            TrafficStats.clearRateLimitCaches()
+        }
         val snapshotAfterDownload = StatsSnapshot(context, internalInterfaceName)
         val (expectedDownloadLower, expectedDownloadUpper) = getExpectedStatsBounds(
             TEST_DOWNLOAD_SIZE,
@@ -236,6 +242,9 @@
         )
 
         // Verify upload data usage accounting.
+        if (SdkLevel.isAtLeastV()) {
+            TrafficStats.clearRateLimitCaches()
+        }
         val snapshotAfterUpload = StatsSnapshot(context, internalInterfaceName)
         val (expectedUploadLower, expectedUploadUpper) = getExpectedStatsBounds(
             TEST_UPLOAD_SIZE,
@@ -299,7 +308,8 @@
         val buf = ByteArray(DEFAULT_BUFFER_SIZE)
 
         httpServer.addResponse(
-            TestHttpServer.Request(path, NanoHTTPD.Method.POST), NanoHTTPD.Response.Status.OK,
+            TestHttpServer.Request(path, NanoHTTPD.Method.POST),
+            NanoHTTPD.Response.Status.OK,
             content = getRandomString(downloadSize)
         )
         var httpConnection: HttpURLConnection? = null
@@ -349,15 +359,19 @@
     ) {
         operator fun plus(other: BareStats): BareStats {
             return BareStats(
-                this.rxBytes + other.rxBytes, this.rxPackets + other.rxPackets,
-                this.txBytes + other.txBytes, this.txPackets + other.txPackets
+                this.rxBytes + other.rxBytes,
+                this.rxPackets + other.rxPackets,
+                this.txBytes + other.txBytes,
+                this.txPackets + other.txPackets
             )
         }
 
         operator fun minus(other: BareStats): BareStats {
             return BareStats(
-                this.rxBytes - other.rxBytes, this.rxPackets - other.rxPackets,
-                this.txBytes - other.txBytes, this.txPackets - other.txPackets
+                this.rxBytes - other.rxBytes,
+                this.rxPackets - other.rxPackets,
+                this.txBytes - other.txBytes,
+                this.txPackets - other.txPackets
             )
         }
 
@@ -405,8 +419,12 @@
         private fun getUidDetail(iface: String, tag: Int): BareStats {
             return getNetworkStatsThat(iface, tag) { nsm, template ->
                 nsm.queryDetailsForUidTagState(
-                    template, Long.MIN_VALUE, Long.MAX_VALUE,
-                    Process.myUid(), tag, Bucket.STATE_ALL
+                    template,
+                    Long.MIN_VALUE,
+                    Long.MAX_VALUE,
+                    Process.myUid(),
+                    tag,
+                    Bucket.STATE_ALL
                 )
             }
         }
@@ -498,28 +516,36 @@
         assertInRange(
             "Unexpected iface traffic stats",
             after.iface,
-            before.trafficStatsIface, after.trafficStatsIface,
-            lower, upper
+            before.trafficStatsIface,
+            after.trafficStatsIface,
+            lower,
+            upper
         )
         // Uid traffic stats are counted in both direction because the external network
         // traffic is also attributed to the test uid.
         assertInRange(
             "Unexpected uid traffic stats",
             after.iface,
-            before.trafficStatsUid, after.trafficStatsUid,
-            lower + lower.reverse(), upper + upper.reverse()
+            before.trafficStatsUid,
+            after.trafficStatsUid,
+            lower + lower.reverse(),
+            upper + upper.reverse()
         )
         assertInRange(
             "Unexpected non-tagged summary stats",
             after.iface,
-            before.statsSummary, after.statsSummary,
-            lower, upper
+            before.statsSummary,
+            after.statsSummary,
+            lower,
+            upper
         )
         assertInRange(
             "Unexpected non-tagged uid stats",
             after.iface,
-            before.statsUid, after.statsUid,
-            lower, upper
+            before.statsUid,
+            after.statsUid,
+            lower,
+            upper
         )
     }
 
@@ -546,14 +572,16 @@
         assertInRange(
             "Unexpected tagged summary stats",
             after.iface,
-            before.taggedSummary, after.taggedSummary,
+            before.taggedSummary,
+            after.taggedSummary,
             lower,
             upper
         )
         assertInRange(
             "Unexpected tagged uid stats: ${Process.myUid()}",
             after.iface,
-            before.taggedUid, after.taggedUid,
+            before.taggedUid,
+            after.taggedUid,
             lower,
             upper
         )
@@ -570,7 +598,8 @@
     ) {
         // Passing the value after operation and the value before operation to dump the actual
         // numbers if it fails.
-        assertTrue(checkInRange(before, after, lower, upper),
+        assertTrue(
+            checkInRange(before, after, lower, upper),
             "$tag on $iface: $after - $before is not within range [$lower, $upper]"
         )
     }
diff --git a/tests/integration/src/com/android/server/net/integrationtests/ServiceManagerWrapperIntegrationTest.kt b/tests/integration/src/com/android/server/net/integrationtests/ServiceManagerWrapperIntegrationTest.kt
new file mode 100644
index 0000000..7e00ed2
--- /dev/null
+++ b/tests/integration/src/com/android/server/net/integrationtests/ServiceManagerWrapperIntegrationTest.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.server.net.integrationtests
+
+import android.content.Context
+import android.os.Build
+import com.android.server.ServiceManagerWrapper
+import com.android.testutils.ConnectivityModuleTest
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import com.android.testutils.DevSdkIgnoreRunner
+import kotlin.test.assertNotNull
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Integration tests for {@link ServiceManagerWrapper}. */
+@RunWith(DevSdkIgnoreRunner::class)
+@IgnoreUpTo(Build.VERSION_CODES.S)
+@ConnectivityModuleTest
+class ServiceManagerWrapperIntegrationTest {
+    @Test
+    fun testWaitForService_successFullyRetrievesConnectivityServiceBinder() {
+        assertNotNull(ServiceManagerWrapper.waitForService(Context.CONNECTIVITY_SERVICE))
+    }
+}
diff --git a/tests/unit/java/android/net/nsd/NsdServiceInfoTest.kt b/tests/unit/java/android/net/nsd/NsdServiceInfoTest.kt
new file mode 100644
index 0000000..8f86f06
--- /dev/null
+++ b/tests/unit/java/android/net/nsd/NsdServiceInfoTest.kt
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.nsd
+
+import android.os.Build
+import androidx.test.filters.SmallTest
+import com.android.testutils.ConnectivityModuleTest
+import com.android.testutils.DevSdkIgnoreRule
+import com.android.testutils.DevSdkIgnoreRunner
+import kotlin.test.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Unit tests for {@link NsdServiceInfo}. */
+@SmallTest
+@ConnectivityModuleTest
+@RunWith(DevSdkIgnoreRunner::class)
+@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
+class NsdServiceInfoTest {
+    @Test
+    fun testToString_txtRecord() {
+        val info = NsdServiceInfo().apply {
+            this.setAttribute("abc", byteArrayOf(0xff.toByte(), 0xfe.toByte()))
+            this.setAttribute("def", null as String?)
+            this.setAttribute("ghi", "猫")
+            this.setAttribute("jkl", byteArrayOf(0, 0x21))
+            this.setAttribute("mno", "Hey Tom! It's you?.~{}")
+        }
+
+        val infoStr = info.toString()
+
+        assertTrue(
+            infoStr.contains("txtRecord: " +
+                "{abc=0xFFFE, def=(null), ghi=0xE78CAB, jkl=0x0021, mno=Hey Tom! It's you?.~{}}"),
+            infoStr)
+    }
+}
diff --git a/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
index 3f6e88d..aa28e5a 100644
--- a/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
+++ b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
@@ -231,8 +231,10 @@
         val clientId = 99
         val transactionId = 100
         val durationMs = 10L
+        val sentQueryCount = 10
         val metrics = NetworkNsdReportedMetrics(clientId, deps)
-        metrics.reportServiceResolutionStop(true /* isLegacy */, transactionId, durationMs)
+        metrics.reportServiceResolutionStop(
+                true /* isLegacy */, transactionId, durationMs, sentQueryCount)
 
         val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
         verify(deps).statsWrite(eventCaptor.capture())
@@ -243,6 +245,7 @@
             assertEquals(NsdEventType.NET_RESOLVE, it.type)
             assertEquals(MdnsQueryResult.MQR_SERVICE_RESOLUTION_STOP, it.queryResult)
             assertEquals(durationMs, it.eventDurationMillisec)
+            assertEquals(sentQueryCount, it.sentQueryCount)
         }
     }
 
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index f41d7b2..ce49533 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -172,6 +172,7 @@
 import static com.android.server.ConnectivityServiceTestUtils.transportToLegacyType;
 import static com.android.server.NetworkAgentWrapper.CallbackType.OnQosCallbackRegister;
 import static com.android.server.NetworkAgentWrapper.CallbackType.OnQosCallbackUnregister;
+import static com.android.server.connectivity.ConnectivityFlags.BACKGROUND_FIREWALL_CHAIN;
 import static com.android.server.connectivity.ConnectivityFlags.INGRESS_TO_VPN_ADDRESS_FILTERING;
 import static com.android.testutils.Cleanup.testAndCleanup;
 import static com.android.testutils.ConcurrentUtils.await;
@@ -391,6 +392,7 @@
 import com.android.internal.util.WakeupMessage;
 import com.android.internal.util.test.BroadcastInterceptingContext;
 import com.android.internal.util.test.FakeSettingsProvider;
+import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.ArrayTrackRecord;
 import com.android.net.module.util.BaseNetdUnsolicitedEventListener;
 import com.android.net.module.util.CollectionUtils;
@@ -2169,6 +2171,8 @@
                     return true;
                 case INGRESS_TO_VPN_ADDRESS_FILTERING:
                     return true;
+                case BACKGROUND_FIREWALL_CHAIN:
+                    return true;
                 default:
                     return super.isFeatureNotChickenedOut(context, name);
             }
@@ -10486,7 +10490,10 @@
         doTestSetUidFirewallRule(FIREWALL_CHAIN_POWERSAVE, FIREWALL_RULE_DENY);
         doTestSetUidFirewallRule(FIREWALL_CHAIN_RESTRICTED, FIREWALL_RULE_DENY);
         doTestSetUidFirewallRule(FIREWALL_CHAIN_LOW_POWER_STANDBY, FIREWALL_RULE_DENY);
-        doTestSetUidFirewallRule(FIREWALL_CHAIN_BACKGROUND, FIREWALL_RULE_DENY);
+        if (SdkLevel.isAtLeastV()) {
+            // FIREWALL_CHAIN_BACKGROUND is only available on V+.
+            doTestSetUidFirewallRule(FIREWALL_CHAIN_BACKGROUND, FIREWALL_RULE_DENY);
+        }
         doTestSetUidFirewallRule(FIREWALL_CHAIN_OEM_DENY_1, FIREWALL_RULE_ALLOW);
         doTestSetUidFirewallRule(FIREWALL_CHAIN_OEM_DENY_2, FIREWALL_RULE_ALLOW);
         doTestSetUidFirewallRule(FIREWALL_CHAIN_OEM_DENY_3, FIREWALL_RULE_ALLOW);
@@ -10494,16 +10501,19 @@
 
     @Test @IgnoreUpTo(SC_V2)
     public void testSetFirewallChainEnabled() throws Exception {
-        final List<Integer> firewallChains = Arrays.asList(
+        final List<Integer> firewallChains = new ArrayList<>(Arrays.asList(
                 FIREWALL_CHAIN_DOZABLE,
                 FIREWALL_CHAIN_STANDBY,
                 FIREWALL_CHAIN_POWERSAVE,
                 FIREWALL_CHAIN_RESTRICTED,
                 FIREWALL_CHAIN_LOW_POWER_STANDBY,
-                FIREWALL_CHAIN_BACKGROUND,
                 FIREWALL_CHAIN_OEM_DENY_1,
                 FIREWALL_CHAIN_OEM_DENY_2,
-                FIREWALL_CHAIN_OEM_DENY_3);
+                FIREWALL_CHAIN_OEM_DENY_3));
+        if (SdkLevel.isAtLeastV()) {
+            // FIREWALL_CHAIN_BACKGROUND is only available on V+.
+            firewallChains.add(FIREWALL_CHAIN_BACKGROUND);
+        }
         for (final int chain: firewallChains) {
             mCm.setFirewallChainEnabled(chain, true /* enabled */);
             verify(mBpfNetMaps).setChildChain(chain, true /* enable */);
@@ -10550,7 +10560,10 @@
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_POWERSAVE, allowlist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_RESTRICTED, allowlist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_LOW_POWER_STANDBY, allowlist);
-        doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_BACKGROUND, allowlist);
+        if (SdkLevel.isAtLeastV()) {
+            // FIREWALL_CHAIN_BACKGROUND is only available on V+.
+            doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_BACKGROUND, allowlist);
+        }
 
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_STANDBY, denylist);
         doTestSetFirewallChainEnabledCloseSocket(FIREWALL_CHAIN_OEM_DENY_1, denylist);
@@ -10572,7 +10585,10 @@
         doTestReplaceFirewallChain(FIREWALL_CHAIN_POWERSAVE);
         doTestReplaceFirewallChain(FIREWALL_CHAIN_RESTRICTED);
         doTestReplaceFirewallChain(FIREWALL_CHAIN_LOW_POWER_STANDBY);
-        doTestReplaceFirewallChain(FIREWALL_CHAIN_BACKGROUND);
+        if (SdkLevel.isAtLeastV()) {
+            // FIREWALL_CHAIN_BACKGROUND is only available on V+.
+            doTestReplaceFirewallChain(FIREWALL_CHAIN_BACKGROUND);
+        }
         doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_1);
         doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_2);
         doTestReplaceFirewallChain(FIREWALL_CHAIN_OEM_DENY_3);
@@ -17434,11 +17450,12 @@
         }
 
         mWiFiAgent.disconnect();
-        waitForIdle();
 
         if (expectUnavailable) {
+            testFactory.expectRequestRemove();
             testFactory.assertRequestCountEquals(0);
         } else {
+            testFactory.expectRequestAdd();
             testFactory.assertRequestCountEquals(1);
         }
 
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index 881de56..5731d01 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -521,6 +521,56 @@
     }
 
     @Test
+    @EnableCompatChanges(ENABLE_PLATFORM_MDNS_BACKEND)
+    public void testDiscoverOnTetheringDownstream_DiscoveryManager() throws Exception {
+        final NsdManager client = connectClient(mService);
+        final DiscoveryListener discListener = mock(DiscoveryListener.class);
+        client.discoverServices(SERVICE_TYPE, PROTOCOL, discListener);
+        waitForIdle();
+
+        final ArgumentCaptor<MdnsServiceBrowserListener> discoverListenerCaptor =
+                ArgumentCaptor.forClass(MdnsServiceBrowserListener.class);
+        final InOrder discManagerOrder = inOrder(mDiscoveryManager);
+        final String serviceTypeWithLocalDomain = SERVICE_TYPE + ".local";
+        discManagerOrder.verify(mDiscoveryManager).registerListener(eq(serviceTypeWithLocalDomain),
+                discoverListenerCaptor.capture(), any());
+
+        final int interfaceIdx = 123;
+        final MdnsServiceInfo mockServiceInfo = new MdnsServiceInfo(
+                SERVICE_NAME, /* serviceInstanceName */
+                serviceTypeWithLocalDomain.split("\\."), /* serviceType */
+                List.of(), /* subtypes */
+                new String[] {"android", "local"}, /* hostName */
+                12345, /* port */
+                List.of(IPV4_ADDRESS),
+                List.of(IPV6_ADDRESS),
+                List.of(), /* textStrings */
+                List.of(), /* textEntries */
+                interfaceIdx, /* interfaceIndex */
+                null /* network */,
+                Instant.MAX /* expirationTime */);
+
+        // Verify service is found with the interface index
+        discoverListenerCaptor.getValue().onServiceNameDiscovered(
+                mockServiceInfo, false /* isServiceFromCache */);
+        final ArgumentCaptor<NsdServiceInfo> foundInfoCaptor =
+                ArgumentCaptor.forClass(NsdServiceInfo.class);
+        verify(discListener, timeout(TIMEOUT_MS)).onServiceFound(foundInfoCaptor.capture());
+        final NsdServiceInfo foundInfo = foundInfoCaptor.getValue();
+        assertNull(foundInfo.getNetwork());
+        assertEquals(interfaceIdx, foundInfo.getInterfaceIndex());
+
+        // Using the returned service info to resolve or register callback uses the interface index
+        client.resolveService(foundInfo, mock(ResolveListener.class));
+        client.registerServiceInfoCallback(foundInfo, Runnable::run,
+                mock(ServiceInfoCallback.class));
+        waitForIdle();
+
+        discManagerOrder.verify(mDiscoveryManager, times(2)).registerListener(any(), any(), argThat(
+                o -> o.getNetwork() == null && o.getInterfaceIndex() == interfaceIdx));
+    }
+
+    @Test
     @DisableCompatChanges(ENABLE_PLATFORM_MDNS_BACKEND)
     @DevSdkIgnoreRule.IgnoreAfter(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
     public void testDiscoverOnBlackholeNetwork() throws Exception {
@@ -854,7 +904,7 @@
                 request.getServiceName().equals(ns.getServiceName())
                         && request.getServiceType().equals(ns.getServiceType())));
         verify(mMetrics).reportServiceResolutionStop(
-                true /* isLegacy */, resolveId, 10L /* durationMs */);
+                true /* isLegacy */, resolveId, 10L /* durationMs */, 0 /* sentQueryCount */);
     }
 
     @Test
@@ -928,7 +978,7 @@
                 request.getServiceName().equals(ns.getServiceName())
                         && request.getServiceType().equals(ns.getServiceType())));
         verify(mMetrics).reportServiceResolutionStop(
-                true /* isLegacy */, getAddrId, 10L /* durationMs */);
+                true /* isLegacy */, getAddrId, 10L /* durationMs */,  0 /* sentQueryCount */);
     }
 
     private void verifyUpdatedServiceInfo(NsdServiceInfo info, String serviceName,
@@ -1629,20 +1679,23 @@
         // Subtypes are not used for resolution, only for discovery
         assertEquals(Collections.emptyList(), optionsCaptor.getValue().getSubtypes());
 
+        final MdnsListener listener = listenerCaptor.getValue();
+        // Callbacks for query sent.
+        listener.onDiscoveryQuerySent(Collections.emptyList(), 1 /* transactionId */);
+
         doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
         client.stopServiceResolution(resolveListener);
         waitForIdle();
 
         // Verify the listener has been unregistered.
-        final MdnsListener listener = listenerCaptor.getValue();
         verify(mDiscoveryManager, timeout(TIMEOUT_MS))
                 .unregisterListener(eq(constructedServiceType), eq(listener));
         verify(resolveListener, timeout(TIMEOUT_MS)).onResolutionStopped(argThat(ns ->
                 request.getServiceName().equals(ns.getServiceName())
                         && request.getServiceType().equals(ns.getServiceType())));
         verify(mSocketProvider, timeout(CLEANUP_DELAY_MS + TIMEOUT_MS)).requestStopWhenInactive();
-        verify(mMetrics).reportServiceResolutionStop(
-                false /* isLegacy */, listener.mTransactionId, 10L /* durationMs */);
+        verify(mMetrics).reportServiceResolutionStop(false /* isLegacy */, listener.mTransactionId,
+                10L /* durationMs */, 1 /* sentQueryCount */);
     }
 
     @Test
diff --git a/tests/unit/java/com/android/server/connectivity/SatelliteAccessControllerTest.kt b/tests/unit/java/com/android/server/connectivity/SatelliteAccessControllerTest.kt
index 193078b..7885325 100644
--- a/tests/unit/java/com/android/server/connectivity/SatelliteAccessControllerTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/SatelliteAccessControllerTest.kt
@@ -18,17 +18,22 @@
 import android.Manifest
 import android.app.role.OnRoleHoldersChangedListener
 import android.app.role.RoleManager
+import android.content.BroadcastReceiver
 import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
 import android.content.pm.ApplicationInfo
 import android.content.pm.PackageManager
-import android.content.pm.UserInfo
 import android.os.Build
 import android.os.Handler
+import android.os.Looper
 import android.os.UserHandle
+import android.os.UserManager
 import android.util.ArraySet
-import com.android.server.makeMockUserManager
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
 import com.android.testutils.DevSdkIgnoreRunner
+import java.util.concurrent.Executor
+import java.util.function.Consumer
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -36,30 +41,32 @@
 import org.mockito.ArgumentMatchers.any
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.eq
+import org.mockito.ArgumentMatchers.isNull
 import org.mockito.Mockito.doReturn
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.never
+import org.mockito.Mockito.timeout
+import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
-import java.util.concurrent.Executor
-import java.util.function.Consumer
 
-private const val USER = 0
-val USER_INFO = UserInfo(USER, "" /* name */, UserInfo.FLAG_PRIMARY)
-val USER_HANDLE = UserHandle(USER)
 private const val PRIMARY_USER = 0
 private const val SECONDARY_USER = 10
 private val PRIMARY_USER_HANDLE = UserHandle.of(PRIMARY_USER)
 private val SECONDARY_USER_HANDLE = UserHandle.of(SECONDARY_USER)
+
 // sms app names
 private const val SMS_APP1 = "sms_app_1"
 private const val SMS_APP2 = "sms_app_2"
+
 // sms app ids
 private const val SMS_APP_ID1 = 100
 private const val SMS_APP_ID2 = 101
+
 // UID for app1 and app2 on primary user
 // These app could become default sms app for user1
 private val PRIMARY_USER_SMS_APP_UID1 = UserHandle.getUid(PRIMARY_USER, SMS_APP_ID1)
 private val PRIMARY_USER_SMS_APP_UID2 = UserHandle.getUid(PRIMARY_USER, SMS_APP_ID2)
+
 // UID for app1 and app2 on secondary user
 // These app could become default sms app for user2
 private val SECONDARY_USER_SMS_APP_UID1 = UserHandle.getUid(SECONDARY_USER, SMS_APP_ID1)
@@ -69,154 +76,259 @@
 @IgnoreUpTo(Build.VERSION_CODES.TIRAMISU)
 class SatelliteAccessControllerTest {
     private val context = mock(Context::class.java)
-    private val mPackageManager = mock(PackageManager::class.java)
-    private val mHandler = mock(Handler::class.java)
-    private val mRoleManager =
-        mock(SatelliteAccessController.Dependencies::class.java)
+    private val primaryUserContext = mock(Context::class.java)
+    private val secondaryUserContext = mock(Context::class.java)
+    private val mPackageManagerPrimaryUser = mock(PackageManager::class.java)
+    private val mPackageManagerSecondaryUser = mock(PackageManager::class.java)
+    private val mDeps = mock(SatelliteAccessController.Dependencies::class.java)
     private val mCallback = mock(Consumer::class.java) as Consumer<Set<Int>>
-    private val mSatelliteAccessController =
-        SatelliteAccessController(context, mRoleManager, mCallback, mHandler)
+    private val userManager = mock(UserManager::class.java)
+    private val mHandler = Handler(Looper.getMainLooper())
+    private var mSatelliteAccessController =
+        SatelliteAccessController(context, mDeps, mCallback, mHandler)
     private lateinit var mRoleHolderChangedListener: OnRoleHoldersChangedListener
+    private lateinit var mUserRemovedReceiver: BroadcastReceiver
+
+    private fun <T> mockService(name: String, clazz: Class<T>, service: T) {
+        doReturn(name).`when`(context).getSystemServiceName(clazz)
+        doReturn(service).`when`(context).getSystemService(name)
+        if (context.getSystemService(clazz) == null) {
+            // Test is using mockito-extended
+            doReturn(service).`when`(context).getSystemService(clazz)
+        }
+    }
+
     @Before
     @Throws(PackageManager.NameNotFoundException::class)
     fun setup() {
-        makeMockUserManager(USER_INFO, USER_HANDLE)
-        doReturn(context).`when`(context).createContextAsUser(any(), anyInt())
-        doReturn(mPackageManager).`when`(context).packageManager
+        doReturn(emptyList<UserHandle>()).`when`(userManager).getUserHandles(true)
+        mockService(Context.USER_SERVICE, UserManager::class.java, userManager)
 
-        doReturn(PackageManager.PERMISSION_GRANTED)
-            .`when`(mPackageManager)
-            .checkPermission(Manifest.permission.SATELLITE_COMMUNICATION, SMS_APP1)
-        doReturn(PackageManager.PERMISSION_GRANTED)
-            .`when`(mPackageManager)
-            .checkPermission(Manifest.permission.SATELLITE_COMMUNICATION, SMS_APP2)
+        doReturn(primaryUserContext).`when`(context).createContextAsUser(PRIMARY_USER_HANDLE, 0)
+        doReturn(mPackageManagerPrimaryUser).`when`(primaryUserContext).packageManager
 
-        // Initialise default message application primary user package1
+        doReturn(secondaryUserContext).`when`(context).createContextAsUser(SECONDARY_USER_HANDLE, 0)
+        doReturn(mPackageManagerSecondaryUser).`when`(secondaryUserContext).packageManager
+
+        for (app in listOf(SMS_APP1, SMS_APP2)) {
+            doReturn(PackageManager.PERMISSION_GRANTED)
+                .`when`(mPackageManagerPrimaryUser)
+                .checkPermission(Manifest.permission.SATELLITE_COMMUNICATION, app)
+            doReturn(PackageManager.PERMISSION_GRANTED)
+                .`when`(mPackageManagerSecondaryUser)
+                .checkPermission(Manifest.permission.SATELLITE_COMMUNICATION, app)
+        }
+
+        // Initialise message application primary user package1
         val applicationInfo1 = ApplicationInfo()
         applicationInfo1.uid = PRIMARY_USER_SMS_APP_UID1
         doReturn(applicationInfo1)
-            .`when`(mPackageManager)
+            .`when`(mPackageManagerPrimaryUser)
             .getApplicationInfo(eq(SMS_APP1), anyInt())
 
-        // Initialise default message application primary user package2
+        // Initialise message application primary user package2
         val applicationInfo2 = ApplicationInfo()
         applicationInfo2.uid = PRIMARY_USER_SMS_APP_UID2
         doReturn(applicationInfo2)
-            .`when`(mPackageManager)
+            .`when`(mPackageManagerPrimaryUser)
             .getApplicationInfo(eq(SMS_APP2), anyInt())
 
-        // Get registered listener using captor
-        val listenerCaptor = ArgumentCaptor.forClass(
-            OnRoleHoldersChangedListener::class.java
-        )
-        mSatelliteAccessController.start()
-        verify(mRoleManager).addOnRoleHoldersChangedListenerAsUser(
-            any(Executor::class.java), listenerCaptor.capture(), any(UserHandle::class.java))
-        mRoleHolderChangedListener = listenerCaptor.value
+        // Initialise message application secondary user package1
+        val applicationInfo3 = ApplicationInfo()
+        applicationInfo3.uid = SECONDARY_USER_SMS_APP_UID1
+        doReturn(applicationInfo3)
+            .`when`(mPackageManagerSecondaryUser)
+            .getApplicationInfo(eq(SMS_APP1), anyInt())
+
+        // Initialise message application secondary user package2
+        val applicationInfo4 = ApplicationInfo()
+        applicationInfo4.uid = SECONDARY_USER_SMS_APP_UID2
+        doReturn(applicationInfo4)
+            .`when`(mPackageManagerSecondaryUser)
+            .getApplicationInfo(eq(SMS_APP2), anyInt())
     }
 
     @Test
     fun test_onRoleHoldersChanged_SatelliteFallbackUid_Changed_SingleUser() {
-        doReturn(listOf<String>()).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
+        startSatelliteAccessController()
+        doReturn(listOf<String>()).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback, never()).accept(any())
 
         // check DEFAULT_MESSAGING_APP1 is available as satellite network fallback uid
         doReturn(listOf(SMS_APP1))
-            .`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
+            .`when`(mDeps).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback).accept(setOf(PRIMARY_USER_SMS_APP_UID1))
 
         // check SMS_APP2 is available as satellite network Fallback uid
-        doReturn(listOf(SMS_APP2)).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
+        doReturn(listOf(SMS_APP2)).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback).accept(setOf(PRIMARY_USER_SMS_APP_UID2))
 
         // check no uid is available as satellite network fallback uid
-        doReturn(listOf<String>()).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
+        doReturn(listOf<String>()).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback).accept(ArraySet())
     }
 
     @Test
     fun test_onRoleHoldersChanged_NoSatelliteCommunicationPermission() {
-        doReturn(listOf<Any>()).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
+        startSatelliteAccessController()
+        doReturn(listOf<Any>()).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback, never()).accept(any())
 
         // check DEFAULT_MESSAGING_APP1 is not available as satellite network fallback uid
         // since satellite communication permission not available.
         doReturn(PackageManager.PERMISSION_DENIED)
-            .`when`(mPackageManager)
+            .`when`(mPackageManagerPrimaryUser)
             .checkPermission(Manifest.permission.SATELLITE_COMMUNICATION, SMS_APP1)
         doReturn(listOf(SMS_APP1))
-            .`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
+            .`when`(mDeps).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback, never()).accept(any())
     }
 
     @Test
     fun test_onRoleHoldersChanged_RoleSms_NotAvailable() {
+        startSatelliteAccessController()
         doReturn(listOf(SMS_APP1))
-            .`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
-        mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_BROWSER,
-            PRIMARY_USER_HANDLE)
+            .`when`(mDeps).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
+        mRoleHolderChangedListener.onRoleHoldersChanged(
+            RoleManager.ROLE_BROWSER,
+            PRIMARY_USER_HANDLE
+        )
         verify(mCallback, never()).accept(any())
     }
 
     @Test
     fun test_onRoleHoldersChanged_SatelliteNetworkFallbackUid_Changed_multiUser() {
-        doReturn(listOf<String>()).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
+        startSatelliteAccessController()
+        doReturn(listOf<String>()).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback, never()).accept(any())
 
         // check SMS_APP1 is available as satellite network fallback uid at primary user
         doReturn(listOf(SMS_APP1))
-            .`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
+            .`when`(mDeps).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback).accept(setOf(PRIMARY_USER_SMS_APP_UID1))
 
         // check SMS_APP2 is available as satellite network fallback uid at primary user
-        doReturn(listOf(SMS_APP2)).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
+        doReturn(listOf(SMS_APP2)).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
         verify(mCallback).accept(setOf(PRIMARY_USER_SMS_APP_UID2))
 
         // check SMS_APP1 is available as satellite network fallback uid at secondary user
-        val applicationInfo1 = ApplicationInfo()
-        applicationInfo1.uid = SECONDARY_USER_SMS_APP_UID1
-        doReturn(applicationInfo1).`when`(mPackageManager)
-            .getApplicationInfo(eq(SMS_APP1), anyInt())
-        doReturn(listOf(SMS_APP1)).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            SECONDARY_USER_HANDLE)
+        doReturn(listOf(SMS_APP1)).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            SECONDARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, SECONDARY_USER_HANDLE)
         verify(mCallback).accept(setOf(PRIMARY_USER_SMS_APP_UID2, SECONDARY_USER_SMS_APP_UID1))
 
         // check no uid is available as satellite network fallback uid at primary user
-        doReturn(listOf<String>()).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
-        mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS,
-            PRIMARY_USER_HANDLE)
+        doReturn(listOf<String>()).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
+        mRoleHolderChangedListener.onRoleHoldersChanged(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
         verify(mCallback).accept(setOf(SECONDARY_USER_SMS_APP_UID1))
 
         // check SMS_APP2 is available as satellite network fallback uid at secondary user
-        applicationInfo1.uid = SECONDARY_USER_SMS_APP_UID2
-        doReturn(applicationInfo1).`when`(mPackageManager)
-            .getApplicationInfo(eq(SMS_APP2), anyInt())
         doReturn(listOf(SMS_APP2))
-            .`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS, SECONDARY_USER_HANDLE)
+            .`when`(mDeps).getRoleHoldersAsUser(RoleManager.ROLE_SMS, SECONDARY_USER_HANDLE)
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, SECONDARY_USER_HANDLE)
         verify(mCallback).accept(setOf(SECONDARY_USER_SMS_APP_UID2))
 
         // check no uid is available as satellite network fallback uid at secondary user
-        doReturn(listOf<String>()).`when`(mRoleManager).getRoleHoldersAsUser(RoleManager.ROLE_SMS,
-            SECONDARY_USER_HANDLE)
+        doReturn(listOf<String>()).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            SECONDARY_USER_HANDLE
+        )
         mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, SECONDARY_USER_HANDLE)
         verify(mCallback).accept(ArraySet())
     }
+
+    @Test
+    fun test_SatelliteFallbackUidCallback_OnUserRemoval() {
+        startSatelliteAccessController()
+        // check SMS_APP2 is available as satellite network fallback uid at primary user
+        doReturn(listOf(SMS_APP2)).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            PRIMARY_USER_HANDLE
+        )
+        mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
+        verify(mCallback).accept(setOf(PRIMARY_USER_SMS_APP_UID2))
+
+        // check SMS_APP1 is available as satellite network fallback uid at secondary user
+        doReturn(listOf(SMS_APP1)).`when`(mDeps).getRoleHoldersAsUser(
+            RoleManager.ROLE_SMS,
+            SECONDARY_USER_HANDLE
+        )
+        mRoleHolderChangedListener.onRoleHoldersChanged(RoleManager.ROLE_SMS, SECONDARY_USER_HANDLE)
+        verify(mCallback).accept(setOf(PRIMARY_USER_SMS_APP_UID2, SECONDARY_USER_SMS_APP_UID1))
+
+        val userRemovalIntent = Intent(Intent.ACTION_USER_REMOVED)
+        userRemovalIntent.putExtra(Intent.EXTRA_USER, SECONDARY_USER_HANDLE)
+        mUserRemovedReceiver.onReceive(context, userRemovalIntent)
+        verify(mCallback, times(2)).accept(setOf(PRIMARY_USER_SMS_APP_UID2))
+    }
+
+    @Test
+    fun testOnStartUpCallbackSatelliteFallbackUidWithExistingUsers() {
+        doReturn(
+            listOf(PRIMARY_USER_HANDLE)
+        ).`when`(userManager).getUserHandles(true)
+        doReturn(listOf(SMS_APP1))
+            .`when`(mDeps).getRoleHoldersAsUser(RoleManager.ROLE_SMS, PRIMARY_USER_HANDLE)
+        // At start up, SatelliteAccessController must call CS callback with existing users'
+        // default messaging apps uids.
+        startSatelliteAccessController()
+        verify(mCallback, timeout(500)).accept(setOf(PRIMARY_USER_SMS_APP_UID1))
+    }
+
+    private fun startSatelliteAccessController() {
+        mSatelliteAccessController.start()
+        // Get registered listener using captor
+        val listenerCaptor = ArgumentCaptor.forClass(OnRoleHoldersChangedListener::class.java)
+        verify(mDeps).addOnRoleHoldersChangedListenerAsUser(
+            any(Executor::class.java),
+            listenerCaptor.capture(),
+            any(UserHandle::class.java)
+        )
+        mRoleHolderChangedListener = listenerCaptor.value
+
+        // Get registered receiver using captor
+        val userRemovedReceiverCaptor = ArgumentCaptor.forClass(BroadcastReceiver::class.java)
+        verify(context).registerReceiver(
+            userRemovedReceiverCaptor.capture(),
+            any(IntentFilter::class.java),
+            isNull(),
+            any(Handler::class.java)
+        )
+         mUserRemovedReceiver = userRemovedReceiverCaptor.value
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
index 5251e2a..b5c0132 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -18,6 +18,8 @@
 
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
@@ -65,8 +67,9 @@
     private static final String SERVICE_TYPE_2 = "_test._tcp.local";
     private static final Network NETWORK_1 = Mockito.mock(Network.class);
     private static final Network NETWORK_2 = Mockito.mock(Network.class);
+    private static final int INTERFACE_INDEX_NULL_NETWORK = 123;
     private static final SocketKey SOCKET_KEY_NULL_NETWORK =
-            new SocketKey(null /* network */, 999 /* interfaceIndex */);
+            new SocketKey(null /* network */, INTERFACE_INDEX_NULL_NETWORK);
     private static final SocketKey SOCKET_KEY_NETWORK_1 =
             new SocketKey(NETWORK_1, 998 /* interfaceIndex */);
     private static final SocketKey SOCKET_KEY_NETWORK_2 =
@@ -97,6 +100,8 @@
     private HandlerThread thread;
     private Handler handler;
 
+    private int createdServiceTypeClientCount;
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -106,11 +111,13 @@
         handler = new Handler(thread.getLooper());
         doReturn(thread.getLooper()).when(socketClient).getLooper();
         doReturn(true).when(socketClient).supportsRequestingSpecificNetworks();
+        createdServiceTypeClientCount = 0;
         discoveryManager = new MdnsDiscoveryManager(executorProvider, socketClient,
                 sharedLog, MdnsFeatureFlags.newBuilder().build()) {
                     @Override
                     MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType,
                             @NonNull SocketKey socketKey) {
+                        createdServiceTypeClientCount++;
                         final Pair<String, SocketKey> perSocketServiceType =
                                 Pair.create(serviceType, socketKey);
                         if (perSocketServiceType.equals(PER_SOCKET_SERVICE_TYPE_1_NULL_NETWORK)) {
@@ -128,6 +135,7 @@
                                 PER_SOCKET_SERVICE_TYPE_2_NETWORK_2)) {
                             return mockServiceTypeClientType2Network2;
                         }
+                        fail("Unexpected perSocketServiceType: " + perSocketServiceType);
                         return null;
                     }
                 };
@@ -324,7 +332,6 @@
 
         // Receive a response, it should be processed on the client.
         final MdnsPacket response = createMdnsPacket(SERVICE_TYPE_1);
-        final int ifIndex = 1;
         runOnHandler(() -> discoveryManager.onResponseReceived(response, SOCKET_KEY_NULL_NETWORK));
         verify(mockServiceTypeClientType1NullNetwork).processResponse(
                 response, SOCKET_KEY_NULL_NETWORK);
@@ -350,6 +357,39 @@
         verify(socketClient, never()).stopDiscovery();
     }
 
+    @Test
+    public void testInterfaceIndexRequested_OnlyUsesSelectedInterface() throws IOException {
+        final MdnsSearchOptions searchOptions =
+                MdnsSearchOptions.newBuilder()
+                        .setNetwork(null /* network */)
+                        .setInterfaceIndex(INTERFACE_INDEX_NULL_NETWORK)
+                        .build();
+
+        final SocketCreationCallback callback = expectSocketCreationCallback(
+                SERVICE_TYPE_1, mockListenerOne, searchOptions);
+        final SocketKey unusedIfaceKey = new SocketKey(null, INTERFACE_INDEX_NULL_NETWORK + 1);
+        final SocketKey matchingIfaceWithNetworkKey =
+                new SocketKey(Mockito.mock(Network.class), INTERFACE_INDEX_NULL_NETWORK);
+        runOnHandler(() -> {
+            callback.onSocketCreated(unusedIfaceKey);
+            callback.onSocketCreated(matchingIfaceWithNetworkKey);
+            callback.onSocketCreated(SOCKET_KEY_NULL_NETWORK);
+            callback.onSocketCreated(SOCKET_KEY_NETWORK_1);
+        });
+        // Only the client for INTERFACE_INDEX_NULL_NETWORK is created
+        verify(mockServiceTypeClientType1NullNetwork).startSendAndReceive(
+                mockListenerOne, searchOptions);
+        assertEquals(1, createdServiceTypeClientCount);
+
+        runOnHandler(() -> {
+            callback.onSocketDestroyed(SOCKET_KEY_NETWORK_1);
+            callback.onSocketDestroyed(SOCKET_KEY_NULL_NETWORK);
+            callback.onSocketDestroyed(matchingIfaceWithNetworkKey);
+            callback.onSocketDestroyed(unusedIfaceKey);
+        });
+        verify(mockServiceTypeClientType1NullNetwork).notifySocketDestroyed();
+    }
+
     private MdnsPacket createMdnsPacket(String serviceType) {
         final String[] type = TextUtils.split(serviceType, "\\.");
         final ArrayList<String> name = new ArrayList<>(type.length + 1);
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiserTest.kt b/tests/unit/java/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiserTest.kt
index 69fec85..629ac67 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiserTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiserTest.kt
@@ -18,7 +18,6 @@
 
 import android.net.InetAddresses.parseNumericAddress
 import android.net.LinkAddress
-import android.net.nsd.NsdManager
 import android.net.nsd.NsdServiceInfo
 import android.os.Build
 import android.os.HandlerThread
@@ -55,6 +54,7 @@
 import org.mockito.Mockito.never
 import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
+import org.mockito.Mockito.inOrder
 
 private const val LOG_TAG = "testlogtag"
 private const val TIMEOUT_MS = 10_000L
@@ -65,6 +65,7 @@
 
 private const val TEST_SERVICE_ID_1 = 42
 private const val TEST_SERVICE_ID_DUPLICATE = 43
+private const val TEST_SERVICE_ID_2 = 44
 private val TEST_SERVICE_1 = NsdServiceInfo().apply {
     serviceType = "_testservice._tcp"
     serviceName = "MyTestService"
@@ -78,6 +79,13 @@
     port = 12345
 }
 
+private val TEST_SERVICE_1_CUSTOM_HOST = NsdServiceInfo().apply {
+    serviceType = "_testservice._tcp"
+    serviceName = "MyTestService"
+    hostname = "MyTestHost"
+    port = 12345
+}
+
 @RunWith(DevSdkIgnoreRunner::class)
 @IgnoreUpTo(Build.VERSION_CODES.S_V2)
 class MdnsInterfaceAdvertiserTest {
@@ -183,6 +191,93 @@
     }
 
     @Test
+    fun testAddRemoveServiceWithCustomHost_restartProbingForProbingServices() {
+        val customHost1 = NsdServiceInfo().apply {
+            hostname = "MyTestHost"
+            hostAddresses = listOf(
+                    parseNumericAddress("192.0.2.23"),
+                    parseNumericAddress("2001:db8::1"))
+        }
+        addServiceAndFinishProbing(TEST_SERVICE_ID_1, customHost1)
+        addServiceAndFinishProbing(TEST_SERVICE_ID_2, TEST_SERVICE_1_CUSTOM_HOST)
+        repository.setServiceProbing(TEST_SERVICE_ID_2)
+        val probingInfo = mock(ProbingInfo::class.java)
+        doReturn("MyTestHost")
+                .`when`(repository).getHostnameForServiceId(TEST_SERVICE_ID_1)
+        doReturn(TEST_SERVICE_ID_2).`when`(probingInfo).serviceId
+        doReturn(listOf(probingInfo))
+                .`when`(repository).restartProbingForHostname("MyTestHost")
+        val inOrder = inOrder(prober, announcer)
+
+        // Remove the custom host: the custom host's announcement is stopped and the probing
+        // services which use that hostname are re-announced.
+        advertiser.removeService(TEST_SERVICE_ID_1)
+
+        inOrder.verify(prober).stop(TEST_SERVICE_ID_1)
+        inOrder.verify(announcer).stop(TEST_SERVICE_ID_1)
+        inOrder.verify(prober).stop(TEST_SERVICE_ID_2)
+        inOrder.verify(prober).startProbing(probingInfo)
+    }
+
+    @Test
+    fun testAddRemoveServiceWithCustomHost_restartAnnouncingForProbedServices() {
+        val customHost1 = NsdServiceInfo().apply {
+            hostname = "MyTestHost"
+            hostAddresses = listOf(
+                    parseNumericAddress("192.0.2.23"),
+                    parseNumericAddress("2001:db8::1"))
+        }
+        addServiceAndFinishProbing(TEST_SERVICE_ID_1, customHost1)
+        val announcementInfo =
+                addServiceAndFinishProbing(TEST_SERVICE_ID_2, TEST_SERVICE_1_CUSTOM_HOST)
+        doReturn("MyTestHost")
+                .`when`(repository).getHostnameForServiceId(TEST_SERVICE_ID_1)
+        doReturn(listOf(announcementInfo))
+                .`when`(repository).restartAnnouncingForHostname("MyTestHost")
+        val inOrder = inOrder(prober, announcer)
+
+        // Remove the custom host: the custom host's announcement is stopped and the probed services
+        // which use that hostname are re-announced.
+        advertiser.removeService(TEST_SERVICE_ID_1)
+
+        inOrder.verify(prober).stop(TEST_SERVICE_ID_1)
+        inOrder.verify(announcer).stop(TEST_SERVICE_ID_1)
+        inOrder.verify(announcer).stop(TEST_SERVICE_ID_2)
+        inOrder.verify(announcer).startSending(TEST_SERVICE_ID_2, announcementInfo, 0L /* initialDelayMs */)
+    }
+
+    @Test
+    fun testAddMoreAddressesForCustomHost_restartAnnouncingForProbedServices() {
+        val customHost = NsdServiceInfo().apply {
+            hostname = "MyTestHost"
+            hostAddresses = listOf(
+                parseNumericAddress("192.0.2.23"),
+                parseNumericAddress("2001:db8::1"))
+        }
+        doReturn("MyTestHost")
+            .`when`(repository).getHostnameForServiceId(TEST_SERVICE_ID_1)
+        doReturn("MyTestHost")
+            .`when`(repository).getHostnameForServiceId(TEST_SERVICE_ID_2)
+        val announcementInfo1 =
+            addServiceAndFinishProbing(TEST_SERVICE_ID_1, TEST_SERVICE_1_CUSTOM_HOST)
+
+        val probingInfo2 = addServiceAndStartProbing(TEST_SERVICE_ID_2, customHost)
+        val announcementInfo2 = AnnouncementInfo(TEST_SERVICE_ID_2, emptyList(), emptyList())
+        doReturn(announcementInfo2).`when`(repository).onProbingSucceeded(probingInfo2)
+        doReturn(listOf(announcementInfo1, announcementInfo2))
+            .`when`(repository).restartAnnouncingForHostname("MyTestHost")
+        probeCb.onFinished(probingInfo2)
+
+        val inOrder = inOrder(prober, announcer)
+
+        inOrder.verify(announcer)
+            .startSending(TEST_SERVICE_ID_2, announcementInfo2, 0L /* initialDelayMs */)
+        inOrder.verify(announcer).stop(TEST_SERVICE_ID_1)
+        inOrder.verify(announcer)
+            .startSending(TEST_SERVICE_ID_1, announcementInfo1, 0L /* initialDelayMs */)
+    }
+
+    @Test
     fun testDoubleRemove() {
         addServiceAndFinishProbing(TEST_SERVICE_ID_1, TEST_SERVICE_1)
 
@@ -422,8 +517,8 @@
         verify(prober, never()).startProbing(any())
     }
 
-    private fun addServiceAndFinishProbing(serviceId: Int, serviceInfo: NsdServiceInfo):
-            AnnouncementInfo {
+    private fun addServiceAndStartProbing(serviceId: Int, serviceInfo: NsdServiceInfo):
+            ProbingInfo {
         val testProbingInfo = mock(ProbingInfo::class.java)
         doReturn(serviceId).`when`(testProbingInfo).serviceId
         doReturn(testProbingInfo).`when`(repository).setServiceProbing(serviceId)
@@ -432,8 +527,15 @@
         verify(repository).addService(serviceId, serviceInfo, null /* ttl */)
         verify(prober).startProbing(testProbingInfo)
 
+        return testProbingInfo
+    }
+
+    private fun addServiceAndFinishProbing(serviceId: Int, serviceInfo: NsdServiceInfo):
+            AnnouncementInfo {
+        val testProbingInfo = addServiceAndStartProbing(serviceId, serviceInfo)
+
         // Simulate probing success: continues to announcing
-        val testAnnouncementInfo = mock(AnnouncementInfo::class.java)
+        val testAnnouncementInfo = AnnouncementInfo(serviceId, emptyList(), emptyList())
         doReturn(testAnnouncementInfo).`when`(repository).onProbingSucceeded(testProbingInfo)
         probeCb.onFinished(testProbingInfo)
         return testAnnouncementInfo
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
index 9474464..fb3d183 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
@@ -23,6 +23,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.timeout;
@@ -47,6 +48,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
+import org.mockito.InOrder;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -55,6 +57,7 @@
 import java.net.DatagramPacket;
 import java.net.NetworkInterface;
 import java.net.SocketException;
+import java.util.ArrayList;
 import java.util.List;
 
 @RunWith(DevSdkIgnoreRunner.class)
@@ -154,7 +157,7 @@
         verify(mSocketCreationCallback).onSocketCreated(tetherSocketKey2);
 
         // Send packet to IPv4 with mSocketKey and verify sending has been called.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv4Packet), mSocketKey,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket).send(ipv4Packet);
@@ -162,7 +165,7 @@
         verify(tetherIfaceSock2, never()).send(any());
 
         // Send packet to IPv4 with onlyUseIpv6OnIpv6OnlyNetworks = true, the packet will be sent.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv4Packet), mSocketKey,
                 true /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, times(2)).send(ipv4Packet);
@@ -170,7 +173,7 @@
         verify(tetherIfaceSock2, never()).send(any());
 
         // Send packet to IPv6 with tetherSocketKey1 and verify sending has been called.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv6Packet, tetherSocketKey1,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv6Packet), tetherSocketKey1,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, never()).send(ipv6Packet);
@@ -180,7 +183,7 @@
         // Send packet to IPv6 with onlyUseIpv6OnIpv6OnlyNetworks = true, the packet will not be
         // sent. Therefore, the tetherIfaceSock1.send() and tetherIfaceSock2.send() are still be
         // called once.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv6Packet, tetherSocketKey1,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv6Packet), tetherSocketKey1,
                 true /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, never()).send(ipv6Packet);
@@ -266,7 +269,7 @@
         verify(mSocketCreationCallback).onSocketCreated(socketKey3);
 
         // Send IPv4 packet on the mSocketKey and verify sending has been called.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv4Packet), mSocketKey,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket).send(ipv4Packet);
@@ -295,7 +298,7 @@
         verify(socketCreationCb2).onSocketCreated(socketKey3);
 
         // Send IPv4 packet on socket2 and verify sending to the socket2 only.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, socketKey2,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv4Packet), socketKey2,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         // ipv4Packet still sent only once on mSocket: times(1) matches the packet sent earlier on
@@ -309,7 +312,7 @@
         verify(mProvider, timeout(DEFAULT_TIMEOUT)).unrequestSocket(callback2);
 
         // Send IPv4 packet again and verify it's still sent a second time
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, socketKey2,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv4Packet), socketKey2,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(socket2, times(2)).send(ipv4Packet);
@@ -320,7 +323,7 @@
         verify(mProvider, timeout(DEFAULT_TIMEOUT)).unrequestSocket(callback);
 
         // Send IPv4 packet and verify no more sending.
-        mSocketClient.sendPacketRequestingMulticastResponse(ipv4Packet, mSocketKey,
+        mSocketClient.sendPacketRequestingMulticastResponse(List.of(ipv4Packet), mSocketKey,
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mSocket, times(1)).send(ipv4Packet);
@@ -407,4 +410,31 @@
         verify(creationCallback3).onSocketDestroyed(mSocketKey);
         verify(creationCallback3, never()).onSocketDestroyed(socketKey2);
     }
+
+    @Test
+    public void testSendPacketWithMultipleDatagramPacket() throws IOException {
+        final SocketCallback callback = expectSocketCallback();
+        final List<DatagramPacket> packets = new ArrayList<>();
+        for (int i = 0; i < 10; i++) {
+            packets.add(new DatagramPacket(new byte[10 + i] /* buff */, 0 /* offset */,
+                    10 + i /* length */, MdnsConstants.IPV4_SOCKET_ADDR));
+        }
+        doReturn(true).when(mSocket).hasJoinedIpv4();
+        doReturn(true).when(mSocket).hasJoinedIpv6();
+        doReturn(createEmptyNetworkInterface()).when(mSocket).getInterface();
+
+        // Notify socket created
+        callback.onSocketCreated(mSocketKey, mSocket, List.of());
+        verify(mSocketCreationCallback).onSocketCreated(mSocketKey);
+
+        // Send packets to IPv4 with mSocketKey then verify sending has been called and the
+        // sequence is correct.
+        mSocketClient.sendPacketRequestingMulticastResponse(packets, mSocketKey,
+                false /* onlyUseIpv6OnIpv6OnlyNetworks */);
+        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        InOrder inOrder = inOrder(mSocket);
+        for (int i = 0; i < 10; i++) {
+            inOrder.verify(mSocket).send(packets.get(i));
+        }
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsRecordRepositoryTest.kt b/tests/unit/java/com/android/server/connectivity/mdns/MdnsRecordRepositoryTest.kt
index c69b1e1..271cc65 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsRecordRepositoryTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsRecordRepositoryTest.kt
@@ -24,6 +24,7 @@
 import com.android.server.connectivity.mdns.MdnsAnnouncer.AnnouncementInfo
 import com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_HOST
 import com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_SERVICE
+import com.android.server.connectivity.mdns.MdnsProber.ProbingInfo
 import com.android.server.connectivity.mdns.MdnsRecord.TYPE_A
 import com.android.server.connectivity.mdns.MdnsRecord.TYPE_AAAA
 import com.android.server.connectivity.mdns.MdnsRecord.TYPE_PTR
@@ -51,6 +52,10 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.eq
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.verify
 
 private const val TEST_SERVICE_ID_1 = 42
 private const val TEST_SERVICE_ID_2 = 43
@@ -112,6 +117,14 @@
     port = TEST_PORT
 }
 
+private val TEST_SERVICE_CUSTOM_HOST_NO_ADDRESSES = NsdServiceInfo().apply {
+    hostname = "TestHost"
+    hostAddresses = listOf()
+    serviceType = "_testservice._tcp"
+    serviceName = "TestService"
+    port = TEST_PORT
+}
+
 @RunWith(DevSdkIgnoreRunner::class)
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
 class MdnsRecordRepositoryTest {
@@ -1676,6 +1689,127 @@
         assertEquals(0, reply.additionalAnswers.size)
         assertEquals(knownAnswers, reply.knownAnswers)
     }
+
+    @Test
+    fun testRestartProbingForHostname() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.initWithService(TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1,
+                setOf(TEST_SUBTYPE, TEST_SUBTYPE2))
+        repository.addService(TEST_SERVICE_CUSTOM_HOST_ID_1,
+                TEST_SERVICE_CUSTOM_HOST_NO_ADDRESSES, null)
+        repository.setServiceProbing(TEST_SERVICE_CUSTOM_HOST_ID_1)
+        repository.removeService(TEST_CUSTOM_HOST_ID_1)
+
+        val probingInfos = repository.restartProbingForHostname("TestHost")
+
+        assertEquals(1, probingInfos.size)
+        val probingInfo = probingInfos.get(0)
+        assertEquals(TEST_SERVICE_CUSTOM_HOST_ID_1, probingInfo.serviceId)
+        val packet = probingInfo.getPacket(0)
+        assertEquals(0, packet.transactionId)
+        assertEquals(MdnsConstants.FLAGS_QUERY, packet.flags)
+        assertEquals(0, packet.answers.size)
+        assertEquals(0, packet.additionalRecords.size)
+        assertEquals(1, packet.questions.size)
+        val serviceName = arrayOf("TestService", "_testservice", "_tcp", "local")
+        assertEquals(MdnsAnyRecord(serviceName, false /* unicast */), packet.questions[0])
+        assertThat(packet.authorityRecords).containsExactly(
+                MdnsServiceRecord(
+                        serviceName,
+                        0L /* receiptTimeMillis */,
+                        false /* cacheFlush */,
+                        SHORT_TTL /* ttlMillis */,
+                        0 /* servicePriority */,
+                        0 /* serviceWeight */,
+                        TEST_PORT,
+                        TEST_CUSTOM_HOST_1_NAME))
+    }
+
+    @Test
+    fun testRestartAnnouncingForHostname() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.initWithService(TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1,
+                setOf(TEST_SUBTYPE, TEST_SUBTYPE2))
+        repository.addServiceAndFinishProbing(TEST_SERVICE_CUSTOM_HOST_ID_1,
+                TEST_SERVICE_CUSTOM_HOST_NO_ADDRESSES)
+        repository.removeService(TEST_CUSTOM_HOST_ID_1)
+
+        val announcementInfos = repository.restartAnnouncingForHostname("TestHost")
+
+        assertEquals(1, announcementInfos.size)
+        val announcementInfo = announcementInfos.get(0)
+        assertEquals(TEST_SERVICE_CUSTOM_HOST_ID_1, announcementInfo.serviceId)
+        val packet = announcementInfo.getPacket(0)
+        assertEquals(0, packet.transactionId)
+        assertEquals(0x8400 /* response, authoritative */, packet.flags)
+        assertEquals(0, packet.questions.size)
+        assertEquals(0, packet.authorityRecords.size)
+        val serviceName = arrayOf("TestService", "_testservice", "_tcp", "local")
+        val serviceType = arrayOf("_testservice", "_tcp", "local")
+        val v4AddrRev = getReverseDnsAddress(TEST_ADDRESSES[0].address)
+        val v6Addr1Rev = getReverseDnsAddress(TEST_ADDRESSES[1].address)
+        val v6Addr2Rev = getReverseDnsAddress(TEST_ADDRESSES[2].address)
+        assertThat(packet.answers).containsExactly(
+                MdnsPointerRecord(
+                        serviceType,
+                        0L /* receiptTimeMillis */,
+                        // Not a unique name owned by the announcer, so cacheFlush=false
+                        false /* cacheFlush */,
+                        4500000L /* ttlMillis */,
+                        serviceName),
+                MdnsServiceRecord(
+                        serviceName,
+                        0L /* receiptTimeMillis */,
+                        true /* cacheFlush */,
+                        120000L /* ttlMillis */,
+                        0 /* servicePriority */,
+                        0 /* serviceWeight */,
+                        TEST_PORT /* servicePort */,
+                        TEST_CUSTOM_HOST_1_NAME),
+                MdnsTextRecord(
+                        serviceName,
+                        0L /* receiptTimeMillis */,
+                        true /* cacheFlush */,
+                        4500000L /* ttlMillis */,
+                        emptyList() /* entries */),
+                MdnsPointerRecord(
+                        arrayOf("_services", "_dns-sd", "_udp", "local"),
+                        0L /* receiptTimeMillis */,
+                        false /* cacheFlush */,
+                        4500000L /* ttlMillis */,
+                        serviceType))
+        assertThat(packet.additionalRecords).containsExactly(
+                MdnsNsecRecord(v4AddrRev,
+                        0L /* receiptTimeMillis */,
+                        true /* cacheFlush */,
+                        120000L /* ttlMillis */,
+                        v4AddrRev,
+                        intArrayOf(TYPE_PTR)),
+                MdnsNsecRecord(TEST_HOSTNAME,
+                        0L /* receiptTimeMillis */,
+                        true /* cacheFlush */,
+                        120000L /* ttlMillis */,
+                        TEST_HOSTNAME,
+                        intArrayOf(TYPE_A, TYPE_AAAA)),
+                MdnsNsecRecord(v6Addr1Rev,
+                        0L /* receiptTimeMillis */,
+                        true /* cacheFlush */,
+                        120000L /* ttlMillis */,
+                        v6Addr1Rev,
+                        intArrayOf(TYPE_PTR)),
+                MdnsNsecRecord(v6Addr2Rev,
+                        0L /* receiptTimeMillis */,
+                        true /* cacheFlush */,
+                        120000L /* ttlMillis */,
+                        v6Addr2Rev,
+                        intArrayOf(TYPE_PTR)),
+                MdnsNsecRecord(serviceName,
+                        0L /* receiptTimeMillis */,
+                        true /* cacheFlush */,
+                        4500000L /* ttlMillis */,
+                        serviceName,
+                        intArrayOf(TYPE_TXT, TYPE_SRV)))
+    }
 }
 
 private fun MdnsRecordRepository.initWithService(
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
index 2eb9440..44fa55c 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -162,59 +162,59 @@
             expectedIPv6Packets[i] = new DatagramPacket(buf, 0 /* offset */, 5 /* length */,
                     MdnsConstants.getMdnsIPv6Address(), MdnsConstants.MDNS_PORT);
         }
-        when(mockDeps.getDatagramPacketFromMdnsPacket(
-                any(), any(MdnsPacket.class), eq(IPV4_ADDRESS)))
-                .thenReturn(expectedIPv4Packets[0])
-                .thenReturn(expectedIPv4Packets[1])
-                .thenReturn(expectedIPv4Packets[2])
-                .thenReturn(expectedIPv4Packets[3])
-                .thenReturn(expectedIPv4Packets[4])
-                .thenReturn(expectedIPv4Packets[5])
-                .thenReturn(expectedIPv4Packets[6])
-                .thenReturn(expectedIPv4Packets[7])
-                .thenReturn(expectedIPv4Packets[8])
-                .thenReturn(expectedIPv4Packets[9])
-                .thenReturn(expectedIPv4Packets[10])
-                .thenReturn(expectedIPv4Packets[11])
-                .thenReturn(expectedIPv4Packets[12])
-                .thenReturn(expectedIPv4Packets[13])
-                .thenReturn(expectedIPv4Packets[14])
-                .thenReturn(expectedIPv4Packets[15])
-                .thenReturn(expectedIPv4Packets[16])
-                .thenReturn(expectedIPv4Packets[17])
-                .thenReturn(expectedIPv4Packets[18])
-                .thenReturn(expectedIPv4Packets[19])
-                .thenReturn(expectedIPv4Packets[20])
-                .thenReturn(expectedIPv4Packets[21])
-                .thenReturn(expectedIPv4Packets[22])
-                .thenReturn(expectedIPv4Packets[23]);
+        when(mockDeps.getDatagramPacketsFromMdnsPacket(
+                any(), any(MdnsPacket.class), eq(IPV4_ADDRESS), anyBoolean()))
+                .thenReturn(List.of(expectedIPv4Packets[0]))
+                .thenReturn(List.of(expectedIPv4Packets[1]))
+                .thenReturn(List.of(expectedIPv4Packets[2]))
+                .thenReturn(List.of(expectedIPv4Packets[3]))
+                .thenReturn(List.of(expectedIPv4Packets[4]))
+                .thenReturn(List.of(expectedIPv4Packets[5]))
+                .thenReturn(List.of(expectedIPv4Packets[6]))
+                .thenReturn(List.of(expectedIPv4Packets[7]))
+                .thenReturn(List.of(expectedIPv4Packets[8]))
+                .thenReturn(List.of(expectedIPv4Packets[9]))
+                .thenReturn(List.of(expectedIPv4Packets[10]))
+                .thenReturn(List.of(expectedIPv4Packets[11]))
+                .thenReturn(List.of(expectedIPv4Packets[12]))
+                .thenReturn(List.of(expectedIPv4Packets[13]))
+                .thenReturn(List.of(expectedIPv4Packets[14]))
+                .thenReturn(List.of(expectedIPv4Packets[15]))
+                .thenReturn(List.of(expectedIPv4Packets[16]))
+                .thenReturn(List.of(expectedIPv4Packets[17]))
+                .thenReturn(List.of(expectedIPv4Packets[18]))
+                .thenReturn(List.of(expectedIPv4Packets[19]))
+                .thenReturn(List.of(expectedIPv4Packets[20]))
+                .thenReturn(List.of(expectedIPv4Packets[21]))
+                .thenReturn(List.of(expectedIPv4Packets[22]))
+                .thenReturn(List.of(expectedIPv4Packets[23]));
 
-        when(mockDeps.getDatagramPacketFromMdnsPacket(
-                any(), any(MdnsPacket.class), eq(IPV6_ADDRESS)))
-                .thenReturn(expectedIPv6Packets[0])
-                .thenReturn(expectedIPv6Packets[1])
-                .thenReturn(expectedIPv6Packets[2])
-                .thenReturn(expectedIPv6Packets[3])
-                .thenReturn(expectedIPv6Packets[4])
-                .thenReturn(expectedIPv6Packets[5])
-                .thenReturn(expectedIPv6Packets[6])
-                .thenReturn(expectedIPv6Packets[7])
-                .thenReturn(expectedIPv6Packets[8])
-                .thenReturn(expectedIPv6Packets[9])
-                .thenReturn(expectedIPv6Packets[10])
-                .thenReturn(expectedIPv6Packets[11])
-                .thenReturn(expectedIPv6Packets[12])
-                .thenReturn(expectedIPv6Packets[13])
-                .thenReturn(expectedIPv6Packets[14])
-                .thenReturn(expectedIPv6Packets[15])
-                .thenReturn(expectedIPv6Packets[16])
-                .thenReturn(expectedIPv6Packets[17])
-                .thenReturn(expectedIPv6Packets[18])
-                .thenReturn(expectedIPv6Packets[19])
-                .thenReturn(expectedIPv6Packets[20])
-                .thenReturn(expectedIPv6Packets[21])
-                .thenReturn(expectedIPv6Packets[22])
-                .thenReturn(expectedIPv6Packets[23]);
+        when(mockDeps.getDatagramPacketsFromMdnsPacket(
+                any(), any(MdnsPacket.class), eq(IPV6_ADDRESS), anyBoolean()))
+                .thenReturn(List.of(expectedIPv6Packets[0]))
+                .thenReturn(List.of(expectedIPv6Packets[1]))
+                .thenReturn(List.of(expectedIPv6Packets[2]))
+                .thenReturn(List.of(expectedIPv6Packets[3]))
+                .thenReturn(List.of(expectedIPv6Packets[4]))
+                .thenReturn(List.of(expectedIPv6Packets[5]))
+                .thenReturn(List.of(expectedIPv6Packets[6]))
+                .thenReturn(List.of(expectedIPv6Packets[7]))
+                .thenReturn(List.of(expectedIPv6Packets[8]))
+                .thenReturn(List.of(expectedIPv6Packets[9]))
+                .thenReturn(List.of(expectedIPv6Packets[10]))
+                .thenReturn(List.of(expectedIPv6Packets[11]))
+                .thenReturn(List.of(expectedIPv6Packets[12]))
+                .thenReturn(List.of(expectedIPv6Packets[13]))
+                .thenReturn(List.of(expectedIPv6Packets[14]))
+                .thenReturn(List.of(expectedIPv6Packets[15]))
+                .thenReturn(List.of(expectedIPv6Packets[16]))
+                .thenReturn(List.of(expectedIPv6Packets[17]))
+                .thenReturn(List.of(expectedIPv6Packets[18]))
+                .thenReturn(List.of(expectedIPv6Packets[19]))
+                .thenReturn(List.of(expectedIPv6Packets[20]))
+                .thenReturn(List.of(expectedIPv6Packets[21]))
+                .thenReturn(List.of(expectedIPv6Packets[22]))
+                .thenReturn(List.of(expectedIPv6Packets[23]));
 
         thread = new HandlerThread("MdnsServiceTypeClientTests");
         thread.start();
@@ -694,23 +694,23 @@
                 .addSubtype("subtype1").build();
         final MdnsSearchOptions searchOptions2 = MdnsSearchOptions.newBuilder()
                 .addSubtype("subtype2").build();
-        doCallRealMethod().when(mockDeps).getDatagramPacketFromMdnsPacket(
-                any(), any(MdnsPacket.class), any(InetSocketAddress.class));
+        doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
+                any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
         startSendAndReceive(mockListenerOne, searchOptions1);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
 
         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient, mockDeps);
 
         // Verify the query asks for subtype1
-        final ArgumentCaptor<DatagramPacket> subtype1QueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> subtype1QueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
                 subtype1QueryCaptor.capture(),
                 eq(socketKey), eq(false));
 
         final MdnsPacket subtype1Query = MdnsPacket.parse(
-                new MdnsPacketReader(subtype1QueryCaptor.getValue()));
+                new MdnsPacketReader(subtype1QueryCaptor.getValue().get(0)));
 
         assertEquals(2, subtype1Query.questions.size());
         assertTrue(hasQuestion(subtype1Query, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
@@ -722,8 +722,8 @@
         inOrder.verify(mockDeps).removeMessages(any(), eq(EVENT_START_QUERYTASK));
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
 
-        final ArgumentCaptor<DatagramPacket> combinedSubtypesQueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> combinedSubtypesQueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
                 combinedSubtypesQueryCaptor.capture(),
                 eq(socketKey), eq(false));
@@ -731,7 +731,7 @@
         inOrder.verify(mockDeps).sendMessageDelayed(any(), any(), anyLong());
 
         final MdnsPacket combinedSubtypesQuery = MdnsPacket.parse(
-                new MdnsPacketReader(combinedSubtypesQueryCaptor.getValue()));
+                new MdnsPacketReader(combinedSubtypesQueryCaptor.getValue().get(0)));
 
         assertEquals(3, combinedSubtypesQuery.questions.size());
         assertTrue(hasQuestion(combinedSubtypesQuery, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
@@ -747,15 +747,15 @@
         dispatchMessage();
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
 
-        final ArgumentCaptor<DatagramPacket> subtype2QueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> subtype2QueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 subtype2QueryCaptor.capture(),
                 eq(socketKey), eq(false));
 
         final MdnsPacket subtype2Query = MdnsPacket.parse(
-                new MdnsPacketReader(subtype2QueryCaptor.getValue()));
+                new MdnsPacketReader(subtype2QueryCaptor.getValue().get(0)));
 
         assertEquals(2, subtype2Query.questions.size());
         assertTrue(hasQuestion(subtype2Query, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
@@ -1201,8 +1201,8 @@
         final MdnsSearchOptions resolveOptions2 = MdnsSearchOptions.newBuilder()
                 .setResolveInstanceName(instanceName).build();
 
-        doCallRealMethod().when(mockDeps).getDatagramPacketFromMdnsPacket(
-                any(), any(MdnsPacket.class), any(InetSocketAddress.class));
+        doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
+                any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
 
         startSendAndReceive(mockListenerOne, resolveOptions1);
         startSendAndReceive(mockListenerTwo, resolveOptions2);
@@ -1210,8 +1210,8 @@
         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
 
         // Verify a query for SRV/TXT was sent, but no PTR query
-        final ArgumentCaptor<DatagramPacket> srvTxtQueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> srvTxtQueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
@@ -1223,7 +1223,7 @@
         verify(mockListenerTwo).onDiscoveryQuerySent(any(), anyInt());
 
         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
-                new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
+                new MdnsPacketReader(srvTxtQueryCaptor.getValue().get(0)));
 
         final String[] serviceName = getTestServiceName(instanceName);
         assertEquals(1, srvTxtQueryPacket.questions.size());
@@ -1255,8 +1255,8 @@
 
         // Expect a query for A/AAAA
         dispatchMessage();
-        final ArgumentCaptor<DatagramPacket> addressQueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> addressQueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 addressQueryCaptor.capture(),
@@ -1266,7 +1266,7 @@
         verify(mockListenerTwo, times(2)).onDiscoveryQuerySent(any(), anyInt());
 
         final MdnsPacket addressQueryPacket = MdnsPacket.parse(
-                new MdnsPacketReader(addressQueryCaptor.getValue()));
+                new MdnsPacketReader(addressQueryCaptor.getValue().get(0)));
         assertEquals(2, addressQueryPacket.questions.size());
         assertTrue(hasQuestion(addressQueryPacket, MdnsRecord.TYPE_A, hostname));
         assertTrue(hasQuestion(addressQueryPacket, MdnsRecord.TYPE_AAAA, hostname));
@@ -1316,15 +1316,15 @@
         final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
                 .setResolveInstanceName(instanceName).build();
 
-        doCallRealMethod().when(mockDeps).getDatagramPacketFromMdnsPacket(
-                any(), any(MdnsPacket.class), any(InetSocketAddress.class));
+        doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
+                any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
 
         startSendAndReceive(mockListenerOne, resolveOptions);
         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
 
         // Get the query for SRV/TXT
-        final ArgumentCaptor<DatagramPacket> srvTxtQueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> srvTxtQueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
@@ -1334,7 +1334,7 @@
         assertNotNull(delayMessage);
 
         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
-                new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
+                new MdnsPacketReader(srvTxtQueryCaptor.getValue().get(0)));
 
         final String[] serviceName = getTestServiceName(instanceName);
         assertTrue(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_ANY, serviceName));
@@ -1378,8 +1378,8 @@
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
 
         // Expect a renewal query
-        final ArgumentCaptor<DatagramPacket> renewalQueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> renewalQueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         // Second and later sends are sent as "expect multicast response" queries
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 renewalQueryCaptor.capture(),
@@ -1388,7 +1388,7 @@
         assertNotNull(delayMessage);
         inOrder.verify(mockListenerOne).onDiscoveryQuerySent(any(), anyInt());
         final MdnsPacket renewalPacket = MdnsPacket.parse(
-                new MdnsPacketReader(renewalQueryCaptor.getValue()));
+                new MdnsPacketReader(renewalQueryCaptor.getValue().get(0)));
         assertTrue(hasQuestion(renewalPacket, MdnsRecord.TYPE_ANY, serviceName));
         inOrder.verifyNoMoreInteractions();
 
@@ -1937,14 +1937,14 @@
                 serviceCache,
                 MdnsFeatureFlags.newBuilder().setIsQueryWithKnownAnswerEnabled(true).build());
 
-        doCallRealMethod().when(mockDeps).getDatagramPacketFromMdnsPacket(
-                any(), any(MdnsPacket.class), any(InetSocketAddress.class));
+        doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
+                any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
 
         startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
 
-        final ArgumentCaptor<DatagramPacket> queryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> queryCaptor =
+                ArgumentCaptor.forClass(List.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
@@ -1953,7 +1953,7 @@
         assertNotNull(delayMessage);
 
         final MdnsPacket queryPacket = MdnsPacket.parse(
-                new MdnsPacketReader(queryCaptor.getValue()));
+                new MdnsPacketReader(queryCaptor.getValue().get(0)));
         assertTrue(hasQuestion(queryPacket, MdnsRecord.TYPE_PTR));
 
         // Process a response
@@ -1981,14 +1981,14 @@
 
         // Expect a query with known answers
         dispatchMessage();
-        final ArgumentCaptor<DatagramPacket> knownAnswersQueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> knownAnswersQueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 knownAnswersQueryCaptor.capture(), eq(socketKey), eq(false));
 
         final MdnsPacket knownAnswersQueryPacket = MdnsPacket.parse(
-                new MdnsPacketReader(knownAnswersQueryCaptor.getValue()));
+                new MdnsPacketReader(knownAnswersQueryCaptor.getValue().get(0)));
         assertTrue(hasQuestion(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
         assertTrue(hasAnswer(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
         assertFalse(hasAnswer(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, subtypeLabels));
@@ -2001,16 +2001,16 @@
                 serviceCache,
                 MdnsFeatureFlags.newBuilder().setIsQueryWithKnownAnswerEnabled(true).build());
 
-        doCallRealMethod().when(mockDeps).getDatagramPacketFromMdnsPacket(
-                any(), any(MdnsPacket.class), any(InetSocketAddress.class));
+        doCallRealMethod().when(mockDeps).getDatagramPacketsFromMdnsPacket(
+                any(), any(MdnsPacket.class), any(InetSocketAddress.class), anyBoolean());
 
         final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
                 .addSubtype("subtype").build();
         startSendAndReceive(mockListenerOne, options);
         InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
 
-        final ArgumentCaptor<DatagramPacket> queryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> queryCaptor =
+                ArgumentCaptor.forClass(List.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         // Send twice for IPv4 and IPv6
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingUnicastResponse(
@@ -2019,7 +2019,7 @@
         assertNotNull(delayMessage);
 
         final MdnsPacket queryPacket = MdnsPacket.parse(
-                new MdnsPacketReader(queryCaptor.getValue()));
+                new MdnsPacketReader(queryCaptor.getValue().get(0)));
         final String[] subtypeLabels = Stream.concat(Stream.of("_subtype", "_sub"),
                 Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
         assertTrue(hasQuestion(queryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
@@ -2048,14 +2048,14 @@
 
         // Expect a query with known answers
         dispatchMessage();
-        final ArgumentCaptor<DatagramPacket> knownAnswersQueryCaptor =
-                ArgumentCaptor.forClass(DatagramPacket.class);
+        final ArgumentCaptor<List<DatagramPacket>> knownAnswersQueryCaptor =
+                ArgumentCaptor.forClass(List.class);
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         inOrder.verify(mockSocketClient, times(2)).sendPacketRequestingMulticastResponse(
                 knownAnswersQueryCaptor.capture(), eq(socketKey), eq(false));
 
         final MdnsPacket knownAnswersQueryPacket = MdnsPacket.parse(
-                new MdnsPacketReader(knownAnswersQueryCaptor.getValue()));
+                new MdnsPacketReader(knownAnswersQueryCaptor.getValue().get(0)));
         assertTrue(hasQuestion(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
         assertTrue(hasQuestion(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, subtypeLabels));
         assertTrue(hasAnswer(knownAnswersQueryPacket, MdnsRecord.TYPE_PTR, SERVICE_TYPE_LABELS));
@@ -2083,17 +2083,21 @@
         currentThreadExecutor.getAndClearLastScheduledRunnable().run();
         if (expectsUnicastResponse) {
             verify(mockSocketClient).sendPacketRequestingUnicastResponse(
-                    expectedIPv4Packets[index], socketKey, false);
+                    argThat(pkts -> pkts.get(0).equals(expectedIPv4Packets[index])),
+                    eq(socketKey), eq(false));
             if (multipleSocketDiscovery) {
                 verify(mockSocketClient).sendPacketRequestingUnicastResponse(
-                        expectedIPv6Packets[index], socketKey, false);
+                        argThat(pkts -> pkts.get(0).equals(expectedIPv6Packets[index])),
+                        eq(socketKey), eq(false));
             }
         } else {
             verify(mockSocketClient).sendPacketRequestingMulticastResponse(
-                    expectedIPv4Packets[index], socketKey, false);
+                    argThat(pkts -> pkts.get(0).equals(expectedIPv4Packets[index])),
+                    eq(socketKey), eq(false));
             if (multipleSocketDiscovery) {
                 verify(mockSocketClient).sendPacketRequestingMulticastResponse(
-                        expectedIPv6Packets[index], socketKey, false);
+                        argThat(pkts -> pkts.get(0).equals(expectedIPv6Packets[index])),
+                        eq(socketKey), eq(false));
             }
         }
         verify(mockDeps, times(index + 1))
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
index 7ced1cb..1989ed3 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
@@ -27,6 +27,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
@@ -53,6 +54,7 @@
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
+import org.mockito.InOrder;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.mockito.invocation.InvocationOnMock;
@@ -60,6 +62,8 @@
 import java.io.IOException;
 import java.net.DatagramPacket;
 import java.net.InetSocketAddress;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -234,7 +238,7 @@
 
         // Sends a packet.
         DatagramPacket packet = getTestDatagramPacket();
-        mdnsClient.sendPacketRequestingMulticastResponse(packet,
+        mdnsClient.sendPacketRequestingMulticastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         // mockMulticastSocket.send() will be called on another thread. If we verify it immediately,
         // it may not be called yet. So timeout is added.
@@ -242,7 +246,7 @@
         verify(mockUnicastSocket, timeout(TIMEOUT).times(0)).send(packet);
 
         // Verify the packet is sent by the unicast socket.
-        mdnsClient.sendPacketRequestingUnicastResponse(packet,
+        mdnsClient.sendPacketRequestingUnicastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         verify(mockMulticastSocket, timeout(TIMEOUT).times(1)).send(packet);
         verify(mockUnicastSocket, timeout(TIMEOUT).times(1)).send(packet);
@@ -287,7 +291,7 @@
 
         // Sends a packet.
         DatagramPacket packet = getTestDatagramPacket();
-        mdnsClient.sendPacketRequestingMulticastResponse(packet,
+        mdnsClient.sendPacketRequestingMulticastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         // mockMulticastSocket.send() will be called on another thread. If we verify it immediately,
         // it may not be called yet. So timeout is added.
@@ -295,7 +299,7 @@
         verify(mockUnicastSocket, timeout(TIMEOUT).times(0)).send(packet);
 
         // Verify the packet is sent by the multicast socket as well.
-        mdnsClient.sendPacketRequestingUnicastResponse(packet,
+        mdnsClient.sendPacketRequestingUnicastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         verify(mockMulticastSocket, timeout(TIMEOUT).times(2)).send(packet);
         verify(mockUnicastSocket, timeout(TIMEOUT).times(0)).send(packet);
@@ -354,7 +358,7 @@
     public void testStopDiscovery_queueIsCleared() throws IOException {
         mdnsClient.startDiscovery();
         mdnsClient.stopDiscovery();
-        mdnsClient.sendPacketRequestingMulticastResponse(getTestDatagramPacket(),
+        mdnsClient.sendPacketRequestingMulticastResponse(List.of(getTestDatagramPacket()),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
 
         synchronized (mdnsClient.multicastPacketQueue) {
@@ -366,7 +370,7 @@
     public void testSendPacket_afterDiscoveryStops() throws IOException {
         mdnsClient.startDiscovery();
         mdnsClient.stopDiscovery();
-        mdnsClient.sendPacketRequestingMulticastResponse(getTestDatagramPacket(),
+        mdnsClient.sendPacketRequestingMulticastResponse(List.of(getTestDatagramPacket()),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
 
         synchronized (mdnsClient.multicastPacketQueue) {
@@ -380,7 +384,7 @@
         //MdnsConfigsFlagsImpl.mdnsPacketQueueMaxSize.override(2L);
         mdnsClient.startDiscovery();
         for (int i = 0; i < 100; i++) {
-            mdnsClient.sendPacketRequestingMulticastResponse(getTestDatagramPacket(),
+            mdnsClient.sendPacketRequestingMulticastResponse(List.of(getTestDatagramPacket()),
                     false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         }
 
@@ -478,9 +482,9 @@
 
         mdnsClient.startDiscovery();
         DatagramPacket packet = getTestDatagramPacket();
-        mdnsClient.sendPacketRequestingUnicastResponse(packet,
+        mdnsClient.sendPacketRequestingUnicastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
-        mdnsClient.sendPacketRequestingMulticastResponse(packet,
+        mdnsClient.sendPacketRequestingMulticastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
 
         // Wait for the timer to be triggered.
@@ -511,9 +515,9 @@
         assertFalse(mdnsClient.receivedUnicastResponse);
         assertFalse(mdnsClient.cannotReceiveMulticastResponse.get());
 
-        mdnsClient.sendPacketRequestingUnicastResponse(packet,
+        mdnsClient.sendPacketRequestingUnicastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
-        mdnsClient.sendPacketRequestingMulticastResponse(packet,
+        mdnsClient.sendPacketRequestingMulticastResponse(List.of(packet),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */);
         Thread.sleep(MdnsConfigs.checkMulticastResponseIntervalMs() * 2);
 
@@ -570,6 +574,26 @@
                 .onResponseReceived(any(), argThat(key -> key.getInterfaceIndex() == -1));
     }
 
+    @Test
+    public void testSendPacketWithMultipleDatagramPacket() throws IOException {
+        mdnsClient.startDiscovery();
+        final List<DatagramPacket> packets = new ArrayList<>();
+        for (int i = 0; i < 10; i++) {
+            packets.add(new DatagramPacket(new byte[10 + i] /* buff */, 0 /* offset */,
+                    10 + i /* length */, MdnsConstants.IPV4_SOCKET_ADDR));
+        }
+
+        // Sends packets.
+        mdnsClient.sendPacketRequestingMulticastResponse(packets,
+                false /* onlyUseIpv6OnIpv6OnlyNetworks */);
+        InOrder inOrder = inOrder(mockMulticastSocket);
+        for (int i = 0; i < 10; i++) {
+            // mockMulticastSocket.send() will be called on another thread. If we verify it
+            // immediately, it may not be called yet. So timeout is added.
+            inOrder.verify(mockMulticastSocket, timeout(TIMEOUT)).send(packets.get(i));
+        }
+    }
+
     private DatagramPacket getTestDatagramPacket() {
         return new DatagramPacket(buf, 0, 5,
                 new InetSocketAddress(MdnsConstants.getMdnsIPv4Address(), 5353 /* port */));
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/util/MdnsUtilsTest.kt b/tests/unit/java/com/android/server/connectivity/mdns/util/MdnsUtilsTest.kt
index f705bcb..009205e 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/util/MdnsUtilsTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/util/MdnsUtilsTest.kt
@@ -17,6 +17,13 @@
 package com.android.server.connectivity.mdns.util
 
 import android.os.Build
+import com.android.server.connectivity.mdns.MdnsConstants
+import com.android.server.connectivity.mdns.MdnsConstants.FLAG_TRUNCATED
+import com.android.server.connectivity.mdns.MdnsPacket
+import com.android.server.connectivity.mdns.MdnsPacketReader
+import com.android.server.connectivity.mdns.MdnsPointerRecord
+import com.android.server.connectivity.mdns.MdnsRecord
+import com.android.server.connectivity.mdns.util.MdnsUtils.createQueryDatagramPackets
 import com.android.server.connectivity.mdns.util.MdnsUtils.equalsDnsLabelIgnoreDnsCase
 import com.android.server.connectivity.mdns.util.MdnsUtils.equalsIgnoreDnsCase
 import com.android.server.connectivity.mdns.util.MdnsUtils.toDnsLabelsLowerCase
@@ -24,6 +31,8 @@
 import com.android.server.connectivity.mdns.util.MdnsUtils.truncateServiceName
 import com.android.testutils.DevSdkIgnoreRule
 import com.android.testutils.DevSdkIgnoreRunner
+import java.net.DatagramPacket
+import kotlin.test.assertContentEquals
 import org.junit.Assert.assertArrayEquals
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
@@ -43,19 +52,27 @@
         assertEquals("ţést", toDnsLowerCase("ţést"))
         // Unicode characters 0x10000 (𐀀), 0x10001 (𐀁), 0x10041 (𐁁)
         // Note the last 2 bytes of 0x10041 are identical to 'A', but it should remain unchanged.
-        assertEquals("test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- ",
-                toDnsLowerCase("Test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- "))
+        assertEquals(
+            "test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- ",
+                toDnsLowerCase("Test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- ")
+        )
         // Also test some characters where the first surrogate is not \ud800
-        assertEquals("test: >\ud83c\udff4\udb40\udc67\udb40\udc62\udb40" +
+        assertEquals(
+            "test: >\ud83c\udff4\udb40\udc67\udb40\udc62\udb40" +
                 "\udc77\udb40\udc6c\udb40\udc73\udb40\udc7f<",
-                toDnsLowerCase("Test: >\ud83c\udff4\udb40\udc67\udb40\udc62\udb40" +
-                        "\udc77\udb40\udc6c\udb40\udc73\udb40\udc7f<"))
+                toDnsLowerCase(
+                    "Test: >\ud83c\udff4\udb40\udc67\udb40\udc62\udb40" +
+                        "\udc77\udb40\udc6c\udb40\udc73\udb40\udc7f<"
+                )
+        )
     }
 
     @Test
     fun testToDnsLabelsLowerCase() {
-        assertArrayEquals(arrayOf("test", "tÉst", "ţést"),
-            toDnsLabelsLowerCase(arrayOf("TeSt", "TÉST", "ţést")))
+        assertArrayEquals(
+            arrayOf("test", "tÉst", "ţést"),
+            toDnsLabelsLowerCase(arrayOf("TeSt", "TÉST", "ţést"))
+        )
     }
 
     @Test
@@ -67,13 +84,17 @@
         assertFalse(equalsIgnoreDnsCase("ŢÉST", "ţést"))
         // Unicode characters 0x10000 (𐀀), 0x10001 (𐀁), 0x10041 (𐁁)
         // Note the last 2 bytes of 0x10041 are identical to 'A', but it should remain unchanged.
-        assertTrue(equalsIgnoreDnsCase("test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- ",
-                "Test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- "))
+        assertTrue(equalsIgnoreDnsCase(
+            "test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- ",
+                "Test: -->\ud800\udc00 \ud800\udc01 \ud800\udc41<-- "
+        ))
         // Also test some characters where the first surrogate is not \ud800
-        assertTrue(equalsIgnoreDnsCase("test: >\ud83c\udff4\udb40\udc67\udb40\udc62\udb40" +
+        assertTrue(equalsIgnoreDnsCase(
+            "test: >\ud83c\udff4\udb40\udc67\udb40\udc62\udb40" +
                 "\udc77\udb40\udc6c\udb40\udc73\udb40\udc7f<",
                 "Test: >\ud83c\udff4\udb40\udc67\udb40\udc62\udb40" +
-                        "\udc77\udb40\udc6c\udb40\udc73\udb40\udc7f<"))
+                        "\udc77\udb40\udc6c\udb40\udc73\udb40\udc7f<"
+        ))
     }
 
     @Test
@@ -92,14 +113,84 @@
 
     @Test
     fun testTypeEqualsOrIsSubtype() {
-        assertTrue(MdnsUtils.typeEqualsOrIsSubtype(arrayOf("_type", "_tcp", "local"),
-            arrayOf("_type", "_TCP", "local")))
-        assertTrue(MdnsUtils.typeEqualsOrIsSubtype(arrayOf("_type", "_tcp", "local"),
-            arrayOf("a", "_SUB", "_type", "_TCP", "local")))
-        assertFalse(MdnsUtils.typeEqualsOrIsSubtype(arrayOf("_sub", "_type", "_tcp", "local"),
-                arrayOf("_type", "_TCP", "local")))
+        assertTrue(MdnsUtils.typeEqualsOrIsSubtype(
+            arrayOf("_type", "_tcp", "local"),
+            arrayOf("_type", "_TCP", "local")
+        ))
+        assertTrue(MdnsUtils.typeEqualsOrIsSubtype(
+            arrayOf("_type", "_tcp", "local"),
+            arrayOf("a", "_SUB", "_type", "_TCP", "local")
+        ))
+        assertFalse(MdnsUtils.typeEqualsOrIsSubtype(
+            arrayOf("_sub", "_type", "_tcp", "local"),
+                arrayOf("_type", "_TCP", "local")
+        ))
         assertFalse(MdnsUtils.typeEqualsOrIsSubtype(
                 arrayOf("a", "_other", "_type", "_tcp", "local"),
-                arrayOf("a", "_SUB", "_type", "_TCP", "local")))
+                arrayOf("a", "_SUB", "_type", "_TCP", "local")
+        ))
+    }
+
+    @Test
+    fun testCreateQueryDatagramPackets() {
+        // Question data bytes:
+        // Name label(17)(duplicated labels) + PTR type(2) + cacheFlush(2) = 21
+        //
+        // Known answers data bytes:
+        // Name label(17)(duplicated labels) + PTR type(2) + cacheFlush(2) + receiptTimeMillis(4)
+        // + Data length(2) + Pointer data(18)(duplicated labels) = 45
+        val questions = mutableListOf<MdnsRecord>()
+        val knownAnswers = mutableListOf<MdnsRecord>()
+        for (i in 1..100) {
+            questions.add(MdnsPointerRecord(arrayOf("_testservice$i", "_tcp", "local"), false))
+            knownAnswers.add(MdnsPointerRecord(
+                    arrayOf("_testservice$i", "_tcp", "local"),
+                    0L,
+                    false,
+                    4_500_000L,
+                    arrayOf("MyTestService$i", "_testservice$i", "_tcp", "local")
+            ))
+        }
+        // MdnsPacket data bytes:
+        // Questions(21 * 100) + Answers(45 * 100) = 6600 -> at least 5 packets
+        val query = MdnsPacket(
+                MdnsConstants.FLAGS_QUERY,
+                questions as List<MdnsRecord>,
+                knownAnswers as List<MdnsRecord>,
+                emptyList(),
+                emptyList()
+        )
+        // Expect the oversize MdnsPacket to be separated into 5 DatagramPackets.
+        val bufferSize = 1500
+        val packets = createQueryDatagramPackets(
+                ByteArray(bufferSize),
+                query,
+                MdnsConstants.IPV4_SOCKET_ADDR
+        )
+        assertEquals(5, packets.size)
+        assertTrue(packets.all { packet -> packet.length < bufferSize })
+
+        val mdnsPacket = createMdnsPacketFromMultipleDatagramPackets(packets)
+        assertEquals(query.flags, mdnsPacket.flags)
+        assertContentEquals(query.questions, mdnsPacket.questions)
+        assertContentEquals(query.answers, mdnsPacket.answers)
+    }
+
+    private fun createMdnsPacketFromMultipleDatagramPackets(
+            packets: List<DatagramPacket>
+    ): MdnsPacket {
+        var flags = 0
+        val questions = mutableListOf<MdnsRecord>()
+        val answers = mutableListOf<MdnsRecord>()
+        for ((index, packet) in packets.withIndex()) {
+            val mdnsPacket = MdnsPacket.parse(MdnsPacketReader(packet))
+            if (index != packets.size - 1) {
+                assertTrue((mdnsPacket.flags and FLAG_TRUNCATED) == FLAG_TRUNCATED)
+            }
+            flags = mdnsPacket.flags
+            questions.addAll(mdnsPacket.questions)
+            answers.addAll(mdnsPacket.answers)
+        }
+        return MdnsPacket(flags, questions, answers, emptyList(), emptyList())
     }
 }
diff --git a/tests/unit/java/com/android/server/connectivityservice/CSFirewallChainTest.kt b/tests/unit/java/com/android/server/connectivityservice/CSFirewallChainTest.kt
new file mode 100644
index 0000000..16de4da
--- /dev/null
+++ b/tests/unit/java/com/android/server/connectivityservice/CSFirewallChainTest.kt
@@ -0,0 +1,157 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server
+
+import android.net.ConnectivityManager
+import android.os.Build
+import androidx.test.filters.SmallTest
+import com.android.server.connectivity.ConnectivityFlags.BACKGROUND_FIREWALL_CHAIN
+import com.android.testutils.DevSdkIgnoreRule
+import com.android.testutils.DevSdkIgnoreRule.IgnoreAfter
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import com.android.testutils.DevSdkIgnoreRunner
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyBoolean
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Mockito.any
+import org.mockito.Mockito.clearInvocations
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+
+@RunWith(DevSdkIgnoreRunner::class)
+@SmallTest
+@IgnoreUpTo(Build.VERSION_CODES.S_V2)
+class CSFirewallChainTest : CSTest() {
+    @get:Rule
+    val ignoreRule = DevSdkIgnoreRule()
+
+    // Tests for setFirewallChainEnabled on FIREWALL_CHAIN_BACKGROUND
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, false)])
+    fun setFirewallChainEnabled_backgroundChainDisabled() {
+        verifySetFirewallChainEnabledOnBackgroundDoesNothing()
+    }
+
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, true)])
+    @IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    fun setFirewallChainEnabled_backgroundChainEnabled_afterU() {
+        cm.setFirewallChainEnabled(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, true)
+        verify(bpfNetMaps).setChildChain(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, true)
+
+        clearInvocations(bpfNetMaps)
+
+        cm.setFirewallChainEnabled(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, false)
+        verify(bpfNetMaps).setChildChain(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, false)
+    }
+
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, true)])
+    @IgnoreAfter(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    fun setFirewallChainEnabled_backgroundChainEnabled_uptoU() {
+        verifySetFirewallChainEnabledOnBackgroundDoesNothing()
+    }
+
+    private fun verifySetFirewallChainEnabledOnBackgroundDoesNothing() {
+        cm.setFirewallChainEnabled(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, true)
+        verify(bpfNetMaps, never()).setChildChain(anyInt(), anyBoolean())
+
+        cm.setFirewallChainEnabled(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, false)
+        verify(bpfNetMaps, never()).setChildChain(anyInt(), anyBoolean())
+    }
+
+    // Tests for replaceFirewallChain on FIREWALL_CHAIN_BACKGROUND
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, false)])
+    fun replaceFirewallChain_backgroundChainDisabled() {
+        verifyReplaceFirewallChainOnBackgroundDoesNothing()
+    }
+
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, true)])
+    @IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    fun replaceFirewallChain_backgroundChainEnabled_afterU() {
+        val uids = intArrayOf(53, 42, 79)
+        cm.replaceFirewallChain(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uids)
+        verify(bpfNetMaps).replaceUidChain(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uids)
+    }
+
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, true)])
+    @IgnoreAfter(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    fun replaceFirewallChain_backgroundChainEnabled_uptoU() {
+        verifyReplaceFirewallChainOnBackgroundDoesNothing()
+    }
+
+    private fun verifyReplaceFirewallChainOnBackgroundDoesNothing() {
+        val uids = intArrayOf(53, 42, 79)
+        cm.replaceFirewallChain(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uids)
+        verify(bpfNetMaps, never()).replaceUidChain(anyInt(), any(IntArray::class.java))
+    }
+
+    // Tests for setUidFirewallRule on FIREWALL_CHAIN_BACKGROUND
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, false)])
+    fun setUidFirewallRule_backgroundChainDisabled() {
+        verifySetUidFirewallRuleOnBackgroundDoesNothing()
+    }
+
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, true)])
+    @IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    fun setUidFirewallRule_backgroundChainEnabled_afterU() {
+        val uid = 2345
+
+        cm.setUidFirewallRule(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uid,
+            ConnectivityManager.FIREWALL_RULE_DEFAULT)
+        verify(bpfNetMaps).setUidRule(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uid,
+            ConnectivityManager.FIREWALL_RULE_DENY)
+
+        clearInvocations(bpfNetMaps)
+
+        cm.setUidFirewallRule(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uid,
+            ConnectivityManager.FIREWALL_RULE_DENY)
+        verify(bpfNetMaps).setUidRule(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uid,
+            ConnectivityManager.FIREWALL_RULE_DENY)
+
+        clearInvocations(bpfNetMaps)
+
+        cm.setUidFirewallRule(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uid,
+            ConnectivityManager.FIREWALL_RULE_ALLOW)
+        verify(bpfNetMaps).setUidRule(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uid,
+            ConnectivityManager.FIREWALL_RULE_ALLOW)
+    }
+
+    @Test
+    @FeatureFlags(flags = [Flag(BACKGROUND_FIREWALL_CHAIN, true)])
+    @IgnoreAfter(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    fun setUidFirewallRule_backgroundChainEnabled_uptoU() {
+        verifySetUidFirewallRuleOnBackgroundDoesNothing()
+    }
+
+    private fun verifySetUidFirewallRuleOnBackgroundDoesNothing() {
+        val uid = 2345
+
+        listOf(ConnectivityManager.FIREWALL_RULE_DEFAULT, ConnectivityManager.FIREWALL_RULE_ALLOW,
+            ConnectivityManager.FIREWALL_RULE_DENY).forEach { rule ->
+            cm.setUidFirewallRule(ConnectivityManager.FIREWALL_CHAIN_BACKGROUND, uid, rule)
+            verify(bpfNetMaps, never()).setUidRule(anyInt(), anyInt(), anyInt())
+        }
+    }
+}
diff --git a/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt b/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt
index e8664c1..bb7fb51 100644
--- a/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/CSIngressDiscardRuleTests.kt
@@ -30,6 +30,7 @@
 import android.net.VpnTransportInfo
 import android.os.Build
 import androidx.test.filters.SmallTest
+import com.android.server.connectivity.ConnectivityFlags
 import com.android.testutils.DevSdkIgnoreRule
 import com.android.testutils.DevSdkIgnoreRunner
 import com.android.testutils.RecorderCallback.CallbackEntry.LinkPropertiesChanged
@@ -56,7 +57,9 @@
                         TYPE_VPN_SERVICE,
                         "MySession12345",
                         false /* bypassable */,
-                        false /* longLivedTcpConnectionsExpensive */))
+                        false /* longLivedTcpConnectionsExpensive */
+                )
+        )
         .build()
 
 private fun wifiNc() = NetworkCapabilities.Builder()
@@ -286,4 +289,19 @@
         waitForIdle()
         verify(bpfNetMaps).removeIngressDiscardRule(IPV6_ADDRESS)
     }
+
+    @Test @FeatureFlags([Flag(ConnectivityFlags.INGRESS_TO_VPN_ADDRESS_FILTERING, false)])
+    fun testVpnIngressDiscardRule_FeatureDisabled() {
+        val nr = nr(TRANSPORT_VPN)
+        val cb = TestableNetworkCallback()
+        cm.registerNetworkCallback(nr, cb)
+        val nc = vpnNc()
+        val lp = lp(VPN_IFNAME, IPV6_LINK_ADDRESS, LOCAL_IPV6_LINK_ADDRRESS)
+        val agent = Agent(nc = nc, lp = lp)
+        agent.connect()
+        cb.expectAvailableCallbacks(agent.network, validated = false)
+
+        // IngressDiscardRule should not be added since feature is disabled
+        verify(bpfNetMaps, never()).setIngressDiscardRule(any(), any())
+    }
 }
diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
index 5c4617b..3b06ad0 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
@@ -79,11 +79,15 @@
 import java.util.concurrent.TimeUnit
 import java.util.function.BiConsumer
 import java.util.function.Consumer
+import kotlin.annotation.AnnotationRetention.RUNTIME
+import kotlin.annotation.AnnotationTarget.FUNCTION
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 import kotlin.test.fail
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
+import org.junit.rules.TestName
 import org.mockito.AdditionalAnswers.delegatesTo
 import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.doReturn
@@ -126,6 +130,9 @@
 // TODO (b/272685721) : make ConnectivityServiceTest smaller and faster by moving the setup
 // parts into this class and moving the individual tests to multiple separate classes.
 open class CSTest {
+    @get:Rule
+    val testNameRule = TestName()
+
     companion object {
         val CSTestExecutor = Executors.newSingleThreadExecutor()
     }
@@ -154,9 +161,9 @@
         it[ConnectivityService.ALLOW_SYSUI_CONNECTIVITY_REPORTS] = true
         it[ConnectivityService.ALLOW_SATALLITE_NETWORK_FALLBACK] = true
         it[ConnectivityFlags.INGRESS_TO_VPN_ADDRESS_FILTERING] = true
+        it[ConnectivityFlags.BACKGROUND_FIREWALL_CHAIN] = true
     }
-    fun enableFeature(f: String) = enabledFeatures.set(f, true)
-    fun disableFeature(f: String) = enabledFeatures.set(f, false)
+    fun setFeatureEnabled(flag: String, enabled: Boolean) = enabledFeatures.set(flag, enabled)
 
     // When adding new members, consider if it's not better to build the object in CSTestHelpers
     // to keep this file clean of implementation details. Generally, CSTestHelpers should only
@@ -201,8 +208,32 @@
     lateinit var cm: ConnectivityManager
     lateinit var csHandler: Handler
 
+    // Tests can use this annotation to set flag values before constructing ConnectivityService
+    // e.g. @FeatureFlags([Flag(flagName1, true/false), Flag(flagName2, true/false)])
+    @Retention(RUNTIME)
+    @Target(FUNCTION)
+    annotation class FeatureFlags(val flags: Array<Flag>)
+
+    @Retention(RUNTIME)
+    @Target(FUNCTION)
+    annotation class Flag(val name: String, val enabled: Boolean)
+
     @Before
     fun setUp() {
+        // Set feature flags before constructing ConnectivityService
+        val testMethodName = testNameRule.methodName
+        try {
+            val testMethod = this::class.java.getMethod(testMethodName)
+            val featureFlags = testMethod.getAnnotation(FeatureFlags::class.java)
+            if (featureFlags != null) {
+                for (flag in featureFlags.flags) {
+                    setFeatureEnabled(flag.name, flag.enabled)
+                }
+            }
+        } catch (ignored: NoSuchMethodException) {
+            // This is expected for parameterized tests
+        }
+
         alarmHandlerThread = HandlerThread("TestAlarmManager").also { it.start() }
         alarmManager = makeMockAlarmManager(alarmHandlerThread)
         service = makeConnectivityService(context, netd, deps).also { it.systemReadyInternal() }
diff --git a/tests/unit/java/com/android/server/ethernet/EthernetNetworkFactoryTest.java b/tests/unit/java/com/android/server/ethernet/EthernetNetworkFactoryTest.java
index 949e0c2..70d4ad8 100644
--- a/tests/unit/java/com/android/server/ethernet/EthernetNetworkFactoryTest.java
+++ b/tests/unit/java/com/android/server/ethernet/EthernetNetworkFactoryTest.java
@@ -554,4 +554,22 @@
         mNetworkOfferCallback.onNetworkNeeded(createDefaultRequest());
         verify(mIpClient, never()).startProvisioning(any());
     }
+
+    @Test
+    public void testGetMacAddressProvisionedInterface() throws Exception {
+        initEthernetNetworkFactory();
+        createAndVerifyProvisionedInterface(TEST_IFACE);
+
+        final String result = mNetFactory.getHwAddress(TEST_IFACE);
+        assertEquals(HW_ADDR, result);
+    }
+
+    @Test
+    public void testGetMacAddressForNonExistingInterface() {
+        initEthernetNetworkFactory();
+
+        final String result = mNetFactory.getHwAddress(TEST_IFACE);
+        // No interface exists due to not calling createAndVerifyProvisionedInterface(...).
+        assertNull(result);
+    }
 }
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java b/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
index 0bbc34c..e62ac74 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsObserversTest.java
@@ -47,6 +47,7 @@
 import android.net.NetworkTemplate;
 import android.os.HandlerThread;
 import android.os.IBinder;
+import android.os.Looper;
 import android.os.Process;
 import android.os.UserHandle;
 import android.telephony.TelephonyManager;
@@ -126,7 +127,13 @@
 
         mObserverHandlerThread = new HandlerThread("NetworkStatsObserversTest");
         mObserverHandlerThread.start();
-        mStatsObservers = new NetworkStatsObservers(mObserverHandlerThread.getLooper());
+        final Looper observerLooper = mObserverHandlerThread.getLooper();
+        mStatsObservers = new NetworkStatsObservers() {
+            @Override
+            protected Looper getHandlerLooperLocked() {
+                return observerLooper;
+            }
+        };
 
         mActiveIfaces = new ArrayMap<>();
         mActiveUidIfaces = new ArrayMap<>();
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 3d7ad66..6425daa 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -67,11 +67,14 @@
 import static com.android.server.net.NetworkStatsEventLogger.POLL_REASON_RAT_CHANGED;
 import static com.android.server.net.NetworkStatsEventLogger.PollEvent.pollReasonNameOf;
 import static com.android.server.net.NetworkStatsService.ACTION_NETWORK_STATS_POLL;
+import static com.android.server.net.NetworkStatsService.DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS;
+import static com.android.server.net.NetworkStatsService.DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES;
 import static com.android.server.net.NetworkStatsService.NETSTATS_FASTDATAINPUT_FALLBACKS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_FASTDATAINPUT_SUCCESSES_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_ATTEMPTS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_FALLBACKS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_SUCCESSES_COUNTER_NAME;
+import static com.android.server.net.NetworkStatsService.TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.junit.Assert.assertEquals;
@@ -116,6 +119,7 @@
 import android.net.TestNetworkSpecifier;
 import android.net.TetherStatsParcel;
 import android.net.TetheringManager;
+import android.net.TrafficStats;
 import android.net.UnderlyingNetworkInfo;
 import android.net.netstats.provider.INetworkStatsProviderCallback;
 import android.net.wifi.WifiInfo;
@@ -123,7 +127,9 @@
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.IBinder;
+import android.os.Looper;
 import android.os.PowerManager;
+import android.os.Process;
 import android.os.SimpleClock;
 import android.provider.Settings;
 import android.system.ErrnoException;
@@ -158,12 +164,15 @@
 import com.android.testutils.HandlerUtils;
 import com.android.testutils.TestBpfMap;
 import com.android.testutils.TestableNetworkStatsProviderBinder;
+import com.android.testutils.com.android.testutils.SetFeatureFlagsRule;
+import com.android.testutils.com.android.testutils.SetFeatureFlagsRule.FeatureFlag;
 
 import libcore.testing.io.TestIoUtils;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -188,6 +197,7 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
 
 /**
  * Tests for {@link NetworkStatsService}.
@@ -201,6 +211,7 @@
 // NetworkStatsService is not updatable before T, so tests do not need to be backwards compatible
 @DevSdkIgnoreRule.IgnoreUpTo(SC_V2)
 public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
+
     private static final String TAG = "NetworkStatsServiceTest";
 
     private static final long TEST_START = 1194220800000L;
@@ -292,7 +303,18 @@
     private String mCompareStatsResult = null;
     private @Mock Resources mResources;
     private Boolean mIsDebuggable;
+    private HandlerThread mObserverHandlerThread;
     final TestDependencies mDeps = new TestDependencies();
+    final HashMap<String, Boolean> mFeatureFlags = new HashMap<>();
+
+    // This will set feature flags from @FeatureFlag annotations
+    // into the map before setUp() runs.
+    @Rule
+    public final SetFeatureFlagsRule mSetFeatureFlagsRule =
+            new SetFeatureFlagsRule((name, enabled) -> {
+                mFeatureFlags.put(name, enabled);
+                return null;
+            });
 
     private class MockContext extends BroadcastInterceptingContext {
         private final Context mBaseContext;
@@ -375,13 +397,24 @@
                 powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
 
         mHandlerThread = new HandlerThread("NetworkStatsServiceTest-HandlerThread");
+        // Create a separate thread for observers to run on. This thread cannot be the same
+        // as the handler thread, because the observer callback is fired on this thread, and
+        // it should not be blocked by client code. Additionally, creating the observers
+        // object requires a looper, which can only be obtained after a thread has been started.
+        mObserverHandlerThread = new HandlerThread("NetworkStatsServiceTest-ObserversThread");
+        mObserverHandlerThread.start();
+        final Looper observerLooper = mObserverHandlerThread.getLooper();
+        final NetworkStatsObservers statsObservers = new NetworkStatsObservers() {
+            @Override
+            protected Looper getHandlerLooperLocked() {
+                return observerLooper;
+            }
+        };
         mService = new NetworkStatsService(mServiceContext, mNetd, mAlarmManager, wakeLock,
-                mClock, mSettings, mStatsFactory, mDeps);
+                mClock, mSettings, mStatsFactory, statsObservers, mDeps);
 
         mElapsedRealtime = 0L;
 
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
         // Verify that system ready fetches realtime stats
@@ -417,6 +450,7 @@
 
     class TestDependencies extends NetworkStatsService.Dependencies {
         private int mCompareStatsInvocation = 0;
+        private NetworkStats.Entry mMockedTrafficStatsNativeStat = null;
 
         @Override
         public File getLegacyStatsDir() {
@@ -558,6 +592,43 @@
         public boolean supportEventLogger(@NonNull Context cts) {
             return true;
         }
+
+        @Override
+        public boolean supportTrafficStatsRateLimitCache(Context ctx) {
+            return mFeatureFlags.getOrDefault(TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, false);
+        }
+
+        @Override
+        public int getTrafficStatsRateLimitCacheExpiryDuration() {
+            return DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS;
+        }
+
+        @Override
+        public int getTrafficStatsRateLimitCacheMaxEntries() {
+            return DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES;
+        }
+
+        @Nullable
+        @Override
+        public NetworkStats.Entry nativeGetTotalStat() {
+            return mMockedTrafficStatsNativeStat;
+        }
+
+        @Nullable
+        @Override
+        public NetworkStats.Entry nativeGetIfaceStat(String iface) {
+            return mMockedTrafficStatsNativeStat;
+        }
+
+        @Nullable
+        @Override
+        public NetworkStats.Entry nativeGetUidStat(int uid) {
+            return mMockedTrafficStatsNativeStat;
+        }
+
+        public void setNativeStat(NetworkStats.Entry entry) {
+            mMockedTrafficStatsNativeStat = entry;
+        }
     }
 
     @After
@@ -574,6 +645,10 @@
             mHandlerThread.quitSafely();
             mHandlerThread.join();
         }
+        if (mObserverHandlerThread != null) {
+            mObserverHandlerThread.quitSafely();
+            mObserverHandlerThread.join();
+        }
     }
 
     private void initWifiStats(NetworkStateSnapshot snapshot) throws Exception {
@@ -710,8 +785,6 @@
         assertStatsFilesExist(true);
 
         // boot through serviceReady() again
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
 
         mService.systemReady();
@@ -910,7 +983,16 @@
     }
 
     @Test
-    public void testMobileStatsByRatType() throws Exception {
+    public void testMobileStatsByRatTypeForSatellite() throws Exception {
+        doTestMobileStatsByRatType(new NetworkStateSnapshot[]{buildSatelliteMobileState(IMSI_1)});
+    }
+
+    @Test
+    public void testMobileStatsByRatTypeForCellular() throws Exception {
+        doTestMobileStatsByRatType(new NetworkStateSnapshot[]{buildMobileState(IMSI_1)});
+    }
+
+    private void doTestMobileStatsByRatType(NetworkStateSnapshot[] states) throws Exception {
         final NetworkTemplate template3g = new NetworkTemplate.Builder(MATCH_MOBILE)
                 .setRatType(TelephonyManager.NETWORK_TYPE_UMTS)
                 .setMeteredness(METERED_YES).build();
@@ -920,8 +1002,6 @@
         final NetworkTemplate template5g = new NetworkTemplate.Builder(MATCH_MOBILE)
                 .setRatType(TelephonyManager.NETWORK_TYPE_NR)
                 .setMeteredness(METERED_YES).build();
-        final NetworkStateSnapshot[] states =
-                new NetworkStateSnapshot[]{buildMobileState(IMSI_1)};
 
         // 3G network comes online.
         mockNetworkStatsSummary(buildEmptyStats());
@@ -935,7 +1015,7 @@
         incrementCurrentTime(MINUTE_IN_MILLIS);
         mockNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 1)
                 .addEntry(new NetworkStats.Entry(TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE,
-                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 12L, 18L, 14L, 1L, 0L)));
+                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 12L, 18L, 14L, 1L, 0L)));
         forcePollAndWaitForIdle();
 
         // Verify 3g templates gets stats.
@@ -950,7 +1030,7 @@
         mockNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 1)
                 // Append more traffic on existing 3g stats entry.
                 .addEntry(new NetworkStats.Entry(TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE,
-                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 16L, 22L, 17L, 2L, 0L))
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 16L, 22L, 17L, 2L, 0L))
                 // Add entry that is new on 4g.
                 .addEntry(new NetworkStats.Entry(TEST_IFACE, UID_RED, SET_FOREGROUND, TAG_NONE,
                         METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 33L, 27L, 8L, 10L, 1L)));
@@ -1352,6 +1432,57 @@
     }
 
     @Test
+    public void testGetUidStatsForTransportWithCellularAndSatellite() throws Exception {
+        // Setup satellite mobile network and Cellular mobile network
+        mockDefaultSettings();
+        mockNetworkStatsUidDetail(buildEmptyStats());
+
+        final NetworkStateSnapshot mobileState = buildStateOfTransport(
+                NetworkCapabilities.TRANSPORT_CELLULAR, TYPE_MOBILE,
+                TEST_IFACE2, IMSI_1, null /* wifiNetworkKey */,
+                false /* isTemporarilyNotMetered */, false /* isRoaming */);
+
+        final NetworkStateSnapshot[] states = new NetworkStateSnapshot[]{mobileState,
+                buildSatelliteMobileState(IMSI_1)};
+        mService.notifyNetworkStatus(NETWORKS_MOBILE, states, getActiveIface(states),
+                new UnderlyingNetworkInfo[0]);
+        setMobileRatTypeAndWaitForIdle(TelephonyManager.NETWORK_TYPE_LTE);
+
+        // mock traffic on satellite network
+        final NetworkStats.Entry entrySatellite = new NetworkStats.Entry(
+                TEST_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 80L, 5L, 70L, 15L, 1L);
+
+        // mock traffic on cellular network
+        final NetworkStats.Entry entryCellular = new NetworkStats.Entry(
+                TEST_IFACE2, UID_RED, SET_DEFAULT, TAG_NONE, METERED_NO, ROAMING_NO,
+                DEFAULT_NETWORK_NO, 100L, 15L, 150L, 15L, 1L);
+
+        final TetherStatsParcel[] emptyTetherStats = {};
+        // The interfaces that expect to be used to query the stats.
+        final String[] mobileIfaces = {TEST_IFACE, TEST_IFACE2};
+        incrementCurrentTime(HOUR_IN_MILLIS);
+        mockDefaultSettings();
+        mockNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 2)
+                .insertEntry(entrySatellite).insertEntry(entryCellular), emptyTetherStats,
+                mobileIfaces);
+        // with getUidStatsForTransport(TRANSPORT_CELLULAR) return stats of both cellular
+        // and satellite
+        final NetworkStats mobileStats = mService.getUidStatsForTransport(
+                NetworkCapabilities.TRANSPORT_CELLULAR);
+
+        // The iface field of the returned stats should be null because getUidStatsForTransport
+        // clears the interface field before it returns the result.
+        assertValues(mobileStats, null /* iface */, UID_RED, SET_DEFAULT, TAG_NONE,
+                METERED_NO, ROAMING_NO, METERED_NO, 180L, 20L, 220L, 30L, 2L);
+
+        // getUidStatsForTransport(TRANSPORT_SATELLITE) is not supported
+        assertThrows(IllegalArgumentException.class,
+                () -> mService.getUidStatsForTransport(NetworkCapabilities.TRANSPORT_SATELLITE));
+
+    }
+
+    @Test
     public void testForegroundBackground() throws Exception {
         // pretend that network comes online
         mockDefaultSettings();
@@ -2034,8 +2165,6 @@
                 getLegacyCollection(PREFIX_UID_TAG, true /* includeTags */));
 
         // Mock zero usage and boot through serviceReady(), verify there is no imported data.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
         assertStatsFilesExist(false);
@@ -2047,8 +2176,6 @@
         assertStatsFilesExist(false);
 
         // Boot through systemReady() again.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
 
@@ -2122,8 +2249,6 @@
                 getLegacyCollection(PREFIX_UID_TAG, true /* includeTags */));
 
         // Mock zero usage and boot through serviceReady(), verify there is no imported data.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
         assertStatsFilesExist(false);
@@ -2135,8 +2260,6 @@
         assertStatsFilesExist(false);
 
         // Boot through systemReady() again.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
 
@@ -2288,6 +2411,68 @@
         assertUidTotal(sTemplateWifi, UID_GREEN, 64L, 3L, 1024L, 8L, 0);
     }
 
+    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false)
+    @Test
+    public void testTrafficStatsRateLimitCache_disabled() throws Exception {
+        doTestTrafficStatsRateLimitCache(false /* cacheEnabled */);
+    }
+
+    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @Test
+    public void testTrafficStatsRateLimitCache_enabled() throws Exception {
+        doTestTrafficStatsRateLimitCache(true /* cacheEnabled */);
+    }
+
+    private void doTestTrafficStatsRateLimitCache(boolean cacheEnabled) throws Exception {
+        mockDefaultSettings();
+        // Calling uid is not injected into the service, use the real uid to pass the caller check.
+        final int myUid = Process.myUid();
+        mockTrafficStatsValues(64L, 3L, 1024L, 8L);
+        assertTrafficStatsValues(TEST_IFACE, myUid, 64L, 3L, 1024L, 8L);
+
+        // Verify the values are cached.
+        incrementCurrentTime(DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS / 2);
+        mockTrafficStatsValues(65L, 8L, 1055L, 9L);
+        if (cacheEnabled) {
+            assertTrafficStatsValues(TEST_IFACE, myUid, 64L, 3L, 1024L, 8L);
+        } else {
+            assertTrafficStatsValues(TEST_IFACE, myUid, 65L, 8L, 1055L, 9L);
+        }
+
+        // Verify the values are updated after cache expiry.
+        incrementCurrentTime(DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS);
+        assertTrafficStatsValues(TEST_IFACE, myUid, 65L, 8L, 1055L, 9L);
+    }
+
+    private void mockTrafficStatsValues(long rxBytes, long rxPackets,
+            long txBytes, long txPackets) {
+        // In practice, keys and operations are not used and filled with default values when
+        // returned by JNI layer.
+        final NetworkStats.Entry entry = new NetworkStats.Entry(IFACE_ALL, UID_ALL, SET_DEFAULT,
+                TAG_NONE, METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO,
+                rxBytes, rxPackets, txBytes, txPackets, 0L);
+        mDeps.setNativeStat(entry);
+    }
+
+    // Assert for 3 different API return values respectively.
+    private void assertTrafficStatsValues(String iface, int uid, long rxBytes, long rxPackets,
+            long txBytes, long txPackets) {
+        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
+                (type) -> mService.getTotalStats(type));
+        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
+                (type) -> mService.getIfaceStats(iface, type));
+        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
+                (type) -> mService.getUidStats(uid, type));
+    }
+
+    private void assertTrafficStatsValuesThat(long rxBytes, long rxPackets, long txBytes,
+            long txPackets, Function<Integer, Long> fetcher) {
+        assertEquals(rxBytes, (long) fetcher.apply(TrafficStats.TYPE_RX_BYTES));
+        assertEquals(rxPackets, (long) fetcher.apply(TrafficStats.TYPE_RX_PACKETS));
+        assertEquals(txBytes, (long) fetcher.apply(TrafficStats.TYPE_TX_BYTES));
+        assertEquals(txPackets, (long) fetcher.apply(TrafficStats.TYPE_TX_PACKETS));
+    }
+
     private void assertShouldRunComparison(boolean expected, boolean isDebuggable) {
         assertEquals("shouldRunComparison (debuggable=" + isDebuggable + "): ",
                 expected, mService.shouldRunComparison());
@@ -2352,6 +2537,8 @@
     }
 
     private void prepareForSystemReady() throws Exception {
+        mockDefaultSettings();
+        mockNetworkStatsUidDetail(buildEmptyStats());
         mockNetworkStatsSummary(buildEmptyStats());
     }
 
@@ -2509,6 +2696,12 @@
                 false /* isTemporarilyNotMetered */, false /* isRoaming */);
     }
 
+    private static NetworkStateSnapshot buildSatelliteMobileState(String subscriberId) {
+        return buildStateOfTransport(NetworkCapabilities.TRANSPORT_SATELLITE, TYPE_MOBILE,
+                TEST_IFACE, subscriberId, null /* wifiNetworkKey */,
+                false /* isTemporarilyNotMetered */, false /* isRoaming */);
+    }
+
     private static NetworkStateSnapshot buildTestState(@NonNull String iface,
             @Nullable String wifiNetworkKey) {
         return buildStateOfTransport(NetworkCapabilities.TRANSPORT_TEST, TYPE_TEST,
diff --git a/thread/framework/java/android/net/thread/ActiveOperationalDataset.java b/thread/framework/java/android/net/thread/ActiveOperationalDataset.java
index 4a7d3a7..22457f5 100644
--- a/thread/framework/java/android/net/thread/ActiveOperationalDataset.java
+++ b/thread/framework/java/android/net/thread/ActiveOperationalDataset.java
@@ -18,7 +18,7 @@
 
 import static com.android.internal.util.Preconditions.checkArgument;
 import static com.android.internal.util.Preconditions.checkState;
-import static com.android.net.module.util.HexDump.dumpHexString;
+import static com.android.net.module.util.HexDump.toHexString;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
@@ -610,7 +610,7 @@
         sb.append("{networkName=")
                 .append(getNetworkName())
                 .append(", extendedPanId=")
-                .append(dumpHexString(getExtendedPanId()))
+                .append(toHexString(getExtendedPanId()))
                 .append(", panId=")
                 .append(getPanId())
                 .append(", channel=")
@@ -1109,7 +1109,7 @@
             sb.append("{rotation=")
                     .append(mRotationTimeHours)
                     .append(", flags=")
-                    .append(dumpHexString(mFlags))
+                    .append(toHexString(mFlags))
                     .append("}");
             return sb.toString();
         }
diff --git a/thread/service/Android.bp b/thread/service/Android.bp
index 6e2fac1..a82a499 100644
--- a/thread/service/Android.bp
+++ b/thread/service/Android.bp
@@ -45,6 +45,9 @@
         "modules-utils-shell-command-handler",
         "net-utils-device-common",
         "net-utils-device-common-netlink",
+        // The required dependency net-utils-device-common-struct-base is in the classpath via
+        // framework-connectivity
+        "net-utils-device-common-struct",
         "ot-daemon-aidl-java",
     ],
     apex_available: ["com.android.tethering"],
diff --git a/thread/service/java/com/android/server/thread/NsdPublisher.java b/thread/service/java/com/android/server/thread/NsdPublisher.java
index 3c7a72b..2c14f1d 100644
--- a/thread/service/java/com/android/server/thread/NsdPublisher.java
+++ b/thread/service/java/com/android/server/thread/NsdPublisher.java
@@ -21,6 +21,7 @@
 import android.annotation.NonNull;
 import android.content.Context;
 import android.net.InetAddresses;
+import android.net.nsd.DiscoveryRequest;
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
 import android.os.Handler;
@@ -31,15 +32,18 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.thread.openthread.DnsTxtAttribute;
+import com.android.server.thread.openthread.INsdDiscoverServiceCallback;
 import com.android.server.thread.openthread.INsdPublisher;
+import com.android.server.thread.openthread.INsdResolveServiceCallback;
 import com.android.server.thread.openthread.INsdStatusReceiver;
 
+import java.net.Inet6Address;
 import java.net.InetAddress;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Deque;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Executor;
 
 /**
@@ -50,14 +54,6 @@
  *
  * <p>All the data members of this class MUST be accessed in the {@code mHandler}'s Thread except
  * {@code mHandler} itself.
- *
- * <p>TODO: b/323300118 - Remove the following mechanism when the race condition in NsdManager is
- * fixed.
- *
- * <p>There's always only one running registration job at any timepoint. All other pending jobs are
- * queued in {@code mRegistrationJobs}. When a registration job is complete (i.e. the according
- * method in {@link NsdManager.RegistrationListener} is called), it will start the next registration
- * job in the queue.
  */
 public final class NsdPublisher extends INsdPublisher.Stub {
     // TODO: b/321883491 - specify network for mDNS operations
@@ -66,7 +62,8 @@
     private final Handler mHandler;
     private final Executor mExecutor;
     private final SparseArray<RegistrationListener> mRegistrationListeners = new SparseArray<>(0);
-    private final Deque<Runnable> mRegistrationJobs = new ArrayDeque<>();
+    private final SparseArray<DiscoveryListener> mDiscoveryListeners = new SparseArray<>(0);
+    private final SparseArray<ServiceInfoListener> mServiceInfoListeners = new SparseArray<>(0);
 
     @VisibleForTesting
     public NsdPublisher(NsdManager nsdManager, Handler handler) {
@@ -89,13 +86,9 @@
             List<DnsTxtAttribute> txt,
             INsdStatusReceiver receiver,
             int listenerId) {
-        postRegistrationJob(
-                () -> {
-                    NsdServiceInfo serviceInfo =
-                            buildServiceInfoForService(
-                                    hostname, name, type, subTypeList, port, txt);
-                    registerInternal(serviceInfo, receiver, listenerId, "service");
-                });
+        NsdServiceInfo serviceInfo =
+                buildServiceInfoForService(hostname, name, type, subTypeList, port, txt);
+        mHandler.post(() -> registerInternal(serviceInfo, receiver, listenerId, "service"));
     }
 
     private static NsdServiceInfo buildServiceInfoForService(
@@ -124,11 +117,8 @@
     @Override
     public void registerHost(
             String name, List<String> addresses, INsdStatusReceiver receiver, int listenerId) {
-        postRegistrationJob(
-                () -> {
-                    NsdServiceInfo serviceInfo = buildServiceInfoForHost(name, addresses);
-                    registerInternal(serviceInfo, receiver, listenerId, "host");
-                });
+        NsdServiceInfo serviceInfo = buildServiceInfoForHost(name, addresses);
+        mHandler.post(() -> registerInternal(serviceInfo, receiver, listenerId, "host"));
     }
 
     private static NsdServiceInfo buildServiceInfoForHost(
@@ -170,7 +160,7 @@
     }
 
     public void unregister(INsdStatusReceiver receiver, int listenerId) {
-        postRegistrationJob(() -> unregisterInternal(receiver, listenerId));
+        mHandler.post(() -> unregisterInternal(receiver, listenerId));
     }
 
     public void unregisterInternal(INsdStatusReceiver receiver, int listenerId) {
@@ -197,6 +187,110 @@
         mNsdManager.unregisterService(registrationListener);
     }
 
+    @Override
+    public void discoverService(String type, INsdDiscoverServiceCallback callback, int listenerId) {
+        mHandler.post(() -> discoverServiceInternal(type, callback, listenerId));
+    }
+
+    private void discoverServiceInternal(
+            String type, INsdDiscoverServiceCallback callback, int listenerId) {
+        checkOnHandlerThread();
+        Log.i(
+                TAG,
+                "Discovering services."
+                        + " Listener ID: "
+                        + listenerId
+                        + ", service type: "
+                        + type);
+
+        DiscoveryListener listener = new DiscoveryListener(listenerId, type, callback);
+        mDiscoveryListeners.append(listenerId, listener);
+        DiscoveryRequest discoveryRequest =
+                new DiscoveryRequest.Builder(type).setNetwork(null).build();
+        mNsdManager.discoverServices(discoveryRequest, mExecutor, listener);
+    }
+
+    @Override
+    public void stopServiceDiscovery(int listenerId) {
+        mHandler.post(() -> stopServiceDiscoveryInternal(listenerId));
+    }
+
+    private void stopServiceDiscoveryInternal(int listenerId) {
+        checkOnHandlerThread();
+
+        DiscoveryListener listener = mDiscoveryListeners.get(listenerId);
+        if (listener == null) {
+            Log.w(
+                    TAG,
+                    "Failed to stop service discovery. Listener ID "
+                            + listenerId
+                            + ". The listener is null.");
+            return;
+        }
+
+        Log.i(TAG, "Stopping service discovery. Listener: " + listener);
+        mNsdManager.stopServiceDiscovery(listener);
+    }
+
+    @Override
+    public void resolveService(
+            String name, String type, INsdResolveServiceCallback callback, int listenerId) {
+        mHandler.post(() -> resolveServiceInternal(name, type, callback, listenerId));
+    }
+
+    private void resolveServiceInternal(
+            String name, String type, INsdResolveServiceCallback callback, int listenerId) {
+        checkOnHandlerThread();
+
+        NsdServiceInfo serviceInfo = new NsdServiceInfo();
+        serviceInfo.setServiceName(name);
+        serviceInfo.setServiceType(type);
+        serviceInfo.setNetwork(null);
+        Log.i(
+                TAG,
+                "Resolving service."
+                        + " Listener ID: "
+                        + listenerId
+                        + ", service name: "
+                        + name
+                        + ", service type: "
+                        + type);
+
+        ServiceInfoListener listener = new ServiceInfoListener(serviceInfo, listenerId, callback);
+        mServiceInfoListeners.append(listenerId, listener);
+        mNsdManager.registerServiceInfoCallback(serviceInfo, mExecutor, listener);
+    }
+
+    @Override
+    public void stopServiceResolution(int listenerId) {
+        mHandler.post(() -> stopServiceResolutionInternal(listenerId));
+    }
+
+    private void stopServiceResolutionInternal(int listenerId) {
+        checkOnHandlerThread();
+
+        ServiceInfoListener listener = mServiceInfoListeners.get(listenerId);
+        if (listener == null) {
+            Log.w(
+                    TAG,
+                    "Failed to stop service resolution. Listener ID: "
+                            + listenerId
+                            + ". The listener is null.");
+            return;
+        }
+
+        Log.i(TAG, "Stopping service resolution. Listener: " + listener);
+
+        try {
+            mNsdManager.unregisterServiceInfoCallback(listener);
+        } catch (IllegalArgumentException e) {
+            Log.w(
+                    TAG,
+                    "Failed to stop the service resolution because it's already stopped. Listener: "
+                            + listener);
+        }
+    }
+
     private void checkOnHandlerThread() {
         if (mHandler.getLooper().getThread() != Thread.currentThread()) {
             throw new IllegalStateException(
@@ -226,7 +320,6 @@
             }
         }
         mRegistrationListeners.clear();
-        mRegistrationJobs.clear();
     }
 
     /** On ot-daemon died, reset. */
@@ -234,39 +327,6 @@
         reset();
     }
 
-    // TODO: b/323300118 - Remove this mechanism when the race condition in NsdManager is fixed.
-    /** Fetch the first job from the queue and run it. See the class doc for more details. */
-    private void peekAndRun() {
-        if (mRegistrationJobs.isEmpty()) {
-            return;
-        }
-        Runnable job = mRegistrationJobs.getFirst();
-        job.run();
-    }
-
-    // TODO: b/323300118 - Remove this mechanism when the race condition in NsdManager is fixed.
-    /**
-     * Pop the first job from the queue and run the next job. See the class doc for more details.
-     */
-    private void popAndRunNext() {
-        if (mRegistrationJobs.isEmpty()) {
-            Log.i(TAG, "No registration jobs when trying to pop and run next.");
-            return;
-        }
-        mRegistrationJobs.removeFirst();
-        peekAndRun();
-    }
-
-    private void postRegistrationJob(Runnable registrationJob) {
-        mHandler.post(
-                () -> {
-                    mRegistrationJobs.addLast(registrationJob);
-                    if (mRegistrationJobs.size() == 1) {
-                        peekAndRun();
-                    }
-                });
-    }
-
     private final class RegistrationListener implements NsdManager.RegistrationListener {
         private final NsdServiceInfo mServiceInfo;
         private final int mListenerId;
@@ -304,7 +364,6 @@
             } catch (RemoteException ignored) {
                 // do nothing if the client is dead
             }
-            popAndRunNext();
         }
 
         @Override
@@ -326,7 +385,6 @@
                     // do nothing if the client is dead
                 }
             }
-            popAndRunNext();
         }
 
         @Override
@@ -344,7 +402,6 @@
             } catch (RemoteException ignored) {
                 // do nothing if the client is dead
             }
-            popAndRunNext();
         }
 
         @Override
@@ -365,7 +422,168 @@
                 }
             }
             mRegistrationListeners.remove(mListenerId);
-            popAndRunNext();
+        }
+    }
+
+    private final class DiscoveryListener implements NsdManager.DiscoveryListener {
+        private final int mListenerId;
+        private final String mType;
+        private final INsdDiscoverServiceCallback mDiscoverServiceCallback;
+
+        DiscoveryListener(
+                int listenerId,
+                @NonNull String type,
+                @NonNull INsdDiscoverServiceCallback discoverServiceCallback) {
+            mListenerId = listenerId;
+            mType = type;
+            mDiscoverServiceCallback = discoverServiceCallback;
+        }
+
+        @Override
+        public void onStartDiscoveryFailed(String serviceType, int errorCode) {
+            Log.e(
+                    TAG,
+                    "Failed to start service discovery."
+                            + " Error code: "
+                            + errorCode
+                            + ", listener: "
+                            + this);
+            mDiscoveryListeners.remove(mListenerId);
+        }
+
+        @Override
+        public void onStopDiscoveryFailed(String serviceType, int errorCode) {
+            Log.e(
+                    TAG,
+                    "Failed to stop service discovery."
+                            + " Error code: "
+                            + errorCode
+                            + ", listener: "
+                            + this);
+            mDiscoveryListeners.remove(mListenerId);
+        }
+
+        @Override
+        public void onDiscoveryStarted(String serviceType) {
+            Log.i(TAG, "Started service discovery. Listener: " + this);
+        }
+
+        @Override
+        public void onDiscoveryStopped(String serviceType) {
+            Log.i(TAG, "Stopped service discovery. Listener: " + this);
+            mDiscoveryListeners.remove(mListenerId);
+        }
+
+        @Override
+        public void onServiceFound(NsdServiceInfo serviceInfo) {
+            Log.i(TAG, "Found service: " + serviceInfo);
+            try {
+                mDiscoverServiceCallback.onServiceDiscovered(
+                        serviceInfo.getServiceName(), mType, true);
+            } catch (RemoteException e) {
+                // do nothing if the client is dead
+            }
+        }
+
+        @Override
+        public void onServiceLost(NsdServiceInfo serviceInfo) {
+            Log.i(TAG, "Lost service: " + serviceInfo);
+            try {
+                mDiscoverServiceCallback.onServiceDiscovered(
+                        serviceInfo.getServiceName(), mType, false);
+            } catch (RemoteException e) {
+                // do nothing if the client is dead
+            }
+        }
+
+        @Override
+        public String toString() {
+            return "ID: " + mListenerId + ", type: " + mType;
+        }
+    }
+
+    private final class ServiceInfoListener implements NsdManager.ServiceInfoCallback {
+        private final String mName;
+        private final String mType;
+        private final INsdResolveServiceCallback mResolveServiceCallback;
+        private final int mListenerId;
+
+        ServiceInfoListener(
+                @NonNull NsdServiceInfo serviceInfo,
+                int listenerId,
+                @NonNull INsdResolveServiceCallback resolveServiceCallback) {
+            mName = serviceInfo.getServiceName();
+            mType = serviceInfo.getServiceType();
+            mListenerId = listenerId;
+            mResolveServiceCallback = resolveServiceCallback;
+        }
+
+        @Override
+        public void onServiceInfoCallbackRegistrationFailed(int errorCode) {
+            Log.e(
+                    TAG,
+                    "Failed to register service info callback."
+                            + " Listener ID: "
+                            + mListenerId
+                            + ", error: "
+                            + errorCode
+                            + ", service name: "
+                            + mName
+                            + ", service type: "
+                            + mType);
+        }
+
+        @Override
+        public void onServiceUpdated(@NonNull NsdServiceInfo serviceInfo) {
+            Log.i(
+                    TAG,
+                    "Service is resolved. "
+                            + " Listener ID: "
+                            + mListenerId
+                            + ", serviceInfo: "
+                            + serviceInfo);
+            List<String> addresses = new ArrayList<>();
+            for (InetAddress address : serviceInfo.getHostAddresses()) {
+                if (address instanceof Inet6Address) {
+                    addresses.add(address.getHostAddress());
+                }
+            }
+            List<DnsTxtAttribute> txtList = new ArrayList<>();
+            for (Map.Entry<String, byte[]> entry : serviceInfo.getAttributes().entrySet()) {
+                DnsTxtAttribute attribute = new DnsTxtAttribute();
+                attribute.name = entry.getKey();
+                attribute.value = Arrays.copyOf(entry.getValue(), entry.getValue().length);
+                txtList.add(attribute);
+            }
+            // TODO: b/329018320 - Use the serviceInfo.getExpirationTime to derive TTL.
+            int ttlSeconds = 10;
+            try {
+                mResolveServiceCallback.onServiceResolved(
+                        serviceInfo.getHostname(),
+                        serviceInfo.getServiceName(),
+                        serviceInfo.getServiceType(),
+                        serviceInfo.getPort(),
+                        addresses,
+                        txtList,
+                        ttlSeconds);
+
+            } catch (RemoteException e) {
+                // do nothing if the client is dead
+            }
+        }
+
+        @Override
+        public void onServiceLost() {}
+
+        @Override
+        public void onServiceInfoCallbackUnregistered() {
+            Log.i(TAG, "The service info callback is unregistered. Listener: " + this);
+            mServiceInfoListeners.remove(mListenerId);
+        }
+
+        @Override
+        public String toString() {
+            return "ID: " + mListenerId + ", service name: " + mName + ", service type: " + mType;
         }
     }
 }
diff --git a/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java b/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java
index 5f4627f..0559499 100644
--- a/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java
+++ b/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java
@@ -73,7 +73,6 @@
 import android.content.res.Resources;
 import android.net.ConnectivityManager;
 import android.net.InetAddresses;
-import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.LocalNetworkConfig;
 import android.net.LocalNetworkInfo;
@@ -106,7 +105,6 @@
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.RemoteException;
-import android.os.SystemClock;
 import android.os.UserManager;
 import android.util.Log;
 import android.util.SparseArray;
@@ -129,8 +127,6 @@
 
 import java.io.IOException;
 import java.net.Inet6Address;
-import java.net.InetAddress;
-import java.net.UnknownHostException;
 import java.nio.charset.StandardCharsets;
 import java.security.SecureRandom;
 import java.time.Instant;
@@ -165,6 +161,9 @@
     // Note that this regex allows "XX:XX-XX" as well but we don't need to be a strict checker
     private static final String OUI_REGEX = "^([0-9A-Fa-f]{2}[:-]?){2}([0-9A-Fa-f]{2})$";
 
+    // The channel mask that indicates all channels from channel 11 to channel 24
+    private static final int CHANNEL_MASK_11_TO_24 = 0x1FFF800;
+
     // Below member fields can be accessed from both the binder and handler threads
 
     private final Context mContext;
@@ -182,6 +181,7 @@
     private final NsdPublisher mNsdPublisher;
     private final OtDaemonCallbackProxy mOtDaemonCallbackProxy = new OtDaemonCallbackProxy();
     private final ConnectivityResources mResources;
+    private final Supplier<String> mCountryCodeSupplier;
 
     // This should not be directly used for calling IOtDaemon APIs because ot-daemon may die and
     // {@code mOtDaemon} will be set to {@code null}. Instead, use {@code getOtDaemon()}
@@ -215,7 +215,8 @@
             ThreadPersistentSettings persistentSettings,
             NsdPublisher nsdPublisher,
             UserManager userManager,
-            ConnectivityResources resources) {
+            ConnectivityResources resources,
+            Supplier<String> countryCodeSupplier) {
         mContext = context;
         mHandler = handler;
         mNetworkProvider = networkProvider;
@@ -230,10 +231,13 @@
         mNsdPublisher = nsdPublisher;
         mUserManager = userManager;
         mResources = resources;
+        mCountryCodeSupplier = countryCodeSupplier;
     }
 
     public static ThreadNetworkControllerService newInstance(
-            Context context, ThreadPersistentSettings persistentSettings) {
+            Context context,
+            ThreadPersistentSettings persistentSettings,
+            Supplier<String> countryCodeSupplier) {
         HandlerThread handlerThread = new HandlerThread("ThreadHandlerThread");
         handlerThread.start();
         Handler handler = new Handler(handlerThread.getLooper());
@@ -251,39 +255,8 @@
                 persistentSettings,
                 NsdPublisher.newInstance(context, handler),
                 context.getSystemService(UserManager.class),
-                new ConnectivityResources(context));
-    }
-
-    private static Inet6Address bytesToInet6Address(byte[] addressBytes) {
-        try {
-            return (Inet6Address) Inet6Address.getByAddress(addressBytes);
-        } catch (UnknownHostException e) {
-            // This is unlikely to happen unless the Thread daemon is critically broken
-            return null;
-        }
-    }
-
-    private static InetAddress addressInfoToInetAddress(Ipv6AddressInfo addressInfo) {
-        return bytesToInet6Address(addressInfo.address);
-    }
-
-    private static LinkAddress newLinkAddress(Ipv6AddressInfo addressInfo) {
-        long deprecationTimeMillis =
-                addressInfo.isPreferred
-                        ? LinkAddress.LIFETIME_PERMANENT
-                        : SystemClock.elapsedRealtime();
-
-        InetAddress address = addressInfoToInetAddress(addressInfo);
-
-        // flags and scope will be adjusted automatically depending on the address and
-        // its lifetimes.
-        return new LinkAddress(
-                address,
-                addressInfo.prefixLength,
-                0 /* flags */,
-                0 /* scope */,
-                deprecationTimeMillis,
-                LinkAddress.LIFETIME_PERMANENT /* expirationTime */);
+                new ConnectivityResources(context),
+                countryCodeSupplier);
     }
 
     private NetworkRequest newUpstreamNetworkRequest() {
@@ -346,8 +319,9 @@
                 mTunIfController.getTunFd(),
                 isEnabled(),
                 mNsdPublisher,
-                getMeshcopTxtAttributes(mResources.get()));
-        otDaemon.registerStateCallback(mOtDaemonCallbackProxy, -1);
+                getMeshcopTxtAttributes(mResources.get()),
+                mOtDaemonCallbackProxy,
+                mCountryCodeSupplier.get());
         otDaemon.asBinder().linkToDeath(() -> mHandler.post(this::onOtDaemonDied), 0);
         mOtDaemon = otDaemon;
         return mOtDaemon;
@@ -405,7 +379,10 @@
     public void initialize() {
         mHandler.post(
                 () -> {
-                    Log.d(TAG, "Initializing Thread system service...");
+                    Log.d(
+                            TAG,
+                            "Initializing Thread system service: Thread is "
+                                    + (isEnabled() ? "enabled" : "disabled"));
                     try {
                         mTunIfController.createTunInterface();
                     } catch (IOException e) {
@@ -490,6 +467,8 @@
             return;
         }
 
+        Log.i(TAG, "Set Thread enabled: " + isEnabled + ", persist: " + persist);
+
         if (persist) {
             // The persistent setting keeps the desired enabled state, thus it's set regardless
             // the otDaemon set enabled state operation succeeded or not, so that it can recover
@@ -687,6 +666,7 @@
                 new NetworkCapabilities.Builder()
                         .addTransportType(NetworkCapabilities.TRANSPORT_THREAD)
                         .addCapability(NetworkCapabilities.NET_CAPABILITY_LOCAL_NETWORK)
+                        .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_METERED)
                         .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED)
                         .build();
         final NetworkScore score =
@@ -813,6 +793,14 @@
 
     private static int selectChannel(
             int supportedChannelMask, int preferredChannelMask, Random random) {
+        // Due to radio hardware performance reasons, many Thread radio chips need to reduce their
+        // transmit power on edge channels to pass regulatory RF certification. Thread edge channel
+        // 25 and 26 are not preferred here.
+        //
+        // If users want to use channel 25 or 26, they can change the channel via the method
+        // ActiveOperationalDataset.Builder(activeOperationalDataset).setChannel(channel).build().
+        preferredChannelMask = preferredChannelMask & CHANNEL_MASK_11_TO_24;
+
         // If the preferred channel mask is not empty, select a random channel from it, otherwise
         // choose one from the supported channel mask.
         preferredChannelMask = preferredChannelMask & supportedChannelMask;
@@ -1159,20 +1147,10 @@
         }
     }
 
-    private void handleAddressChanged(Ipv6AddressInfo addressInfo, boolean isAdded) {
+    private void handleAddressChanged(List<Ipv6AddressInfo> addressInfoList) {
         checkOnHandlerThread();
-        InetAddress address = addressInfoToInetAddress(addressInfo);
-        if (address.isMulticastAddress()) {
-            Log.i(TAG, "Ignoring multicast address " + address.getHostAddress());
-            return;
-        }
 
-        LinkAddress linkAddress = newLinkAddress(addressInfo);
-        if (isAdded) {
-            mTunIfController.addAddress(linkAddress);
-        } else {
-            mTunIfController.removeAddress(linkAddress);
-        }
+        mTunIfController.updateAddresses(addressInfoList);
 
         // The OT daemon can send link property updates before the networkAgent is
         // registered
@@ -1358,8 +1336,11 @@
                 return;
             }
 
+            final int deviceRole = mState.deviceRole;
+            mState = null;
+
             // If this device is already STOPPED or DETACHED, do nothing
-            if (!ThreadNetworkController.isAttached(mState.deviceRole)) {
+            if (!ThreadNetworkController.isAttached(deviceRole)) {
                 return;
             }
 
@@ -1518,8 +1499,8 @@
         }
 
         @Override
-        public void onAddressChanged(Ipv6AddressInfo addressInfo, boolean isAdded) {
-            mHandler.post(() -> handleAddressChanged(addressInfo, isAdded));
+        public void onAddressChanged(List<Ipv6AddressInfo> addressInfoList) {
+            mHandler.post(() -> handleAddressChanged(addressInfoList));
         }
 
         @Override
diff --git a/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java b/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java
index ffa7b44..a194114 100644
--- a/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java
+++ b/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java
@@ -16,6 +16,8 @@
 
 package com.android.server.thread;
 
+import static com.android.server.thread.ThreadPersistentSettings.THREAD_COUNTRY_CODE;
+
 import android.annotation.Nullable;
 import android.annotation.StringDef;
 import android.annotation.TargetApi;
@@ -83,6 +85,7 @@
                 COUNTRY_CODE_SOURCE_TELEPHONY,
                 COUNTRY_CODE_SOURCE_TELEPHONY_LAST,
                 COUNTRY_CODE_SOURCE_WIFI,
+                COUNTRY_CODE_SOURCE_SETTINGS,
             })
     private @interface CountryCodeSource {}
 
@@ -93,6 +96,7 @@
     private static final String COUNTRY_CODE_SOURCE_TELEPHONY = "Telephony";
     private static final String COUNTRY_CODE_SOURCE_TELEPHONY_LAST = "TelephonyLast";
     private static final String COUNTRY_CODE_SOURCE_WIFI = "Wifi";
+    private static final String COUNTRY_CODE_SOURCE_SETTINGS = "Settings";
 
     private static final CountryCodeInfo DEFAULT_COUNTRY_CODE_INFO =
             new CountryCodeInfo(DEFAULT_COUNTRY_CODE, COUNTRY_CODE_SOURCE_DEFAULT);
@@ -107,6 +111,7 @@
     private final SubscriptionManager mSubscriptionManager;
     private final Map<Integer, TelephonyCountryCodeSlotInfo> mTelephonyCountryCodeSlotInfoMap =
             new ArrayMap();
+    private final ThreadPersistentSettings mPersistentSettings;
 
     @Nullable private CountryCodeInfo mCurrentCountryCodeInfo;
     @Nullable private CountryCodeInfo mLocationCountryCodeInfo;
@@ -215,7 +220,8 @@
             Context context,
             TelephonyManager telephonyManager,
             SubscriptionManager subscriptionManager,
-            @Nullable String oemCountryCode) {
+            @Nullable String oemCountryCode,
+            ThreadPersistentSettings persistentSettings) {
         mLocationManager = locationManager;
         mThreadNetworkControllerService = threadNetworkControllerService;
         mGeocoder = geocoder;
@@ -224,14 +230,19 @@
         mContext = context;
         mTelephonyManager = telephonyManager;
         mSubscriptionManager = subscriptionManager;
+        mPersistentSettings = persistentSettings;
 
         if (oemCountryCode != null) {
             mOemCountryCodeInfo = new CountryCodeInfo(oemCountryCode, COUNTRY_CODE_SOURCE_OEM);
         }
+
+        mCurrentCountryCodeInfo = pickCountryCode();
     }
 
     public static ThreadNetworkCountryCode newInstance(
-            Context context, ThreadNetworkControllerService controllerService) {
+            Context context,
+            ThreadNetworkControllerService controllerService,
+            ThreadPersistentSettings persistentSettings) {
         return new ThreadNetworkCountryCode(
                 context.getSystemService(LocationManager.class),
                 controllerService,
@@ -241,7 +252,8 @@
                 context,
                 context.getSystemService(TelephonyManager.class),
                 context.getSystemService(SubscriptionManager.class),
-                ThreadNetworkProperties.country_code().orElse(null));
+                ThreadNetworkProperties.country_code().orElse(null),
+                persistentSettings);
     }
 
     /** Sets up this country code module to listen to location country code changes. */
@@ -485,6 +497,11 @@
             return mLocationCountryCodeInfo;
         }
 
+        String settingsCountryCode = mPersistentSettings.get(THREAD_COUNTRY_CODE);
+        if (settingsCountryCode != null) {
+            return new CountryCodeInfo(settingsCountryCode, COUNTRY_CODE_SOURCE_SETTINGS);
+        }
+
         if (mOemCountryCodeInfo != null) {
             return mOemCountryCodeInfo;
         }
@@ -498,6 +515,8 @@
             public void onSuccess() {
                 synchronized ("ThreadNetworkCountryCode.this") {
                     mCurrentCountryCodeInfo = countryCodeInfo;
+                    mPersistentSettings.put(
+                            THREAD_COUNTRY_CODE.key, countryCodeInfo.getCountryCode());
                 }
             }
 
@@ -536,10 +555,9 @@
                 newOperationReceiver(countryCodeInfo));
     }
 
-    /** Returns the current country code or {@code null} if no country code is set. */
-    @Nullable
+    /** Returns the current country code. */
     public synchronized String getCountryCode() {
-        return (mCurrentCountryCodeInfo != null) ? mCurrentCountryCodeInfo.getCountryCode() : null;
+        return mCurrentCountryCodeInfo.getCountryCode();
     }
 
     /**
diff --git a/thread/service/java/com/android/server/thread/ThreadNetworkService.java b/thread/service/java/com/android/server/thread/ThreadNetworkService.java
index 37c1cf1..30c67ca 100644
--- a/thread/service/java/com/android/server/thread/ThreadNetworkService.java
+++ b/thread/service/java/com/android/server/thread/ThreadNetworkService.java
@@ -60,13 +60,16 @@
         if (phase == SystemService.PHASE_SYSTEM_SERVICES_READY) {
             mPersistentSettings.initialize();
             mControllerService =
-                    ThreadNetworkControllerService.newInstance(mContext, mPersistentSettings);
+                    ThreadNetworkControllerService.newInstance(
+                            mContext, mPersistentSettings, () -> mCountryCode.getCountryCode());
+            mCountryCode =
+                    ThreadNetworkCountryCode.newInstance(
+                            mContext, mControllerService, mPersistentSettings);
             mControllerService.initialize();
         } else if (phase == SystemService.PHASE_BOOT_COMPLETED) {
             // Country code initialization is delayed to the BOOT_COMPLETED phase because it will
             // call into Wi-Fi and Telephony service whose country code module is ready after
             // PHASE_ACTIVITY_MANAGER_READY and PHASE_THIRD_PARTY_APPS_CAN_START
-            mCountryCode = ThreadNetworkCountryCode.newInstance(mContext, mControllerService);
             mCountryCode.initialize();
             mShellCommand =
                     new ThreadNetworkShellCommand(requireNonNull(mControllerService), mCountryCode);
diff --git a/thread/service/java/com/android/server/thread/ThreadPersistentSettings.java b/thread/service/java/com/android/server/thread/ThreadPersistentSettings.java
index 5cb53fe..8aaff60 100644
--- a/thread/service/java/com/android/server/thread/ThreadPersistentSettings.java
+++ b/thread/service/java/com/android/server/thread/ThreadPersistentSettings.java
@@ -61,7 +61,10 @@
 
     /******** Thread persistent setting keys ***************/
     /** Stores the Thread feature toggle state, true for enabled and false for disabled. */
-    public static final Key<Boolean> THREAD_ENABLED = new Key<>("Thread_enabled", true);
+    public static final Key<Boolean> THREAD_ENABLED = new Key<>("thread_enabled", true);
+
+    /** Stores the Thread country code, null if no country code is stored. */
+    public static final Key<String> THREAD_COUNTRY_CODE = new Key<>("thread_country_code", null);
 
     /******** Thread persistent setting keys ***************/
 
@@ -123,7 +126,9 @@
     private <T> T getObject(String key, T defaultValue) {
         Object value;
         synchronized (mLock) {
-            if (defaultValue instanceof Boolean) {
+            if (defaultValue == null) {
+                value = mSettings.getString(key, null);
+            } else if (defaultValue instanceof Boolean) {
                 value = mSettings.getBoolean(key, (Boolean) defaultValue);
             } else if (defaultValue instanceof Integer) {
                 value = mSettings.getInt(key, (Integer) defaultValue);
diff --git a/thread/service/java/com/android/server/thread/TunInterfaceController.java b/thread/service/java/com/android/server/thread/TunInterfaceController.java
index b29a54f..97cdd55 100644
--- a/thread/service/java/com/android/server/thread/TunInterfaceController.java
+++ b/thread/service/java/com/android/server/thread/TunInterfaceController.java
@@ -29,12 +29,19 @@
 import android.system.OsConstants;
 import android.util.Log;
 
+import com.android.net.module.util.LinkPropertiesUtils.CompareResult;
 import com.android.net.module.util.netlink.NetlinkUtils;
 import com.android.net.module.util.netlink.RtNetlinkAddressMessage;
+import com.android.server.thread.openthread.Ipv6AddressInfo;
 
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.io.InterruptedIOException;
+import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.ArrayList;
+import java.util.List;
 
 /** Controller for virtual/tunnel network interfaces. */
 public class TunInterfaceController {
@@ -178,6 +185,38 @@
         }
     }
 
+    public void updateAddresses(List<Ipv6AddressInfo> addressInfoList) {
+        final List<LinkAddress> newLinkAddresses = new ArrayList<>();
+        boolean hasActiveOmrAddress = false;
+
+        for (Ipv6AddressInfo addressInfo : addressInfoList) {
+            if (addressInfo.isActiveOmr) {
+                hasActiveOmrAddress = true;
+                break;
+            }
+        }
+
+        for (Ipv6AddressInfo addressInfo : addressInfoList) {
+            InetAddress address = addressInfoToInetAddress(addressInfo);
+            if (address.isMulticastAddress()) {
+                // TODO: Logging here will create repeated logs for a single multicast address, and
+                // it currently is not mandatory for debugging. Add log for ignored multicast
+                // address when necessary.
+                continue;
+            }
+            newLinkAddresses.add(newLinkAddress(addressInfo, hasActiveOmrAddress));
+        }
+
+        final CompareResult<LinkAddress> addressDiff =
+                new CompareResult<>(mLinkProperties.getAllLinkAddresses(), newLinkAddresses);
+        for (LinkAddress linkAddress : addressDiff.removed) {
+            removeAddress(linkAddress);
+        }
+        for (LinkAddress linkAddress : addressDiff.added) {
+            addAddress(linkAddress);
+        }
+    }
+
     private RouteInfo getRouteForAddress(LinkAddress linkAddress) {
         return new RouteInfo(
                 new IpPrefix(linkAddress.getAddress(), linkAddress.getPrefixLength()),
@@ -195,4 +234,44 @@
             Log.e(TAG, "Failed to set Thread TUN interface down");
         }
     }
+
+    private static InetAddress addressInfoToInetAddress(Ipv6AddressInfo addressInfo) {
+        return bytesToInet6Address(addressInfo.address);
+    }
+
+    private static Inet6Address bytesToInet6Address(byte[] addressBytes) {
+        try {
+            return (Inet6Address) Inet6Address.getByAddress(addressBytes);
+        } catch (UnknownHostException e) {
+            // This is unlikely to happen unless the Thread daemon is critically broken
+            return null;
+        }
+    }
+
+    private static LinkAddress newLinkAddress(
+            Ipv6AddressInfo addressInfo, boolean hasActiveOmrAddress) {
+        // Mesh-local addresses and OMR address have the same scope, to distinguish them we set
+        // mesh-local addresses as deprecated when there is an active OMR address.
+        // For OMR address and link-local address we only use the value isPreferred set by
+        // ot-daemon.
+        boolean isPreferred = addressInfo.isPreferred;
+        if (addressInfo.isMeshLocal && hasActiveOmrAddress) {
+            isPreferred = false;
+        }
+
+        final long deprecationTimeMillis =
+                isPreferred ? LinkAddress.LIFETIME_PERMANENT : SystemClock.elapsedRealtime();
+
+        final InetAddress address = addressInfoToInetAddress(addressInfo);
+
+        // flags and scope will be adjusted automatically depending on the address and
+        // its lifetimes.
+        return new LinkAddress(
+                address,
+                addressInfo.prefixLength,
+                0 /* flags */,
+                0 /* scope */,
+                deprecationTimeMillis,
+                LinkAddress.LIFETIME_PERMANENT /* expirationTime */);
+    }
 }
diff --git a/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java b/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
index 0591c87..0e7f3be 100644
--- a/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
+++ b/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
@@ -17,6 +17,12 @@
 package android.net.thread.cts;
 
 import static android.Manifest.permission.ACCESS_NETWORK_STATE;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_LOCAL_NETWORK;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VPN;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED;
 import static android.net.thread.ThreadNetworkController.DEVICE_ROLE_CHILD;
 import static android.net.thread.ThreadNetworkController.DEVICE_ROLE_LEADER;
 import static android.net.thread.ThreadNetworkController.DEVICE_ROLE_ROUTER;
@@ -73,6 +79,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 
@@ -743,17 +750,6 @@
     }
 
     @Test
-    public void setEnabled_toggleAfterJoin_joinsThreadNetworkAgain() throws Exception {
-        joinRandomizedDatasetAndWait(mController);
-
-        setEnabledAndWait(mController, false);
-        assertThat(getDeviceRole(mController)).isEqualTo(DEVICE_ROLE_STOPPED);
-        setEnabledAndWait(mController, true);
-
-        runAsShell(ACCESS_NETWORK_STATE, () -> waitForAttachedState(mController));
-    }
-
-    @Test
     public void setEnabled_enableFollowedByDisable_allSucceed() throws Exception {
         joinRandomizedDatasetAndWait(mController);
         CompletableFuture<Void> setFuture1 = new CompletableFuture<>();
@@ -805,6 +801,20 @@
 
         assertThat(isAttached(mController)).isTrue();
         assertThat(networkFuture.get(NETWORK_CALLBACK_TIMEOUT_MILLIS, MILLISECONDS)).isNotNull();
+        NetworkCapabilities caps =
+                runAsShell(
+                        ACCESS_NETWORK_STATE, () -> cm.getNetworkCapabilities(networkFuture.get()));
+        assertThat(caps).isNotNull();
+        assertThat(caps.hasTransport(NetworkCapabilities.TRANSPORT_THREAD)).isTrue();
+        assertThat(caps.getCapabilities())
+                .asList()
+                .containsAtLeast(
+                        NET_CAPABILITY_LOCAL_NETWORK,
+                        NET_CAPABILITY_NOT_METERED,
+                        NET_CAPABILITY_NOT_RESTRICTED,
+                        NET_CAPABILITY_NOT_VCN_MANAGED,
+                        NET_CAPABILITY_NOT_VPN,
+                        NET_CAPABILITY_TRUSTED);
     }
 
     private void grantPermissions(String... permissions) {
@@ -837,6 +847,7 @@
     }
 
     @Test
+    @Ignore("b/333649897: Enable this when it's not flaky at all")
     public void meshcopService_joinedNetwork_discoveredHasNetwork() throws Exception {
         setUpTestNetwork();
 
@@ -864,11 +875,12 @@
     @Test
     public void meshcopService_threadDisabled_notDiscovered() throws Exception {
         setUpTestNetwork();
-
         CompletableFuture<NsdServiceInfo> serviceLostFuture = new CompletableFuture<>();
         NsdManager.DiscoveryListener listener =
                 discoverForServiceLost(MESHCOP_SERVICE_TYPE, serviceLostFuture);
+
         setEnabledAndWait(mController, false);
+
         try {
             serviceLostFuture.get(SERVICE_LOST_TIMEOUT_MILLIS, MILLISECONDS);
         } catch (InterruptedException | ExecutionException | TimeoutException ignored) {
@@ -877,7 +889,6 @@
         } finally {
             mNsdManager.stopServiceDiscovery(listener);
         }
-
         assertThrows(
                 TimeoutException.class,
                 () -> discoverService(MESHCOP_SERVICE_TYPE, SERVICE_LOST_TIMEOUT_MILLIS));
@@ -1112,7 +1123,12 @@
                         serviceInfoFuture.complete(serviceInfo);
                     }
                 };
-        mNsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD, listener);
+        mNsdManager.discoverServices(
+                serviceType,
+                NsdManager.PROTOCOL_DNS_SD,
+                mTestNetworkTracker.getNetwork(),
+                mExecutor,
+                listener);
         try {
             serviceInfoFuture.get(timeoutMilliseconds, MILLISECONDS);
         } finally {
@@ -1131,7 +1147,12 @@
                         serviceInfoFuture.complete(serviceInfo);
                     }
                 };
-        mNsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD, listener);
+        mNsdManager.discoverServices(
+                serviceType,
+                NsdManager.PROTOCOL_DNS_SD,
+                mTestNetworkTracker.getNetwork(),
+                mExecutor,
+                listener);
         return listener;
     }
 
diff --git a/thread/tests/integration/Android.bp b/thread/tests/integration/Android.bp
index 9677ec5..94985b1 100644
--- a/thread/tests/integration/Android.bp
+++ b/thread/tests/integration/Android.bp
@@ -30,6 +30,7 @@
         "net-tests-utils",
         "net-utils-device-common",
         "net-utils-device-common-bpf",
+        "net-utils-device-common-struct-base",
         "testables",
         "ThreadNetworkTestUtils",
         "truth",
diff --git a/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java b/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java
index 5fe4325..8c63d37 100644
--- a/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java
+++ b/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java
@@ -18,6 +18,7 @@
 
 import static android.Manifest.permission.MANAGE_TEST_NETWORKS;
 import static android.net.thread.utils.IntegrationTestUtils.JOIN_TIMEOUT;
+import static android.net.thread.utils.IntegrationTestUtils.getIpv6LinkAddresses;
 import static android.net.thread.utils.IntegrationTestUtils.isExpectedIcmpv6Packet;
 import static android.net.thread.utils.IntegrationTestUtils.isFromIpv6Source;
 import static android.net.thread.utils.IntegrationTestUtils.isInMulticastGroup;
@@ -33,6 +34,7 @@
 import static com.android.testutils.TestPermissionUtil.runAsShell;
 
 import static com.google.common.io.BaseEncoding.base16;
+import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -43,6 +45,8 @@
 
 import android.content.Context;
 import android.net.InetAddresses;
+import android.net.IpPrefix;
+import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.MacAddress;
 import android.net.thread.utils.FullThreadDevice;
@@ -55,6 +59,7 @@
 import android.net.thread.utils.ThreadNetworkControllerWrapper;
 import android.os.Handler;
 import android.os.HandlerThread;
+import android.os.SystemClock;
 
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.filters.LargeTest;
@@ -126,11 +131,12 @@
         mFtds = new ArrayList<>();
 
         setUpInfraNetwork();
+        mController.setEnabledAndWait(true);
         mController.joinAndWait(DEFAULT_DATASET);
 
         // Creates a infra network device.
         mInfraNetworkReader = newPacketReader(mInfraNetworkTracker.getTestIface(), mHandler);
-        startInfraDevice();
+        startInfraDeviceAndWaitForOnLinkAddr();
 
         // Create Ftds
         for (int i = 0; i < NUM_FTD; ++i) {
@@ -185,7 +191,7 @@
          * </pre>
          */
 
-        startInfraDevice();
+        startInfraDeviceAndWaitForOnLinkAddr();
         FullThreadDevice ftd = mFtds.get(0);
         startFtdChild(ftd);
 
@@ -195,6 +201,36 @@
     }
 
     @Test
+    public void unicastRouting_afterInfraNetworkSwitchInfraDevicePingThreadDeviceOmr_replyReceived()
+            throws Exception {
+        /*
+         * <pre>
+         * Topology:
+         *                 infra network                       Thread
+         * infra device -------------------- Border Router -------------- Full Thread device
+         *                                   (Cuttlefish)
+         * </pre>
+         */
+
+        FullThreadDevice ftd = mFtds.get(0);
+        startFtdChild(ftd);
+        Inet6Address ftdOmr = ftd.getOmrAddress();
+        // Create a new infra network and let Thread prefer it
+        TestNetworkTracker oldInfraNetworkTracker = mInfraNetworkTracker;
+        try {
+            setUpInfraNetwork();
+            mInfraNetworkReader = newPacketReader(mInfraNetworkTracker.getTestIface(), mHandler);
+            startInfraDeviceAndWaitForOnLinkAddr();
+
+            mInfraDevice.sendEchoRequest(ftd.getOmrAddress());
+
+            assertNotNull(pollForPacketOnInfraNetwork(ICMPV6_ECHO_REPLY_TYPE, ftdOmr));
+        } finally {
+            runAsShell(MANAGE_TEST_NETWORKS, () -> oldInfraNetworkTracker.teardown());
+        }
+    }
+
+    @Test
     public void unicastRouting_borderRouterSendsUdpToThreadDevice_datagramReceived()
             throws Exception {
         /*
@@ -221,6 +257,21 @@
     }
 
     @Test
+    public void unicastRouting_meshLocalAddressesAreNotPreferred() throws Exception {
+        // When BR is enabled, there will be OMR address, so the mesh-local addresses are expected
+        // to be deprecated.
+        List<LinkAddress> linkAddresses = getIpv6LinkAddresses("thread-wpan");
+        IpPrefix meshLocalPrefix = DEFAULT_DATASET.getMeshLocalPrefix();
+
+        for (LinkAddress address : linkAddresses) {
+            if (meshLocalPrefix.contains(address.getAddress())) {
+                assertThat(address.getDeprecationTime()).isAtMost(SystemClock.elapsedRealtime());
+                assertThat(address.isPreferred()).isFalse();
+            }
+        }
+    }
+
+    @Test
     @RequiresIpv6MulticastRouting
     public void multicastRouting_ftdSubscribedMulticastAddress_infraLinkJoinsMulticastGroup()
             throws Exception {
@@ -528,7 +579,7 @@
         tearDownInfraNetwork();
         setUpInfraNetwork();
         mInfraNetworkReader = newPacketReader(mInfraNetworkTracker.getTestIface(), mHandler);
-        startInfraDevice();
+        startInfraDeviceAndWaitForOnLinkAddr();
 
         mInfraDevice.sendEchoRequest(GROUP_ADDR_SCOPE_5);
 
@@ -555,7 +606,7 @@
         tearDownInfraNetwork();
         setUpInfraNetwork();
         mInfraNetworkReader = newPacketReader(mInfraNetworkTracker.getTestIface(), mHandler);
-        startInfraDevice();
+        startInfraDeviceAndWaitForOnLinkAddr();
 
         ftd.ping(GROUP_ADDR_SCOPE_4);
 
@@ -587,7 +638,7 @@
         assertNotNull(ftdOmr);
     }
 
-    private void startInfraDevice() throws Exception {
+    private void startInfraDeviceAndWaitForOnLinkAddr() throws Exception {
         mInfraDevice =
                 new InfraNetworkDevice(MacAddress.fromString("1:2:3:4:5:6"), mInfraNetworkReader);
         mInfraDevice.runSlaac(Duration.ofSeconds(60));
diff --git a/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java b/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java
index 39a1671..e10f134 100644
--- a/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java
+++ b/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java
@@ -17,6 +17,7 @@
 package android.net.thread;
 
 import static android.net.InetAddresses.parseNumericAddress;
+import static android.net.nsd.NsdManager.PROTOCOL_DNS_SD;
 import static android.net.thread.utils.IntegrationTestUtils.JOIN_TIMEOUT;
 import static android.net.thread.utils.IntegrationTestUtils.SERVICE_DISCOVERY_TIMEOUT;
 import static android.net.thread.utils.IntegrationTestUtils.discoverForServiceLost;
@@ -37,6 +38,7 @@
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
 import android.net.thread.utils.FullThreadDevice;
+import android.net.thread.utils.OtDaemonController;
 import android.net.thread.utils.TapTestNetworkTracker;
 import android.net.thread.utils.ThreadFeatureCheckerRule;
 import android.net.thread.utils.ThreadFeatureCheckerRule.RequiresSimulationThreadDevice;
@@ -65,6 +67,7 @@
 import java.util.Map;
 import java.util.Random;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeoutException;
 
 /** Integration test cases for Service Discovery feature. */
@@ -72,7 +75,6 @@
 @RequiresThreadFeature
 @RequiresSimulationThreadDevice
 @LargeTest
-@Ignore("TODO: b/328527773 - enable the test when it's stable")
 public class ServiceDiscoveryTest {
     private static final String TAG = ServiceDiscoveryTest.class.getSimpleName();
     private static final int NUM_FTD = 3;
@@ -96,15 +98,17 @@
     private final Context mContext = ApplicationProvider.getApplicationContext();
     private final ThreadNetworkControllerWrapper mController =
             ThreadNetworkControllerWrapper.newInstance(mContext);
-
+    private final OtDaemonController mOtCtl = new OtDaemonController();
     private HandlerThread mHandlerThread;
     private NsdManager mNsdManager;
     private TapTestNetworkTracker mTestNetworkTracker;
-    private List<FullThreadDevice> mFtds;
+    private final List<FullThreadDevice> mFtds = new ArrayList<>();
+    private final List<RegistrationListener> mRegistrationListeners = new ArrayList<>();
 
     @Before
     public void setUp() throws Exception {
-
+        mOtCtl.factoryReset();
+        mController.setEnabledAndWait(true);
         mController.joinAndWait(DEFAULT_DATASET);
         mNsdManager = mContext.getSystemService(NsdManager.class);
 
@@ -117,7 +121,6 @@
 
         // Create the FTDs in setUp() so that the FTDs can be safely released in tearDown().
         // Don't create new FTDs in test cases.
-        mFtds = new ArrayList<>();
         for (int i = 0; i < NUM_FTD; ++i) {
             FullThreadDevice ftd = new FullThreadDevice(10 + i /* node ID */);
             ftd.autoStartSrpClient();
@@ -127,6 +130,9 @@
 
     @After
     public void tearDown() throws Exception {
+        for (RegistrationListener listener : mRegistrationListeners) {
+            unregisterService(listener);
+        }
         for (FullThreadDevice ftd : mFtds) {
             // Clear registered SRP hosts and services
             if (ftd.isSrpHostRegistered()) {
@@ -259,6 +265,7 @@
     }
 
     @Test
+    @Ignore("TODO: b/333806992 - Enable when it's not flaky at all")
     public void advertisingProxy_srpClientUnregistersService_serviceIsNotDiscoverableByMdns()
             throws Exception {
         /*
@@ -314,6 +321,177 @@
         assertThat(txtMap.get("mn")).isEqualTo("Thread Border Router".getBytes(UTF_8));
     }
 
+    @Test
+    @Ignore("TODO: b/332452386 - Enable this test case when it handles the multi-client case well")
+    public void discoveryProxy_multipleClientsBrowseAndResolveServiceOverMdns() throws Exception {
+        /*
+         * <pre>
+         * Topology:
+         *                    Thread
+         *  Border Router -------------- Full Thread device
+         *  (Cuttlefish)
+         * </pre>
+         */
+
+        RegistrationListener listener = new RegistrationListener();
+        NsdServiceInfo info = new NsdServiceInfo();
+        info.setServiceType("_testservice._tcp");
+        info.setServiceName("test-service");
+        info.setPort(12345);
+        info.setHostname("testhost");
+        info.setHostAddresses(List.of(parseNumericAddress("2001::1")));
+        info.setAttribute("key1", bytes(0x01, 0x02));
+        info.setAttribute("key2", bytes(0x03));
+        registerService(info, listener);
+        mRegistrationListeners.add(listener);
+        for (int i = 0; i < NUM_FTD; ++i) {
+            FullThreadDevice ftd = mFtds.get(i);
+            ftd.joinNetwork(DEFAULT_DATASET);
+            ftd.waitForStateAnyOf(List.of("router", "child"), JOIN_TIMEOUT);
+            ftd.setDnsServerAddress(mOtCtl.getMlEid().getHostAddress());
+        }
+        final ArrayList<NsdServiceInfo> browsedServices = new ArrayList<>();
+        final ArrayList<NsdServiceInfo> resolvedServices = new ArrayList<>();
+        final ArrayList<Thread> threads = new ArrayList<>();
+        for (int i = 0; i < NUM_FTD; ++i) {
+            browsedServices.add(null);
+            resolvedServices.add(null);
+        }
+        for (int i = 0; i < NUM_FTD; ++i) {
+            final FullThreadDevice ftd = mFtds.get(i);
+            final int index = i;
+            Runnable task =
+                    () -> {
+                        browsedServices.set(
+                                index,
+                                ftd.browseService("_testservice._tcp.default.service.arpa."));
+                        resolvedServices.set(
+                                index,
+                                ftd.resolveService(
+                                        "test-service", "_testservice._tcp.default.service.arpa."));
+                    };
+            threads.add(new Thread(task));
+        }
+        for (Thread thread : threads) {
+            thread.start();
+        }
+        for (Thread thread : threads) {
+            thread.join();
+        }
+
+        for (int i = 0; i < NUM_FTD; ++i) {
+            NsdServiceInfo browsedService = browsedServices.get(i);
+            assertThat(browsedService.getServiceName()).isEqualTo("test-service");
+            assertThat(browsedService.getPort()).isEqualTo(12345);
+
+            NsdServiceInfo resolvedService = resolvedServices.get(i);
+            assertThat(resolvedService.getServiceName()).isEqualTo("test-service");
+            assertThat(resolvedService.getPort()).isEqualTo(12345);
+            assertThat(resolvedService.getHostname()).isEqualTo("testhost.default.service.arpa.");
+            assertThat(resolvedService.getHostAddresses())
+                    .containsExactly(parseNumericAddress("2001::1"));
+            assertThat(resolvedService.getAttributes())
+                    .comparingValuesUsing(BYTE_ARRAY_EQUALITY)
+                    .containsExactly("key1", bytes(0x01, 0x02), "key2", bytes(3));
+        }
+    }
+
+    @Test
+    public void discoveryProxy_browseAndResolveServiceAtSrpServer() throws Exception {
+        /*
+         * <pre>
+         * Topology:
+         *                    Thread
+         *  Border Router -------+------ SRP client
+         *  (Cuttlefish)         |
+         *                       +------ DNS client
+         *
+         * </pre>
+         */
+        FullThreadDevice srpClient = mFtds.get(0);
+        srpClient.joinNetwork(DEFAULT_DATASET);
+        srpClient.waitForStateAnyOf(List.of("router", "child"), JOIN_TIMEOUT);
+        srpClient.setSrpHostname("my-host");
+        srpClient.setSrpHostAddresses(List.of((Inet6Address) parseNumericAddress("2001::1")));
+        srpClient.addSrpService(
+                "my-service",
+                "_test._udp",
+                List.of("_sub1"),
+                12345 /* port */,
+                Map.of("key1", bytes(0x01, 0x02), "key2", bytes(0x03)));
+
+        FullThreadDevice dnsClient = mFtds.get(1);
+        dnsClient.joinNetwork(DEFAULT_DATASET);
+        dnsClient.waitForStateAnyOf(List.of("router", "child"), JOIN_TIMEOUT);
+        dnsClient.setDnsServerAddress(mOtCtl.getMlEid().getHostAddress());
+
+        NsdServiceInfo browsedService = dnsClient.browseService("_test._udp.default.service.arpa.");
+        assertThat(browsedService.getServiceName()).isEqualTo("my-service");
+        assertThat(browsedService.getPort()).isEqualTo(12345);
+        assertThat(browsedService.getHostname()).isEqualTo("my-host.default.service.arpa.");
+        assertThat(browsedService.getHostAddresses())
+                .containsExactly(parseNumericAddress("2001::1"));
+        assertThat(browsedService.getAttributes())
+                .comparingValuesUsing(BYTE_ARRAY_EQUALITY)
+                .containsExactly("key1", bytes(0x01, 0x02), "key2", bytes(3));
+
+        NsdServiceInfo resolvedService =
+                dnsClient.resolveService("my-service", "_test._udp.default.service.arpa.");
+        assertThat(resolvedService.getServiceName()).isEqualTo("my-service");
+        assertThat(resolvedService.getPort()).isEqualTo(12345);
+        assertThat(resolvedService.getHostname()).isEqualTo("my-host.default.service.arpa.");
+        assertThat(resolvedService.getHostAddresses())
+                .containsExactly(parseNumericAddress("2001::1"));
+        assertThat(resolvedService.getAttributes())
+                .comparingValuesUsing(BYTE_ARRAY_EQUALITY)
+                .containsExactly("key1", bytes(0x01, 0x02), "key2", bytes(3));
+    }
+
+    private void registerService(NsdServiceInfo serviceInfo, RegistrationListener listener)
+            throws InterruptedException, ExecutionException, TimeoutException {
+        mNsdManager.registerService(serviceInfo, PROTOCOL_DNS_SD, listener);
+        listener.waitForRegistered();
+    }
+
+    private void unregisterService(RegistrationListener listener)
+            throws InterruptedException, ExecutionException, TimeoutException {
+        mNsdManager.unregisterService(listener);
+        listener.waitForUnregistered();
+    }
+
+    private static class RegistrationListener implements NsdManager.RegistrationListener {
+        private final CompletableFuture<Void> mRegisteredFuture = new CompletableFuture<>();
+        private final CompletableFuture<Void> mUnRegisteredFuture = new CompletableFuture<>();
+
+        RegistrationListener() {}
+
+        @Override
+        public void onRegistrationFailed(NsdServiceInfo serviceInfo, int errorCode) {}
+
+        @Override
+        public void onUnregistrationFailed(NsdServiceInfo serviceInfo, int errorCode) {}
+
+        @Override
+        public void onServiceRegistered(NsdServiceInfo serviceInfo) {
+            mRegisteredFuture.complete(null);
+        }
+
+        @Override
+        public void onServiceUnregistered(NsdServiceInfo serviceInfo) {
+            mUnRegisteredFuture.complete(null);
+        }
+
+        public void waitForRegistered()
+                throws InterruptedException, ExecutionException, TimeoutException {
+            mRegisteredFuture.get(SERVICE_DISCOVERY_TIMEOUT.toMillis(), MILLISECONDS);
+        }
+
+        public void waitForUnregistered()
+                throws InterruptedException, ExecutionException, TimeoutException {
+            mUnRegisteredFuture.get(SERVICE_DISCOVERY_TIMEOUT.toMillis(), MILLISECONDS);
+        }
+    }
+
     private static byte[] bytes(int... byteInts) {
         byte[] bytes = new byte[byteInts.length];
         for (int i = 0; i < byteInts.length; ++i) {
diff --git a/thread/tests/integration/src/android/net/thread/ThreadIntegrationTest.java b/thread/tests/integration/src/android/net/thread/ThreadIntegrationTest.java
index 4a006cf..1410d41 100644
--- a/thread/tests/integration/src/android/net/thread/ThreadIntegrationTest.java
+++ b/thread/tests/integration/src/android/net/thread/ThreadIntegrationTest.java
@@ -21,13 +21,19 @@
 import static android.net.thread.ThreadNetworkController.DEVICE_ROLE_STOPPED;
 import static android.net.thread.utils.IntegrationTestUtils.CALLBACK_TIMEOUT;
 import static android.net.thread.utils.IntegrationTestUtils.RESTART_JOIN_TIMEOUT;
+import static android.net.thread.utils.IntegrationTestUtils.getIpv6LinkAddresses;
+import static android.net.thread.utils.IntegrationTestUtils.waitFor;
 
 import static com.android.compatibility.common.util.SystemUtil.runShellCommand;
+import static com.android.compatibility.common.util.SystemUtil.runShellCommandOrThrow;
 
 import static com.google.common.io.BaseEncoding.base16;
 import static com.google.common.truth.Truth.assertThat;
 
 import android.content.Context;
+import android.net.IpPrefix;
+import android.net.LinkAddress;
+import android.net.thread.utils.FullThreadDevice;
 import android.net.thread.utils.OtDaemonController;
 import android.net.thread.utils.ThreadFeatureCheckerRule;
 import android.net.thread.utils.ThreadFeatureCheckerRule.RequiresThreadFeature;
@@ -44,14 +50,28 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.IOException;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
 import java.net.Inet6Address;
+import java.net.InetAddress;
+import java.time.Duration;
+import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 /** Tests for E2E Android Thread integration with ot-daemon, ConnectivityService, etc.. */
 @LargeTest
 @RequiresThreadFeature
 @RunWith(AndroidJUnit4.class)
 public class ThreadIntegrationTest {
+    // The byte[] buffer size for UDP tests
+    private static final int UDP_BUFFER_SIZE = 1024;
+
+    // The maximum time for OT addresses to be propagated to the TUN interface "thread-wpan"
+    private static final Duration TUN_ADDR_UPDATE_TIMEOUT = Duration.ofSeconds(1);
+
     // A valid Thread Active Operational Dataset generated from OpenThread CLI "dataset init new".
     private static final byte[] DEFAULT_DATASET_TLVS =
             base16().decode(
@@ -65,24 +85,32 @@
 
     @Rule public final ThreadFeatureCheckerRule mThreadRule = new ThreadFeatureCheckerRule();
 
+    private ExecutorService mExecutor;
     private final Context mContext = ApplicationProvider.getApplicationContext();
     private final ThreadNetworkControllerWrapper mController =
             ThreadNetworkControllerWrapper.newInstance(mContext);
     private OtDaemonController mOtCtl;
+    private FullThreadDevice mFtd;
 
     @Before
     public void setUp() throws Exception {
+        mExecutor = Executors.newSingleThreadExecutor();
         mOtCtl = new OtDaemonController();
         mController.leaveAndWait();
 
         // TODO: b/323301831 - This is a workaround to avoid unnecessary delay to re-form a network
         mOtCtl.factoryReset();
+
+        mFtd = new FullThreadDevice(10 /* nodeId */);
     }
 
     @After
     public void tearDown() throws Exception {
         mController.setTestNetworkAsUpstreamAndWait(null);
         mController.leaveAndWait();
+
+        mFtd.destroy();
+        mExecutor.shutdownNow();
     }
 
     @Test
@@ -97,13 +125,16 @@
     }
 
     @Test
-    public void otDaemonRestart_JoinedNetworkAndStopped_autoRejoined() throws Exception {
+    public void otDaemonRestart_JoinedNetworkAndStopped_autoRejoinedAndTunIfStateConsistent()
+            throws Exception {
         mController.joinAndWait(DEFAULT_DATASET);
 
         runShellCommand("stop ot-daemon");
 
         mController.waitForRole(DEVICE_ROLE_DETACHED, CALLBACK_TIMEOUT);
         mController.waitForRole(DEVICE_ROLE_LEADER, RESTART_JOIN_TIMEOUT);
+        assertThat(mOtCtl.isInterfaceUp()).isTrue();
+        assertThat(runShellCommand("ifconfig thread-wpan")).contains("UP POINTOPOINT RUNNING");
     }
 
     @Test
@@ -120,23 +151,112 @@
         mController.joinAndWait(DEFAULT_DATASET);
 
         mOtCtl.factoryReset();
-        String ifconfig = runShellCommand("ifconfig thread-wpan");
 
+        String ifconfig = runShellCommand("ifconfig thread-wpan");
         assertThat(ifconfig).doesNotContain("inet6 addr");
     }
 
+    // TODO (b/323300829): add test for removing an OT address
     @Test
     public void tunInterface_joinedNetwork_otAddressesAddedToTunInterface() throws Exception {
         mController.joinAndWait(DEFAULT_DATASET);
 
-        String ifconfig = runShellCommand("ifconfig thread-wpan");
         List<Inet6Address> otAddresses = mOtCtl.getAddresses();
         assertThat(otAddresses).isNotEmpty();
-        for (Inet6Address otAddress : otAddresses) {
-            assertThat(ifconfig).contains(otAddress.getHostAddress());
+        // TODO: it's cleaner to have a retry() method to retry failed asserts in given delay so
+        // that we can write assertThat() in the Predicate
+        waitFor(
+                () -> {
+                    String ifconfig = runShellCommand("ifconfig thread-wpan");
+                    return otAddresses.stream()
+                            .allMatch(addr -> ifconfig.contains(addr.getHostAddress()));
+                },
+                TUN_ADDR_UPDATE_TIMEOUT);
+    }
+
+    @Test
+    public void otDaemonRestart_latestCountryCodeIsSetToOtDaemon() throws Exception {
+        runThreadCommand("force-country-code enabled CN");
+
+        runShellCommand("stop ot-daemon");
+        // TODO(b/323331973): the sleep is needed to workaround the race conditions
+        SystemClock.sleep(200);
+        mController.waitForRole(DEVICE_ROLE_STOPPED, CALLBACK_TIMEOUT);
+
+        assertThat(mOtCtl.getCountryCode()).isEqualTo("CN");
+    }
+
+    @Test
+    public void udp_appStartEchoServer_endDeviceUdpEchoSuccess() throws Exception {
+        // Topology:
+        //   Test App ------ thread-wpan ------ End Device
+
+        mController.joinAndWait(DEFAULT_DATASET);
+        startFtdChild(mFtd, DEFAULT_DATASET);
+        final Inet6Address serverAddress = mOtCtl.getMeshLocalAddresses().get(0);
+        final int serverPort = 9527;
+
+        mExecutor.execute(() -> startUdpEchoServerAndWait(serverAddress, serverPort));
+        mFtd.udpOpen();
+        mFtd.udpSend("Hello,Thread", serverAddress, serverPort);
+        String udpReply = mFtd.udpReceive();
+
+        assertThat(udpReply).isEqualTo("Hello,Thread");
+    }
+
+    @Test
+    public void joinNetworkWithBrDisabled_meshLocalAddressesArePreferred() throws Exception {
+        // When BR feature is disabled, there is no OMR address, so the mesh-local addresses are
+        // expected to be preferred.
+        mOtCtl.executeCommand("br disable");
+        mController.joinAndWait(DEFAULT_DATASET);
+
+        IpPrefix meshLocalPrefix = DEFAULT_DATASET.getMeshLocalPrefix();
+        List<LinkAddress> linkAddresses = getIpv6LinkAddresses("thread-wpan");
+        for (LinkAddress address : linkAddresses) {
+            if (meshLocalPrefix.contains(address.getAddress())) {
+                assertThat(address.getDeprecationTime())
+                        .isGreaterThan(SystemClock.elapsedRealtime());
+                assertThat(address.isPreferred()).isTrue();
+            }
         }
+
+        mOtCtl.executeCommand("br enable");
     }
 
     // TODO (b/323300829): add more tests for integration with linux platform and
     // ConnectivityService
+
+    private static String runThreadCommand(String cmd) {
+        return runShellCommandOrThrow("cmd thread_network " + cmd);
+    }
+
+    private void startFtdChild(FullThreadDevice ftd, ActiveOperationalDataset activeDataset)
+            throws Exception {
+        ftd.factoryReset();
+        ftd.joinNetwork(activeDataset);
+        ftd.waitForStateAnyOf(List.of("router", "child"), Duration.ofSeconds(8));
+    }
+
+    /**
+     * Starts a UDP echo server and replies to the first UDP message.
+     *
+     * <p>This method exits when the first UDP message is received and the reply is sent
+     */
+    private void startUdpEchoServerAndWait(InetAddress serverAddress, int serverPort) {
+        try (var udpServerSocket = new DatagramSocket(serverPort, serverAddress)) {
+            DatagramPacket recvPacket =
+                    new DatagramPacket(new byte[UDP_BUFFER_SIZE], UDP_BUFFER_SIZE);
+            udpServerSocket.receive(recvPacket);
+            byte[] sendBuffer = Arrays.copyOf(recvPacket.getData(), recvPacket.getData().length);
+            udpServerSocket.send(
+                    new DatagramPacket(
+                            sendBuffer,
+                            sendBuffer.length,
+                            (Inet6Address) recvPacket.getAddress(),
+                            recvPacket.getPort()));
+        } catch (IOException e) {
+            throw new IllegalStateException(e);
+        }
+    }
 }
diff --git a/thread/tests/integration/src/android/net/thread/ThreadNetworkControllerTest.java b/thread/tests/integration/src/android/net/thread/ThreadNetworkControllerTest.java
index 496ec9f..ba04348 100644
--- a/thread/tests/integration/src/android/net/thread/ThreadNetworkControllerTest.java
+++ b/thread/tests/integration/src/android/net/thread/ThreadNetworkControllerTest.java
@@ -27,6 +27,7 @@
 import static org.junit.Assert.assertThrows;
 
 import android.content.Context;
+import android.net.thread.utils.ThreadFeatureCheckerRule;
 import android.net.thread.utils.ThreadFeatureCheckerRule.RequiresThreadFeature;
 import android.os.OutcomeReceiver;
 import android.util.SparseIntArray;
@@ -37,6 +38,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -64,6 +66,8 @@
                 }
             };
 
+    @Rule public final ThreadFeatureCheckerRule mThreadRule = new ThreadFeatureCheckerRule();
+
     private final Context mContext = ApplicationProvider.getApplicationContext();
     private ExecutorService mExecutor;
     private ThreadNetworkController mController;
diff --git a/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java b/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
index 600b662..5e70f6c 100644
--- a/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
+++ b/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
@@ -20,10 +20,9 @@
 
 import static com.google.common.io.BaseEncoding.base16;
 
-import static org.junit.Assert.fail;
-
 import android.net.InetAddresses;
 import android.net.IpPrefix;
+import android.net.nsd.NsdServiceInfo;
 import android.net.thread.ActiveOperationalDataset;
 
 import com.google.errorprone.annotations.FormatMethod;
@@ -34,6 +33,7 @@
 import java.io.InputStreamReader;
 import java.io.OutputStreamWriter;
 import java.net.Inet6Address;
+import java.net.InetAddress;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -216,6 +216,11 @@
         return matcher.group(4);
     }
 
+    /** Sends a UDP message to given IPv6 address and port. */
+    public void udpSend(String message, Inet6Address serverAddr, int serverPort) {
+        executeCommand("udp send %s %d %s", serverAddr.getHostAddress(), serverPort, message);
+    }
+
     /** Enables the SRP client and run in autostart mode. */
     public void autoStartSrpClient() {
         executeCommand("srp client autostart enable");
@@ -327,6 +332,55 @@
         return false;
     }
 
+    /** Sets the DNS server address. */
+    public void setDnsServerAddress(String address) {
+        executeCommand("dns config " + address);
+    }
+
+    /** Returns the first browsed service instance of {@code serviceType}. */
+    public NsdServiceInfo browseService(String serviceType) {
+        // CLI output:
+        // DNS browse response for _testservice._tcp.default.service.arpa.
+        // test-service
+        //    Port:12345, Priority:0, Weight:0, TTL:10
+        //    Host:testhost.default.service.arpa.
+        //    HostAddress:2001:0:0:0:0:0:0:1 TTL:10
+        //    TXT:[key1=0102, key2=03] TTL:10
+
+        List<String> lines = executeCommand("dns browse " + serviceType);
+        NsdServiceInfo info = new NsdServiceInfo();
+        info.setServiceName(lines.get(1));
+        info.setServiceType(serviceType);
+        info.setPort(DnsServiceCliOutputParser.parsePort(lines.get(2)));
+        info.setHostname(DnsServiceCliOutputParser.parseHostname(lines.get(3)));
+        info.setHostAddresses(List.of(DnsServiceCliOutputParser.parseHostAddress(lines.get(4))));
+        DnsServiceCliOutputParser.parseTxtIntoServiceInfo(lines.get(5), info);
+
+        return info;
+    }
+
+    /** Returns the resolved service instance. */
+    public NsdServiceInfo resolveService(String serviceName, String serviceType) {
+        // CLI output:
+        // DNS service resolution response for test-service for service
+        // _test._tcp.default.service.arpa.
+        // Port:12345, Priority:0, Weight:0, TTL:10
+        // Host:Android.default.service.arpa.
+        // HostAddress:2001:0:0:0:0:0:0:1 TTL:10
+        // TXT:[key1=0102, key2=03] TTL:10
+
+        List<String> lines = executeCommand("dns service %s %s", serviceName, serviceType);
+        NsdServiceInfo info = new NsdServiceInfo();
+        info.setServiceName(serviceName);
+        info.setServiceType(serviceType);
+        info.setPort(DnsServiceCliOutputParser.parsePort(lines.get(1)));
+        info.setHostname(DnsServiceCliOutputParser.parseHostname(lines.get(2)));
+        info.setHostAddresses(List.of(DnsServiceCliOutputParser.parseHostAddress(lines.get(3))));
+        DnsServiceCliOutputParser.parseTxtIntoServiceInfo(lines.get(4), info);
+
+        return info;
+    }
+
     /** Runs the "factoryreset" command on the device. */
     public void factoryReset() {
         try {
@@ -423,7 +477,7 @@
                 break;
             }
             if (line.startsWith("Error")) {
-                fail("ot-cli-ftd reported an error: " + line);
+                throw new IOException("ot-cli-ftd reported an error: " + line);
             }
             if (!line.startsWith("> ")) {
                 result.add(line);
@@ -454,4 +508,45 @@
     private static String toHexString(byte[] bytes) {
         return base16().encode(bytes);
     }
+
+    private static final class DnsServiceCliOutputParser {
+        /** Returns the first match in the input of a given regex pattern. */
+        private static Matcher firstMatchOf(String input, String regex) {
+            Matcher matcher = Pattern.compile(regex).matcher(input);
+            matcher.find();
+            return matcher;
+        }
+
+        // Example: "Port:12345"
+        private static int parsePort(String line) {
+            return Integer.parseInt(firstMatchOf(line, "Port:(\\d+)").group(1));
+        }
+
+        // Example: "Host:Android.default.service.arpa."
+        private static String parseHostname(String line) {
+            return firstMatchOf(line, "Host:(.+)").group(1);
+        }
+
+        // Example: "HostAddress:2001:0:0:0:0:0:0:1"
+        private static InetAddress parseHostAddress(String line) {
+            return InetAddresses.parseNumericAddress(
+                    firstMatchOf(line, "HostAddress:([^ ]+)").group(1));
+        }
+
+        // Example: "TXT:[key1=0102, key2=03]"
+        private static void parseTxtIntoServiceInfo(String line, NsdServiceInfo serviceInfo) {
+            String txtString = firstMatchOf(line, "TXT:\\[([^\\]]+)\\]").group(1);
+            for (String txtEntry : txtString.split(",")) {
+                String[] nameAndValue = txtEntry.trim().split("=");
+                String name = nameAndValue[0];
+                String value = nameAndValue[1];
+                byte[] bytes = new byte[value.length() / 2];
+                for (int i = 0; i < value.length(); i += 2) {
+                    byte b = (byte) ((value.charAt(i) - '0') << 4 | (value.charAt(i + 1) - '0'));
+                    bytes[i / 2] = b;
+                }
+                serviceInfo.setAttribute(name, bytes);
+            }
+        }
+    }
 }
diff --git a/thread/tests/integration/src/android/net/thread/utils/IntegrationTestUtils.java b/thread/tests/integration/src/android/net/thread/utils/IntegrationTestUtils.java
index 2237e65..9be9566 100644
--- a/thread/tests/integration/src/android/net/thread/utils/IntegrationTestUtils.java
+++ b/thread/tests/integration/src/android/net/thread/utils/IntegrationTestUtils.java
@@ -25,6 +25,8 @@
 
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
+import android.net.InetAddresses;
+import android.net.LinkAddress;
 import android.net.TestNetworkInterface;
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
@@ -85,7 +87,7 @@
      */
     public static void waitFor(Supplier<Boolean> condition, Duration timeout)
             throws TimeoutException {
-        final long intervalMills = 1000;
+        final long intervalMills = 500;
         final long timeoutMills = timeout.toMillis();
 
         for (long i = 0; i < timeoutMills; i += intervalMills) {
@@ -291,6 +293,20 @@
         return false;
     }
 
+    public static List<LinkAddress> getIpv6LinkAddresses(String interfaceName) throws IOException {
+        List<LinkAddress> addresses = new ArrayList<>();
+        final String cmd = " ip -6 addr show dev " + interfaceName;
+        final String output = runShellCommandOrThrow(cmd);
+
+        for (final String line : output.split("\\n")) {
+            if (line.contains("inet6")) {
+                addresses.add(parseAddressLine(line));
+            }
+        }
+
+        return addresses;
+    }
+
     /** Return the first discovered service of {@code serviceType}. */
     public static NsdServiceInfo discoverService(NsdManager nsdManager, String serviceType)
             throws Exception {
@@ -392,4 +408,29 @@
         @Override
         public void onServiceInfoCallbackUnregistered() {}
     }
+
+    /**
+     * Parses a line of output from "ip -6 addr show" into a {@link LinkAddress}.
+     *
+     * <p>Example line: "inet6 2001:db8:1:1::1/64 scope global deprecated"
+     */
+    private static LinkAddress parseAddressLine(String line) {
+        String[] parts = line.trim().split("\\s+");
+        String addressString = parts[1];
+        String[] pieces = addressString.split("/", 2);
+        int prefixLength = Integer.parseInt(pieces[1]);
+        final InetAddress address = InetAddresses.parseNumericAddress(pieces[0]);
+        long deprecationTimeMillis =
+                line.contains("deprecated")
+                        ? SystemClock.elapsedRealtime()
+                        : LinkAddress.LIFETIME_PERMANENT;
+
+        return new LinkAddress(
+                address,
+                prefixLength,
+                0 /* flags */,
+                0 /* scope */,
+                deprecationTimeMillis,
+                LinkAddress.LIFETIME_PERMANENT /* expirationTime */);
+    }
 }
diff --git a/thread/tests/integration/src/android/net/thread/utils/OtDaemonController.java b/thread/tests/integration/src/android/net/thread/utils/OtDaemonController.java
index 4a06fe8..b3175fd 100644
--- a/thread/tests/integration/src/android/net/thread/utils/OtDaemonController.java
+++ b/thread/tests/integration/src/android/net/thread/utils/OtDaemonController.java
@@ -16,13 +16,16 @@
 
 package android.net.thread.utils;
 
+import android.annotation.Nullable;
 import android.net.InetAddresses;
+import android.net.IpPrefix;
 import android.os.SystemClock;
 
 import com.android.compatibility.common.util.SystemUtil;
 
 import java.net.Inet6Address;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -53,16 +56,68 @@
 
     /** Returns the list of IPv6 addresses on ot-daemon. */
     public List<Inet6Address> getAddresses() {
-        String output = executeCommand("ipaddr");
-        return Arrays.asList(output.split("\n")).stream()
-                .map(String::trim)
-                .filter(str -> !str.equals("Done"))
+        return executeCommandAndParse("ipaddr").stream()
                 .map(addr -> InetAddresses.parseNumericAddress(addr))
                 .map(inetAddr -> (Inet6Address) inetAddr)
                 .toList();
     }
 
+    /** Returns {@code true} if the Thread interface is up. */
+    public boolean isInterfaceUp() {
+        String output = executeCommand("ifconfig");
+        return output.contains("up");
+    }
+
+    /** Returns the ML-EID of the device. */
+    public Inet6Address getMlEid() {
+        String addressStr = executeCommandAndParse("ipaddr mleid").get(0);
+        return (Inet6Address) InetAddresses.parseNumericAddress(addressStr);
+    }
+
+    /** Returns the country code on ot-daemon. */
+    public String getCountryCode() {
+        return executeCommandAndParse("region").get(0);
+    }
+
+    /**
+     * Returns the list of IPv6 Mesh-Local addresses on ot-daemon.
+     *
+     * <p>The return List can be empty if no Mesh-Local prefix exists.
+     */
+    public List<Inet6Address> getMeshLocalAddresses() {
+        IpPrefix meshLocalPrefix = getMeshLocalPrefix();
+        if (meshLocalPrefix == null) {
+            return Collections.emptyList();
+        }
+        return getAddresses().stream().filter(addr -> meshLocalPrefix.contains(addr)).toList();
+    }
+
+    /**
+     * Returns the Mesh-Local prefix or {@code null} if none exists (e.g. the Active Dataset is not
+     * set).
+     */
+    @Nullable
+    public IpPrefix getMeshLocalPrefix() {
+        List<IpPrefix> prefixes =
+                executeCommandAndParse("prefix meshlocal").stream()
+                        .map(prefix -> new IpPrefix(prefix))
+                        .toList();
+        return prefixes.isEmpty() ? null : prefixes.get(0);
+    }
+
     public String executeCommand(String cmd) {
         return SystemUtil.runShellCommand(OT_CTL + " " + cmd);
     }
+
+    /**
+     * Executes a ot-ctl command and parse the output to a list of strings.
+     *
+     * <p>The trailing "Done" in the command output will be dropped.
+     */
+    public List<String> executeCommandAndParse(String cmd) {
+        return Arrays.asList(executeCommand(cmd).split("\n")).stream()
+                .map(String::trim)
+                .filter(str -> !str.equals("Done"))
+                .toList();
+    }
 }
diff --git a/thread/tests/unit/src/com/android/server/thread/NsdPublisherTest.java b/thread/tests/unit/src/com/android/server/thread/NsdPublisherTest.java
index d860166..8886c73 100644
--- a/thread/tests/unit/src/com/android/server/thread/NsdPublisherTest.java
+++ b/thread/tests/unit/src/com/android/server/thread/NsdPublisherTest.java
@@ -23,6 +23,7 @@
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.spy;
@@ -30,24 +31,30 @@
 import static org.mockito.Mockito.verify;
 
 import android.net.InetAddresses;
+import android.net.nsd.DiscoveryRequest;
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
 import android.os.Handler;
 import android.os.test.TestLooper;
 
 import com.android.server.thread.openthread.DnsTxtAttribute;
+import com.android.server.thread.openthread.INsdDiscoverServiceCallback;
+import com.android.server.thread.openthread.INsdResolveServiceCallback;
 import com.android.server.thread.openthread.INsdStatusReceiver;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
+import org.mockito.ArgumentMatcher;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.net.InetAddress;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executor;
 
@@ -57,6 +64,8 @@
 
     @Mock private INsdStatusReceiver mRegistrationReceiver;
     @Mock private INsdStatusReceiver mUnregistrationReceiver;
+    @Mock private INsdDiscoverServiceCallback mDiscoverServiceCallback;
+    @Mock private INsdResolveServiceCallback mResolveServiceCallback;
 
     private TestLooper mTestLooper;
     private NsdPublisher mNsdPublisher;
@@ -469,6 +478,165 @@
     }
 
     @Test
+    public void discoverService_serviceDiscovered() throws Exception {
+        prepareTest();
+
+        mNsdPublisher.discoverService("_test._tcp", mDiscoverServiceCallback, 10 /* listenerId */);
+        mTestLooper.dispatchAll();
+        ArgumentCaptor<NsdManager.DiscoveryListener> discoveryListenerArgumentCaptor =
+                ArgumentCaptor.forClass(NsdManager.DiscoveryListener.class);
+        verify(mMockNsdManager, times(1))
+                .discoverServices(
+                        eq(new DiscoveryRequest.Builder(PROTOCOL_DNS_SD, "_test._tcp").build()),
+                        any(Executor.class),
+                        discoveryListenerArgumentCaptor.capture());
+        NsdManager.DiscoveryListener actualDiscoveryListener =
+                discoveryListenerArgumentCaptor.getValue();
+        NsdServiceInfo serviceInfo = new NsdServiceInfo();
+        serviceInfo.setServiceName("test");
+        serviceInfo.setServiceType(null);
+        actualDiscoveryListener.onServiceFound(serviceInfo);
+        mTestLooper.dispatchAll();
+
+        verify(mDiscoverServiceCallback, times(1))
+                .onServiceDiscovered("test", "_test._tcp", true /* isFound */);
+    }
+
+    @Test
+    public void discoverService_serviceLost() throws Exception {
+        prepareTest();
+
+        mNsdPublisher.discoverService("_test._tcp", mDiscoverServiceCallback, 10 /* listenerId */);
+        mTestLooper.dispatchAll();
+        ArgumentCaptor<NsdManager.DiscoveryListener> discoveryListenerArgumentCaptor =
+                ArgumentCaptor.forClass(NsdManager.DiscoveryListener.class);
+        verify(mMockNsdManager, times(1))
+                .discoverServices(
+                        eq(new DiscoveryRequest.Builder(PROTOCOL_DNS_SD, "_test._tcp").build()),
+                        any(Executor.class),
+                        discoveryListenerArgumentCaptor.capture());
+        NsdManager.DiscoveryListener actualDiscoveryListener =
+                discoveryListenerArgumentCaptor.getValue();
+        NsdServiceInfo serviceInfo = new NsdServiceInfo();
+        serviceInfo.setServiceName("test");
+        serviceInfo.setServiceType(null);
+        actualDiscoveryListener.onServiceLost(serviceInfo);
+        mTestLooper.dispatchAll();
+
+        verify(mDiscoverServiceCallback, times(1))
+                .onServiceDiscovered("test", "_test._tcp", false /* isFound */);
+    }
+
+    @Test
+    public void stopServiceDiscovery() {
+        prepareTest();
+
+        mNsdPublisher.discoverService("_test._tcp", mDiscoverServiceCallback, 10 /* listenerId */);
+        mTestLooper.dispatchAll();
+        ArgumentCaptor<NsdManager.DiscoveryListener> discoveryListenerArgumentCaptor =
+                ArgumentCaptor.forClass(NsdManager.DiscoveryListener.class);
+        verify(mMockNsdManager, times(1))
+                .discoverServices(
+                        eq(new DiscoveryRequest.Builder(PROTOCOL_DNS_SD, "_test._tcp").build()),
+                        any(Executor.class),
+                        discoveryListenerArgumentCaptor.capture());
+        NsdManager.DiscoveryListener actualDiscoveryListener =
+                discoveryListenerArgumentCaptor.getValue();
+        NsdServiceInfo serviceInfo = new NsdServiceInfo();
+        serviceInfo.setServiceName("test");
+        serviceInfo.setServiceType(null);
+        actualDiscoveryListener.onServiceFound(serviceInfo);
+        mNsdPublisher.stopServiceDiscovery(10 /* listenerId */);
+        mTestLooper.dispatchAll();
+
+        verify(mMockNsdManager, times(1)).stopServiceDiscovery(actualDiscoveryListener);
+    }
+
+    @Test
+    public void resolveService_serviceResolved() throws Exception {
+        prepareTest();
+
+        mNsdPublisher.resolveService(
+                "test", "_test._tcp", mResolveServiceCallback, 10 /* listenerId */);
+        mTestLooper.dispatchAll();
+        ArgumentCaptor<NsdServiceInfo> serviceInfoArgumentCaptor =
+                ArgumentCaptor.forClass(NsdServiceInfo.class);
+        ArgumentCaptor<NsdManager.ServiceInfoCallback> serviceInfoCallbackArgumentCaptor =
+                ArgumentCaptor.forClass(NsdManager.ServiceInfoCallback.class);
+        verify(mMockNsdManager, times(1))
+                .registerServiceInfoCallback(
+                        serviceInfoArgumentCaptor.capture(),
+                        any(Executor.class),
+                        serviceInfoCallbackArgumentCaptor.capture());
+        assertThat(serviceInfoArgumentCaptor.getValue().getServiceName()).isEqualTo("test");
+        assertThat(serviceInfoArgumentCaptor.getValue().getServiceType()).isEqualTo("_test._tcp");
+        NsdServiceInfo serviceInfo = new NsdServiceInfo();
+        serviceInfo.setServiceName("test");
+        serviceInfo.setServiceType("_test._tcp");
+        serviceInfo.setPort(12345);
+        serviceInfo.setHostname("test-host");
+        serviceInfo.setHostAddresses(
+                List.of(
+                        InetAddress.parseNumericAddress("2001::1"),
+                        InetAddress.parseNumericAddress("2001::2")));
+        serviceInfo.setAttribute("key1", new byte[] {(byte) 0x01, (byte) 0x02});
+        serviceInfo.setAttribute("key2", new byte[] {(byte) 0x03});
+        serviceInfoCallbackArgumentCaptor.getValue().onServiceUpdated(serviceInfo);
+        mTestLooper.dispatchAll();
+
+        verify(mResolveServiceCallback, times(1))
+                .onServiceResolved(
+                        eq("test-host"),
+                        eq("test"),
+                        eq("_test._tcp"),
+                        eq(12345),
+                        eq(List.of("2001::1", "2001::2")),
+                        argThat(
+                                new TxtMatcher(
+                                        List.of(
+                                                makeTxtAttribute("key1", List.of(0x01, 0x02)),
+                                                makeTxtAttribute("key2", List.of(0x03))))),
+                        anyInt());
+    }
+
+    @Test
+    public void stopServiceResolution() throws Exception {
+        prepareTest();
+
+        mNsdPublisher.resolveService(
+                "test", "_test._tcp", mResolveServiceCallback, 10 /* listenerId */);
+        mTestLooper.dispatchAll();
+        ArgumentCaptor<NsdServiceInfo> serviceInfoArgumentCaptor =
+                ArgumentCaptor.forClass(NsdServiceInfo.class);
+        ArgumentCaptor<NsdManager.ServiceInfoCallback> serviceInfoCallbackArgumentCaptor =
+                ArgumentCaptor.forClass(NsdManager.ServiceInfoCallback.class);
+        verify(mMockNsdManager, times(1))
+                .registerServiceInfoCallback(
+                        serviceInfoArgumentCaptor.capture(),
+                        any(Executor.class),
+                        serviceInfoCallbackArgumentCaptor.capture());
+        assertThat(serviceInfoArgumentCaptor.getValue().getServiceName()).isEqualTo("test");
+        assertThat(serviceInfoArgumentCaptor.getValue().getServiceType()).isEqualTo("_test._tcp");
+        NsdServiceInfo serviceInfo = new NsdServiceInfo();
+        serviceInfo.setServiceName("test");
+        serviceInfo.setServiceType("_test._tcp");
+        serviceInfo.setPort(12345);
+        serviceInfo.setHostname("test-host");
+        serviceInfo.setHostAddresses(
+                List.of(
+                        InetAddress.parseNumericAddress("2001::1"),
+                        InetAddress.parseNumericAddress("2001::2")));
+        serviceInfo.setAttribute("key1", new byte[] {(byte) 0x01, (byte) 0x02});
+        serviceInfo.setAttribute("key2", new byte[] {(byte) 0x03});
+        serviceInfoCallbackArgumentCaptor.getValue().onServiceUpdated(serviceInfo);
+        mNsdPublisher.stopServiceResolution(10 /* listenerId */);
+        mTestLooper.dispatchAll();
+
+        verify(mMockNsdManager, times(1))
+                .unregisterServiceInfoCallback(serviceInfoCallbackArgumentCaptor.getValue());
+    }
+
+    @Test
     public void reset_unregisterAll() {
         prepareTest();
 
@@ -582,6 +750,30 @@
         return addresses;
     }
 
+    private static class TxtMatcher implements ArgumentMatcher<List<DnsTxtAttribute>> {
+        private final List<DnsTxtAttribute> mAttributes;
+
+        TxtMatcher(List<DnsTxtAttribute> attributes) {
+            mAttributes = attributes;
+        }
+
+        @Override
+        public boolean matches(List<DnsTxtAttribute> argument) {
+            if (argument.size() != mAttributes.size()) {
+                return false;
+            }
+            for (int i = 0; i < argument.size(); ++i) {
+                if (!Objects.equals(argument.get(i).name, mAttributes.get(i).name)) {
+                    return false;
+                }
+                if (!Arrays.equals(argument.get(i).value, mAttributes.get(i).value)) {
+                    return false;
+                }
+            }
+            return true;
+        }
+    }
+
     // @Before and @Test run in different threads. NsdPublisher requires the jobs are run on the
     // thread looper, so TestLooper needs to be created inside each test case to install the
     // correct looper.
diff --git a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkControllerServiceTest.java b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkControllerServiceTest.java
index 151ed5b..493058f 100644
--- a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkControllerServiceTest.java
+++ b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkControllerServiceTest.java
@@ -25,6 +25,7 @@
 import static android.net.thread.ThreadNetworkManager.DISALLOW_THREAD_NETWORK;
 import static android.net.thread.ThreadNetworkManager.PERMISSION_THREAD_NETWORK_PRIVILEGED;
 
+import static com.android.server.thread.ThreadNetworkCountryCode.DEFAULT_COUNTRY_CODE;
 import static com.android.server.thread.openthread.IOtDaemon.ErrorCode.OT_ERROR_INVALID_STATE;
 
 import static com.google.common.io.BaseEncoding.base16;
@@ -78,7 +79,9 @@
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
+import org.mockito.InOrder;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 import java.util.concurrent.CompletableFuture;
@@ -112,7 +115,11 @@
     private static final String DEFAULT_NETWORK_NAME = "thread-wpan0";
     private static final int OT_ERROR_NONE = 0;
     private static final int DEFAULT_SUPPORTED_CHANNEL_MASK = 0x07FFF800; // from channel 11 to 26
-    private static final int DEFAULT_PREFERRED_CHANNEL_MASK = 0x00000800; // channel 11
+
+    // The DEFAULT_PREFERRED_CHANNEL_MASK is the ot-daemon preferred channel mask. Channel 25 and
+    // 26 are not preferred by dataset. The ThreadNetworkControllerService will only select channel
+    // 11 when it creates randomized dataset.
+    private static final int DEFAULT_PREFERRED_CHANNEL_MASK = 0x06000800; // channel 11, 25 and 26
     private static final int DEFAULT_SELECTED_CHANNEL = 11;
     private static final byte[] DEFAULT_SUPPORTED_CHANNEL_MASK_ARRAY = base16().decode("001FFFE0");
 
@@ -180,7 +187,8 @@
                         mMockPersistentSettings,
                         mMockNsdPublisher,
                         mMockUserManager,
-                        mConnectivityResources);
+                        mConnectivityResources,
+                        () -> DEFAULT_COUNTRY_CODE);
         mService.setTestNetworkAgent(mMockNetworkAgent);
     }
 
@@ -489,4 +497,23 @@
         assertThat(mFakeOtDaemon.isInitialized()).isTrue();
         verify(mockJoinReceiver, times(1)).onSuccess();
     }
+
+    @Test
+    public void onOtDaemonDied_joinedNetwork_interfaceStateBackToUp() throws Exception {
+        mService.initialize();
+        final IOperationReceiver mockReceiver = mock(IOperationReceiver.class);
+        mService.join(DEFAULT_ACTIVE_DATASET, mockReceiver);
+        mTestLooper.dispatchAll();
+        mTestLooper.moveTimeForward(FakeOtDaemon.JOIN_DELAY.toMillis() + 100);
+        mTestLooper.dispatchAll();
+
+        Mockito.reset(mMockInfraIfController);
+        mFakeOtDaemon.terminate();
+        mTestLooper.dispatchAll();
+
+        verify(mMockTunIfController, times(1)).onOtDaemonDied();
+        InOrder inOrder = Mockito.inOrder(mMockTunIfController);
+        inOrder.verify(mMockTunIfController, times(1)).setInterfaceUp(false);
+        inOrder.verify(mMockTunIfController, times(1)).setInterfaceUp(true);
+    }
 }
diff --git a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
index 5ca6511..ca9741d 100644
--- a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
+++ b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
@@ -19,6 +19,7 @@
 import static android.net.thread.ThreadNetworkException.ERROR_INTERNAL_ERROR;
 
 import static com.android.server.thread.ThreadNetworkCountryCode.DEFAULT_COUNTRY_CODE;
+import static com.android.server.thread.ThreadPersistentSettings.THREAD_COUNTRY_CODE;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -104,6 +105,7 @@
     @Mock List<SubscriptionInfo> mSubscriptionInfoList;
     @Mock SubscriptionInfo mSubscriptionInfo0;
     @Mock SubscriptionInfo mSubscriptionInfo1;
+    @Mock ThreadPersistentSettings mPersistentSettings;
 
     private ThreadNetworkCountryCode mThreadNetworkCountryCode;
     private boolean mErrorSetCountryCode;
@@ -164,7 +166,8 @@
                 mContext,
                 mTelephonyManager,
                 mSubscriptionManager,
-                oemCountryCode);
+                oemCountryCode,
+                mPersistentSettings);
     }
 
     private static Address newAddress(String countryCode) {
@@ -450,6 +453,14 @@
     }
 
     @Test
+    public void settingsCountryCode_settingsCountryCodeIsActive_settingsCountryCodeIsUsed() {
+        when(mPersistentSettings.get(THREAD_COUNTRY_CODE)).thenReturn(TEST_COUNTRY_CODE_CN);
+        mThreadNetworkCountryCode.initialize();
+
+        assertThat(mThreadNetworkCountryCode.getCountryCode()).isEqualTo(TEST_COUNTRY_CODE_CN);
+    }
+
+    @Test
     public void dump_allCountryCodeInfoAreDumped() {
         StringWriter stringWriter = new StringWriter();
         PrintWriter printWriter = new PrintWriter(stringWriter);
diff --git a/thread/tests/unit/src/com/android/server/thread/ThreadPersistentSettingsTest.java b/thread/tests/unit/src/com/android/server/thread/ThreadPersistentSettingsTest.java
index 9406a2f..7d2fe91 100644
--- a/thread/tests/unit/src/com/android/server/thread/ThreadPersistentSettingsTest.java
+++ b/thread/tests/unit/src/com/android/server/thread/ThreadPersistentSettingsTest.java
@@ -16,6 +16,7 @@
 
 package com.android.server.thread;
 
+import static com.android.server.thread.ThreadPersistentSettings.THREAD_COUNTRY_CODE;
 import static com.android.server.thread.ThreadPersistentSettings.THREAD_ENABLED;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -54,6 +55,8 @@
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class ThreadPersistentSettingsTest {
+    private static final String TEST_COUNTRY_CODE = "CN";
+
     @Mock private AtomicFile mAtomicFile;
     @Mock Resources mResources;
     @Mock ConnectivityResources mConnectivityResources;
@@ -131,6 +134,28 @@
         verify(mAtomicFile).finishWrite(any());
     }
 
+    @Test
+    public void put_ThreadCountryCodeString_returnsString() throws Exception {
+        mThreadPersistentSettings.put(THREAD_COUNTRY_CODE.key, TEST_COUNTRY_CODE);
+
+        assertThat(mThreadPersistentSettings.get(THREAD_COUNTRY_CODE)).isEqualTo(TEST_COUNTRY_CODE);
+
+        // Confirm that file writes have been triggered.
+        verify(mAtomicFile).startWrite();
+        verify(mAtomicFile).finishWrite(any());
+    }
+
+    @Test
+    public void put_ThreadCountryCodeNull_returnsNull() throws Exception {
+        mThreadPersistentSettings.put(THREAD_COUNTRY_CODE.key, null);
+
+        assertThat(mThreadPersistentSettings.get(THREAD_COUNTRY_CODE)).isNull();
+
+        // Confirm that file writes have been triggered.
+        verify(mAtomicFile).startWrite();
+        verify(mAtomicFile).finishWrite(any());
+    }
+
     private byte[] createXmlForParsing(String key, Boolean value) throws Exception {
         PersistableBundle bundle = new PersistableBundle();
         ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
diff --git a/thread/tests/utils/Android.bp b/thread/tests/utils/Android.bp
index 24e9bb9..726ec9d 100644
--- a/thread/tests/utils/Android.bp
+++ b/thread/tests/utils/Android.bp
@@ -27,6 +27,7 @@
         "net-tests-utils",
         "net-utils-device-common",
         "net-utils-device-common-bpf",
+        "net-utils-device-common-struct-base",
     ],
     srcs: [
         "src/**/*.java",
diff --git a/thread/tests/utils/src/android/net/thread/utils/TapTestNetworkTracker.java b/thread/tests/utils/src/android/net/thread/utils/TapTestNetworkTracker.java
index 43f177d..b586a19 100644
--- a/thread/tests/utils/src/android/net/thread/utils/TapTestNetworkTracker.java
+++ b/thread/tests/utils/src/android/net/thread/utils/TapTestNetworkTracker.java
@@ -62,6 +62,7 @@
     private final Looper mLooper;
     private TestNetworkInterface mInterface;
     private TestableNetworkAgent mAgent;
+    private Network mNetwork;
     private final TestableNetworkCallback mNetworkCallback;
     private final ConnectivityManager mConnectivityManager;
 
@@ -91,6 +92,11 @@
         return mInterface.getInterfaceName();
     }
 
+    /** Returns the {@link android.net.Network} of the test network. */
+    public Network getNetwork() {
+        return mNetwork;
+    }
+
     private void setUpTestNetwork() throws Exception {
         mInterface = mContext.getSystemService(TestNetworkManager.class).createTapInterface();
 
@@ -105,13 +111,13 @@
                         newNetworkCapabilities(),
                         lp,
                         new NetworkAgentConfig.Builder().build());
-        final Network network = mAgent.register();
+        mNetwork = mAgent.register();
         mAgent.markConnected();
 
         PollingCheck.check(
                 "No usable address on interface",
                 TIMEOUT.toMillis(),
-                () -> hasUsableAddress(network, getInterfaceName()));
+                () -> hasUsableAddress(mNetwork, getInterfaceName()));
 
         lp.setLinkAddresses(makeLinkAddresses());
         mAgent.sendLinkProperties(lp);