diff --git a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
index 735bdc6..18ceb24 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
+++ b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
@@ -22,10 +22,16 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
 import android.net.http.HttpEngine;
+import android.net.http.HttpException;
+import android.net.http.InlineExecutionProhibitedException;
+import android.net.http.UploadDataProvider;
+import android.net.http.UploadDataSink;
 import android.net.http.UrlRequest;
 import android.net.http.UrlRequest.Status;
 import android.net.http.UrlResponseInfo;
@@ -43,8 +49,18 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
 @RunWith(AndroidJUnit4.class)
 public class UrlRequestTest {
+    private static final Executor DIRECT_EXECUTOR = Runnable::run;
+
     private TestUrlRequestCallback mCallback;
     private HttpCtsTestServer mTestServer;
     private HttpEngine mHttpEngine;
@@ -128,8 +144,203 @@
     }
 
     @Test
-    public void testUrlRequestFail_FailedCalled() throws Exception {
+    public void testUrlRequestFail_FailedCalled() {
         createUrlRequestBuilder("http://0.0.0.0:0/").build().start();
         mCallback.expectCallback(ResponseStep.ON_FAILED);
     }
+
+    @Test
+    public void testUrlRequest_directExecutor_allowed() throws InterruptedException {
+        TestUrlRequestCallback callback = new TestUrlRequestCallback();
+        callback.setAllowDirectExecutor(true);
+        UrlRequest.Builder builder = mHttpEngine.newUrlRequestBuilder(
+                mTestServer.getEchoBodyUrl(), callback, DIRECT_EXECUTOR);
+        UploadDataProvider dataProvider = InMemoryUploadDataProvider.fromUtf8String("test");
+        builder.setUploadDataProvider(dataProvider, DIRECT_EXECUTOR);
+        builder.addHeader("Content-Type", "text/plain;charset=UTF-8");
+        builder.setAllowDirectExecutor(true);
+        builder.build().start();
+        callback.blockForDone();
+
+        if (callback.mOnErrorCalled) {
+            throw new AssertionError("Expected no exception", callback.mError);
+        }
+
+        assertEquals(200, callback.mResponseInfo.getHttpStatusCode());
+        assertEquals("test", callback.mResponseAsString);
+    }
+
+    @Test
+    public void testUrlRequest_directExecutor_disallowed_uploadDataProvider() throws Exception {
+        TestUrlRequestCallback callback = new TestUrlRequestCallback();
+        // This applies just locally to the test callback, not to SUT
+        callback.setAllowDirectExecutor(true);
+
+        UrlRequest.Builder builder = mHttpEngine.newUrlRequestBuilder(
+                mTestServer.getEchoBodyUrl(), callback, Executors.newSingleThreadExecutor());
+        UploadDataProvider dataProvider = InMemoryUploadDataProvider.fromUtf8String("test");
+
+        builder.setUploadDataProvider(dataProvider, DIRECT_EXECUTOR)
+                .addHeader("Content-Type", "text/plain;charset=UTF-8")
+                .build()
+                .start();
+        callback.blockForDone();
+
+        assertTrue(callback.mOnErrorCalled);
+        assertTrue(callback.mError.getCause() instanceof InlineExecutionProhibitedException);
+    }
+
+    @Test
+    public void testUrlRequest_directExecutor_disallowed_responseCallback() throws Exception {
+        TestUrlRequestCallback callback = new TestUrlRequestCallback();
+        // This applies just locally to the test callback, not to SUT
+        callback.setAllowDirectExecutor(true);
+
+        UrlRequest.Builder builder = mHttpEngine.newUrlRequestBuilder(
+                mTestServer.getEchoBodyUrl(), callback, DIRECT_EXECUTOR);
+        UploadDataProvider dataProvider = InMemoryUploadDataProvider.fromUtf8String("test");
+
+        builder.setUploadDataProvider(dataProvider, Executors.newSingleThreadExecutor())
+                .addHeader("Content-Type", "text/plain;charset=UTF-8")
+                .build()
+                .start();
+        callback.blockForDone();
+
+        assertTrue(callback.mOnErrorCalled);
+        assertTrue(callback.mError.getCause() instanceof InlineExecutionProhibitedException);
+    }
+
+    @Test
+    public void testUrlRequest_nonDirectByteBuffer() throws Exception {
+        BlockingQueue<HttpException> onFailedException = new ArrayBlockingQueue<>(1);
+
+        UrlRequest request =
+                mHttpEngine
+                        .newUrlRequestBuilder(
+                                mTestServer.getSuccessUrl(),
+                                new StubUrlRequestCallback() {
+                                    @Override
+                                    public void onResponseStarted(
+                                            UrlRequest request, UrlResponseInfo info) {
+                                        // note: allocate, not allocateDirect
+                                        request.read(ByteBuffer.allocate(1024));
+                                    }
+
+                                    @Override
+                                    public void onFailed(
+                                            UrlRequest request,
+                                            UrlResponseInfo info,
+                                            HttpException error) {
+                                        onFailedException.add(error);
+                                    }
+                                },
+                                Executors.newSingleThreadExecutor())
+                        .build();
+        request.start();
+
+        HttpException e = onFailedException.poll(5, TimeUnit.SECONDS);
+        assertNotNull(e);
+        assertTrue(e.getCause() instanceof IllegalArgumentException);
+        assertTrue(e.getCause().getMessage().contains("direct"));
+    }
+
+    @Test
+    public void testUrlRequest_fullByteBuffer() throws Exception {
+        BlockingQueue<HttpException> onFailedException = new ArrayBlockingQueue<>(1);
+
+        UrlRequest request =
+                mHttpEngine
+                        .newUrlRequestBuilder(
+                                mTestServer.getSuccessUrl(),
+                                new StubUrlRequestCallback() {
+                                    @Override
+                                    public void onResponseStarted(
+                                            UrlRequest request, UrlResponseInfo info) {
+                                        ByteBuffer bb = ByteBuffer.allocateDirect(1024);
+                                        bb.position(bb.limit());
+                                        request.read(bb);
+                                    }
+
+                                    @Override
+                                    public void onFailed(
+                                            UrlRequest request,
+                                            UrlResponseInfo info,
+                                            HttpException error) {
+                                        onFailedException.add(error);
+                                    }
+                                },
+                                Executors.newSingleThreadExecutor())
+                        .build();
+        request.start();
+
+        HttpException e = onFailedException.poll(5, TimeUnit.SECONDS);
+        assertNotNull(e);
+        assertTrue(e.getCause() instanceof IllegalArgumentException);
+        assertTrue(e.getCause().getMessage().contains("full"));
+    }
+
+    private static class StubUrlRequestCallback extends UrlRequest.Callback {
+
+        @Override
+        public void onRedirectReceived(
+                UrlRequest request, UrlResponseInfo info, String newLocationUrl) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void onResponseStarted(UrlRequest request, UrlResponseInfo info) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void onReadCompleted(
+                UrlRequest request, UrlResponseInfo info, ByteBuffer byteBuffer) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void onSucceeded(UrlRequest request, UrlResponseInfo info) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void onFailed(UrlRequest request, UrlResponseInfo info, HttpException error) {
+            throw new UnsupportedOperationException(error);
+        }
+    }
+
+    private static class InMemoryUploadDataProvider extends UploadDataProvider {
+        private final byte[] mBody;
+        private int mNextChunkStartIndex = 0;
+
+        private InMemoryUploadDataProvider(byte[] body) {
+            this.mBody = body;
+        }
+
+        static InMemoryUploadDataProvider fromUtf8String(String body) {
+            return new InMemoryUploadDataProvider(body.getBytes(StandardCharsets.UTF_8));
+        }
+
+        @Override
+        public long getLength() {
+            return mBody.length;
+        }
+
+        @Override
+        public void read(UploadDataSink uploadDataSink, ByteBuffer byteBuffer) {
+            if (mNextChunkStartIndex >= getLength()) {
+                throw new IllegalStateException("Body of known length is exhausted");
+            }
+            int nextChunkSize =
+                    Math.min(mBody.length - mNextChunkStartIndex, byteBuffer.remaining());
+            byteBuffer.put(mBody, mNextChunkStartIndex, nextChunkSize);
+            mNextChunkStartIndex += nextChunkSize;
+            uploadDataSink.onReadSucceeded(false);
+        }
+
+        @Override
+        public void rewind(UploadDataSink uploadDataSink) {
+            mNextChunkStartIndex = 0;
+        }
+    }
 }
diff --git a/Cronet/tools/import/import_cronet.sh b/Cronet/tools/import/import_cronet.sh
index 7642914..eb82551 100755
--- a/Cronet/tools/import/import_cronet.sh
+++ b/Cronet/tools/import/import_cronet.sh
@@ -19,40 +19,70 @@
 #  Environment:
 #   ANDROID_BUILD_TOP: path the root of the current Android directory.
 #  Arguments:
-#   -l: The last revision that was imported.
-#   -n: The new revision to import.
+#   -l rev: The last revision that was imported.
+#  Optional Arguments:
+#   -n rev: The new revision to import.
+#   -f: Force copybara to ignore a failure to find the last imported revision.
 
-OPTSTRING=l:n:
+OPTSTRING=fl:n:
 
 usage() {
     cat <<EOF
-Usage: import_cronet.sh -l last-rev -n new-rev
+Usage: import_cronet.sh -n new-rev [-l last-rev] [-f]
 EOF
     exit 1
 }
 
 #######################################
