diff --git a/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java b/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java
index 53c80ae..13a7a22 100644
--- a/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java
+++ b/Tethering/src/com/android/networkstack/tethering/OffloadHardwareInterface.java
@@ -16,6 +16,7 @@
 
 package com.android.networkstack.tethering;
 
+import static com.android.net.module.util.netlink.NetlinkUtils.SOCKET_RECV_BUFSIZE;
 import static com.android.net.module.util.netlink.StructNlMsgHdr.NLM_F_DUMP;
 import static com.android.net.module.util.netlink.StructNlMsgHdr.NLM_F_REQUEST;
 
@@ -198,7 +199,8 @@
         public NativeHandle createConntrackSocket(final int groups) {
             final FileDescriptor fd;
             try {
-                fd = NetlinkUtils.netlinkSocketForProto(OsConstants.NETLINK_NETFILTER);
+                fd = NetlinkUtils.netlinkSocketForProto(OsConstants.NETLINK_NETFILTER,
+                        SOCKET_RECV_BUFSIZE);
             } catch (ErrnoException e) {
                 mLog.e("Unable to create conntrack socket " + e);
                 return null;
diff --git a/Tethering/src/com/android/networkstack/tethering/Tethering.java b/Tethering/src/com/android/networkstack/tethering/Tethering.java
index 552b105..873961a 100644
--- a/Tethering/src/com/android/networkstack/tethering/Tethering.java
+++ b/Tethering/src/com/android/networkstack/tethering/Tethering.java
@@ -368,6 +368,7 @@
 
         // Load tethering configuration.
         updateConfiguration();
+        mConfig.readEnableSyncSM(mContext);
         // It is OK for the configuration to be passed to the PrivateAddressCoordinator at
         // construction time because the only part of the configuration it uses is
         // shouldEnableWifiP2pDedicatedIp(), and currently do not support changing that.
diff --git a/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java b/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java
index 0678525..298940e 100644
--- a/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java
+++ b/Tethering/src/com/android/networkstack/tethering/TetheringConfiguration.java
@@ -132,6 +132,9 @@
 
     public static final String TETHER_FORCE_RANDOM_PREFIX_BASE_SELECTION =
             "tether_force_random_prefix_base_selection";
+
+    public static final String TETHER_ENABLE_SYNC_SM = "tether_enable_sync_sm";
+
     /**
      * Default value that used to periodic polls tether offload stats from tethering offload HAL
      * to make the data warnings work.
@@ -139,7 +142,7 @@
     public static final int DEFAULT_TETHER_OFFLOAD_POLL_INTERVAL_MS = 5000;
 
     /** A flag for using synchronous or asynchronous state machine. */
-    public static final boolean USE_SYNC_SM = false;
+    public static boolean USE_SYNC_SM = false;
 
     public final String[] tetherableUsbRegexs;
     public final String[] tetherableWifiRegexs;
@@ -385,6 +388,16 @@
         return mRandomPrefixBase;
     }
 
+    /**
+     * Check whether sync SM is enabled then set it to USE_SYNC_SM. This should be called once
+     * when tethering is created. Otherwise if the flag is pushed while tethering is enabled,
+     * then it's possible for some IpServer(s) running the new sync state machine while others
+     * use the async state machine.
+     */
+    public void readEnableSyncSM(final Context ctx) {
+        USE_SYNC_SM = mDeps.isFeatureEnabled(ctx, TETHER_ENABLE_SYNC_SM);
+    }
+
     /** Does the dumping.*/
     public void dump(PrintWriter pw) {
         pw.print("activeDataSubId: ");
@@ -438,6 +451,9 @@
 
         pw.print("mRandomPrefixBase: ");
         pw.println(mRandomPrefixBase);
+
+        pw.print("USE_SYNC_SM: ");
+        pw.println(USE_SYNC_SM);
     }
 
     /** Returns the string representation of this object.*/
diff --git a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
index c232697..2933a44 100644
--- a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
+++ b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
@@ -865,12 +865,12 @@
         // trigger the reselection, the total test time may over test suite 1 minmute timeout.
         // Probably need to disable/restore all internet networks in a common place of test
         // process. Currently, EthernetTetheringTest is part of CTS test which needs wifi network
-        // connection if device has wifi feature. CtsNetUtils#toggleWifi() checks wifi connection
-        // during the toggling process.
-        // See Tethering#chooseUpstreamType, CtsNetUtils#toggleWifi.
+        // connection if device has wifi feature.
+        // See Tethering#chooseUpstreamType
         // TODO: toggle cellular network if the device has no WIFI feature.
         Log.d(TAG, "Toggle WIFI to retry upstream selection");
-        sCtsNetUtils.toggleWifi();
+        sCtsNetUtils.disableWifi();
+        sCtsNetUtils.ensureWifiConnected();
 
         // Wait for expected upstream.
         final CompletableFuture<Network> future = new CompletableFuture<>();
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringConfigurationTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringConfigurationTest.java
index aa322dc..dd51c7a 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringConfigurationTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringConfigurationTest.java
@@ -754,4 +754,27 @@
                 new TetheringConfiguration(mMockContext, mLog, INVALID_SUBSCRIPTION_ID, mDeps);
         assertEquals(p2pLeasesSubnetPrefixLength, p2pCfg.getP2pLeasesSubnetPrefixLength());
     }
+
+    private void setTetherEnableSyncSMFlagEnabled(Boolean enabled) {
+        mDeps.setFeatureEnabled(TetheringConfiguration.TETHER_ENABLE_SYNC_SM, enabled);
+        new TetheringConfiguration(
+                mMockContext, mLog, INVALID_SUBSCRIPTION_ID, mDeps).readEnableSyncSM(mMockContext);
+    }
+
+    private void assertEnableSyncSM(boolean value) {
+        assertEquals(value, TetheringConfiguration.USE_SYNC_SM);
+    }
+
+    @Test
+    public void testEnableSyncSMFlag() throws Exception {
+        // Test default disabled
+        setTetherEnableSyncSMFlagEnabled(null);
+        assertEnableSyncSM(false);
+
+        setTetherEnableSyncSMFlagEnabled(true);
+        assertEnableSyncSM(true);
+
+        setTetherEnableSyncSMFlagEnabled(false);
+        assertEnableSyncSM(false);
+    }
 }
diff --git a/framework-t/api/current.txt b/framework-t/api/current.txt
index 60a88c0..7cd3d4f 100644
--- a/framework-t/api/current.txt
+++ b/framework-t/api/current.txt
@@ -210,9 +210,26 @@
 
 package android.net.nsd {
 
+  @FlaggedApi("com.android.net.flags.nsd_subtypes_support_enabled") public final class DiscoveryRequest implements android.os.Parcelable {
+    method public int describeContents();
+    method @Nullable public android.net.Network getNetwork();
+    method @NonNull public String getServiceType();
+    method @Nullable public String getSubtype();
+    method public void writeToParcel(@NonNull android.os.Parcel, int);
+    field @NonNull public static final android.os.Parcelable.Creator<android.net.nsd.DiscoveryRequest> CREATOR;
+  }
+
+  public static final class DiscoveryRequest.Builder {
+    ctor public DiscoveryRequest.Builder(@NonNull String);
+    method @NonNull public android.net.nsd.DiscoveryRequest build();
+    method @NonNull public android.net.nsd.DiscoveryRequest.Builder setNetwork(@Nullable android.net.Network);
+    method @NonNull public android.net.nsd.DiscoveryRequest.Builder setSubtype(@Nullable String);
+  }
+
   public final class NsdManager {
     method public void discoverServices(String, int, android.net.nsd.NsdManager.DiscoveryListener);
     method public void discoverServices(@NonNull String, int, @Nullable android.net.Network, @NonNull java.util.concurrent.Executor, @NonNull android.net.nsd.NsdManager.DiscoveryListener);
+    method @FlaggedApi("com.android.net.flags.nsd_subtypes_support_enabled") public void discoverServices(@NonNull android.net.nsd.DiscoveryRequest, @NonNull java.util.concurrent.Executor, @NonNull android.net.nsd.NsdManager.DiscoveryListener);
     method @RequiresPermission(android.Manifest.permission.ACCESS_NETWORK_STATE) public void discoverServices(@NonNull String, int, @NonNull android.net.NetworkRequest, @NonNull java.util.concurrent.Executor, @NonNull android.net.nsd.NsdManager.DiscoveryListener);
     method public void registerService(android.net.nsd.NsdServiceInfo, int, android.net.nsd.NsdManager.RegistrationListener);
     method public void registerService(@NonNull android.net.nsd.NsdServiceInfo, int, @NonNull java.util.concurrent.Executor, @NonNull android.net.nsd.NsdManager.RegistrationListener);
diff --git a/framework-t/api/system-current.txt b/framework-t/api/system-current.txt
index d346af3..8251f85 100644
--- a/framework-t/api/system-current.txt
+++ b/framework-t/api/system-current.txt
@@ -500,6 +500,7 @@
     method @RequiresPermission(allOf={android.Manifest.permission.ACCESS_NETWORK_STATE, "android.permission.THREAD_NETWORK_PRIVILEGED"}) public void registerOperationalDatasetCallback(@NonNull java.util.concurrent.Executor, @NonNull android.net.thread.ThreadNetworkController.OperationalDatasetCallback);
     method @RequiresPermission(android.Manifest.permission.ACCESS_NETWORK_STATE) public void registerStateCallback(@NonNull java.util.concurrent.Executor, @NonNull android.net.thread.ThreadNetworkController.StateCallback);
     method @RequiresPermission("android.permission.THREAD_NETWORK_PRIVILEGED") public void scheduleMigration(@NonNull android.net.thread.PendingOperationalDataset, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<java.lang.Void,android.net.thread.ThreadNetworkException>);
+    method @RequiresPermission("android.permission.THREAD_NETWORK_PRIVILEGED") public void setEnabled(boolean, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<java.lang.Void,android.net.thread.ThreadNetworkException>);
     method @RequiresPermission(allOf={android.Manifest.permission.ACCESS_NETWORK_STATE, "android.permission.THREAD_NETWORK_PRIVILEGED"}) public void unregisterOperationalDatasetCallback(@NonNull android.net.thread.ThreadNetworkController.OperationalDatasetCallback);
     method @RequiresPermission(android.Manifest.permission.ACCESS_NETWORK_STATE) public void unregisterStateCallback(@NonNull android.net.thread.ThreadNetworkController.StateCallback);
     field public static final int DEVICE_ROLE_CHILD = 2; // 0x2
@@ -507,6 +508,9 @@
     field public static final int DEVICE_ROLE_LEADER = 4; // 0x4
     field public static final int DEVICE_ROLE_ROUTER = 3; // 0x3
     field public static final int DEVICE_ROLE_STOPPED = 0; // 0x0
+    field public static final int STATE_DISABLED = 0; // 0x0
+    field public static final int STATE_DISABLING = 2; // 0x2
+    field public static final int STATE_ENABLED = 1; // 0x1
     field public static final int THREAD_VERSION_1_3 = 4; // 0x4
   }
 
@@ -518,6 +522,7 @@
   public static interface ThreadNetworkController.StateCallback {
     method public void onDeviceRoleChanged(int);
     method public default void onPartitionIdChanged(long);
+    method public default void onThreadEnableStateChanged(int);
   }
 
   @FlaggedApi("com.android.net.thread.flags.thread_enabled") public class ThreadNetworkException extends java.lang.Exception {
@@ -530,6 +535,7 @@
     field public static final int ERROR_REJECTED_BY_PEER = 8; // 0x8
     field public static final int ERROR_RESOURCE_EXHAUSTED = 10; // 0xa
     field public static final int ERROR_RESPONSE_BAD_FORMAT = 9; // 0x9
+    field public static final int ERROR_THREAD_DISABLED = 12; // 0xc
     field public static final int ERROR_TIMEOUT = 3; // 0x3
     field public static final int ERROR_UNAVAILABLE = 4; // 0x4
     field public static final int ERROR_UNKNOWN = 11; // 0xb
diff --git a/framework-t/src/android/app/usage/NetworkStatsManager.java b/framework-t/src/android/app/usage/NetworkStatsManager.java
index d139544..7fa0661 100644
--- a/framework-t/src/android/app/usage/NetworkStatsManager.java
+++ b/framework-t/src/android/app/usage/NetworkStatsManager.java
@@ -510,6 +510,27 @@
      * Query network usage statistics details for a given uid.
      * This may take a long time, and apps should avoid calling this on their main thread.
      *
+     * @param networkType As defined in {@link ConnectivityManager}, e.g.
+     *            {@link ConnectivityManager#TYPE_MOBILE}, {@link ConnectivityManager#TYPE_WIFI}
+     *            etc.
+     * @param subscriberId If applicable, the subscriber id of the network interface.
+     *                     <p>Starting with API level 29, the {@code subscriberId} is guarded by
+     *                     additional restrictions. Calling apps that do not meet the new
+     *                     requirements to access the {@code subscriberId} can provide a {@code
+     *                     null} value when querying for the mobile network type to receive usage
+     *                     for all mobile networks. For additional details see {@link
+     *                     TelephonyManager#getSubscriberId()}.
+     *                     <p>Starting with API level 31, calling apps can provide a
+     *                     {@code subscriberId} with wifi network type to receive usage for
+     *                     wifi networks which is under the given subscription if applicable.
+     *                     Otherwise, pass {@code null} when querying all wifi networks.
+     * @param startTime Start of period. Defined in terms of "Unix time", see
+     *            {@link java.lang.System#currentTimeMillis}.
+     * @param endTime End of period. Defined in terms of "Unix time", see
+     *            {@link java.lang.System#currentTimeMillis}.
+     * @param uid UID of app
+     * @return Statistics which is described above.
+     * @throws SecurityException if permissions are insufficient to read network statistics.
      * @see #queryDetailsForUidTagState(int, String, long, long, int, int, int)
      */
     @NonNull
diff --git a/framework-t/src/android/net/nsd/DiscoveryRequest.java b/framework-t/src/android/net/nsd/DiscoveryRequest.java
new file mode 100644
index 0000000..b0b71ea
--- /dev/null
+++ b/framework-t/src/android/net/nsd/DiscoveryRequest.java
@@ -0,0 +1,247 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.nsd;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.net.Network;
+import android.os.Parcel;
+import android.os.Parcelable;
+import android.text.TextUtils;
+
+import java.util.Objects;
+
+/**
+ * Encapsulates parameters for {@link NsdManager#discoverServices}.
+ */
+@FlaggedApi(NsdManager.Flags.NSD_SUBTYPES_SUPPORT_ENABLED)
+public final class DiscoveryRequest implements Parcelable {
+    private final int mProtocolType;
+
+    @NonNull
+    private final String mServiceType;
+
+    @Nullable
+    private final String mSubtype;
+
+    @Nullable
+    private final Network mNetwork;
+
+    // TODO: add mDiscoveryConfig for more fine-grained discovery behavior control
+
+    @NonNull
+    public static final Creator<DiscoveryRequest> CREATOR =
+            new Creator<>() {
+                @Override
+                public DiscoveryRequest createFromParcel(Parcel in) {
+                    int protocolType = in.readInt();
+                    String serviceType = in.readString();
+                    String subtype = in.readString();
+                    Network network =
+                            in.readParcelable(Network.class.getClassLoader(), Network.class);
+                    return new DiscoveryRequest(protocolType, serviceType, subtype, network);
+                }
+
+                @Override
+                public DiscoveryRequest[] newArray(int size) {
+                    return new DiscoveryRequest[size];
+                }
+            };
+
+    private DiscoveryRequest(int protocolType, @NonNull String serviceType,
+            @Nullable String subtype, @Nullable Network network) {
+        mProtocolType = protocolType;
+        mServiceType = serviceType;
+        mSubtype = subtype;
+        mNetwork = network;
+    }
+
+    /**
+     * Returns the service type in format of dot-joint string of two labels.
+     *
+     * For example, "_ipp._tcp" for internet printer and "_matter._tcp" for <a
+     * href="https://csa-iot.org/all-solutions/matter">Matter</a> operational device.
+     */
+    @NonNull
+    public String getServiceType() {
+        return mServiceType;
+    }
+
+    /**
+     * Returns the subtype without the trailing "._sub" label or {@code null} if no subtype is
+     * specified.
+     *
+     * For example, the return value will be "_printer" for subtype "_printer._sub".
+     */
+    @Nullable
+    public String getSubtype() {
+        return mSubtype;
+    }
+
+    /**
+     * Returns the service discovery protocol.
+     *
+     * @hide
+     */
+    public int getProtocolType() {
+        return mProtocolType;
+    }
+
+    /**
+     * Returns the {@link Network} on which the query should be sent or {@code null} if no
+     * network is specified.
+     */
+    @Nullable
+    public Network getNetwork() {
+        return mNetwork;
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append(", protocolType: ").append(mProtocolType)
+            .append(", serviceType: ").append(mServiceType)
+            .append(", subtype: ").append(mSubtype)
+            .append(", network: ").append(mNetwork);
+        return sb.toString();
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        } else if (!(other instanceof DiscoveryRequest)) {
+            return false;
+        } else {
+            DiscoveryRequest otherRequest = (DiscoveryRequest) other;
+            return mProtocolType == otherRequest.mProtocolType
+                    && Objects.equals(mServiceType, otherRequest.mServiceType)
+                    && Objects.equals(mSubtype, otherRequest.mSubtype)
+                    && Objects.equals(mNetwork, otherRequest.mNetwork);
+        }
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(mProtocolType, mServiceType, mSubtype, mNetwork);
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        dest.writeInt(mProtocolType);
+        dest.writeString(mServiceType);
+        dest.writeString(mSubtype);
+        dest.writeParcelable(mNetwork, flags);
+    }
+
+    /** The builder for creating new {@link DiscoveryRequest} objects. */
+    public static final class Builder {
+        private final int mProtocolType;
+
+        @NonNull
+        private String mServiceType;
+
+        @Nullable
+        private String mSubtype;
+
+        @Nullable
+        private Network mNetwork;
+
+        /**
+         * Creates a new default {@link Builder} object with given service type.
+         *
+         * @throws IllegalArgumentException if {@code serviceType} is {@code null} or an empty
+         * string
+         */
+        public Builder(@NonNull String serviceType) {
+            this(NsdManager.PROTOCOL_DNS_SD, serviceType);
+        }
+
+        /** @hide */
+        public Builder(int protocolType, @NonNull String serviceType) {
+            NsdManager.checkProtocol(protocolType);
+            mProtocolType = protocolType;
+            setServiceType(serviceType);
+        }
+
+        /**
+         * Sets the service type to be discovered or {@code null} if no services should be queried.
+         *
+         * The {@code serviceType} must be a dot-joint string of two labels. For example,
+         * "_ipp._tcp" for internet printer. Additionally, the first label must start with
+         * underscore ('_') and the second label must be either "_udp" or "_tcp". Otherwise, {@link
+         * NsdManager#discoverServices} will fail with {@link NsdManager#FAILURE_BAD_PARAMETER}.
+         *
+         * @throws IllegalArgumentException if {@code serviceType} is {@code null} or an empty
+         * string
+         *
+         * @hide
+         */
+        @NonNull
+        public Builder setServiceType(@NonNull String serviceType) {
+            if (TextUtils.isEmpty(serviceType)) {
+                throw new IllegalArgumentException("Service type cannot be empty");
+            }
+            mServiceType = serviceType;
+            return this;
+        }
+
+        /**
+         * Sets the optional subtype of the services to be discovered.
+         *
+         * If a non-empty {@code subtype} is specified, it must start with underscore ('_') and
+         * have the trailing "._sub" removed. Otherwise, {@link NsdManager#discoverServices} will
+         * fail with {@link NsdManager#FAILURE_BAD_PARAMETER}. For example, {@code subtype} should
+         * be "_printer" for DNS name "_printer._sub._http._tcp". In this case, only services with
+         * this {@code subtype} will be queried, rather than all services of the base service type.
+         *
+         * Note that a non-empty service type must be specified with {@link #setServiceType} if a
+         * non-empty subtype is specified by this method.
+         */
+        @NonNull
+        public Builder setSubtype(@Nullable String subtype) {
+            mSubtype = subtype;
+            return this;
+        }
+
+        /**
+         * Sets the {@link Network} on which the discovery queries should be sent.
+         *
+         * @param network the discovery network or {@code null} if the query should be sent on
+         * all supported networks
+         */
+        @NonNull
+        public Builder setNetwork(@Nullable Network network) {
+            mNetwork = network;
+            return this;
+        }
+
+        /**
+         * Creates a new {@link DiscoveryRequest} object.
+         */
+        @NonNull
+        public DiscoveryRequest build() {
+            return new DiscoveryRequest(mProtocolType, mServiceType, mSubtype, mNetwork);
+        }
+    }
+}
diff --git a/framework-t/src/android/net/nsd/INsdManagerCallback.aidl b/framework-t/src/android/net/nsd/INsdManagerCallback.aidl
index d89bfa9..55820ec 100644
--- a/framework-t/src/android/net/nsd/INsdManagerCallback.aidl
+++ b/framework-t/src/android/net/nsd/INsdManagerCallback.aidl
@@ -17,6 +17,7 @@
 package android.net.nsd;
 
 import android.os.Messenger;
+import android.net.nsd.DiscoveryRequest;
 import android.net.nsd.NsdServiceInfo;
 
 /**
@@ -24,7 +25,7 @@
  * @hide
  */
 oneway interface INsdManagerCallback {
-    void onDiscoverServicesStarted(int listenerKey, in NsdServiceInfo info);
+    void onDiscoverServicesStarted(int listenerKey, in DiscoveryRequest discoveryRequest);
     void onDiscoverServicesFailed(int listenerKey, int error);
     void onServiceFound(int listenerKey, in NsdServiceInfo info);
     void onServiceLost(int listenerKey, in NsdServiceInfo info);
diff --git a/framework-t/src/android/net/nsd/INsdServiceConnector.aidl b/framework-t/src/android/net/nsd/INsdServiceConnector.aidl
index b03eb29..9a31278 100644
--- a/framework-t/src/android/net/nsd/INsdServiceConnector.aidl
+++ b/framework-t/src/android/net/nsd/INsdServiceConnector.aidl
@@ -17,6 +17,7 @@
 package android.net.nsd;
 
 import android.net.nsd.AdvertisingRequest;
+import android.net.nsd.DiscoveryRequest;
 import android.net.nsd.INsdManagerCallback;
 import android.net.nsd.IOffloadEngine;
 import android.net.nsd.NsdServiceInfo;
@@ -30,7 +31,7 @@
 interface INsdServiceConnector {
     void registerService(int listenerKey, in AdvertisingRequest advertisingRequest);
     void unregisterService(int listenerKey);
-    void discoverServices(int listenerKey, in NsdServiceInfo serviceInfo);
+    void discoverServices(int listenerKey, in DiscoveryRequest discoveryRequest);
     void stopDiscovery(int listenerKey);
     void resolveService(int listenerKey, in NsdServiceInfo serviceInfo);
     void startDaemon();
@@ -39,4 +40,4 @@
     void unregisterServiceInfoCallback(int listenerKey);
     void registerOffloadEngine(String ifaceName, in IOffloadEngine cb, long offloadCapabilities, long offloadType);
     void unregisterOffloadEngine(in IOffloadEngine cb);
-}
\ No newline at end of file
+}
diff --git a/framework-t/src/android/net/nsd/NsdManager.java b/framework-t/src/android/net/nsd/NsdManager.java
index b4f2be9..27b4955 100644
--- a/framework-t/src/android/net/nsd/NsdManager.java
+++ b/framework-t/src/android/net/nsd/NsdManager.java
@@ -57,6 +57,7 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.Executor;
 import java.util.regex.Matcher;
@@ -158,6 +159,8 @@
                 "com.android.net.flags.nsd_subtypes_support_enabled";
         static final String ADVERTISE_REQUEST_API =
                 "com.android.net.flags.advertise_request_api";
+        static final String NSD_CUSTOM_HOSTNAME_ENABLED =
+                "com.android.net.flags.nsd_custom_hostname_enabled";
     }
 
     /**
@@ -360,6 +363,8 @@
     @GuardedBy("mMapLock")
     private final SparseArray<NsdServiceInfo> mServiceMap = new SparseArray<>();
     @GuardedBy("mMapLock")
+    private final SparseArray<DiscoveryRequest> mDiscoveryMap = new SparseArray<>();
+    @GuardedBy("mMapLock")
     private final SparseArray<Executor> mExecutorMap = new SparseArray<>();
     private final Object mMapLock = new Object();
     // Map of listener key sent by client -> per-network discovery tracker
@@ -715,6 +720,12 @@
             mServHandler.sendMessage(mServHandler.obtainMessage(message, 0, listenerKey, info));
         }
 
+        private void sendDiscoveryRequest(
+                int message, int listenerKey, DiscoveryRequest discoveryRequest) {
+            mServHandler.sendMessage(
+                    mServHandler.obtainMessage(message, 0, listenerKey, discoveryRequest));
+        }
+
         private void sendError(int message, int listenerKey, int error) {
             mServHandler.sendMessage(mServHandler.obtainMessage(message, error, listenerKey));
         }
@@ -724,8 +735,8 @@
         }
 
         @Override
-        public void onDiscoverServicesStarted(int listenerKey, NsdServiceInfo info) {
-            sendInfo(DISCOVER_SERVICES_STARTED, listenerKey, info);
+        public void onDiscoverServicesStarted(int listenerKey, DiscoveryRequest discoveryRequest) {
+            sendDiscoveryRequest(DISCOVER_SERVICES_STARTED, listenerKey, discoveryRequest);
         }
 
         @Override
@@ -1003,10 +1014,12 @@
             final Object obj = message.obj;
             final Object listener;
             final NsdServiceInfo ns;
+            final DiscoveryRequest discoveryRequest;
             final Executor executor;
             synchronized (mMapLock) {
                 listener = mListenerMap.get(key);
                 ns = mServiceMap.get(key);
+                discoveryRequest = mDiscoveryMap.get(key);
                 executor = mExecutorMap.get(key);
             }
             if (listener == null) {
@@ -1014,17 +1027,22 @@
                 return;
             }
             if (DBG) {
-                Log.d(TAG, "received " + nameOf(what) + " for key " + key + ", service " + ns);
+                if (discoveryRequest != null) {
+                    Log.d(TAG, "received " + nameOf(what) + " for key " + key + ", discovery "
+                            + discoveryRequest);
+                } else {
+                    Log.d(TAG, "received " + nameOf(what) + " for key " + key + ", service " + ns);
+                }
             }
             switch (what) {
                 case DISCOVER_SERVICES_STARTED:
-                    final String s = getNsdServiceInfoType((NsdServiceInfo) obj);
+                    final String s = getNsdServiceInfoType((DiscoveryRequest) obj);
                     executor.execute(() -> ((DiscoveryListener) listener).onDiscoveryStarted(s));
                     break;
                 case DISCOVER_SERVICES_FAILED:
                     removeListener(key);
                     executor.execute(() -> ((DiscoveryListener) listener).onStartDiscoveryFailed(
-                            getNsdServiceInfoType(ns), errorCode));
+                            getNsdServiceInfoType(discoveryRequest), errorCode));
                     break;
                 case SERVICE_FOUND:
                     executor.execute(() -> ((DiscoveryListener) listener).onServiceFound(
@@ -1039,12 +1057,12 @@
                     // the effect for the client is indistinguishable from STOP_DISCOVERY_SUCCEEDED
                     removeListener(key);
                     executor.execute(() -> ((DiscoveryListener) listener).onStopDiscoveryFailed(
-                            getNsdServiceInfoType(ns), errorCode));
+                            getNsdServiceInfoType(discoveryRequest), errorCode));
                     break;
                 case STOP_DISCOVERY_SUCCEEDED:
                     removeListener(key);
                     executor.execute(() -> ((DiscoveryListener) listener).onDiscoveryStopped(
-                            getNsdServiceInfoType(ns)));
+                            getNsdServiceInfoType(discoveryRequest)));
                     break;
                 case REGISTER_SERVICE_FAILED:
                     removeListener(key);
@@ -1117,21 +1135,33 @@
         return mListenerKey;
     }
 
-    // Assert that the listener is not in the map, then add it and returns its key
-    private int putListener(Object listener, Executor e, NsdServiceInfo s) {
-        checkListener(listener);
-        final int key;
+    private int putListener(Object listener, Executor e, NsdServiceInfo serviceInfo) {
         synchronized (mMapLock) {
-            int valueIndex = mListenerMap.indexOfValue(listener);
+            return putListener(listener, e, mServiceMap, serviceInfo);
+        }
+    }
+
+    private int putListener(Object listener, Executor e, DiscoveryRequest discoveryRequest) {
+        synchronized (mMapLock) {
+            return putListener(listener, e, mDiscoveryMap, discoveryRequest);
+        }
+    }
+
+    // Assert that the listener is not in the map, then add it and returns its key
+    private <T> int putListener(Object listener, Executor e, SparseArray<T> map, T value) {
+        synchronized (mMapLock) {
+            checkListener(listener);
+            final int key;
+            final int valueIndex = mListenerMap.indexOfValue(listener);
             if (valueIndex != -1) {
                 throw new IllegalArgumentException("listener already in use");
             }
             key = nextListenerKey();
             mListenerMap.put(key, listener);
-            mServiceMap.put(key, s);
+            map.put(key, value);
             mExecutorMap.put(key, e);
+            return key;
         }
-        return key;
     }
 
     private int updateRegisteredListener(Object listener, Executor e, NsdServiceInfo s) {
@@ -1148,6 +1178,7 @@
         synchronized (mMapLock) {
             mListenerMap.remove(key);
             mServiceMap.remove(key);
+            mDiscoveryMap.remove(key);
             mExecutorMap.remove(key);
         }
     }
@@ -1163,9 +1194,9 @@
         }
     }
 
-    private static String getNsdServiceInfoType(NsdServiceInfo s) {
-        if (s == null) return "?";
-        return s.getServiceType();
+    private static String getNsdServiceInfoType(DiscoveryRequest r) {
+        if (r == null) return "?";
+        return r.getServiceType();
     }
 
     /**
@@ -1208,7 +1239,7 @@
      */
     public void registerService(@NonNull NsdServiceInfo serviceInfo, int protocolType,
             @NonNull Executor executor, @NonNull RegistrationListener listener) {
-        checkServiceInfo(serviceInfo);
+        checkServiceInfoForRegistration(serviceInfo);
         checkProtocol(protocolType);
         final AdvertisingRequest.Builder builder = new AdvertisingRequest.Builder(serviceInfo,
                 protocolType);
@@ -1267,7 +1298,10 @@
      * @return Type and comma-separated list of subtypes, or null if invalid format.
      */
     @Nullable
-    private static Pair<String, String> getTypeAndSubtypes(@NonNull String typeWithSubtype) {
+    private static Pair<String, String> getTypeAndSubtypes(@Nullable String typeWithSubtype) {
+        if (typeWithSubtype == null) {
+            return null;
+        }
         final Matcher matcher = Pattern.compile(TYPE_REGEX).matcher(typeWithSubtype);
         if (!matcher.matches()) return null;
         // Reject specifications using leading subtypes with a dot
@@ -1298,10 +1332,7 @@
             @NonNull RegistrationListener listener) {
         final NsdServiceInfo serviceInfo = advertisingRequest.getServiceInfo();
         final int protocolType = advertisingRequest.getProtocolType();
-        if (serviceInfo.getPort() <= 0) {
-            throw new IllegalArgumentException("Invalid port number");
-        }
-        checkServiceInfo(serviceInfo);
+        checkServiceInfoForRegistration(serviceInfo);
         checkProtocol(protocolType);
         final int key;
         // For update only request, the old listener has to be reused
@@ -1406,15 +1437,44 @@
         if (TextUtils.isEmpty(serviceType)) {
             throw new IllegalArgumentException("Service type cannot be empty");
         }
-        checkProtocol(protocolType);
+        DiscoveryRequest request = new DiscoveryRequest.Builder(protocolType, serviceType)
+                .setNetwork(network).build();
+        discoverServices(request, executor, listener);
+    }
 
-        NsdServiceInfo s = new NsdServiceInfo();
-        s.setServiceType(serviceType);
-        s.setNetwork(network);
-
-        int key = putListener(listener, executor, s);
+    /**
+     * Initiates service discovery to browse for instances of a service type. Service discovery
+     * consumes network bandwidth and will continue until the application calls
+     * {@link #stopServiceDiscovery}.
+     *
+     * <p> The function call immediately returns after sending a request to start service
+     * discovery to the framework. The application is notified of a success to initiate
+     * discovery through the callback {@link DiscoveryListener#onDiscoveryStarted} or a failure
+     * through {@link DiscoveryListener#onStartDiscoveryFailed}.
+     *
+     * <p> Upon successful start, application is notified when a service is found with
+     * {@link DiscoveryListener#onServiceFound} or when a service is lost with
+     * {@link DiscoveryListener#onServiceLost}.
+     *
+     * <p> Upon failure to start, service discovery is not active and application does
+     * not need to invoke {@link #stopServiceDiscovery}
+     *
+     * <p> The application should call {@link #stopServiceDiscovery} when discovery of this
+     * service type is no longer required, and/or whenever the application is paused or
+     * stopped.
+     *
+     * @param discoveryRequest the {@link DiscoveryRequest} object which specifies the discovery
+     * parameters such as service type, subtype and network
+     * @param executor Executor to run listener callbacks with
+     * @param listener  The listener notifies of a successful discovery and is used
+     * to stop discovery on this serviceType through a call on {@link #stopServiceDiscovery}.
+     */
+    @FlaggedApi(Flags.NSD_SUBTYPES_SUPPORT_ENABLED)
+    public void discoverServices(@NonNull DiscoveryRequest discoveryRequest,
+            @NonNull Executor executor, @NonNull DiscoveryListener listener) {
+        int key = putListener(listener, executor, discoveryRequest);
         try {
-            mService.discoverServices(key, s);
+            mService.discoverServices(key, discoveryRequest);
         } catch (RemoteException e) {
             e.rethrowFromSystemServer();
         }
@@ -1465,12 +1525,10 @@
             throw new IllegalArgumentException("Service type cannot be empty");
         }
         Objects.requireNonNull(networkRequest, "NetworkRequest cannot be null");
-        checkProtocol(protocolType);
+        DiscoveryRequest discoveryRequest =
+                new DiscoveryRequest.Builder(protocolType, serviceType).build();
 
-        NsdServiceInfo s = new NsdServiceInfo();
-        s.setServiceType(serviceType);
-
-        final int baseListenerKey = putListener(listener, executor, s);
+        final int baseListenerKey = putListener(listener, executor, discoveryRequest);
 
         final PerNetworkDiscoveryTracker discoveryInfo = new PerNetworkDiscoveryTracker(
                 serviceType, protocolType, executor, listener);
@@ -1551,7 +1609,7 @@
     @Deprecated
     public void resolveService(@NonNull NsdServiceInfo serviceInfo,
             @NonNull Executor executor, @NonNull ResolveListener listener) {
-        checkServiceInfo(serviceInfo);
+        checkServiceInfoForResolution(serviceInfo);
         int key = putListener(listener, executor, serviceInfo);
         try {
             mService.resolveService(key, serviceInfo);
@@ -1602,9 +1660,10 @@
      * @param executor Executor to run callbacks with
      * @param listener to receive callback upon service update
      */
+    // TODO: use {@link DiscoveryRequest} to specify the service to be subscribed
     public void registerServiceInfoCallback(@NonNull NsdServiceInfo serviceInfo,
             @NonNull Executor executor, @NonNull ServiceInfoCallback listener) {
-        checkServiceInfo(serviceInfo);
+        checkServiceInfoForResolution(serviceInfo);
         int key = putListener(listener, executor, serviceInfo);
         try {
             mService.registerServiceInfoCallback(key, serviceInfo);
@@ -1643,13 +1702,13 @@
         Objects.requireNonNull(listener, "listener cannot be null");
     }
 
-    private static void checkProtocol(int protocolType) {
+    static void checkProtocol(int protocolType) {
         if (protocolType != PROTOCOL_DNS_SD) {
             throw new IllegalArgumentException("Unsupported protocol");
         }
     }
 
-    private static void checkServiceInfo(NsdServiceInfo serviceInfo) {
+    private static void checkServiceInfoForResolution(NsdServiceInfo serviceInfo) {
         Objects.requireNonNull(serviceInfo, "NsdServiceInfo cannot be null");
         if (TextUtils.isEmpty(serviceInfo.getServiceName())) {
             throw new IllegalArgumentException("Service name cannot be empty");
@@ -1658,4 +1717,46 @@
             throw new IllegalArgumentException("Service type cannot be empty");
         }
     }
+
+    /**
+     * Check if the {@link NsdServiceInfo} is valid for registration.
+     *
+     * The following can be registered:
+     * - A service with an optional host.
+     * - A hostname with addresses.
+     *
+     * Note that:
+     * - When registering a service, the service name, service type and port must be specified. If
+     *   hostname is specified, the host addresses can optionally be specified.
+     * - When registering a host without a service, the addresses must be specified.
+     *
+     * @hide
+     */
+    public static void checkServiceInfoForRegistration(NsdServiceInfo serviceInfo) {
+        Objects.requireNonNull(serviceInfo, "NsdServiceInfo cannot be null");
+        boolean hasServiceName = !TextUtils.isEmpty(serviceInfo.getServiceName());
+        boolean hasServiceType = !TextUtils.isEmpty(serviceInfo.getServiceType());
+        boolean hasHostname = !TextUtils.isEmpty(serviceInfo.getHostname());
+        boolean hasHostAddresses = !CollectionUtils.isEmpty(serviceInfo.getHostAddresses());
+
+        if (serviceInfo.getPort() < 0) {
+            throw new IllegalArgumentException("Invalid port");
+        }
+
+        if (hasServiceType || hasServiceName || (serviceInfo.getPort() > 0)) {
+            if (!(hasServiceType && hasServiceName && (serviceInfo.getPort() > 0))) {
+                throw new IllegalArgumentException(
+                        "The service type, service name or port is missing");
+            }
+        }
+
+        if (!hasServiceType && !hasHostname) {
+            throw new IllegalArgumentException("No service or host specified in NsdServiceInfo");
+        }
+
+        if (!hasServiceType && hasHostname && !hasHostAddresses) {
+            // TODO: b/317946010 - This may be allowed when it supports registering KEY RR.
+            throw new IllegalArgumentException("No host addresses specified in NsdServiceInfo");
+        }
+    }
 }
diff --git a/framework-t/src/android/net/nsd/NsdServiceInfo.java b/framework-t/src/android/net/nsd/NsdServiceInfo.java
index ac4ea23..146d4ca 100644
--- a/framework-t/src/android/net/nsd/NsdServiceInfo.java
+++ b/framework-t/src/android/net/nsd/NsdServiceInfo.java
@@ -49,8 +49,10 @@
 
     private static final String TAG = "NsdServiceInfo";
 
+    @Nullable
     private String mServiceName;
 
+    @Nullable
     private String mServiceType;
 
     private final Set<String> mSubtypes;
@@ -59,6 +61,9 @@
 
     private final List<InetAddress> mHostAddresses;
 
+    @Nullable
+    private String mHostname;
+
     private int mPort;
 
     @Nullable
@@ -90,6 +95,7 @@
         mSubtypes = new ArraySet<>(other.getSubtypes());
         mTxtRecord = new ArrayMap<>(other.mTxtRecord);
         mHostAddresses = new ArrayList<>(other.getHostAddresses());
+        mHostname = other.getHostname();
         mPort = other.getPort();
         mNetwork = other.getNetwork();
         mInterfaceIndex = other.getInterfaceIndex();
@@ -169,6 +175,43 @@
     }
 
     /**
+     * Get the hostname.
+     *
+     * <p>When a service is resolved, it returns the hostname of the resolved service . The top
+     * level domain ".local." is omitted.
+     *
+     * <p>For example, it returns "MyHost" when the service's hostname is "MyHost.local.".
+     *
+     * @hide
+     */
+//    @FlaggedApi(NsdManager.Flags.NSD_CUSTOM_HOSTNAME_ENABLED)
+    @Nullable
+    public String getHostname() {
+        return mHostname;
+    }
+
+    /**
+     * Set a custom hostname for this service instance for registration.
+     *
+     * <p>A hostname must be in ".local." domain. The ".local." must be omitted when calling this
+     * method.
+     *
+     * <p>For example, you should call setHostname("MyHost") to use the hostname "MyHost.local.".
+     *
+     * <p>If a hostname is set with this method, the addresses set with {@link #setHostAddresses}
+     * will be registered with the hostname.
+     *
+     * <p>If the hostname is null (which is the default for a new {@link NsdServiceInfo}), a random
+     * hostname is used and the addresses of this device will be registered.
+     *
+     * @hide
+     */
+//    @FlaggedApi(NsdManager.Flags.NSD_CUSTOM_HOSTNAME_ENABLED)
+    public void setHostname(@Nullable String hostname) {
+        mHostname = hostname;
+    }
+
+    /**
      * Unpack txt information from a base-64 encoded byte array.
      *
      * @param txtRecordsRawBytes The raw base64 encoded byte array.
@@ -454,6 +497,7 @@
                 .append(", type: ").append(mServiceType)
                 .append(", subtypes: ").append(TextUtils.join(", ", mSubtypes))
                 .append(", hostAddresses: ").append(TextUtils.join(", ", mHostAddresses))
+                .append(", hostname: ").append(mHostname)
                 .append(", port: ").append(mPort)
                 .append(", network: ").append(mNetwork);
 
@@ -494,6 +538,7 @@
         for (InetAddress address : mHostAddresses) {
             InetAddressUtils.parcelInetAddress(dest, address, flags);
         }
+        dest.writeString(mHostname);
     }
 
     /** Implement the Parcelable interface */
@@ -523,6 +568,7 @@
                 for (int i = 0; i < size; i++) {
                     info.mHostAddresses.add(InetAddressUtils.unparcelInetAddress(in));
                 }
+                info.mHostname = in.readString();
                 return info;
             }
 
diff --git a/framework/aidl-export/android/net/nsd/DiscoveryRequest.aidl b/framework/aidl-export/android/net/nsd/DiscoveryRequest.aidl
new file mode 100644
index 0000000..481a066
--- /dev/null
+++ b/framework/aidl-export/android/net/nsd/DiscoveryRequest.aidl
@@ -0,0 +1,19 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.nsd;
+
+@JavaOnlyStableParcelable parcelable DiscoveryRequest;
diff --git a/framework/src/android/net/BpfNetMapsUtils.java b/framework/src/android/net/BpfNetMapsUtils.java
index 11d610c..0be30bb 100644
--- a/framework/src/android/net/BpfNetMapsUtils.java
+++ b/framework/src/android/net/BpfNetMapsUtils.java
@@ -125,13 +125,11 @@
         return sj.toString();
     }
 
-    public static final boolean PRE_T = !SdkLevel.isAtLeastT();
-
     /**
      * Throw UnsupportedOperationException if SdkLevel is before T.
      */
     public static void throwIfPreT(final String msg) {
-        if (PRE_T) {
+        if (!SdkLevel.isAtLeastT()) {
             throw new UnsupportedOperationException(msg);
         }
     }
diff --git a/service-t/jni/com_android_server_net_NetworkStatsService.cpp b/service-t/jni/com_android_server_net_NetworkStatsService.cpp
index 48ac993..c999398 100644
--- a/service-t/jni/com_android_server_net_NetworkStatsService.cpp
+++ b/service-t/jni/com_android_server_net_NetworkStatsService.cpp
@@ -52,8 +52,14 @@
         return nullptr;
     }
 
+    // Find the constructor.
+    jmethodID constructorID = env->GetMethodID(gEntryClass, "<init>", "()V");
+    if (constructorID == nullptr) {
+        return nullptr;
+    }
+
     // Create a new instance of the Java class
-    jobject result = env->AllocObject(gEntryClass);
+    jobject result = env->NewObject(gEntryClass, constructorID);
     if (result == nullptr) {
         return nullptr;
     }
diff --git a/service-t/lint-baseline.xml b/service-t/lint-baseline.xml
index 7f05b8c..e4b92d6 100644
--- a/service-t/lint-baseline.xml
+++ b/service-t/lint-baseline.xml
@@ -3,17 +3,6 @@
 
     <issue
         id="NewApi"
-        message="Call requires API level 31 (current min is 30): `BpfMap`"
-        errorLine1="                return new BpfMap&lt;&gt;(IFACE_INDEX_NAME_MAP_PATH, BpfMap.BPF_F_RDWR,"
-        errorLine2="                       ~~~~~~~~~~~~">
-        <location
-            file="packages/modules/Connectivity/service-t/src/com/android/server/net/BpfInterfaceMapUpdater.java"
-            line="69"
-            column="24"/>
-    </issue>
-
-    <issue
-        id="NewApi"
         message="Call requires API level 31 (current min is 30): `new android.net.EthernetNetworkSpecifier`"
         errorLine1="                .setNetworkSpecifier(new EthernetNetworkSpecifier(ifaceName))"
         errorLine2="                                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~">
diff --git a/service-t/native/libs/libnetworkstats/BpfNetworkStats.cpp b/service-t/native/libs/libnetworkstats/BpfNetworkStats.cpp
index 8a58e56..d3e331e 100644
--- a/service-t/native/libs/libnetworkstats/BpfNetworkStats.cpp
+++ b/service-t/native/libs/libnetworkstats/BpfNetworkStats.cpp
@@ -50,6 +50,15 @@
     return ifaceStatsMap;
 }
 
+Result<IfaceValue> ifindex2name(const uint32_t ifindex) {
+    Result<IfaceValue> v = getIfaceIndexNameMap().readValue(ifindex);
+    if (v.ok()) return v;
+    IfaceValue iv = {};
+    if (!if_indextoname(ifindex, iv.name)) return v;
+    getIfaceIndexNameMap().writeValue(ifindex, iv, BPF_ANY);
+    return iv;
+}
+
 void bpfRegisterIface(const char* iface) {
     if (!iface) return;
     if (strlen(iface) >= sizeof(IfaceValue)) return;
@@ -78,19 +87,19 @@
 
 int bpfGetIfaceStatsInternal(const char* iface, StatsValue* stats,
                              const BpfMapRO<uint32_t, StatsValue>& ifaceStatsMap,
-                             const BpfMapRO<uint32_t, IfaceValue>& ifaceNameMap) {
+                             const IfIndexToNameFunc ifindex2name) {
     *stats = {};
     int64_t unknownIfaceBytesTotal = 0;
     const auto processIfaceStats =
-            [iface, stats, &ifaceNameMap, &unknownIfaceBytesTotal](
+            [iface, stats, ifindex2name, &unknownIfaceBytesTotal](
                     const uint32_t& key,
                     const BpfMapRO<uint32_t, StatsValue>& ifaceStatsMap) -> Result<void> {
-        IfaceValue ifname;
-        if (getIfaceNameFromMap(ifaceNameMap, ifaceStatsMap, key, ifname, key,
-                                &unknownIfaceBytesTotal)) {
+        Result<IfaceValue> ifname = ifindex2name(key);
+        if (!ifname.ok()) {
+            maybeLogUnknownIface(key, ifaceStatsMap, key, &unknownIfaceBytesTotal);
             return Result<void>();
         }
-        if (!iface || !strcmp(iface, ifname.name)) {
+        if (!iface || !strcmp(iface, ifname.value().name)) {
             Result<StatsValue> statsEntry = ifaceStatsMap.readValue(key);
             if (!statsEntry.ok()) {
                 return statsEntry.error();
@@ -104,7 +113,7 @@
 }
 
 int bpfGetIfaceStats(const char* iface, StatsValue* stats) {
-    return bpfGetIfaceStatsInternal(iface, stats, getIfaceStatsMap(), getIfaceIndexNameMap());
+    return bpfGetIfaceStatsInternal(iface, stats, getIfaceStatsMap(), ifindex2name);
 }
 
 int bpfGetIfIndexStatsInternal(uint32_t ifindex, StatsValue* stats,
@@ -138,22 +147,22 @@
 
 int parseBpfNetworkStatsDetailInternal(std::vector<stats_line>& lines,
                                        const BpfMapRO<StatsKey, StatsValue>& statsMap,
-                                       const BpfMapRO<uint32_t, IfaceValue>& ifaceMap) {
+                                       const IfIndexToNameFunc ifindex2name) {
     int64_t unknownIfaceBytesTotal = 0;
     const auto processDetailUidStats =
-            [&lines, &unknownIfaceBytesTotal, &ifaceMap](
+            [&lines, &unknownIfaceBytesTotal, &ifindex2name](
                     const StatsKey& key,
                     const BpfMapRO<StatsKey, StatsValue>& statsMap) -> Result<void> {
-        IfaceValue ifname;
-        if (getIfaceNameFromMap(ifaceMap, statsMap, key.ifaceIndex, ifname, key,
-                                &unknownIfaceBytesTotal)) {
+        Result<IfaceValue> ifname = ifindex2name(key.ifaceIndex);
+        if (!ifname.ok()) {
+            maybeLogUnknownIface(key.ifaceIndex, statsMap, key, &unknownIfaceBytesTotal);
             return Result<void>();
         }
         Result<StatsValue> statsEntry = statsMap.readValue(key);
         if (!statsEntry.ok()) {
             return base::ResultError(statsEntry.error().message(), statsEntry.error().code());
         }
-        stats_line newLine = populateStatsEntry(key, statsEntry.value(), ifname);
+        stats_line newLine = populateStatsEntry(key, statsEntry.value(), ifname.value());
         lines.push_back(newLine);
         if (newLine.tag) {
             // account tagged traffic in the untagged stats (for historical reasons?)
@@ -212,7 +221,7 @@
     // TODO: the above comment feels like it may be obsolete / out of date,
     // since we no longer swap the map via netd binder rpc - though we do
     // still swap it.
-    int ret = parseBpfNetworkStatsDetailInternal(*lines, *inactiveStatsMap, getIfaceIndexNameMap());
+    int ret = parseBpfNetworkStatsDetailInternal(*lines, *inactiveStatsMap, ifindex2name);
     if (ret) {
         ALOGE("parse detail network stats failed: %s", strerror(errno));
         return ret;
@@ -229,13 +238,14 @@
 
 int parseBpfNetworkStatsDevInternal(std::vector<stats_line>& lines,
                                     const BpfMapRO<uint32_t, StatsValue>& statsMap,
-                                    const BpfMapRO<uint32_t, IfaceValue>& ifaceMap) {
+                                    const IfIndexToNameFunc ifindex2name) {
     int64_t unknownIfaceBytesTotal = 0;
-    const auto processDetailIfaceStats = [&lines, &unknownIfaceBytesTotal, &ifaceMap, &statsMap](
+    const auto processDetailIfaceStats = [&lines, &unknownIfaceBytesTotal, ifindex2name, &statsMap](
                                              const uint32_t& key, const StatsValue& value,
                                              const BpfMapRO<uint32_t, StatsValue>&) {
-        IfaceValue ifname;
-        if (getIfaceNameFromMap(ifaceMap, statsMap, key, ifname, key, &unknownIfaceBytesTotal)) {
+        Result<IfaceValue> ifname = ifindex2name(key);
+        if (!ifname.ok()) {
+            maybeLogUnknownIface(key, statsMap, key, &unknownIfaceBytesTotal);
             return Result<void>();
         }
         StatsKey fakeKey = {
@@ -243,7 +253,7 @@
                 .tag = (uint32_t)TAG_NONE,
                 .counterSet = (uint32_t)SET_ALL,
         };
-        lines.push_back(populateStatsEntry(fakeKey, value, ifname));
+        lines.push_back(populateStatsEntry(fakeKey, value, ifname.value()));
         return Result<void>();
     };
     Result<void> res = statsMap.iterateWithValue(processDetailIfaceStats);
@@ -258,7 +268,7 @@
 }
 
 int parseBpfNetworkStatsDev(std::vector<stats_line>* lines) {
-    return parseBpfNetworkStatsDevInternal(*lines, getIfaceStatsMap(), getIfaceIndexNameMap());
+    return parseBpfNetworkStatsDevInternal(*lines, getIfaceStatsMap(), ifindex2name);
 }
 
 void groupNetworkStats(std::vector<stats_line>& lines) {
diff --git a/service-t/native/libs/libnetworkstats/BpfNetworkStatsTest.cpp b/service-t/native/libs/libnetworkstats/BpfNetworkStatsTest.cpp
index 2c01904..484c166 100644
--- a/service-t/native/libs/libnetworkstats/BpfNetworkStatsTest.cpp
+++ b/service-t/native/libs/libnetworkstats/BpfNetworkStatsTest.cpp
@@ -77,6 +77,10 @@
     BpfMap<uint32_t, IfaceValue> mFakeIfaceIndexNameMap;
     BpfMap<uint32_t, StatsValue> mFakeIfaceStatsMap;
 
+    IfIndexToNameFunc mIfIndex2Name = [this](const uint32_t ifindex){
+        return mFakeIfaceIndexNameMap.readValue(ifindex);
+    };
+
     void SetUp() {
         ASSERT_EQ(0, setrlimitForTest());
 
@@ -228,7 +232,7 @@
     populateFakeStats(TEST_UID1, 0, IFACE_INDEX1, TEST_COUNTERSET0, value1, mFakeStatsMap);
     populateFakeStats(TEST_UID1, 0, IFACE_INDEX2, TEST_COUNTERSET1, value1, mFakeStatsMap);
     populateFakeStats(TEST_UID2, 0, IFACE_INDEX3, TEST_COUNTERSET1, value1, mFakeStatsMap);
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((unsigned long)3, lines.size());
 }
 
@@ -256,16 +260,15 @@
     EXPECT_RESULT_OK(mFakeIfaceStatsMap.writeValue(ifaceStatsKey, value1, BPF_ANY));
 
     StatsValue result1 = {};
-    ASSERT_EQ(0, bpfGetIfaceStatsInternal(IFACE_NAME1, &result1, mFakeIfaceStatsMap,
-                                          mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0,
+              bpfGetIfaceStatsInternal(IFACE_NAME1, &result1, mFakeIfaceStatsMap, mIfIndex2Name));
     expectStatsEqual(value1, result1);
     StatsValue result2 = {};
-    ASSERT_EQ(0, bpfGetIfaceStatsInternal(IFACE_NAME2, &result2, mFakeIfaceStatsMap,
-                                          mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0,
+              bpfGetIfaceStatsInternal(IFACE_NAME2, &result2, mFakeIfaceStatsMap, mIfIndex2Name));
     expectStatsEqual(value2, result2);
     StatsValue totalResult = {};
-    ASSERT_EQ(0, bpfGetIfaceStatsInternal(NULL, &totalResult, mFakeIfaceStatsMap,
-                                          mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, bpfGetIfaceStatsInternal(NULL, &totalResult, mFakeIfaceStatsMap, mIfIndex2Name));
     StatsValue totalValue = {
             .rxPackets = TEST_PACKET0 * 2 + TEST_PACKET1,
             .rxBytes = TEST_BYTES0 * 2 + TEST_BYTES1,
@@ -304,7 +307,7 @@
                       mFakeStatsMap);
     populateFakeStats(TEST_UID2, TEST_TAG, IFACE_INDEX1, TEST_COUNTERSET0, value1, mFakeStatsMap);
     std::vector<stats_line> lines;
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((unsigned long)7, lines.size());
 }
 
@@ -324,7 +327,7 @@
     populateFakeStats(TEST_UID1, 0, IFACE_INDEX1, TEST_COUNTERSET1, value1, mFakeStatsMap);
     populateFakeStats(TEST_UID2, 0, IFACE_INDEX1, TEST_COUNTERSET0, value1, mFakeStatsMap);
     std::vector<stats_line> lines;
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((unsigned long)4, lines.size());
 }
 
@@ -352,18 +355,20 @@
             .counterSet = TEST_COUNTERSET0,
             .ifaceIndex = ifaceIndex,
     };
-    IfaceValue ifname;
     int64_t unknownIfaceBytesTotal = 0;
-    ASSERT_EQ(-ENODEV, getIfaceNameFromMap(mFakeIfaceIndexNameMap, mFakeStatsMap, ifaceIndex,
-                                           ifname, curKey, &unknownIfaceBytesTotal));
+    ASSERT_EQ(false, mFakeIfaceIndexNameMap.readValue(ifaceIndex).ok());
+    maybeLogUnknownIface(ifaceIndex, mFakeStatsMap, curKey, &unknownIfaceBytesTotal);
+
     ASSERT_EQ(((int64_t)(TEST_BYTES0 * 20 + TEST_BYTES1 * 20)), unknownIfaceBytesTotal);
     curKey.ifaceIndex = IFACE_INDEX2;
-    ASSERT_EQ(-ENODEV, getIfaceNameFromMap(mFakeIfaceIndexNameMap, mFakeStatsMap, ifaceIndex,
-                                           ifname, curKey, &unknownIfaceBytesTotal));
+
+    ASSERT_EQ(false, mFakeIfaceIndexNameMap.readValue(ifaceIndex).ok());
+    maybeLogUnknownIface(ifaceIndex, mFakeStatsMap, curKey, &unknownIfaceBytesTotal);
+
     ASSERT_EQ(-1, unknownIfaceBytesTotal);
     std::vector<stats_line> lines;
     // TODO: find a way to test the total of unknown Iface Bytes go above limit.
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((unsigned long)1, lines.size());
     expectStatsLineEqual(value1, IFACE_NAME1, TEST_UID1, TEST_COUNTERSET0, 0, lines.front());
 }
@@ -394,8 +399,7 @@
     ifaceStatsKey = IFACE_INDEX4;
     EXPECT_RESULT_OK(mFakeIfaceStatsMap.writeValue(ifaceStatsKey, value2, BPF_ANY));
     std::vector<stats_line> lines;
-    ASSERT_EQ(0,
-              parseBpfNetworkStatsDevInternal(lines, mFakeIfaceStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDevInternal(lines, mFakeIfaceStatsMap, mIfIndex2Name));
     ASSERT_EQ((unsigned long)4, lines.size());
 
     expectStatsLineEqual(value1, IFACE_NAME1, UID_ALL, SET_ALL, TAG_NONE, lines[0]);
@@ -439,13 +443,13 @@
     std::vector<stats_line> lines;
 
     // Test empty stats.
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((size_t) 0, lines.size());
     lines.clear();
 
     // Test 1 line stats.
     populateFakeStats(TEST_UID1, TEST_TAG, IFACE_INDEX1, TEST_COUNTERSET0, value1, mFakeStatsMap);
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((size_t) 2, lines.size());  // TEST_TAG != 0 -> 1 entry becomes 2 lines
     expectStatsLineEqual(value1, IFACE_NAME1, TEST_UID1, TEST_COUNTERSET0, 0, lines[0]);
     expectStatsLineEqual(value1, IFACE_NAME1, TEST_UID1, TEST_COUNTERSET0, TEST_TAG, lines[1]);
@@ -457,7 +461,7 @@
     populateFakeStats(TEST_UID1, TEST_TAG + 1, IFACE_INDEX1, TEST_COUNTERSET0, value2,
                       mFakeStatsMap);
     populateFakeStats(TEST_UID2, TEST_TAG, IFACE_INDEX1, TEST_COUNTERSET0, value1, mFakeStatsMap);
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((size_t) 9, lines.size());
     lines.clear();
 
@@ -465,7 +469,7 @@
     populateFakeStats(TEST_UID1, TEST_TAG, IFACE_INDEX3, TEST_COUNTERSET0, value1, mFakeStatsMap);
     populateFakeStats(TEST_UID2, TEST_TAG, IFACE_INDEX3, TEST_COUNTERSET0, value1, mFakeStatsMap);
 
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((size_t) 9, lines.size());
 
     // Verify Sorted & Grouped.
@@ -490,8 +494,7 @@
     ifaceStatsKey = IFACE_INDEX3;
     EXPECT_RESULT_OK(mFakeIfaceStatsMap.writeValue(ifaceStatsKey, value1, BPF_ANY));
 
-    ASSERT_EQ(0,
-              parseBpfNetworkStatsDevInternal(lines, mFakeIfaceStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDevInternal(lines, mFakeIfaceStatsMap, mIfIndex2Name));
     ASSERT_EQ((size_t) 2, lines.size());
 
     expectStatsLineEqual(value3, IFACE_NAME1, UID_ALL, SET_ALL, TAG_NONE, lines[0]);
@@ -532,7 +535,7 @@
     // TODO: Mutate counterSet and enlarge TEST_MAP_SIZE if overflow on counterSet is possible.
 
     std::vector<stats_line> lines;
-    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mFakeIfaceIndexNameMap));
+    ASSERT_EQ(0, parseBpfNetworkStatsDetailInternal(lines, mFakeStatsMap, mIfIndex2Name));
     ASSERT_EQ((size_t) 12, lines.size());
 
     // Uid 0 first
diff --git a/service-t/native/libs/libnetworkstats/include/netdbpf/BpfNetworkStats.h b/service-t/native/libs/libnetworkstats/include/netdbpf/BpfNetworkStats.h
index 0f28c1d..59eb195 100644
--- a/service-t/native/libs/libnetworkstats/include/netdbpf/BpfNetworkStats.h
+++ b/service-t/native/libs/libnetworkstats/include/netdbpf/BpfNetworkStats.h
@@ -55,36 +55,27 @@
 bool operator==(const stats_line& lhs, const stats_line& rhs);
 bool operator<(const stats_line& lhs, const stats_line& rhs);
 
+// This mirrors BpfMap.h's:
+//   Result<Value> readValue(const Key key) const
+// for a BpfMap<uint32_t, IfaceValue>
+using IfIndexToNameFunc = std::function<Result<IfaceValue>(const uint32_t)>;
+
 // For test only
 int bpfGetUidStatsInternal(uid_t uid, StatsValue* stats,
                            const BpfMapRO<uint32_t, StatsValue>& appUidStatsMap);
 // For test only
 int bpfGetIfaceStatsInternal(const char* iface, StatsValue* stats,
                              const BpfMapRO<uint32_t, StatsValue>& ifaceStatsMap,
-                             const BpfMapRO<uint32_t, IfaceValue>& ifaceNameMap);
+                             const IfIndexToNameFunc ifindex2name);
 // For test only
 int bpfGetIfIndexStatsInternal(uint32_t ifindex, StatsValue* stats,
                                const BpfMapRO<uint32_t, StatsValue>& ifaceStatsMap);
 // For test only
 int parseBpfNetworkStatsDetailInternal(std::vector<stats_line>& lines,
                                        const BpfMapRO<StatsKey, StatsValue>& statsMap,
-                                       const BpfMapRO<uint32_t, IfaceValue>& ifaceMap);
+                                       const IfIndexToNameFunc ifindex2name);
 // For test only
 int cleanStatsMapInternal(const base::unique_fd& cookieTagMap, const base::unique_fd& tagStatsMap);
-// For test only
-template <class Key>
-int getIfaceNameFromMap(const BpfMapRO<uint32_t, IfaceValue>& ifaceMap,
-                        const BpfMapRO<Key, StatsValue>& statsMap,
-                        uint32_t ifaceIndex, IfaceValue& ifname,
-                        const Key& curKey, int64_t* unknownIfaceBytesTotal) {
-    auto iface = ifaceMap.readValue(ifaceIndex);
-    if (!iface.ok()) {
-        maybeLogUnknownIface(ifaceIndex, statsMap, curKey, unknownIfaceBytesTotal);
-        return -ENODEV;
-    }
-    ifname = iface.value();
-    return 0;
-}
 
 template <class Key>
 void maybeLogUnknownIface(int ifaceIndex, const BpfMapRO<Key, StatsValue>& statsMap,
@@ -112,7 +103,7 @@
 // For test only
 int parseBpfNetworkStatsDevInternal(std::vector<stats_line>& lines,
                                     const BpfMapRO<uint32_t, StatsValue>& statsMap,
-                                    const BpfMapRO<uint32_t, IfaceValue>& ifaceMap);
+                                    const IfIndexToNameFunc ifindex2name);
 
 void bpfRegisterIface(const char* iface);
 int bpfGetUidStats(uid_t uid, StatsValue* stats);
diff --git a/service-t/src/com/android/server/IpSecXfrmController.java b/service-t/src/com/android/server/IpSecXfrmController.java
index c8abd40..3cfbf83 100644
--- a/service-t/src/com/android/server/IpSecXfrmController.java
+++ b/service-t/src/com/android/server/IpSecXfrmController.java
@@ -15,6 +15,7 @@
  */
 package com.android.server;
 
+import static com.android.net.module.util.netlink.NetlinkUtils.SOCKET_RECV_BUFSIZE;
 import static com.android.net.module.util.netlink.xfrm.XfrmNetlinkMessage.IPPROTO_ESP;
 import static com.android.net.module.util.netlink.xfrm.XfrmNetlinkMessage.NETLINK_XFRM;
 import static com.android.net.module.util.netlink.xfrm.XfrmNetlinkMessage.XFRM_MSG_NEWSA;
@@ -106,7 +107,8 @@
     public static class Dependencies {
         /** Get a new XFRM netlink socket and connect it */
         public FileDescriptor newNetlinkSocket() throws ErrnoException, SocketException {
-            final FileDescriptor fd = NetlinkUtils.netlinkSocketForProto(NETLINK_XFRM);
+            final FileDescriptor fd =
+                    NetlinkUtils.netlinkSocketForProto(NETLINK_XFRM, SOCKET_RECV_BUFSIZE);
             NetlinkUtils.connectToKernel(fd);
             return fd;
         }
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 62a2ae5..ac794a1 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -35,6 +35,8 @@
 import static com.android.server.connectivity.mdns.MdnsAdvertiser.AdvertiserMetrics;
 import static com.android.server.connectivity.mdns.MdnsConstants.NO_PACKET;
 import static com.android.server.connectivity.mdns.MdnsRecord.MAX_LABEL_LENGTH;
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE;
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.PASSIVE_QUERY_MODE;
 import static com.android.server.connectivity.mdns.util.MdnsUtils.Clock;
 
 import android.annotation.NonNull;
@@ -54,6 +56,7 @@
 import android.net.mdns.aidl.RegistrationInfo;
 import android.net.mdns.aidl.ResolutionInfo;
 import android.net.nsd.AdvertisingRequest;
+import android.net.nsd.DiscoveryRequest;
 import android.net.nsd.INsdManager;
 import android.net.nsd.INsdManagerCallback;
 import android.net.nsd.INsdServiceConnector;
@@ -250,6 +253,8 @@
 
     private final RemoteCallbackList<IOffloadEngine> mOffloadEngines =
             new RemoteCallbackList<>();
+    @NonNull
+    private final MdnsFeatureFlags mMdnsFeatureFlags;
 
     private static class OffloadEngineInfo {
         @NonNull final String mInterfaceName;
@@ -742,8 +747,8 @@
                 switch (msg.what) {
                     case NsdManager.DISCOVER_SERVICES: {
                         if (DBG) Log.d(TAG, "Discover services");
-                        final ListenerArgs args = (ListenerArgs) msg.obj;
-                        clientInfo = mClients.get(args.connector);
+                        final DiscoveryArgs discoveryArgs = (DiscoveryArgs) msg.obj;
+                        clientInfo = mClients.get(discoveryArgs.connector);
                         // If the binder death notification for a INsdManagerCallback was received
                         // before any calls are received by NsdService, the clientInfo would be
                         // cleared and cause NPE. Add a null check here to prevent this corner case.
@@ -758,10 +763,10 @@
                             break;
                         }
 
-                        final NsdServiceInfo info = args.serviceInfo;
+                        final DiscoveryRequest discoveryRequest = discoveryArgs.discoveryRequest;
                         transactionId = getUniqueId();
                         final Pair<String, List<String>> typeAndSubtype =
-                                parseTypeAndSubtype(info.getServiceType());
+                                parseTypeAndSubtype(discoveryRequest.getServiceType());
                         final String serviceType = typeAndSubtype == null
                                 ? null : typeAndSubtype.first;
                         if (clientInfo.mUseJavaBackend
@@ -773,41 +778,55 @@
                                 break;
                             }
 
+                            String subtype = discoveryRequest.getSubtype();
+                            if (subtype == null && !typeAndSubtype.second.isEmpty()) {
+                                subtype = typeAndSubtype.second.get(0);
+                            }
+
+                            if (subtype != null && !checkSubtypeLabel(subtype)) {
+                                clientInfo.onDiscoverServicesFailedImmediately(clientRequestId,
+                                        NsdManager.FAILURE_BAD_PARAMETERS, false /* isLegacy */);
+                                break;
+                            }
+
                             final String listenServiceType = serviceType + ".local";
                             maybeStartMonitoringSockets();
                             final MdnsListener listener = new DiscoveryListener(clientRequestId,
                                     transactionId, listenServiceType);
                             final MdnsSearchOptions.Builder optionsBuilder =
                                     MdnsSearchOptions.newBuilder()
-                                            .setNetwork(info.getNetwork())
+                                            .setNetwork(discoveryRequest.getNetwork())
                                             .setRemoveExpiredService(true)
-                                            .setIsPassiveMode(true);
-                            if (!typeAndSubtype.second.isEmpty()) {
-                                // The parsing ensures subtype starts with an underscore.
+                                            .setQueryMode(
+                                                    mMdnsFeatureFlags.isAggressiveQueryModeEnabled()
+                                                            ? AGGRESSIVE_QUERY_MODE
+                                                            : PASSIVE_QUERY_MODE);
+                            if (subtype != null) {
+                                // checkSubtypeLabels() ensures that subtypes start with '_' but
                                 // MdnsSearchOptions expects the underscore to not be present.
-                                optionsBuilder.addSubtype(
-                                        typeAndSubtype.second.get(0).substring(1));
+                                optionsBuilder.addSubtype(subtype.substring(1));
                             }
                             mMdnsDiscoveryManager.registerListener(
                                     listenServiceType, listener, optionsBuilder.build());
                             final ClientRequest request = storeDiscoveryManagerRequestMap(
                                     clientRequestId, transactionId, listener, clientInfo,
-                                    info.getNetwork());
-                            clientInfo.onDiscoverServicesStarted(clientRequestId, info, request);
+                                    discoveryRequest.getNetwork());
+                            clientInfo.onDiscoverServicesStarted(
+                                    clientRequestId, discoveryRequest, request);
                             clientInfo.log("Register a DiscoveryListener " + transactionId
                                     + " for service type:" + listenServiceType);
                         } else {
                             maybeStartDaemon();
-                            if (discoverServices(transactionId, info)) {
+                            if (discoverServices(transactionId, discoveryRequest)) {
                                 if (DBG) {
                                     Log.d(TAG, "Discover " + msg.arg2 + " " + transactionId
-                                            + info.getServiceType());
+                                            + discoveryRequest.getServiceType());
                                 }
                                 final ClientRequest request = storeLegacyRequestMap(clientRequestId,
                                         transactionId, clientInfo, msg.what,
                                         mClock.elapsedRealtime());
                                 clientInfo.onDiscoverServicesStarted(
-                                        clientRequestId, info, request);
+                                        clientRequestId, discoveryRequest, request);
                             } else {
                                 stopServiceDiscovery(transactionId);
                                 clientInfo.onDiscoverServicesFailedImmediately(clientRequestId,
@@ -882,10 +901,18 @@
                                 serviceType);
                         final String registerServiceType = typeSubtype == null
                                 ? null : typeSubtype.first;
+                        final String hostname = serviceInfo.getHostname();
+                        // Keep compatible with the legacy behavior: It's allowed to set host
+                        // addresses for a service registration although the host addresses
+                        // won't be registered. To register the addresses for a host, the
+                        // hostname must be specified.
+                        if (hostname == null) {
+                            serviceInfo.setHostAddresses(Collections.emptyList());
+                        }
                         if (clientInfo.mUseJavaBackend
                                 || mDeps.isMdnsAdvertiserEnabled(mContext)
                                 || useAdvertiserForType(registerServiceType)) {
-                            if (registerServiceType == null) {
+                            if (serviceType != null && registerServiceType == null) {
                                 Log.e(TAG, "Invalid service type: " + serviceType);
                                 clientInfo.onRegisterServiceFailedImmediately(clientRequestId,
                                         NsdManager.FAILURE_INTERNAL_ERROR, false /* isLegacy */);
@@ -908,14 +935,25 @@
                             } else {
                                 transactionId = getUniqueId();
                             }
-                            serviceInfo.setServiceType(registerServiceType);
-                            serviceInfo.setServiceName(truncateServiceName(
-                                    serviceInfo.getServiceName()));
+
+                            if (registerServiceType != null) {
+                                serviceInfo.setServiceType(registerServiceType);
+                                serviceInfo.setServiceName(
+                                        truncateServiceName(serviceInfo.getServiceName()));
+                            }
+
+                            if (!checkHostname(hostname)) {
+                                clientInfo.onRegisterServiceFailedImmediately(clientRequestId,
+                                        NsdManager.FAILURE_BAD_PARAMETERS, false /* isLegacy */);
+                                break;
+                            }
 
                             Set<String> subtypes = new ArraySet<>(serviceInfo.getSubtypes());
-                            for (String subType: typeSubtype.second) {
-                                if (!TextUtils.isEmpty(subType)) {
-                                    subtypes.add(subType);
+                            if (typeSubtype != null && typeSubtype.second != null) {
+                                for (String subType : typeSubtype.second) {
+                                    if (!TextUtils.isEmpty(subType)) {
+                                        subtypes.add(subType);
+                                    }
                                 }
                             }
                             subtypes = dedupSubtypeLabels(subtypes);
@@ -932,7 +970,7 @@
                                     MdnsAdvertisingOptions.newBuilder().setIsOnlyUpdate(
                                             isUpdateOnly).build();
                             mAdvertiser.addOrUpdateService(transactionId, serviceInfo,
-                                    mdnsAdvertisingOptions);
+                                    mdnsAdvertisingOptions, clientInfo.mUid);
                             storeAdvertiserRequestMap(clientRequestId, transactionId, clientInfo,
                                     serviceInfo.getNetwork());
                         } else {
@@ -1031,7 +1069,9 @@
                                     transactionId, resolveServiceType);
                             final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
                                     .setNetwork(info.getNetwork())
-                                    .setIsPassiveMode(true)
+                                    .setQueryMode(mMdnsFeatureFlags.isAggressiveQueryModeEnabled()
+                                            ? AGGRESSIVE_QUERY_MODE
+                                            : PASSIVE_QUERY_MODE)
                                     .setResolveInstanceName(info.getServiceName())
                                     .setRemoveExpiredService(true)
                                     .build();
@@ -1129,7 +1169,9 @@
                                 transactionId, resolveServiceType);
                         final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
                                 .setNetwork(info.getNetwork())
-                                .setIsPassiveMode(true)
+                                .setQueryMode(mMdnsFeatureFlags.isAggressiveQueryModeEnabled()
+                                        ? AGGRESSIVE_QUERY_MODE
+                                        : PASSIVE_QUERY_MODE)
                                 .setResolveInstanceName(info.getServiceName())
                                 .setRemoveExpiredService(true)
                                 .build();
@@ -1522,6 +1564,7 @@
                                 Log.e(TAG, "Invalid attribute", e);
                             }
                         }
+                        info.setHostname(getHostname(serviceInfo));
                         final List<InetAddress> addresses = getInetAddresses(serviceInfo);
                         if (addresses.size() != 0) {
                             info.setHostAddresses(addresses);
@@ -1558,6 +1601,7 @@
                             }
                         }
 
+                        info.setHostname(getHostname(serviceInfo));
                         final List<InetAddress> addresses = getInetAddresses(serviceInfo);
                         info.setHostAddresses(addresses);
                         clientInfo.onServiceUpdated(clientRequestId, info, request);
@@ -1604,6 +1648,16 @@
         return addresses;
     }
 
+    @NonNull
+    private static String getHostname(@NonNull MdnsServiceInfo serviceInfo) {
+        String[] hostname = serviceInfo.getHostName();
+        // Strip the "local" top-level domain.
+        if (hostname.length >= 2 && hostname[hostname.length - 1].equals("local")) {
+            hostname = Arrays.copyOf(hostname, hostname.length - 1);
+        }
+        return String.join(".", hostname);
+    }
+
     private static void setServiceNetworkForCallback(NsdServiceInfo info, int netId, int ifaceIdx) {
         switch (netId) {
             case NETID_UNSET:
@@ -1689,6 +1743,21 @@
         return new Pair<>(queryType, Collections.emptyList());
     }
 
+    /**
+     * Checks if the hostname is valid.
+     *
+     * <p>For now NsdService only allows single-label hostnames conforming to RFC 1035. In other
+     * words, the hostname should be at most 63 characters long and it only contains letters, digits
+     * and hyphens.
+     */
+    public static boolean checkHostname(@Nullable String hostname) {
+        if (hostname == null) {
+            return true;
+        }
+        String HOSTNAME_REGEX = "^[a-zA-Z]([a-zA-Z0-9-_]{0,61}[a-zA-Z0-9])?$";
+        return Pattern.compile(HOSTNAME_REGEX).matcher(hostname).matches();
+    }
+
     /** Returns {@code true} if {@code subtype} is a valid DNS-SD subtype label. */
     private static boolean checkSubtypeLabel(String subtype) {
         return Pattern.compile("^" + TYPE_SUBTYPE_LABEL_REGEX + "$").matcher(subtype).matches();
@@ -1728,7 +1797,7 @@
         am.addOnUidImportanceListener(new UidImportanceListener(handler),
                 mRunningAppActiveImportanceCutoff);
 
-        final MdnsFeatureFlags flags = new MdnsFeatureFlags.Builder()
+        mMdnsFeatureFlags = new MdnsFeatureFlags.Builder()
                 .setIsMdnsOffloadFeatureEnabled(mDeps.isTetheringFeatureNotChickenedOut(
                         mContext, MdnsFeatureFlags.NSD_FORCE_DISABLE_MDNS_OFFLOAD))
                 .setIncludeInetAddressRecordsInProbing(mDeps.isFeatureEnabled(
@@ -1741,18 +1810,21 @@
                         mContext, MdnsFeatureFlags.NSD_KNOWN_ANSWER_SUPPRESSION))
                 .setIsUnicastReplyEnabled(mDeps.isFeatureEnabled(
                         mContext, MdnsFeatureFlags.NSD_UNICAST_REPLY_ENABLED))
+                .setIsAggressiveQueryModeEnabled(mDeps.isFeatureEnabled(
+                        mContext, MdnsFeatureFlags.NSD_AGGRESSIVE_QUERY_MODE))
                 .setOverrideProvider(flag -> mDeps.isFeatureEnabled(
                         mContext, FORCE_ENABLE_FLAG_FOR_TEST_PREFIX + flag))
                 .build();
         mMdnsSocketClient =
                 new MdnsMultinetworkSocketClient(handler.getLooper(), mMdnsSocketProvider,
-                        LOGGER.forSubComponent("MdnsMultinetworkSocketClient"), flags);
+                        LOGGER.forSubComponent("MdnsMultinetworkSocketClient"), mMdnsFeatureFlags);
         mMdnsDiscoveryManager = deps.makeMdnsDiscoveryManager(new ExecutorProvider(),
-                mMdnsSocketClient, LOGGER.forSubComponent("MdnsDiscoveryManager"), flags);
+                mMdnsSocketClient, LOGGER.forSubComponent("MdnsDiscoveryManager"),
+                mMdnsFeatureFlags);
         handler.post(() -> mMdnsSocketClient.setCallback(mMdnsDiscoveryManager));
         mAdvertiser = deps.makeMdnsAdvertiser(handler.getLooper(), mMdnsSocketProvider,
-                new AdvertiserCallback(), LOGGER.forSubComponent("MdnsAdvertiser"), flags,
-                mContext);
+                new AdvertiserCallback(), LOGGER.forSubComponent("MdnsAdvertiser"),
+                mMdnsFeatureFlags, mContext);
         mClock = deps.makeClock();
     }
 
@@ -2018,9 +2090,10 @@
             final int clientRequestId = getClientRequestIdOrLog(clientInfo, transactionId);
             if (clientRequestId < 0) return;
 
-            // onRegisterServiceSucceeded only has the service name in its info. This aligns with
-            // historical behavior.
+            // onRegisterServiceSucceeded only has the service name and hostname in its info. This
+            // aligns with historical behavior.
             final NsdServiceInfo cbInfo = new NsdServiceInfo(registeredInfo.getServiceName(), null);
+            cbInfo.setHostname(registeredInfo.getHostname());
             final ClientRequest request = clientInfo.mClientRequests.get(clientRequestId);
             clientInfo.onRegisterServiceSucceeded(clientRequestId, cbInfo, request);
         }
@@ -2115,12 +2188,22 @@
         }
     }
 
+    private static final class DiscoveryArgs {
+        public final NsdServiceConnector connector;
+        public final DiscoveryRequest discoveryRequest;
+        DiscoveryArgs(NsdServiceConnector connector, DiscoveryRequest discoveryRequest) {
+            this.connector = connector;
+            this.discoveryRequest = discoveryRequest;
+        }
+    }
+
     private class NsdServiceConnector extends INsdServiceConnector.Stub
             implements IBinder.DeathRecipient  {
 
         @Override
         public void registerService(int listenerKey, AdvertisingRequest advertisingRequest)
                 throws RemoteException {
+            NsdManager.checkServiceInfoForRegistration(advertisingRequest.getServiceInfo());
             mNsdStateMachine.sendMessage(mNsdStateMachine.obtainMessage(
                     NsdManager.REGISTER_SERVICE, 0, listenerKey,
                     new AdvertisingArgs(this, advertisingRequest)
@@ -2135,10 +2218,10 @@
         }
 
         @Override
-        public void discoverServices(int listenerKey, NsdServiceInfo serviceInfo) {
+        public void discoverServices(int listenerKey, DiscoveryRequest discoveryRequest) {
             mNsdStateMachine.sendMessage(mNsdStateMachine.obtainMessage(
                     NsdManager.DISCOVER_SERVICES, 0, listenerKey,
-                    new ListenerArgs(this, serviceInfo)));
+                    new DiscoveryArgs(this, discoveryRequest)));
         }
 
         @Override
@@ -2277,15 +2360,15 @@
         return mMDnsManager.stopOperation(transactionId);
     }
 
-    private boolean discoverServices(int transactionId, NsdServiceInfo serviceInfo) {
+    private boolean discoverServices(int transactionId, DiscoveryRequest discoveryRequest) {
         if (mMDnsManager == null) {
             Log.wtf(TAG, "discoverServices: mMDnsManager is null");
             return false;
         }
 
-        final String type = serviceInfo.getServiceType();
-        final int discoverInterface = getNetworkInterfaceIndex(serviceInfo);
-        if (serviceInfo.getNetwork() != null && discoverInterface == IFACE_IDX_ANY) {
+        final String type = discoveryRequest.getServiceType();
+        final int discoverInterface = getNetworkInterfaceIndex(discoveryRequest);
+        if (discoveryRequest.getNetwork() != null && discoverInterface == IFACE_IDX_ANY) {
             Log.e(TAG, "Interface to discover service on not found");
             return false;
         }
@@ -2335,7 +2418,26 @@
             }
             return IFACE_IDX_ANY;
         }
+        return getNetworkInterfaceIndex(network);
+    }
 
+    /**
+     * Returns the interface to use to discover a service on a specific network, or {@link
+     * IFACE_IDX_ANY} if no network is specified.
+     */
+    private int getNetworkInterfaceIndex(DiscoveryRequest discoveryRequest) {
+        final Network network = discoveryRequest.getNetwork();
+        if (network == null) {
+            return IFACE_IDX_ANY;
+        }
+        return getNetworkInterfaceIndex(network);
+    }
+
+    /**
+     * Returns the interface of a specific network, or {@link IFACE_IDX_ANY} if no interface is
+     * associated with {@code network}.
+     */
+    private int getNetworkInterfaceIndex(@NonNull Network network) {
         String interfaceName = getNetworkInterfaceName(network);
         if (interfaceName == null) {
             return IFACE_IDX_ANY;
@@ -2710,12 +2812,12 @@
                     && !(request instanceof AdvertiserClientRequest);
         }
 
-        void onDiscoverServicesStarted(int listenerKey, NsdServiceInfo info,
+        void onDiscoverServicesStarted(int listenerKey, DiscoveryRequest discoveryRequest,
                 ClientRequest request) {
             mMetrics.reportServiceDiscoveryStarted(
                     isLegacyClientRequest(request), request.mTransactionId);
             try {
-                mCb.onDiscoverServicesStarted(listenerKey, info);
+                mCb.onDiscoverServicesStarted(listenerKey, discoveryRequest);
             } catch (RemoteException e) {
                 Log.e(TAG, "Error calling onDiscoverServicesStarted", e);
             }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java b/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java
index b2af93c..0b60572 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java
@@ -17,6 +17,8 @@
 package com.android.server.connectivity.mdns;
 
 import static com.android.server.connectivity.mdns.MdnsConstants.NO_PACKET;
+import static com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_HOST;
+import static com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_SERVICE;
 import static com.android.server.connectivity.mdns.MdnsRecord.MAX_LABEL_LENGTH;
 
 import android.annotation.NonNull;
@@ -31,6 +33,7 @@
 import android.net.nsd.OffloadServiceInfo;
 import android.os.Build;
 import android.os.Looper;
+import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.Log;
 import android.util.SparseArray;
@@ -151,7 +154,9 @@
                 mSharedLog.wtf("Register succeeded for unknown registration");
                 return;
             }
-            if (mMdnsFeatureFlags.mIsMdnsOffloadFeatureEnabled) {
+            if (mMdnsFeatureFlags.mIsMdnsOffloadFeatureEnabled
+                    // TODO: Enable offload when the serviceInfo contains a custom host.
+                    && TextUtils.isEmpty(registration.getServiceInfo().getHostname())) {
                 final String interfaceName = advertiser.getSocketInterfaceName();
                 final List<OffloadServiceInfoWrapper> existingOffloadServiceInfoWrappers =
                         mInterfaceOffloadServices.computeIfAbsent(interfaceName,
@@ -179,8 +184,11 @@
         }
 
         @Override
-        public void onServiceConflict(@NonNull MdnsInterfaceAdvertiser advertiser, int serviceId) {
-            mSharedLog.i("Found conflict, restarted probing for service " + serviceId);
+        public void onServiceConflict(@NonNull MdnsInterfaceAdvertiser advertiser, int serviceId,
+                int conflictType) {
+            mSharedLog.i("Found conflict, restarted probing for service "
+                    + serviceId + " "
+                    + conflictType);
 
             final Registration registration = mRegistrations.get(serviceId);
             if (registration == null) return;
@@ -205,10 +213,22 @@
                 return;
             }
 
-            // Conflict was found during probing; rename once to find a name that has no conflict
-            registration.updateForConflict(
-                    registration.makeNewServiceInfoForConflict(1 /* renameCount */),
-                    1 /* renameCount */);
+            if ((conflictType & CONFLICT_SERVICE) != 0) {
+                // Service conflict was found during probing; rename once to find a name that has no
+                // conflict
+                registration.updateForServiceConflict(
+                        registration.makeNewServiceInfoForServiceConflict(1 /* renameCount */),
+                        1 /* renameCount */);
+            }
+
+            if ((conflictType & CONFLICT_HOST) != 0) {
+                // Host conflict was found during probing; rename once to find a name that has no
+                // conflict
+                registration.updateForHostConflict(
+                        registration.makeNewServiceInfoForHostConflict(1 /* renameCount */),
+                        1 /* renameCount */);
+            }
+
             registration.mConflictDuringProbingCount++;
 
             // Keep renaming if the new name conflicts in local registrations
@@ -231,23 +251,54 @@
         }
     };
 
-    private boolean hasAnyConflict(
+    private boolean hasAnyServiceConflict(
             @NonNull BiPredicate<Network, InterfaceAdvertiserRequest> applicableAdvertiserFilter,
-            @NonNull NsdServiceInfo newInfo) {
-        return any(mAdvertiserRequests, (network, adv) ->
-                applicableAdvertiserFilter.test(network, adv) && adv.hasConflict(newInfo));
+            @NonNull NsdServiceInfo newInfo,
+            @NonNull Registration originalRegistration) {
+        return any(
+                mAdvertiserRequests,
+                (network, adv) ->
+                        applicableAdvertiserFilter.test(network, adv)
+                                && adv.hasServiceConflict(newInfo, originalRegistration));
+    }
+
+    private boolean hasAnyHostConflict(
+            @NonNull BiPredicate<Network, InterfaceAdvertiserRequest> applicableAdvertiserFilter,
+            @NonNull NsdServiceInfo newInfo,
+            int clientUid) {
+        // Check if it conflicts with custom hosts.
+        if (any(
+                mAdvertiserRequests,
+                (network, adv) ->
+                        applicableAdvertiserFilter.test(network, adv)
+                                && adv.hasHostConflict(newInfo, clientUid))) {
+            return true;
+        }
+        // Check if it conflicts with the default hostname.
+        return MdnsUtils.equalsIgnoreDnsCase(newInfo.getHostname(), mDeviceHostName[0]);
     }
 
     private void updateRegistrationUntilNoConflict(
             @NonNull BiPredicate<Network, InterfaceAdvertiserRequest> applicableAdvertiserFilter,
             @NonNull Registration registration) {
-        int renameCount = 0;
         NsdServiceInfo newInfo = registration.getServiceInfo();
-        while (hasAnyConflict(applicableAdvertiserFilter, newInfo)) {
-            renameCount++;
-            newInfo = registration.makeNewServiceInfoForConflict(renameCount);
+
+        int renameServiceCount = 0;
+        while (hasAnyServiceConflict(applicableAdvertiserFilter, newInfo, registration)) {
+            renameServiceCount++;
+            newInfo = registration.makeNewServiceInfoForServiceConflict(renameServiceCount);
         }
-        registration.updateForConflict(newInfo, renameCount);
+        registration.updateForServiceConflict(newInfo, renameServiceCount);
+
+        if (!TextUtils.isEmpty(registration.getServiceInfo().getHostname())) {
+            int renameHostCount = 0;
+            while (hasAnyHostConflict(
+                    applicableAdvertiserFilter, newInfo, registration.mClientUid)) {
+                renameHostCount++;
+                newInfo = registration.makeNewServiceInfoForHostConflict(renameHostCount);
+            }
+            registration.updateForHostConflict(newInfo, renameHostCount);
+        }
     }
 
     private void maybeSendOffloadStop(final String interfaceName, int serviceId) {
@@ -326,17 +377,34 @@
 
         /**
          * Return whether using the proposed new {@link NsdServiceInfo} to add a registration would
-         * cause a conflict in this {@link InterfaceAdvertiserRequest}.
+         * cause a conflict of the service in this {@link InterfaceAdvertiserRequest}.
          */
-        boolean hasConflict(@NonNull NsdServiceInfo newInfo) {
-            return getConflictingService(newInfo) >= 0;
+        boolean hasServiceConflict(
+                @NonNull NsdServiceInfo newInfo, @NonNull Registration originalRegistration) {
+            return getConflictingRegistrationDueToService(newInfo, originalRegistration) >= 0;
         }
 
         /**
-         * Get the ID of a conflicting service, or -1 if none.
+         * Return whether using the proposed new {@link NsdServiceInfo} to add a registration would
+         * cause a conflict of the host in this {@link InterfaceAdvertiserRequest}.
+         *
+         * @param clientUid UID of the user who wants to advertise the serviceInfo.
          */
-        int getConflictingService(@NonNull NsdServiceInfo info) {
+        boolean hasHostConflict(@NonNull NsdServiceInfo newInfo, int clientUid) {
+            return getConflictingRegistrationDueToHost(newInfo, clientUid) >= 0;
+        }
+
+        /** Get the ID of a conflicting registration due to service, or -1 if none. */
+        int getConflictingRegistrationDueToService(
+                @NonNull NsdServiceInfo info, @NonNull Registration originalRegistration) {
+            if (TextUtils.isEmpty(info.getServiceName())) {
+                return -1;
+            }
             for (int i = 0; i < mPendingRegistrations.size(); i++) {
+                // Never conflict with itself
+                if (mPendingRegistrations.valueAt(i) == originalRegistration) {
+                    continue;
+                }
                 final NsdServiceInfo other = mPendingRegistrations.valueAt(i).getServiceInfo();
                 if (MdnsUtils.equalsIgnoreDnsCase(info.getServiceName(), other.getServiceName())
                         && MdnsUtils.equalsIgnoreDnsCase(info.getServiceType(),
@@ -348,9 +416,34 @@
         }
 
         /**
+         * Get the ID of a conflicting registration due to host, or -1 if none.
+         *
+         * <p>It's valid that multiple registrations from the same user are using the same hostname.
+         *
+         * <p>If there's already another registration with the same hostname requested by another
+         * user, this is considered a conflict.
+         */
+        int getConflictingRegistrationDueToHost(@NonNull NsdServiceInfo info, int clientUid) {
+            if (TextUtils.isEmpty(info.getHostname())) {
+                return -1;
+            }
+            for (int i = 0; i < mPendingRegistrations.size(); i++) {
+                final Registration otherRegistration = mPendingRegistrations.valueAt(i);
+                final NsdServiceInfo otherInfo = otherRegistration.getServiceInfo();
+                if (clientUid != otherRegistration.mClientUid
+                        && MdnsUtils.equalsIgnoreDnsCase(
+                                info.getHostname(), otherInfo.getHostname())) {
+                    return mPendingRegistrations.keyAt(i);
+                }
+            }
+            return -1;
+        }
+
+        /**
          * Add a service to advertise.
          *
-         * Conflicts must be checked via {@link #getConflictingService} before attempting to add.
+         * <p>Conflicts must be checked via {@link #getConflictingRegistrationDueToService} and
+         * {@link #getConflictingRegistrationDueToHost} before attempting to add.
          */
         void addService(int id, @NonNull Registration registration) {
             mPendingRegistrations.put(id, registration);
@@ -484,27 +577,35 @@
     }
 
     private static class Registration {
-        @NonNull
-        final String mOriginalName;
+        @Nullable
+        final String mOriginalServiceName;
+        @Nullable
+        final String mOriginalHostname;
         boolean mNotifiedRegistrationSuccess;
-        private int mConflictCount;
+        private int mServiceNameConflictCount;
+        private int mHostnameConflictCount;
         @NonNull
         private NsdServiceInfo mServiceInfo;
+        final int mClientUid;
         int mConflictDuringProbingCount;
         int mConflictAfterProbingCount;
 
-        private Registration(@NonNull NsdServiceInfo serviceInfo) {
-            this.mOriginalName = serviceInfo.getServiceName();
+
+        private Registration(@NonNull NsdServiceInfo serviceInfo, int clientUid) {
+            this.mOriginalServiceName = serviceInfo.getServiceName();
+            this.mOriginalHostname = serviceInfo.getHostname();
             this.mServiceInfo = serviceInfo;
+            this.mClientUid = clientUid;
         }
 
-        /**
-         * Matches between the NsdServiceInfo in the Registration and the provided argument.
-         */
-        public boolean matches(@Nullable NsdServiceInfo newInfo) {
-            return Objects.equals(newInfo.getServiceName(), mOriginalName) && Objects.equals(
-                    newInfo.getServiceType(), mServiceInfo.getServiceType()) && Objects.equals(
-                    newInfo.getNetwork(), mServiceInfo.getNetwork());
+        /** Check if the new {@link NsdServiceInfo} doesn't update any data other than subtypes. */
+        public boolean isSubtypeOnlyUpdate(@NonNull NsdServiceInfo newInfo) {
+            return Objects.equals(newInfo.getServiceName(), mOriginalServiceName)
+                    && Objects.equals(newInfo.getServiceType(), mServiceInfo.getServiceType())
+                    && newInfo.getPort() == mServiceInfo.getPort()
+                    && Objects.equals(newInfo.getHostname(), mOriginalHostname)
+                    && Objects.equals(newInfo.getHostAddresses(), mServiceInfo.getHostAddresses())
+                    && Objects.equals(newInfo.getNetwork(), mServiceInfo.getNetwork());
         }
 
         /**
@@ -521,8 +622,19 @@
          * @param newInfo New service info to use.
          * @param renameCount How many renames were done before reaching the current name.
          */
-        private void updateForConflict(@NonNull NsdServiceInfo newInfo, int renameCount) {
-            mConflictCount += renameCount;
+        private void updateForServiceConflict(@NonNull NsdServiceInfo newInfo, int renameCount) {
+            mServiceNameConflictCount += renameCount;
+            mServiceInfo = newInfo;
+        }
+
+        /**
+         * Update the registration to use a different host name, after a conflict was found.
+         *
+         * @param newInfo New service info to use.
+         * @param renameCount How many renames were done before reaching the current name.
+         */
+        private void updateForHostConflict(@NonNull NsdServiceInfo newInfo, int renameCount) {
+            mHostnameConflictCount += renameCount;
             mServiceInfo = newInfo;
         }
 
@@ -538,7 +650,7 @@
          * @param renameCount How much to increase the number suffix for this conflict.
          */
         @NonNull
-        public NsdServiceInfo makeNewServiceInfoForConflict(int renameCount) {
+        public NsdServiceInfo makeNewServiceInfoForServiceConflict(int renameCount) {
             // In case of conflict choose a different service name. After the first conflict use
             // "Name (2)", then "Name (3)" etc.
             // TODO: use a hidden method in NsdServiceInfo once MdnsAdvertiser is moved to service-t
@@ -547,13 +659,40 @@
             return newInfo;
         }
 
+        /**
+         * Make a new hostname for the registration, after a conflict was found.
+         *
+         * <p>If a name conflict was found during probing or because different advertising requests
+         * used the same name, the registration is attempted again with a new name (here using a
+         * number suffix, -1, -2, etc). Registration success is notified once probing succeeds with
+         * a new name.
+         *
+         * @param renameCount How much to increase the number suffix for this conflict.
+         */
+        @NonNull
+        public NsdServiceInfo makeNewServiceInfoForHostConflict(int renameCount) {
+            // In case of conflict choose a different hostname. After the first conflict use
+            // "Name-2", then "Name-3" etc.
+            final NsdServiceInfo newInfo = new NsdServiceInfo(mServiceInfo);
+            newInfo.setHostname(getUpdatedHostname(renameCount));
+            return newInfo;
+        }
+
         private String getUpdatedServiceName(int renameCount) {
-            final String suffix = " (" + (mConflictCount + renameCount + 1) + ")";
-            final String truncatedServiceName = MdnsUtils.truncateServiceName(mOriginalName,
+            final String suffix = " (" + (mServiceNameConflictCount + renameCount + 1) + ")";
+            final String truncatedServiceName = MdnsUtils.truncateServiceName(mOriginalServiceName,
                     MAX_LABEL_LENGTH - suffix.length());
             return truncatedServiceName + suffix;
         }
 
+        private String getUpdatedHostname(int renameCount) {
+            final String suffix = "-" + (mHostnameConflictCount + renameCount + 1);
+            final String truncatedHostname =
+                    MdnsUtils.truncateServiceName(
+                            mOriginalHostname, MAX_LABEL_LENGTH - suffix.length());
+            return truncatedHostname + suffix;
+        }
+
         @NonNull
         public NsdServiceInfo getServiceInfo() {
             return mServiceInfo;
@@ -681,9 +820,10 @@
      * @param id A unique ID for the service.
      * @param service The service info to advertise.
      * @param advertisingOptions The advertising options.
+     * @param clientUid The UID who wants to advertise the service.
      */
     public void addOrUpdateService(int id, NsdServiceInfo service,
-            MdnsAdvertisingOptions advertisingOptions) {
+            MdnsAdvertisingOptions advertisingOptions, int clientUid) {
         checkThread();
         final Registration existingRegistration = mRegistrations.get(id);
         final Network network = service.getNetwork();
@@ -695,7 +835,7 @@
                 mCb.onRegisterServiceFailed(id, NsdManager.FAILURE_INTERNAL_ERROR);
                 return;
             }
-            if (!(existingRegistration.matches(service))) {
+            if (!(existingRegistration.isSubtypeOnlyUpdate(service))) {
                 mSharedLog.e("Update request can only update subType, serviceInfo: " + service
                         + ", existing serviceInfo: " + existingRegistration.getServiceInfo());
                 mCb.onRegisterServiceFailed(id, NsdManager.FAILURE_INTERNAL_ERROR);
@@ -715,7 +855,7 @@
             }
             mSharedLog.i("Adding service " + service + " with ID " + id + " and subtypes "
                     + subtypes + " advertisingOptions " + advertisingOptions);
-            registration = new Registration(service);
+            registration = new Registration(service, clientUid);
             final BiPredicate<Network, InterfaceAdvertiserRequest> checkConflictFilter;
             if (network == null) {
                 // If registering on all networks, no advertiser must have conflicts
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java b/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
index 9466162..fe9bbba 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsFeatureFlags.java
@@ -57,6 +57,11 @@
      */
     public static final String NSD_UNICAST_REPLY_ENABLED = "nsd_unicast_reply_enabled";
 
+    /**
+     * A feature flag to control whether the aggressive query mode should be enabled.
+     */
+    public static final String NSD_AGGRESSIVE_QUERY_MODE = "nsd_aggressive_query_mode";
+
     // Flag for offload feature
     public final boolean mIsMdnsOffloadFeatureEnabled;
 
@@ -75,6 +80,9 @@
     // Flag to enable replying unicast to queries requesting unicast replies
     public final boolean mIsUnicastReplyEnabled;
 
+    // Flag for aggressive query mode
+    public final boolean mIsAggressiveQueryModeEnabled;
+
     @Nullable
     private final FlagOverrideProvider mOverrideProvider;
 
@@ -103,6 +111,13 @@
     }
 
     /**
+     * Indicates whether {@link #NSD_AGGRESSIVE_QUERY_MODE} is enabled, including for testing.
+     */
+    public boolean isAggressiveQueryModeEnabled() {
+        return mIsAggressiveQueryModeEnabled || isForceEnabledForTest(NSD_AGGRESSIVE_QUERY_MODE);
+    }
+
+    /**
      * The constructor for {@link MdnsFeatureFlags}.
      */
     public MdnsFeatureFlags(boolean isOffloadFeatureEnabled,
@@ -111,6 +126,7 @@
             boolean isLabelCountLimitEnabled,
             boolean isKnownAnswerSuppressionEnabled,
             boolean isUnicastReplyEnabled,
+            boolean isAggressiveQueryModeEnabled,
             @Nullable FlagOverrideProvider overrideProvider) {
         mIsMdnsOffloadFeatureEnabled = isOffloadFeatureEnabled;
         mIncludeInetAddressRecordsInProbing = includeInetAddressRecordsInProbing;
@@ -118,6 +134,7 @@
         mIsLabelCountLimitEnabled = isLabelCountLimitEnabled;
         mIsKnownAnswerSuppressionEnabled = isKnownAnswerSuppressionEnabled;
         mIsUnicastReplyEnabled = isUnicastReplyEnabled;
+        mIsAggressiveQueryModeEnabled = isAggressiveQueryModeEnabled;
         mOverrideProvider = overrideProvider;
     }
 
@@ -136,6 +153,7 @@
         private boolean mIsLabelCountLimitEnabled;
         private boolean mIsKnownAnswerSuppressionEnabled;
         private boolean mIsUnicastReplyEnabled;
+        private boolean mIsAggressiveQueryModeEnabled;
         private FlagOverrideProvider mOverrideProvider;
 
         /**
@@ -148,6 +166,7 @@
             mIsLabelCountLimitEnabled = true; // Default enabled.
             mIsKnownAnswerSuppressionEnabled = false;
             mIsUnicastReplyEnabled = true;
+            mIsAggressiveQueryModeEnabled = false;
             mOverrideProvider = null;
         }
 
@@ -224,6 +243,16 @@
         }
 
         /**
+         * Set whether the aggressive query mode is enabled.
+         *
+         * @see #NSD_AGGRESSIVE_QUERY_MODE
+         */
+        public Builder setIsAggressiveQueryModeEnabled(boolean isAggressiveQueryModeEnabled) {
+            mIsAggressiveQueryModeEnabled = isAggressiveQueryModeEnabled;
+            return this;
+        }
+
+        /**
          * Builds a {@link MdnsFeatureFlags} with the arguments supplied to this builder.
          */
         public MdnsFeatureFlags build() {
@@ -233,6 +262,7 @@
                     mIsLabelCountLimitEnabled,
                     mIsKnownAnswerSuppressionEnabled,
                     mIsUnicastReplyEnabled,
+                    mIsAggressiveQueryModeEnabled,
                     mOverrideProvider);
         }
     }
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 730bd7e..aa51c41 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
@@ -22,10 +22,12 @@
 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;
@@ -37,6 +39,7 @@
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -44,6 +47,9 @@
  */
 @RequiresApi(Build.VERSION_CODES.TIRAMISU)
 public class MdnsInterfaceAdvertiser implements MulticastPacketReader.PacketHandler {
+    public static final int CONFLICT_SERVICE = 1 << 0;
+    public static final int CONFLICT_HOST = 1 << 1;
+
     private static final boolean DBG = MdnsAdvertiser.DBG;
     @VisibleForTesting
     public static final long EXIT_ANNOUNCEMENT_DELAY_MS = 100L;
@@ -85,10 +91,15 @@
         /**
          * Called by the advertiser when a conflict was found, during or after probing.
          *
-         * If a conflict is found during probing, the {@link #renameServiceForConflict} must be
+         * <p>If a conflict is found during probing, the {@link #renameServiceForConflict} must be
          * called to restart probing and attempt registration with a different name.
+         *
+         * <p>{@code conflictType} is a bitmap telling which part of the service is conflicting. See
+         * {@link MdnsInterfaceAdvertiser#CONFLICT_SERVICE} and {@link
+         * MdnsInterfaceAdvertiser#CONFLICT_HOST}.
          */
-        void onServiceConflict(@NonNull MdnsInterfaceAdvertiser advertiser, int serviceId);
+        void onServiceConflict(
+                @NonNull MdnsInterfaceAdvertiser advertiser, int serviceId, int conflictType);
 
         /**
          * Called by the advertiser when it destroyed itself.
@@ -336,6 +347,7 @@
         final MdnsProber.ProbingInfo probingInfo = mRecordRepository.setServiceProbing(serviceId);
         if (probingInfo == null) return false;
 
+        mAnnouncer.stop(serviceId);
         mProber.restartForConflict(probingInfo);
         return true;
     }
@@ -384,8 +396,16 @@
                     + packet.additionalRecords.size() + " additional from " + srcCopy);
         }
 
-        for (int conflictServiceId : mRecordRepository.getConflictingServices(packet)) {
-            mCbHandler.post(() -> mCb.onServiceConflict(this, conflictServiceId));
+        Map<Integer, Integer> conflictingServices =
+                mRecordRepository.getConflictingServices(packet);
+
+        for (Map.Entry<Integer, Integer> entry : conflictingServices.entrySet()) {
+            int serviceId = entry.getKey();
+            int conflictType = entry.getValue();
+            mCbHandler.post(
+                    () -> {
+                        mCb.onServiceConflict(this, serviceId, conflictType);
+                    });
         }
 
         // Even in case of conflict, add replies for other services. But in general conflicts would
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 78c3082..fb45454 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsRecordRepository.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsRecordRepository.java
@@ -19,6 +19,8 @@
 import static com.android.server.connectivity.mdns.MdnsConstants.IPV4_SOCKET_ADDR;
 import static com.android.server.connectivity.mdns.MdnsConstants.IPV6_SOCKET_ADDR;
 import static com.android.server.connectivity.mdns.MdnsConstants.NO_PACKET;
+import static com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_HOST;
+import static com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_SERVICE;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -28,6 +30,8 @@
 import android.os.Build;
 import android.os.Looper;
 import android.os.SystemClock;
+import android.text.TextUtils;
+import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.SparseArray;
 
@@ -54,6 +58,8 @@
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 
 /**
  * A repository of records advertised through {@link MdnsInterfaceAdvertiser}.
@@ -158,11 +164,13 @@
         public final List<RecordInfo<?>> allRecords;
         @NonNull
         public final List<RecordInfo<MdnsPointerRecord>> ptrRecords;
-        @NonNull
+        @Nullable
         public final RecordInfo<MdnsServiceRecord> srvRecord;
-        @NonNull
+        @Nullable
         public final RecordInfo<MdnsTextRecord> txtRecord;
         @NonNull
+        public final List<RecordInfo<MdnsInetAddressRecord>> addressRecords;
+        @NonNull
         public final NsdServiceInfo serviceInfo;
 
         /**
@@ -202,65 +210,96 @@
                 int repliedServiceCount, int sentPacketCount, boolean exiting, boolean isProbing) {
             this.serviceInfo = serviceInfo;
 
-            final String[] serviceType = splitServiceType(serviceInfo);
-            final String[] serviceName = splitFullyQualifiedName(serviceInfo, serviceType);
+            final boolean hasService = !TextUtils.isEmpty(serviceInfo.getServiceType());
+            final boolean hasCustomHost = !TextUtils.isEmpty(serviceInfo.getHostname());
+            final String[] hostname =
+                    hasCustomHost
+                            ? new String[] {serviceInfo.getHostname(), LOCAL_TLD}
+                            : deviceHostname;
+            final ArrayList<RecordInfo<?>> allRecords = new ArrayList<>(5);
 
-            // Service PTR records
-            ptrRecords = new ArrayList<>(serviceInfo.getSubtypes().size() + 1);
-            ptrRecords.add(new RecordInfo<>(
-                    serviceInfo,
-                    new MdnsPointerRecord(
-                            serviceType,
-                            0L /* receiptTimeMillis */,
-                            false /* cacheFlush */,
-                            NON_NAME_RECORDS_TTL_MILLIS,
-                            serviceName),
-                    true /* sharedName */));
-            for (String subtype : serviceInfo.getSubtypes()) {
+            if (hasService) {
+                final String[] serviceType = splitServiceType(serviceInfo);
+                final String[] serviceName = splitFullyQualifiedName(serviceInfo, serviceType);
+                // Service PTR records
+                ptrRecords = new ArrayList<>(serviceInfo.getSubtypes().size() + 1);
                 ptrRecords.add(new RecordInfo<>(
-                    serviceInfo,
-                    new MdnsPointerRecord(
-                            MdnsUtils.constructFullSubtype(serviceType, subtype),
-                            0L /* receiptTimeMillis */,
-                            false /* cacheFlush */,
-                            NON_NAME_RECORDS_TTL_MILLIS,
-                            serviceName),
-                    true /* sharedName */));
+                        serviceInfo,
+                        new MdnsPointerRecord(
+                                serviceType,
+                                0L /* receiptTimeMillis */,
+                                false /* cacheFlush */,
+                                NON_NAME_RECORDS_TTL_MILLIS,
+                                serviceName),
+                        true /* sharedName */));
+                for (String subtype : serviceInfo.getSubtypes()) {
+                    ptrRecords.add(new RecordInfo<>(
+                            serviceInfo,
+                            new MdnsPointerRecord(
+                                    MdnsUtils.constructFullSubtype(serviceType, subtype),
+                                    0L /* receiptTimeMillis */,
+                                    false /* cacheFlush */,
+                                    NON_NAME_RECORDS_TTL_MILLIS,
+                                    serviceName),
+                            true /* sharedName */));
+                }
+
+                srvRecord = new RecordInfo<>(
+                        serviceInfo,
+                        new MdnsServiceRecord(serviceName,
+                                0L /* receiptTimeMillis */,
+                                true /* cacheFlush */,
+                                NAME_RECORDS_TTL_MILLIS,
+                                0 /* servicePriority */, 0 /* serviceWeight */,
+                                serviceInfo.getPort(),
+                                hostname),
+                        false /* sharedName */);
+
+                txtRecord = new RecordInfo<>(
+                        serviceInfo,
+                        new MdnsTextRecord(serviceName,
+                                0L /* receiptTimeMillis */,
+                                // Service name is verified unique after probing
+                                true /* cacheFlush */,
+                                NON_NAME_RECORDS_TTL_MILLIS,
+                                attrsToTextEntries(serviceInfo.getAttributes())),
+                        false /* sharedName */);
+
+                allRecords.addAll(ptrRecords);
+                allRecords.add(srvRecord);
+                allRecords.add(txtRecord);
+                // Service type enumeration record (RFC6763 9.)
+                allRecords.add(new RecordInfo<>(
+                        serviceInfo,
+                        new MdnsPointerRecord(
+                                DNS_SD_SERVICE_TYPE,
+                                0L /* receiptTimeMillis */,
+                                false /* cacheFlush */,
+                                NON_NAME_RECORDS_TTL_MILLIS,
+                                serviceType),
+                        true /* sharedName */));
+            } else {
+                ptrRecords = Collections.emptyList();
+                srvRecord = null;
+                txtRecord = null;
             }
 
-            srvRecord = new RecordInfo<>(
-                    serviceInfo,
-                    new MdnsServiceRecord(serviceName,
-                            0L /* receiptTimeMillis */,
-                            true /* cacheFlush */,
-                            NAME_RECORDS_TTL_MILLIS, 0 /* servicePriority */, 0 /* serviceWeight */,
-                            serviceInfo.getPort(),
-                            deviceHostname),
-                    false /* sharedName */);
-
-            txtRecord = new RecordInfo<>(
-                    serviceInfo,
-                    new MdnsTextRecord(serviceName,
-                            0L /* receiptTimeMillis */,
-                            true /* cacheFlush */, // Service name is verified unique after probing
-                            NON_NAME_RECORDS_TTL_MILLIS,
-                            attrsToTextEntries(serviceInfo.getAttributes())),
-                    false /* sharedName */);
-
-            final ArrayList<RecordInfo<?>> allRecords = new ArrayList<>(5);
-            allRecords.addAll(ptrRecords);
-            allRecords.add(srvRecord);
-            allRecords.add(txtRecord);
-            // Service type enumeration record (RFC6763 9.)
-            allRecords.add(new RecordInfo<>(
-                    serviceInfo,
-                    new MdnsPointerRecord(
-                            DNS_SD_SERVICE_TYPE,
-                            0L /* receiptTimeMillis */,
-                            false /* cacheFlush */,
-                            NON_NAME_RECORDS_TTL_MILLIS,
-                            serviceType),
-                    true /* sharedName */));
+            if (hasCustomHost) {
+                addressRecords = new ArrayList<>(serviceInfo.getHostAddresses().size());
+                for (InetAddress address : serviceInfo.getHostAddresses()) {
+                    addressRecords.add(new RecordInfo<>(
+                                    serviceInfo,
+                                    new MdnsInetAddressRecord(hostname,
+                                            0L /* receiptTimeMillis */,
+                                            true /* cacheFlush */,
+                                            NAME_RECORDS_TTL_MILLIS,
+                                            address),
+                                    false /* sharedName */));
+                }
+                allRecords.addAll(addressRecords);
+            } else {
+                addressRecords = Collections.emptyList();
+            }
 
             this.allRecords = Collections.unmodifiableList(allRecords);
             this.repliedServiceCount = repliedServiceCount;
@@ -368,32 +407,38 @@
     /**
      * @return The ID of the service identified by its name, or -1 if none.
      */
-    private int getServiceByName(@NonNull String serviceName) {
+    private int getServiceByName(@Nullable String serviceName) {
+        if (TextUtils.isEmpty(serviceName)) {
+            return -1;
+        }
         for (int i = 0; i < mServices.size(); i++) {
             final ServiceRegistration registration = mServices.valueAt(i);
-            if (MdnsUtils.equalsIgnoreDnsCase(serviceName,
-                    registration.serviceInfo.getServiceName())) {
+            if (MdnsUtils.equalsIgnoreDnsCase(
+                    serviceName, registration.serviceInfo.getServiceName())) {
                 return mServices.keyAt(i);
             }
         }
         return -1;
     }
 
-    private MdnsProber.ProbingInfo makeProbingInfo(int serviceId,
-            @NonNull MdnsServiceRecord srvRecord,
-            @NonNull List<MdnsInetAddressRecord> inetAddressRecords) {
+    private MdnsProber.ProbingInfo makeProbingInfo(
+            int serviceId, ServiceRegistration registration) {
         final List<MdnsRecord> probingRecords = new ArrayList<>();
         // Probe with cacheFlush cleared; it is set when announcing, as it was verified unique:
         // RFC6762 10.2
-        probingRecords.add(new MdnsServiceRecord(srvRecord.getName(),
-                0L /* receiptTimeMillis */,
-                false /* cacheFlush */,
-                srvRecord.getTtl(),
-                srvRecord.getServicePriority(), srvRecord.getServiceWeight(),
-                srvRecord.getServicePort(),
-                srvRecord.getServiceHost()));
+        if (registration.srvRecord != null) {
+            MdnsServiceRecord srvRecord = registration.srvRecord.record;
+            probingRecords.add(new MdnsServiceRecord(srvRecord.getName(),
+                    0L /* receiptTimeMillis */,
+                    false /* cacheFlush */,
+                    srvRecord.getTtl(),
+                    srvRecord.getServicePriority(), srvRecord.getServiceWeight(),
+                    srvRecord.getServicePort(),
+                    srvRecord.getServiceHost()));
+        }
 
-        for (MdnsInetAddressRecord inetAddressRecord : inetAddressRecords) {
+        for (MdnsInetAddressRecord inetAddressRecord :
+                makeProbingInetAddressRecords(registration.serviceInfo)) {
             probingRecords.add(new MdnsInetAddressRecord(inetAddressRecord.getName(),
                     0L /* receiptTimeMillis */,
                     false /* cacheFlush */,
@@ -510,6 +555,9 @@
     public MdnsReplyInfo getReply(MdnsPacket packet, InetSocketAddress src) {
         final long now = SystemClock.elapsedRealtime();
 
+        // TODO: b/322142420 - Set<RecordInfo<?>> may contain duplicate records wrapped in different
+        // RecordInfo<?>s when custom host is enabled.
+
         // Use LinkedHashSet for preserving the insert order of the RRs, so that RRs of the same
         // service or host are grouped together (which is more developer-friendly).
         final Set<RecordInfo<?>> answerInfo = new LinkedHashSet<>();
@@ -520,8 +568,10 @@
         for (MdnsRecord question : packet.questions) {
             // Add answers from general records
             if (addReplyFromService(question, mGeneralRecords, null /* servicePtrRecord */,
-                    null /* serviceSrvRecord */, null /* serviceTxtRecord */, replyUnicastEnabled,
-                    now, answerInfo, additionalAnswerInfo, Collections.emptyList())) {
+                    null /* serviceSrvRecord */, null /* serviceTxtRecord */,
+                    null /* hostname */,
+                    replyUnicastEnabled, now, answerInfo, additionalAnswerInfo,
+                    Collections.emptyList())) {
                 replyUnicast &= question.isUnicastReplyRequested();
             }
 
@@ -530,7 +580,9 @@
                 final ServiceRegistration registration = mServices.valueAt(i);
                 if (registration.exiting || registration.isProbing) continue;
                 if (addReplyFromService(question, registration.allRecords, registration.ptrRecords,
-                        registration.srvRecord, registration.txtRecord, replyUnicastEnabled, now,
+                        registration.srvRecord, registration.txtRecord,
+                        registration.serviceInfo.getHostname(),
+                        replyUnicastEnabled, now,
                         answerInfo, additionalAnswerInfo, packet.answers)) {
                     replyUnicast &= question.isUnicastReplyRequested();
                     registration.repliedServiceCount++;
@@ -548,7 +600,12 @@
         final List<MdnsRecord> additionalAnswerRecords =
                 new ArrayList<>(additionalAnswerInfo.size());
         for (RecordInfo<?> info : additionalAnswerInfo) {
-            additionalAnswerRecords.add(info.record);
+            // Different RecordInfos may contain the same record.
+            // For example, when there are multiple services referring to the same custom host,
+            // there are multiple RecordInfos containing the same address record.
+            if (!additionalAnswerRecords.contains(info.record)) {
+                additionalAnswerRecords.add(info.record);
+            }
         }
 
         // RFC6762 6.1: negative responses
@@ -618,7 +675,10 @@
             if (!replyUnicast) {
                 info.lastAdvertisedTimeMs = info.lastSentTimeMs;
             }
-            answerRecords.add(info.record);
+            // Different RecordInfos may the contain the same record
+            if (!answerRecords.contains(info.record)) {
+                answerRecords.add(info.record);
+            }
         }
 
         return new MdnsReplyInfo(answerRecords, additionalAnswerRecords, delayMs, dest, src,
@@ -642,6 +702,7 @@
             @Nullable List<RecordInfo<MdnsPointerRecord>> servicePtrRecords,
             @Nullable RecordInfo<MdnsServiceRecord> serviceSrvRecord,
             @Nullable RecordInfo<MdnsTextRecord> serviceTxtRecord,
+            @Nullable String hostname,
             boolean replyUnicastEnabled, long now, @NonNull Set<RecordInfo<?>> answerInfo,
             @NonNull Set<RecordInfo<?>> additionalAnswerInfo,
             @NonNull List<MdnsRecord> knownAnswerRecords) {
@@ -735,11 +796,7 @@
 
         // RFC6763 12.1&.2: if including PTR or SRV record, include the address records it names
         if (hasDnsSdPtrRecordAnswer || hasDnsSdSrvRecordAnswer) {
-            for (RecordInfo<?> record : mGeneralRecords) {
-                if (record.record instanceof MdnsInetAddressRecord) {
-                    additionalAnswerInfo.add(record);
-                }
-            }
+            additionalAnswerInfo.addAll(getInetAddressRecordsForHostname(hostname));
         }
         return true;
     }
@@ -853,29 +910,46 @@
             MdnsProber.ProbingInfo probeSuccessInfo)
             throws IOException {
 
-        final ServiceRegistration registration = mServices.get(probeSuccessInfo.getServiceId());
-        if (registration == null) throw new IOException(
-                "Service is not registered: " + probeSuccessInfo.getServiceId());
+        int serviceId = probeSuccessInfo.getServiceId();
+        final ServiceRegistration registration = mServices.get(serviceId);
+        if (registration == null) {
+            throw new IOException("Service is not registered: " + serviceId);
+        }
         registration.setProbing(false);
 
-        final ArrayList<MdnsRecord> answers = new ArrayList<>();
+        final Set<MdnsRecord> answersSet = new LinkedHashSet<>();
         final ArrayList<MdnsRecord> additionalAnswers = new ArrayList<>();
 
-        // Interface address records in general records
-        for (RecordInfo<?> record : mGeneralRecords) {
-            answers.add(record.record);
+        // When using default host, add interface address records from general records
+        if (TextUtils.isEmpty(registration.serviceInfo.getHostname())) {
+            for (RecordInfo<?> record : mGeneralRecords) {
+                answersSet.add(record.record);
+            }
+        } else {
+            // TODO: b/321617573 - include PTR records for addresses
+            // The custom host may have more addresses in other registrations
+            forEachActiveServiceRegistrationWithHostname(
+                    registration.serviceInfo.getHostname(),
+                    (id, otherRegistration) -> {
+                        if (otherRegistration.isProbing) {
+                            return;
+                        }
+                        for (RecordInfo<?> addressRecordInfo : otherRegistration.addressRecords) {
+                            answersSet.add(addressRecordInfo.record);
+                        }
+                    });
         }
 
         // All service records
         for (RecordInfo<?> info : registration.allRecords) {
-            answers.add(info.record);
+            answersSet.add(info.record);
         }
 
         addNsecRecordsForUniqueNames(additionalAnswers,
                 mGeneralRecords.iterator(), registration.allRecords.iterator());
 
-        return new MdnsAnnouncer.AnnouncementInfo(probeSuccessInfo.getServiceId(),
-                answers, additionalAnswers);
+        return new MdnsAnnouncer.AnnouncementInfo(
+                probeSuccessInfo.getServiceId(), new ArrayList<>(answersSet), additionalAnswers);
     }
 
     /**
@@ -894,8 +968,13 @@
         for (RecordInfo<MdnsPointerRecord> ptrRecord : registration.ptrRecords) {
             answers.add(ptrRecord.record);
         }
-        answers.add(registration.srvRecord.record);
-        answers.add(registration.txtRecord.record);
+        if (registration.srvRecord != null) {
+            answers.add(registration.srvRecord.record);
+        }
+        if (registration.txtRecord != null) {
+            answers.add(registration.txtRecord.record);
+        }
+        // TODO: Support custom host. It currently only supports default host.
         for (RecordInfo<?> record : mGeneralRecords) {
             if (record.record instanceof MdnsInetAddressRecord) {
                 answers.add(record.record);
@@ -910,70 +989,181 @@
                 Collections.emptyList() /* additionalRecords */);
     }
 
+    /** Check if the record is in any service registration */
+    private boolean hasInetAddressRecord(@NonNull MdnsInetAddressRecord record) {
+        for (int i = 0; i < mServices.size(); i++) {
+            final ServiceRegistration registration = mServices.valueAt(i);
+            if (registration.exiting) continue;
+
+            for (RecordInfo<MdnsInetAddressRecord> localRecord : registration.addressRecords) {
+                if (Objects.equals(localRecord.record, record)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     /**
      * Get the service IDs of services conflicting with a received packet.
+     *
+     * <p>It returns a Map of service ID => conflict type. Conflict type is a bitmap telling which
+     * part of the service is conflicting. See {@link MdnsInterfaceAdvertiser#CONFLICT_SERVICE} and
+     * {@link MdnsInterfaceAdvertiser#CONFLICT_HOST}.
      */
-    public Set<Integer> getConflictingServices(MdnsPacket packet) {
+    public Map<Integer, Integer> getConflictingServices(MdnsPacket packet) {
         // Avoid allocating a new set for each incoming packet: use an empty set by default.
-        Set<Integer> conflicting = Collections.emptySet();
+        Map<Integer, Integer> conflicting = Collections.emptyMap();
         for (MdnsRecord record : packet.answers) {
             for (int i = 0; i < mServices.size(); i++) {
                 final ServiceRegistration registration = mServices.valueAt(i);
                 if (registration.exiting) continue;
 
-                // Only look for conflicts in service name, as a different service name can be used
-                // if there is a conflict, but there is nothing actionable if any other conflict
-                // happens. In fact probing is only done for the service name in the SRV record.
-                // This means only SRV and TXT records need to be checked.
-                final RecordInfo<MdnsServiceRecord> srvRecord = registration.srvRecord;
-                if (!MdnsUtils.equalsDnsLabelIgnoreDnsCase(record.getName(),
-                        srvRecord.record.getName())) {
-                    continue;
+                int conflictType = 0;
+
+                if (conflictForService(record, registration)) {
+                    conflictType |= CONFLICT_SERVICE;
                 }
 
-                // As per RFC6762 9., it's fine if the "conflict" is an identical record with same
-                // data.
-                if (record instanceof MdnsServiceRecord) {
-                    final MdnsServiceRecord local = srvRecord.record;
-                    final MdnsServiceRecord other = (MdnsServiceRecord) record;
-                    // Note "equals" does not consider TTL or receipt time, as intended here
-                    if (Objects.equals(local, other)) {
-                        continue;
+                if (conflictForHost(record, registration)) {
+                    conflictType |= CONFLICT_HOST;
+                }
+
+                if (conflictType != 0) {
+                    if (conflicting.isEmpty()) {
+                        // Conflict was found: use a mutable set
+                        conflicting = new ArrayMap<>();
                     }
+                    final int serviceId = mServices.keyAt(i);
+                    conflicting.put(serviceId, conflictType);
                 }
-
-                if (record instanceof MdnsTextRecord) {
-                    final MdnsTextRecord local = registration.txtRecord.record;
-                    final MdnsTextRecord other = (MdnsTextRecord) record;
-                    if (Objects.equals(local, other)) {
-                        continue;
-                    }
-                }
-
-                if (conflicting.size() == 0) {
-                    // Conflict was found: use a mutable set
-                    conflicting = new ArraySet<>();
-                }
-                final int serviceId = mServices.keyAt(i);
-                conflicting.add(serviceId);
             }
         }
 
         return conflicting;
     }
 
-    private List<MdnsInetAddressRecord> makeProbingInetAddressRecords() {
-        final List<MdnsInetAddressRecord> records = new ArrayList<>();
-        if (mMdnsFeatureFlags.mIncludeInetAddressRecordsInProbing) {
-            for (RecordInfo<?> record : mGeneralRecords) {
-                if (record.record instanceof MdnsInetAddressRecord) {
-                    records.add((MdnsInetAddressRecord) record.record);
-                }
+
+    private static boolean conflictForService(
+            @NonNull MdnsRecord record, @NonNull ServiceRegistration registration) {
+        if (registration.srvRecord == null) {
+            return false;
+        }
+
+        final RecordInfo<MdnsServiceRecord> srvRecord = registration.srvRecord;
+        if (!MdnsUtils.equalsDnsLabelIgnoreDnsCase(record.getName(), srvRecord.record.getName())) {
+            return false;
+        }
+
+        // As per RFC6762 9., it's fine if the "conflict" is an identical record with same
+        // data.
+        if (record instanceof MdnsServiceRecord) {
+            final MdnsServiceRecord local = srvRecord.record;
+            final MdnsServiceRecord other = (MdnsServiceRecord) record;
+            // Note "equals" does not consider TTL or receipt time, as intended here
+            if (Objects.equals(local, other)) {
+                return false;
             }
         }
+
+        if (record instanceof MdnsTextRecord) {
+            final MdnsTextRecord local = registration.txtRecord.record;
+            final MdnsTextRecord other = (MdnsTextRecord) record;
+            if (Objects.equals(local, other)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    private boolean conflictForHost(
+            @NonNull MdnsRecord record, @NonNull ServiceRegistration registration) {
+        // Only custom hosts are checked. When using the default host, the hostname is derived from
+        // a UUID and it's supposed to be unique.
+        if (registration.serviceInfo.getHostname() == null) {
+            return false;
+        }
+
+        // The record's name cannot be registered by NsdManager so it's not a conflict.
+        if (record.getName().length != 2 || !record.getName()[1].equals(LOCAL_TLD)) {
+            return false;
+        }
+
+        // Different names. There won't be a conflict.
+        if (!MdnsUtils.equalsIgnoreDnsCase(
+                record.getName()[0], registration.serviceInfo.getHostname())) {
+            return false;
+        }
+
+        // If this registration has any address record and there's no identical record in the
+        // repository, it's a conflict. There will be no conflict if no registration has addresses
+        // for that hostname.
+        if (record instanceof MdnsInetAddressRecord) {
+            if (!registration.addressRecords.isEmpty()) {
+                return !hasInetAddressRecord((MdnsInetAddressRecord) record);
+            }
+        }
+
+        return false;
+    }
+
+    private List<RecordInfo<MdnsInetAddressRecord>> getInetAddressRecordsForHostname(
+            @Nullable String hostname) {
+        List<RecordInfo<MdnsInetAddressRecord>> records = new ArrayList<>();
+        if (TextUtils.isEmpty(hostname)) {
+            forEachAddressRecord(mGeneralRecords, records::add);
+        } else {
+            forEachActiveServiceRegistrationWithHostname(
+                    hostname,
+                    (id, service) -> {
+                        if (service.isProbing) return;
+                        records.addAll(service.addressRecords);
+                    });
+        }
         return records;
     }
 
+    private List<MdnsInetAddressRecord> makeProbingInetAddressRecords(
+            @NonNull NsdServiceInfo serviceInfo) {
+        final List<MdnsInetAddressRecord> records = new ArrayList<>();
+        if (TextUtils.isEmpty(serviceInfo.getHostname())) {
+            if (mMdnsFeatureFlags.mIncludeInetAddressRecordsInProbing) {
+                forEachAddressRecord(mGeneralRecords, r -> records.add(r.record));
+            }
+        } else {
+            forEachActiveServiceRegistrationWithHostname(
+                    serviceInfo.getHostname(),
+                    (id, service) -> {
+                        for (RecordInfo<MdnsInetAddressRecord> recordInfo :
+                                service.addressRecords) {
+                            records.add(recordInfo.record);
+                        }
+                    });
+        }
+        return records;
+    }
+
+    private static void forEachAddressRecord(
+            List<RecordInfo<?>> records, Consumer<RecordInfo<MdnsInetAddressRecord>> consumer) {
+        for (RecordInfo<?> record : records) {
+            if (record.record instanceof MdnsInetAddressRecord) {
+                consumer.accept((RecordInfo<MdnsInetAddressRecord>) record);
+            }
+        }
+    }
+
+    private void forEachActiveServiceRegistrationWithHostname(
+            @NonNull String hostname, BiConsumer<Integer, ServiceRegistration> consumer) {
+        for (int i = 0; i < mServices.size(); ++i) {
+            int id = mServices.keyAt(i);
+            ServiceRegistration service = mServices.valueAt(i);
+            if (service.exiting) continue;
+            if (MdnsUtils.equalsIgnoreDnsCase(service.serviceInfo.getHostname(), hostname)) {
+                consumer.accept(id, service);
+            }
+        }
+    }
+
     /**
      * (Re)set a service to the probing state.
      * @return The {@link MdnsProber.ProbingInfo} to send for probing.
@@ -984,8 +1174,8 @@
         if (registration == null) return null;
 
         registration.setProbing(true);
-        return makeProbingInfo(
-                serviceId, registration.srvRecord.record, makeProbingInetAddressRecords());
+
+        return makeProbingInfo(serviceId, registration);
     }
 
     /**
@@ -1021,8 +1211,7 @@
         final ServiceRegistration newService = new ServiceRegistration(mDeviceHostname, newInfo,
                 existing.repliedServiceCount, existing.sentPacketCount);
         mServices.put(serviceId, newService);
-        return makeProbingInfo(
-                serviceId, newService.srvRecord.record, makeProbingInetAddressRecords());
+        return makeProbingInfo(serviceId, newService);
     }
 
     /**
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 63835d9..086094b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSearchOptions.java
@@ -39,6 +39,15 @@
  * @hide
  */
 public class MdnsSearchOptions implements Parcelable {
+    // Passive query mode scans less frequently in order to conserve battery and produce less
+    // network traffic.
+    public static final int PASSIVE_QUERY_MODE = 0;
+    // Active query mode scans frequently.
+    public static final int ACTIVE_QUERY_MODE = 1;
+    // Aggressive query mode scans more frequently than the active mode at first, and sends both
+    // unicast and multicast queries simultaneously, but in long sessions it eventually sends as
+    // many queries as the PASSIVE mode.
+    public static final int AGGRESSIVE_QUERY_MODE = 2;
 
     /** @hide */
     public static final Parcelable.Creator<MdnsSearchOptions> CREATOR =
@@ -47,7 +56,7 @@
                 public MdnsSearchOptions createFromParcel(Parcel source) {
                     return new MdnsSearchOptions(
                             source.createStringArrayList(),
-                            source.readInt() == 1,
+                            source.readInt(),
                             source.readInt() == 1,
                             source.readParcelable(null),
                             source.readString(),
@@ -64,7 +73,7 @@
     private final List<String> subtypes;
     @Nullable
     private final String resolveInstanceName;
-    private final boolean isPassiveMode;
+    private final int queryMode;
     private final boolean onlyUseIpv6OnIpv6OnlyNetworks;
     private final int numOfQueriesBeforeBackoff;
     private final boolean removeExpiredService;
@@ -74,7 +83,7 @@
     /** Parcelable constructs for a {@link MdnsSearchOptions}. */
     MdnsSearchOptions(
             List<String> subtypes,
-            boolean isPassiveMode,
+            int queryMode,
             boolean removeExpiredService,
             @Nullable Network network,
             @Nullable String resolveInstanceName,
@@ -84,7 +93,7 @@
         if (subtypes != null) {
             this.subtypes.addAll(subtypes);
         }
-        this.isPassiveMode = isPassiveMode;
+        this.queryMode = queryMode;
         this.onlyUseIpv6OnIpv6OnlyNetworks = onlyUseIpv6OnIpv6OnlyNetworks;
         this.numOfQueriesBeforeBackoff = numOfQueriesBeforeBackoff;
         this.removeExpiredService = removeExpiredService;
@@ -111,11 +120,10 @@
     }
 
     /**
-     * @return {@code true} if the passive mode is used. The passive mode scans less frequently in
-     * order to conserve battery and produce less network traffic.
+     * @return the current query mode.
      */
-    public boolean isPassiveMode() {
-        return isPassiveMode;
+    public int getQueryMode() {
+        return queryMode;
     }
 
     /**
@@ -166,7 +174,7 @@
     @Override
     public void writeToParcel(Parcel out, int flags) {
         out.writeStringList(subtypes);
-        out.writeInt(isPassiveMode ? 1 : 0);
+        out.writeInt(queryMode);
         out.writeInt(removeExpiredService ? 1 : 0);
         out.writeParcelable(mNetwork, 0);
         out.writeString(resolveInstanceName);
@@ -177,7 +185,7 @@
     /** A builder to create {@link MdnsSearchOptions}. */
     public static final class Builder {
         private final Set<String> subtypes;
-        private boolean isPassiveMode = true;
+        private int queryMode = PASSIVE_QUERY_MODE;
         private boolean onlyUseIpv6OnIpv6OnlyNetworks = false;
         private int numOfQueriesBeforeBackoff = 3;
         private boolean removeExpiredService;
@@ -212,14 +220,12 @@
         }
 
         /**
-         * Sets if the passive mode scan should be used. The passive mode scans less frequently in
-         * order to conserve battery and produce less network traffic.
+         * Sets which query mode should be used.
          *
-         * @param isPassiveMode If set to {@code true}, passive mode will be used. If set to {@code
-         *                      false}, active mode will be used.
+         * @param queryMode the query mode should be used.
          */
-        public Builder setIsPassiveMode(boolean isPassiveMode) {
-            this.isPassiveMode = isPassiveMode;
+        public Builder setQueryMode(int queryMode) {
+            this.queryMode = queryMode;
             return this;
         }
 
@@ -276,7 +282,7 @@
         public MdnsSearchOptions build() {
             return new MdnsSearchOptions(
                     new ArrayList<>(subtypes),
-                    isPassiveMode,
+                    queryMode,
                     removeExpiredService,
                     mNetwork,
                     resolveInstanceName,
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 df0a040..3a69d67 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -81,7 +81,7 @@
                     notifyRemovedServiceToListeners(previousResponse, "Service record expired");
                 }
             };
-    private final ArrayMap<MdnsServiceBrowserListener, MdnsSearchOptions> listeners =
+    private final ArrayMap<MdnsServiceBrowserListener, ListenerInfo> listeners =
             new ArrayMap<>();
     private final boolean removeServiceAfterTtlExpires =
             MdnsConfigs.removeServiceAfterTtlExpires();
@@ -95,6 +95,32 @@
     private long currentSessionId = 0;
     private long lastSentTime;
 
+    private static class ListenerInfo {
+        @NonNull
+        final MdnsSearchOptions searchOptions;
+        final Set<String> discoveredServiceNames;
+
+        ListenerInfo(@NonNull MdnsSearchOptions searchOptions,
+                @Nullable ListenerInfo previousInfo) {
+            this.searchOptions = searchOptions;
+            this.discoveredServiceNames = previousInfo == null
+                    ? MdnsUtils.newSet() : previousInfo.discoveredServiceNames;
+        }
+
+        /**
+         * Set the given service name as discovered.
+         *
+         * @return true if the service name was not discovered before.
+         */
+        boolean setServiceDiscovered(@NonNull String serviceName) {
+            return discoveredServiceNames.add(MdnsUtils.toDnsLowerCase(serviceName));
+        }
+
+        void unsetServiceDiscovered(@NonNull String serviceName) {
+            discoveredServiceNames.remove(MdnsUtils.toDnsLowerCase(serviceName));
+        }
+    }
+
     private class QueryTaskHandler extends Handler {
         QueryTaskHandler(Looper looper) {
             super(looper);
@@ -311,12 +337,16 @@
         ensureRunningOnHandlerThread(handler);
         this.searchOptions = searchOptions;
         boolean hadReply = false;
-        if (listeners.put(listener, searchOptions) == null) {
+        final ListenerInfo existingInfo = listeners.get(listener);
+        final ListenerInfo listenerInfo = new ListenerInfo(searchOptions, existingInfo);
+        listeners.put(listener, listenerInfo);
+        if (existingInfo == null) {
             for (MdnsResponse existingResponse : serviceCache.getCachedServices(cacheKey)) {
                 if (!responseMatchesOptions(existingResponse, searchOptions)) continue;
                 final MdnsServiceInfo info =
                         buildMdnsServiceInfoFromResponse(existingResponse, serviceTypeLabels);
                 listener.onServiceNameDiscovered(info, true /* isServiceFromCache */);
+                listenerInfo.setServiceDiscovered(info.getServiceInstanceName());
                 if (existingResponse.isComplete()) {
                     listener.onServiceFound(info, true /* isServiceFromCache */);
                     hadReply = true;
@@ -330,7 +360,7 @@
         // interested anymore.
         final QueryTaskConfig taskConfig = new QueryTaskConfig(
                 searchOptions.getSubtypes(),
-                searchOptions.isPassiveMode(),
+                searchOptions.getQueryMode(),
                 searchOptions.onlyUseIpv6OnIpv6OnlyNetworks(),
                 searchOptions.numOfQueriesBeforeBackoff(),
                 socketKey);
@@ -480,9 +510,10 @@
     private void notifyRemovedServiceToListeners(@NonNull MdnsResponse response,
             @NonNull String message) {
         for (int i = 0; i < listeners.size(); i++) {
-            if (!responseMatchesOptions(response, listeners.valueAt(i))) continue;
+            if (!responseMatchesOptions(response, listeners.valueAt(i).searchOptions)) continue;
             final MdnsServiceBrowserListener listener = listeners.keyAt(i);
             if (response.getServiceInstanceName() != null) {
+                listeners.valueAt(i).unsetServiceDiscovered(response.getServiceInstanceName());
                 final MdnsServiceInfo serviceInfo = buildMdnsServiceInfoFromResponse(
                         response, serviceTypeLabels);
                 if (response.isComplete()) {
@@ -511,10 +542,9 @@
         final MdnsResponse currentResponse =
                 serviceCache.getCachedService(serviceInstanceName, cacheKey);
 
-        boolean newServiceFound = false;
+        final boolean newInCache = currentResponse == null;
         boolean serviceBecomesComplete = false;
-        if (currentResponse == null) {
-            newServiceFound = true;
+        if (newInCache) {
             if (serviceInstanceName != null) {
                 serviceCache.addOrUpdateService(cacheKey, response);
             }
@@ -525,25 +555,28 @@
             serviceBecomesComplete = !before && after;
         }
         sharedLog.i(String.format(
-                "Handling response from service: %s, newServiceFound: %b, serviceBecomesComplete:"
+                "Handling response from service: %s, newInCache: %b, serviceBecomesComplete:"
                         + " %b, responseIsComplete: %b",
-                serviceInstanceName, newServiceFound, serviceBecomesComplete,
+                serviceInstanceName, newInCache, serviceBecomesComplete,
                 response.isComplete()));
         MdnsServiceInfo serviceInfo =
                 buildMdnsServiceInfoFromResponse(response, serviceTypeLabels);
 
         for (int i = 0; i < listeners.size(); i++) {
-            if (!responseMatchesOptions(response, listeners.valueAt(i))) continue;
+            // If a service stops matching the options (currently can only happen if it loses a
+            // subtype), service lost callbacks should also be sent; this is not done today as
+            // only expiration of SRV records is used, not PTR records used for subtypes, so
+            // services never lose PTR record subtypes.
+            if (!responseMatchesOptions(response, listeners.valueAt(i).searchOptions)) continue;
             final MdnsServiceBrowserListener listener = listeners.keyAt(i);
+            final ListenerInfo listenerInfo = listeners.valueAt(i);
+            final boolean newServiceFound = listenerInfo.setServiceDiscovered(serviceInstanceName);
             if (newServiceFound) {
                 sharedLog.log("onServiceNameDiscovered: " + serviceInfo);
                 listener.onServiceNameDiscovered(serviceInfo, false /* isServiceFromCache */);
             }
 
             if (response.isComplete()) {
-                // There is a bug here: the newServiceFound is global right now. The state needs
-                // to be per listener because of the  responseMatchesOptions() filter.
-                // Otherwise, it won't handle the subType update properly.
                 if (newServiceFound || serviceBecomesComplete) {
                     sharedLog.log("onServiceFound: " + serviceInfo);
                     listener.onServiceFound(serviceInfo, false /* isServiceFromCache */);
@@ -579,7 +612,7 @@
     private List<MdnsResponse> makeResponsesForResolve(@NonNull SocketKey socketKey) {
         final List<MdnsResponse> resolveResponses = new ArrayList<>();
         for (int i = 0; i < listeners.size(); i++) {
-            final String resolveName = listeners.valueAt(i).getResolveInstanceName();
+            final String resolveName = listeners.valueAt(i).searchOptions.getResolveInstanceName();
             if (resolveName == null) {
                 continue;
             }
diff --git a/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
index 19282b0..10a71a2 100644
--- a/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
+++ b/service-t/src/com/android/server/connectivity/mdns/QueryTaskConfig.java
@@ -16,6 +16,9 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE;
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.PASSIVE_QUERY_MODE;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 
@@ -33,19 +36,29 @@
 
     private static final int INITIAL_TIME_BETWEEN_BURSTS_MS =
             (int) MdnsConfigs.initialTimeBetweenBurstsMs();
-    private static final int TIME_BETWEEN_BURSTS_MS = (int) MdnsConfigs.timeBetweenBurstsMs();
+    private static final int MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS =
+            (int) MdnsConfigs.timeBetweenBurstsMs();
     private static final int QUERIES_PER_BURST = (int) MdnsConfigs.queriesPerBurst();
     private static final int TIME_BETWEEN_QUERIES_IN_BURST_MS =
             (int) MdnsConfigs.timeBetweenQueriesInBurstMs();
     private static final int QUERIES_PER_BURST_PASSIVE_MODE =
             (int) MdnsConfigs.queriesPerBurstPassive();
     private static final int UNSIGNED_SHORT_MAX_VALUE = 65536;
+    @VisibleForTesting
+    // RFC 6762 5.2: The interval between the first two queries MUST be at least one second.
+    static final int INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS = 1000;
+    @VisibleForTesting
+    // Basically this tries to send one query per typical DTIM interval 100ms, to maximize the
+    // chances that a query will be received if devices are using a DTIM multiplier (in which case
+    // they only listen once every [multiplier] DTIM intervals).
+    static final int TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS = 100;
+    static final int MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS = 60000;
     // The following fields are used by QueryTask so we need to test them.
     @VisibleForTesting
     final List<String> subtypes;
     private final boolean alwaysAskForUnicastResponse =
             MdnsConfigs.alwaysAskForUnicastResponseInEachBurst();
-    private final boolean usePassiveMode;
+    private final int queryMode;
     final boolean onlyUseIpv6OnIpv6OnlyNetworks;
     private final int numOfQueriesBeforeBackoff;
     @VisibleForTesting
@@ -66,7 +79,7 @@
             int queriesPerBurst, int timeBetweenBurstsInMs,
             long delayUntilNextTaskWithoutBackoffMs) {
         this.subtypes = new ArrayList<>(other.subtypes);
-        this.usePassiveMode = other.usePassiveMode;
+        this.queryMode = other.queryMode;
         this.onlyUseIpv6OnIpv6OnlyNetworks = other.onlyUseIpv6OnIpv6OnlyNetworks;
         this.numOfQueriesBeforeBackoff = other.numOfQueriesBeforeBackoff;
         this.transactionId = transactionId;
@@ -79,12 +92,13 @@
         this.queryCount = queryCount;
         this.socketKey = other.socketKey;
     }
+
     QueryTaskConfig(@NonNull Collection<String> subtypes,
-            boolean usePassiveMode,
+            int queryMode,
             boolean onlyUseIpv6OnIpv6OnlyNetworks,
             int numOfQueriesBeforeBackoff,
             @Nullable SocketKey socketKey) {
-        this.usePassiveMode = usePassiveMode;
+        this.queryMode = queryMode;
         this.onlyUseIpv6OnIpv6OnlyNetworks = onlyUseIpv6OnIpv6OnlyNetworks;
         this.numOfQueriesBeforeBackoff = numOfQueriesBeforeBackoff;
         this.subtypes = new ArrayList<>(subtypes);
@@ -94,21 +108,59 @@
         this.expectUnicastResponse = true;
         this.isFirstBurst = true;
         // Config the scan frequency based on the scan mode.
-        if (this.usePassiveMode) {
+        if (this.queryMode == AGGRESSIVE_QUERY_MODE) {
+            this.timeBetweenBurstsInMs = INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS;
+            this.delayUntilNextTaskWithoutBackoffMs =
+                    TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS;
+        } else if (this.queryMode == PASSIVE_QUERY_MODE) {
             // In passive scan mode, sends a single burst of QUERIES_PER_BURST queries, and then
             // in each TIME_BETWEEN_BURSTS interval, sends QUERIES_PER_BURST_PASSIVE_MODE
             // queries.
-            this.timeBetweenBurstsInMs = TIME_BETWEEN_BURSTS_MS;
+            this.timeBetweenBurstsInMs = MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS;
+            this.delayUntilNextTaskWithoutBackoffMs = TIME_BETWEEN_QUERIES_IN_BURST_MS;
         } else {
             // In active scan mode, sends a burst of QUERIES_PER_BURST queries,
             // TIME_BETWEEN_QUERIES_IN_BURST_MS apart, then waits for the scan interval, and
             // then repeats. The scan interval starts as INITIAL_TIME_BETWEEN_BURSTS_MS and
             // doubles until it maxes out at TIME_BETWEEN_BURSTS_MS.
             this.timeBetweenBurstsInMs = INITIAL_TIME_BETWEEN_BURSTS_MS;
+            this.delayUntilNextTaskWithoutBackoffMs = TIME_BETWEEN_QUERIES_IN_BURST_MS;
         }
         this.socketKey = socketKey;
         this.queryCount = 0;
-        this.delayUntilNextTaskWithoutBackoffMs = TIME_BETWEEN_QUERIES_IN_BURST_MS;
+    }
+
+    long getDelayUntilNextTaskWithoutBackoff(boolean isFirstQueryInBurst,
+            boolean isLastQueryInBurst) {
+        if (isFirstQueryInBurst && queryMode == AGGRESSIVE_QUERY_MODE) {
+            return 0;
+        }
+        if (isLastQueryInBurst) {
+            return timeBetweenBurstsInMs;
+        }
+        return queryMode == AGGRESSIVE_QUERY_MODE
+                ? TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS
+                : TIME_BETWEEN_QUERIES_IN_BURST_MS;
+    }
+
+    boolean getNextExpectUnicastResponse(boolean isLastQueryInBurst) {
+        if (!isLastQueryInBurst) {
+            return false;
+        }
+        if (queryMode == AGGRESSIVE_QUERY_MODE) {
+            return true;
+        }
+        return alwaysAskForUnicastResponse;
+    }
+
+    int getNextTimeBetweenBurstsMs(boolean isLastQueryInBurst) {
+        if (!isLastQueryInBurst) {
+            return timeBetweenBurstsInMs;
+        }
+        final int maxTimeBetweenBursts = queryMode == AGGRESSIVE_QUERY_MODE
+                ? MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS : MAX_TIME_BETWEEN_ACTIVE_PASSIVE_BURSTS_MS;
+        return timeBetweenBurstsInMs < maxTimeBetweenBursts
+                ? Math.min(timeBetweenBurstsInMs * 2, maxTimeBetweenBursts) : timeBetweenBurstsInMs;
     }
 
     /**
@@ -120,43 +172,28 @@
         if (newTransactionId > UNSIGNED_SHORT_MAX_VALUE) {
             newTransactionId = 1;
         }
-        boolean newExpectUnicastResponse = false;
         boolean newIsFirstBurst = isFirstBurst;
         int newQueriesPerBurst = queriesPerBurst;
         int newBurstCounter = burstCounter + 1;
-        long newDelayUntilNextTaskWithoutBackoffMs = delayUntilNextTaskWithoutBackoffMs;
-        int newTimeBetweenBurstsInMs = timeBetweenBurstsInMs;
-        // Only the first query expects uni-cast response.
-        if (newBurstCounter == queriesPerBurst) {
+        final boolean isFirstQueryInBurst = newBurstCounter == 1;
+        final boolean isLastQueryInBurst = newBurstCounter == queriesPerBurst;
+        if (isLastQueryInBurst) {
             newBurstCounter = 0;
-
-            if (alwaysAskForUnicastResponse) {
-                newExpectUnicastResponse = true;
-            }
-            // In passive scan mode, sends a single burst of QUERIES_PER_BURST queries, and
-            // then in each TIME_BETWEEN_BURSTS interval, sends QUERIES_PER_BURST_PASSIVE_MODE
-            // queries.
             if (isFirstBurst) {
                 newIsFirstBurst = false;
-                if (usePassiveMode) {
+                // In passive scan mode, sends a single burst of QUERIES_PER_BURST queries, and
+                // then in each TIME_BETWEEN_BURSTS interval, sends QUERIES_PER_BURST_PASSIVE_MODE
+                // queries.
+                if (queryMode == PASSIVE_QUERY_MODE) {
                     newQueriesPerBurst = QUERIES_PER_BURST_PASSIVE_MODE;
                 }
             }
-            // In active scan mode, sends a burst of QUERIES_PER_BURST queries,
-            // TIME_BETWEEN_QUERIES_IN_BURST_MS apart, then waits for the scan interval, and
-            // then repeats. The scan interval starts as INITIAL_TIME_BETWEEN_BURSTS_MS and
-            // doubles until it maxes out at TIME_BETWEEN_BURSTS_MS.
-            newDelayUntilNextTaskWithoutBackoffMs = timeBetweenBurstsInMs;
-            if (timeBetweenBurstsInMs < TIME_BETWEEN_BURSTS_MS) {
-                newTimeBetweenBurstsInMs = Math.min(timeBetweenBurstsInMs * 2,
-                        TIME_BETWEEN_BURSTS_MS);
-            }
-        } else {
-            newDelayUntilNextTaskWithoutBackoffMs = TIME_BETWEEN_QUERIES_IN_BURST_MS;
         }
+
         return new QueryTaskConfig(this, newQueryCount, newTransactionId,
-                newExpectUnicastResponse, newIsFirstBurst, newBurstCounter, newQueriesPerBurst,
-                newTimeBetweenBurstsInMs, newDelayUntilNextTaskWithoutBackoffMs);
+                getNextExpectUnicastResponse(isLastQueryInBurst), newIsFirstBurst, newBurstCounter,
+                newQueriesPerBurst, getNextTimeBetweenBurstsMs(isLastQueryInBurst),
+                getDelayUntilNextTaskWithoutBackoff(isFirstQueryInBurst, isLastQueryInBurst));
     }
 
     /**
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 8fc8114..d553210 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
@@ -86,7 +86,10 @@
     /**
      * Compare two strings by DNS case-insensitive lowercase.
      */
-    public static boolean equalsIgnoreDnsCase(@NonNull String a, @NonNull String b) {
+    public static boolean equalsIgnoreDnsCase(@Nullable String a, @Nullable String b) {
+        if (a == null || b == null) {
+            return a == null && b == null;
+        }
         if (a.length() != b.length()) return false;
         for (int i = 0; i < a.length(); i++) {
             if (toDnsLowerCase(a.charAt(i)) != toDnsLowerCase(b.charAt(i))) {
diff --git a/service-t/src/com/android/server/net/BpfInterfaceMapHelper.java b/service-t/src/com/android/server/net/BpfInterfaceMapHelper.java
new file mode 100644
index 0000000..3c95b8e
--- /dev/null
+++ b/service-t/src/com/android/server/net/BpfInterfaceMapHelper.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2022 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;
+
+import android.os.Build;
+import android.system.ErrnoException;
+import android.util.IndentingPrintWriter;
+import android.util.Log;
+
+import androidx.annotation.RequiresApi;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.net.module.util.BpfDump;
+import com.android.net.module.util.BpfMap;
+import com.android.net.module.util.IBpfMap;
+import com.android.net.module.util.Struct.S32;
+
+/**
+ * Monitor interface added (without removed) and right interface name and its index to bpf map.
+ */
+@RequiresApi(Build.VERSION_CODES.TIRAMISU)
+public class BpfInterfaceMapHelper {
+    private static final String TAG = BpfInterfaceMapHelper.class.getSimpleName();
+    // This is current path but may be changed soon.
+    private static final String IFACE_INDEX_NAME_MAP_PATH =
+            "/sys/fs/bpf/netd_shared/map_netd_iface_index_name_map";
+    private final IBpfMap<S32, InterfaceMapValue> mIndexToIfaceBpfMap;
+
+    public BpfInterfaceMapHelper() {
+        this(new Dependencies());
+    }
+
+    @VisibleForTesting
+    public BpfInterfaceMapHelper(Dependencies deps) {
+        mIndexToIfaceBpfMap = deps.getInterfaceMap();
+    }
+
+    /**
+     * Dependencies of BpfInerfaceMapUpdater, for injection in tests.
+     */
+    @VisibleForTesting
+    public static class Dependencies {
+        /** Create BpfMap for updating interface and index mapping. */
+        public IBpfMap<S32, InterfaceMapValue> getInterfaceMap() {
+            try {
+                return new BpfMap<>(IFACE_INDEX_NAME_MAP_PATH,
+                    S32.class, InterfaceMapValue.class);
+            } catch (ErrnoException e) {
+                Log.e(TAG, "Cannot create interface map: " + e);
+                return null;
+            }
+        }
+    }
+
+    /** get interface name by interface index from bpf map */
+    public String getIfNameByIndex(final int index) {
+        try {
+            final InterfaceMapValue value = mIndexToIfaceBpfMap.getValue(new S32(index));
+            if (value == null) {
+                Log.e(TAG, "No if name entry for index " + index);
+                return null;
+            }
+            return value.getInterfaceNameString();
+        } catch (ErrnoException e) {
+            Log.e(TAG, "Failed to get entry for index " + index + ": " + e);
+            return null;
+        }
+    }
+
+    /**
+     * Dump BPF map
+     *
+     * @param pw print writer
+     */
+    public void dump(final IndentingPrintWriter pw) {
+        pw.println("BPF map status:");
+        pw.increaseIndent();
+        BpfDump.dumpMapStatus(mIndexToIfaceBpfMap, pw, "IfaceIndexNameMap",
+                IFACE_INDEX_NAME_MAP_PATH);
+        pw.decreaseIndent();
+        pw.println("BPF map content:");
+        pw.increaseIndent();
+        BpfDump.dumpMap(mIndexToIfaceBpfMap, pw, "IfaceIndexNameMap",
+                (key, value) -> "ifaceIndex=" + key.val
+                        + " ifaceName=" + value.getInterfaceNameString());
+        pw.decreaseIndent();
+    }
+}
diff --git a/service-t/src/com/android/server/net/BpfInterfaceMapUpdater.java b/service-t/src/com/android/server/net/BpfInterfaceMapUpdater.java
deleted file mode 100644
index 59de2c4..0000000
--- a/service-t/src/com/android/server/net/BpfInterfaceMapUpdater.java
+++ /dev/null
@@ -1,179 +0,0 @@
-/*
- * Copyright (C) 2022 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;
-
-import android.os.Build;
-import android.content.Context;
-import android.net.INetd;
-import android.os.Handler;
-import android.os.IBinder;
-import android.os.RemoteException;
-import android.os.ServiceSpecificException;
-import android.system.ErrnoException;
-import android.util.IndentingPrintWriter;
-import android.util.Log;
-
-import androidx.annotation.RequiresApi;
-
-import com.android.internal.annotations.VisibleForTesting;
-import com.android.net.module.util.BaseNetdUnsolicitedEventListener;
-import com.android.net.module.util.BpfDump;
-import com.android.net.module.util.BpfMap;
-import com.android.net.module.util.IBpfMap;
-import com.android.net.module.util.InterfaceParams;
-import com.android.net.module.util.Struct.S32;
-
-/**
- * Monitor interface added (without removed) and right interface name and its index to bpf map.
- */
-@RequiresApi(Build.VERSION_CODES.TIRAMISU)
-public class BpfInterfaceMapUpdater {
-    private static final String TAG = BpfInterfaceMapUpdater.class.getSimpleName();
-    // This is current path but may be changed soon.
-    private static final String IFACE_INDEX_NAME_MAP_PATH =
-            "/sys/fs/bpf/netd_shared/map_netd_iface_index_name_map";
-    private final IBpfMap<S32, InterfaceMapValue> mIndexToIfaceBpfMap;
-    private final INetd mNetd;
-    private final Handler mHandler;
-    private final Dependencies mDeps;
-
-    public BpfInterfaceMapUpdater(Context ctx, Handler handler) {
-        this(ctx, handler, new Dependencies());
-    }
-
-    @VisibleForTesting
-    public BpfInterfaceMapUpdater(Context ctx, Handler handler, Dependencies deps) {
-        mDeps = deps;
-        mIndexToIfaceBpfMap = deps.getInterfaceMap();
-        mNetd = deps.getINetd(ctx);
-        mHandler = handler;
-    }
-
-    /**
-     * Dependencies of BpfInerfaceMapUpdater, for injection in tests.
-     */
-    @VisibleForTesting
-    public static class Dependencies {
-        /** Create BpfMap for updating interface and index mapping. */
-        public IBpfMap<S32, InterfaceMapValue> getInterfaceMap() {
-            try {
-                return new BpfMap<>(IFACE_INDEX_NAME_MAP_PATH,
-                    S32.class, InterfaceMapValue.class);
-            } catch (ErrnoException e) {
-                Log.e(TAG, "Cannot create interface map: " + e);
-                return null;
-            }
-        }
-
-        /** Get InterfaceParams for giving interface name. */
-        public InterfaceParams getInterfaceParams(String ifaceName) {
-            return InterfaceParams.getByName(ifaceName);
-        }
-
-        /** Get INetd binder object. */
-        public INetd getINetd(Context ctx) {
-            return INetd.Stub.asInterface((IBinder) ctx.getSystemService(Context.NETD_SERVICE));
-        }
-    }
-
-    /**
-     * Start listening interface update event.
-     * Query current interface names before listening.
-     */
-    public void start() {
-        mHandler.post(() -> {
-            if (mIndexToIfaceBpfMap == null) {
-                Log.wtf(TAG, "Fail to start: Null bpf map");
-                return;
-            }
-
-            try {
-                // TODO: use a NetlinkMonitor and listen for RTM_NEWLINK messages instead.
-                mNetd.registerUnsolicitedEventListener(new InterfaceChangeObserver());
-            } catch (RemoteException e) {
-                Log.wtf(TAG, "Unable to register netd UnsolicitedEventListener, " + e);
-            }
-
-            final String[] ifaces;
-            try {
-                // TODO: use a netlink dump to get the current interface list.
-                ifaces = mNetd.interfaceGetList();
-            } catch (RemoteException | ServiceSpecificException e) {
-                Log.wtf(TAG, "Unable to query interface names by netd, " + e);
-                return;
-            }
-
-            for (String ifaceName : ifaces) {
-                addInterface(ifaceName);
-            }
-        });
-    }
-
-    private void addInterface(String ifaceName) {
-        final InterfaceParams iface = mDeps.getInterfaceParams(ifaceName);
-        if (iface == null) {
-            Log.e(TAG, "Unable to get InterfaceParams for " + ifaceName);
-            return;
-        }
-
-        try {
-            mIndexToIfaceBpfMap.updateEntry(new S32(iface.index), new InterfaceMapValue(ifaceName));
-        } catch (ErrnoException e) {
-            Log.e(TAG, "Unable to update entry for " + ifaceName + ", " + e);
-        }
-    }
-
-    private class InterfaceChangeObserver extends BaseNetdUnsolicitedEventListener {
-        @Override
-        public void onInterfaceAdded(String ifName) {
-            mHandler.post(() -> addInterface(ifName));
-        }
-    }
-
-    /** get interface name by interface index from bpf map */
-    public String getIfNameByIndex(final int index) {
-        try {
-            final InterfaceMapValue value = mIndexToIfaceBpfMap.getValue(new S32(index));
-            if (value == null) {
-                Log.e(TAG, "No if name entry for index " + index);
-                return null;
-            }
-            return value.getInterfaceNameString();
-        } catch (ErrnoException e) {
-            Log.e(TAG, "Failed to get entry for index " + index + ": " + e);
-            return null;
-        }
-    }
-
-    /**
-     * Dump BPF map
-     *
-     * @param pw print writer
-     */
-    public void dump(final IndentingPrintWriter pw) {
-        pw.println("BPF map status:");
-        pw.increaseIndent();
-        BpfDump.dumpMapStatus(mIndexToIfaceBpfMap, pw, "IfaceIndexNameMap",
-                IFACE_INDEX_NAME_MAP_PATH);
-        pw.decreaseIndent();
-        pw.println("BPF map content:");
-        pw.increaseIndent();
-        BpfDump.dumpMap(mIndexToIfaceBpfMap, pw, "IfaceIndexNameMap",
-                (key, value) -> "ifaceIndex=" + key.val
-                        + " ifaceName=" + value.getInterfaceNameString());
-        pw.decreaseIndent();
-    }
-}
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 7b24315..ec10158 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -476,7 +476,7 @@
     private final LocationPermissionChecker mLocationPermissionChecker;
 
     @NonNull
-    private final BpfInterfaceMapUpdater mInterfaceMapUpdater;
+    private final BpfInterfaceMapHelper mInterfaceMapHelper;
 
     @Nullable
     private final SkDestroyListener mSkDestroyListener;
@@ -628,8 +628,7 @@
         mContentObserver = mDeps.makeContentObserver(mHandler, mSettings,
                 mNetworkStatsSubscriptionsMonitor);
         mLocationPermissionChecker = mDeps.makeLocationPermissionChecker(mContext);
-        mInterfaceMapUpdater = mDeps.makeBpfInterfaceMapUpdater(mContext, mHandler);
-        mInterfaceMapUpdater.start();
+        mInterfaceMapHelper = mDeps.makeBpfInterfaceMapHelper();
         mUidCounterSetMap = mDeps.getUidCounterSetMap();
         mCookieTagMap = mDeps.getCookieTagMap();
         mStatsMapA = mDeps.getStatsMapA();
@@ -798,11 +797,10 @@
             return new LocationPermissionChecker(context);
         }
 
-        /** Create BpfInterfaceMapUpdater to update bpf interface map. */
+        /** Create BpfInterfaceMapHelper to update bpf interface map. */
         @NonNull
-        public BpfInterfaceMapUpdater makeBpfInterfaceMapUpdater(
-                @NonNull Context ctx, @NonNull Handler handler) {
-            return new BpfInterfaceMapUpdater(ctx, handler);
+        public BpfInterfaceMapHelper makeBpfInterfaceMapHelper() {
+            return new BpfInterfaceMapHelper();
         }
 
         /** Get counter sets map for each UID. */
@@ -2889,9 +2887,9 @@
             }
 
             pw.println();
-            pw.println("InterfaceMapUpdater:");
+            pw.println("InterfaceMapHelper:");
             pw.increaseIndent();
-            mInterfaceMapUpdater.dump(pw);
+            mInterfaceMapHelper.dump(pw);
             pw.decreaseIndent();
 
             pw.println();
@@ -3038,7 +3036,7 @@
         BpfDump.dumpMap(statsMap, pw, mapName,
                 "ifaceIndex ifaceName tag_hex uid_int cnt_set rxBytes rxPackets txBytes txPackets",
                 (key, value) -> {
-                    final String ifName = mInterfaceMapUpdater.getIfNameByIndex(key.ifaceIndex);
+                    final String ifName = mInterfaceMapHelper.getIfNameByIndex(key.ifaceIndex);
                     return key.ifaceIndex + " "
                             + (ifName != null ? ifName : "unknown") + " "
                             + "0x" + Long.toHexString(key.tag) + " "
@@ -3056,7 +3054,7 @@
         BpfDump.dumpMap(mIfaceStatsMap, pw, "mIfaceStatsMap",
                 "ifaceIndex ifaceName rxBytes rxPackets txBytes txPackets",
                 (key, value) -> {
-                    final String ifName = mInterfaceMapUpdater.getIfNameByIndex(key.val);
+                    final String ifName = mInterfaceMapHelper.getIfNameByIndex(key.val);
                     return key.val + " "
                             + (ifName != null ? ifName : "unknown") + " "
                             + value.rxBytes + " "
diff --git a/service/ServiceConnectivityResources/OWNERS b/service/ServiceConnectivityResources/OWNERS
new file mode 100644
index 0000000..df41ff2
--- /dev/null
+++ b/service/ServiceConnectivityResources/OWNERS
@@ -0,0 +1,2 @@
+per-file res/values/config_thread.xml = file:platform/packages/modules/Connectivity:main:/thread/OWNERS
+per-file res/values/overlayable.xml = file:platform/packages/modules/Connectivity:main:/thread/OWNERS
diff --git a/service/lint-baseline.xml b/service/lint-baseline.xml
index a091a2b..b09589c 100644
--- a/service/lint-baseline.xml
+++ b/service/lint-baseline.xml
@@ -3,72 +3,6 @@
 
     <issue
         id="NewApi"
-        message="Call requires API level 31 (current min is 30): `BpfMap`"
-        errorLine1="            return new BpfMap&lt;&gt;("
-        errorLine2="                   ~~~~~~~~~~~~">
-        <location
-            file="packages/modules/Connectivity/service/src/com/android/server/BpfNetMaps.java"
-            line="189"
-            column="20"/>
-    </issue>
-
-    <issue
-        id="NewApi"
-        message="Call requires API level 31 (current min is 30): `BpfMap`"
-        errorLine1="            return new BpfMap&lt;&gt;("
-        errorLine2="                   ~~~~~~~~~~~~">
-        <location
-            file="packages/modules/Connectivity/service/src/com/android/server/BpfNetMaps.java"
-            line="198"
-            column="20"/>
-    </issue>
-
-    <issue
-        id="NewApi"
-        message="Call requires API level 31 (current min is 30): `BpfMap`"
-        errorLine1="            return new BpfMap&lt;&gt;("
-        errorLine2="                   ~~~~~~~~~~~~">
-        <location
-            file="packages/modules/Connectivity/service/src/com/android/server/BpfNetMaps.java"
-            line="207"
-            column="20"/>
-    </issue>
-
-    <issue
-        id="NewApi"
-        message="Call requires API level 31 (current min is 30): `BpfMap`"
-        errorLine1="            return new BpfMap&lt;&gt;(COOKIE_TAG_MAP_PATH, BpfMap.BPF_F_RDWR,"
-        errorLine2="                   ~~~~~~~~~~~~">
-        <location
-            file="packages/modules/Connectivity/service/src/com/android/server/BpfNetMaps.java"
-            line="216"
-            column="20"/>
-    </issue>
-
-    <issue
-        id="NewApi"
-        message="Call requires API level 31 (current min is 30): `BpfMap`"
-        errorLine1="            return new BpfMap&lt;&gt;("
-        errorLine2="                   ~~~~~~~~~~~~">
-        <location
-            file="packages/modules/Connectivity/service/src/com/android/server/BpfNetMaps.java"
-            line="225"
-            column="20"/>
-    </issue>
-
-    <issue
-        id="NewApi"
-        message="Call requires API level 31 (current min is 30): `BpfMap`"
-        errorLine1="            return new BpfMap&lt;&gt;(INGRESS_DISCARD_MAP_PATH, BpfMap.BPF_F_RDWR,"
-        errorLine2="                   ~~~~~~~~~~~~">
-        <location
-            file="packages/modules/Connectivity/service/src/com/android/server/BpfNetMaps.java"
-            line="234"
-            column="20"/>
-    </issue>
-
-    <issue
-        id="NewApi"
         message="Call requires API level 31 (current min is 30): `BpfBitmap`"
         errorLine1="                return new BpfBitmap(BLOCKED_PORTS_MAP_PATH);"
         errorLine2="                       ~~~~~~~~~~~~~">
diff --git a/service/proguard.flags b/service/proguard.flags
index cf25f05..ed9a65f 100644
--- a/service/proguard.flags
+++ b/service/proguard.flags
@@ -15,3 +15,7 @@
     static final % EVENT_*;
 }
 
+# b/313539492 Keep the onLocalNetworkInfoChanged method in classes extending Connectivity.NetworkCallback.
+-keepclassmembers class * extends **android.net.ConnectivityManager$NetworkCallback {
+    public void onLocalNetworkInfoChanged(**android.net.Network, **android.net.LocalNetworkInfo);
+}
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index aa40285..a7fddd0 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -31,7 +31,6 @@
 import static android.net.BpfNetMapsConstants.UID_OWNER_MAP_PATH;
 import static android.net.BpfNetMapsConstants.UID_PERMISSION_MAP_PATH;
 import static android.net.BpfNetMapsConstants.UID_RULES_CONFIGURATION_KEY;
-import static android.net.BpfNetMapsUtils.PRE_T;
 import static android.net.BpfNetMapsUtils.getMatchByFirewallChain;
 import static android.net.BpfNetMapsUtils.isFirewallAllowList;
 import static android.net.BpfNetMapsUtils.matchToString;
@@ -68,6 +67,7 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.modules.utils.BackgroundThread;
+import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.BpfDump;
 import com.android.net.module.util.BpfMap;
 import com.android.net.module.util.IBpfMap;
@@ -95,7 +95,7 @@
  */
 public class BpfNetMaps {
     static {
-        if (!PRE_T) {
+        if (SdkLevel.isAtLeastT()) {
             System.loadLibrary("service-connectivity");
         }
     }
@@ -184,6 +184,7 @@
         sIngressDiscardMap = ingressDiscardMap;
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static IBpfMap<S32, U32> getConfigurationMap() {
         try {
             return new BpfMap<>(
@@ -193,6 +194,7 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static IBpfMap<S32, UidOwnerValue> getUidOwnerMap() {
         try {
             return new BpfMap<>(
@@ -202,6 +204,7 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static IBpfMap<S32, U8> getUidPermissionMap() {
         try {
             return new BpfMap<>(
@@ -211,6 +214,7 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static IBpfMap<CookieTagMapKey, CookieTagMapValue> getCookieTagMap() {
         try {
             return new BpfMap<>(COOKIE_TAG_MAP_PATH,
@@ -220,6 +224,7 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static IBpfMap<S32, U8> getDataSaverEnabledMap() {
         try {
             return new BpfMap<>(
@@ -229,6 +234,7 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static IBpfMap<IngressDiscardKey, IngressDiscardValue> getIngressDiscardMap() {
         try {
             return new BpfMap<>(INGRESS_DISCARD_MAP_PATH,
@@ -238,6 +244,7 @@
         }
     }
 
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static void initBpfMaps() {
         if (sConfigurationMap == null) {
             sConfigurationMap = getConfigurationMap();
@@ -295,6 +302,7 @@
      * Initializes the class if it is not already initialized. This method will open maps but not
      * cause any other effects. This method may be called multiple times on any thread.
      */
+    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     private static synchronized void ensureInitialized(final Context context) {
         if (sInitialized) return;
         initBpfMaps();
@@ -348,7 +356,7 @@
     public BpfNetMaps(final Context context) {
         this(context, null);
 
-        if (PRE_T) throw new IllegalArgumentException("BpfNetMaps need to use netd before T");
+        if (!SdkLevel.isAtLeastT()) throw new IllegalArgumentException("BpfNetMaps need to use netd before T");
     }
 
     public BpfNetMaps(final Context context, final INetd netd) {
@@ -357,7 +365,7 @@
 
     @VisibleForTesting
     public BpfNetMaps(final Context context, final INetd netd, final Dependencies deps) {
-        if (!PRE_T) {
+        if (SdkLevel.isAtLeastT()) {
             ensureInitialized(context);
         }
         mNetd = netd;
@@ -371,7 +379,7 @@
     }
 
     private void throwIfPreT(final String msg) {
-        if (PRE_T) {
+        if (!SdkLevel.isAtLeastT()) {
             throw new UnsupportedOperationException(msg);
         }
     }
@@ -712,7 +720,7 @@
      *                                  cause of the failure.
      */
     public void addUidInterfaceRules(final String ifName, final int[] uids) throws RemoteException {
-        if (PRE_T) {
+        if (!SdkLevel.isAtLeastT()) {
             mNetd.firewallAddUidInterfaceRules(ifName, uids);
             return;
         }
@@ -750,7 +758,7 @@
      *                                  cause of the failure.
      */
     public void removeUidInterfaceRules(final int[] uids) throws RemoteException {
-        if (PRE_T) {
+        if (!SdkLevel.isAtLeastT()) {
             mNetd.firewallRemoveUidInterfaceRules(uids);
             return;
         }
@@ -829,7 +837,7 @@
      * @throws RemoteException when netd has crashed.
      */
     public void setNetPermForUids(final int permissions, final int[] uids) throws RemoteException {
-        if (PRE_T) {
+        if (!SdkLevel.isAtLeastT()) {
             mNetd.trafficSetNetPermForUids(permissions, uids);
             return;
         }
@@ -1019,7 +1027,7 @@
     @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void dump(final IndentingPrintWriter pw, final FileDescriptor fd, boolean verbose)
             throws IOException, ServiceSpecificException {
-        if (PRE_T) {
+        if (!SdkLevel.isAtLeastT()) {
             throw new ServiceSpecificException(
                     EOPNOTSUPP, "dumpsys connectivity trafficcontroller dump not available on pre-T"
                     + " devices, use dumpsys netd trafficcontroller instead.");
diff --git a/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java b/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
index b980c7d..fecaa09 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/InetDiagMessage.java
@@ -31,6 +31,7 @@
 import static com.android.net.module.util.netlink.NetlinkConstants.stringForProtocol;
 import static com.android.net.module.util.netlink.NetlinkUtils.DEFAULT_RECV_BUFSIZE;
 import static com.android.net.module.util.netlink.NetlinkUtils.IO_TIMEOUT_MS;
+import static com.android.net.module.util.netlink.NetlinkUtils.SOCKET_RECV_BUFSIZE;
 import static com.android.net.module.util.netlink.NetlinkUtils.TCP_ALIVE_STATE_FILTER;
 import static com.android.net.module.util.netlink.NetlinkUtils.connectToKernel;
 import static com.android.net.module.util.netlink.StructNlMsgHdr.NLM_F_DUMP;
@@ -281,7 +282,7 @@
         int uid = INVALID_UID;
         FileDescriptor fd = null;
         try {
-            fd = NetlinkUtils.netlinkSocketForProto(NETLINK_INET_DIAG);
+            fd = NetlinkUtils.netlinkSocketForProto(NETLINK_INET_DIAG, SOCKET_RECV_BUFSIZE);
             connectToKernel(fd);
             uid = lookupUid(protocol, local, remote, fd);
         } catch (ErrnoException | SocketException | IllegalArgumentException
diff --git a/staticlibs/device/com/android/net/module/util/netlink/NetlinkUtils.java b/staticlibs/device/com/android/net/module/util/netlink/NetlinkUtils.java
index 7c2be2c..61d347c 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/NetlinkUtils.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/NetlinkUtils.java
@@ -29,6 +29,7 @@
 import static android.system.OsConstants.SO_RCVBUF;
 import static android.system.OsConstants.SO_RCVTIMEO;
 import static android.system.OsConstants.SO_SNDTIMEO;
+
 import static com.android.net.module.util.netlink.NetlinkConstants.hexify;
 import static com.android.net.module.util.netlink.NetlinkConstants.NLMSG_DONE;
 import static com.android.net.module.util.netlink.NetlinkConstants.RTNL_FAMILY_IP6MR;
@@ -49,6 +50,7 @@
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.net.Inet6Address;
+import java.net.InetAddress;
 import java.net.SocketException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -56,7 +58,6 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
 
 /**
  * Utilities for netlink related class that may not be able to fit into a specific class.
@@ -85,6 +86,7 @@
 
     public static final int DEFAULT_RECV_BUFSIZE = 8 * 1024;
     public static final int SOCKET_RECV_BUFSIZE = 64 * 1024;
+    public static final int SOCKET_DUMP_RECV_BUFSIZE = 128 * 1024;
 
     /**
      * Return whether the input ByteBuffer contains enough remaining bytes for
@@ -159,7 +161,7 @@
      */
     public static void sendOneShotKernelMessage(int nlProto, byte[] msg) throws ErrnoException {
         final String errPrefix = "Error in NetlinkSocket.sendOneShotKernelMessage";
-        final FileDescriptor fd = netlinkSocketForProto(nlProto);
+        final FileDescriptor fd = netlinkSocketForProto(nlProto, SOCKET_RECV_BUFSIZE);
 
         try {
             connectToKernel(fd);
@@ -177,19 +179,19 @@
     }
 
     /**
-     * Send an RTM_NEWADDR message to kernel to add or update an IPv6 address.
+     * Send an RTM_NEWADDR message to kernel to add or update an IP address.
      *
      * @param ifIndex interface index.
-     * @param ip IPv6 address to be added.
-     * @param prefixlen IPv6 address prefix length.
-     * @param flags IPv6 address flags.
-     * @param scope IPv6 address scope.
-     * @param preferred The preferred lifetime of IPv6 address.
-     * @param valid The valid lifetime of IPv6 address.
+     * @param ip IP address to be added.
+     * @param prefixlen IP address prefix length.
+     * @param flags IP address flags.
+     * @param scope IP address scope.
+     * @param preferred The preferred lifetime of IP address.
+     * @param valid The valid lifetime of IP address.
      */
-    public static boolean sendRtmNewAddressRequest(int ifIndex, @NonNull final Inet6Address ip,
+    public static boolean sendRtmNewAddressRequest(int ifIndex, @NonNull final InetAddress ip,
             short prefixlen, int flags, byte scope, long preferred, long valid) {
-        Objects.requireNonNull(ip, "IPv6 address to be added should not be null.");
+        Objects.requireNonNull(ip, "IP address to be added should not be null.");
         final byte[] msg = RtNetlinkAddressMessage.newRtmNewAddressMessage(1 /* seqNo*/, ip,
                 prefixlen, flags, scope, ifIndex, preferred, valid);
         try {
@@ -225,12 +227,18 @@
     /**
      * Create netlink socket with the given netlink protocol type.
      *
+     * @param nlProto the netlink protocol
+     * @param bufferSize the receive buffer size to set when the value is not 0
+     *
      * @return fd the fileDescriptor of the socket.
      * @throws ErrnoException if the FileDescriptor not connect to be created successfully
      */
-    public static FileDescriptor netlinkSocketForProto(int nlProto) throws ErrnoException {
-        final FileDescriptor fd = Os.socket(AF_NETLINK, SOCK_DGRAM, nlProto);
-        Os.setsockoptInt(fd, SOL_SOCKET, SO_RCVBUF, SOCKET_RECV_BUFSIZE);
+    public static FileDescriptor netlinkSocketForProto(int nlProto, int bufferSize)
+            throws ErrnoException {
+        final FileDescriptor fd = Os.socket(AF_NETLINK, SOCK_DGRAM | SOCK_CLOEXEC, nlProto);
+        if (bufferSize > 0) {
+            Os.setsockoptInt(fd, SOL_SOCKET, SO_RCVBUF, bufferSize);
+        }
         return fd;
     }
 
@@ -238,7 +246,7 @@
      * Construct a netlink inet_diag socket.
      */
     public static FileDescriptor createNetLinkInetDiagSocket() throws ErrnoException {
-        return Os.socket(AF_NETLINK, SOCK_DGRAM | SOCK_CLOEXEC, NETLINK_INET_DIAG);
+        return netlinkSocketForProto(NETLINK_INET_DIAG, 0);
     }
 
     /**
@@ -373,7 +381,7 @@
             Consumer<T> func)
             throws SocketException, InterruptedIOException, ErrnoException {
         // Create socket
-        final FileDescriptor fd = netlinkSocketForProto(nlFamily);
+        final FileDescriptor fd = netlinkSocketForProto(nlFamily, SOCKET_DUMP_RECV_BUFSIZE);
         try {
             getAndProcessNetlinkDumpMessagesWithFd(fd, dumpRequestMessage, nlFamily,
                     msgClass, func);
diff --git a/staticlibs/device/com/android/net/module/util/netlink/RtNetlinkAddressMessage.java b/staticlibs/device/com/android/net/module/util/netlink/RtNetlinkAddressMessage.java
index cbe0ab0..4846df7 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/RtNetlinkAddressMessage.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/RtNetlinkAddressMessage.java
@@ -16,6 +16,7 @@
 
 package com.android.net.module.util.netlink;
 
+import static com.android.net.module.util.Inet4AddressUtils.getBroadcastAddress;
 import static com.android.net.module.util.netlink.StructNlMsgHdr.NLM_F_ACK;
 import static com.android.net.module.util.netlink.StructNlMsgHdr.NLM_F_REPLACE;
 import static com.android.net.module.util.netlink.StructNlMsgHdr.NLM_F_REQUEST;
@@ -28,6 +29,7 @@
 
 import com.android.net.module.util.HexDump;
 
+import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
@@ -48,6 +50,8 @@
  */
 public class RtNetlinkAddressMessage extends NetlinkMessage {
     public static final short IFA_ADDRESS        = 1;
+    public static final short IFA_LOCAL          = 2;
+    public static final short IFA_BROADCAST      = 4;
     public static final short IFA_CACHEINFO      = 6;
     public static final short IFA_FLAGS          = 8;
 
@@ -71,6 +75,7 @@
         mIfacacheInfo = structIfacacheInfo;
         mFlags = flags;
     }
+
     private RtNetlinkAddressMessage(@NonNull StructNlMsgHdr header) {
         this(header, null, null, null, 0);
     }
@@ -158,6 +163,24 @@
         // still be packed to ByteBuffer even if the flag is 0.
         final StructNlAttr flags = new StructNlAttr(IFA_FLAGS, mFlags);
         flags.pack(byteBuffer);
+
+        // Add the required IFA_LOCAL and IFA_BROADCAST attributes for IPv4 addresses. The IFA_LOCAL
+        // attribute represents the local address, which is equivalent to IFA_ADDRESS on a normally
+        // configured broadcast interface, however, for PPP interfaces, IFA_ADDRESS indicates the
+        // destination address and the local address is provided in the IFA_LOCAL attribute. If the
+        // IFA_LOCAL attribute is not present in the RTM_NEWADDR message, the kernel replies with an
+        // error netlink message with invalid parameters. IFA_BROADCAST is also required, otherwise
+        // the broadcast on the interface is 0.0.0.0. See include/uapi/linux/if_addr.h for details.
+        // For IPv6 addresses, the IFA_ADDRESS attribute applies and introduces no ambiguity.
+        if (mIpAddress instanceof Inet4Address) {
+            final StructNlAttr localAddress = new StructNlAttr(IFA_LOCAL, mIpAddress);
+            localAddress.pack(byteBuffer);
+
+            final Inet4Address broadcast =
+                    getBroadcastAddress((Inet4Address) mIpAddress, mIfaddrmsg.prefixLen);
+            final StructNlAttr broadcastAddress = new StructNlAttr(IFA_BROADCAST, broadcast);
+            broadcastAddress.pack(byteBuffer);
+        }
     }
 
     /**
@@ -184,7 +207,7 @@
                 0 /* tstamp */);
         msg.mFlags = flags;
 
-        final byte[] bytes = new byte[msg.getRequiredSpace()];
+        final byte[] bytes = new byte[msg.getRequiredSpace(family)];
         nlmsghdr.nlmsg_len = bytes.length;
         final ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
         byteBuffer.order(ByteOrder.nativeOrder());
@@ -237,7 +260,7 @@
     // RtNetlinkAddressMessage, e.g. RTM_DELADDR sent from user space to kernel to delete an
     // IP address only requires IFA_ADDRESS attribute. The caller should check if these attributes
     // are necessary to carry when constructing a RtNetlinkAddressMessage.
-    private int getRequiredSpace() {
+    private int getRequiredSpace(int family) {
         int spaceRequired = StructNlMsgHdr.STRUCT_SIZE + StructIfaddrMsg.STRUCT_SIZE;
         // IFA_ADDRESS attr
         spaceRequired += NetlinkConstants.alignedLengthOf(
@@ -247,6 +270,14 @@
                 StructNlAttr.NLA_HEADERLEN + StructIfacacheInfo.STRUCT_SIZE);
         // IFA_FLAGS "u32" attr
         spaceRequired += StructNlAttr.NLA_HEADERLEN + 4;
+        if (family == OsConstants.AF_INET) {
+            // IFA_LOCAL attr
+            spaceRequired += NetlinkConstants.alignedLengthOf(
+                    StructNlAttr.NLA_HEADERLEN + mIpAddress.getAddress().length);
+            // IFA_BROADCAST attr
+            spaceRequired += NetlinkConstants.alignedLengthOf(
+                    StructNlAttr.NLA_HEADERLEN + mIpAddress.getAddress().length);
+        }
         return spaceRequired;
     }
 
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NetlinkUtilsTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NetlinkUtilsTest.java
index 0958f11..b644a53 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NetlinkUtilsTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/NetlinkUtilsTest.java
@@ -21,6 +21,8 @@
 import static android.system.OsConstants.AF_UNSPEC;
 import static android.system.OsConstants.EACCES;
 import static android.system.OsConstants.NETLINK_ROUTE;
+import static android.system.OsConstants.SOL_SOCKET;
+import static android.system.OsConstants.SO_RCVBUF;
 import static com.android.net.module.util.netlink.NetlinkConstants.RTNL_FAMILY_IP6MR;
 import static com.android.net.module.util.netlink.NetlinkUtils.DEFAULT_RECV_BUFSIZE;
 import static com.android.net.module.util.netlink.StructNlMsgHdr.NLM_F_DUMP;
@@ -33,6 +35,8 @@
 import static org.junit.Assume.assumeFalse;
 
 import android.content.Context;
+import android.net.util.SocketUtils;
+import android.os.Build;
 import android.system.ErrnoException;
 import android.system.NetlinkSocketAddress;
 import android.system.Os;
@@ -43,6 +47,7 @@
 
 import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.Struct;
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 
 import libcore.io.IoUtils;
 
@@ -204,4 +209,23 @@
             assertNotNull("Route doesn't contain destination: " + route, route.getDestination());
         }
     }
+
+    @Test @IgnoreUpTo(Build.VERSION_CODES.R) // getsockoptInt requires > R
+    public void testNetlinkSocketForProto_defaultBufferSize() throws Exception {
+        final FileDescriptor fd = NetlinkUtils.netlinkSocketForProto(NETLINK_ROUTE, 0);
+        final int bufferSize = Os.getsockoptInt(fd, SOL_SOCKET, SO_RCVBUF) / 2;
+
+        assertTrue("bufferSize: " + bufferSize, bufferSize > 0); // whatever the default value is
+        SocketUtils.closeSocket(fd);
+    }
+
+    @Test @IgnoreUpTo(Build.VERSION_CODES.R) // getsockoptInt requires > R
+    public void testNetlinkSocketForProto_setBufferSize() throws Exception {
+        final FileDescriptor fd = NetlinkUtils.netlinkSocketForProto(NETLINK_ROUTE,
+                8000);
+        final int bufferSize = Os.getsockoptInt(fd, SOL_SOCKET, SO_RCVBUF) / 2;
+
+        assertEquals(8000, bufferSize);
+        SocketUtils.closeSocket(fd);
+    }
 }
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/RtNetlinkAddressMessageTest.java b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/RtNetlinkAddressMessageTest.java
index 01126d2..1d08525 100644
--- a/staticlibs/tests/unit/src/com/android/net/module/util/netlink/RtNetlinkAddressMessageTest.java
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/netlink/RtNetlinkAddressMessageTest.java
@@ -42,6 +42,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -179,6 +180,57 @@
     }
 
     @Test
+    public void testCreateRtmNewAddressMessage_IPv4Address() {
+        // Hexadecimal representation of our created packet.
+        final String expectedNewAddressHex =
+                // struct nlmsghdr
+                "4c000000"      // length = 76
+                + "1400"        // type = 20 (RTM_NEWADDR)
+                + "0501"        // flags = NLM_F_ACK | NLM_F_REQUEST | NLM_F_REPLACE
+                + "01000000"    // seqno = 1
+                + "00000000"    // pid = 0 (send to kernel)
+                // struct IfaddrMsg
+                + "02"          // family = inet
+                + "18"          // prefix len = 24
+                + "00"          // flags = 0
+                + "00"          // scope = RT_SCOPE_UNIVERSE
+                + "14000000"    // ifindex = 20
+                // struct nlattr: IFA_ADDRESS
+                + "0800"        // len
+                + "0100"        // type
+                + "C0A80491"    // IPv4 address = 192.168.4.145
+                // struct nlattr: IFA_CACHEINFO
+                + "1400"        // len
+                + "0600"        // type
+                + "C0A80000"    // preferred = 43200s
+                + "C0A80000"    // valid = 43200s
+                + "00000000"    // cstamp
+                + "00000000"    // tstamp
+                // struct nlattr: IFA_FLAGS
+                + "0800"        // len
+                + "0800"        // type
+                + "00000000"    // flags = 0
+                // struct nlattr: IFA_LOCAL
+                + "0800"        // len
+                + "0200"        // type
+                + "C0A80491"    // local address = 192.168.4.145
+                // struct nlattr: IFA_BROADCAST
+                + "0800"        // len
+                + "0400"        // type
+                + "C0A804FF";   // broadcast address = 192.168.4.255
+        final byte[] expectedNewAddress =
+                HexEncoding.decode(expectedNewAddressHex.toCharArray(), false);
+
+        final Inet4Address ipAddress =
+                (Inet4Address) InetAddresses.parseNumericAddress("192.168.4.145");
+        final byte[] bytes = RtNetlinkAddressMessage.newRtmNewAddressMessage(1 /* seqno */,
+                ipAddress, (short) 24 /* prefix len */, 0 /* flags */,
+                (byte) RT_SCOPE_UNIVERSE /* scope */, 20 /* ifindex */,
+                (long) 0xA8C0 /* preferred */, (long) 0xA8C0 /* valid */);
+        assertArrayEquals(expectedNewAddress, bytes);
+    }
+
+    @Test
     public void testCreateRtmDelAddressMessage() {
         // Hexadecimal representation of our created packet.
         final String expectedDelAddressHex =
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt b/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
index 10accd4..69fdbf8 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
@@ -31,6 +31,7 @@
 import org.junit.runner.notification.Failure
 import org.junit.runner.notification.RunNotifier
 import org.junit.runners.Parameterized
+import org.mockito.Mockito
 
 /**
  * A runner that can skip tests based on the development SDK as defined in [DevSdkIgnoreRule].
@@ -124,6 +125,9 @@
             notifier.fireTestFailure(Failure(leakMonitorDesc,
                     IllegalStateException("Unexpected thread changes: $threadsDiff")))
         }
+        // Clears up internal state of all inline mocks.
+        // TODO: Call clearInlineMocks() at the end of each test.
+        Mockito.framework().clearInlineMocks()
         notifier.fireTestFinished(leakMonitorDesc)
     }
 
diff --git a/tests/common/java/android/net/nsd/NsdServiceInfoTest.java b/tests/common/java/android/net/nsd/NsdServiceInfoTest.java
index 79c4980..8e89037 100644
--- a/tests/common/java/android/net/nsd/NsdServiceInfoTest.java
+++ b/tests/common/java/android/net/nsd/NsdServiceInfoTest.java
@@ -119,6 +119,7 @@
         fullInfo.setSubtypes(Set.of("_thread", "_matter"));
         fullInfo.setPort(4242);
         fullInfo.setHostAddresses(List.of(IPV4_ADDRESS));
+        fullInfo.setHostname("home");
         fullInfo.setNetwork(new Network(123));
         fullInfo.setInterfaceIndex(456);
         checkParcelable(fullInfo);
@@ -134,6 +135,7 @@
         attributedInfo.setServiceType("_kitten._tcp");
         attributedInfo.setPort(4242);
         attributedInfo.setHostAddresses(List.of(IPV6_ADDRESS, IPV4_ADDRESS));
+        attributedInfo.setHostname("home");
         attributedInfo.setAttribute("color", "pink");
         attributedInfo.setAttribute("sound", (new String("にゃあ")).getBytes("UTF-8"));
         attributedInfo.setAttribute("adorable", (String) null);
@@ -169,6 +171,7 @@
         assertEquals(original.getServiceName(), result.getServiceName());
         assertEquals(original.getServiceType(), result.getServiceType());
         assertEquals(original.getHost(), result.getHost());
+        assertEquals(original.getHostname(), result.getHostname());
         assertTrue(original.getPort() == result.getPort());
         assertEquals(original.getNetwork(), result.getNetwork());
         assertEquals(original.getInterfaceIndex(), result.getInterfaceIndex());
diff --git a/tests/cts/net/Android.bp b/tests/cts/net/Android.bp
index 3d53d6c..f0e0ae8 100644
--- a/tests/cts/net/Android.bp
+++ b/tests/cts/net/Android.bp
@@ -54,6 +54,7 @@
         "junit",
         "junit-params",
         "modules-utils-build",
+        "net-tests-utils",
         "net-utils-framework-common",
         "truth",
         "TetheringIntegrationTestsBaseLib",
diff --git a/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java b/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java
index 8d68c5f..af1af43 100644
--- a/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java
+++ b/tests/cts/net/api23Test/src/android/net/cts/api23test/ConnectivityManagerApi23Test.java
@@ -65,7 +65,7 @@
         }
         ConnectivityReceiver.prepare();
 
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
 
         // The connectivity broadcast has been sent; push through a terminal broadcast
         // to wait for in the receive to confirm it didn't see the connectivity change.
@@ -88,7 +88,7 @@
                 .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK));
         Thread.sleep(200);
 
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
 
         Intent getConnectivityCount = new Intent(GET_WIFI_CONNECTIVITY_ACTION_COUNT);
         assertEquals(2, sendOrderedBroadcastAndReturnResultCode(
@@ -106,7 +106,7 @@
         filter.addAction(ConnectivityManager.CONNECTIVITY_ACTION);
         mContext.registerReceiver(receiver, filter);
 
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
         Intent finalIntent = new Intent(ConnectivityReceiver.FINAL_ACTION);
         finalIntent.setClass(mContext, ConnectivityReceiver.class);
         mContext.sendBroadcast(finalIntent);
diff --git a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
index 466514c..3e5d0ba 100644
--- a/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/BatteryStatsManagerTest.java
@@ -199,7 +199,8 @@
         Log.d(TAG, "Generate traffic on wifi network.");
         generateNetworkTraffic(wifiNetwork, url);
         // Wifi battery stats are updated when wifi on.
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.disableWifi();
+        mCtsNetUtils.ensureWifiConnected();
 
         // Check wifi battery stats are updated.
         runAsShell(UPDATE_DEVICE_STATS,
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 544f300..2646b60 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -1353,9 +1353,7 @@
     public void testToggleWifiConnectivityAction() throws Exception {
         assumeTrue(mPackageManager.hasSystemFeature(FEATURE_WIFI));
 
-        // toggleWifi calls connectToWifi and disconnectFromWifi, which both wait for
-        // CONNECTIVITY_ACTION broadcasts.
-        mCtsNetUtils.toggleWifi();
+        mCtsNetUtils.reconnectWifiAndWaitForConnectivityAction();
     }
 
     /** Verify restricted networks cannot be requested. */
diff --git a/tests/cts/net/src/android/net/cts/DiscoveryRequestTest.kt b/tests/cts/net/src/android/net/cts/DiscoveryRequestTest.kt
new file mode 100644
index 0000000..909a5bc
--- /dev/null
+++ b/tests/cts/net/src/android/net/cts/DiscoveryRequestTest.kt
@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.cts
+
+import android.net.Network
+import android.net.nsd.DiscoveryRequest
+import android.os.Build
+import androidx.test.filters.SmallTest
+import com.android.testutils.ConnectivityModuleTest
+import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import com.android.testutils.DevSdkIgnoreRunner
+import com.android.testutils.assertParcelingIsLossless
+import com.android.testutils.assertThrows
+import kotlin.test.assertEquals
+import kotlin.test.assertNotEquals
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** CTS tests for {@link DiscoveryRequest}. */
+@IgnoreUpTo(Build.VERSION_CODES.S_V2)
+@RunWith(DevSdkIgnoreRunner::class)
+@SmallTest
+@ConnectivityModuleTest
+class DiscoveryRequestTest {
+    @Test
+    fun testParcelingIsLossLess() {
+        val requestWithNullFields =
+                DiscoveryRequest.Builder("_ipps._tcp").build()
+        val requestWithAllFields =
+                DiscoveryRequest.Builder("_ipps._tcp")
+                                .setSubtype("_xyz")
+                                .setNetwork(Network(1))
+                                .build()
+
+        assertParcelingIsLossless(requestWithNullFields)
+        assertParcelingIsLossless(requestWithAllFields)
+    }
+
+    @Test
+    fun testBuilder_success() {
+        val request = DiscoveryRequest.Builder("_ipps._tcp")
+                                      .setSubtype("_xyz")
+                                      .setNetwork(Network(1))
+                                      .build()
+
+        assertEquals("_ipps._tcp", request.serviceType)
+        assertEquals("_xyz", request.subtype)
+        assertEquals(Network(1), request.network)
+    }
+
+    @Test
+    fun testBuilderConstructor_emptyServiceType_throwsIllegalArgument() {
+        assertThrows(IllegalArgumentException::class.java) {
+            DiscoveryRequest.Builder("")
+        }
+    }
+
+    @Test
+    fun testEquality() {
+        val request1 = DiscoveryRequest.Builder("_ipps._tcp").build()
+        val request2 = DiscoveryRequest.Builder("_ipps._tcp").build()
+        val request3 = DiscoveryRequest.Builder("_ipps._tcp")
+                .setSubtype("_xyz")
+                .setNetwork(Network(1))
+                .build()
+        val request4 = DiscoveryRequest.Builder("_ipps._tcp")
+                .setSubtype("_xyz")
+                .setNetwork(Network(1))
+                .build()
+
+        assertEquals(request1, request2)
+        assertEquals(request3, request4)
+        assertNotEquals(request1, request3)
+        assertNotEquals(request2, request4)
+    }
+}
diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index e433442..9aa3c84 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -53,6 +53,7 @@
 import android.net.cts.NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent.ServiceUpdatedLost
 import android.net.cts.NsdServiceInfoCallbackRecord.ServiceInfoCallbackEvent.UnregisterCallbackSucceeded
 import android.net.cts.util.CtsNetUtils
+import android.net.nsd.DiscoveryRequest
 import android.net.nsd.NsdManager
 import android.net.nsd.NsdServiceInfo
 import android.net.nsd.OffloadEngine
@@ -113,6 +114,7 @@
 import kotlin.math.min
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
+import kotlin.test.assertNotEquals
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 import kotlin.test.fail
@@ -125,6 +127,7 @@
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import kotlin.test.assertNotEquals
 
 private const val TAG = "NsdManagerTest"
 private const val TIMEOUT_MS = 2000L
@@ -160,7 +163,11 @@
     private val cm by lazy { context.getSystemService(ConnectivityManager::class.java)!! }
     private val serviceName = "NsdTest%09d".format(Random().nextInt(1_000_000_000))
     private val serviceName2 = "NsdTest%09d".format(Random().nextInt(1_000_000_000))
+    private val serviceName3 = "NsdTest%09d".format(Random().nextInt(1_000_000_000))
     private val serviceType = "_nmt%09d._tcp".format(Random().nextInt(1_000_000_000))
+    private val serviceType2 = "_nmt%09d._tcp".format(Random().nextInt(1_000_000_000))
+    private val customHostname = "NsdTestHost%09d".format(Random().nextInt(1_000_000_000))
+    private val customHostname2 = "NsdTestHost%09d".format(Random().nextInt(1_000_000_000))
     private val handlerThread = HandlerThread(NsdManagerTest::class.java.simpleName)
     private val ctsNetUtils by lazy{ CtsNetUtils(context) }
 
@@ -1042,9 +1049,11 @@
             nsdManager.discoverServices("_subtype1.$serviceType",
                     NsdManager.PROTOCOL_DNS_SD,
                     testNetwork1.network, Executor { it.run() }, subtype1DiscoveryRecord)
-            nsdManager.discoverServices("_subtype2.$serviceType",
-                    NsdManager.PROTOCOL_DNS_SD,
-                    testNetwork1.network, Executor { it.run() }, subtype2DiscoveryRecord)
+
+            nsdManager.discoverServices(
+                    DiscoveryRequest.Builder(serviceType).setSubtype("_subtype2")
+                            .setNetwork(testNetwork1.network).build(),
+                    Executor { it.run() }, subtype2DiscoveryRecord)
 
             val info1 = subtype1DiscoveryRecord.waitForServiceDiscovered(
                     serviceName, serviceType, testNetwork1.network)
@@ -1099,6 +1108,28 @@
     }
 
     @Test
+    fun testSubtypeDiscovery_typeMatchButSubtypeNotMatch_notDiscovered() {
+        val si1 = makeTestServiceInfo(network = testNetwork1.network).apply {
+            serviceType += ",_subtype1"
+        }
+        val registrationRecord = NsdRegistrationRecord()
+        val subtype2DiscoveryRecord = NsdDiscoveryRecord()
+        tryTest {
+            registerService(registrationRecord, si1)
+            val request = DiscoveryRequest.Builder(serviceType)
+                    .setSubtype("_subtype2").setNetwork(testNetwork1.network).build()
+            nsdManager.discoverServices(request, { it.run() }, subtype2DiscoveryRecord)
+            subtype2DiscoveryRecord.expectCallback<DiscoveryStarted>()
+            subtype2DiscoveryRecord.assertNoCallback(timeoutMs = 2000)
+        } cleanupStep {
+            nsdManager.stopServiceDiscovery(subtype2DiscoveryRecord)
+            subtype2DiscoveryRecord.expectCallback<DiscoveryStopped>()
+        } cleanup {
+            nsdManager.unregisterService(registrationRecord)
+        }
+    }
+
+    @Test
     fun testSubtypeAdvertising_tooManySubtypes_returnsFailureBadParameters() {
         val si = makeTestServiceInfo(network = testNetwork1.network)
         // Sets 101 subtypes in total
@@ -1163,6 +1194,83 @@
     }
 
     @Test
+    fun testRegisterServiceWithCustomHostAndAddresses_conflictDuringProbing_hostRenamed() {
+        val si = makeTestServiceInfo(testNetwork1.network).apply {
+            hostname = customHostname
+            hostAddresses = listOf(
+                    parseNumericAddress("192.0.2.24"),
+                    parseNumericAddress("2001:db8::3"))
+        }
+
+        val packetReader = TapPacketReader(Handler(handlerThread.looper),
+                testNetwork1.iface.fileDescriptor.fileDescriptor, 1500 /* maxPacketSize */)
+        packetReader.startAsyncForTest()
+        handlerThread.waitForIdle(TIMEOUT_MS)
+
+        // Register service on testNetwork1
+        val registrationRecord = NsdRegistrationRecord()
+        nsdManager.registerService(si, NsdManager.PROTOCOL_DNS_SD, { it.run() },
+                registrationRecord)
+
+        tryTest {
+            assertNotNull(packetReader.pollForProbe(serviceName, serviceType),
+                    "Did not find a probe for the service")
+            packetReader.sendResponse(buildConflictingAnnouncementForCustomHost())
+
+            // Registration must use an updated hostname to avoid the conflict
+            val cb = registrationRecord.expectCallback<ServiceRegistered>(REGISTRATION_TIMEOUT_MS)
+            // Service name is not renamed because there's no conflict on the service name.
+            assertEquals(serviceName, cb.serviceInfo.serviceName)
+            val hostname = cb.serviceInfo.hostname ?: fail("Missing hostname")
+            hostname.let {
+                assertTrue("Unexpected registered hostname: $it",
+                        it.startsWith(customHostname) && it != customHostname)
+            }
+        } cleanupStep {
+            nsdManager.unregisterService(registrationRecord)
+            registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
+        }
+    }
+
+    @Test
+    fun testRegisterServiceWithCustomHostNoAddresses_noConflictDuringProbing_notRenamed() {
+        val si = makeTestServiceInfo(testNetwork1.network).apply {
+            hostname = customHostname
+        }
+
+        val packetReader = TapPacketReader(Handler(handlerThread.looper),
+                testNetwork1.iface.fileDescriptor.fileDescriptor, 1500 /* maxPacketSize */)
+        packetReader.startAsyncForTest()
+        handlerThread.waitForIdle(TIMEOUT_MS)
+
+        // Register service on testNetwork1
+        val registrationRecord = NsdRegistrationRecord()
+        nsdManager.registerService(si, NsdManager.PROTOCOL_DNS_SD, { it.run() },
+                registrationRecord)
+
+        tryTest {
+            assertNotNull(packetReader.pollForProbe(serviceName, serviceType),
+                    "Did not find a probe for the service")
+            // Not a conflict because no record is registered for the hostname
+            packetReader.sendResponse(buildConflictingAnnouncementForCustomHost())
+
+            // Registration is not renamed because there's no conflict
+            val cb = registrationRecord.expectCallback<ServiceRegistered>(REGISTRATION_TIMEOUT_MS)
+            assertEquals(serviceName, cb.serviceInfo.serviceName)
+            assertEquals(customHostname, cb.serviceInfo.hostname)
+        } cleanupStep {
+            nsdManager.unregisterService(registrationRecord)
+            registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
+        }
+    }
+
+    @Test
     fun testRegisterWithConflictAfterProbing() {
         // This test requires shims supporting T+ APIs (NsdServiceInfo.network)
         assumeTrue(TestUtils.shouldTestTApis())
@@ -1237,6 +1345,121 @@
         }
     }
 
+    @Test
+    fun testRegisterServiceWithCustomHostAndAddresses_conflictAfterProbing_hostRenamed() {
+        val si = makeTestServiceInfo(testNetwork1.network).apply {
+            hostname = customHostname
+            hostAddresses = listOf(
+                    parseNumericAddress("192.0.2.24"),
+                    parseNumericAddress("2001:db8::3"))
+        }
+
+        // Register service on testNetwork1
+        val registrationRecord = NsdRegistrationRecord()
+        val discoveryRecord = NsdDiscoveryRecord()
+        val registeredService = registerService(registrationRecord, si)
+        val packetReader = TapPacketReader(
+                Handler(handlerThread.looper),
+                testNetwork1.iface.fileDescriptor.fileDescriptor, 1500 /* maxPacketSize */)
+        packetReader.startAsyncForTest()
+        handlerThread.waitForIdle(TIMEOUT_MS)
+
+        tryTest {
+            repeat(3) {
+                assertNotNull(packetReader.pollForAdvertisement(serviceName, serviceType),
+                        "Expect 3 announcements sent after initial probing")
+            }
+
+            assertEquals(si.serviceName, registeredService.serviceName)
+            assertEquals(si.hostname, registeredService.hostname)
+
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network, { it.run() }, discoveryRecord)
+            val discoveredInfo = discoveryRecord.waitForServiceDiscovered(
+                    si.serviceName, serviceType)
+
+            // Send a conflicting announcement
+            val conflictingAnnouncement = buildConflictingAnnouncementForCustomHost()
+            packetReader.sendResponse(conflictingAnnouncement)
+
+            // Expect to see probes (RFC6762 9., service is reset to probing state)
+            assertNotNull(packetReader.pollForProbe(serviceName, serviceType),
+                    "Probe not received within timeout after conflict")
+
+            // Send the conflicting packet again to reply to the probe
+            packetReader.sendResponse(conflictingAnnouncement)
+
+            val newRegistration =
+                    registrationRecord
+                            .expectCallbackEventually<ServiceRegistered>(REGISTRATION_TIMEOUT_MS) {
+                                it.serviceInfo.serviceName == serviceName
+                                        && it.serviceInfo.hostname.let { hostname ->
+                                    hostname != null
+                                            && hostname.startsWith(customHostname)
+                                            && hostname != customHostname
+                                }
+                            }
+
+            val resolvedInfo = resolveService(discoveredInfo)
+            assertEquals(newRegistration.serviceInfo.serviceName, resolvedInfo.serviceName)
+            assertEquals(newRegistration.serviceInfo.hostname, resolvedInfo.hostname)
+
+            discoveryRecord.assertNoCallback()
+        } cleanupStep {
+            nsdManager.stopServiceDiscovery(discoveryRecord)
+            discoveryRecord.expectCallback<DiscoveryStopped>()
+        } cleanupStep {
+            nsdManager.unregisterService(registrationRecord)
+            registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
+        }
+    }
+
+    @Test
+    fun testRegisterServiceWithCustomHostNoAddresses_noConflictAfterProbing_notRenamed() {
+        val si = makeTestServiceInfo(testNetwork1.network).apply {
+            hostname = customHostname
+        }
+
+        // Register service on testNetwork1
+        val registrationRecord = NsdRegistrationRecord()
+        val discoveryRecord = NsdDiscoveryRecord()
+        val registeredService = registerService(registrationRecord, si)
+        val packetReader = TapPacketReader(Handler(handlerThread.looper),
+                testNetwork1.iface.fileDescriptor.fileDescriptor, 1500 /* maxPacketSize */)
+        packetReader.startAsyncForTest()
+        handlerThread.waitForIdle(TIMEOUT_MS)
+
+        tryTest {
+            assertNotNull(packetReader.pollForAdvertisement(serviceName, serviceType),
+                    "No announcements sent after initial probing")
+
+            assertEquals(si.serviceName, registeredService.serviceName)
+            assertEquals(si.hostname, registeredService.hostname)
+
+            // Send a conflicting announcement
+            val conflictingAnnouncement = buildConflictingAnnouncementForCustomHost()
+            packetReader.sendResponse(conflictingAnnouncement)
+
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network, { it.run() }, discoveryRecord)
+
+            // The service is not renamed
+            discoveryRecord.waitForServiceDiscovered(si.serviceName, serviceType)
+        } cleanupStep {
+            nsdManager.stopServiceDiscovery(discoveryRecord)
+            discoveryRecord.expectCallback<DiscoveryStopped>()
+        } cleanupStep {
+            nsdManager.unregisterService(registrationRecord)
+            registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
+        }
+    }
+
     // Test that even if only a PTR record is received as a reply when discovering, without the
     // SRV, TXT, address records as recommended (but not mandated) by RFC 6763 12, the service can
     // still be discovered.
@@ -1421,6 +1644,212 @@
         return Inet6Address.getByAddress(addrBytes) as Inet6Address
     }
 
+    @Test
+    fun testAdvertisingAndDiscovery_servicesWithCustomHost_customHostAddressesFound() {
+        val hostAddresses1 = listOf(
+                parseNumericAddress("192.0.2.23"),
+                parseNumericAddress("2001:db8::1"),
+                parseNumericAddress("2001:db8::2"))
+        val hostAddresses2 = listOf(
+                parseNumericAddress("192.0.2.24"),
+                parseNumericAddress("2001:db8::3"))
+        val si1 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.serviceName = serviceName
+            it.serviceType = serviceType
+            it.port = TEST_PORT
+            it.hostname = customHostname
+            it.hostAddresses = hostAddresses1
+        }
+        val si2 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.serviceName = serviceName2
+            it.serviceType = serviceType
+            it.port = TEST_PORT + 1
+            it.hostname = customHostname2
+            it.hostAddresses = hostAddresses2
+        }
+        val registrationRecord1 = NsdRegistrationRecord()
+        val registrationRecord2 = NsdRegistrationRecord()
+
+        val discoveryRecord1 = NsdDiscoveryRecord()
+        val discoveryRecord2 = NsdDiscoveryRecord()
+        tryTest {
+            registerService(registrationRecord1, si1)
+
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network, Executor { it.run() }, discoveryRecord1)
+
+            val discoveredInfo = discoveryRecord1.waitForServiceDiscovered(
+                    serviceName, serviceType, testNetwork1.network)
+            val resolvedInfo = resolveService(discoveredInfo)
+
+            assertEquals(TEST_PORT, resolvedInfo.port)
+            assertEquals(si1.hostname, resolvedInfo.hostname)
+            assertAddressEquals(hostAddresses1, resolvedInfo.hostAddresses)
+
+            registerService(registrationRecord2, si2)
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network, Executor { it.run() }, discoveryRecord2)
+
+            val discoveredInfo2 = discoveryRecord2.waitForServiceDiscovered(
+                    serviceName2, serviceType, testNetwork1.network)
+            val resolvedInfo2 = resolveService(discoveredInfo2)
+
+            assertEquals(TEST_PORT + 1, resolvedInfo2.port)
+            assertEquals(si2.hostname, resolvedInfo2.hostname)
+            assertAddressEquals(hostAddresses2, resolvedInfo2.hostAddresses)
+        } cleanupStep {
+            nsdManager.stopServiceDiscovery(discoveryRecord1)
+            nsdManager.stopServiceDiscovery(discoveryRecord2)
+
+            discoveryRecord1.expectCallbackEventually<DiscoveryStopped>()
+            discoveryRecord2.expectCallbackEventually<DiscoveryStopped>()
+        } cleanup {
+            nsdManager.unregisterService(registrationRecord1)
+            nsdManager.unregisterService(registrationRecord2)
+        }
+    }
+
+    @Test
+    fun testAdvertisingAndDiscovery_multipleRegistrationsForSameCustomHost_unionOfAddressesFound() {
+        val hostAddresses1 = listOf(
+                parseNumericAddress("192.0.2.23"),
+                parseNumericAddress("2001:db8::1"),
+                parseNumericAddress("2001:db8::2"))
+        val hostAddresses2 = listOf(
+                parseNumericAddress("192.0.2.24"),
+                parseNumericAddress("2001:db8::3"))
+        val hostAddresses3 = listOf(
+                parseNumericAddress("2001:db8::3"),
+                parseNumericAddress("2001:db8::5"))
+        val si1 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.hostname = customHostname
+            it.hostAddresses = hostAddresses1
+        }
+        val si2 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.serviceName = serviceName
+            it.serviceType = serviceType
+            it.port = TEST_PORT
+            it.hostname = customHostname
+            it.hostAddresses = hostAddresses2
+        }
+        val si3 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.serviceName = serviceName3
+            it.serviceType = serviceType
+            it.port = TEST_PORT + 1
+            it.hostname = customHostname
+            it.hostAddresses = hostAddresses3
+        }
+
+        val registrationRecord1 = NsdRegistrationRecord()
+        val registrationRecord2 = NsdRegistrationRecord()
+        val registrationRecord3 = NsdRegistrationRecord()
+
+        val discoveryRecord = NsdDiscoveryRecord()
+        tryTest {
+            registerService(registrationRecord1, si1)
+            registerService(registrationRecord2, si2)
+
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network, Executor { it.run() }, discoveryRecord)
+
+            val discoveredInfo1 = discoveryRecord.waitForServiceDiscovered(
+                    serviceName, serviceType, testNetwork1.network)
+            val resolvedInfo1 = resolveService(discoveredInfo1)
+
+            assertEquals(TEST_PORT, resolvedInfo1.port)
+            assertEquals(si1.hostname, resolvedInfo1.hostname)
+            assertAddressEquals(
+                    hostAddresses1 + hostAddresses2,
+                    resolvedInfo1.hostAddresses)
+
+            registerService(registrationRecord3, si3)
+
+            val discoveredInfo2 = discoveryRecord.waitForServiceDiscovered(
+                    serviceName3, serviceType, testNetwork1.network)
+            val resolvedInfo2 = resolveService(discoveredInfo2)
+
+            assertEquals(TEST_PORT + 1, resolvedInfo2.port)
+            assertEquals(si2.hostname, resolvedInfo2.hostname)
+            assertAddressEquals(
+                    hostAddresses1 + hostAddresses2 + hostAddresses3,
+                    resolvedInfo2.hostAddresses)
+        } cleanupStep {
+            nsdManager.stopServiceDiscovery(discoveryRecord)
+
+            discoveryRecord.expectCallbackEventually<DiscoveryStopped>()
+        } cleanup {
+            nsdManager.unregisterService(registrationRecord1)
+            nsdManager.unregisterService(registrationRecord2)
+            nsdManager.unregisterService(registrationRecord3)
+        }
+    }
+
+    @Test
+    fun testAdvertisingAndDiscovery_servicesWithTheSameCustomHostAddressOmitted_addressesFound() {
+        val hostAddresses = listOf(
+                parseNumericAddress("192.0.2.23"),
+                parseNumericAddress("2001:db8::1"),
+                parseNumericAddress("2001:db8::2"))
+        val si1 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.serviceType = serviceType
+            it.serviceName = serviceName
+            it.port = TEST_PORT
+            it.hostname = customHostname
+            it.hostAddresses = hostAddresses
+        }
+        val si2 = NsdServiceInfo().also {
+            it.network = testNetwork1.network
+            it.serviceType = serviceType
+            it.serviceName = serviceName2
+            it.port = TEST_PORT + 1
+            it.hostname = customHostname
+        }
+
+        val registrationRecord1 = NsdRegistrationRecord()
+        val registrationRecord2 = NsdRegistrationRecord()
+
+        val discoveryRecord = NsdDiscoveryRecord()
+        tryTest {
+            registerService(registrationRecord1, si1)
+
+            nsdManager.discoverServices(serviceType, NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network, Executor { it.run() }, discoveryRecord)
+
+            val discoveredInfo1 = discoveryRecord.waitForServiceDiscovered(
+                    serviceName, serviceType, testNetwork1.network)
+            val resolvedInfo1 = resolveService(discoveredInfo1)
+
+            assertEquals(serviceName, discoveredInfo1.serviceName)
+            assertEquals(TEST_PORT, resolvedInfo1.port)
+            assertEquals(si1.hostname, resolvedInfo1.hostname)
+            assertAddressEquals(hostAddresses, resolvedInfo1.hostAddresses)
+
+            registerService(registrationRecord2, si2)
+
+            val discoveredInfo2 = discoveryRecord.waitForServiceDiscovered(
+                    serviceName2, serviceType, testNetwork1.network)
+            val resolvedInfo2 = resolveService(discoveredInfo2)
+
+            assertEquals(serviceName2, discoveredInfo2.serviceName)
+            assertEquals(TEST_PORT + 1, resolvedInfo2.port)
+            assertEquals(si2.hostname, resolvedInfo2.hostname)
+            assertAddressEquals(hostAddresses, resolvedInfo2.hostAddresses)
+        } cleanupStep {
+            nsdManager.stopServiceDiscovery(discoveryRecord)
+
+            discoveryRecord.expectCallback<DiscoveryStopped>()
+        } cleanup {
+            nsdManager.unregisterService(registrationRecord1)
+            nsdManager.unregisterService(registrationRecord2)
+        }
+    }
+
     private fun buildConflictingAnnouncement(): ByteBuffer {
         /*
         Generated with:
@@ -1437,6 +1866,22 @@
         return buildMdnsPacket(mdnsPayload)
     }
 
+    private fun buildConflictingAnnouncementForCustomHost(): ByteBuffer {
+        /*
+        Generated with scapy:
+        raw(DNS(rd=0, qr=1, aa=1, qd = None, an =
+            DNSRR(rrname='NsdTestHost123456789.local', type=28, rclass=1, ttl=120,
+                    rdata='2001:db8::321')
+        )).hex()
+         */
+        val mdnsPayload = HexDump.hexStringToByteArray("000084000000000100000000144e7364" +
+                "54657374486f7374313233343536373839056c6f63616c00001c000100000078001020010db80000" +
+                "00000000000000000321")
+        replaceCustomHostnameWithTestSuffix(mdnsPayload)
+
+        return buildMdnsPacket(mdnsPayload)
+    }
+
     /**
      * Replaces occurrences of "NsdTest123456789" and "_nmt123456789" in mDNS payload with the
      * actual random name and type that are used by the test.
@@ -1453,6 +1898,19 @@
         replaceAll(packetBuffer, testPacketTypePrefix, encodedTypePrefix)
     }
 
+    /**
+     * Replaces occurrences of "NsdTestHost123456789" in mDNS payload with the
+     * actual random host name that are used by the test.
+     */
+    private fun replaceCustomHostnameWithTestSuffix(mdnsPayload: ByteArray) {
+        // Test custom hostnames have consistent length and are always ASCII
+        val testPacketName = "NsdTestHost123456789".encodeToByteArray()
+        val encodedHostname = customHostname.encodeToByteArray()
+
+        val packetBuffer = ByteBuffer.wrap(mdnsPayload)
+        replaceAll(packetBuffer, testPacketName, encodedHostname)
+    }
+
     private tailrec fun replaceAll(buffer: ByteBuffer, source: ByteArray, replacement: ByteArray) {
         assertEquals(source.size, replacement.size)
         val index = buffer.array().indexOf(source)
@@ -1551,3 +2009,9 @@
     if (this == null) return ""
     return String(this, StandardCharsets.UTF_8)
 }
+
+private fun assertAddressEquals(expected: List<InetAddress>, actual: List<InetAddress>) {
+    // No duplicate addresses in the actual address list
+    assertEquals(actual.toSet().size, actual.size)
+    assertEquals(expected.toSet(), actual.toSet())
+}
\ No newline at end of file
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 96330e2..3d828a4 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
@@ -173,21 +173,39 @@
         return cb;
     }
 
-    // Toggle WiFi twice, leaving it in the state it started in
-    public void toggleWifi() throws Exception {
-        if (mWifiManager.isWifiEnabled()) {
-            Network wifiNetwork = getWifiNetwork();
-            // Ensure system default network is WIFI because it's expected in disconnectFromWifi()
-            expectNetworkIsSystemDefault(wifiNetwork);
-            disconnectFromWifi(wifiNetwork);
-            connectToWifi();
-        } else {
-            connectToWifi();
-            Network wifiNetwork = getWifiNetwork();
-            // Ensure system default network is WIFI because it's expected in disconnectFromWifi()
-            expectNetworkIsSystemDefault(wifiNetwork);
-            disconnectFromWifi(wifiNetwork);
+    /**
+     * Toggle Wi-Fi off and on, waiting for the {@link ConnectivityManager#CONNECTIVITY_ACTION}
+     * broadcast in both cases.
+     */
+    public void reconnectWifiAndWaitForConnectivityAction() throws Exception {
+        assertTrue(mWifiManager.isWifiEnabled());
+        Network wifiNetwork = getWifiNetwork();
+        // Ensure system default network is WIFI because it's expected in disconnectFromWifi()
+        expectNetworkIsSystemDefault(wifiNetwork);
+        disconnectFromWifi(wifiNetwork, true /* expectLegacyBroadcast */);
+        connectToWifi(true /* expectLegacyBroadcast */);
+    }
+
+    /**
+     * Turn Wi-Fi off, then back on and make sure it connects, if it is supported.
+     */
+    public void reconnectWifiIfSupported() throws Exception {
+        if (!mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_WIFI)) {
+            return;
         }
+        disableWifi();
+        ensureWifiConnected();
+    }
+
+    /**
+     * Turn cell data off, then back on and make sure it connects, if it is supported.
+     */
+    public void reconnectCellIfSupported() throws Exception {
+        if (!mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_TELEPHONY)) {
+            return;
+        }
+        setMobileDataEnabled(false);
+        setMobileDataEnabled(true);
     }
 
     public Network expectNetworkIsSystemDefault(Network network)
diff --git a/tests/unit/java/android/net/NetworkUtilsTest.java b/tests/unit/java/android/net/NetworkUtilsTest.java
index 2bf2211..e453c02 100644
--- a/tests/unit/java/android/net/NetworkUtilsTest.java
+++ b/tests/unit/java/android/net/NetworkUtilsTest.java
@@ -159,10 +159,9 @@
         return timeval;
     }
 
-    @Test
-    public void testSetSockOptBytes() throws ErrnoException {
-        final FileDescriptor sock = Os.socket(AF_INET6, SOCK_DGRAM, IPPROTO_ICMPV6);
-        final StructTimeval writeTimeval = StructTimeval.fromMillis(1200);
+    private void testSetSockOptBytes(FileDescriptor sock, long timeValMillis)
+            throws ErrnoException {
+        final StructTimeval writeTimeval = StructTimeval.fromMillis(timeValMillis);
         byte[] timeval = getTimevalBytes(writeTimeval);
         final StructTimeval readTimeval;
 
@@ -170,6 +169,16 @@
         readTimeval = Os.getsockoptTimeval(sock, SOL_SOCKET, SO_RCVTIMEO);
 
         assertEquals(writeTimeval, readTimeval);
+    }
+
+    @Test
+    public void testSetSockOptBytes() throws ErrnoException {
+        final FileDescriptor sock = Os.socket(AF_INET6, SOCK_DGRAM, IPPROTO_ICMPV6);
+
+        testSetSockOptBytes(sock, 3000);
+
+        testSetSockOptBytes(sock, 5000);
+
         SocketUtils.closeSocketQuietly(sock);
     }
 
diff --git a/tests/unit/java/android/net/nsd/NsdManagerTest.java b/tests/unit/java/android/net/nsd/NsdManagerTest.java
index 461ead8..951675c 100644
--- a/tests/unit/java/android/net/nsd/NsdManagerTest.java
+++ b/tests/unit/java/android/net/nsd/NsdManagerTest.java
@@ -52,6 +52,9 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.net.InetAddress;
+import java.util.List;
+
 @DevSdkIgnoreRunner.MonitorThreadLeak
 @RunWith(DevSdkIgnoreRunner.class)
 @SmallTest
@@ -285,6 +288,7 @@
     private void doTestDiscoverService() throws Exception {
         NsdManager manager = mManager;
 
+        DiscoveryRequest request1 = new DiscoveryRequest.Builder("a_type").build();
         NsdServiceInfo reply1 = new NsdServiceInfo("a_name", "a_type");
         NsdServiceInfo reply2 = new NsdServiceInfo("another_name", "a_type");
         NsdServiceInfo reply3 = new NsdServiceInfo("a_third_name", "a_type");
@@ -305,7 +309,7 @@
         int key2 = getRequestKey(req ->
                 verify(mServiceConn, times(2)).discoverServices(req.capture(), any()));
 
-        mCallback.onDiscoverServicesStarted(key2, reply1);
+        mCallback.onDiscoverServicesStarted(key2, request1);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStarted("a_type");
 
 
@@ -345,7 +349,7 @@
         int key3 = getRequestKey(req ->
                 verify(mServiceConn, times(3)).discoverServices(req.capture(), any()));
 
-        mCallback.onDiscoverServicesStarted(key3, reply1);
+        mCallback.onDiscoverServicesStarted(key3, request1);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStarted("a_type");
 
         // Client unregisters immediately, it fails
@@ -369,6 +373,9 @@
         NsdManager.RegistrationListener listener1 = mock(NsdManager.RegistrationListener.class);
         NsdManager.DiscoveryListener listener2 = mock(NsdManager.DiscoveryListener.class);
         NsdManager.ResolveListener listener3 = mock(NsdManager.ResolveListener.class);
+        NsdManager.RegistrationListener listener4 = mock(NsdManager.RegistrationListener.class);
+        NsdManager.RegistrationListener listener5 = mock(NsdManager.RegistrationListener.class);
+        NsdManager.RegistrationListener listener6 = mock(NsdManager.RegistrationListener.class);
 
         NsdServiceInfo invalidService = new NsdServiceInfo(null, null);
         NsdServiceInfo validService = new NsdServiceInfo("a_name", "_a_type._tcp");
@@ -378,6 +385,7 @@
                 "_a_type._tcp,_sub1,_s2");
         NsdServiceInfo otherSubtypeUpdate = new NsdServiceInfo("a_name", "_a_type._tcp,_sub1,_s3");
         NsdServiceInfo dotSyntaxSubtypeUpdate = new NsdServiceInfo("a_name", "_sub1._a_type._tcp");
+
         validService.setPort(2222);
         otherServiceWithSubtype.setPort(2222);
         validServiceDuplicate.setPort(2222);
@@ -385,6 +393,33 @@
         otherSubtypeUpdate.setPort(2222);
         dotSyntaxSubtypeUpdate.setPort(2222);
 
+        NsdServiceInfo invalidMissingHostnameWithAddresses = new NsdServiceInfo(null, null);
+        invalidMissingHostnameWithAddresses.setHostAddresses(
+                List.of(
+                        InetAddress.parseNumericAddress("192.168.82.14"),
+                        InetAddress.parseNumericAddress("2001::1")));
+
+        NsdServiceInfo validCustomHostWithAddresses = new NsdServiceInfo(null, null);
+        validCustomHostWithAddresses.setHostname("a_host");
+        validCustomHostWithAddresses.setHostAddresses(
+                List.of(
+                        InetAddress.parseNumericAddress("192.168.82.14"),
+                        InetAddress.parseNumericAddress("2001::1")));
+
+        NsdServiceInfo validServiceWithCustomHostAndAddresses =
+                new NsdServiceInfo("a_name", "_a_type._tcp");
+        validServiceWithCustomHostAndAddresses.setPort(2222);
+        validServiceWithCustomHostAndAddresses.setHostname("a_host");
+        validServiceWithCustomHostAndAddresses.setHostAddresses(
+                List.of(
+                        InetAddress.parseNumericAddress("192.168.82.14"),
+                        InetAddress.parseNumericAddress("2001::1")));
+
+        NsdServiceInfo validServiceWithCustomHostNoAddresses =
+                new NsdServiceInfo("a_name", "_a_type._tcp");
+        validServiceWithCustomHostNoAddresses.setPort(2222);
+        validServiceWithCustomHostNoAddresses.setHostname("a_host");
+
         // Service registration
         //  - invalid arguments
         mustFail(() -> { manager.unregisterService(null); });
@@ -393,6 +428,8 @@
         mustFail(() -> { manager.registerService(invalidService, PROTOCOL, listener1); });
         mustFail(() -> { manager.registerService(validService, -1, listener1); });
         mustFail(() -> { manager.registerService(validService, PROTOCOL, null); });
+        mustFail(() -> {
+            manager.registerService(invalidMissingHostnameWithAddresses, PROTOCOL, listener1); });
         manager.registerService(validService, PROTOCOL, listener1);
         //  - update without subtype is not allowed
         mustFail(() -> { manager.registerService(validServiceDuplicate, PROTOCOL, listener1); });
@@ -414,6 +451,15 @@
         // TODO: make listener immediately reusable
         //mustFail(() -> { manager.unregisterService(listener1); });
         //manager.registerService(validService, PROTOCOL, listener1);
+        //  - registering a custom host without a service is valid
+        manager.registerService(validCustomHostWithAddresses, PROTOCOL, listener4);
+        manager.unregisterService(listener4);
+        //  - registering a service with a custom host is valid
+        manager.registerService(validServiceWithCustomHostAndAddresses, PROTOCOL, listener5);
+        manager.unregisterService(listener5);
+        //  - registering a service with a custom host with no addresses is valid
+        manager.registerService(validServiceWithCustomHostNoAddresses, PROTOCOL, listener6);
+        manager.unregisterService(listener6);
 
         // Discover service
         //  - invalid arguments
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index a17197e..b60f0b4 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -1139,7 +1139,7 @@
         verify(mAdvertiser).addOrUpdateService(anyInt(), argThat(s ->
                 "Instance".equals(s.getServiceName())
                         && SERVICE_TYPE.equals(s.getServiceType())
-                        && s.getSubtypes().equals(Set.of("_subtype"))), any());
+                        && s.getSubtypes().equals(Set.of("_subtype"))), any(), anyInt());
 
         final DiscoveryListener discListener = mock(DiscoveryListener.class);
         client.discoverServices(typeWithSubtype, PROTOCOL, network, Runnable::run, discListener);
@@ -1246,7 +1246,7 @@
 
         final ArgumentCaptor<Integer> serviceIdCaptor = ArgumentCaptor.forClass(Integer.class);
         verify(mAdvertiser).addOrUpdateService(serviceIdCaptor.capture(),
-                argThat(info -> matches(info, regInfo)), any());
+                argThat(info -> matches(info, regInfo)), any(), anyInt());
 
         client.unregisterService(regListenerWithoutFeature);
         waitForIdle();
@@ -1275,7 +1275,7 @@
         service1.setHostAddresses(List.of(parseNumericAddress("2001:db8::123")));
         service1.setPort(1234);
         final NsdServiceInfo service2 = new NsdServiceInfo(SERVICE_NAME, "_type2._tcp");
-        service2.setHostAddresses(List.of(parseNumericAddress("2001:db8::123")));
+        service1.setHostAddresses(List.of(parseNumericAddress("2001:db8::123")));
         service2.setPort(1234);
 
         client.discoverServices(service1.getServiceType(),
@@ -1307,9 +1307,9 @@
 
         // The advertiser is enabled for _type2 but not _type1
         verify(mAdvertiser, never()).addOrUpdateService(anyInt(),
-                argThat(info -> matches(info, service1)), any());
+                argThat(info -> matches(info, service1)), any(), anyInt());
         verify(mAdvertiser).addOrUpdateService(anyInt(), argThat(info -> matches(info, service2)),
-                any());
+                any(), anyInt());
     }
 
     @Test
@@ -1334,7 +1334,7 @@
         verify(mSocketProvider).startMonitoringSockets();
         final ArgumentCaptor<Integer> idCaptor = ArgumentCaptor.forClass(Integer.class);
         verify(mAdvertiser).addOrUpdateService(idCaptor.capture(), argThat(info ->
-                matches(info, regInfo)), any());
+                matches(info, regInfo)), any(), anyInt());
 
         // Verify onServiceRegistered callback
         final MdnsAdvertiser.AdvertiserCallback cb = cbCaptor.getValue();
@@ -1382,7 +1382,7 @@
 
         client.registerService(regInfo, NsdManager.PROTOCOL_DNS_SD, Runnable::run, regListener);
         waitForIdle();
-        verify(mAdvertiser, never()).addOrUpdateService(anyInt(), any(), any());
+        verify(mAdvertiser, never()).addOrUpdateService(anyInt(), any(), any(), anyInt());
 
         verify(regListener, timeout(TIMEOUT_MS)).onRegistrationFailed(
                 argThat(info -> matches(info, regInfo)), eq(FAILURE_INTERNAL_ERROR));
@@ -1411,8 +1411,12 @@
         waitForIdle();
         final ArgumentCaptor<Integer> idCaptor = ArgumentCaptor.forClass(Integer.class);
         // Service name is truncated to 63 characters
-        verify(mAdvertiser).addOrUpdateService(idCaptor.capture(),
-                argThat(info -> info.getServiceName().equals("a".repeat(63))), any());
+        verify(mAdvertiser)
+                .addOrUpdateService(
+                        idCaptor.capture(),
+                        argThat(info -> info.getServiceName().equals("a".repeat(63))),
+                        any(),
+                        anyInt());
 
         // Verify onServiceRegistered callback
         final MdnsAdvertiser.AdvertiserCallback cb = cbCaptor.getValue();
@@ -1510,7 +1514,7 @@
         client.registerService(regInfo, NsdManager.PROTOCOL_DNS_SD, Runnable::run, regListener);
         waitForIdle();
         verify(mSocketProvider).startMonitoringSockets();
-        verify(mAdvertiser).addOrUpdateService(anyInt(), any(), any());
+        verify(mAdvertiser).addOrUpdateService(anyInt(), any(), any(), anyInt());
 
         // Verify the discovery uses MdnsDiscoveryManager
         final DiscoveryListener discListener = mock(DiscoveryListener.class);
@@ -1543,7 +1547,7 @@
         client.registerService(regInfo, NsdManager.PROTOCOL_DNS_SD, Runnable::run, regListener);
         waitForIdle();
         verify(mSocketProvider).startMonitoringSockets();
-        verify(mAdvertiser).addOrUpdateService(anyInt(), any(), any());
+        verify(mAdvertiser).addOrUpdateService(anyInt(), any(), any(), anyInt());
 
         final Network wifiNetwork1 = new Network(123);
         final Network wifiNetwork2 = new Network(124);
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt b/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt
index 5c04362..f753c93 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsAdvertiserTest.kt
@@ -32,6 +32,7 @@
 import com.android.net.module.util.SharedLog
 import com.android.server.connectivity.ConnectivityResources
 import com.android.server.connectivity.mdns.MdnsAdvertiser.AdvertiserCallback
+import com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_SERVICE
 import com.android.server.connectivity.mdns.MdnsSocketProvider.SocketCallback
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
 import com.android.testutils.DevSdkIgnoreRunner
@@ -76,6 +77,7 @@
 private const val TEST_SUBTYPE2 = "_subtype2"
 private val TEST_INTERFACE1 = "test_iface1"
 private val TEST_INTERFACE2 = "test_iface2"
+private val TEST_CLIENT_UID_1 = 10010
 private val TEST_OFFLOAD_PACKET1 = byteArrayOf(0x01, 0x02, 0x03)
 private val TEST_OFFLOAD_PACKET2 = byteArrayOf(0x02, 0x03, 0x04)
 private val DEFAULT_ADVERTISING_OPTION = MdnsAdvertisingOptions.getDefaultOptions()
@@ -227,7 +229,7 @@
         val advertiser =
             MdnsAdvertiser(thread.looper, socketProvider, cb, mockDeps, sharedlog, flags, context)
         postSync { advertiser.addOrUpdateService(SERVICE_ID_1, SERVICE_1,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
 
         val socketCbCaptor = ArgumentCaptor.forClass(SocketCallback::class.java)
         verify(socketProvider).requestSocket(eq(TEST_NETWORK_1), socketCbCaptor.capture())
@@ -254,7 +256,10 @@
         verify(cb).onOffloadStartOrUpdate(eq(TEST_INTERFACE1), eq(OFFLOAD_SERVICEINFO_NO_SUBTYPE))
 
         // Service is conflicted.
-        postSync { intAdvCbCaptor.value.onServiceConflict(mockInterfaceAdvertiser1, SERVICE_ID_1) }
+        postSync {
+            intAdvCbCaptor.value
+                    .onServiceConflict(mockInterfaceAdvertiser1, SERVICE_ID_1, CONFLICT_SERVICE)
+        }
 
         // Verify the metrics data
         doReturn(25).`when`(mockInterfaceAdvertiser1).getServiceRepliedRequestsCount(SERVICE_ID_1)
@@ -289,7 +294,7 @@
         val advertiser =
             MdnsAdvertiser(thread.looper, socketProvider, cb, mockDeps, sharedlog, flags, context)
         postSync { advertiser.addOrUpdateService(SERVICE_ID_1, ALL_NETWORKS_SERVICE_SUBTYPE,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
 
         val socketCbCaptor = ArgumentCaptor.forClass(SocketCallback::class.java)
         verify(socketProvider).requestSocket(eq(ALL_NETWORKS_SERVICE_SUBTYPE.network),
@@ -327,9 +332,18 @@
                 argThat { it.matches(ALL_NETWORKS_SERVICE_SUBTYPE) })
 
         // Services are conflicted.
-        postSync { intAdvCbCaptor1.value.onServiceConflict(mockInterfaceAdvertiser1, SERVICE_ID_1) }
-        postSync { intAdvCbCaptor1.value.onServiceConflict(mockInterfaceAdvertiser1, SERVICE_ID_1) }
-        postSync { intAdvCbCaptor2.value.onServiceConflict(mockInterfaceAdvertiser2, SERVICE_ID_1) }
+        postSync {
+            intAdvCbCaptor1.value
+                    .onServiceConflict(mockInterfaceAdvertiser1, SERVICE_ID_1, CONFLICT_SERVICE)
+        }
+        postSync {
+            intAdvCbCaptor1.value
+                    .onServiceConflict(mockInterfaceAdvertiser1, SERVICE_ID_1, CONFLICT_SERVICE)
+        }
+        postSync {
+            intAdvCbCaptor2.value
+                    .onServiceConflict(mockInterfaceAdvertiser2, SERVICE_ID_1, CONFLICT_SERVICE)
+        }
 
         // Verify the metrics data
         doReturn(10).`when`(mockInterfaceAdvertiser1).getServiceRepliedRequestsCount(SERVICE_ID_1)
@@ -361,18 +375,19 @@
         val advertiser =
             MdnsAdvertiser(thread.looper, socketProvider, cb, mockDeps, sharedlog, flags, context)
         postSync {
-            advertiser.addOrUpdateService(SERVICE_ID_1, SERVICE_1, DEFAULT_ADVERTISING_OPTION)
+            advertiser.addOrUpdateService(SERVICE_ID_1, SERVICE_1, DEFAULT_ADVERTISING_OPTION,
+                    TEST_CLIENT_UID_1)
             advertiser.addOrUpdateService(SERVICE_ID_2,
                 NsdServiceInfo("TestService2", "_PRIORITYTEST._udp").apply {
                     port = 12345
                     hostAddresses = listOf(TEST_ADDR)
-                }, DEFAULT_ADVERTISING_OPTION)
+                }, DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1)
             advertiser.addOrUpdateService(
                 SERVICE_ID_3,
                 NsdServiceInfo("TestService3", "_notprioritized._tcp").apply {
                     port = 12345
                     hostAddresses = listOf(TEST_ADDR)
-                }, DEFAULT_ADVERTISING_OPTION)
+                }, DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1)
         }
 
         val socketCbCaptor = ArgumentCaptor.forClass(SocketCallback::class.java)
@@ -419,7 +434,7 @@
         val advertiser =
             MdnsAdvertiser(thread.looper, socketProvider, cb, mockDeps, sharedlog, flags, context)
         postSync { advertiser.addOrUpdateService(SERVICE_ID_1, SERVICE_1,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
 
         val oneNetSocketCbCaptor = ArgumentCaptor.forClass(SocketCallback::class.java)
         verify(socketProvider).requestSocket(eq(TEST_NETWORK_1), oneNetSocketCbCaptor.capture())
@@ -427,18 +442,18 @@
 
         // Register a service with the same name on all networks (name conflict)
         postSync { advertiser.addOrUpdateService(SERVICE_ID_2, ALL_NETWORKS_SERVICE,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
         val allNetSocketCbCaptor = ArgumentCaptor.forClass(SocketCallback::class.java)
         verify(socketProvider).requestSocket(eq(null), allNetSocketCbCaptor.capture())
         val allNetSocketCb = allNetSocketCbCaptor.value
 
         postSync { advertiser.addOrUpdateService(LONG_SERVICE_ID_1, LONG_SERVICE_1,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
         postSync { advertiser.addOrUpdateService(LONG_SERVICE_ID_2, LONG_ALL_NETWORKS_SERVICE,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
 
         postSync { advertiser.addOrUpdateService(CASE_INSENSITIVE_TEST_SERVICE_ID,
-                ALL_NETWORKS_SERVICE_2, DEFAULT_ADVERTISING_OPTION) }
+                ALL_NETWORKS_SERVICE_2, DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
 
         // Callbacks for matching network and all networks both get the socket
         postSync {
@@ -508,7 +523,7 @@
                 MdnsAdvertiser(thread.looper, socketProvider, cb, mockDeps, sharedlog, flags,
                     context)
         postSync { advertiser.addOrUpdateService(SERVICE_ID_1, ALL_NETWORKS_SERVICE,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
 
         val socketCbCaptor = ArgumentCaptor.forClass(SocketCallback::class.java)
         verify(socketProvider).requestSocket(eq(null), socketCbCaptor.capture())
@@ -523,16 +538,17 @@
 
         // Update with serviceId that is not registered yet should fail
         postSync { advertiser.addOrUpdateService(SERVICE_ID_2, ALL_NETWORKS_SERVICE_SUBTYPE,
-                updateOptions) }
+                updateOptions, TEST_CLIENT_UID_1) }
         verify(cb).onRegisterServiceFailed(SERVICE_ID_2, NsdManager.FAILURE_INTERNAL_ERROR)
 
         // Update service with different NsdServiceInfo should fail
-        postSync { advertiser.addOrUpdateService(SERVICE_ID_1, SERVICE_1_SUBTYPE, updateOptions) }
+        postSync { advertiser.addOrUpdateService(SERVICE_ID_1, SERVICE_1_SUBTYPE, updateOptions,
+                TEST_CLIENT_UID_1) }
         verify(cb).onRegisterServiceFailed(SERVICE_ID_1, NsdManager.FAILURE_INTERNAL_ERROR)
 
         // Update service with same NsdServiceInfo but different subType should succeed
         postSync { advertiser.addOrUpdateService(SERVICE_ID_1, ALL_NETWORKS_SERVICE_SUBTYPE,
-                updateOptions) }
+                updateOptions, TEST_CLIENT_UID_1) }
         verify(mockInterfaceAdvertiser1).updateService(eq(SERVICE_ID_1), eq(setOf(TEST_SUBTYPE)))
 
         // Newly created MdnsInterfaceAdvertiser will get addService() call.
@@ -547,7 +563,7 @@
             MdnsAdvertiser(thread.looper, socketProvider, cb, mockDeps, sharedlog, flags, context)
         verify(mockDeps, times(1)).generateHostname()
         postSync { advertiser.addOrUpdateService(SERVICE_ID_1, SERVICE_1,
-                DEFAULT_ADVERTISING_OPTION) }
+                DEFAULT_ADVERTISING_OPTION, TEST_CLIENT_UID_1) }
         postSync { advertiser.removeService(SERVICE_ID_1) }
         verify(mockDeps, times(2)).generateHostname()
     }
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 0e5cc50..0637ad1 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiserTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiserTest.kt
@@ -18,6 +18,7 @@
 
 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
@@ -26,6 +27,7 @@
 import com.android.server.connectivity.mdns.MdnsAnnouncer.AnnouncementInfo
 import com.android.server.connectivity.mdns.MdnsAnnouncer.BaseAnnouncementInfo
 import com.android.server.connectivity.mdns.MdnsAnnouncer.ExitAnnouncementInfo
+import com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.CONFLICT_SERVICE
 import com.android.server.connectivity.mdns.MdnsInterfaceAdvertiser.EXIT_ANNOUNCEMENT_DELAY_MS
 import com.android.server.connectivity.mdns.MdnsPacketRepeater.PacketRepeaterCallback
 import com.android.server.connectivity.mdns.MdnsProber.ProbingInfo
@@ -347,7 +349,8 @@
     @Test
     fun testConflict() {
         addServiceAndFinishProbing(TEST_SERVICE_ID_1, TEST_SERVICE_1)
-        doReturn(setOf(TEST_SERVICE_ID_1)).`when`(repository).getConflictingServices(any())
+        doReturn(mapOf(TEST_SERVICE_ID_1 to CONFLICT_SERVICE))
+                .`when`(repository).getConflictingServices(any())
 
         // Reply obtained with:
         // scapy.raw(scapy.DNS(
@@ -373,7 +376,7 @@
         }
 
         thread.waitForIdle(TIMEOUT_MS)
-        verify(cb).onServiceConflict(advertiser, TEST_SERVICE_ID_1)
+        verify(cb).onServiceConflict(advertiser, TEST_SERVICE_ID_1, CONFLICT_SERVICE)
     }
 
     @Test
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 06f12fe..fd8d98b 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsRecordRepositoryTest.kt
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsRecordRepositoryTest.kt
@@ -22,6 +22,8 @@
 import android.os.Build
 import android.os.HandlerThread
 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.MdnsRecord.TYPE_A
 import com.android.server.connectivity.mdns.MdnsRecord.TYPE_AAAA
 import com.android.server.connectivity.mdns.MdnsRecord.TYPE_PTR
@@ -52,6 +54,9 @@
 private const val TEST_SERVICE_ID_1 = 42
 private const val TEST_SERVICE_ID_2 = 43
 private const val TEST_SERVICE_ID_3 = 44
+private const val TEST_CUSTOM_HOST_ID_1 = 45
+private const val TEST_CUSTOM_HOST_ID_2 = 46
+private const val TEST_SERVICE_CUSTOM_HOST_ID_1 = 48
 private const val TEST_PORT = 12345
 private const val TEST_SUBTYPE = "_subtype"
 private const val TEST_SUBTYPE2 = "_subtype2"
@@ -86,6 +91,26 @@
     port = TEST_PORT
 }
 
+private val TEST_CUSTOM_HOST_1 = NsdServiceInfo().apply {
+    hostname = "TestHost"
+    hostAddresses = listOf(parseNumericAddress("2001:db8::1"), parseNumericAddress("2001:db8::2"))
+}
+
+private val TEST_CUSTOM_HOST_1_NAME = arrayOf("TestHost", "local")
+
+private val TEST_CUSTOM_HOST_2 = NsdServiceInfo().apply {
+    hostname = "OtherTestHost"
+    hostAddresses = listOf(parseNumericAddress("2001:db8::3"), parseNumericAddress("2001:db8::4"))
+}
+
+private val TEST_SERVICE_CUSTOM_HOST_1 = NsdServiceInfo().apply {
+    hostname = "TestHost"
+    hostAddresses = listOf(parseNumericAddress("2001:db8::1"))
+    serviceType = "_testservice._tcp"
+    serviceName = "TestService"
+    port = TEST_PORT
+}
+
 @RunWith(DevSdkIgnoreRunner::class)
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
 class MdnsRecordRepositoryTest {
@@ -569,6 +594,92 @@
             ), reply.additionalAnswers)
     }
 
+
+    @Test
+    fun testGetReply_ptrQuestionForServiceWithCustomHost_customHostUsedInAdditionalAnswers() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.initWithService(TEST_SERVICE_CUSTOM_HOST_ID_1, TEST_SERVICE_CUSTOM_HOST_1,
+                setOf(TEST_SUBTYPE, TEST_SUBTYPE2))
+        val src = InetSocketAddress(parseNumericAddress("fe80::1234"), 5353)
+        val serviceName = arrayOf("TestService", "_testservice", "_tcp", "local")
+
+        val query = makeQuery(TYPE_PTR to arrayOf("_testservice", "_tcp", "local"))
+        val reply = repository.getReply(query, src)
+
+        assertNotNull(reply)
+        assertEquals(listOf(
+                MdnsPointerRecord(
+                        arrayOf("_testservice", "_tcp", "local"),
+                        0L, false, LONG_TTL, serviceName)),
+                reply.answers)
+        assertEquals(listOf(
+                MdnsTextRecord(serviceName, 0L, true, LONG_TTL, listOf()),
+                MdnsServiceRecord(serviceName, 0L, true, SHORT_TTL,
+                        0, 0, TEST_PORT, TEST_CUSTOM_HOST_1_NAME),
+                MdnsInetAddressRecord(
+                        TEST_CUSTOM_HOST_1_NAME, 0L, true, SHORT_TTL,
+                        parseNumericAddress("2001:db8::1")),
+                MdnsNsecRecord(serviceName, 0L, true, LONG_TTL, serviceName /* nextDomain */,
+                        intArrayOf(TYPE_TXT, TYPE_SRV)),
+                MdnsNsecRecord(TEST_CUSTOM_HOST_1_NAME, 0L, true, SHORT_TTL,
+                        TEST_CUSTOM_HOST_1_NAME /* nextDomain */,
+                        intArrayOf(TYPE_AAAA)),
+        ), reply.additionalAnswers)
+    }
+
+    @Test
+    fun testGetReply_ptrQuestionForServicesWithSameCustomHost_customHostUsedInAdditionalAnswers() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        val serviceWithCustomHost1 = NsdServiceInfo().apply {
+            hostname = "TestHost"
+            hostAddresses = listOf(
+                    parseNumericAddress("2001:db8::1"),
+                    parseNumericAddress("192.0.2.1"))
+            serviceType = "_testservice._tcp"
+            serviceName = "TestService1"
+            port = TEST_PORT
+        }
+        val serviceWithCustomHost2 = NsdServiceInfo().apply {
+            hostname = "TestHost"
+            hostAddresses = listOf(
+                    parseNumericAddress("2001:db8::1"),
+                    parseNumericAddress("2001:db8::3"))
+        }
+        repository.addServiceAndFinishProbing(TEST_SERVICE_ID_1, serviceWithCustomHost1)
+        repository.addServiceAndFinishProbing(TEST_SERVICE_ID_2, serviceWithCustomHost2)
+        val src = InetSocketAddress(parseNumericAddress("fe80::1234"), 5353)
+        val serviceName = arrayOf("TestService1", "_testservice", "_tcp", "local")
+
+        val query = makeQuery(TYPE_PTR to arrayOf("_testservice", "_tcp", "local"))
+        val reply = repository.getReply(query, src)
+
+        assertNotNull(reply)
+        assertEquals(listOf(
+                MdnsPointerRecord(
+                        arrayOf("_testservice", "_tcp", "local"),
+                        0L, false, LONG_TTL, serviceName)),
+                reply.answers)
+        assertEquals(listOf(
+                MdnsTextRecord(serviceName, 0L, true, LONG_TTL, listOf()),
+                MdnsServiceRecord(serviceName, 0L, true, SHORT_TTL,
+                        0, 0, TEST_PORT, TEST_CUSTOM_HOST_1_NAME),
+                MdnsInetAddressRecord(
+                        TEST_CUSTOM_HOST_1_NAME, 0L, true, SHORT_TTL,
+                        parseNumericAddress("2001:db8::1")),
+                MdnsInetAddressRecord(
+                        TEST_CUSTOM_HOST_1_NAME, 0L, true, SHORT_TTL,
+                        parseNumericAddress("192.0.2.1")),
+                MdnsInetAddressRecord(
+                        TEST_CUSTOM_HOST_1_NAME, 0L, true, SHORT_TTL,
+                        parseNumericAddress("2001:db8::3")),
+                MdnsNsecRecord(serviceName, 0L, true, LONG_TTL, serviceName /* nextDomain */,
+                        intArrayOf(TYPE_TXT, TYPE_SRV)),
+                MdnsNsecRecord(TEST_CUSTOM_HOST_1_NAME, 0L, true, SHORT_TTL,
+                        TEST_CUSTOM_HOST_1_NAME /* nextDomain */,
+                        intArrayOf(TYPE_A, TYPE_AAAA)),
+        ), reply.additionalAnswers)
+    }
+
     @Test
     fun testGetReply_singleSubtypePtrQuestion_returnsSrvTxtAddressNsecRecords() {
         val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
@@ -708,6 +819,90 @@
     }
 
     @Test
+    fun testGetReply_AAAAQuestionForCustomHost_returnsAAAARecords() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.initWithService(
+                TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1, subtypes = setOf(),
+                listOf(LinkAddress(parseNumericAddress("192.0.2.111"), 24)))
+        repository.addService(TEST_CUSTOM_HOST_ID_2, TEST_CUSTOM_HOST_2)
+        val src = InetSocketAddress(parseNumericAddress("fe80::123"), 5353)
+
+        val query = makeQuery(TYPE_AAAA to TEST_CUSTOM_HOST_1_NAME)
+        val reply = repository.getReply(query, src)
+
+        assertNotNull(reply)
+        assertEquals(listOf(
+                MdnsInetAddressRecord(TEST_CUSTOM_HOST_1_NAME,
+                        0, false, LONG_TTL, parseNumericAddress("2001:db8::1")),
+                MdnsInetAddressRecord(TEST_CUSTOM_HOST_1_NAME,
+                        0, false, LONG_TTL, parseNumericAddress("2001:db8::2"))),
+                reply.answers)
+        assertEquals(
+                listOf(MdnsNsecRecord(TEST_CUSTOM_HOST_1_NAME,
+                        0L, true, SHORT_TTL,
+                        TEST_CUSTOM_HOST_1_NAME /* nextDomain */,
+                        intArrayOf(TYPE_AAAA))),
+                reply.additionalAnswers)
+    }
+
+
+    @Test
+    fun testGetReply_AAAAQuestionForCustomHostInMultipleRegistrations_returnsAAAARecords() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+
+        repository.addServiceAndFinishProbing(TEST_CUSTOM_HOST_ID_1, NsdServiceInfo().apply {
+            hostname = "TestHost"
+            hostAddresses = listOf(
+                    parseNumericAddress("2001:db8::1"),
+                    parseNumericAddress("2001:db8::2"))
+        })
+        repository.addServiceAndFinishProbing(TEST_CUSTOM_HOST_ID_2, NsdServiceInfo().apply {
+            hostname = "TestHost"
+            hostAddresses = listOf(
+                    parseNumericAddress("2001:db8::1"),
+                    parseNumericAddress("2001:db8::3"))
+        })
+        val src = InetSocketAddress(parseNumericAddress("fe80::123"), 5353)
+
+        val query = makeQuery(TYPE_AAAA to TEST_CUSTOM_HOST_1_NAME)
+        val reply = repository.getReply(query, src)
+
+        assertNotNull(reply)
+        assertEquals(listOf(
+                MdnsInetAddressRecord(TEST_CUSTOM_HOST_1_NAME,
+                        0, false, LONG_TTL, parseNumericAddress("2001:db8::1")),
+                MdnsInetAddressRecord(TEST_CUSTOM_HOST_1_NAME,
+                        0, false, LONG_TTL, parseNumericAddress("2001:db8::2")),
+                MdnsInetAddressRecord(TEST_CUSTOM_HOST_1_NAME,
+                        0, false, LONG_TTL, parseNumericAddress("2001:db8::3"))),
+                reply.answers)
+        assertEquals(
+                listOf(MdnsNsecRecord(TEST_CUSTOM_HOST_1_NAME,
+                        0L, true, SHORT_TTL,
+                        TEST_CUSTOM_HOST_1_NAME /* nextDomain */,
+                        intArrayOf(TYPE_AAAA))),
+                reply.additionalAnswers)
+    }
+
+    @Test
+    fun testGetReply_customHostRemoved_noAnswerToAAAAQuestion() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.initWithService(
+                TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1, subtypes = setOf(),
+                listOf(LinkAddress(parseNumericAddress("192.0.2.111"), 24)))
+        repository.addService(TEST_SERVICE_CUSTOM_HOST_ID_1, TEST_SERVICE_CUSTOM_HOST_1)
+        repository.removeService(TEST_CUSTOM_HOST_ID_1)
+        repository.removeService(TEST_SERVICE_CUSTOM_HOST_ID_1)
+
+        val src = InetSocketAddress(parseNumericAddress("fe80::123"), 5353)
+
+        val query = makeQuery(TYPE_AAAA to TEST_CUSTOM_HOST_1_NAME)
+        val reply = repository.getReply(query, src)
+
+        assertNull(reply)
+    }
+
+    @Test
     fun testGetReply_ptrAndSrvQuestions_doesNotReturnSrvRecordInAdditionalAnswerSection() {
         val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
         repository.initWithService(TEST_SERVICE_ID_1, TEST_SERVICE_1, setOf(TEST_SUBTYPE))
@@ -815,7 +1010,10 @@
                 emptyList() /* authorityRecords */,
                 emptyList() /* additionalRecords */)
 
-        assertEquals(setOf(TEST_SERVICE_ID_1, TEST_SERVICE_ID_2),
+        assertEquals(
+                mapOf(
+                        TEST_SERVICE_ID_1 to CONFLICT_SERVICE,
+                        TEST_SERVICE_ID_2 to CONFLICT_SERVICE),
                 repository.getConflictingServices(packet))
     }
 
@@ -843,8 +1041,131 @@
             emptyList() /* authorityRecords */,
             emptyList() /* additionalRecords */)
 
-        assertEquals(setOf(TEST_SERVICE_ID_1, TEST_SERVICE_ID_2),
-            repository.getConflictingServices(packet))
+        assertEquals(
+                mapOf(TEST_SERVICE_ID_1 to CONFLICT_SERVICE,
+                        TEST_SERVICE_ID_2 to CONFLICT_SERVICE),
+                repository.getConflictingServices(packet))
+    }
+
+    @Test
+    fun testGetConflictingServices_customHosts_differentAddresses() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.addService(TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1)
+        repository.addService(TEST_CUSTOM_HOST_ID_2, TEST_CUSTOM_HOST_2)
+
+        val packet = MdnsPacket(
+                0, /* flags */
+                emptyList(), /* questions */
+                listOf(
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::5")),
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::6")),
+                ) /* answers */,
+                emptyList() /* authorityRecords */,
+                emptyList() /* additionalRecords */)
+
+        assertEquals(mapOf(TEST_CUSTOM_HOST_ID_1 to CONFLICT_HOST),
+                repository.getConflictingServices(packet))
+    }
+
+    @Test
+    fun testGetConflictingServices_customHosts_moreAddressesThanUs_conflict() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.addService(TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1)
+        repository.addService(TEST_CUSTOM_HOST_ID_2, TEST_CUSTOM_HOST_2)
+
+        val packet = MdnsPacket(
+                0, /* flags */
+                emptyList(), /* questions */
+                listOf(
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::1")),
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::2")),
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::3")),
+                ) /* answers */,
+                emptyList() /* authorityRecords */,
+                emptyList() /* additionalRecords */)
+
+        assertEquals(mapOf(TEST_CUSTOM_HOST_ID_1 to CONFLICT_HOST),
+                repository.getConflictingServices(packet))
+    }
+
+    @Test
+    fun testGetConflictingServices_customHostsReplyHasFewerAddressesThanUs_noConflict() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.addService(TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1)
+        repository.addService(TEST_CUSTOM_HOST_ID_2, TEST_CUSTOM_HOST_2)
+
+        val packet = MdnsPacket(
+                0, /* flags */
+                emptyList(), /* questions */
+                listOf(
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::2")),
+                ) /* answers */,
+                emptyList() /* authorityRecords */,
+                emptyList() /* additionalRecords */)
+
+        assertEquals(emptyMap(),
+                repository.getConflictingServices(packet))
+    }
+
+    @Test
+    fun testGetConflictingServices_customHostsReplyHasIdenticalHosts_noConflict() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.addService(TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1)
+        repository.addService(TEST_CUSTOM_HOST_ID_2, TEST_CUSTOM_HOST_2)
+
+        val packet = MdnsPacket(
+                0, /* flags */
+                emptyList(), /* questions */
+                listOf(
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::1")),
+                        MdnsInetAddressRecord(arrayOf("TestHost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::2")),
+                ) /* answers */,
+                emptyList() /* authorityRecords */,
+                emptyList() /* additionalRecords */)
+
+        assertEquals(emptyMap(),
+                repository.getConflictingServices(packet))
+    }
+
+
+    @Test
+    fun testGetConflictingServices_customHostsCaseInsensitiveReplyHasIdenticalHosts_noConflict() {
+        val repository = MdnsRecordRepository(thread.looper, deps, TEST_HOSTNAME, makeFlags())
+        repository.addService(TEST_CUSTOM_HOST_ID_1, TEST_CUSTOM_HOST_1)
+        repository.addService(TEST_CUSTOM_HOST_ID_2, TEST_CUSTOM_HOST_2)
+
+        val packet = MdnsPacket(
+                0, /* flags */
+                emptyList(), /* questions */
+                listOf(
+                        MdnsInetAddressRecord(arrayOf("TESTHOST", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::1")),
+                        MdnsInetAddressRecord(arrayOf("testhost", "local"),
+                                0L /* receiptTimeMillis */, true /* cacheFlush */,
+                                0L /* ttlMillis */, parseNumericAddress("2001:db8::2")),
+                ) /* answers */,
+                emptyList() /* authorityRecords */,
+                emptyList() /* additionalRecords */)
+
+        assertEquals(emptyMap(),
+                repository.getConflictingServices(packet))
     }
 
     @Test
@@ -873,7 +1194,7 @@
                 emptyList() /* additionalRecords */)
 
         // Above records are identical to the actual registrations: no conflict
-        assertEquals(emptySet(), repository.getConflictingServices(packet))
+        assertEquals(emptyMap(), repository.getConflictingServices(packet))
     }
 
     @Test
@@ -902,7 +1223,7 @@
                 emptyList() /* additionalRecords */)
 
         // Above records are identical to the actual registrations: no conflict
-        assertEquals(emptySet(), repository.getConflictingServices(packet))
+        assertEquals(emptyMap(), repository.getConflictingServices(packet))
     }
 
     @Test
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 7a2e4bf..df23da4 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -16,7 +16,13 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.ACTIVE_QUERY_MODE;
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.AGGRESSIVE_QUERY_MODE;
+import static com.android.server.connectivity.mdns.MdnsSearchOptions.PASSIVE_QUERY_MODE;
 import static com.android.server.connectivity.mdns.MdnsServiceTypeClient.EVENT_START_QUERYTASK;
+import static com.android.server.connectivity.mdns.QueryTaskConfig.INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS;
+import static com.android.server.connectivity.mdns.QueryTaskConfig.MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS;
+import static com.android.server.connectivity.mdns.QueryTaskConfig.TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.junit.Assert.assertArrayEquals;
@@ -145,8 +151,8 @@
         MockitoAnnotations.initMocks(this);
         doReturn(TEST_ELAPSED_REALTIME).when(mockDecoderClock).elapsedRealtime();
 
-        expectedIPv4Packets = new DatagramPacket[16];
-        expectedIPv6Packets = new DatagramPacket[16];
+        expectedIPv4Packets = new DatagramPacket[24];
+        expectedIPv6Packets = new DatagramPacket[24];
         socketKey = new SocketKey(mockNetwork, INTERFACE_INDEX);
 
         for (int i = 0; i < expectedIPv4Packets.length; ++i) {
@@ -171,7 +177,15 @@
                 .thenReturn(expectedIPv4Packets[12])
                 .thenReturn(expectedIPv4Packets[13])
                 .thenReturn(expectedIPv4Packets[14])
-                .thenReturn(expectedIPv4Packets[15]);
+                .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(mockPacketWriter.getPacket(IPV6_ADDRESS))
                 .thenReturn(expectedIPv6Packets[0])
@@ -189,7 +203,15 @@
                 .thenReturn(expectedIPv6Packets[12])
                 .thenReturn(expectedIPv6Packets[13])
                 .thenReturn(expectedIPv6Packets[14])
-                .thenReturn(expectedIPv6Packets[15]);
+                .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]);
 
         thread = new HandlerThread("MdnsServiceTypeClientTests");
         thread.start();
@@ -267,8 +289,8 @@
 
     @Test
     public void sendQueries_activeScanMode() {
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(false).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // Always try to remove the task.
         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
@@ -319,8 +341,8 @@
 
     @Test
     public void sendQueries_reentry_activeScanMode() {
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(false).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // Always try to remove the task.
         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
@@ -333,7 +355,7 @@
                 MdnsSearchOptions.newBuilder()
                         .addSubtype(SUBTYPE)
                         .addSubtype("_subtype2")
-                        .setIsPassiveMode(false)
+                        .setQueryMode(ACTIVE_QUERY_MODE)
                         .build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // The previous scheduled task should be canceled.
@@ -353,8 +375,8 @@
 
     @Test
     public void sendQueries_passiveScanMode() {
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(true).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // Always try to remove the task.
         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
@@ -380,8 +402,10 @@
     @Test
     public void sendQueries_activeScanWithQueryBackoff() {
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(
-                        false).setNumOfQueriesBeforeBackoff(11).build();
+                MdnsSearchOptions.newBuilder()
+                        .addSubtype(SUBTYPE)
+                        .setQueryMode(ACTIVE_QUERY_MODE)
+                        .setNumOfQueriesBeforeBackoff(11).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // Always try to remove the task.
         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
@@ -439,8 +463,10 @@
     @Test
     public void sendQueries_passiveScanWithQueryBackoff() {
         MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(
-                        true).setNumOfQueriesBeforeBackoff(3).build();
+                MdnsSearchOptions.newBuilder()
+                        .addSubtype(SUBTYPE)
+                        .setQueryMode(PASSIVE_QUERY_MODE)
+                        .setNumOfQueriesBeforeBackoff(3).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // Always try to remove the task.
         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
@@ -497,8 +523,8 @@
 
     @Test
     public void sendQueries_reentry_passiveScanMode() {
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(true).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // Always try to remove the task.
         verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
@@ -511,7 +537,7 @@
                 MdnsSearchOptions.newBuilder()
                         .addSubtype(SUBTYPE)
                         .addSubtype("_subtype2")
-                        .setIsPassiveMode(true)
+                        .setQueryMode(PASSIVE_QUERY_MODE)
                         .build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // The previous scheduled task should be canceled.
@@ -533,10 +559,10 @@
     @Ignore("MdnsConfigs is not configurable currently.")
     public void testQueryTaskConfig_alwaysAskForUnicastResponse() {
         //MdnsConfigsFlagsImpl.alwaysAskForUnicastResponseInEachBurst.override(true);
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(false).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
         QueryTaskConfig config = new QueryTaskConfig(
-                searchOptions.getSubtypes(), searchOptions.isPassiveMode(),
+                searchOptions.getSubtypes(), searchOptions.getQueryMode(),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */, 3 /* numOfQueriesBeforeBackoff */,
                 socketKey);
 
@@ -564,10 +590,10 @@
 
     @Test
     public void testQueryTaskConfig_askForUnicastInFirstQuery() {
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(false).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(ACTIVE_QUERY_MODE).build();
         QueryTaskConfig config = new QueryTaskConfig(
-                searchOptions.getSubtypes(), searchOptions.isPassiveMode(),
+                searchOptions.getSubtypes(), searchOptions.getQueryMode(),
                 false /* onlyUseIpv6OnIpv6OnlyNetworks */, 3 /* numOfQueriesBeforeBackoff */,
                 socketKey);
 
@@ -595,8 +621,8 @@
 
     @Test
     public void testIfPreviousTaskIsCanceledWhenNewSessionStarts() {
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(true).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
@@ -605,7 +631,7 @@
                 MdnsSearchOptions.newBuilder()
                         .addSubtype(SUBTYPE)
                         .addSubtype("_subtype2")
-                        .setIsPassiveMode(true)
+                        .setQueryMode(PASSIVE_QUERY_MODE)
                         .build();
         startSendAndReceive(mockListenerOne, searchOptions);
 
@@ -624,8 +650,8 @@
     @Ignore("MdnsConfigs is not configurable currently.")
     public void testIfPreviousTaskIsCanceledWhenSessionStops() {
         //MdnsConfigsFlagsImpl.shouldCancelScanTaskWhenFutureIsNull.override(true);
-        MdnsSearchOptions searchOptions =
-                MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE).setIsPassiveMode(true).build();
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(PASSIVE_QUERY_MODE).build();
         startSendAndReceive(mockListenerOne, searchOptions);
         // Change the sutypes and start a new session.
         stopSendAndReceive(mockListenerOne);
@@ -1492,6 +1518,91 @@
     }
 
     @Test
+    public void testProcessResponse_SubtypeChange() {
+        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
+                serviceCache);
+
+        final String matchingInstance = "instance1";
+        final String subtype = "_subtype";
+        final String ipV4Address = "192.0.2.0";
+        final String ipV6Address = "2001:db8::";
+
+        final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
+                .addSubtype("othersub").build();
+
+        startSendAndReceive(mockListenerOne, options);
+
+        // Complete response from instanceName
+        final MdnsPacket packetWithoutSubtype = createResponse(
+                matchingInstance, ipV4Address, 5353, SERVICE_TYPE_LABELS,
+                Collections.emptyMap() /* textAttributes */, TEST_TTL);
+        final MdnsPointerRecord originalPtr = (MdnsPointerRecord) CollectionUtils.findFirst(
+                packetWithoutSubtype.answers, r -> r instanceof MdnsPointerRecord);
+
+        // Add a subtype PTR record
+        final ArrayList<MdnsRecord> newAnswers = new ArrayList<>(packetWithoutSubtype.answers);
+        newAnswers.add(new MdnsPointerRecord(
+                // PTR should be _subtype._sub._type._tcp.local -> instance1._type._tcp.local
+                Stream.concat(Stream.of(subtype, "_sub"), Arrays.stream(SERVICE_TYPE_LABELS))
+                        .toArray(String[]::new),
+                originalPtr.getReceiptTime(), originalPtr.getCacheFlush(), originalPtr.getTtl(),
+                originalPtr.getPointer()));
+        processResponse(new MdnsPacket(
+                packetWithoutSubtype.flags,
+                packetWithoutSubtype.questions,
+                newAnswers,
+                packetWithoutSubtype.authorityRecords,
+                packetWithoutSubtype.additionalRecords), socketKey);
+
+        // The subtype does not match
+        final InOrder inOrder = inOrder(mockListenerOne);
+        inOrder.verify(mockListenerOne, never()).onServiceNameDiscovered(any(), anyBoolean());
+
+        // Add another matching subtype
+        newAnswers.add(new MdnsPointerRecord(
+                // PTR should be _subtype._sub._type._tcp.local -> instance1._type._tcp.local
+                Stream.concat(Stream.of("_othersub", "_sub"), Arrays.stream(SERVICE_TYPE_LABELS))
+                        .toArray(String[]::new),
+                originalPtr.getReceiptTime(), originalPtr.getCacheFlush(), originalPtr.getTtl(),
+                originalPtr.getPointer()));
+        processResponse(new MdnsPacket(
+                packetWithoutSubtype.flags,
+                packetWithoutSubtype.questions,
+                newAnswers,
+                packetWithoutSubtype.authorityRecords,
+                packetWithoutSubtype.additionalRecords), socketKey);
+
+        final ArgumentMatcher<MdnsServiceInfo> subtypeInstanceMatcher = info ->
+                info.getServiceInstanceName().equals(matchingInstance)
+                        && info.getSubtypes().equals(List.of("_subtype", "_othersub"));
+
+        // Service found callbacks are sent now
+        inOrder.verify(mockListenerOne).onServiceNameDiscovered(
+                argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
+        inOrder.verify(mockListenerOne).onServiceFound(
+                argThat(subtypeInstanceMatcher), eq(false) /* isServiceFromCache */);
+
+        // Address update: update callbacks are sent
+        processResponse(createResponse(
+                matchingInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), TEST_TTL), socketKey);
+
+        inOrder.verify(mockListenerOne).onServiceUpdated(argThat(info ->
+                subtypeInstanceMatcher.matches(info)
+                        && info.getIpv4Addresses().equals(List.of(ipV4Address))
+                        && info.getIpv6Addresses().equals(List.of(ipV6Address))));
+
+        // Goodbye: service removed callbacks are sent
+        processResponse(createResponse(
+                matchingInstance, ipV6Address, 5353, SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), 0L /* ttl */), socketKey);
+
+        inOrder.verify(mockListenerOne).onServiceRemoved(matchServiceName(matchingInstance));
+        inOrder.verify(mockListenerOne).onServiceNameRemoved(matchServiceName(matchingInstance));
+    }
+
+    @Test
     public void testNotifySocketDestroyed() throws Exception {
         client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
                 mockDecoderClock, socketKey, mockSharedLog, thread.getLooper(), mockDeps,
@@ -1666,6 +1777,118 @@
                 socketKey);
     }
 
+    private int getBetweenBurstTime(int burstCounter, int currentBetweenTime, int maxBetweenTime,
+            int initialBetweenTime) {
+        return currentBetweenTime < maxBetweenTime
+                ? Math.min(initialBetweenTime * (int) Math.pow(2, burstCounter), maxBetweenTime)
+                : currentBetweenTime;
+    }
+
+    @Test
+    public void sendQueries_aggressiveScanMode() {
+        final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(AGGRESSIVE_QUERY_MODE).build();
+        startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        int burstCounter = 0;
+        int betweenBurstTime = 0;
+        for (int i = 0; i < expectedIPv4Packets.length; i += 3) {
+            verifyAndSendQuery(i, betweenBurstTime, /* expectsUnicastResponse= */ true);
+            verifyAndSendQuery(i + 1, /* timeInMs= */ 0, /* expectsUnicastResponse= */ false);
+            verifyAndSendQuery(i + 2, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
+                    /* expectsUnicastResponse= */ false);
+            betweenBurstTime = getBetweenBurstTime(burstCounter, betweenBurstTime,
+                    MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS,
+                    INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS);
+            burstCounter++;
+        }
+        // Verify that Task is not removed before stopSendAndReceive was called.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        // Stop sending packets.
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+    }
+
+    @Test
+    public void sendQueries_reentry_aggressiveScanMode() {
+        final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE).setQueryMode(AGGRESSIVE_QUERY_MODE).build();
+        startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        // First burst, first query is sent.
+        verifyAndSendQuery(0, /* timeInMs= */ 0, /* expectsUnicastResponse= */ true);
+
+        // After the first query is sent, change the subtypes, and restart.
+        final MdnsSearchOptions searchOptions2 = MdnsSearchOptions.newBuilder().addSubtype(SUBTYPE)
+                .addSubtype("_subtype2").setQueryMode(AGGRESSIVE_QUERY_MODE).build();
+        startSendAndReceive(mockListenerOne, searchOptions2);
+        // The previous scheduled task should be canceled.
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        // Queries should continue to be sent.
+        verifyAndSendQuery(1, /* timeInMs= */ 0, /* expectsUnicastResponse= */ true);
+        verifyAndSendQuery(2, /* timeInMs= */ 0, /* expectsUnicastResponse= */ false);
+        verifyAndSendQuery(3, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
+                /* expectsUnicastResponse= */ false);
+
+        // Stop sending packets.
+        stopSendAndReceive(mockListenerOne);
+        verify(mockDeps, times(3)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+    }
+
+    @Test
+    public void sendQueries_blendScanWithQueryBackoff() {
+        final int numOfQueriesBeforeBackoff = 11;
+        final MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder()
+                .addSubtype(SUBTYPE)
+                .setQueryMode(AGGRESSIVE_QUERY_MODE)
+                .setNumOfQueriesBeforeBackoff(numOfQueriesBeforeBackoff)
+                .build();
+        startSendAndReceive(mockListenerOne, searchOptions);
+        // Always try to remove the task.
+        verify(mockDeps, times(1)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+
+        int burstCounter = 0;
+        int betweenBurstTime = 0;
+        for (int i = 0; i < numOfQueriesBeforeBackoff; i += 3) {
+            verifyAndSendQuery(i, betweenBurstTime, /* expectsUnicastResponse= */ true);
+            verifyAndSendQuery(i + 1, /* timeInMs= */ 0, /* expectsUnicastResponse= */ false);
+            verifyAndSendQuery(i + 2, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
+                    /* expectsUnicastResponse= */ false);
+            betweenBurstTime = getBetweenBurstTime(burstCounter, betweenBurstTime,
+                    MAX_TIME_BETWEEN_AGGRESSIVE_BURSTS_MS,
+                    INITIAL_AGGRESSIVE_TIME_BETWEEN_BURSTS_MS);
+            burstCounter++;
+        }
+        // In backoff mode, the current scheduled task will be canceled and reschedule if the
+        // 0.8 * smallestRemainingTtl is larger than time to next run.
+        long currentTime = TEST_TTL / 2 + TEST_ELAPSED_REALTIME;
+        doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
+        doReturn(true).when(mockDeps).hasMessages(any(), eq(EVENT_START_QUERYTASK));
+        processResponse(createResponse(
+                "service-instance-1", "192.0.2.123", 5353,
+                SERVICE_TYPE_LABELS,
+                Collections.emptyMap(), TEST_TTL), socketKey);
+        verify(mockDeps, times(2)).removeMessages(any(), eq(EVENT_START_QUERYTASK));
+        assertNotNull(delayMessage);
+        verifyAndSendQuery(12 /* index */, (long) (TEST_TTL / 2 * 0.8) /* timeInMs */,
+                true /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                14 /* scheduledCount */);
+        currentTime += (long) (TEST_TTL / 2 * 0.8);
+        doReturn(currentTime).when(mockDecoderClock).elapsedRealtime();
+        verifyAndSendQuery(13 /* index */, 0 /* timeInMs */,
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                15 /* scheduledCount */);
+        verifyAndSendQuery(14 /* index */, TIME_BETWEEN_RETRANSMISSION_QUERIES_IN_BURST_MS,
+                false /* expectsUnicastResponse */, true /* multipleSocketDiscovery */,
+                16 /* scheduledCount */);
+    }
+
     private static MdnsServiceInfo matchServiceName(String name) {
         return argThat(info -> info.getServiceInstanceName().equals(name));
     }
diff --git a/tests/unit/java/com/android/server/net/BpfInterfaceMapHelperTest.java b/tests/unit/java/com/android/server/net/BpfInterfaceMapHelperTest.java
new file mode 100644
index 0000000..7b3bea3
--- /dev/null
+++ b/tests/unit/java/com/android/server/net/BpfInterfaceMapHelperTest.java
@@ -0,0 +1,113 @@
+/*
+ * Copyright (C) 2022 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;
+
+import static android.system.OsConstants.EPERM;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
+
+import android.os.Build;
+import android.system.ErrnoException;
+import android.util.IndentingPrintWriter;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.net.module.util.BaseNetdUnsolicitedEventListener;
+import com.android.net.module.util.IBpfMap;
+import com.android.net.module.util.Struct.S32;
+import com.android.testutils.DevSdkIgnoreRule;
+import com.android.testutils.DevSdkIgnoreRunner;
+import com.android.testutils.TestBpfMap;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+
+import java.io.PrintWriter;
+import java.io.StringWriter;
+
+@SmallTest
+@RunWith(DevSdkIgnoreRunner.class)
+@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
+public final class BpfInterfaceMapHelperTest {
+    private static final int TEST_INDEX = 1;
+    private static final int TEST_INDEX2 = 2;
+    private static final String TEST_INTERFACE_NAME = "test1";
+    private static final String TEST_INTERFACE_NAME2 = "test2";
+
+    private BaseNetdUnsolicitedEventListener mListener;
+    private BpfInterfaceMapHelper mUpdater;
+    private IBpfMap<S32, InterfaceMapValue> mBpfMap =
+            spy(new TestBpfMap<>(S32.class, InterfaceMapValue.class));
+
+    private class TestDependencies extends BpfInterfaceMapHelper.Dependencies {
+        @Override
+        public IBpfMap<S32, InterfaceMapValue> getInterfaceMap() {
+            return mBpfMap;
+        }
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        mUpdater = new BpfInterfaceMapHelper(new TestDependencies());
+    }
+
+    @Test
+    public void testGetIfNameByIndex() throws Exception {
+        mBpfMap.updateEntry(new S32(TEST_INDEX), new InterfaceMapValue(TEST_INTERFACE_NAME));
+        assertEquals(TEST_INTERFACE_NAME, mUpdater.getIfNameByIndex(TEST_INDEX));
+    }
+
+    @Test
+    public void testGetIfNameByIndexNoEntry() {
+        assertNull(mUpdater.getIfNameByIndex(TEST_INDEX));
+    }
+
+    @Test
+    public void testGetIfNameByIndexException() throws Exception {
+        doThrow(new ErrnoException("", EPERM)).when(mBpfMap).getValue(new S32(TEST_INDEX));
+        assertNull(mUpdater.getIfNameByIndex(TEST_INDEX));
+    }
+
+    private void assertDumpContains(final String dump, final String message) {
+        assertTrue(String.format("dump(%s) does not contain '%s'", dump, message),
+                dump.contains(message));
+    }
+
+    private String getDump() {
+        final StringWriter sw = new StringWriter();
+        mUpdater.dump(new IndentingPrintWriter(new PrintWriter(sw), " "));
+        return sw.toString();
+    }
+
+    @Test
+    public void testDump() throws ErrnoException {
+        mBpfMap.updateEntry(new S32(TEST_INDEX), new InterfaceMapValue(TEST_INTERFACE_NAME));
+        mBpfMap.updateEntry(new S32(TEST_INDEX2), new InterfaceMapValue(TEST_INTERFACE_NAME2));
+
+        final String dump = getDump();
+        assertDumpContains(dump, "IfaceIndexNameMap: OK");
+        assertDumpContains(dump, "ifaceIndex=1 ifaceName=test1");
+        assertDumpContains(dump, "ifaceIndex=2 ifaceName=test2");
+    }
+}
diff --git a/tests/unit/java/com/android/server/net/BpfInterfaceMapUpdaterTest.java b/tests/unit/java/com/android/server/net/BpfInterfaceMapUpdaterTest.java
deleted file mode 100644
index c730856..0000000
--- a/tests/unit/java/com/android/server/net/BpfInterfaceMapUpdaterTest.java
+++ /dev/null
@@ -1,174 +0,0 @@
-/*
- * Copyright (C) 2022 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;
-
-import static android.system.OsConstants.EPERM;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.doThrow;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.when;
-
-import android.content.Context;
-import android.net.INetd;
-import android.net.MacAddress;
-import android.os.Build;
-import android.os.Handler;
-import android.os.test.TestLooper;
-import android.system.ErrnoException;
-import android.util.IndentingPrintWriter;
-
-import androidx.test.filters.SmallTest;
-
-import com.android.net.module.util.BaseNetdUnsolicitedEventListener;
-import com.android.net.module.util.IBpfMap;
-import com.android.net.module.util.InterfaceParams;
-import com.android.net.module.util.Struct.S32;
-import com.android.testutils.DevSdkIgnoreRule;
-import com.android.testutils.DevSdkIgnoreRunner;
-import com.android.testutils.TestBpfMap;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
-import java.io.PrintWriter;
-import java.io.StringWriter;
-
-@SmallTest
-@RunWith(DevSdkIgnoreRunner.class)
-@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
-public final class BpfInterfaceMapUpdaterTest {
-    private static final int TEST_INDEX = 1;
-    private static final int TEST_INDEX2 = 2;
-    private static final String TEST_INTERFACE_NAME = "test1";
-    private static final String TEST_INTERFACE_NAME2 = "test2";
-
-    private final TestLooper mLooper = new TestLooper();
-    private BaseNetdUnsolicitedEventListener mListener;
-    private BpfInterfaceMapUpdater mUpdater;
-    private IBpfMap<S32, InterfaceMapValue> mBpfMap =
-            spy(new TestBpfMap<>(S32.class, InterfaceMapValue.class));
-    @Mock private INetd mNetd;
-    @Mock private Context mContext;
-
-    private class TestDependencies extends BpfInterfaceMapUpdater.Dependencies {
-        @Override
-        public IBpfMap<S32, InterfaceMapValue> getInterfaceMap() {
-            return mBpfMap;
-        }
-
-        @Override
-        public InterfaceParams getInterfaceParams(String ifaceName) {
-            if (ifaceName.equals(TEST_INTERFACE_NAME)) {
-                return new InterfaceParams(TEST_INTERFACE_NAME, TEST_INDEX,
-                        MacAddress.ALL_ZEROS_ADDRESS);
-            } else if (ifaceName.equals(TEST_INTERFACE_NAME2)) {
-                return new InterfaceParams(TEST_INTERFACE_NAME2, TEST_INDEX2,
-                        MacAddress.ALL_ZEROS_ADDRESS);
-            }
-
-            return null;
-        }
-
-        @Override
-        public INetd getINetd(Context ctx) {
-            return mNetd;
-        }
-    }
-
-    @Before
-    public void setUp() throws Exception {
-        MockitoAnnotations.initMocks(this);
-        when(mNetd.interfaceGetList()).thenReturn(new String[] {TEST_INTERFACE_NAME});
-        mUpdater = new BpfInterfaceMapUpdater(mContext, new Handler(mLooper.getLooper()),
-                new TestDependencies());
-    }
-
-    private void verifyStartUpdater() throws Exception {
-        mUpdater.start();
-        mLooper.dispatchAll();
-        final ArgumentCaptor<BaseNetdUnsolicitedEventListener> listenerCaptor =
-                ArgumentCaptor.forClass(BaseNetdUnsolicitedEventListener.class);
-        verify(mNetd).registerUnsolicitedEventListener(listenerCaptor.capture());
-        mListener = listenerCaptor.getValue();
-        verify(mBpfMap).updateEntry(eq(new S32(TEST_INDEX)),
-                eq(new InterfaceMapValue(TEST_INTERFACE_NAME)));
-    }
-
-    @Test
-    public void testUpdateInterfaceMap() throws Exception {
-        verifyStartUpdater();
-
-        mListener.onInterfaceAdded(TEST_INTERFACE_NAME2);
-        mLooper.dispatchAll();
-        verify(mBpfMap).updateEntry(eq(new S32(TEST_INDEX2)),
-                eq(new InterfaceMapValue(TEST_INTERFACE_NAME2)));
-
-        // Check that when onInterfaceRemoved is called, nothing happens.
-        mListener.onInterfaceRemoved(TEST_INTERFACE_NAME);
-        mLooper.dispatchAll();
-        verifyNoMoreInteractions(mBpfMap);
-    }
-
-    @Test
-    public void testGetIfNameByIndex() throws Exception {
-        mBpfMap.updateEntry(new S32(TEST_INDEX), new InterfaceMapValue(TEST_INTERFACE_NAME));
-        assertEquals(TEST_INTERFACE_NAME, mUpdater.getIfNameByIndex(TEST_INDEX));
-    }
-
-    @Test
-    public void testGetIfNameByIndexNoEntry() {
-        assertNull(mUpdater.getIfNameByIndex(TEST_INDEX));
-    }
-
-    @Test
-    public void testGetIfNameByIndexException() throws Exception {
-        doThrow(new ErrnoException("", EPERM)).when(mBpfMap).getValue(new S32(TEST_INDEX));
-        assertNull(mUpdater.getIfNameByIndex(TEST_INDEX));
-    }
-
-    private void assertDumpContains(final String dump, final String message) {
-        assertTrue(String.format("dump(%s) does not contain '%s'", dump, message),
-                dump.contains(message));
-    }
-
-    private String getDump() {
-        final StringWriter sw = new StringWriter();
-        mUpdater.dump(new IndentingPrintWriter(new PrintWriter(sw), " "));
-        return sw.toString();
-    }
-
-    @Test
-    public void testDump() throws ErrnoException {
-        mBpfMap.updateEntry(new S32(TEST_INDEX), new InterfaceMapValue(TEST_INTERFACE_NAME));
-        mBpfMap.updateEntry(new S32(TEST_INDEX2), new InterfaceMapValue(TEST_INTERFACE_NAME2));
-
-        final String dump = getDump();
-        assertDumpContains(dump, "IfaceIndexNameMap: OK");
-        assertDumpContains(dump, "ifaceIndex=1 ifaceName=test1");
-        assertDumpContains(dump, "ifaceIndex=2 ifaceName=test2");
-    }
-}
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index a5fee5b..3ed51bc 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -254,7 +254,7 @@
     private @Mock AlarmManager mAlarmManager;
     @Mock
     private NetworkStatsSubscriptionsMonitor mNetworkStatsSubscriptionsMonitor;
-    private @Mock BpfInterfaceMapUpdater mBpfInterfaceMapUpdater;
+    private @Mock BpfInterfaceMapHelper mBpfInterfaceMapHelper;
     private HandlerThread mHandlerThread;
     @Mock
     private LocationPermissionChecker mLocationPermissionChecker;
@@ -519,9 +519,8 @@
         }
 
         @Override
-        public BpfInterfaceMapUpdater makeBpfInterfaceMapUpdater(
-                @NonNull Context ctx, @NonNull Handler handler) {
-            return mBpfInterfaceMapUpdater;
+        public BpfInterfaceMapHelper makeBpfInterfaceMapHelper() {
+            return mBpfInterfaceMapHelper;
         }
 
         @Override
@@ -2764,13 +2763,13 @@
 
     @Test
     public void testDumpStatsMap() throws ErrnoException {
-        doReturn("wlan0").when(mBpfInterfaceMapUpdater).getIfNameByIndex(10 /* index */);
+        doReturn("wlan0").when(mBpfInterfaceMapHelper).getIfNameByIndex(10 /* index */);
         doTestDumpStatsMap("wlan0");
     }
 
     @Test
     public void testDumpStatsMapUnknownInterface() throws ErrnoException {
-        doReturn(null).when(mBpfInterfaceMapUpdater).getIfNameByIndex(10 /* index */);
+        doReturn(null).when(mBpfInterfaceMapHelper).getIfNameByIndex(10 /* index */);
         doTestDumpStatsMap("unknown");
     }
 
@@ -2785,13 +2784,13 @@
 
     @Test
     public void testDumpIfaceStatsMap() throws Exception {
-        doReturn("wlan0").when(mBpfInterfaceMapUpdater).getIfNameByIndex(10 /* index */);
+        doReturn("wlan0").when(mBpfInterfaceMapHelper).getIfNameByIndex(10 /* index */);
         doTestDumpIfaceStatsMap("wlan0");
     }
 
     @Test
     public void testDumpIfaceStatsMapUnknownInterface() throws Exception {
-        doReturn(null).when(mBpfInterfaceMapUpdater).getIfNameByIndex(10 /* index */);
+        doReturn(null).when(mBpfInterfaceMapHelper).getIfNameByIndex(10 /* index */);
         doTestDumpIfaceStatsMap("unknown");
     }
 
diff --git a/thread/framework/java/android/net/thread/IStateCallback.aidl b/thread/framework/java/android/net/thread/IStateCallback.aidl
index d7cbda9..9d0a571 100644
--- a/thread/framework/java/android/net/thread/IStateCallback.aidl
+++ b/thread/framework/java/android/net/thread/IStateCallback.aidl
@@ -22,4 +22,5 @@
 oneway interface IStateCallback {
     void onDeviceRoleChanged(int deviceRole);
     void onPartitionIdChanged(long partitionId);
+    void onThreadEnableStateChanged(int enabledState);
 }
diff --git a/thread/framework/java/android/net/thread/IThreadNetworkController.aidl b/thread/framework/java/android/net/thread/IThreadNetworkController.aidl
index a9da8d6..485e25d 100644
--- a/thread/framework/java/android/net/thread/IThreadNetworkController.aidl
+++ b/thread/framework/java/android/net/thread/IThreadNetworkController.aidl
@@ -42,4 +42,6 @@
 
     int getThreadVersion();
     void createRandomizedDataset(String networkName, IActiveOperationalDatasetReceiver receiver);
+
+    void setEnabled(boolean enabled, in IOperationReceiver receiver);
 }
diff --git a/thread/framework/java/android/net/thread/ThreadNetworkController.java b/thread/framework/java/android/net/thread/ThreadNetworkController.java
index 7242ed7..db761a3 100644
--- a/thread/framework/java/android/net/thread/ThreadNetworkController.java
+++ b/thread/framework/java/android/net/thread/ThreadNetworkController.java
@@ -68,6 +68,15 @@
     /** The device is a Thread Leader. */
     public static final int DEVICE_ROLE_LEADER = 4;
 
+    /** The Thread radio is disabled. */
+    public static final int STATE_DISABLED = 0;
+
+    /** The Thread radio is enabled. */
+    public static final int STATE_ENABLED = 1;
+
+    /** The Thread radio is being disabled. */
+    public static final int STATE_DISABLING = 2;
+
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
     @IntDef({
@@ -79,6 +88,13 @@
     })
     public @interface DeviceRole {}
 
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(
+            prefix = {"STATE_"},
+            value = {STATE_DISABLED, STATE_ENABLED, STATE_DISABLING})
+    public @interface EnabledState {}
+
     /** Thread standard version 1.3. */
     public static final int THREAD_VERSION_1_3 = 4;
 
@@ -106,6 +122,40 @@
         mControllerService = controllerService;
     }
 
+    /**
+     * Enables/Disables the radio of this ThreadNetworkController. The requested enabled state will
+     * be persistent and survives device reboots.
+     *
+     * <p>When Thread is in {@code STATE_DISABLED}, {@link ThreadNetworkController} APIs which
+     * require the Thread radio will fail with error code {@link
+     * ThreadNetworkException#ERROR_THREAD_DISABLED}. When Thread is in {@code STATE_DISABLING},
+     * {@link ThreadNetworkController} APIs that return a {@link ThreadNetworkException} will fail
+     * with error code {@link ThreadNetworkException#ERROR_BUSY}.
+     *
+     * <p>On success, {@link OutcomeReceiver#onResult} of {@code receiver} is called. It indicates
+     * the operation has completed. But there maybe subsequent calls to update the enabled state,
+     * callers of this method should use {@link #registerStateCallback} to subscribe to the Thread
+     * enabled state changes.
+     *
+     * <p>On failure, {@link OutcomeReceiver#onError} of {@code receiver} will be invoked with a
+     * specific error in {@link ThreadNetworkException#ERROR_}.
+     *
+     * @param enabled {@code true} for enabling Thread
+     * @param executor the executor to execute {@code receiver}
+     * @param receiver the receiver to receive result of this operation
+     */
+    @RequiresPermission("android.permission.THREAD_NETWORK_PRIVILEGED")
+    public void setEnabled(
+            boolean enabled,
+            @NonNull @CallbackExecutor Executor executor,
+            @NonNull OutcomeReceiver<Void, ThreadNetworkException> receiver) {
+        try {
+            mControllerService.setEnabled(enabled, new OperationReceiverProxy(executor, receiver));
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
     /** Returns the Thread version this device is operating on. */
     @ThreadVersion
     public int getThreadVersion() {
@@ -170,6 +220,16 @@
          * @param partitionId the new Thread partition ID
          */
         default void onPartitionIdChanged(long partitionId) {}
+
+        /**
+         * The Thread enabled state has changed.
+         *
+         * <p>The Thread enabled state can be set with {@link setEnabled}, it may also be updated by
+         * airplane mode or admin control.
+         *
+         * @param enabledState the new Thread enabled state
+         */
+        default void onThreadEnableStateChanged(@EnabledState int enabledState) {}
     }
 
     private static final class StateCallbackProxy extends IStateCallback.Stub {
@@ -200,6 +260,16 @@
                 Binder.restoreCallingIdentity(identity);
             }
         }
+
+        @Override
+        public void onThreadEnableStateChanged(@EnabledState int enabled) {
+            final long identity = Binder.clearCallingIdentity();
+            try {
+                mExecutor.execute(() -> mCallback.onThreadEnableStateChanged(enabled));
+            } finally {
+                Binder.restoreCallingIdentity(identity);
+            }
+        }
     }
 
     /**
diff --git a/thread/framework/java/android/net/thread/ThreadNetworkException.java b/thread/framework/java/android/net/thread/ThreadNetworkException.java
index 4fd445b..23ed53e 100644
--- a/thread/framework/java/android/net/thread/ThreadNetworkException.java
+++ b/thread/framework/java/android/net/thread/ThreadNetworkException.java
@@ -48,6 +48,7 @@
         ERROR_RESPONSE_BAD_FORMAT,
         ERROR_RESOURCE_EXHAUSTED,
         ERROR_UNKNOWN,
+        ERROR_THREAD_DISABLED,
     })
     public @interface ErrorCode {}
 
@@ -124,11 +125,18 @@
     public static final int ERROR_RESOURCE_EXHAUSTED = 10;
 
     /**
-     * The operation failed because of an unknown error in the system. This typically indicates
-     * that the caller doesn't understand error codes added in newer Android versions.
+     * The operation failed because of an unknown error in the system. This typically indicates that
+     * the caller doesn't understand error codes added in newer Android versions.
      */
     public static final int ERROR_UNKNOWN = 11;
 
+    /**
+     * The operation failed because the Thread radio is disabled by {@link
+     * ThreadNetworkController#setEnabled}, airplane mode or device admin. The caller should retry
+     * only after Thread is enabled.
+     */
+    public static final int ERROR_THREAD_DISABLED = 12;
+
     private final int mErrorCode;
 
     /** Creates a new {@link ThreadNetworkException} object with given error code and message. */
diff --git a/thread/scripts/make-pretty.sh b/thread/scripts/make-pretty.sh
index e4bd459..c176bfa 100755
--- a/thread/scripts/make-pretty.sh
+++ b/thread/scripts/make-pretty.sh
@@ -3,5 +3,7 @@
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 
 GOOGLE_JAVA_FORMAT=$SCRIPT_DIR/../../../../../prebuilts/tools/common/google-java-format/google-java-format
+ANDROID_BP_FORMAT=$SCRIPT_DIR/../../../../../prebuilts/build-tools/linux-x86/bin/bpfmt
 
 $GOOGLE_JAVA_FORMAT --aosp -i $(find $SCRIPT_DIR/../ -name "*.java")
+$ANDROID_BP_FORMAT -w $(find $SCRIPT_DIR/../ -name "*.bp")
diff --git a/thread/service/Android.bp b/thread/service/Android.bp
index 69295cc..0132235 100644
--- a/thread/service/Android.bp
+++ b/thread/service/Android.bp
@@ -47,9 +47,6 @@
         "ot-daemon-aidl-java",
     ],
     apex_available: ["com.android.tethering"],
-    optimize: {
-        proguard_flags_files: ["proguard.flags"],
-    },
 }
 
 cc_library_shared {
diff --git a/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java b/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java
index 1c51c42..7b9f290 100644
--- a/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java
+++ b/thread/service/java/com/android/server/thread/ThreadNetworkControllerService.java
@@ -27,6 +27,9 @@
 import static android.net.thread.ActiveOperationalDataset.MESH_LOCAL_PREFIX_FIRST_BYTE;
 import static android.net.thread.ActiveOperationalDataset.SecurityPolicy.DEFAULT_ROTATION_TIME_HOURS;
 import static android.net.thread.ThreadNetworkController.DEVICE_ROLE_DETACHED;
+import static android.net.thread.ThreadNetworkController.STATE_DISABLED;
+import static android.net.thread.ThreadNetworkController.STATE_DISABLING;
+import static android.net.thread.ThreadNetworkController.STATE_ENABLED;
 import static android.net.thread.ThreadNetworkController.THREAD_VERSION_1_3;
 import static android.net.thread.ThreadNetworkException.ERROR_ABORTED;
 import static android.net.thread.ThreadNetworkException.ERROR_BUSY;
@@ -35,6 +38,7 @@
 import static android.net.thread.ThreadNetworkException.ERROR_REJECTED_BY_PEER;
 import static android.net.thread.ThreadNetworkException.ERROR_RESOURCE_EXHAUSTED;
 import static android.net.thread.ThreadNetworkException.ERROR_RESPONSE_BAD_FORMAT;
+import static android.net.thread.ThreadNetworkException.ERROR_THREAD_DISABLED;
 import static android.net.thread.ThreadNetworkException.ERROR_TIMEOUT;
 import static android.net.thread.ThreadNetworkException.ERROR_UNSUPPORTED_CHANNEL;
 import static android.net.thread.ThreadNetworkManager.PERMISSION_THREAD_NETWORK_PRIVILEGED;
@@ -48,7 +52,11 @@
 import static com.android.server.thread.openthread.IOtDaemon.ErrorCode.OT_ERROR_REASSEMBLY_TIMEOUT;
 import static com.android.server.thread.openthread.IOtDaemon.ErrorCode.OT_ERROR_REJECTED;
 import static com.android.server.thread.openthread.IOtDaemon.ErrorCode.OT_ERROR_RESPONSE_TIMEOUT;
+import static com.android.server.thread.openthread.IOtDaemon.ErrorCode.OT_ERROR_THREAD_DISABLED;
 import static com.android.server.thread.openthread.IOtDaemon.ErrorCode.OT_ERROR_UNSUPPORTED_CHANNEL;
+import static com.android.server.thread.openthread.IOtDaemon.OT_STATE_DISABLED;
+import static com.android.server.thread.openthread.IOtDaemon.OT_STATE_DISABLING;
+import static com.android.server.thread.openthread.IOtDaemon.OT_STATE_ENABLED;
 import static com.android.server.thread.openthread.IOtDaemon.TUN_IF_NAME;
 
 import android.Manifest.permission;
@@ -160,6 +168,7 @@
     private UpstreamNetworkCallback mUpstreamNetworkCallback;
     private TestNetworkSpecifier mUpstreamTestNetworkSpecifier;
     private final HashMap<Network, String> mNetworkToInterface;
+    private final ThreadPersistentSettings mPersistentSettings;
 
     private BorderRouterConfigurationParcel mBorderRouterConfig;
 
@@ -171,7 +180,8 @@
             Supplier<IOtDaemon> otDaemonSupplier,
             ConnectivityManager connectivityManager,
             TunInterfaceController tunIfController,
-            InfraInterfaceController infraIfController) {
+            InfraInterfaceController infraIfController,
+            ThreadPersistentSettings persistentSettings) {
         mContext = context;
         mHandler = handler;
         mNetworkProvider = networkProvider;
@@ -182,9 +192,11 @@
         mUpstreamNetworkRequest = newUpstreamNetworkRequest();
         mNetworkToInterface = new HashMap<Network, String>();
         mBorderRouterConfig = new BorderRouterConfigurationParcel();
+        mPersistentSettings = persistentSettings;
     }
 
-    public static ThreadNetworkControllerService newInstance(Context context) {
+    public static ThreadNetworkControllerService newInstance(
+            Context context, ThreadPersistentSettings persistentSettings) {
         HandlerThread handlerThread = new HandlerThread("ThreadHandlerThread");
         handlerThread.start();
         NetworkProvider networkProvider =
@@ -197,7 +209,8 @@
                 () -> IOtDaemon.Stub.asInterface(ServiceManagerWrapper.waitForService("ot_daemon")),
                 context.getSystemService(ConnectivityManager.class),
                 new TunInterfaceController(TUN_IF_NAME),
-                new InfraInterfaceController());
+                new InfraInterfaceController(),
+                persistentSettings);
     }
 
     private static Inet6Address bytesToInet6Address(byte[] addressBytes) {
@@ -273,7 +286,9 @@
         if (otDaemon == null) {
             throw new RemoteException("Internal error: failed to start OT daemon");
         }
-        otDaemon.initialize(mTunIfController.getTunFd());
+        otDaemon.initialize(
+                mTunIfController.getTunFd(),
+                mPersistentSettings.get(ThreadPersistentSettings.THREAD_ENABLED));
         otDaemon.registerStateCallback(mOtDaemonCallbackProxy, -1);
         otDaemon.asBinder().linkToDeath(() -> mHandler.post(this::onOtDaemonDied), 0);
         mOtDaemon = otDaemon;
@@ -308,6 +323,26 @@
                 });
     }
 
+    public void setEnabled(@NonNull boolean isEnabled, @NonNull IOperationReceiver receiver) {
+        enforceAllPermissionsGranted(PERMISSION_THREAD_NETWORK_PRIVILEGED);
+
+        mHandler.post(() -> setEnabledInternal(isEnabled, new OperationReceiverWrapper(receiver)));
+    }
+
+    private void setEnabledInternal(
+            @NonNull boolean isEnabled, @Nullable OperationReceiverWrapper receiver) {
+        // 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
+        // to the desired value after reboot.
+        mPersistentSettings.put(ThreadPersistentSettings.THREAD_ENABLED.key, isEnabled);
+        try {
+            getOtDaemon().setThreadEnabled(isEnabled, newOtStatusReceiver(receiver));
+        } catch (RemoteException e) {
+            Log.e(TAG, "otDaemon.setThreadEnabled failed", e);
+            receiver.onError(ERROR_INTERNAL_ERROR, "Thread stack error");
+        }
+    }
+
     private void requestUpstreamNetwork() {
         if (mUpstreamNetworkCallback != null) {
             throw new AssertionError("The upstream network request is already there.");
@@ -658,6 +693,8 @@
                 return ERROR_REJECTED_BY_PEER;
             case OT_ERROR_UNSUPPORTED_CHANNEL:
                 return ERROR_UNSUPPORTED_CHANNEL;
+            case OT_ERROR_THREAD_DISABLED:
+                return ERROR_THREAD_DISABLED;
             default:
                 return ERROR_INTERNAL_ERROR;
         }
@@ -1001,6 +1038,15 @@
             }
         }
 
+        private void notifyThreadEnabledUpdated(IStateCallback callback, int enabledState) {
+            try {
+                callback.onThreadEnableStateChanged(enabledState);
+                Log.i(TAG, "onThreadEnableStateChanged " + enabledState);
+            } catch (RemoteException ignored) {
+                // do nothing if the client is dead
+            }
+        }
+
         public void unregisterStateCallback(IStateCallback callback) {
             checkOnHandlerThread();
             if (!mStateCallbacks.containsKey(callback)) {
@@ -1065,6 +1111,31 @@
         }
 
         @Override
+        public void onThreadEnabledChanged(int state) {
+            mHandler.post(() -> onThreadEnabledChangedInternal(state));
+        }
+
+        private void onThreadEnabledChangedInternal(int state) {
+            checkOnHandlerThread();
+            for (IStateCallback callback : mStateCallbacks.keySet()) {
+                notifyThreadEnabledUpdated(callback, otStateToAndroidState(state));
+            }
+        }
+
+        private static int otStateToAndroidState(int state) {
+            switch (state) {
+                case OT_STATE_ENABLED:
+                    return STATE_ENABLED;
+                case OT_STATE_DISABLED:
+                    return STATE_DISABLED;
+                case OT_STATE_DISABLING:
+                    return STATE_DISABLING;
+                default:
+                    throw new IllegalArgumentException("Unknown ot state " + state);
+            }
+        }
+
+        @Override
         public void onStateChanged(OtDaemonState newState, long listenerId) {
             mHandler.post(() -> onStateChangedInternal(newState, listenerId));
         }
diff --git a/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java b/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java
index b7b6233..ffa7b44 100644
--- a/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java
+++ b/thread/service/java/com/android/server/thread/ThreadNetworkCountryCode.java
@@ -31,6 +31,7 @@
 import android.net.wifi.WifiManager;
 import android.net.wifi.WifiManager.ActiveCountryCodeChangedCallback;
 import android.os.Build;
+import android.sysprop.ThreadNetworkProperties;
 import android.telephony.SubscriptionInfo;
 import android.telephony.SubscriptionManager;
 import android.telephony.TelephonyManager;
@@ -53,8 +54,8 @@
 
 /**
  * Provide functions for making changes to Thread Network country code. This Country Code is from
- * location, WiFi or telephony configuration. This class sends Country Code to Thread Network native
- * layer.
+ * location, WiFi, telephony or OEM configuration. This class sends Country Code to Thread Network
+ * native layer.
  *
  * <p>This class is thread-safe.
  */
@@ -77,6 +78,7 @@
             value = {
                 COUNTRY_CODE_SOURCE_DEFAULT,
                 COUNTRY_CODE_SOURCE_LOCATION,
+                COUNTRY_CODE_SOURCE_OEM,
                 COUNTRY_CODE_SOURCE_OVERRIDE,
                 COUNTRY_CODE_SOURCE_TELEPHONY,
                 COUNTRY_CODE_SOURCE_TELEPHONY_LAST,
@@ -86,6 +88,7 @@
 
     private static final String COUNTRY_CODE_SOURCE_DEFAULT = "Default";
     private static final String COUNTRY_CODE_SOURCE_LOCATION = "Location";
+    private static final String COUNTRY_CODE_SOURCE_OEM = "Oem";
     private static final String COUNTRY_CODE_SOURCE_OVERRIDE = "Override";
     private static final String COUNTRY_CODE_SOURCE_TELEPHONY = "Telephony";
     private static final String COUNTRY_CODE_SOURCE_TELEPHONY_LAST = "TelephonyLast";
@@ -111,6 +114,7 @@
     @Nullable private CountryCodeInfo mWifiCountryCodeInfo;
     @Nullable private CountryCodeInfo mTelephonyCountryCodeInfo;
     @Nullable private CountryCodeInfo mTelephonyLastCountryCodeInfo;
+    @Nullable private CountryCodeInfo mOemCountryCodeInfo;
 
     /** Container class to store Thread country code information. */
     private static final class CountryCodeInfo {
@@ -118,13 +122,35 @@
         @CountryCodeSource private String mSource;
         private final Instant mUpdatedTimestamp;
 
+        /**
+         * Constructs a new {@code CountryCodeInfo} from the given country code, country code source
+         * and country coode created time.
+         *
+         * @param countryCode a String representation of the country code as defined in ISO 3166.
+         * @param countryCodeSource a String representation of country code source.
+         * @param instant a Instant representation of the time when the country code was created.
+         * @throws IllegalArgumentException if {@code countryCode} contains invalid country code.
+         */
         public CountryCodeInfo(
                 String countryCode, @CountryCodeSource String countryCodeSource, Instant instant) {
+            if (!isValidCountryCode(countryCode)) {
+                throw new IllegalArgumentException("Country code is invalid: " + countryCode);
+            }
+
             mCountryCode = countryCode;
             mSource = countryCodeSource;
             mUpdatedTimestamp = instant;
         }
 
+        /**
+         * Constructs a new {@code CountryCodeInfo} from the given country code, country code
+         * source. The updated timestamp of the country code will be set to the time when {@code
+         * CountryCodeInfo} was constructed.
+         *
+         * @param countryCode a String representation of the country code as defined in ISO 3166.
+         * @param countryCodeSource a String representation of country code source.
+         * @throws IllegalArgumentException if {@code countryCode} contains invalid country code.
+         */
         public CountryCodeInfo(String countryCode, @CountryCodeSource String countryCodeSource) {
             this(countryCode, countryCodeSource, Instant.now());
         }
@@ -188,7 +214,8 @@
             WifiManager wifiManager,
             Context context,
             TelephonyManager telephonyManager,
-            SubscriptionManager subscriptionManager) {
+            SubscriptionManager subscriptionManager,
+            @Nullable String oemCountryCode) {
         mLocationManager = locationManager;
         mThreadNetworkControllerService = threadNetworkControllerService;
         mGeocoder = geocoder;
@@ -197,6 +224,10 @@
         mContext = context;
         mTelephonyManager = telephonyManager;
         mSubscriptionManager = subscriptionManager;
+
+        if (oemCountryCode != null) {
+            mOemCountryCodeInfo = new CountryCodeInfo(oemCountryCode, COUNTRY_CODE_SOURCE_OEM);
+        }
     }
 
     public static ThreadNetworkCountryCode newInstance(
@@ -209,7 +240,8 @@
                 context.getSystemService(WifiManager.class),
                 context,
                 context.getSystemService(TelephonyManager.class),
-                context.getSystemService(SubscriptionManager.class));
+                context.getSystemService(SubscriptionManager.class),
+                ThreadNetworkProperties.country_code().orElse(null));
     }
 
     /** Sets up this country code module to listen to location country code changes. */
@@ -278,7 +310,14 @@
         public void onActiveCountryCodeChanged(String countryCode) {
             Log.d(TAG, "Wifi country code is changed to " + countryCode);
             synchronized ("ThreadNetworkCountryCode.this") {
-                mWifiCountryCodeInfo = new CountryCodeInfo(countryCode, COUNTRY_CODE_SOURCE_WIFI);
+                if (isValidCountryCode(countryCode)) {
+                    mWifiCountryCodeInfo =
+                            new CountryCodeInfo(countryCode, COUNTRY_CODE_SOURCE_WIFI);
+                } else {
+                    Log.w(TAG, "WiFi country code " + countryCode + " is invalid");
+                    mWifiCountryCodeInfo = null;
+                }
+
                 updateCountryCode(false /* forceUpdate */);
             }
         }
@@ -418,6 +457,9 @@
      *       number will be used.
      *   <li>5. Location country code - Country code retrieved from LocationManager passive location
      *       provider.
+     *   <li>6. OEM country code - Country code retrieved from the system property
+     *       `threadnetwork.country_code`.
+     *   <li>7. Default country code `WW`.
      * </ul>
      *
      * @return the selected country code information.
@@ -443,6 +485,10 @@
             return mLocationCountryCodeInfo;
         }
 
+        if (mOemCountryCodeInfo != null) {
+            return mOemCountryCodeInfo;
+        }
+
         return DEFAULT_COUNTRY_CODE_INFO;
     }
 
@@ -531,13 +577,14 @@
     /** Dumps the current state of this ThreadNetworkCountryCode object. */
     public synchronized void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("---- Dump of ThreadNetworkCountryCode begin ----");
-        pw.println("mOverrideCountryCodeInfo: " + mOverrideCountryCodeInfo);
+        pw.println("mOverrideCountryCodeInfo        : " + mOverrideCountryCodeInfo);
         pw.println("mTelephonyCountryCodeSlotInfoMap: " + mTelephonyCountryCodeSlotInfoMap);
-        pw.println("mTelephonyCountryCodeInfo: " + mTelephonyCountryCodeInfo);
-        pw.println("mWifiCountryCodeInfo: " + mWifiCountryCodeInfo);
-        pw.println("mTelephonyLastCountryCodeInfo: " + mTelephonyLastCountryCodeInfo);
-        pw.println("mLocationCountryCodeInfo: " + mLocationCountryCodeInfo);
-        pw.println("mCurrentCountryCodeInfo: " + mCurrentCountryCodeInfo);
+        pw.println("mTelephonyCountryCodeInfo       : " + mTelephonyCountryCodeInfo);
+        pw.println("mWifiCountryCodeInfo            : " + mWifiCountryCodeInfo);
+        pw.println("mTelephonyLastCountryCodeInfo   : " + mTelephonyLastCountryCodeInfo);
+        pw.println("mLocationCountryCodeInfo        : " + mLocationCountryCodeInfo);
+        pw.println("mOemCountryCodeInfo             : " + mOemCountryCodeInfo);
+        pw.println("mCurrentCountryCodeInfo         : " + mCurrentCountryCodeInfo);
         pw.println("---- Dump of ThreadNetworkCountryCode end ------");
     }
 }
diff --git a/thread/service/java/com/android/server/thread/ThreadNetworkService.java b/thread/service/java/com/android/server/thread/ThreadNetworkService.java
index 53f2d4f..5cf27f7 100644
--- a/thread/service/java/com/android/server/thread/ThreadNetworkService.java
+++ b/thread/service/java/com/android/server/thread/ThreadNetworkService.java
@@ -18,16 +18,21 @@
 
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 
+import static com.android.net.module.util.DeviceConfigUtils.TETHERING_MODULE_NAME;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.content.ApexEnvironment;
 import android.content.Context;
 import android.net.thread.IThreadNetworkController;
 import android.net.thread.IThreadNetworkManager;
 import android.os.Binder;
 import android.os.ParcelFileDescriptor;
+import android.util.AtomicFile;
 
 import com.android.server.SystemService;
 
+import java.io.File;
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.Collections;
@@ -40,11 +45,18 @@
     private final Context mContext;
     @Nullable private ThreadNetworkCountryCode mCountryCode;
     @Nullable private ThreadNetworkControllerService mControllerService;
+    private final ThreadPersistentSettings mPersistentSettings;
     @Nullable private ThreadNetworkShellCommand mShellCommand;
 
     /** Creates a new {@link ThreadNetworkService} object. */
     public ThreadNetworkService(Context context) {
         mContext = context;
+        mPersistentSettings =
+                new ThreadPersistentSettings(
+                        new AtomicFile(
+                                new File(
+                                        getOrCreateThreadnetworkDir(),
+                                        ThreadPersistentSettings.FILE_NAME)));
     }
 
     /**
@@ -54,7 +66,9 @@
      */
     public void onBootPhase(int phase) {
         if (phase == SystemService.PHASE_SYSTEM_SERVICES_READY) {
-            mControllerService = ThreadNetworkControllerService.newInstance(mContext);
+            mPersistentSettings.initialize();
+            mControllerService =
+                    ThreadNetworkControllerService.newInstance(mContext, mPersistentSettings);
             mControllerService.initialize();
         } else if (phase == SystemService.PHASE_BOOT_COMPLETED) {
             // Country code initialization is delayed to the BOOT_COMPLETED phase because it will
@@ -109,4 +123,19 @@
 
         pw.println();
     }
+
+    /** Get device protected storage dir for the tethering apex. */
+    private static File getOrCreateThreadnetworkDir() {
+        final File threadnetworkDir;
+        final File apexDataDir =
+                ApexEnvironment.getApexEnvironment(TETHERING_MODULE_NAME)
+                        .getDeviceProtectedDataDir();
+        threadnetworkDir = new File(apexDataDir, "thread");
+
+        if (threadnetworkDir.exists() || threadnetworkDir.mkdirs()) {
+            return threadnetworkDir;
+        }
+        throw new IllegalStateException(
+                "Cannot write into thread network data directory: " + threadnetworkDir);
+    }
 }
diff --git a/thread/service/java/com/android/server/thread/TunInterfaceController.java b/thread/service/java/com/android/server/thread/TunInterfaceController.java
index 7223b2a..619a510 100644
--- a/thread/service/java/com/android/server/thread/TunInterfaceController.java
+++ b/thread/service/java/com/android/server/thread/TunInterfaceController.java
@@ -60,7 +60,8 @@
     public void createTunInterface() throws IOException {
         mParcelTunFd = ParcelFileDescriptor.adoptFd(nativeCreateTunInterface(mIfName, MTU));
         try {
-            mNetlinkSocket = NetlinkUtils.netlinkSocketForProto(OsConstants.NETLINK_ROUTE);
+            mNetlinkSocket =
+                    NetlinkUtils.netlinkSocketForProto(OsConstants.NETLINK_ROUTE, 0);
         } catch (ErrnoException e) {
             throw new IOException("Failed to create netlink socket", e);
         }
diff --git a/thread/service/proguard.flags b/thread/service/proguard.flags
deleted file mode 100644
index 5028982..0000000
--- a/thread/service/proguard.flags
+++ /dev/null
@@ -1,4 +0,0 @@
-# Ensure the callback methods are not stripped
--keepclassmembers class **.ThreadNetworkControllerService$ThreadNetworkCallback {
-    *;
-}
diff --git a/thread/tests/cts/Android.bp b/thread/tests/cts/Android.bp
index 3cf31e5..2f38bfd 100644
--- a/thread/tests/cts/Android.bp
+++ b/thread/tests/cts/Android.bp
@@ -45,7 +45,7 @@
     libs: [
         "android.test.base",
         "android.test.runner",
-        "framework-connectivity-module-api-stubs-including-flagged"
+        "framework-connectivity-module-api-stubs-including-flagged",
     ],
     // Test coverage system runs on different devices. Need to
     // compile for all architectures.
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 7a6c9aa..aab4b2e 100644
--- a/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
+++ b/thread/tests/cts/src/android/net/thread/cts/ThreadNetworkControllerTest.java
@@ -16,11 +16,19 @@
 
 package android.net.thread.cts;
 
+import static android.Manifest.permission.ACCESS_NETWORK_STATE;
+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;
 import static android.net.thread.ThreadNetworkController.DEVICE_ROLE_STOPPED;
+import static android.net.thread.ThreadNetworkController.STATE_DISABLED;
+import static android.net.thread.ThreadNetworkController.STATE_DISABLING;
+import static android.net.thread.ThreadNetworkController.STATE_ENABLED;
 import static android.net.thread.ThreadNetworkController.THREAD_VERSION_1_3;
 import static android.net.thread.ThreadNetworkException.ERROR_ABORTED;
 import static android.net.thread.ThreadNetworkException.ERROR_FAILED_PRECONDITION;
 import static android.net.thread.ThreadNetworkException.ERROR_REJECTED_BY_PEER;
+import static android.net.thread.ThreadNetworkException.ERROR_THREAD_DISABLED;
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
@@ -29,12 +37,13 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeNotNull;
 
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
-import android.Manifest.permission;
 import android.content.Context;
 import android.net.ConnectivityManager;
 import android.net.Network;
@@ -54,11 +63,11 @@
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.filters.LargeTest;
 
+import com.android.net.module.util.ArrayTrackRecord;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import com.android.testutils.DevSdkIgnoreRunner;
-
-import com.google.common.util.concurrent.SettableFuture;
+import com.android.testutils.FunctionalUtils.ThrowingRunnable;
 
 import org.junit.After;
 import org.junit.Before;
@@ -68,9 +77,11 @@
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -81,52 +92,714 @@
 @IgnoreUpTo(Build.VERSION_CODES.TIRAMISU) // Thread is available on only U+
 public class ThreadNetworkControllerTest {
     private static final int JOIN_TIMEOUT_MILLIS = 30 * 1000;
+    private static final int LEAVE_TIMEOUT_MILLIS = 2_000;
+    private static final int MIGRATION_TIMEOUT_MILLIS = 40 * 1_000;
     private static final int NETWORK_CALLBACK_TIMEOUT_MILLIS = 10 * 1000;
-    private static final int CALLBACK_TIMEOUT_MILLIS = 1000;
-    private static final String PERMISSION_THREAD_NETWORK_PRIVILEGED =
+    private static final int CALLBACK_TIMEOUT_MILLIS = 1_000;
+    private static final int ENABLED_TIMEOUT_MILLIS = 2_000;
+    private static final String THREAD_NETWORK_PRIVILEGED =
             "android.permission.THREAD_NETWORK_PRIVILEGED";
 
     @Rule public DevSdkIgnoreRule mIgnoreRule = new DevSdkIgnoreRule();
 
     private final Context mContext = ApplicationProvider.getApplicationContext();
     private ExecutorService mExecutor;
-    private ThreadNetworkManager mManager;
+    private ThreadNetworkController mController;
 
     private Set<String> mGrantedPermissions;
 
     @Before
-    public void setUp() {
-        mExecutor = Executors.newSingleThreadExecutor();
-        mManager = mContext.getSystemService(ThreadNetworkManager.class);
+    public void setUp() throws Exception {
+
         mGrantedPermissions = new HashSet<String>();
+        mExecutor = Executors.newSingleThreadExecutor();
+        ThreadNetworkManager manager = mContext.getSystemService(ThreadNetworkManager.class);
+        if (manager != null) {
+            mController = manager.getAllThreadNetworkControllers().get(0);
+        }
 
         // TODO: we will also need it in tearDown(), it's better to have a Rule to skip
         // tests if a feature is not available.
-        assumeNotNull(mManager);
+        assumeNotNull(mController);
+
+        setEnabledAndWait(mController, true);
     }
 
     @After
     public void tearDown() throws Exception {
-        if (mManager != null) {
-            leaveAndWait();
+        if (mController != null) {
+            grantPermissions(THREAD_NETWORK_PRIVILEGED);
+            CompletableFuture<Void> future = new CompletableFuture<>();
+            mController.leave(mExecutor, future::complete);
+            future.get(LEAVE_TIMEOUT_MILLIS, MILLISECONDS);
+        }
+        dropAllPermissions();
+    }
+
+    @Test
+    public void getThreadVersion_returnsAtLeastThreadVersion1P3() {
+        assertThat(mController.getThreadVersion()).isAtLeast(THREAD_VERSION_1_3);
+    }
+
+    @Test
+    public void registerStateCallback_permissionsGranted_returnsCurrentStates() throws Exception {
+        CompletableFuture<Integer> deviceRole = new CompletableFuture<>();
+        StateCallback callback = deviceRole::complete;
+
+        try {
+            runAsShell(
+                    ACCESS_NETWORK_STATE,
+                    () -> mController.registerStateCallback(mExecutor, callback));
+
+            assertThat(deviceRole.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS))
+                    .isEqualTo(DEVICE_ROLE_STOPPED);
+        } finally {
+            runAsShell(ACCESS_NETWORK_STATE, () -> mController.unregisterStateCallback(callback));
+        }
+    }
+
+    @Test
+    public void registerStateCallback_returnsUpdatedEnabledStates() throws Exception {
+        CompletableFuture<Void> setFuture1 = new CompletableFuture<>();
+        CompletableFuture<Void> setFuture2 = new CompletableFuture<>();
+        EnabledStateListener listener = new EnabledStateListener(mController);
+
+        try {
+            runAsShell(
+                    THREAD_NETWORK_PRIVILEGED,
+                    () -> {
+                        mController.setEnabled(false, mExecutor, newOutcomeReceiver(setFuture1));
+                    });
+            setFuture1.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+
+            runAsShell(
+                    THREAD_NETWORK_PRIVILEGED,
+                    () -> {
+                        mController.setEnabled(true, mExecutor, newOutcomeReceiver(setFuture2));
+                    });
+            setFuture2.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+
+            listener.expectThreadEnabledState(STATE_ENABLED);
+            listener.expectThreadEnabledState(STATE_DISABLING);
+            listener.expectThreadEnabledState(STATE_DISABLED);
+            listener.expectThreadEnabledState(STATE_ENABLED);
+        } finally {
+            listener.unregisterStateCallback();
+        }
+    }
+
+    @Test
+    public void registerStateCallback_noPermissions_throwsSecurityException() throws Exception {
+        dropAllPermissions();
+
+        assertThrows(
+                SecurityException.class,
+                () -> mController.registerStateCallback(mExecutor, role -> {}));
+    }
+
+    @Test
+    public void registerStateCallback_alreadyRegistered_throwsIllegalArgumentException()
+            throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE);
+        CompletableFuture<Integer> deviceRole = new CompletableFuture<>();
+        StateCallback callback = role -> deviceRole.complete(role);
+
+        mController.registerStateCallback(mExecutor, callback);
+
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> mController.registerStateCallback(mExecutor, callback));
+    }
+
+    @Test
+    public void unregisterStateCallback_noPermissions_throwsSecurityException() throws Exception {
+        CompletableFuture<Integer> deviceRole = new CompletableFuture<>();
+        StateCallback callback = role -> deviceRole.complete(role);
+        runAsShell(
+                ACCESS_NETWORK_STATE, () -> mController.registerStateCallback(mExecutor, callback));
+
+        try {
             dropAllPermissions();
+            assertThrows(
+                    SecurityException.class, () -> mController.unregisterStateCallback(callback));
+        } finally {
+            runAsShell(ACCESS_NETWORK_STATE, () -> mController.unregisterStateCallback(callback));
         }
     }
 
-    private List<ThreadNetworkController> getAllControllers() {
-        return mManager.getAllThreadNetworkControllers();
+    @Test
+    public void unregisterStateCallback_callbackRegistered_success() throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE);
+        CompletableFuture<Integer> deviceRole = new CompletableFuture<>();
+        StateCallback callback = role -> deviceRole.complete(role);
+
+        assertDoesNotThrow(() -> mController.registerStateCallback(mExecutor, callback));
+        mController.unregisterStateCallback(callback);
     }
 
-    private void leaveAndWait() throws Exception {
-        grantPermissions(PERMISSION_THREAD_NETWORK_PRIVILEGED);
+    @Test
+    public void unregisterStateCallback_callbackNotRegistered_throwsIllegalArgumentException()
+            throws Exception {
+        CompletableFuture<Integer> deviceRole = new CompletableFuture<>();
+        StateCallback callback = role -> deviceRole.complete(role);
 
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<Void> future = SettableFuture.create();
-            controller.leave(mExecutor, future::set);
-            future.get();
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> mController.unregisterStateCallback(callback));
+    }
+
+    @Test
+    public void unregisterStateCallback_alreadyUnregistered_throwsIllegalArgumentException()
+            throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE);
+        CompletableFuture<Integer> deviceRole = new CompletableFuture<>();
+        StateCallback callback = deviceRole::complete;
+        mController.registerStateCallback(mExecutor, callback);
+        mController.unregisterStateCallback(callback);
+
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> mController.unregisterStateCallback(callback));
+    }
+
+    @Test
+    public void registerOperationalDatasetCallback_permissionsGranted_returnsCurrentStates()
+            throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE, THREAD_NETWORK_PRIVILEGED);
+        CompletableFuture<ActiveOperationalDataset> activeFuture = new CompletableFuture<>();
+        CompletableFuture<PendingOperationalDataset> pendingFuture = new CompletableFuture<>();
+        var callback = newDatasetCallback(activeFuture, pendingFuture);
+
+        try {
+            mController.registerOperationalDatasetCallback(mExecutor, callback);
+
+            assertThat(activeFuture.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS)).isNull();
+            assertThat(pendingFuture.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS)).isNull();
+        } finally {
+            mController.unregisterOperationalDatasetCallback(callback);
         }
     }
 
+    @Test
+    public void registerOperationalDatasetCallback_noPermissions_throwsSecurityException()
+            throws Exception {
+        dropAllPermissions();
+        CompletableFuture<ActiveOperationalDataset> activeFuture = new CompletableFuture<>();
+        CompletableFuture<PendingOperationalDataset> pendingFuture = new CompletableFuture<>();
+        var callback = newDatasetCallback(activeFuture, pendingFuture);
+
+        assertThrows(
+                SecurityException.class,
+                () -> mController.registerOperationalDatasetCallback(mExecutor, callback));
+    }
+
+    @Test
+    public void unregisterOperationalDatasetCallback_callbackRegistered_success() throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE, THREAD_NETWORK_PRIVILEGED);
+        CompletableFuture<ActiveOperationalDataset> activeFuture = new CompletableFuture<>();
+        CompletableFuture<PendingOperationalDataset> pendingFuture = new CompletableFuture<>();
+        var callback = newDatasetCallback(activeFuture, pendingFuture);
+        mController.registerOperationalDatasetCallback(mExecutor, callback);
+
+        assertDoesNotThrow(() -> mController.unregisterOperationalDatasetCallback(callback));
+    }
+
+    @Test
+    public void unregisterOperationalDatasetCallback_noPermissions_throwsSecurityException()
+            throws Exception {
+        CompletableFuture<ActiveOperationalDataset> activeFuture = new CompletableFuture<>();
+        CompletableFuture<PendingOperationalDataset> pendingFuture = new CompletableFuture<>();
+        var callback = newDatasetCallback(activeFuture, pendingFuture);
+        runAsShell(
+                ACCESS_NETWORK_STATE,
+                THREAD_NETWORK_PRIVILEGED,
+                () -> mController.registerOperationalDatasetCallback(mExecutor, callback));
+
+        try {
+            dropAllPermissions();
+            assertThrows(
+                    SecurityException.class,
+                    () -> mController.unregisterOperationalDatasetCallback(callback));
+        } finally {
+            runAsShell(
+                    ACCESS_NETWORK_STATE,
+                    THREAD_NETWORK_PRIVILEGED,
+                    () -> mController.unregisterOperationalDatasetCallback(callback));
+        }
+    }
+
+    private static <V> OutcomeReceiver<V, ThreadNetworkException> newOutcomeReceiver(
+            CompletableFuture<V> future) {
+        return new OutcomeReceiver<V, ThreadNetworkException>() {
+            @Override
+            public void onResult(V result) {
+                future.complete(result);
+            }
+
+            @Override
+            public void onError(ThreadNetworkException e) {
+                future.completeExceptionally(e);
+            }
+        };
+    }
+
+    @Test
+    public void join_withPrivilegedPermission_success() throws Exception {
+        ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", mController);
+        CompletableFuture<Void> joinFuture = new CompletableFuture<>();
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> mController.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture)));
+        joinFuture.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
+
+        assertThat(isAttached(mController)).isTrue();
+        assertThat(getActiveOperationalDataset(mController)).isEqualTo(activeDataset);
+    }
+
+    @Test
+    public void join_withoutPrivilegedPermission_throwsSecurityException() throws Exception {
+        dropAllPermissions();
+        ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", mController);
+
+        assertThrows(
+                SecurityException.class, () -> mController.join(activeDataset, mExecutor, v -> {}));
+    }
+
+    @Test
+    public void join_threadDisabled_failsWithErrorThreadDisabled() throws Exception {
+        setEnabledAndWait(mController, false);
+        ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", mController);
+        CompletableFuture<Void> joinFuture = new CompletableFuture<>();
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> mController.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture)));
+
+        var thrown =
+                assertThrows(
+                        ExecutionException.class,
+                        () -> joinFuture.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS));
+        var threadException = (ThreadNetworkException) thrown.getCause();
+        assertThat(threadException.getErrorCode()).isEqualTo(ERROR_THREAD_DISABLED);
+    }
+
+    @Test
+    public void join_concurrentRequests_firstOneIsAborted() throws Exception {
+        final byte[] KEY_1 = new byte[] {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
+        final byte[] KEY_2 = new byte[] {2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2};
+        ActiveOperationalDataset activeDataset1 =
+                new ActiveOperationalDataset.Builder(newRandomizedDataset("TestNet", mController))
+                        .setNetworkKey(KEY_1)
+                        .build();
+        ActiveOperationalDataset activeDataset2 =
+                new ActiveOperationalDataset.Builder(activeDataset1).setNetworkKey(KEY_2).build();
+        CompletableFuture<Void> joinFuture1 = new CompletableFuture<>();
+        CompletableFuture<Void> joinFuture2 = new CompletableFuture<>();
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> {
+                    mController.join(activeDataset1, mExecutor, newOutcomeReceiver(joinFuture1));
+                    mController.join(activeDataset2, mExecutor, newOutcomeReceiver(joinFuture2));
+                });
+
+        var thrown =
+                assertThrows(
+                        ExecutionException.class,
+                        () -> joinFuture1.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS));
+        var threadException = (ThreadNetworkException) thrown.getCause();
+        assertThat(threadException.getErrorCode()).isEqualTo(ERROR_ABORTED);
+        joinFuture2.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
+        assertThat(isAttached(mController)).isTrue();
+        assertThat(getActiveOperationalDataset(mController)).isEqualTo(activeDataset2);
+    }
+
+    @Test
+    public void leave_withPrivilegedPermission_success() throws Exception {
+        CompletableFuture<Void> leaveFuture = new CompletableFuture<>();
+        joinRandomizedDatasetAndWait(mController);
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> mController.leave(mExecutor, newOutcomeReceiver(leaveFuture)));
+        leaveFuture.get(LEAVE_TIMEOUT_MILLIS, MILLISECONDS);
+
+        assertThat(getDeviceRole(mController)).isEqualTo(DEVICE_ROLE_STOPPED);
+    }
+
+    @Test
+    public void leave_withoutPrivilegedPermission_throwsSecurityException() {
+        dropAllPermissions();
+
+        assertThrows(SecurityException.class, () -> mController.leave(mExecutor, v -> {}));
+    }
+
+    @Test
+    public void leave_threadDisabled_success() throws Exception {
+        setEnabledAndWait(mController, false);
+        CompletableFuture<Void> leaveFuture = new CompletableFuture<>();
+
+        leave(mController, newOutcomeReceiver(leaveFuture));
+        leaveFuture.get(LEAVE_TIMEOUT_MILLIS, MILLISECONDS);
+
+        assertThat(getDeviceRole(mController)).isEqualTo(DEVICE_ROLE_STOPPED);
+    }
+
+    @Test
+    public void leave_concurrentRequests_bothSuccess() throws Exception {
+        CompletableFuture<Void> leaveFuture1 = new CompletableFuture<>();
+        CompletableFuture<Void> leaveFuture2 = new CompletableFuture<>();
+        joinRandomizedDatasetAndWait(mController);
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> {
+                    mController.leave(mExecutor, newOutcomeReceiver(leaveFuture1));
+                    mController.leave(mExecutor, newOutcomeReceiver(leaveFuture2));
+                });
+
+        leaveFuture1.get(LEAVE_TIMEOUT_MILLIS, MILLISECONDS);
+        leaveFuture2.get(LEAVE_TIMEOUT_MILLIS, MILLISECONDS);
+        assertThat(getDeviceRole(mController)).isEqualTo(DEVICE_ROLE_STOPPED);
+    }
+
+    @Test
+    public void scheduleMigration_withPrivilegedPermission_newDatasetApplied() throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE, THREAD_NETWORK_PRIVILEGED);
+        ActiveOperationalDataset activeDataset1 =
+                new ActiveOperationalDataset.Builder(newRandomizedDataset("TestNet", mController))
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(1L, 0, false))
+                        .setExtendedPanId(new byte[] {1, 1, 1, 1, 1, 1, 1, 1})
+                        .build();
+        ActiveOperationalDataset activeDataset2 =
+                new ActiveOperationalDataset.Builder(activeDataset1)
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(2L, 0, false))
+                        .setNetworkName("ThreadNet2")
+                        .build();
+        PendingOperationalDataset pendingDataset2 =
+                new PendingOperationalDataset(
+                        activeDataset2,
+                        OperationalDatasetTimestamp.fromInstant(Instant.now()),
+                        Duration.ofSeconds(30));
+        CompletableFuture<Void> joinFuture = new CompletableFuture<>();
+        CompletableFuture<Void> migrateFuture = new CompletableFuture<>();
+        mController.join(activeDataset1, mExecutor, newOutcomeReceiver(joinFuture));
+        joinFuture.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
+
+        mController.scheduleMigration(
+                pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture));
+        migrateFuture.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
+
+        CompletableFuture<Boolean> dataset2IsApplied = new CompletableFuture<>();
+        CompletableFuture<Boolean> pendingDatasetIsRemoved = new CompletableFuture<>();
+        OperationalDatasetCallback datasetCallback =
+                new OperationalDatasetCallback() {
+                    @Override
+                    public void onActiveOperationalDatasetChanged(
+                            ActiveOperationalDataset activeDataset) {
+                        if (activeDataset.equals(activeDataset2)) {
+                            dataset2IsApplied.complete(true);
+                        }
+                    }
+
+                    @Override
+                    public void onPendingOperationalDatasetChanged(
+                            PendingOperationalDataset pendingDataset) {
+                        if (pendingDataset == null) {
+                            pendingDatasetIsRemoved.complete(true);
+                        }
+                    }
+                };
+        mController.registerOperationalDatasetCallback(directExecutor(), datasetCallback);
+        try {
+            assertThat(dataset2IsApplied.get(MIGRATION_TIMEOUT_MILLIS, MILLISECONDS)).isTrue();
+            assertThat(pendingDatasetIsRemoved.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS)).isTrue();
+        } finally {
+            mController.unregisterOperationalDatasetCallback(datasetCallback);
+        }
+    }
+
+    @Test
+    public void scheduleMigration_whenNotAttached_failWithPreconditionError() throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE, THREAD_NETWORK_PRIVILEGED);
+        PendingOperationalDataset pendingDataset =
+                new PendingOperationalDataset(
+                        newRandomizedDataset("TestNet", mController),
+                        OperationalDatasetTimestamp.fromInstant(Instant.now()),
+                        Duration.ofSeconds(30));
+        CompletableFuture<Void> migrateFuture = new CompletableFuture<>();
+
+        mController.scheduleMigration(pendingDataset, mExecutor, newOutcomeReceiver(migrateFuture));
+
+        ThreadNetworkException thrown =
+                (ThreadNetworkException)
+                        assertThrows(ExecutionException.class, migrateFuture::get).getCause();
+        assertThat(thrown.getErrorCode()).isEqualTo(ERROR_FAILED_PRECONDITION);
+    }
+
+    @Test
+    public void scheduleMigration_secondRequestHasSmallerTimestamp_rejectedByLeader()
+            throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE, THREAD_NETWORK_PRIVILEGED);
+        final ActiveOperationalDataset activeDataset =
+                new ActiveOperationalDataset.Builder(newRandomizedDataset("testNet", mController))
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(1L, 0, false))
+                        .build();
+        ActiveOperationalDataset activeDataset1 =
+                new ActiveOperationalDataset.Builder(activeDataset)
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(2L, 0, false))
+                        .setNetworkName("testNet1")
+                        .build();
+        PendingOperationalDataset pendingDataset1 =
+                new PendingOperationalDataset(
+                        activeDataset1,
+                        new OperationalDatasetTimestamp(100, 0, false),
+                        Duration.ofSeconds(30));
+        ActiveOperationalDataset activeDataset2 =
+                new ActiveOperationalDataset.Builder(activeDataset)
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(3L, 0, false))
+                        .setNetworkName("testNet2")
+                        .build();
+        PendingOperationalDataset pendingDataset2 =
+                new PendingOperationalDataset(
+                        activeDataset2,
+                        new OperationalDatasetTimestamp(20, 0, false),
+                        Duration.ofSeconds(30));
+        CompletableFuture<Void> joinFuture = new CompletableFuture<>();
+        CompletableFuture<Void> migrateFuture1 = new CompletableFuture<>();
+        CompletableFuture<Void> migrateFuture2 = new CompletableFuture<>();
+        mController.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture));
+        joinFuture.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
+
+        mController.scheduleMigration(
+                pendingDataset1, mExecutor, newOutcomeReceiver(migrateFuture1));
+        migrateFuture1.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
+        mController.scheduleMigration(
+                pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture2));
+
+        ThreadNetworkException thrown =
+                (ThreadNetworkException)
+                        assertThrows(ExecutionException.class, migrateFuture2::get).getCause();
+        assertThat(thrown.getErrorCode()).isEqualTo(ERROR_REJECTED_BY_PEER);
+    }
+
+    @Test
+    public void scheduleMigration_secondRequestHasLargerTimestamp_newDatasetApplied()
+            throws Exception {
+        grantPermissions(ACCESS_NETWORK_STATE, THREAD_NETWORK_PRIVILEGED);
+        final ActiveOperationalDataset activeDataset =
+                new ActiveOperationalDataset.Builder(newRandomizedDataset("validName", mController))
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(1L, 0, false))
+                        .build();
+        ActiveOperationalDataset activeDataset1 =
+                new ActiveOperationalDataset.Builder(activeDataset)
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(2L, 0, false))
+                        .setNetworkName("testNet1")
+                        .build();
+        PendingOperationalDataset pendingDataset1 =
+                new PendingOperationalDataset(
+                        activeDataset1,
+                        new OperationalDatasetTimestamp(100, 0, false),
+                        Duration.ofSeconds(30));
+        ActiveOperationalDataset activeDataset2 =
+                new ActiveOperationalDataset.Builder(activeDataset)
+                        .setActiveTimestamp(new OperationalDatasetTimestamp(3L, 0, false))
+                        .setNetworkName("testNet2")
+                        .build();
+        PendingOperationalDataset pendingDataset2 =
+                new PendingOperationalDataset(
+                        activeDataset2,
+                        new OperationalDatasetTimestamp(200, 0, false),
+                        Duration.ofSeconds(30));
+        CompletableFuture<Void> joinFuture = new CompletableFuture<>();
+        CompletableFuture<Void> migrateFuture1 = new CompletableFuture<>();
+        CompletableFuture<Void> migrateFuture2 = new CompletableFuture<>();
+        mController.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture));
+        joinFuture.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
+
+        mController.scheduleMigration(
+                pendingDataset1, mExecutor, newOutcomeReceiver(migrateFuture1));
+        migrateFuture1.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
+        mController.scheduleMigration(
+                pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture2));
+        migrateFuture2.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
+
+        CompletableFuture<Boolean> dataset2IsApplied = new CompletableFuture<>();
+        CompletableFuture<Boolean> pendingDatasetIsRemoved = new CompletableFuture<>();
+        OperationalDatasetCallback datasetCallback =
+                new OperationalDatasetCallback() {
+                    @Override
+                    public void onActiveOperationalDatasetChanged(
+                            ActiveOperationalDataset activeDataset) {
+                        if (activeDataset.equals(activeDataset2)) {
+                            dataset2IsApplied.complete(true);
+                        }
+                    }
+
+                    @Override
+                    public void onPendingOperationalDatasetChanged(
+                            PendingOperationalDataset pendingDataset) {
+                        if (pendingDataset == null) {
+                            pendingDatasetIsRemoved.complete(true);
+                        }
+                    }
+                };
+        mController.registerOperationalDatasetCallback(directExecutor(), datasetCallback);
+        try {
+            assertThat(dataset2IsApplied.get(MIGRATION_TIMEOUT_MILLIS, MILLISECONDS)).isTrue();
+            assertThat(pendingDatasetIsRemoved.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS)).isTrue();
+        } finally {
+            mController.unregisterOperationalDatasetCallback(datasetCallback);
+        }
+    }
+
+    @Test
+    public void scheduleMigration_threadDisabled_failsWithErrorThreadDisabled() throws Exception {
+        ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", mController);
+        PendingOperationalDataset pendingDataset =
+                new PendingOperationalDataset(
+                        activeDataset,
+                        OperationalDatasetTimestamp.fromInstant(Instant.now()),
+                        Duration.ofSeconds(30));
+        joinRandomizedDatasetAndWait(mController);
+        CompletableFuture<Void> migrationFuture = new CompletableFuture<>();
+
+        setEnabledAndWait(mController, false);
+
+        scheduleMigration(mController, pendingDataset, newOutcomeReceiver(migrationFuture));
+
+        ThreadNetworkException thrown =
+                (ThreadNetworkException)
+                        assertThrows(ExecutionException.class, migrationFuture::get).getCause();
+        assertThat(thrown.getErrorCode()).isEqualTo(ERROR_THREAD_DISABLED);
+    }
+
+    @Test
+    public void createRandomizedDataset_wrongNetworkNameLength_throwsIllegalArgumentException() {
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> mController.createRandomizedDataset("", mExecutor, dataset -> {}));
+
+        assertThrows(
+                IllegalArgumentException.class,
+                () ->
+                        mController.createRandomizedDataset(
+                                "ANetNameIs17Bytes", mExecutor, dataset -> {}));
+    }
+
+    @Test
+    public void createRandomizedDataset_validNetworkName_success() throws Exception {
+        ActiveOperationalDataset dataset = newRandomizedDataset("validName", mController);
+
+        assertThat(dataset.getNetworkName()).isEqualTo("validName");
+        assertThat(dataset.getPanId()).isLessThan(0xffff);
+        assertThat(dataset.getChannelMask().size()).isAtLeast(1);
+        assertThat(dataset.getExtendedPanId()).hasLength(8);
+        assertThat(dataset.getNetworkKey()).hasLength(16);
+        assertThat(dataset.getPskc()).hasLength(16);
+        assertThat(dataset.getMeshLocalPrefix().getPrefixLength()).isEqualTo(64);
+        assertThat(dataset.getMeshLocalPrefix().getRawAddress()[0]).isEqualTo((byte) 0xfd);
+    }
+
+    @Test
+    public void setEnabled_permissionsGranted_succeeds() throws Exception {
+        CompletableFuture<Void> setFuture1 = new CompletableFuture<>();
+        CompletableFuture<Void> setFuture2 = new CompletableFuture<>();
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> mController.setEnabled(false, mExecutor, newOutcomeReceiver(setFuture1)));
+        setFuture1.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+        waitForEnabledState(mController, booleanToEnabledState(false));
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> mController.setEnabled(true, mExecutor, newOutcomeReceiver(setFuture2)));
+        setFuture2.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+        waitForEnabledState(mController, booleanToEnabledState(true));
+    }
+
+    @Test
+    public void setEnabled_noPermissions_throwsSecurityException() throws Exception {
+        CompletableFuture<Void> setFuture = new CompletableFuture<>();
+        assertThrows(
+                SecurityException.class,
+                () -> mController.setEnabled(false, mExecutor, newOutcomeReceiver(setFuture)));
+    }
+
+    @Test
+    public void setEnabled_disable_leavesThreadNetwork() throws Exception {
+        joinRandomizedDatasetAndWait(mController);
+        setEnabledAndWait(mController, false);
+        assertThat(getDeviceRole(mController)).isEqualTo(DEVICE_ROLE_STOPPED);
+    }
+
+    @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<>();
+        CompletableFuture<Void> setFuture2 = new CompletableFuture<>();
+        EnabledStateListener listener = new EnabledStateListener(mController);
+        listener.expectThreadEnabledState(STATE_ENABLED);
+
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> {
+                    mController.setEnabled(true, mExecutor, newOutcomeReceiver(setFuture1));
+                    mController.setEnabled(false, mExecutor, newOutcomeReceiver(setFuture2));
+                });
+        setFuture1.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+        setFuture2.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+
+        listener.expectThreadEnabledState(STATE_DISABLING);
+        listener.expectThreadEnabledState(STATE_DISABLED);
+        assertThat(getDeviceRole(mController)).isEqualTo(DEVICE_ROLE_STOPPED);
+        // FIXME: this is not called when a exception is thrown after the creation of `listener`
+        listener.unregisterStateCallback();
+    }
+
+    // TODO (b/322437869): add test case to verify when Thread is in DISABLING state, any commands
+    // (join/leave/scheduleMigration/setEnabled) fail with ERROR_BUSY. This is not currently tested
+    // because DISABLING has very short lifecycle, it's not possible to guarantee the command can be
+    // sent before state changes to DISABLED.
+
+    @Test
+    public void threadNetworkCallback_deviceAttached_threadNetworkIsAvailable() throws Exception {
+        CompletableFuture<Network> networkFuture = new CompletableFuture<>();
+        ConnectivityManager cm = mContext.getSystemService(ConnectivityManager.class);
+        NetworkRequest networkRequest =
+                new NetworkRequest.Builder()
+                        .addTransportType(NetworkCapabilities.TRANSPORT_THREAD)
+                        .build();
+        ConnectivityManager.NetworkCallback networkCallback =
+                new ConnectivityManager.NetworkCallback() {
+                    @Override
+                    public void onAvailable(Network network) {
+                        networkFuture.complete(network);
+                    }
+                };
+
+        joinRandomizedDatasetAndWait(mController);
+        runAsShell(
+                ACCESS_NETWORK_STATE,
+                () -> cm.registerNetworkCallback(networkRequest, networkCallback));
+
+        assertThat(isAttached(mController)).isTrue();
+        assertThat(networkFuture.get(NETWORK_CALLBACK_TIMEOUT_MILLIS, MILLISECONDS)).isNotNull();
+    }
+
     private void grantPermissions(String... permissions) {
         for (String permission : permissions) {
             mGrantedPermissions.add(permission);
@@ -142,8 +815,8 @@
 
     private static ActiveOperationalDataset newRandomizedDataset(
             String networkName, ThreadNetworkController controller) throws Exception {
-        SettableFuture<ActiveOperationalDataset> future = SettableFuture.create();
-        controller.createRandomizedDataset(networkName, directExecutor(), future::set);
+        CompletableFuture<ActiveOperationalDataset> future = new CompletableFuture<>();
+        controller.createRandomizedDataset(networkName, directExecutor(), future::complete);
         return future.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
     }
 
@@ -152,642 +825,189 @@
     }
 
     private static int getDeviceRole(ThreadNetworkController controller) throws Exception {
-        SettableFuture<Integer> future = SettableFuture.create();
-        StateCallback callback = future::set;
-        controller.registerStateCallback(directExecutor(), callback);
-        int role = future.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
-        controller.unregisterStateCallback(callback);
-        return role;
+        CompletableFuture<Integer> future = new CompletableFuture<>();
+        StateCallback callback = future::complete;
+        runAsShell(
+                ACCESS_NETWORK_STATE,
+                () -> controller.registerStateCallback(directExecutor(), callback));
+        try {
+            return future.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
+        } finally {
+            runAsShell(ACCESS_NETWORK_STATE, () -> controller.unregisterStateCallback(callback));
+        }
+    }
+
+    private static int waitForAttachedState(ThreadNetworkController controller) throws Exception {
+        List<Integer> attachedRoles = new ArrayList<>();
+        attachedRoles.add(DEVICE_ROLE_CHILD);
+        attachedRoles.add(DEVICE_ROLE_ROUTER);
+        attachedRoles.add(DEVICE_ROLE_LEADER);
+        return waitForStateAnyOf(controller, attachedRoles);
     }
 
     private static int waitForStateAnyOf(
             ThreadNetworkController controller, List<Integer> deviceRoles) throws Exception {
-        SettableFuture<Integer> future = SettableFuture.create();
+        CompletableFuture<Integer> future = new CompletableFuture<>();
         StateCallback callback =
                 newRole -> {
                     if (deviceRoles.contains(newRole)) {
-                        future.set(newRole);
+                        future.complete(newRole);
                     }
                 };
         controller.registerStateCallback(directExecutor(), callback);
-        int role = future.get();
+        int role = future.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
         controller.unregisterStateCallback(callback);
         return role;
     }
 
+    private static void waitForEnabledState(ThreadNetworkController controller, int state)
+            throws Exception {
+        CompletableFuture<Integer> future = new CompletableFuture<>();
+        StateCallback callback =
+                new ThreadNetworkController.StateCallback() {
+                    @Override
+                    public void onDeviceRoleChanged(int r) {}
+
+                    @Override
+                    public void onThreadEnableStateChanged(int enabled) {
+                        if (enabled == state) {
+                            future.complete(enabled);
+                        }
+                    }
+                };
+        runAsShell(
+                ACCESS_NETWORK_STATE,
+                () -> controller.registerStateCallback(directExecutor(), callback));
+        future.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+        runAsShell(ACCESS_NETWORK_STATE, () -> controller.unregisterStateCallback(callback));
+    }
+
+    private void leave(
+            ThreadNetworkController controller,
+            OutcomeReceiver<Void, ThreadNetworkException> receiver) {
+        runAsShell(THREAD_NETWORK_PRIVILEGED, () -> controller.leave(mExecutor, receiver));
+    }
+
+    private void scheduleMigration(
+            ThreadNetworkController controller,
+            PendingOperationalDataset pendingDataset,
+            OutcomeReceiver<Void, ThreadNetworkException> receiver) {
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> controller.scheduleMigration(pendingDataset, mExecutor, receiver));
+    }
+
+    private class EnabledStateListener {
+        private ArrayTrackRecord<Integer> mEnabledStates = new ArrayTrackRecord<>();
+        private final ArrayTrackRecord<Integer>.ReadHead mReadHead = mEnabledStates.newReadHead();
+        ThreadNetworkController mController;
+        StateCallback mCallback =
+                new ThreadNetworkController.StateCallback() {
+                    @Override
+                    public void onDeviceRoleChanged(int r) {}
+
+                    @Override
+                    public void onThreadEnableStateChanged(int enabled) {
+                        mEnabledStates.add(enabled);
+                    }
+                };
+
+        EnabledStateListener(ThreadNetworkController controller) {
+            this.mController = controller;
+            runAsShell(
+                    ACCESS_NETWORK_STATE,
+                    () -> controller.registerStateCallback(mExecutor, mCallback));
+        }
+
+        public void expectThreadEnabledState(int enabled) {
+            assertNotNull(mReadHead.poll(ENABLED_TIMEOUT_MILLIS, e -> (e == enabled)));
+        }
+
+        public void unregisterStateCallback() {
+            runAsShell(ACCESS_NETWORK_STATE, () -> mController.unregisterStateCallback(mCallback));
+        }
+    }
+
+    private int booleanToEnabledState(boolean enabled) {
+        return enabled ? STATE_ENABLED : STATE_DISABLED;
+    }
+
+    private void setEnabledAndWait(ThreadNetworkController controller, boolean enabled)
+            throws Exception {
+        CompletableFuture<Void> setFuture = new CompletableFuture<>();
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> controller.setEnabled(enabled, mExecutor, newOutcomeReceiver(setFuture)));
+        setFuture.get(ENABLED_TIMEOUT_MILLIS, MILLISECONDS);
+        waitForEnabledState(controller, booleanToEnabledState(enabled));
+    }
+
+    private CompletableFuture joinRandomizedDataset(ThreadNetworkController controller)
+            throws Exception {
+        ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", controller);
+        CompletableFuture<Void> joinFuture = new CompletableFuture<>();
+        runAsShell(
+                THREAD_NETWORK_PRIVILEGED,
+                () -> controller.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture)));
+        return joinFuture;
+    }
+
+    private void joinRandomizedDatasetAndWait(ThreadNetworkController controller) throws Exception {
+        CompletableFuture<Void> joinFuture = joinRandomizedDataset(controller);
+        joinFuture.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
+        assertThat(isAttached(controller)).isTrue();
+    }
+
     private static ActiveOperationalDataset getActiveOperationalDataset(
             ThreadNetworkController controller) throws Exception {
-        SettableFuture<ActiveOperationalDataset> future = SettableFuture.create();
-        OperationalDatasetCallback callback = future::set;
-        controller.registerOperationalDatasetCallback(directExecutor(), callback);
-        ActiveOperationalDataset dataset = future.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
-        controller.unregisterOperationalDatasetCallback(callback);
-        return dataset;
+        CompletableFuture<ActiveOperationalDataset> future = new CompletableFuture<>();
+        OperationalDatasetCallback callback = future::complete;
+        runAsShell(
+                ACCESS_NETWORK_STATE,
+                THREAD_NETWORK_PRIVILEGED,
+                () -> controller.registerOperationalDatasetCallback(directExecutor(), callback));
+        try {
+            return future.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
+        } finally {
+            runAsShell(
+                    ACCESS_NETWORK_STATE,
+                    THREAD_NETWORK_PRIVILEGED,
+                    () -> controller.unregisterOperationalDatasetCallback(callback));
+        }
     }
 
     private static PendingOperationalDataset getPendingOperationalDataset(
             ThreadNetworkController controller) throws Exception {
-        SettableFuture<ActiveOperationalDataset> activeFuture = SettableFuture.create();
-        SettableFuture<PendingOperationalDataset> pendingFuture = SettableFuture.create();
+        CompletableFuture<ActiveOperationalDataset> activeFuture = new CompletableFuture<>();
+        CompletableFuture<PendingOperationalDataset> pendingFuture = new CompletableFuture<>();
         controller.registerOperationalDatasetCallback(
                 directExecutor(), newDatasetCallback(activeFuture, pendingFuture));
-        return pendingFuture.get();
+        return pendingFuture.get(CALLBACK_TIMEOUT_MILLIS, MILLISECONDS);
     }
 
     private static OperationalDatasetCallback newDatasetCallback(
-            SettableFuture<ActiveOperationalDataset> activeFuture,
-            SettableFuture<PendingOperationalDataset> pendingFuture) {
+            CompletableFuture<ActiveOperationalDataset> activeFuture,
+            CompletableFuture<PendingOperationalDataset> pendingFuture) {
         return new OperationalDatasetCallback() {
             @Override
             public void onActiveOperationalDatasetChanged(
                     ActiveOperationalDataset activeOpDataset) {
-                activeFuture.set(activeOpDataset);
+                activeFuture.complete(activeOpDataset);
             }
 
             @Override
             public void onPendingOperationalDatasetChanged(
                     PendingOperationalDataset pendingOpDataset) {
-                pendingFuture.set(pendingOpDataset);
+                pendingFuture.complete(pendingOpDataset);
             }
         };
     }
 
-    @Test
-    public void getThreadVersion_returnsAtLeastThreadVersion1P3() {
-        for (ThreadNetworkController controller : getAllControllers()) {
-            assertThat(controller.getThreadVersion()).isAtLeast(THREAD_VERSION_1_3);
+    private static void assertDoesNotThrow(ThrowingRunnable runnable) {
+        try {
+            runnable.run();
+        } catch (Throwable e) {
+            fail("Should not have thrown " + e);
         }
     }
-
-    @Test
-    public void registerStateCallback_permissionsGranted_returnsCurrentStates() throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<Integer> deviceRole = SettableFuture.create();
-            StateCallback callback = deviceRole::set;
-
-            try {
-                controller.registerStateCallback(mExecutor, callback);
-
-                assertThat(deviceRole.get()).isEqualTo(DEVICE_ROLE_STOPPED);
-            } finally {
-                controller.unregisterStateCallback(callback);
-            }
-        }
-    }
-
-    @Test
-    public void registerStateCallback_noPermissions_throwsSecurityException() throws Exception {
-        dropAllPermissions();
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            assertThrows(
-                    SecurityException.class,
-                    () -> controller.registerStateCallback(mExecutor, role -> {}));
-        }
-    }
-
-    @Test
-    public void registerStateCallback_alreadyRegistered_throwsIllegalArgumentException()
-            throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<Integer> deviceRole = SettableFuture.create();
-            StateCallback callback = role -> deviceRole.set(role);
-            controller.registerStateCallback(mExecutor, callback);
-
-            assertThrows(
-                    IllegalArgumentException.class,
-                    () -> controller.registerStateCallback(mExecutor, callback));
-        }
-    }
-
-    @Test
-    public void unregisterStateCallback_noPermissions_throwsSecurityException() throws Exception {
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<Integer> deviceRole = SettableFuture.create();
-            StateCallback callback = role -> deviceRole.set(role);
-            grantPermissions(permission.ACCESS_NETWORK_STATE);
-            controller.registerStateCallback(mExecutor, callback);
-
-            try {
-                dropAllPermissions();
-                assertThrows(
-                        SecurityException.class,
-                        () -> controller.unregisterStateCallback(callback));
-            } finally {
-                grantPermissions(permission.ACCESS_NETWORK_STATE);
-                controller.unregisterStateCallback(callback);
-            }
-        }
-    }
-
-    @Test
-    public void unregisterStateCallback_callbackRegistered_success() throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE);
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<Integer> deviceRole = SettableFuture.create();
-            StateCallback callback = role -> deviceRole.set(role);
-            controller.registerStateCallback(mExecutor, callback);
-
-            controller.unregisterStateCallback(callback);
-        }
-    }
-
-    @Test
-    public void unregisterStateCallback_callbackNotRegistered_throwsIllegalArgumentException()
-            throws Exception {
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<Integer> deviceRole = SettableFuture.create();
-            StateCallback callback = role -> deviceRole.set(role);
-
-            assertThrows(
-                    IllegalArgumentException.class,
-                    () -> controller.unregisterStateCallback(callback));
-        }
-    }
-
-    @Test
-    public void unregisterStateCallback_alreadyUnregistered_throwsIllegalArgumentException()
-            throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE);
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<Integer> deviceRole = SettableFuture.create();
-            StateCallback callback = deviceRole::set;
-            controller.registerStateCallback(mExecutor, callback);
-            controller.unregisterStateCallback(callback);
-
-            assertThrows(
-                    IllegalArgumentException.class,
-                    () -> controller.unregisterStateCallback(callback));
-        }
-    }
-
-    @Test
-    public void registerOperationalDatasetCallback_permissionsGranted_returnsCurrentStates()
-            throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<ActiveOperationalDataset> activeFuture = SettableFuture.create();
-            SettableFuture<PendingOperationalDataset> pendingFuture = SettableFuture.create();
-            var callback = newDatasetCallback(activeFuture, pendingFuture);
-
-            try {
-                controller.registerOperationalDatasetCallback(mExecutor, callback);
-
-                assertThat(activeFuture.get()).isNull();
-                assertThat(pendingFuture.get()).isNull();
-            } finally {
-                controller.unregisterOperationalDatasetCallback(callback);
-            }
-        }
-    }
-
-    @Test
-    public void registerOperationalDatasetCallback_noPermissions_throwsSecurityException()
-            throws Exception {
-        dropAllPermissions();
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<ActiveOperationalDataset> activeFuture = SettableFuture.create();
-            SettableFuture<PendingOperationalDataset> pendingFuture = SettableFuture.create();
-            var callback = newDatasetCallback(activeFuture, pendingFuture);
-
-            assertThrows(
-                    SecurityException.class,
-                    () -> controller.registerOperationalDatasetCallback(mExecutor, callback));
-        }
-    }
-
-    @Test
-    public void unregisterOperationalDatasetCallback_callbackRegistered_success() throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<ActiveOperationalDataset> activeFuture = SettableFuture.create();
-            SettableFuture<PendingOperationalDataset> pendingFuture = SettableFuture.create();
-            var callback = newDatasetCallback(activeFuture, pendingFuture);
-            controller.registerOperationalDatasetCallback(mExecutor, callback);
-
-            controller.unregisterOperationalDatasetCallback(callback);
-        }
-    }
-
-    @Test
-    public void unregisterOperationalDatasetCallback_noPermissions_throwsSecurityException()
-            throws Exception {
-        dropAllPermissions();
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            SettableFuture<ActiveOperationalDataset> activeFuture = SettableFuture.create();
-            SettableFuture<PendingOperationalDataset> pendingFuture = SettableFuture.create();
-            var callback = newDatasetCallback(activeFuture, pendingFuture);
-            grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-            controller.registerOperationalDatasetCallback(mExecutor, callback);
-
-            try {
-                dropAllPermissions();
-                assertThrows(
-                        SecurityException.class,
-                        () -> controller.unregisterOperationalDatasetCallback(callback));
-            } finally {
-                grantPermissions(
-                        permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-                controller.unregisterOperationalDatasetCallback(callback);
-            }
-        }
-    }
-
-    private static <V> OutcomeReceiver<V, ThreadNetworkException> newOutcomeReceiver(
-            SettableFuture<V> future) {
-        return new OutcomeReceiver<V, ThreadNetworkException>() {
-            @Override
-            public void onResult(V result) {
-                future.set(result);
-            }
-
-            @Override
-            public void onError(ThreadNetworkException e) {
-                future.setException(e);
-            }
-        };
-    }
-
-    @Test
-    public void join_withPrivilegedPermission_success() throws Exception {
-        grantPermissions(PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", controller);
-            SettableFuture<Void> joinFuture = SettableFuture.create();
-
-            controller.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture));
-            joinFuture.get();
-
-            grantPermissions(permission.ACCESS_NETWORK_STATE);
-            assertThat(isAttached(controller)).isTrue();
-            assertThat(getActiveOperationalDataset(controller)).isEqualTo(activeDataset);
-        }
-    }
-
-    @Test
-    public void join_withoutPrivilegedPermission_throwsSecurityException() throws Exception {
-        dropAllPermissions();
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", controller);
-
-            assertThrows(
-                    SecurityException.class,
-                    () -> controller.join(activeDataset, mExecutor, v -> {}));
-        }
-    }
-
-    @Test
-    public void join_concurrentRequests_firstOneIsAborted() throws Exception {
-        grantPermissions(PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        final byte[] KEY_1 = new byte[] {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
-        final byte[] KEY_2 = new byte[] {2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2};
-        for (ThreadNetworkController controller : getAllControllers()) {
-            ActiveOperationalDataset activeDataset1 =
-                    new ActiveOperationalDataset.Builder(
-                                    newRandomizedDataset("TestNet", controller))
-                            .setNetworkKey(KEY_1)
-                            .build();
-            ActiveOperationalDataset activeDataset2 =
-                    new ActiveOperationalDataset.Builder(activeDataset1)
-                            .setNetworkKey(KEY_2)
-                            .build();
-            SettableFuture<Void> joinFuture1 = SettableFuture.create();
-            SettableFuture<Void> joinFuture2 = SettableFuture.create();
-
-            controller.join(activeDataset1, mExecutor, newOutcomeReceiver(joinFuture1));
-            controller.join(activeDataset2, mExecutor, newOutcomeReceiver(joinFuture2));
-
-            ThreadNetworkException thrown =
-                    (ThreadNetworkException)
-                            assertThrows(ExecutionException.class, joinFuture1::get).getCause();
-            assertThat(thrown.getErrorCode()).isEqualTo(ERROR_ABORTED);
-            joinFuture2.get();
-            grantPermissions(permission.ACCESS_NETWORK_STATE);
-            assertThat(isAttached(controller)).isTrue();
-            assertThat(getActiveOperationalDataset(controller)).isEqualTo(activeDataset2);
-        }
-    }
-
-    @Test
-    public void leave_withPrivilegedPermission_success() throws Exception {
-        grantPermissions(PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", controller);
-            SettableFuture<Void> joinFuture = SettableFuture.create();
-            SettableFuture<Void> leaveFuture = SettableFuture.create();
-            controller.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture));
-            joinFuture.get();
-
-            controller.leave(mExecutor, newOutcomeReceiver(leaveFuture));
-            leaveFuture.get();
-
-            grantPermissions(permission.ACCESS_NETWORK_STATE);
-            assertThat(getDeviceRole(controller)).isEqualTo(DEVICE_ROLE_STOPPED);
-        }
-    }
-
-    @Test
-    public void leave_withoutPrivilegedPermission_throwsSecurityException() {
-        dropAllPermissions();
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            assertThrows(SecurityException.class, () -> controller.leave(mExecutor, v -> {}));
-        }
-    }
-
-    @Test
-    public void leave_concurrentRequests_bothSuccess() throws Exception {
-        grantPermissions(PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", controller);
-            SettableFuture<Void> joinFuture = SettableFuture.create();
-            SettableFuture<Void> leaveFuture1 = SettableFuture.create();
-            SettableFuture<Void> leaveFuture2 = SettableFuture.create();
-            controller.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture));
-            joinFuture.get();
-
-            controller.leave(mExecutor, newOutcomeReceiver(leaveFuture1));
-            controller.leave(mExecutor, newOutcomeReceiver(leaveFuture2));
-
-            leaveFuture1.get();
-            leaveFuture2.get();
-            grantPermissions(permission.ACCESS_NETWORK_STATE);
-            assertThat(getDeviceRole(controller)).isEqualTo(DEVICE_ROLE_STOPPED);
-        }
-    }
-
-    @Test
-    public void scheduleMigration_withPrivilegedPermission_newDatasetApplied() throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            ActiveOperationalDataset activeDataset1 =
-                    new ActiveOperationalDataset.Builder(
-                                    newRandomizedDataset("TestNet", controller))
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(1L, 0, false))
-                            .setExtendedPanId(new byte[] {1, 1, 1, 1, 1, 1, 1, 1})
-                            .build();
-            ActiveOperationalDataset activeDataset2 =
-                    new ActiveOperationalDataset.Builder(activeDataset1)
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(2L, 0, false))
-                            .setNetworkName("ThreadNet2")
-                            .build();
-            PendingOperationalDataset pendingDataset2 =
-                    new PendingOperationalDataset(
-                            activeDataset2,
-                            OperationalDatasetTimestamp.fromInstant(Instant.now()),
-                            Duration.ofSeconds(30));
-            SettableFuture<Void> joinFuture = SettableFuture.create();
-            SettableFuture<Void> migrateFuture = SettableFuture.create();
-            controller.join(activeDataset1, mExecutor, newOutcomeReceiver(joinFuture));
-            joinFuture.get();
-
-            controller.scheduleMigration(
-                    pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture));
-            migrateFuture.get();
-
-            SettableFuture<Boolean> dataset2IsApplied = SettableFuture.create();
-            SettableFuture<Boolean> pendingDatasetIsRemoved = SettableFuture.create();
-            OperationalDatasetCallback datasetCallback =
-                    new OperationalDatasetCallback() {
-                        @Override
-                        public void onActiveOperationalDatasetChanged(
-                                ActiveOperationalDataset activeDataset) {
-                            if (activeDataset.equals(activeDataset2)) {
-                                dataset2IsApplied.set(true);
-                            }
-                        }
-
-                        @Override
-                        public void onPendingOperationalDatasetChanged(
-                                PendingOperationalDataset pendingDataset) {
-                            if (pendingDataset == null) {
-                                pendingDatasetIsRemoved.set(true);
-                            }
-                        }
-                    };
-            controller.registerOperationalDatasetCallback(directExecutor(), datasetCallback);
-            assertThat(dataset2IsApplied.get()).isTrue();
-            assertThat(pendingDatasetIsRemoved.get()).isTrue();
-            controller.unregisterOperationalDatasetCallback(datasetCallback);
-        }
-    }
-
-    @Test
-    public void scheduleMigration_whenNotAttached_failWithPreconditionError() throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            PendingOperationalDataset pendingDataset =
-                    new PendingOperationalDataset(
-                            newRandomizedDataset("TestNet", controller),
-                            OperationalDatasetTimestamp.fromInstant(Instant.now()),
-                            Duration.ofSeconds(30));
-            SettableFuture<Void> migrateFuture = SettableFuture.create();
-
-            controller.scheduleMigration(
-                    pendingDataset, mExecutor, newOutcomeReceiver(migrateFuture));
-
-            ThreadNetworkException thrown =
-                    (ThreadNetworkException)
-                            assertThrows(ExecutionException.class, migrateFuture::get).getCause();
-            assertThat(thrown.getErrorCode()).isEqualTo(ERROR_FAILED_PRECONDITION);
-        }
-    }
-
-    @Test
-    public void scheduleMigration_secondRequestHasSmallerTimestamp_rejectedByLeader()
-            throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            final ActiveOperationalDataset activeDataset =
-                    new ActiveOperationalDataset.Builder(
-                                    newRandomizedDataset("testNet", controller))
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(1L, 0, false))
-                            .build();
-            ActiveOperationalDataset activeDataset1 =
-                    new ActiveOperationalDataset.Builder(activeDataset)
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(2L, 0, false))
-                            .setNetworkName("testNet1")
-                            .build();
-            PendingOperationalDataset pendingDataset1 =
-                    new PendingOperationalDataset(
-                            activeDataset1,
-                            new OperationalDatasetTimestamp(100, 0, false),
-                            Duration.ofSeconds(30));
-            ActiveOperationalDataset activeDataset2 =
-                    new ActiveOperationalDataset.Builder(activeDataset)
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(3L, 0, false))
-                            .setNetworkName("testNet2")
-                            .build();
-            PendingOperationalDataset pendingDataset2 =
-                    new PendingOperationalDataset(
-                            activeDataset2,
-                            new OperationalDatasetTimestamp(20, 0, false),
-                            Duration.ofSeconds(30));
-            SettableFuture<Void> joinFuture = SettableFuture.create();
-            SettableFuture<Void> migrateFuture1 = SettableFuture.create();
-            SettableFuture<Void> migrateFuture2 = SettableFuture.create();
-            controller.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture));
-            joinFuture.get();
-
-            controller.scheduleMigration(
-                    pendingDataset1, mExecutor, newOutcomeReceiver(migrateFuture1));
-            migrateFuture1.get();
-            controller.scheduleMigration(
-                    pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture2));
-
-            ThreadNetworkException thrown =
-                    (ThreadNetworkException)
-                            assertThrows(ExecutionException.class, migrateFuture2::get).getCause();
-            assertThat(thrown.getErrorCode()).isEqualTo(ERROR_REJECTED_BY_PEER);
-        }
-    }
-
-    @Test
-    public void scheduleMigration_secondRequestHasLargerTimestamp_newDatasetApplied()
-            throws Exception {
-        grantPermissions(permission.ACCESS_NETWORK_STATE, PERMISSION_THREAD_NETWORK_PRIVILEGED);
-
-        for (ThreadNetworkController controller : getAllControllers()) {
-            final ActiveOperationalDataset activeDataset =
-                    new ActiveOperationalDataset.Builder(
-                                    newRandomizedDataset("validName", controller))
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(1L, 0, false))
-                            .build();
-            ActiveOperationalDataset activeDataset1 =
-                    new ActiveOperationalDataset.Builder(activeDataset)
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(2L, 0, false))
-                            .setNetworkName("testNet1")
-                            .build();
-            PendingOperationalDataset pendingDataset1 =
-                    new PendingOperationalDataset(
-                            activeDataset1,
-                            new OperationalDatasetTimestamp(100, 0, false),
-                            Duration.ofSeconds(30));
-            ActiveOperationalDataset activeDataset2 =
-                    new ActiveOperationalDataset.Builder(activeDataset)
-                            .setActiveTimestamp(new OperationalDatasetTimestamp(3L, 0, false))
-                            .setNetworkName("testNet2")
-                            .build();
-            PendingOperationalDataset pendingDataset2 =
-                    new PendingOperationalDataset(
-                            activeDataset2,
-                            new OperationalDatasetTimestamp(200, 0, false),
-                            Duration.ofSeconds(30));
-            SettableFuture<Void> joinFuture = SettableFuture.create();
-            SettableFuture<Void> migrateFuture1 = SettableFuture.create();
-            SettableFuture<Void> migrateFuture2 = SettableFuture.create();
-            controller.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture));
-            joinFuture.get();
-
-            controller.scheduleMigration(
-                    pendingDataset1, mExecutor, newOutcomeReceiver(migrateFuture1));
-            migrateFuture1.get();
-            controller.scheduleMigration(
-                    pendingDataset2, mExecutor, newOutcomeReceiver(migrateFuture2));
-            migrateFuture2.get();
-
-            SettableFuture<Boolean> dataset2IsApplied = SettableFuture.create();
-            SettableFuture<Boolean> pendingDatasetIsRemoved = SettableFuture.create();
-            OperationalDatasetCallback datasetCallback =
-                    new OperationalDatasetCallback() {
-                        @Override
-                        public void onActiveOperationalDatasetChanged(
-                                ActiveOperationalDataset activeDataset) {
-                            if (activeDataset.equals(activeDataset2)) {
-                                dataset2IsApplied.set(true);
-                            }
-                        }
-
-                        @Override
-                        public void onPendingOperationalDatasetChanged(
-                                PendingOperationalDataset pendingDataset) {
-                            if (pendingDataset == null) {
-                                pendingDatasetIsRemoved.set(true);
-                            }
-                        }
-                    };
-            controller.registerOperationalDatasetCallback(directExecutor(), datasetCallback);
-            assertThat(dataset2IsApplied.get()).isTrue();
-            assertThat(pendingDatasetIsRemoved.get()).isTrue();
-            controller.unregisterOperationalDatasetCallback(datasetCallback);
-        }
-    }
-
-    @Test
-    public void createRandomizedDataset_wrongNetworkNameLength_throwsIllegalArgumentException() {
-        for (ThreadNetworkController controller : getAllControllers()) {
-            assertThrows(
-                    IllegalArgumentException.class,
-                    () -> controller.createRandomizedDataset("", mExecutor, dataset -> {}));
-
-            assertThrows(
-                    IllegalArgumentException.class,
-                    () ->
-                            controller.createRandomizedDataset(
-                                    "ANetNameIs17Bytes", mExecutor, dataset -> {}));
-        }
-    }
-
-    @Test
-    public void createRandomizedDataset_validNetworkName_success() throws Exception {
-        for (ThreadNetworkController controller : getAllControllers()) {
-            ActiveOperationalDataset dataset = newRandomizedDataset("validName", controller);
-
-            assertThat(dataset.getNetworkName()).isEqualTo("validName");
-            assertThat(dataset.getPanId()).isLessThan(0xffff);
-            assertThat(dataset.getChannelMask().size()).isAtLeast(1);
-            assertThat(dataset.getExtendedPanId()).hasLength(8);
-            assertThat(dataset.getNetworkKey()).hasLength(16);
-            assertThat(dataset.getPskc()).hasLength(16);
-            assertThat(dataset.getMeshLocalPrefix().getPrefixLength()).isEqualTo(64);
-            assertThat(dataset.getMeshLocalPrefix().getRawAddress()[0]).isEqualTo((byte) 0xfd);
-        }
-    }
-
-    @Test
-    public void threadNetworkCallback_deviceAttached_threadNetworkIsAvailable() throws Exception {
-        ThreadNetworkController controller = mManager.getAllThreadNetworkControllers().get(0);
-        ActiveOperationalDataset activeDataset = newRandomizedDataset("TestNet", controller);
-        SettableFuture<Void> joinFuture = SettableFuture.create();
-        SettableFuture<Network> networkFuture = SettableFuture.create();
-        ConnectivityManager cm = mContext.getSystemService(ConnectivityManager.class);
-        NetworkRequest networkRequest =
-                new NetworkRequest.Builder()
-                        .addTransportType(NetworkCapabilities.TRANSPORT_THREAD)
-                        .build();
-        ConnectivityManager.NetworkCallback networkCallback =
-                new ConnectivityManager.NetworkCallback() {
-                    @Override
-                    public void onAvailable(Network network) {
-                        networkFuture.set(network);
-                    }
-                };
-
-        runAsShell(
-                PERMISSION_THREAD_NETWORK_PRIVILEGED,
-                () -> controller.join(activeDataset, mExecutor, newOutcomeReceiver(joinFuture)));
-        runAsShell(
-                permission.ACCESS_NETWORK_STATE,
-                () -> cm.registerNetworkCallback(networkRequest, networkCallback));
-
-        joinFuture.get(JOIN_TIMEOUT_MILLIS, MILLISECONDS);
-        runAsShell(
-                permission.ACCESS_NETWORK_STATE, () -> assertThat(isAttached(controller)).isTrue());
-        assertThat(networkFuture.get(NETWORK_CALLBACK_TIMEOUT_MILLIS, MILLISECONDS)).isNotNull();
-    }
 }
diff --git a/thread/tests/integration/Android.bp b/thread/tests/integration/Android.bp
index 405fb76..ebd6a4d 100644
--- a/thread/tests/integration/Android.bp
+++ b/thread/tests/integration/Android.bp
@@ -41,11 +41,13 @@
     name: "ThreadNetworkIntegrationTests",
     platform_apis: true,
     manifest: "AndroidManifest.xml",
+    test_config: "AndroidTest.xml",
     defaults: [
         "framework-connectivity-test-defaults",
-        "ThreadNetworkIntegrationTestsDefaults"
+        "ThreadNetworkIntegrationTestsDefaults",
     ],
     test_suites: [
+        "mts-tethering",
         "general-tests",
     ],
     srcs: [
diff --git a/thread/tests/integration/AndroidTest.xml b/thread/tests/integration/AndroidTest.xml
new file mode 100644
index 0000000..ec9b5f3
--- /dev/null
+++ b/thread/tests/integration/AndroidTest.xml
@@ -0,0 +1,44 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    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.
+ -->
+
+<configuration description="Config for Thread integration tests">
+    <option name="test-tag" value="ThreadNetworkIntegrationTests" />
+    <option name="test-suite-tag" value="apct" />
+
+    <!--
+        Only run tests if the device under test is SDK version 34 (Android 14) or above.
+    -->
+    <object type="module_controller"
+            class="com.android.tradefed.testtype.suite.module.Sdk34ModuleController" />
+
+    <!-- Run tests in MTS only if the Tethering Mainline module is installed. -->
+    <object type="module_controller"
+            class="com.android.tradefed.testtype.suite.module.MainlineTestModuleController">
+        <option name="mainline-module-package-name" value="com.google.android.tethering" />
+    </object>
+
+    <!-- Install test -->
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="test-file-name" value="ThreadNetworkIntegrationTests.apk" />
+        <option name="check-min-sdk" value="true" />
+        <option name="cleanup-apks" value="true" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.thread.tests.integration" />
+    </test>
+</configuration>
diff --git a/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java b/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java
index 25f5bd3..5d9f084 100644
--- a/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java
+++ b/thread/tests/integration/src/android/net/thread/BorderRoutingTest.java
@@ -32,9 +32,10 @@
 import static com.android.testutils.TestPermissionUtil.runAsShell;
 
 import static com.google.common.io.BaseEncoding.base16;
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assume.assumeNotNull;
 import static org.junit.Assume.assumeTrue;
 
 import android.content.Context;
@@ -50,8 +51,6 @@
 import com.android.testutils.TapPacketReader;
 import com.android.testutils.TestNetworkTracker;
 
-import com.google.common.util.concurrent.MoreExecutors;
-
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -68,9 +67,7 @@
 public class BorderRoutingTest {
     private static final String TAG = BorderRoutingTest.class.getSimpleName();
     private final Context mContext = ApplicationProvider.getApplicationContext();
-    private final ThreadNetworkManager mThreadNetworkManager =
-            mContext.getSystemService(ThreadNetworkManager.class);
-    private ThreadNetworkController mThreadNetworkController;
+    private ThreadNetworkController mController;
     private HandlerThread mHandlerThread;
     private Handler mHandler;
     private TestNetworkTracker mInfraNetworkTracker;
@@ -88,12 +85,18 @@
 
     @Before
     public void setUp() throws Exception {
+        final ThreadNetworkManager manager = mContext.getSystemService(ThreadNetworkManager.class);
+        if (manager != null) {
+            mController = manager.getAllThreadNetworkControllers().get(0);
+        }
+
+        // Run the tests on only devices where the Thread feature is available
+        assumeNotNull(mController);
+
         mHandlerThread = new HandlerThread(getClass().getSimpleName());
         mHandlerThread.start();
         mHandler = new Handler(mHandlerThread.getLooper());
-        var threadControllers = mThreadNetworkManager.getAllThreadNetworkControllers();
-        assertEquals(threadControllers.size(), 1);
-        mThreadNetworkController = threadControllers.get(0);
+
         mInfraNetworkTracker =
                 runAsShell(
                         MANAGE_TEST_NETWORKS,
@@ -105,27 +108,28 @@
                 NETWORK_SETTINGS,
                 () -> {
                     CountDownLatch latch = new CountDownLatch(1);
-                    mThreadNetworkController.setTestNetworkAsUpstream(
+                    mController.setTestNetworkAsUpstream(
                             mInfraNetworkTracker.getTestIface().getInterfaceName(),
-                            MoreExecutors.directExecutor(),
-                            v -> {
-                                latch.countDown();
-                            });
+                            directExecutor(),
+                            v -> latch.countDown());
                     latch.await();
                 });
     }
 
     @After
     public void tearDown() throws Exception {
+        if (mController == null) {
+            return;
+        }
+
         runAsShell(
                 PERMISSION_THREAD_NETWORK_PRIVILEGED,
                 NETWORK_SETTINGS,
                 () -> {
                     CountDownLatch latch = new CountDownLatch(2);
-                    mThreadNetworkController.setTestNetworkAsUpstream(
-                            null, MoreExecutors.directExecutor(), v -> latch.countDown());
-                    mThreadNetworkController.leave(
-                            MoreExecutors.directExecutor(), v -> latch.countDown());
+                    mController.setTestNetworkAsUpstream(
+                            null, directExecutor(), v -> latch.countDown());
+                    mController.leave(directExecutor(), v -> latch.countDown());
                     latch.await(10, TimeUnit.SECONDS);
                 });
         runAsShell(MANAGE_TEST_NETWORKS, () -> mInfraNetworkTracker.teardown());
@@ -150,12 +154,8 @@
         // BR forms a network.
         runAsShell(
                 PERMISSION_THREAD_NETWORK_PRIVILEGED,
-                () -> {
-                    mThreadNetworkController.join(
-                            DEFAULT_DATASET, MoreExecutors.directExecutor(), result -> {});
-                });
-        waitForStateAnyOf(
-                mThreadNetworkController, List.of(DEVICE_ROLE_LEADER), 30 /* timeoutSeconds */);
+                () -> mController.join(DEFAULT_DATASET, directExecutor(), result -> {}));
+        waitForStateAnyOf(mController, List.of(DEVICE_ROLE_LEADER), 30 /* timeoutSeconds */);
 
         // Creates a Full Thread Device (FTD) and lets it join the network.
         FullThreadDevice ftd = new FullThreadDevice(5 /* node ID */);
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 44a8ab7..1d83abc 100644
--- a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkControllerServiceTest.java
+++ b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkControllerServiceTest.java
@@ -24,6 +24,7 @@
 import static com.google.common.io.BaseEncoding.base16;
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
@@ -85,6 +86,7 @@
     @Mock private TunInterfaceController mMockTunIfController;
     @Mock private ParcelFileDescriptor mMockTunFd;
     @Mock private InfraInterfaceController mMockInfraIfController;
+    @Mock private ThreadPersistentSettings mMockPersistentSettings;
     private Context mContext;
     private TestLooper mTestLooper;
     private FakeOtDaemon mFakeOtDaemon;
@@ -104,6 +106,8 @@
 
         when(mMockTunIfController.getTunFd()).thenReturn(mMockTunFd);
 
+        when(mMockPersistentSettings.get(any())).thenReturn(true);
+
         mService =
                 new ThreadNetworkControllerService(
                         ApplicationProvider.getApplicationContext(),
@@ -112,7 +116,8 @@
                         () -> mFakeOtDaemon,
                         mMockConnectivityManager,
                         mMockTunIfController,
-                        mMockInfraIfController);
+                        mMockInfraIfController,
+                        mMockPersistentSettings);
         mService.setTestNetworkAgent(mMockNetworkAgent);
     }
 
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 17cdd01..5ca6511 100644
--- a/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
+++ b/thread/tests/unit/src/com/android/server/thread/ThreadNetworkCountryCodeTest.java
@@ -22,6 +22,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyDouble;
@@ -55,6 +56,7 @@
 import android.telephony.SubscriptionManager;
 import android.telephony.TelephonyManager;
 
+import androidx.annotation.Nullable;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -70,6 +72,9 @@
 import org.mockito.MockitoAnnotations;
 import org.mockito.stubbing.Answer;
 
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
@@ -80,6 +85,8 @@
 public class ThreadNetworkCountryCodeTest {
     private static final String TEST_COUNTRY_CODE_US = "US";
     private static final String TEST_COUNTRY_CODE_CN = "CN";
+    private static final String TEST_COUNTRY_CODE_INVALID = "INVALID";
+    private static final String TEST_WIFI_DEFAULT_COUNTRY_CODE = "00";
     private static final int TEST_SIM_SLOT_INDEX_0 = 0;
     private static final int TEST_SIM_SLOT_INDEX_1 = 1;
 
@@ -144,16 +151,20 @@
                 .when(mThreadNetworkControllerService)
                 .setCountryCode(any(), any(IOperationReceiver.class));
 
-        mThreadNetworkCountryCode =
-                new ThreadNetworkCountryCode(
-                        mLocationManager,
-                        mThreadNetworkControllerService,
-                        mGeocoder,
-                        mConnectivityResources,
-                        mWifiManager,
-                        mContext,
-                        mTelephonyManager,
-                        mSubscriptionManager);
+        mThreadNetworkCountryCode = newCountryCodeWithOemSource(null);
+    }
+
+    private ThreadNetworkCountryCode newCountryCodeWithOemSource(@Nullable String oemCountryCode) {
+        return new ThreadNetworkCountryCode(
+                mLocationManager,
+                mThreadNetworkControllerService,
+                mGeocoder,
+                mConnectivityResources,
+                mWifiManager,
+                mContext,
+                mTelephonyManager,
+                mSubscriptionManager,
+                oemCountryCode);
     }
 
     private static Address newAddress(String countryCode) {
@@ -163,6 +174,13 @@
     }
 
     @Test
+    public void threadNetworkCountryCode_invalidOemCountryCode_illegalArgumentExceptionIsThrown() {
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> newCountryCodeWithOemSource(TEST_COUNTRY_CODE_INVALID));
+    }
+
+    @Test
     public void initialize_defaultCountryCodeIsUsed() {
         mThreadNetworkCountryCode.initialize();
 
@@ -170,6 +188,15 @@
     }
 
     @Test
+    public void initialize_oemCountryCodeAvailable_oemCountryCodeIsUsed() {
+        mThreadNetworkCountryCode = newCountryCodeWithOemSource(TEST_COUNTRY_CODE_US);
+
+        mThreadNetworkCountryCode.initialize();
+
+        assertThat(mThreadNetworkCountryCode.getCountryCode()).isEqualTo(TEST_COUNTRY_CODE_US);
+    }
+
+    @Test
     public void initialize_locationUseIsDisabled_locationFunctionIsNotCalled() {
         when(mResources.getBoolean(R.bool.config_thread_location_use_for_country_code_enabled))
                 .thenReturn(false);
@@ -233,6 +260,21 @@
     }
 
     @Test
+    public void wifiCountryCode_wifiDefaultCountryCodeIsActive_wifiCountryCodeIsNotUsed() {
+        mThreadNetworkCountryCode.initialize();
+
+        verify(mWifiManager)
+                .registerActiveCountryCodeChangedCallback(
+                        any(), mWifiCountryCodeReceiverCaptor.capture());
+        mWifiCountryCodeReceiverCaptor
+                .getValue()
+                .onActiveCountryCodeChanged(TEST_WIFI_DEFAULT_COUNTRY_CODE);
+
+        assertThat(mThreadNetworkCountryCode.getCountryCode())
+                .isNotEqualTo(TEST_WIFI_DEFAULT_COUNTRY_CODE);
+    }
+
+    @Test
     public void wifiCountryCode_wifiCountryCodeIsInactive_defaultCountryCodeIsUsed() {
         mThreadNetworkCountryCode.initialize();
         verify(mWifiManager)
@@ -406,4 +448,22 @@
                 .setCountryCode(eq(TEST_COUNTRY_CODE_CN), mOperationReceiverCaptor.capture());
         assertThat(mThreadNetworkCountryCode.getCountryCode()).isEqualTo(DEFAULT_COUNTRY_CODE);
     }
+
+    @Test
+    public void dump_allCountryCodeInfoAreDumped() {
+        StringWriter stringWriter = new StringWriter();
+        PrintWriter printWriter = new PrintWriter(stringWriter);
+
+        mThreadNetworkCountryCode.dump(new FileDescriptor(), printWriter, null);
+        String outputString = stringWriter.toString();
+
+        assertThat(outputString).contains("mOverrideCountryCodeInfo");
+        assertThat(outputString).contains("mTelephonyCountryCodeSlotInfoMap");
+        assertThat(outputString).contains("mTelephonyCountryCodeInfo");
+        assertThat(outputString).contains("mWifiCountryCodeInfo");
+        assertThat(outputString).contains("mTelephonyLastCountryCodeInfo");
+        assertThat(outputString).contains("mLocationCountryCodeInfo");
+        assertThat(outputString).contains("mOemCountryCodeInfo");
+        assertThat(outputString).contains("mCurrentCountryCodeInfo");
+    }
 }