+# Create upstream-import branch in external/cronet.
+# Globals:
+#   ANDROID_BUILD_TOP
+# Arguments:
+#   none
+#######################################
+setup_upstream_import_branch() {
+    local git_dir="${ANDROID_BUILD_TOP}/external/cronet"
+    local initial_empty_repo_sha="d1add53d6e90815f363c91d433735556ce79b0d2"
+
+    # Suppress error message if branch already exists.
+    (cd "${git_dir}" && git branch upstream-import "${initial_empty_repo_sha}") 2>/dev/null
+}
+
+#######################################
 # Runs the copybara import of Chromium
 # Globals:
 #   ANDROID_BUILD_TOP
 # Arguments:
-#   last_rev, string
 #   new_rev, string
+#   last_rev, string or empty
+#   force, string or empty
 #######################################
 do_run_copybara() {
-    local _last_rev=$1
-    local _new_rev=$2
+    local _new_rev=$1
+    local _last_rev=$2
+    local _force=$3
+
+    local -a flags
+    flags+=(--git-destination-url="file://${ANDROID_BUILD_TOP}/external/cronet")
+    flags+=(--repo-timeout 3h)
+
+    if [ ! -z "${_force}" ]; then
+        flags+=(--force)
+    fi
+
+    if [ ! -z "${_last_rev}" ]; then
+        flags+=(--last-rev "${_last_rev}")
+    fi
 
     /google/bin/releases/copybara/public/copybara/copybara \
-        --git-destination-url="file://${ANDROID_BUILD_TOP}/external/cronet" \
-        --last-rev "${_last_rev}" \
-        --repo-timeout 3h \
+        "${flags[@]}" \
         "${ANDROID_BUILD_TOP}/packages/modules/Connectivity/Cronet/tools/import/copy.bara.sky" \
         import_cronet "${_new_rev}"
 }
 
 while getopts $OPTSTRING opt; do
     case "${opt}" in
+        f) force=true ;;
         l) last_rev="${OPTARG}" ;;
         n) new_rev="${OPTARG}" ;;
         ?) usage ;;
@@ -60,17 +90,11 @@
     esac
 done
 
-# TODO: Get last-rev from METADATA file.
-# Setting last-rev may only be required for the first commit.
-if [ -z "${last_rev}" ]; then
-    echo "-l argument required"
-    usage
-fi
-
 if [ -z "${new_rev}" ]; then
     echo "-n argument required"
     usage
 fi
 
-do_run_copybara "${last_rev}" "${new_rev}"
+setup_upstream_import_branch
+do_run_copybara "${new_rev}" "${last_rev}" "${force}"
 
diff --git a/framework-t/Android.bp b/framework-t/Android.bp
index 7ef20c5..1a8d46b 100644
--- a/framework-t/Android.bp
+++ b/framework-t/Android.bp
@@ -156,7 +156,7 @@
         "//frameworks/opt/telephony/tests/telephonytests",
         "//packages/modules/CaptivePortalLogin/tests",
         "//packages/modules/Connectivity/Tethering/tests:__subpackages__",
-        "//packages/modules/Connectivity/nearby/tests:__subpackages__",
+        "//packages/modules/Connectivity/nearby:__subpackages__",
         "//packages/modules/Connectivity/tests:__subpackages__",
         "//packages/modules/IPsec/tests/iketests",
         "//packages/modules/NetworkStack/tests:__subpackages__",
diff --git a/framework-t/src/android/net/nsd/NsdManager.java b/framework-t/src/android/net/nsd/NsdManager.java
index b4f8897..36808cf 100644
--- a/framework-t/src/android/net/nsd/NsdManager.java
+++ b/framework-t/src/android/net/nsd/NsdManager.java
@@ -1312,7 +1312,7 @@
      * before registering other callbacks. Upon failure to register a callback for example if
      * it's a duplicated registration, the application is notified through
      * {@link ServiceInfoCallback#onServiceInfoCallbackRegistrationFailed} with
-     * {@link #FAILURE_BAD_PARAMETERS} or {@link #FAILURE_ALREADY_ACTIVE}.
+     * {@link #FAILURE_BAD_PARAMETERS}.
      *
      * @param serviceInfo the service to receive updates for
      * @param executor Executor to run callbacks with
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 7cef58b..381a18a 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -2521,7 +2521,7 @@
     @RequiresPermission(android.Manifest.permission.PACKET_KEEPALIVE_OFFLOAD)
     public @NonNull SocketKeepalive createSocketKeepalive(@NonNull Network network,
             @NonNull Socket socket,
-            @NonNull Executor executor,
+            @NonNull @CallbackExecutor Executor executor,
             @NonNull Callback callback) {
         ParcelFileDescriptor dup;
         try {
@@ -5494,9 +5494,9 @@
      * @return {@code uid} if the connection is found and the app has permission to observe it
      *     (e.g., if it is associated with the calling VPN app's VpnService tunnel) or {@link
      *     android.os.Process#INVALID_UID} if the connection is not found.
-     * @throws {@link SecurityException} if the caller is not the active VpnService for the current
+     * @throws SecurityException if the caller is not the active VpnService for the current
      *     user.
-     * @throws {@link IllegalArgumentException} if an unsupported protocol is requested.
+     * @throws IllegalArgumentException if an unsupported protocol is requested.
      */
     public int getConnectionOwnerUid(
             int protocol, @NonNull InetSocketAddress local, @NonNull InetSocketAddress remote) {
diff --git a/nearby/halfsheet/Android.bp b/nearby/halfsheet/Android.bp
index 2d0d327..8011dc6 100644
--- a/nearby/halfsheet/Android.bp
+++ b/nearby/halfsheet/Android.bp
@@ -27,7 +27,7 @@
     certificate: ":com.android.nearby.halfsheetcertificate",
     libs: [
         "framework-bluetooth",
-        "framework-connectivity-t",
+        "framework-connectivity-t.impl",
         "nearby-service-string",
     ],
     static_libs: [
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 619b64d..8b70a94 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -84,7 +84,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -265,6 +264,35 @@
         }
     }
 
+    private class ServiceInfoListener extends MdnsListener {
+
+        ServiceInfoListener(int clientId, int transactionId, @NonNull NsdServiceInfo reqServiceInfo,
+                @NonNull String listenServiceType) {
+            super(clientId, transactionId, reqServiceInfo, listenServiceType);
+        }
+
+        @Override
+        public void onServiceFound(@NonNull MdnsServiceInfo serviceInfo) {
+            mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
+                    NsdManager.SERVICE_UPDATED,
+                    new MdnsEvent(mClientId, mReqServiceInfo.getServiceType(), serviceInfo));
+        }
+
+        @Override
+        public void onServiceUpdated(@NonNull MdnsServiceInfo serviceInfo) {
+            mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
+                    NsdManager.SERVICE_UPDATED,
+                    new MdnsEvent(mClientId, mReqServiceInfo.getServiceType(), serviceInfo));
+        }
+
+        @Override
+        public void onServiceRemoved(@NonNull MdnsServiceInfo serviceInfo) {
+            mNsdStateMachine.sendMessage(MDNS_DISCOVERY_MANAGER_EVENT, mTransactionId,
+                    NsdManager.SERVICE_UPDATED_LOST,
+                    new MdnsEvent(mClientId, mReqServiceInfo.getServiceType(), serviceInfo));
+        }
+    }
+
     /**
      * Data class of mdns service callback information.
      */
@@ -521,12 +549,6 @@
                 mIdToClientInfoMap.put(globalId, clientInfo);
             }
 
-            private void clearRegisteredServiceInfo(ClientInfo clientInfo) {
-                clientInfo.mRegisteredService = null;
-                clientInfo.mClientIdForServiceUpdates = 0;
-            }
-
-
             /**
              * Truncate a service name to up to 63 UTF-8 bytes.
              *
@@ -832,7 +854,7 @@
                         }
                         break;
                     }
-                    case NsdManager.REGISTER_SERVICE_CALLBACK:
+                    case NsdManager.REGISTER_SERVICE_CALLBACK: {
                         if (DBG) Log.d(TAG, "Register a service callback");
                         args = (ListenerArgs) msg.obj;
                         clientInfo = mClients.get(args.connector);
@@ -844,23 +866,29 @@
                             break;
                         }
 
-                        if (clientInfo.mRegisteredService != null) {
-                            clientInfo.onServiceInfoCallbackRegistrationFailed(
-                                    clientId, NsdManager.FAILURE_ALREADY_ACTIVE);
+                        final NsdServiceInfo info = args.serviceInfo;
+                        id = getUniqueId();
+                        final String serviceType = constructServiceType(info.getServiceType());
+                        if (serviceType == null) {
+                            clientInfo.onServiceInfoCallbackRegistrationFailed(clientId,
+                                    NsdManager.FAILURE_BAD_PARAMETERS);
                             break;
                         }
+                        final String resolveServiceType = serviceType + ".local";
 
-                        maybeStartDaemon();
-                        id = getUniqueId();
-                        if (resolveService(id, args.serviceInfo)) {
-                            clientInfo.mRegisteredService = new NsdServiceInfo();
-                            clientInfo.mClientIdForServiceUpdates = clientId;
-                            storeLegacyRequestMap(clientId, id, clientInfo, msg.what);
-                        } else {
-                            clientInfo.onServiceInfoCallbackRegistrationFailed(
-                                    clientId, NsdManager.FAILURE_BAD_PARAMETERS);
-                        }
+                        maybeStartMonitoringSockets();
+                        final MdnsListener listener =
+                                new ServiceInfoListener(clientId, id, info, resolveServiceType);
+                        final MdnsSearchOptions options = MdnsSearchOptions.newBuilder()
+                                .setNetwork(info.getNetwork())
+                                .setIsPassiveMode(true)
+                                .setResolveInstanceName(info.getServiceName())
+                                .build();
+                        mMdnsDiscoveryManager.registerListener(
+                                resolveServiceType, listener, options);
+                        storeDiscoveryManagerRequestMap(clientId, id, listener, clientInfo);
                         break;
+                    }
                     case NsdManager.UNREGISTER_SERVICE_CALLBACK: {
                         if (DBG) Log.d(TAG, "Unregister a service callback");
                         args = (ListenerArgs) msg.obj;
@@ -875,17 +903,16 @@
 
                         final ClientRequest request = clientInfo.mClientRequests.get(clientId);
                         if (request == null) {
-                            Log.e(TAG, "Unknown client request in STOP_RESOLUTION");
+                            Log.e(TAG, "Unknown client request in UNREGISTER_SERVICE_CALLBACK");
                             break;
                         }
                         id = request.mGlobalId;
-                        removeRequestMap(clientId, id, clientInfo);
-                        if (stopResolveService(id)) {
+                        if (request instanceof DiscoveryManagerRequest) {
+                            stopDiscoveryManagerRequest(request, clientId, id, clientInfo);
                             clientInfo.onServiceInfoCallbackUnregistered(clientId);
                         } else {
-                            Log.e(TAG, "Failed to unregister service info callback");
+                            loge("Unregister failed with non-DiscoveryManagerRequest.");
                         }
-                        clearRegisteredServiceInfo(clientInfo);
                         break;
                     }
                     case MDNS_SERVICE_EVENT:
@@ -904,19 +931,6 @@
                 return HANDLED;
             }
 
-            private void notifyResolveFailedResult(boolean isListenedToUpdates, int clientId,
-                    ClientInfo clientInfo, int error) {
-                if (isListenedToUpdates) {
-                    clientInfo.onServiceInfoCallbackRegistrationFailed(clientId, error);
-                    clearRegisteredServiceInfo(clientInfo);
-                } else {
-                    // The resolve API always returned FAILURE_INTERNAL_ERROR on error; keep it
-                    // for backwards compatibility.
-                    clientInfo.onResolveServiceFailed(clientId, NsdManager.FAILURE_INTERNAL_ERROR);
-                    clientInfo.mResolvedService = null;
-                }
-            }
-
             private boolean handleMDnsServiceEvent(int code, int id, Object obj) {
                 NsdServiceInfo servInfo;
                 ClientInfo clientInfo = mIdToClientInfoMap.get(id);
@@ -973,8 +987,6 @@
                         // found services on the same interface index and their network at the time
                         setServiceNetworkForCallback(servInfo, lostNetId, info.interfaceIdx);
                         clientInfo.onServiceLost(clientId, servInfo);
-                        // TODO: also support registered service lost when not discovering
-                        clientInfo.maybeNotifyRegisteredServiceLost(servInfo);
                         break;
                     }
                     case IMDnsEventListener.SERVICE_DISCOVERY_FAILED:
@@ -1011,11 +1023,7 @@
                         String rest = fullName.substring(index);
                         String type = rest.replace(".local.", "");
 
-                        final boolean isListenedToUpdates =
-                                clientId == clientInfo.mClientIdForServiceUpdates;
-                        final NsdServiceInfo serviceInfo = isListenedToUpdates
-                                ? clientInfo.mRegisteredService : clientInfo.mResolvedService;
-
+                        final NsdServiceInfo serviceInfo = clientInfo.mResolvedService;
                         serviceInfo.setServiceName(name);
                         serviceInfo.setServiceType(type);
                         serviceInfo.setPort(info.port);
@@ -1030,8 +1038,9 @@
                             storeLegacyRequestMap(clientId, id2, clientInfo,
                                     NsdManager.RESOLVE_SERVICE);
                         } else {
-                            notifyResolveFailedResult(isListenedToUpdates, clientId, clientInfo,
-                                    NsdManager.FAILURE_BAD_PARAMETERS);
+                            clientInfo.onResolveServiceFailed(
+                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                            clientInfo.mResolvedService = null;
                         }
                         break;
                     }
@@ -1039,17 +1048,17 @@
                         /* NNN resolveId errorCode */
                         stopResolveService(id);
                         removeRequestMap(clientId, id, clientInfo);
-                        notifyResolveFailedResult(
-                                clientId == clientInfo.mClientIdForServiceUpdates,
-                                clientId, clientInfo, NsdManager.FAILURE_BAD_PARAMETERS);
+                        clientInfo.onResolveServiceFailed(
+                                clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                        clientInfo.mResolvedService = null;
                         break;
                     case IMDnsEventListener.SERVICE_GET_ADDR_FAILED:
                         /* NNN resolveId errorCode */
                         stopGetAddrInfo(id);
                         removeRequestMap(clientId, id, clientInfo);
-                        notifyResolveFailedResult(
-                                clientId == clientInfo.mClientIdForServiceUpdates,
-                                clientId, clientInfo, NsdManager.FAILURE_BAD_PARAMETERS);
+                        clientInfo.onResolveServiceFailed(
+                                clientId, NsdManager.FAILURE_INTERNAL_ERROR);
+                        clientInfo.mResolvedService = null;
                         break;
                     case IMDnsEventListener.SERVICE_GET_ADDR_SUCCESS: {
                         /* NNN resolveId hostname ttl addr interfaceIdx netId */
@@ -1066,38 +1075,19 @@
                         // If the resolved service is on an interface without a network, consider it
                         // as a failure: it would not be usable by apps as they would need
                         // privileged permissions.
-                        if (clientId == clientInfo.mClientIdForServiceUpdates) {
-                            if (netId != NETID_UNSET && serviceHost != null) {
-                                setServiceNetworkForCallback(clientInfo.mRegisteredService,
-                                        netId, info.interfaceIdx);
-                                final List<InetAddress> addresses =
-                                        clientInfo.mRegisteredService.getHostAddresses();
-                                addresses.add(serviceHost);
-                                clientInfo.mRegisteredService.setHostAddresses(addresses);
-                                clientInfo.onServiceUpdated(
-                                        clientId, clientInfo.mRegisteredService);
-                            } else {
-                                stopGetAddrInfo(id);
-                                removeRequestMap(clientId, id, clientInfo);
-                                clearRegisteredServiceInfo(clientInfo);
-                                clientInfo.onServiceInfoCallbackRegistrationFailed(
-                                        clientId, NsdManager.FAILURE_BAD_PARAMETERS);
-                            }
+                        if (netId != NETID_UNSET && serviceHost != null) {
+                            clientInfo.mResolvedService.setHost(serviceHost);
+                            setServiceNetworkForCallback(clientInfo.mResolvedService,
+                                    netId, info.interfaceIdx);
+                            clientInfo.onResolveServiceSucceeded(
+                                    clientId, clientInfo.mResolvedService);
                         } else {
-                            if (netId != NETID_UNSET && serviceHost != null) {
-                                clientInfo.mResolvedService.setHost(serviceHost);
-                                setServiceNetworkForCallback(clientInfo.mResolvedService,
-                                        netId, info.interfaceIdx);
-                                clientInfo.onResolveServiceSucceeded(
-                                        clientId, clientInfo.mResolvedService);
-                            } else {
-                                clientInfo.onResolveServiceFailed(
-                                        clientId, NsdManager.FAILURE_INTERNAL_ERROR);
-                            }
-                            stopGetAddrInfo(id);
-                            removeRequestMap(clientId, id, clientInfo);
-                            clientInfo.mResolvedService = null;
+                            clientInfo.onResolveServiceFailed(
+                                    clientId, NsdManager.FAILURE_INTERNAL_ERROR);
                         }
+                        stopGetAddrInfo(id);
+                        removeRequestMap(clientId, id, clientInfo);
+                        clientInfo.mResolvedService = null;
                         break;
                     }
                     default:
@@ -1195,6 +1185,42 @@
                         stopDiscoveryManagerRequest(request, clientId, transactionId, clientInfo);
                         break;
                     }
+                    case NsdManager.SERVICE_UPDATED: {
+                        final MdnsServiceInfo serviceInfo = event.mMdnsServiceInfo;
+                        info.setPort(serviceInfo.getPort());
+
+                        Map<String, String> attrs = serviceInfo.getAttributes();
+                        for (Map.Entry<String, String> kv : attrs.entrySet()) {
+                            final String key = kv.getKey();
+                            try {
+                                info.setAttribute(key, serviceInfo.getAttributeAsBytes(key));
+                            } catch (IllegalArgumentException e) {
+                                Log.e(TAG, "Invalid attribute", e);
+                            }
+                        }
+
+                        final List<InetAddress> addresses = new ArrayList<>();
+                        for (String ipv4Address : serviceInfo.getIpv4Addresses()) {
+                            try {
+                                addresses.add(InetAddresses.parseNumericAddress(ipv4Address));
+                            } catch (IllegalArgumentException e) {
+                                Log.wtf(TAG, "Invalid ipv4 address", e);
+                            }
+                        }
+                        for (String ipv6Address : serviceInfo.getIpv6Addresses()) {
+                            try {
+                                addresses.add(InetAddresses.parseNumericAddress(ipv6Address));
+                            } catch (IllegalArgumentException e) {
+                                Log.wtf(TAG, "Invalid ipv6 address", e);
+                            }
+                        }
+                        info.setHostAddresses(addresses);
+                        clientInfo.onServiceUpdated(clientId, info);
+                        break;
+                    }
+                    case NsdManager.SERVICE_UPDATED_LOST:
+                        clientInfo.onServiceUpdatedLost(clientId);
+                        break;
                     default:
                         return false;
                 }
@@ -1765,11 +1791,6 @@
         // The target SDK of this client < Build.VERSION_CODES.S
         private boolean mIsPreSClient = false;
 
-        /*** The service that is registered to listen to its updates */
-        private NsdServiceInfo mRegisteredService;
-        /*** The client id that listen to updates */
-        private int mClientIdForServiceUpdates;
-
         private ClientInfo(INsdManagerCallback cb) {
             mCb = cb;
             if (DBG) Log.d(TAG, "New client");
@@ -1864,18 +1885,6 @@
             return -1;
         }
 
-        private void maybeNotifyRegisteredServiceLost(@NonNull NsdServiceInfo info) {
-            if (mRegisteredService == null) return;
-            if (!Objects.equals(mRegisteredService.getServiceName(), info.getServiceName())) return;
-            // Resolved services have a leading dot appended at the beginning of their type, but in
-            // discovered info it's at the end
-            if (!Objects.equals(
-                    mRegisteredService.getServiceType() + ".", "." + info.getServiceType())) {
-                return;
-            }
-            onServiceUpdatedLost(mClientIdForServiceUpdates);
-        }
-
         void onDiscoverServicesStarted(int listenerKey, NsdServiceInfo info) {
             try {
                 mCb.onDiscoverServicesStarted(listenerKey, info);
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
index 67059e7..fb8af8d 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsDiscoveryManager.java
@@ -16,19 +16,25 @@
 
 package com.android.server.connectivity.mdns;
 
+import static com.android.server.connectivity.mdns.MdnsSocketProvider.isNetworkMatched;
+
 import android.Manifest.permission;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.net.Network;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.Log;
+import android.util.Pair;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.connectivity.mdns.util.MdnsLogger;
 
 import java.io.IOException;
-import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
 
 /**
  * This class keeps tracking the set of registered {@link MdnsServiceBrowserListener} instances, and
@@ -42,12 +48,62 @@
     private final ExecutorProvider executorProvider;
     private final MdnsSocketClientBase socketClient;
 
-    private final Map<String, MdnsServiceTypeClient> serviceTypeClients = new ArrayMap<>();
+    @GuardedBy("this")
+    @NonNull private final PerNetworkServiceTypeClients perNetworkServiceTypeClients;
+
+    private static class PerNetworkServiceTypeClients {
+        private final ArrayMap<Pair<String, Network>, MdnsServiceTypeClient> clients =
+                new ArrayMap<>();
+
+        public void put(@NonNull String serviceType, @Nullable Network network,
+                @NonNull MdnsServiceTypeClient client) {
+            final Pair<String, Network> perNetworkServiceType = new Pair<>(serviceType, network);
+            clients.put(perNetworkServiceType, client);
+        }
+
+        @Nullable
+        public MdnsServiceTypeClient get(@NonNull String serviceType, @Nullable Network network) {
+            final Pair<String, Network> perNetworkServiceType = new Pair<>(serviceType, network);
+            return clients.getOrDefault(perNetworkServiceType, null);
+        }
+
+        public List<MdnsServiceTypeClient> getByServiceType(@NonNull String serviceType) {
+            final List<MdnsServiceTypeClient> list = new ArrayList<>();
+            for (int i = 0; i < clients.size(); i++) {
+                final Pair<String, Network> perNetworkServiceType = clients.keyAt(i);
+                if (serviceType.equals(perNetworkServiceType.first)) {
+                    list.add(clients.valueAt(i));
+                }
+            }
+            return list;
+        }
+
+        public List<MdnsServiceTypeClient> getByMatchingNetwork(@Nullable Network network) {
+            final List<MdnsServiceTypeClient> list = new ArrayList<>();
+            for (int i = 0; i < clients.size(); i++) {
+                final Pair<String, Network> perNetworkServiceType = clients.keyAt(i);
+                if (isNetworkMatched(network, perNetworkServiceType.second)) {
+                    list.add(clients.valueAt(i));
+                }
+            }
+            return list;
+        }
+
+        public void remove(@NonNull MdnsServiceTypeClient client) {
+            final int index = clients.indexOfValue(client);
+            clients.removeAt(index);
+        }
+
+        public boolean isEmpty() {
+            return clients.isEmpty();
+        }
+    }
 
     public MdnsDiscoveryManager(@NonNull ExecutorProvider executorProvider,
             @NonNull MdnsSocketClientBase socketClient) {
         this.executorProvider = executorProvider;
         this.socketClient = socketClient;
+        perNetworkServiceTypeClients = new PerNetworkServiceTypeClients();
     }
 
     /**
@@ -67,7 +123,7 @@
         LOGGER.log(
                 "Registering listener for subtypes: %s",
                 TextUtils.join(",", searchOptions.getSubtypes()));
-        if (serviceTypeClients.isEmpty()) {
+        if (perNetworkServiceTypeClients.isEmpty()) {
             // First listener. Starts the socket client.
             try {
                 socketClient.startDiscovery();
@@ -77,16 +133,18 @@
             }
         }
         // Request the network for discovery.
-        socketClient.notifyNetworkRequested(listener, searchOptions.getNetwork());
-
-        // All listeners of the same service types shares the same MdnsServiceTypeClient.
-        MdnsServiceTypeClient serviceTypeClient = serviceTypeClients.get(serviceType);
-        if (serviceTypeClient == null) {
-            serviceTypeClient = createServiceTypeClient(serviceType);
-            serviceTypeClients.put(serviceType, serviceTypeClient);
-        }
-        // TODO(b/264634275): Wait for a socket to be created before sending packets.
-        serviceTypeClient.startSendAndReceive(listener, searchOptions);
+        socketClient.notifyNetworkRequested(listener, searchOptions.getNetwork(), network -> {
+            synchronized (this) {
+                // All listeners of the same service types shares the same MdnsServiceTypeClient.
+                MdnsServiceTypeClient serviceTypeClient =
+                        perNetworkServiceTypeClients.get(serviceType, network);
+                if (serviceTypeClient == null) {
+                    serviceTypeClient = createServiceTypeClient(serviceType, network);
+                    perNetworkServiceTypeClients.put(serviceType, network, serviceTypeClient);
+                }
+                serviceTypeClient.startSendAndReceive(listener, searchOptions);
+            }
+        });
     }
 
     /**
@@ -101,17 +159,21 @@
             @NonNull String serviceType, @NonNull MdnsServiceBrowserListener listener) {
         LOGGER.log("Unregistering listener for service type: %s", serviceType);
         if (DBG) Log.d(TAG, "Unregistering listener for serviceType:" + serviceType);
-        MdnsServiceTypeClient serviceTypeClient = serviceTypeClients.get(serviceType);
-        if (serviceTypeClient == null) {
+        final List<MdnsServiceTypeClient> serviceTypeClients =
+                perNetworkServiceTypeClients.getByServiceType(serviceType);
+        if (serviceTypeClients.isEmpty()) {
             return;
         }
-        if (serviceTypeClient.stopSendAndReceive(listener)) {
-            // No listener is registered for the service type anymore, remove it from the list of
-            // the service type clients.
-            serviceTypeClients.remove(serviceType);
-            if (serviceTypeClients.isEmpty()) {
-                // No discovery request. Stops the socket client.
-                socketClient.stopDiscovery();
+        for (int i = 0; i < serviceTypeClients.size(); i++) {
+            final MdnsServiceTypeClient serviceTypeClient = serviceTypeClients.get(i);
+            if (serviceTypeClient.stopSendAndReceive(listener)) {
+                // No listener is registered for the service type anymore, remove it from the list
+                // of the service type clients.
+                perNetworkServiceTypeClients.remove(serviceTypeClient);
+                if (perNetworkServiceTypeClients.isEmpty()) {
+                    // No discovery request. Stops the socket client.
+                    socketClient.stopDiscovery();
+                }
             }
         }
         // Unrequested the network.
@@ -121,22 +183,26 @@
     @Override
     public synchronized void onResponseReceived(@NonNull MdnsPacket packet,
             int interfaceIndex, Network network) {
-        for (MdnsServiceTypeClient serviceTypeClient : serviceTypeClients.values()) {
+        for (MdnsServiceTypeClient serviceTypeClient
+                : perNetworkServiceTypeClients.getByMatchingNetwork(network)) {
             serviceTypeClient.processResponse(packet, interfaceIndex, network);
         }
     }
 
     @Override
-    public synchronized void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode) {
-        for (MdnsServiceTypeClient serviceTypeClient : serviceTypeClients.values()) {
+    public synchronized void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode,
+            Network network) {
+        for (MdnsServiceTypeClient serviceTypeClient
+                : perNetworkServiceTypeClients.getByMatchingNetwork(network)) {
             serviceTypeClient.onFailedToParseMdnsResponse(receivedPacketNumber, errorCode);
         }
     }
 
     @VisibleForTesting
-    MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType) {
+    MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType,
+            @Nullable Network network) {
         return new MdnsServiceTypeClient(
                 serviceType, socketClient,
-                executorProvider.newServiceTypeClientSchedulerExecutor());
+                executorProvider.newServiceTypeClientSchedulerExecutor(), network);
     }
 }
\ No newline at end of file
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
index 93972d9..5254342 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClient.java
@@ -63,6 +63,12 @@
     }
 
     private class InterfaceSocketCallback implements MdnsSocketProvider.SocketCallback {
+        private final SocketCreationCallback mSocketCreationCallback;
+
+        InterfaceSocketCallback(SocketCreationCallback socketCreationCallback) {
+            mSocketCreationCallback = socketCreationCallback;
+        }
+
         @Override
         public void onSocketCreated(@NonNull Network network,
                 @NonNull MdnsInterfaceSocket socket, @NonNull List<LinkAddress> addresses) {
@@ -76,6 +82,7 @@
             }
             socket.addPacketHandler(handler);
             mActiveNetworkSockets.put(socket, network);
+            mSocketCreationCallback.onSocketCreated(network);
         }
 
         @Override
@@ -114,10 +121,11 @@
      * @param listener the listener for discovery.
      * @param network the target network for discovery. Null means discovery on all possible
      *                interfaces.
+     * @param socketCreationCallback the callback to notify socket creation.
      */
     @Override
     public void notifyNetworkRequested(@NonNull MdnsServiceBrowserListener listener,
-            @Nullable Network network) {
+            @Nullable Network network, @NonNull SocketCreationCallback socketCreationCallback) {
         ensureRunningOnHandlerThread(mHandler);
         InterfaceSocketCallback callback = mRequestedNetworks.get(listener);
         if (callback != null) {
@@ -125,7 +133,7 @@
         }
 
         if (DBG) Log.d(TAG, "notifyNetworkRequested: network=" + network);
-        callback = new InterfaceSocketCallback();
+        callback = new InterfaceSocketCallback(socketCreationCallback);
         mRequestedNetworks.put(listener, callback);
         mSocketProvider.requestSocket(network, callback);
     }
@@ -173,7 +181,7 @@
             if (e.code != MdnsResponseErrorCode.ERROR_NOT_RESPONSE_MESSAGE) {
                 Log.e(TAG, e.getMessage(), e);
                 if (mCallback != null) {
-                    mCallback.onFailedToParseMdnsResponse(packetNumber, e.code);
+                    mCallback.onFailedToParseMdnsResponse(packetNumber, e.code, network);
                 }
             }
             return;
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 df270bb..f87804b 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -57,6 +57,7 @@
     private final MdnsSocketClientBase socketClient;
     private final MdnsResponseDecoder responseDecoder;
     private final ScheduledExecutorService executor;
+    @Nullable private final Network network;
     private final Object lock = new Object();
     private final ArrayMap<MdnsServiceBrowserListener, MdnsSearchOptions> listeners =
             new ArrayMap<>();
@@ -88,8 +89,9 @@
     public MdnsServiceTypeClient(
             @NonNull String serviceType,
             @NonNull MdnsSocketClientBase socketClient,
-            @NonNull ScheduledExecutorService executor) {
-        this(serviceType, socketClient, executor, new MdnsResponseDecoder.Clock());
+            @NonNull ScheduledExecutorService executor,
+            @Nullable Network network) {
+        this(serviceType, socketClient, executor, new MdnsResponseDecoder.Clock(), network);
     }
 
     @VisibleForTesting
@@ -97,13 +99,15 @@
             @NonNull String serviceType,
             @NonNull MdnsSocketClientBase socketClient,
             @NonNull ScheduledExecutorService executor,
-            @NonNull MdnsResponseDecoder.Clock clock) {
+            @NonNull MdnsResponseDecoder.Clock clock,
+            @Nullable Network network) {
         this.serviceType = serviceType;
         this.socketClient = socketClient;
         this.executor = executor;
         this.serviceTypeLabels = TextUtils.split(serviceType, "\\.");
         this.responseDecoder = new MdnsResponseDecoder(clock, serviceTypeLabels);
         this.clock = clock;
+        this.network = network;
     }
 
     private static MdnsServiceInfo buildMdnsServiceInfoFromResponse(
@@ -204,7 +208,9 @@
      */
     public boolean stopSendAndReceive(@NonNull MdnsServiceBrowserListener listener) {
         synchronized (lock) {
-            listeners.remove(listener);
+            if (listeners.remove(listener) == null) {
+                return listeners.isEmpty();
+            }
             if (listeners.isEmpty() && requestTaskFuture != null) {
                 requestTaskFuture.cancel(true);
                 requestTaskFuture = null;
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
index c03e6aa..783b18a 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClient.java
@@ -423,7 +423,7 @@
             LOGGER.w(String.format("Error while decoding %s packet (%d): %d",
                     responseType, packetNumber, e.code));
             if (callback != null) {
-                callback.onFailedToParseMdnsResponse(packetNumber, e.code);
+                callback.onFailedToParseMdnsResponse(packetNumber, e.code, network);
             }
             return e.code;
         }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
index 796dc83..ebafc49 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketClientBase.java
@@ -64,7 +64,9 @@
 
     /*** Notify that the given network is requested for mdns discovery / resolution */
     default void notifyNetworkRequested(@NonNull MdnsServiceBrowserListener listener,
-            @Nullable Network network) { }
+            @Nullable Network network, @NonNull SocketCreationCallback socketCreationCallback) {
+        socketCreationCallback.onSocketCreated(network);
+    }
 
     /*** Notify that the network is unrequested */
     default void notifyNetworkUnrequested(@NonNull MdnsServiceBrowserListener listener) { }
@@ -76,6 +78,13 @@
                 @Nullable Network network);
 
         /*** Parse a mdns response failed */
-        void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode);
+        void onFailedToParseMdnsResponse(int receivedPacketNumber, int errorCode,
+                @Nullable Network network);
+    }
+
+    /*** Callback for requested socket creation  */
+    interface SocketCreationCallback {
+        /*** Notify requested socket is created */
+        void onSocketCreated(@Nullable Network network);
     }
 }
diff --git a/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java b/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
index 6ba0fda..805dd65 100644
--- a/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
+++ b/tests/cts/net/src/android/net/cts/Ikev2VpnTest.java
@@ -60,7 +60,11 @@
 
 import com.android.internal.util.HexDump;
 import com.android.networkstack.apishim.ConstantsShim;
+import com.android.networkstack.apishim.Ikev2VpnProfileBuilderShimImpl;
+import com.android.networkstack.apishim.Ikev2VpnProfileShimImpl;
 import com.android.networkstack.apishim.VpnManagerShimImpl;
+import com.android.networkstack.apishim.common.Ikev2VpnProfileBuilderShim;
+import com.android.networkstack.apishim.common.Ikev2VpnProfileShim;
 import com.android.networkstack.apishim.common.VpnManagerShim;
 import com.android.networkstack.apishim.common.VpnProfileStateShim;
 import com.android.testutils.DevSdkIgnoreRule;
@@ -223,17 +227,28 @@
     }
 
     private Ikev2VpnProfile buildIkev2VpnProfileCommon(
-            @NonNull Ikev2VpnProfile.Builder builder, boolean isRestrictedToTestNetworks,
-            boolean requiresValidation) throws Exception {
+            @NonNull Ikev2VpnProfileBuilderShim builderShim, boolean isRestrictedToTestNetworks,
+            boolean requiresValidation, boolean automaticIpVersionSelectionEnabled,
+            boolean automaticNattKeepaliveTimerEnabled) throws Exception {
 
-        builder.setBypassable(true)
+        builderShim.setBypassable(true)
                 .setAllowedAlgorithms(TEST_ALLOWED_ALGORITHMS)
                 .setProxy(TEST_PROXY_INFO)
                 .setMaxMtu(TEST_MTU)
                 .setMetered(false);
         if (TestUtils.shouldTestTApis()) {
-            builder.setRequiresInternetValidation(requiresValidation);
+            builderShim.setRequiresInternetValidation(requiresValidation);
         }
+
+        if (TestUtils.shouldTestUApis()) {
+            builderShim.setAutomaticIpVersionSelectionEnabled(automaticIpVersionSelectionEnabled);
+            builderShim.setAutomaticNattKeepaliveTimerEnabled(automaticNattKeepaliveTimerEnabled);
+        }
+
+        // Convert shim back to Ikev2VpnProfile.Builder since restrictToTestNetworks is a hidden
+        // method and is not defined in shims.
+        // TODO: replace it in alternative way to remove the hidden method usage
+        final Ikev2VpnProfile.Builder builder = (Ikev2VpnProfile.Builder) builderShim.getBuilder();
         if (isRestrictedToTestNetworks) {
             builder.restrictToTestNetworks();
         }
@@ -249,13 +264,16 @@
                         ? IkeSessionTestUtils.IKE_PARAMS_V6 : IkeSessionTestUtils.IKE_PARAMS_V4,
                         IkeSessionTestUtils.CHILD_PARAMS);
 
-        final Ikev2VpnProfile.Builder builder =
-                new Ikev2VpnProfile.Builder(params)
+        final Ikev2VpnProfileBuilderShim builderShim =
+                Ikev2VpnProfileBuilderShimImpl.newInstance(params)
                         .setRequiresInternetValidation(requiresValidation)
                         .setProxy(TEST_PROXY_INFO)
                         .setMaxMtu(TEST_MTU)
                         .setMetered(false);
-
+        // Convert shim back to Ikev2VpnProfile.Builder since restrictToTestNetworks is a hidden
+        // method and is not defined in shims.
+        // TODO: replace it in alternative way to remove the hidden method usage
+        final Ikev2VpnProfile.Builder builder = (Ikev2VpnProfile.Builder) builderShim.getBuilder();
         if (isRestrictedToTestNetworks) {
             builder.restrictToTestNetworks();
         }
@@ -263,31 +281,35 @@
     }
 
     private Ikev2VpnProfile buildIkev2VpnProfilePsk(@NonNull String remote,
-            boolean isRestrictedToTestNetworks, boolean requiresValidation) throws Exception {
-        final Ikev2VpnProfile.Builder builder =
-                new Ikev2VpnProfile.Builder(remote, TEST_IDENTITY).setAuthPsk(TEST_PSK);
+            boolean isRestrictedToTestNetworks, boolean requiresValidation)
+            throws Exception {
+        final Ikev2VpnProfileBuilderShim builder =
+                Ikev2VpnProfileBuilderShimImpl.newInstance(remote, TEST_IDENTITY)
+                        .setAuthPsk(TEST_PSK);
         return buildIkev2VpnProfileCommon(builder, isRestrictedToTestNetworks,
-                requiresValidation);
+                requiresValidation, false /* automaticIpVersionSelectionEnabled */,
+                false /* automaticNattKeepaliveTimerEnabled */);
     }
 
     private Ikev2VpnProfile buildIkev2VpnProfileUsernamePassword(boolean isRestrictedToTestNetworks)
             throws Exception {
-
-        final Ikev2VpnProfile.Builder builder =
-                new Ikev2VpnProfile.Builder(TEST_SERVER_ADDR_V6, TEST_IDENTITY)
+        final Ikev2VpnProfileBuilderShim builder =
+                Ikev2VpnProfileBuilderShimImpl.newInstance(TEST_SERVER_ADDR_V6, TEST_IDENTITY)
                         .setAuthUsernamePassword(TEST_USER, TEST_PASSWORD, mServerRootCa);
         return buildIkev2VpnProfileCommon(builder, isRestrictedToTestNetworks,
-                false /* requiresValidation */);
+                false /* requiresValidation */, false /* automaticIpVersionSelectionEnabled */,
+                false /* automaticNattKeepaliveTimerEnabled */);
     }
 
     private Ikev2VpnProfile buildIkev2VpnProfileDigitalSignature(boolean isRestrictedToTestNetworks)
             throws Exception {
-        final Ikev2VpnProfile.Builder builder =
-                new Ikev2VpnProfile.Builder(TEST_SERVER_ADDR_V6, TEST_IDENTITY)
+        final Ikev2VpnProfileBuilderShim builder =
+                Ikev2VpnProfileBuilderShimImpl.newInstance(TEST_SERVER_ADDR_V6, TEST_IDENTITY)
                         .setAuthDigitalSignature(
                                 mUserCertKey.cert, mUserCertKey.key, mServerRootCa);
         return buildIkev2VpnProfileCommon(builder, isRestrictedToTestNetworks,
-                false /* requiresValidation */);
+                false /* requiresValidation */, false /* automaticIpVersionSelectionEnabled */,
+                false /* automaticNattKeepaliveTimerEnabled */);
     }
 
     private void checkBasicIkev2VpnProfile(@NonNull Ikev2VpnProfile profile) throws Exception {
@@ -687,6 +709,56 @@
                 true /* testSessionKey */, false /* testIkeTunConnParams */);
     }
 
+    @Test
+    public void testBuildIkev2VpnProfileWithAutomaticNattKeepaliveTimerEnabled() throws Exception {
+        // Cannot use @IgnoreUpTo(Build.VERSION_CODES.TIRAMISU) because this test also requires API
+        // 34 shims, and @IgnoreUpTo does not check that.
+        assumeTrue(TestUtils.shouldTestUApis());
+
+        final Ikev2VpnProfile profileWithDefaultValue = buildIkev2VpnProfilePsk(TEST_SERVER_ADDR_V6,
+                false /* isRestrictedToTestNetworks */, false /* requiresValidation */);
+        final Ikev2VpnProfileShim<Ikev2VpnProfile> shimWithDefaultValue =
+                Ikev2VpnProfileShimImpl.newInstance(profileWithDefaultValue);
+        assertFalse(shimWithDefaultValue.isAutomaticNattKeepaliveTimerEnabled());
+
+        final Ikev2VpnProfileBuilderShim builder =
+                Ikev2VpnProfileBuilderShimImpl.newInstance(TEST_SERVER_ADDR_V6, TEST_IDENTITY)
+                        .setAuthPsk(TEST_PSK);
+        final Ikev2VpnProfile profile = buildIkev2VpnProfileCommon(builder,
+                false /* isRestrictedToTestNetworks */,
+                false /* requiresValidation */,
+                false /* automaticIpVersionSelectionEnabled */,
+                true /* automaticNattKeepaliveTimerEnabled */);
+        final Ikev2VpnProfileShim<Ikev2VpnProfile> shim =
+                Ikev2VpnProfileShimImpl.newInstance(profile);
+        assertTrue(shim.isAutomaticNattKeepaliveTimerEnabled());
+    }
+
+    @Test
+    public void testBuildIkev2VpnProfileWithAutomaticIpVersionSelectionEnabled() throws Exception {
+        // Cannot use @IgnoreUpTo(Build.VERSION_CODES.TIRAMISU) because this test also requires API
+        // 34 shims, and @IgnoreUpTo does not check that.
+        assumeTrue(TestUtils.shouldTestUApis());
+
+        final Ikev2VpnProfile profileWithDefaultValue = buildIkev2VpnProfilePsk(TEST_SERVER_ADDR_V6,
+                false /* isRestrictedToTestNetworks */, false /* requiresValidation */);
+        final Ikev2VpnProfileShim<Ikev2VpnProfile> shimWithDefaultValue =
+                Ikev2VpnProfileShimImpl.newInstance(profileWithDefaultValue);
+        assertFalse(shimWithDefaultValue.isAutomaticIpVersionSelectionEnabled());
+
+        final Ikev2VpnProfileBuilderShim builder =
+                Ikev2VpnProfileBuilderShimImpl.newInstance(TEST_SERVER_ADDR_V6, TEST_IDENTITY)
+                        .setAuthPsk(TEST_PSK);
+        final Ikev2VpnProfile profile = buildIkev2VpnProfileCommon(builder,
+                false /* isRestrictedToTestNetworks */,
+                false /* requiresValidation */,
+                true /* automaticIpVersionSelectionEnabled */,
+                false /* automaticNattKeepaliveTimerEnabled */);
+        final Ikev2VpnProfileShim<Ikev2VpnProfile> shim =
+                Ikev2VpnProfileShimImpl.newInstance(profile);
+        assertTrue(shim.isAutomaticIpVersionSelectionEnabled());
+    }
+
     private static class CertificateAndKey {
         public final X509Certificate cert;
         public final PrivateKey key;
diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index de5e46f..408c546 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -879,7 +879,7 @@
 
     @Test
     fun testRegisterServiceInfoCallback() {
-        // This test requires shims supporting U+ APIs (NsdManager.subscribeService)
+        // This test requires shims supporting U+ APIs (NsdManager.registerServiceInfoCallback)
         assumeTrue(TestUtils.shouldTestUApis())
 
         // Ensure Wi-Fi network connected and get addresses
@@ -909,16 +909,14 @@
             val foundInfo = discoveryRecord.waitForServiceDiscovered(
                     serviceName, wifiNetwork)
 
-            // Subscribe to service and check the addresses are the same as Wi-Fi addresses
+            // Register service callback and check the addresses are the same as Wi-Fi addresses
             nsdShim.registerServiceInfoCallback(nsdManager, foundInfo, { it.run() }, cbRecord)
-            for (i in addresses.indices) {
-                val subscribeCb = cbRecord.expectCallback<ServiceUpdated>()
-                assertEquals(foundInfo.serviceName, subscribeCb.serviceInfo.serviceName)
-                val hostAddresses = subscribeCb.serviceInfo.hostAddresses
-                assertEquals(i + 1, hostAddresses.size)
-                for (hostAddress in hostAddresses) {
-                    assertTrue(addresses.contains(hostAddress))
-                }
+            val serviceInfoCb = cbRecord.expectCallback<ServiceUpdated>()
+            assertEquals(foundInfo.serviceName, serviceInfoCb.serviceInfo.serviceName)
+            val hostAddresses = serviceInfoCb.serviceInfo.hostAddresses
+            assertEquals(addresses.size, hostAddresses.size)
+            for (hostAddress in hostAddresses) {
+                assertTrue(addresses.contains(hostAddress))
             }
         } cleanupStep {
             nsdManager.unregisterService(registrationRecord)
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index 3dc5647..0b48e08 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -101,6 +101,7 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.LinkedList;
 import java.util.List;
@@ -704,119 +705,102 @@
     }
 
     private void verifyUpdatedServiceInfo(NsdServiceInfo info, String serviceName,
-            String serviceType, String address, int port, int interfaceIndex, Network network) {
+            String serviceType, List<InetAddress> address, int port, int interfaceIndex,
+            Network network) {
         assertEquals(serviceName, info.getServiceName());
         assertEquals(serviceType, info.getServiceType());
-        assertTrue(info.getHostAddresses().contains(parseNumericAddress(address)));
+        assertEquals(address, info.getHostAddresses());
         assertEquals(port, info.getPort());
         assertEquals(network, info.getNetwork());
         assertEquals(interfaceIndex, info.getInterfaceIndex());
     }
 
     @Test
-    public void testRegisterAndUnregisterServiceInfoCallback() throws RemoteException {
+    public void testRegisterAndUnregisterServiceInfoCallback() {
         final NsdManager client = connectClient(mService);
         final NsdServiceInfo request = new NsdServiceInfo(SERVICE_NAME, SERVICE_TYPE);
         final NsdManager.ServiceInfoCallback serviceInfoCallback = mock(
                 NsdManager.ServiceInfoCallback.class);
+        final String serviceTypeWithLocalDomain = SERVICE_TYPE + ".local";
+        final Network network = new Network(999);
+        request.setNetwork(network);
         client.registerServiceInfoCallback(request, Runnable::run, serviceInfoCallback);
         waitForIdle();
+        // Verify the registration callback start.
+        final ArgumentCaptor<MdnsServiceBrowserListener> listenerCaptor =
+                ArgumentCaptor.forClass(MdnsServiceBrowserListener.class);
+        verify(mSocketProvider).startMonitoringSockets();
+        verify(mDiscoveryManager).registerListener(eq(serviceTypeWithLocalDomain),
+                listenerCaptor.capture(), argThat(options -> network.equals(options.getNetwork())));
 
-        final IMDnsEventListener eventListener = getEventListener();
-        final ArgumentCaptor<Integer> resolvIdCaptor = ArgumentCaptor.forClass(Integer.class);
-        verify(mMockMDnsM).resolve(resolvIdCaptor.capture(), eq(SERVICE_NAME), eq(SERVICE_TYPE),
-                eq("local.") /* domain */, eq(IFACE_IDX_ANY));
-
-        // Resolve service successfully.
-        final ResolutionInfo resolutionInfo = new ResolutionInfo(
-                resolvIdCaptor.getValue(),
-                IMDnsEventListener.SERVICE_RESOLVED,
-                null /* serviceName */,
-                null /* serviceType */,
-                null /* domain */,
-                SERVICE_FULL_NAME,
-                DOMAIN_NAME,
+        final MdnsServiceBrowserListener listener = listenerCaptor.getValue();
+        final MdnsServiceInfo mdnsServiceInfo = new MdnsServiceInfo(
+                SERVICE_NAME,
+                serviceTypeWithLocalDomain.split("\\."),
+                List.of(), /* subtypes */
+                new String[]{"android", "local"}, /* hostName */
                 PORT,
-                new byte[0] /* txtRecord */,
-                IFACE_IDX_ANY);
-        doReturn(true).when(mMockMDnsM).getServiceAddress(anyInt(), any(), anyInt());
-        eventListener.onServiceResolutionStatus(resolutionInfo);
-        waitForIdle();
+                List.of(IPV4_ADDRESS),
+                List.of(IPV6_ADDRESS),
+                List.of() /* textStrings */,
+                List.of() /* textEntries */,
+                1234,
+                network);
 
-        final ArgumentCaptor<Integer> getAddrIdCaptor = ArgumentCaptor.forClass(Integer.class);
-        verify(mMockMDnsM).getServiceAddress(getAddrIdCaptor.capture(), eq(DOMAIN_NAME),
-                eq(IFACE_IDX_ANY));
-
-        // First address info
-        final String v4Address = "192.0.2.1";
-        final String v6Address = "2001:db8::";
-        final GetAddressInfo addressInfo1 = new GetAddressInfo(
-                getAddrIdCaptor.getValue(),
-                IMDnsEventListener.SERVICE_GET_ADDR_SUCCESS,
-                SERVICE_FULL_NAME,
-                v4Address,
-                IFACE_IDX_ANY,
-                999 /* netId */);
-        eventListener.onGettingServiceAddressStatus(addressInfo1);
-        waitForIdle();
-
+        // Verify onServiceFound callback
+        listener.onServiceFound(mdnsServiceInfo);
         final ArgumentCaptor<NsdServiceInfo> updateInfoCaptor =
                 ArgumentCaptor.forClass(NsdServiceInfo.class);
         verify(serviceInfoCallback, timeout(TIMEOUT_MS).times(1))
                 .onServiceUpdated(updateInfoCaptor.capture());
         verifyUpdatedServiceInfo(updateInfoCaptor.getAllValues().get(0) /* info */, SERVICE_NAME,
-                "." + SERVICE_TYPE, v4Address, PORT, IFACE_IDX_ANY, new Network(999));
+                SERVICE_TYPE,
+                List.of(parseNumericAddress(IPV4_ADDRESS), parseNumericAddress(IPV6_ADDRESS)),
+                PORT, IFACE_IDX_ANY, new Network(999));
 
-        // Second address info
-        final GetAddressInfo addressInfo2 = new GetAddressInfo(
-                getAddrIdCaptor.getValue(),
-                IMDnsEventListener.SERVICE_GET_ADDR_SUCCESS,
-                SERVICE_FULL_NAME,
-                v6Address,
-                IFACE_IDX_ANY,
-                999 /* netId */);
-        eventListener.onGettingServiceAddressStatus(addressInfo2);
-        waitForIdle();
+        // Service addresses changed.
+        final String v4Address = "192.0.2.1";
+        final String v6Address = "2001:db8::1";
+        final MdnsServiceInfo updatedServiceInfo = new MdnsServiceInfo(
+                SERVICE_NAME,
+                serviceTypeWithLocalDomain.split("\\."),
+                List.of(), /* subtypes */
+                new String[]{"android", "local"}, /* hostName */
+                PORT,
+                List.of(v4Address),
+                List.of(v6Address),
+                List.of() /* textStrings */,
+                List.of() /* textEntries */,
+                1234,
+                network);
 
+        // Verify onServiceUpdated callback.
+        listener.onServiceUpdated(updatedServiceInfo);
         verify(serviceInfoCallback, timeout(TIMEOUT_MS).times(2))
                 .onServiceUpdated(updateInfoCaptor.capture());
-        verifyUpdatedServiceInfo(updateInfoCaptor.getAllValues().get(1) /* info */, SERVICE_NAME,
-                "." + SERVICE_TYPE, v6Address, PORT, IFACE_IDX_ANY, new Network(999));
+        verifyUpdatedServiceInfo(updateInfoCaptor.getAllValues().get(2) /* info */, SERVICE_NAME,
+                SERVICE_TYPE,
+                List.of(parseNumericAddress(v4Address), parseNumericAddress(v6Address)),
+                PORT, IFACE_IDX_ANY, new Network(999));
 
+        // Verify service callback unregistration.
         client.unregisterServiceInfoCallback(serviceInfoCallback);
         waitForIdle();
-
         verify(serviceInfoCallback, timeout(TIMEOUT_MS)).onServiceInfoCallbackUnregistered();
     }
 
     @Test
-    public void testRegisterServiceCallbackFailed() throws Exception {
+    public void testRegisterServiceCallbackFailed() {
         final NsdManager client = connectClient(mService);
-        final NsdServiceInfo request = new NsdServiceInfo(SERVICE_NAME, SERVICE_TYPE);
-        final NsdManager.ServiceInfoCallback subscribeListener = mock(
+        final String invalidServiceType = "a_service";
+        final NsdServiceInfo request = new NsdServiceInfo(SERVICE_NAME, invalidServiceType);
+        final NsdManager.ServiceInfoCallback serviceInfoCallback = mock(
                 NsdManager.ServiceInfoCallback.class);
-        client.registerServiceInfoCallback(request, Runnable::run, subscribeListener);
+        client.registerServiceInfoCallback(request, Runnable::run, serviceInfoCallback);
         waitForIdle();
 
-        final IMDnsEventListener eventListener = getEventListener();
-        final ArgumentCaptor<Integer> resolvIdCaptor = ArgumentCaptor.forClass(Integer.class);
-        verify(mMockMDnsM).resolve(resolvIdCaptor.capture(), eq(SERVICE_NAME), eq(SERVICE_TYPE),
-                eq("local.") /* domain */, eq(IFACE_IDX_ANY));
-
-        // Fail to resolve service.
-        final ResolutionInfo resolutionFailedInfo = new ResolutionInfo(
-                resolvIdCaptor.getValue(),
-                IMDnsEventListener.SERVICE_RESOLUTION_FAILED,
-                null /* serviceName */,
-                null /* serviceType */,
-                null /* domain */,
-                null /* serviceFullName */,
-                null /* domainName */,
-                0 /* port */,
-                new byte[0] /* txtRecord */,
-                IFACE_IDX_ANY);
-        eventListener.onServiceResolutionStatus(resolutionFailedInfo);
-        verify(subscribeListener, timeout(TIMEOUT_MS))
+        // Fail to register service callback.
+        verify(serviceInfoCallback, timeout(TIMEOUT_MS))
                 .onServiceInfoCallbackRegistrationFailed(eq(FAILURE_BAD_PARAMETERS));
     }
 
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
index e6b8326..7e7e6a4 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsDiscoveryManagerTests.java
@@ -18,20 +18,25 @@
 
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
-import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.net.Network;
 import android.text.TextUtils;
+import android.util.Pair;
 
+import com.android.server.connectivity.mdns.MdnsSocketClientBase.SocketCreationCallback;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
 
 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;
 
@@ -48,6 +53,10 @@
 
     private static final String SERVICE_TYPE_1 = "_googlecast._tcp.local";
     private static final String SERVICE_TYPE_2 = "_test._tcp.local";
+    private static final Pair<String, Network> PER_NETWORK_SERVICE_TYPE_1 =
+            Pair.create(SERVICE_TYPE_1, null);
+    private static final Pair<String, Network> PER_NETWORK_SERVICE_TYPE_2 =
+            Pair.create(SERVICE_TYPE_2, null);
 
     @Mock private ExecutorProvider executorProvider;
     @Mock private MdnsSocketClientBase socketClient;
@@ -69,10 +78,13 @@
 
         discoveryManager = new MdnsDiscoveryManager(executorProvider, socketClient) {
                     @Override
-                    MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType) {
-                        if (serviceType.equals(SERVICE_TYPE_1)) {
+                    MdnsServiceTypeClient createServiceTypeClient(@NonNull String serviceType,
+                            @Nullable Network network) {
+                        final Pair<String, Network> perNetworkServiceType =
+                                Pair.create(serviceType, network);
+                        if (perNetworkServiceType.equals(PER_NETWORK_SERVICE_TYPE_1)) {
                             return mockServiceTypeClientOne;
-                        } else if (serviceType.equals(SERVICE_TYPE_2)) {
+                        } else if (perNetworkServiceType.equals(PER_NETWORK_SERVICE_TYPE_2)) {
                             return mockServiceTypeClientTwo;
                         }
                         return null;
@@ -80,13 +92,23 @@
                 };
     }
 
+    private void verifyListenerRegistration(String serviceType, MdnsServiceBrowserListener listener,
+            MdnsServiceTypeClient client) throws IOException {
+        final ArgumentCaptor<SocketCreationCallback> callbackCaptor =
+                ArgumentCaptor.forClass(SocketCreationCallback.class);
+        discoveryManager.registerListener(serviceType, listener,
+                MdnsSearchOptions.getDefaultOptions());
+        verify(socketClient).startDiscovery();
+        verify(socketClient).notifyNetworkRequested(
+                eq(listener), any(), callbackCaptor.capture());
+        final SocketCreationCallback callback = callbackCaptor.getValue();
+        callback.onSocketCreated(null /* network */);
+        verify(client).startSendAndReceive(listener, MdnsSearchOptions.getDefaultOptions());
+    }
+
     @Test
     public void registerListener_unregisterListener() throws IOException {
-        discoveryManager.registerListener(
-                SERVICE_TYPE_1, mockListenerOne, MdnsSearchOptions.getDefaultOptions());
-        verify(socketClient).startDiscovery();
-        verify(mockServiceTypeClientOne)
-                .startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        verifyListenerRegistration(SERVICE_TYPE_1, mockListenerOne, mockServiceTypeClientOne);
 
         when(mockServiceTypeClientOne.stopSendAndReceive(mockListenerOne)).thenReturn(true);
         discoveryManager.unregisterListener(SERVICE_TYPE_1, mockListenerOne);
@@ -96,40 +118,30 @@
 
     @Test
     public void registerMultipleListeners() throws IOException {
-        discoveryManager.registerListener(
-                SERVICE_TYPE_1, mockListenerOne, MdnsSearchOptions.getDefaultOptions());
-        verify(socketClient).startDiscovery();
-        verify(mockServiceTypeClientOne)
-                .startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
-
-        discoveryManager.registerListener(
-                SERVICE_TYPE_2, mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
-        verify(mockServiceTypeClientTwo)
-                .startSendAndReceive(mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+        verifyListenerRegistration(SERVICE_TYPE_1, mockListenerOne, mockServiceTypeClientOne);
+        verifyListenerRegistration(SERVICE_TYPE_2, mockListenerTwo, mockServiceTypeClientTwo);
     }
 
     @Test
-    public void onResponseReceived() {
-        discoveryManager.registerListener(
-                SERVICE_TYPE_1, mockListenerOne, MdnsSearchOptions.getDefaultOptions());
-        discoveryManager.registerListener(
-                SERVICE_TYPE_2, mockListenerTwo, MdnsSearchOptions.getDefaultOptions());
+    public void onResponseReceived() throws IOException {
+        verifyListenerRegistration(SERVICE_TYPE_1, mockListenerOne, mockServiceTypeClientOne);
+        verifyListenerRegistration(SERVICE_TYPE_2, mockListenerTwo, mockServiceTypeClientTwo);
 
         MdnsPacket responseForServiceTypeOne = createMdnsPacket(SERVICE_TYPE_1);
         final int ifIndex = 1;
-        final Network network = mock(Network.class);
-        discoveryManager.onResponseReceived(responseForServiceTypeOne, ifIndex, network);
+        discoveryManager.onResponseReceived(responseForServiceTypeOne, ifIndex, null /* network */);
         verify(mockServiceTypeClientOne).processResponse(responseForServiceTypeOne, ifIndex,
-                network);
+                null /* network */);
 
         MdnsPacket responseForServiceTypeTwo = createMdnsPacket(SERVICE_TYPE_2);
-        discoveryManager.onResponseReceived(responseForServiceTypeTwo, ifIndex, network);
+        discoveryManager.onResponseReceived(responseForServiceTypeTwo, ifIndex, null /* network */);
         verify(mockServiceTypeClientTwo).processResponse(responseForServiceTypeTwo, ifIndex,
-                network);
+                null /* network */);
 
         MdnsPacket responseForSubtype = createMdnsPacket("subtype._sub._googlecast._tcp.local");
-        discoveryManager.onResponseReceived(responseForSubtype, ifIndex, network);
-        verify(mockServiceTypeClientOne).processResponse(responseForSubtype, ifIndex, network);
+        discoveryManager.onResponseReceived(responseForSubtype, ifIndex, null /* network */);
+        verify(mockServiceTypeClientOne).processResponse(responseForSubtype, ifIndex,
+                null /* network */);
     }
 
     private MdnsPacket createMdnsPacket(String serviceType) {
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
index 1e322e4..90c43e5 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsMultinetworkSocketClientTest.java
@@ -62,6 +62,7 @@
     @Mock private MdnsInterfaceSocket mSocket;
     @Mock private MdnsServiceBrowserListener mListener;
     @Mock private MdnsSocketClientBase.Callback mCallback;
+    @Mock private MdnsSocketClientBase.SocketCreationCallback mSocketCreationCallback;
     private MdnsMultinetworkSocketClient mSocketClient;
     private Handler mHandler;
 
@@ -78,7 +79,8 @@
     private SocketCallback expectSocketCallback() {
         final ArgumentCaptor<SocketCallback> callbackCaptor =
                 ArgumentCaptor.forClass(SocketCallback.class);
-        mHandler.post(() -> mSocketClient.notifyNetworkRequested(mListener, mNetwork));
+        mHandler.post(() -> mSocketClient.notifyNetworkRequested(
+                mListener, mNetwork, mSocketCreationCallback));
         verify(mProvider, timeout(DEFAULT_TIMEOUT))
                 .requestSocket(eq(mNetwork), callbackCaptor.capture());
         return callbackCaptor.getValue();
@@ -107,6 +109,7 @@
         doReturn(createEmptyNetworkInterface()).when(mSocket).getInterface();
         // Notify socket created
         callback.onSocketCreated(mNetwork, mSocket, List.of());
+        verify(mSocketCreationCallback).onSocketCreated(mNetwork);
 
         // Send packet to IPv4 with target network and verify sending has been called.
         mSocketClient.sendMulticastPacket(ipv4Packet, mNetwork);
@@ -138,6 +141,7 @@
         doReturn(createEmptyNetworkInterface()).when(mSocket).getInterface();
         // Notify socket created
         callback.onSocketCreated(mNetwork, mSocket, List.of());
+        verify(mSocketCreationCallback).onSocketCreated(mNetwork);
 
         final ArgumentCaptor<PacketHandler> handlerCaptor =
                 ArgumentCaptor.forClass(PacketHandler.class);
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 d9fa10f..3e2ea35 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -166,7 +166,7 @@
 
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock) {
+                        mockDecoderClock, mockNetwork) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -701,7 +701,7 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock) {
+                        mockDecoderClock, mockNetwork) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -740,7 +740,7 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock) {
+                        mockDecoderClock, mockNetwork) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -773,7 +773,7 @@
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
-                        mockDecoderClock) {
+                        mockDecoderClock, mockNetwork) {
                     @Override
                     MdnsPacketWriter createMdnsPacketWriter() {
                         return mockPacketWriter;
@@ -897,7 +897,8 @@
 
     @Test
     public void testProcessResponse_Resolve() throws Exception {
-        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor);
+        client = new MdnsServiceTypeClient(
+                SERVICE_TYPE, mockSocketClient, currentThreadExecutor, mockNetwork);
 
         final String instanceName = "service-instance";
         final String[] hostname = new String[] { "testhost "};
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
index 9048686..abb1747 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketClientTests.java
@@ -412,7 +412,8 @@
         mdnsClient.startDiscovery();
 
         verify(mockCallback, timeout(TIMEOUT).atLeast(1))
-                .onFailedToParseMdnsResponse(anyInt(), eq(MdnsResponseErrorCode.ERROR_END_OF_FILE));
+                .onFailedToParseMdnsResponse(
+                        anyInt(), eq(MdnsResponseErrorCode.ERROR_END_OF_FILE), any());
 
         mdnsClient.stopDiscovery();
     }
@@ -433,7 +434,8 @@
         mdnsClient.startDiscovery();
 
         verify(mockCallback, timeout(TIMEOUT).atLeast(1))
-                .onFailedToParseMdnsResponse(1, MdnsResponseErrorCode.ERROR_END_OF_FILE);
+                .onFailedToParseMdnsResponse(
+                        eq(1), eq(MdnsResponseErrorCode.ERROR_END_OF_FILE), any());
 
         mdnsClient.stopDiscovery();
     }