Merge "Added additional logging to bpf handler"
diff --git a/TEST_MAPPING b/TEST_MAPPING
index c1d6a89..0326bf2 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -226,7 +226,12 @@
       ]
     },
     {
-        "name": "CtsTetheringTestLatestSdk[CaptivePortalLoginGoogle.apk+NetworkStackGoogle.apk+com.google.android.resolv.apex+com.google.android.tethering.apex]"
+      "name": "CtsTetheringTestLatestSdk[CaptivePortalLoginGoogle.apk+NetworkStackGoogle.apk+com.google.android.resolv.apex+com.google.android.tethering.apex]",
+      "options": [
+        {
+          "exclude-annotation": "com.android.testutils.NetworkStackModuleTest"
+        }
+      ]
     }
   ],
   "mainline-postsubmit": [
@@ -236,8 +241,13 @@
       "keywords": ["sim"]
     },
     {
-        "name": "CtsTetheringTestLatestSdk[CaptivePortalLoginGoogle.apk+NetworkStackGoogle.apk+com.google.android.resolv.apex+com.google.android.tethering.apex]",
-        "keywords": ["sim"]
+      "name": "CtsTetheringTestLatestSdk[CaptivePortalLoginGoogle.apk+NetworkStackGoogle.apk+com.google.android.resolv.apex+com.google.android.tethering.apex]",
+      "keywords": ["sim"],
+      "options": [
+        {
+          "exclude-annotation": "com.android.testutils.NetworkStackModuleTest"
+        }
+      ]
     }
   ],
   "imports": [
diff --git a/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java b/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
index 4ee5c42..5d57aa5 100644
--- a/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
+++ b/Tethering/tests/integration/src/android/net/EthernetTetheringTest.java
@@ -65,6 +65,7 @@
 import com.android.net.module.util.structs.UdpHeader;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
+import com.android.testutils.NetworkStackModuleTest;
 import com.android.testutils.TapPacketReader;
 
 import org.junit.Rule;
@@ -860,6 +861,8 @@
             (byte) 0x00, (byte) 0x08, (byte) 0x3a, (byte) 0xdf
     };
 
+    // This test requires the update in NetworkStackModule(See b/269692093).
+    @NetworkStackModuleTest
     @Test
     public void testTetherZeroLengthDhcpPacket() throws Exception {
         final TetheringTester tester = initTetheringTester(toList(TEST_IP4_ADDR),
diff --git a/netd/Android.bp b/netd/Android.bp
index 473460d..4325d89 100644
--- a/netd/Android.bp
+++ b/netd/Android.bp
@@ -35,6 +35,9 @@
         "BpfHandler.cpp",
         "NetdUpdatable.cpp",
     ],
+    static_libs: [
+        "libmodules-utils-build",
+    ],
     shared_libs: [
         "libbase",
         "liblog",
diff --git a/netd/BpfHandler.cpp b/netd/BpfHandler.cpp
index 6c25d1b..3984249 100644
--- a/netd/BpfHandler.cpp
+++ b/netd/BpfHandler.cpp
@@ -22,6 +22,7 @@
 #include <inttypes.h>
 
 #include <android-base/unique_fd.h>
+#include <android-modules-utils/sdk_level.h>
 #include <bpf/WaitForProgsLoaded.h>
 #include <log/log.h>
 #include <netdutils/UidConstants.h>
@@ -75,9 +76,11 @@
 }
 
 static Status initPrograms(const char* cg2_path) {
+    if (modules::sdklevel::IsAtLeastU() && !!strcmp(cg2_path, "/sys/fs/cgroup")) abort();
+
     unique_fd cg_fd(open(cg2_path, O_DIRECTORY | O_RDONLY | O_CLOEXEC));
     if (cg_fd == -1) {
-        int ret = errno;
+        const int ret = errno;
         ALOGE("Failed to open the cgroup directory: %s", strerror(ret));
         return statusFromErrno(ret, "Open the cgroup directory failed");
     }
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index c47c572..47a1022 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -619,6 +619,7 @@
                             final MdnsSearchOptions.Builder optionsBuilder =
                                     MdnsSearchOptions.newBuilder()
                                             .setNetwork(info.getNetwork())
+                                            .setRemoveExpiredService(true)
                                             .setIsPassiveMode(true);
                             if (typeAndSubtype.second != null) {
                                 // The parsing ensures subtype starts with an underscore.
@@ -813,6 +814,7 @@
                                     .setNetwork(info.getNetwork())
                                     .setIsPassiveMode(true)
                                     .setResolveInstanceName(info.getServiceName())
+                                    .setRemoveExpiredService(true)
                                     .build();
                             mMdnsDiscoveryManager.registerListener(
                                     resolveServiceType, listener, options);
@@ -906,6 +908,7 @@
                                 .setNetwork(info.getNetwork())
                                 .setIsPassiveMode(true)
                                 .setResolveInstanceName(info.getServiceName())
+                                .setRemoveExpiredService(true)
                                 .build();
                         mMdnsDiscoveryManager.registerListener(
                                 resolveServiceType, listener, options);
diff --git a/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java b/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
index 866ecba..84faf12 100644
--- a/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
+++ b/service-t/src/com/android/server/connectivity/mdns/EnqueueMdnsQueryCallable.java
@@ -24,6 +24,7 @@
 import android.util.Pair;
 
 import com.android.server.connectivity.mdns.util.MdnsLogger;
+import com.android.server.connectivity.mdns.util.MdnsUtils;
 
 import java.io.IOException;
 import java.lang.ref.WeakReference;
@@ -75,6 +76,8 @@
     private final boolean sendDiscoveryQueries;
     @NonNull
     private final List<MdnsResponse> servicesToResolve;
+    @NonNull
+    private final MdnsResponseDecoder.Clock clock;
 
     EnqueueMdnsQueryCallable(
             @NonNull MdnsSocketClientBase requestSender,
@@ -85,7 +88,8 @@
             int transactionId,
             @Nullable Network network,
             boolean sendDiscoveryQueries,
-            @NonNull Collection<MdnsResponse> servicesToResolve) {
+            @NonNull Collection<MdnsResponse> servicesToResolve,
+            @NonNull MdnsResponseDecoder.Clock clock) {
         weakRequestSender = new WeakReference<>(requestSender);
         this.packetWriter = packetWriter;
         serviceTypeLabels = TextUtils.split(serviceType, "\\.");
@@ -95,6 +99,7 @@
         this.network = network;
         this.sendDiscoveryQueries = sendDiscoveryQueries;
         this.servicesToResolve = new ArrayList<>(servicesToResolve);
+        this.clock = clock;
     }
 
     // Incompatible return type for override of Callable#call().
@@ -119,22 +124,24 @@
 
             // List of (name, type) to query
             final ArrayList<Pair<String[], Integer>> missingKnownAnswerRecords = new ArrayList<>();
+            final long now = clock.elapsedRealtime();
             for (MdnsResponse response : servicesToResolve) {
-                // TODO: also send queries to renew record TTL (as per RFC6762 7.1 no need to query
-                // if remaining TTL is more than half the original one, so send the queries if half
-                // the TTL has passed).
-                if (response.isComplete()) continue;
                 final String[] serviceName = response.getServiceName();
                 if (serviceName == null) continue;
-                if (!response.hasTextRecord()) {
+                if (!response.hasTextRecord() || MdnsUtils.isRecordRenewalNeeded(
+                        response.getTextRecord(), now)) {
                     missingKnownAnswerRecords.add(new Pair<>(serviceName, MdnsRecord.TYPE_TXT));
                 }
-                if (!response.hasServiceRecord()) {
+                if (!response.hasServiceRecord() || MdnsUtils.isRecordRenewalNeeded(
+                        response.getServiceRecord(), now)) {
                     missingKnownAnswerRecords.add(new Pair<>(serviceName, MdnsRecord.TYPE_SRV));
                     // The hostname is not yet known, so queries for address records will be sent
                     // the next time the EnqueueMdnsQueryCallable is enqueued if the reply does not
                     // contain them. In practice, advertisers should include the address records
                     // when queried for SRV, although it's not a MUST requirement (RFC6763 12.2).
+                    // TODO: Figure out how to renew the A/AAAA record. Usually A/AAAA record will
+                    //  be included in the response to the SRV record so in high chances there is
+                    //  no need to renew them individually.
                 } else if (!response.hasInet4AddressRecord() && !response.hasInet6AddressRecord()) {
                     final String[] host = response.getServiceRecord().getServiceHost();
                     missingKnownAnswerRecords.add(new Pair<>(host, MdnsRecord.TYPE_A));
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java b/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
index 75c7e6c..761c477 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsConfigs.java
@@ -94,10 +94,6 @@
         return false;
     }
 
-    public static boolean allowSearchOptionsToRemoveExpiredService() {
-        return false;
-    }
-
     public static boolean allowNetworkInterfaceIndexPropagation() {
         return true;
     }
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 c2c0db2..809750d 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsServiceTypeClient.java
@@ -66,9 +66,6 @@
     private final Map<String, MdnsResponse> instanceNameToResponse = new HashMap<>();
     private final boolean removeServiceAfterTtlExpires =
             MdnsConfigs.removeServiceAfterTtlExpires();
-    private final boolean allowSearchOptionsToRemoveExpiredService =
-            MdnsConfigs.allowSearchOptionsToRemoveExpiredService();
-
     private final MdnsResponseDecoder.Clock clock;
 
     @Nullable private MdnsSearchOptions searchOptions;
@@ -374,9 +371,7 @@
         if (removeServiceAfterTtlExpires) {
             return true;
         }
-        return allowSearchOptionsToRemoveExpiredService
-                && searchOptions != null
-                && searchOptions.removeExpiredService();
+        return searchOptions != null && searchOptions.removeExpiredService();
     }
 
     @VisibleForTesting
@@ -537,7 +532,8 @@
                                 config.transactionId,
                                 config.network,
                                 sendDiscoveryQueries,
-                                servicesToResolve)
+                                servicesToResolve,
+                                clock)
                                 .call();
             } catch (RuntimeException e) {
                 sharedLog.e(String.format("Failed to run EnqueueMdnsQueryCallable for subtype: %s",
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java
index 2fa1ae4..dc09bef 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsSocketProvider.java
@@ -25,7 +25,10 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.net.LinkAddress;
@@ -35,6 +38,9 @@
 import android.net.NetworkRequest;
 import android.net.TetheringManager;
 import android.net.TetheringManager.TetheringEventCallback;
+import android.net.wifi.p2p.WifiP2pGroup;
+import android.net.wifi.p2p.WifiP2pInfo;
+import android.net.wifi.p2p.WifiP2pManager;
 import android.os.Handler;
 import android.os.Looper;
 import android.util.ArrayMap;
@@ -51,6 +57,7 @@
 import java.net.SocketException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * The {@link MdnsSocketProvider} manages the multiple sockets for mDns.
@@ -92,6 +99,60 @@
     private final byte[] mPacketReadBuffer = new byte[READ_BUFFER_SIZE];
     private boolean mMonitoringSockets = false;
     private boolean mRequestStop = false;
+    private String mWifiP2pTetherInterface = null;
+
+    private final BroadcastReceiver mIntentReceiver = new BroadcastReceiver() {
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            final String newP2pIface = getWifiP2pInterface(intent);
+
+            if (!mMonitoringSockets || !hasAllNetworksRequest()) {
+                mWifiP2pTetherInterface = newP2pIface;
+                return;
+            }
+
+            // If already serving from the correct interface, nothing to do.
+            if (Objects.equals(mWifiP2pTetherInterface, newP2pIface)) return;
+
+            if (mWifiP2pTetherInterface != null) {
+                if (newP2pIface != null) {
+                    Log.wtf(TAG, "Wifi p2p interface is changed from " + mWifiP2pTetherInterface
+                            + " to " + newP2pIface + " without null broadcast");
+                }
+                // Remove the socket.
+                removeTetherInterfaceSocket(mWifiP2pTetherInterface);
+            }
+
+            // Update mWifiP2pTetherInterface
+            mWifiP2pTetherInterface = newP2pIface;
+
+            // Check whether the socket for wifi p2p interface is created or not.
+            final boolean socketAlreadyExists = mTetherInterfaceSockets.get(newP2pIface) != null;
+            if (newP2pIface != null && !socketAlreadyExists) {
+                // Create a socket for wifi p2p interface.
+                final int ifaceIndex =
+                        mDependencies.getNetworkInterfaceIndexByName(newP2pIface);
+                createSocket(LOCAL_NET, createLPForTetheredInterface(newP2pIface, ifaceIndex));
+            }
+        }
+    };
+
+    @Nullable
+    private static String getWifiP2pInterface(final Intent intent) {
+        final WifiP2pGroup group =
+                intent.getParcelableExtra(WifiP2pManager.EXTRA_WIFI_P2P_GROUP);
+        final WifiP2pInfo p2pInfo =
+                intent.getParcelableExtra(WifiP2pManager.EXTRA_WIFI_P2P_INFO);
+        if (group == null || p2pInfo == null) {
+            return null;
+        }
+
+        if (!p2pInfo.groupFormed) {
+            return null;
+        } else {
+            return group.getInterface();
+        }
+    }
 
     public MdnsSocketProvider(@NonNull Context context, @NonNull Looper looper,
             @NonNull SharedLog sharedLog) {
@@ -138,6 +199,18 @@
 
         mSocketNetlinkMonitor = mDependencies.createSocketNetlinkMonitor(mHandler,
                 mSharedLog.forSubComponent("NetlinkMonitor"), new NetLinkMessageProcessor());
+
+        // Register a intent receiver to listen wifi p2p interface changes.
+        // Note: The wifi p2p interface change is only notified via
+        // TetheringEventCallback#onLocalOnlyInterfacesChanged if the device is the wifi p2p group
+        // owner. In this case, MdnsSocketProvider will receive duplicate interface changes and must
+        // ignore the later notification because the socket has already been created. There is only
+        // one notification from the wifi p2p connection change intent if the device is not the wifi
+        // p2p group owner.
+        final IntentFilter intentFilter =
+                new IntentFilter(WifiP2pManager.WIFI_P2P_CONNECTION_CHANGED_ACTION);
+        mContext.registerReceiver(
+                mIntentReceiver, intentFilter, null /* broadcastPermission */, mHandler);
     }
 
     /**
@@ -376,17 +449,28 @@
         if (!hasAllNetworksRequest()) {
             // Currently, the network for tethering can not be requested, so the sockets for
             // tethering are only created if there is a request for all networks (interfaces).
-            // Therefore, this change can skip if there is no such request.
+            // Therefore, only update the interface list and skip this change if no such request.
             if (DBG) {
                 Log.d(TAG, "Ignore tether interfaces change. There is no request for all"
                         + " networks.");
             }
+            current.clear();
+            current.addAll(updated);
             return;
         }
 
         final CompareResult<String> interfaceDiff = new CompareResult<>(
                 current, updated);
         for (String name : interfaceDiff.added) {
+            // Check if a socket has been created for the interface
+            final SocketInfo socketInfo = mTetherInterfaceSockets.get(name);
+            if (socketInfo != null) {
+                if (DBG) {
+                    mSharedLog.i("Socket is existed for interface:" + name);
+                }
+                continue;
+            }
+
             int ifaceIndex = mDependencies.getNetworkInterfaceIndexByName(name);
             createSocket(LOCAL_NET, createLPForTetheredInterface(name, ifaceIndex));
         }
@@ -580,6 +664,11 @@
             for (String tetheredInterface : mTetheredInterfaces) {
                 retrieveAndNotifySocketFromInterface(tetheredInterface, cb);
             }
+
+            if (mWifiP2pTetherInterface != null
+                    && !mLocalOnlyInterfaces.contains(mWifiP2pTetherInterface)) {
+                retrieveAndNotifySocketFromInterface(mWifiP2pTetherInterface, cb);
+            }
         } else {
             retrieveAndNotifySocketFromNetwork(network, cb);
         }
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 bc94869..eb12b9a 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
@@ -152,4 +152,15 @@
         encoder.encode(CharBuffer.wrap(originalName), out, true /* endOfInput */);
         return new String(out.array(), 0, out.position(), utf8);
     }
+
+    /**
+     * Checks if the MdnsRecord needs to be renewed or not.
+     *
+     * <p>As per RFC6762 7.1 no need to query if remaining TTL is more than half the original one,
+     * so send the queries if half the TTL has passed.
+     */
+    public static boolean isRecordRenewalNeeded(@NonNull MdnsRecord mdnsRecord, final long now) {
+        return mdnsRecord.getTtl() > 0
+                && mdnsRecord.getRemainingTTL(now) <= mdnsRecord.getTtl() / 2;
+    }
 }
\ No newline at end of file
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index c95295c..fad6aaa 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -317,6 +317,7 @@
 import java.io.InterruptedIOException;
 import java.io.PrintWriter;
 import java.io.Writer;
+import java.lang.IllegalArgumentException;
 import java.net.Inet4Address;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
@@ -441,6 +442,8 @@
 
     private final Context mContext;
     private final ConnectivityResources mResources;
+    private final int mWakeUpMark;
+    private final int mWakeUpMask;
     // The Context is created for UserHandle.ALL.
     private final Context mUserAllContext;
     private final Dependencies mDeps;
@@ -1610,6 +1613,29 @@
         mCellularRadioTimesharingCapable =
                 mResources.get().getBoolean(R.bool.config_cellular_radio_timesharing_capable);
 
+        int mark = mResources.get().getInteger(R.integer.config_networkWakeupPacketMark);
+        int mask = mResources.get().getInteger(R.integer.config_networkWakeupPacketMask);
+
+        if (SdkLevel.isAtLeastU()) {
+            // U+ default value of both mark & mask, this is the top bit of the skb->mark,
+            // see //system/netd/include/FwMark.h union Fwmark, field ingress_cpu_wakeup
+            final int defaultUMarkMask = 0x80000000;  // u32
+
+            if ((mark == 0) || (mask == 0)) {
+                // simply treat unset/disabled as the default U value
+                mark = defaultUMarkMask;
+                mask = defaultUMarkMask;
+            }
+            if ((mark != defaultUMarkMask) || (mask != defaultUMarkMask)) {
+                // invalid device overlay settings
+                throw new IllegalArgumentException(
+                        "Bad config_networkWakeupPacketMark/Mask " + mark + "/" + mask);
+            }
+        }
+
+        mWakeUpMark = mark;
+        mWakeUpMask = mask;
+
         mNetd = netd;
         mBpfNetMaps = mDeps.getBpfNetMaps(mContext, netd);
         mHandlerThread = mDeps.makeHandlerThread();
@@ -8088,21 +8114,18 @@
             return;
         }
 
-        int mark = mResources.get().getInteger(R.integer.config_networkWakeupPacketMark);
-        int mask = mResources.get().getInteger(R.integer.config_networkWakeupPacketMask);
-
         // Mask/mark of zero will not detect anything interesting.
         // Don't install rules unless both values are nonzero.
-        if (mark == 0 || mask == 0) {
+        if (mWakeUpMark == 0 || mWakeUpMask == 0) {
             return;
         }
 
         final String prefix = makeNflogPrefix(iface, nai.network.getNetworkHandle());
         try {
             if (add) {
-                mNetd.wakeupAddInterface(iface, prefix, mark, mask);
+                mNetd.wakeupAddInterface(iface, prefix, mWakeUpMark, mWakeUpMask);
             } else {
-                mNetd.wakeupDelInterface(iface, prefix, mark, mask);
+                mNetd.wakeupDelInterface(iface, prefix, mWakeUpMark, mWakeUpMask);
             }
         } catch (Exception e) {
             loge("Exception modifying wakeup packet monitoring: " + e);
@@ -8620,10 +8643,18 @@
     }
 
     private void maybeCloseSockets(NetworkAgentInfo nai, Set<UidRange> ranges,
-            Set<Integer> exemptUids) {
+            UidRangeParcel[] uidRangeParcels, int[] exemptUids) {
         if (nai.isVPN() && !nai.networkAgentConfig.allowBypass) {
             try {
-                mDeps.destroyLiveTcpSockets(UidRange.toIntRanges(ranges), exemptUids);
+                if (mDeps.isAtLeastU()) {
+                    final Set<Integer> exemptUidSet = new ArraySet<>();
+                    for (final int uid: exemptUids) {
+                        exemptUidSet.add(uid);
+                    }
+                    mDeps.destroyLiveTcpSockets(UidRange.toIntRanges(ranges), exemptUidSet);
+                } else {
+                    mNetd.socketDestroy(uidRangeParcels, exemptUids);
+                }
             } catch (Exception e) {
                 loge("Exception in socket destroy: ", e);
             }
@@ -8631,16 +8662,16 @@
     }
 
     private void updateVpnUidRanges(boolean add, NetworkAgentInfo nai, Set<UidRange> uidRanges) {
-        final Set<Integer> exemptUids = new ArraySet<>();
+        int[] exemptUids = new int[2];
         // TODO: Excluding VPN_UID is necessary in order to not to kill the TCP connection used
         // by PPTP. Fix this by making Vpn set the owner UID to VPN_UID instead of system when
         // starting a legacy VPN, and remove VPN_UID here. (b/176542831)
-        exemptUids.add(VPN_UID);
-        exemptUids.add(nai.networkCapabilities.getOwnerUid());
+        exemptUids[0] = VPN_UID;
+        exemptUids[1] = nai.networkCapabilities.getOwnerUid();
         UidRangeParcel[] ranges = toUidRangeStableParcels(uidRanges);
 
         // Close sockets before modifying uid ranges so that RST packets can reach to the server.
-        maybeCloseSockets(nai, uidRanges, exemptUids);
+        maybeCloseSockets(nai, uidRanges, ranges, exemptUids);
         try {
             if (add) {
                 mNetd.networkAddUidRangesParcel(new NativeUidRangeConfig(
@@ -8654,7 +8685,7 @@
                     " on netId " + nai.network.netId + ". " + e);
         }
         // Close sockets that established connection while requesting netd.
-        maybeCloseSockets(nai, uidRanges, exemptUids);
+        maybeCloseSockets(nai, uidRanges, ranges, exemptUids);
     }
 
     private boolean isProxySetOnAnyDefaultNetwork() {
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index f8285ed..62d79a3 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -18,6 +18,7 @@
 
 import static android.net.SocketKeepalive.ERROR_INVALID_SOCKET;
 import static android.net.SocketKeepalive.MIN_INTERVAL_SEC;
+import static android.net.SocketKeepalive.SUCCESS;
 import static android.net.SocketKeepalive.SUCCESS_PAUSED;
 import static android.provider.DeviceConfig.NAMESPACE_TETHERING;
 import static android.system.OsConstants.AF_INET;
@@ -52,6 +53,7 @@
 import android.system.StructTimeval;
 import android.util.LocalLog;
 import android.util.Log;
+import android.util.Pair;
 import android.util.SparseArray;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -381,7 +383,11 @@
             return;
         }
         autoKi.mAutomaticOnOffState = STATE_ENABLED;
-        handleResumeKeepalive(newKi);
+        final int error = handleResumeKeepalive(newKi);
+        if (error != SUCCESS) {
+            // Failed to start the keepalive
+            cleanupAutoOnOffKeepalive(autoKi);
+        }
     }
 
     /**
@@ -402,7 +408,20 @@
      * Forward to KeepaliveTracker.
      */
     public void handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
-        mKeepaliveTracker.handleEventSocketKeepalive(nai, slot, reason);
+        if (mKeepaliveTracker.handleEventSocketKeepalive(nai, slot, reason)) return;
+
+        // The keepalive was stopped and so the autoKi should be cleaned up.
+        final AutomaticOnOffKeepalive autoKi =
+                CollectionUtils.findFirst(
+                        mAutomaticOnOffKeepalives, it -> it.match(nai.network(), slot));
+        if (autoKi == null) {
+            // This may occur when the autoKi gets cleaned up elsewhere (i.e
+            // handleCheckKeepalivesStillValid) while waiting for the network agent to
+            // start the keepalive and the network agent returns an error event.
+            Log.e(TAG, "Attempt cleanup on unknown network, slot");
+            return;
+        }
+        cleanupAutoOnOffKeepalive(autoKi);
     }
 
     /**
@@ -414,6 +433,9 @@
         final List<AutomaticOnOffKeepalive> matches =
                 CollectionUtils.filter(mAutomaticOnOffKeepalives, it -> it.mKi.getNai() == nai);
         for (final AutomaticOnOffKeepalive ki : matches) {
+            if (ki.mAutomaticOnOffState == STATE_SUSPENDED) {
+                mKeepaliveTracker.finalizePausedKeepalive(ki.mKi, reason);
+            }
             cleanupAutoOnOffKeepalive(ki);
         }
     }
@@ -425,9 +447,14 @@
      */
     public void handleStartKeepalive(Message message) {
         final AutomaticOnOffKeepalive autoKi = (AutomaticOnOffKeepalive) message.obj;
+        final int error = mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
+        if (error != SUCCESS) {
+            mEventLog.log("Failed to start keepalive " + autoKi.mCallback + " on "
+                    + autoKi.getNetwork() + " with error " + error);
+            return;
+        }
         mEventLog.log("Start keepalive " + autoKi.mCallback + " on " + autoKi.getNetwork());
         mKeepaliveStatsTracker.onStartKeepalive();
-        mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
 
         // Add automatic on/off request into list to track its life cycle.
         try {
@@ -443,10 +470,22 @@
         }
     }
 
-    private void handleResumeKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
+    /**
+     * Handle resume keepalive with the given KeepaliveInfo
+     *
+     * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
+     */
+    private int handleResumeKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
+        final int error = mKeepaliveTracker.handleStartKeepalive(ki);
+        if (error != SUCCESS) {
+            mEventLog.log("Failed to resume keepalive " + ki.mCallback + " on " + ki.mNai
+                    + " with error " + error);
+            return error;
+        }
         mKeepaliveStatsTracker.onResumeKeepalive();
-        mKeepaliveTracker.handleStartKeepalive(ki);
         mEventLog.log("Resumed successfully keepalive " + ki.mCallback + " on " + ki.mNai);
+
+        return SUCCESS;
     }
 
     private void handlePauseKeepalive(@NonNull final KeepaliveTracker.KeepaliveInfo ki) {
@@ -467,7 +506,7 @@
             final KeepaliveTracker.KeepaliveInfo ki = autoKi.mKi;
             mKeepaliveTracker.handleStopKeepalive(ki.getNai(), ki.getSlot(), reason);
         } else {
-            mKeepaliveTracker.finalizePausedKeepalive(autoKi.mKi);
+            mKeepaliveTracker.finalizePausedKeepalive(autoKi.mKi, reason);
         }
 
         cleanupAutoOnOffKeepalive(autoKi);
@@ -612,7 +651,22 @@
      * Forward to KeepaliveTracker.
      */
     public void handleCheckKeepalivesStillValid(NetworkAgentInfo nai) {
-        mKeepaliveTracker.handleCheckKeepalivesStillValid(nai);
+        ArrayList<Pair<AutomaticOnOffKeepalive, Integer>> invalidKeepalives = null;
+
+        for (final AutomaticOnOffKeepalive autoKi : mAutomaticOnOffKeepalives) {
+            if (!nai.equals(autoKi.mKi.mNai)) continue;
+            final int error = autoKi.mKi.isValid();
+            if (error != SUCCESS) {
+                if (invalidKeepalives == null) {
+                    invalidKeepalives = new ArrayList<>();
+                }
+                invalidKeepalives.add(Pair.create(autoKi, error));
+            }
+        }
+        if (invalidKeepalives == null) return;
+        for (final Pair<AutomaticOnOffKeepalive, Integer> keepaliveAndError : invalidKeepalives) {
+            handleStopKeepalive(keepaliveAndError.first, keepaliveAndError.second);
+        }
     }
 
     @VisibleForTesting
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index cc226ce..941b616 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -54,7 +54,6 @@
 import android.system.ErrnoException;
 import android.system.Os;
 import android.util.Log;
-import android.util.Pair;
 
 import com.android.connectivity.resources.R;
 import com.android.internal.annotations.VisibleForTesting;
@@ -337,7 +336,12 @@
             return SUCCESS;
         }
 
-        private int isValid() {
+        /**
+         * Checks if the keepalive info is valid to start.
+         *
+         * @return SUCCESS if the keepalive is valid and the error reason otherwise.
+         */
+        public int isValid() {
             synchronized (mNai) {
                 int error = checkInterval();
                 if (error == SUCCESS) error = checkLimit();
@@ -348,11 +352,17 @@
             }
         }
 
-        void start(int slot) {
+        /**
+         * Attempt to start the keepalive on the given slot.
+         *
+         * @param slot the slot to start the keepalive on.
+         * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
+         */
+        int start(int slot) {
             // BINDER_DIED can happen if the binder died before the KeepaliveInfo was created and
             // the constructor set the state to BINDER_DIED. If that's the case, the KI is already
             // cleaned up.
-            if (BINDER_DIED == mStartedState) return;
+            if (BINDER_DIED == mStartedState) return BINDER_DIED;
             mSlot = slot;
             int error = isValid();
             if (error == SUCCESS) {
@@ -368,7 +378,7 @@
                             mTcpController.startSocketMonitor(mFd, this, mSlot);
                         } catch (InvalidSocketException e) {
                             handleStopKeepalive(mNai, mSlot, ERROR_INVALID_SOCKET);
-                            return;
+                            return ERROR_INVALID_SOCKET;
                         }
                         final TcpKeepalivePacketData tcpData = (TcpKeepalivePacketData) mPacket;
                         mNai.onAddTcpKeepalivePacketFilter(slot, tcpData);
@@ -377,13 +387,14 @@
                         break;
                     default:
                         Log.wtf(TAG, "Starting keepalive with unknown type: " + mType);
-                        handleStopKeepalive(mNai, mSlot, error);
-                        return;
+                        handleStopKeepalive(mNai, mSlot, ERROR_UNSUPPORTED);
+                        return ERROR_UNSUPPORTED;
                 }
                 mStartedState = STARTING;
+                return SUCCESS;
             } else {
                 handleStopKeepalive(mNai, mSlot, error);
-                return;
+                return error;
             }
         }
 
@@ -444,6 +455,8 @@
             }
         }
 
+        // TODO: This does not clean up the autoKi in AutomaticOnOffKeepaliveTracker and it is not
+        // possible without a big refactor.
         void onFileDescriptorInitiatedStop(final int socketKeepaliveReason) {
             handleStopKeepalive(mNai, mSlot, socketKeepaliveReason);
         }
@@ -486,12 +499,15 @@
 
     /**
      * Handle start keepalives with the message.
+     *
+     * @param ki the keepalive to start.
+     * @return SUCCESS if the keepalive is successfully starting and the error reason otherwise.
      */
-    public void handleStartKeepalive(KeepaliveInfo ki) {
+    public int handleStartKeepalive(KeepaliveInfo ki) {
         NetworkAgentInfo nai = ki.getNai();
         int slot = findFirstFreeSlot(nai);
         mKeepalives.get(nai).put(slot, ki);
-        ki.start(slot);
+        return ki.start(slot);
     }
 
     public void handleStopAllKeepalives(NetworkAgentInfo nai, int reason) {
@@ -593,40 +609,33 @@
     /**
      * Finalize a paused keepalive.
      *
-     * This will simply send the onStopped() callback after checking that this keepalive is
-     * indeed paused.
+     * This will send the appropriate callback after checking that this keepalive is indeed paused.
      *
      * @param ki the keepalive to finalize
+     * @param reason the reason the keepalive is stopped
      */
-    public void finalizePausedKeepalive(@NonNull final KeepaliveInfo ki) {
+    public void finalizePausedKeepalive(@NonNull final KeepaliveInfo ki, int reason) {
         if (SUCCESS_PAUSED != ki.mStopReason) {
             throw new IllegalStateException("Keepalive is not paused");
         }
-        try {
-            ki.mCallback.onStopped();
-        } catch (RemoteException e) {
-            Log.w(TAG, "Discarded onStopped callback while finalizing paused keepalive");
+        if (reason == SUCCESS) {
+            try {
+                ki.mCallback.onStopped();
+            } catch (RemoteException e) {
+                Log.w(TAG, "Discarded onStopped callback while finalizing paused keepalive");
+            }
+        } else {
+            notifyErrorCallback(ki.mCallback, reason);
         }
     }
 
-    public void handleCheckKeepalivesStillValid(NetworkAgentInfo nai) {
-        HashMap <Integer, KeepaliveInfo> networkKeepalives = mKeepalives.get(nai);
-        if (networkKeepalives != null) {
-            ArrayList<Pair<Integer, Integer>> invalidKeepalives = new ArrayList<>();
-            for (int slot : networkKeepalives.keySet()) {
-                int error = networkKeepalives.get(slot).isValid();
-                if (error != SUCCESS) {
-                    invalidKeepalives.add(Pair.create(slot, error));
-                }
-            }
-            for (Pair<Integer, Integer> slotAndError: invalidKeepalives) {
-                handleStopKeepalive(nai, slotAndError.first, slotAndError.second);
-            }
-        }
-    }
-
-    /** Handle keepalive events from lower layer. */
-    public void handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
+    /**
+     * Handle keepalive events from lower layer.
+     *
+     * @return false if the event caused handleStopKeepalive to be called, i.e. the keepalive is
+     *     forced to stop. Otherwise, return true.
+     */
+    public boolean handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
         KeepaliveInfo ki = null;
         try {
             ki = mKeepalives.get(nai).get(slot);
@@ -634,7 +643,7 @@
         if (ki == null) {
             Log.e(TAG, "Event " + NetworkAgent.EVENT_SOCKET_KEEPALIVE + "," + slot + "," + reason
                     + " for unknown keepalive " + slot + " on " + nai.toShortString());
-            return;
+            return true;
         }
 
         // This can be called in a number of situations :
@@ -667,11 +676,13 @@
                     Log.w(TAG, "Discarded " + (ki.mResumed ? "onResumed" : "onStarted")
                             + " callback for slot " + slot);
                 }
+                return true;
             } else {
                 Log.d(TAG, "Failed to start keepalive " + slot + " on " + nai.toShortString()
                         + ": " + reason);
                 // The message indicated some error trying to start: do call handleStopKeepalive.
                 handleStopKeepalive(nai, slot, reason);
+                return false;
             }
         } else if (KeepaliveInfo.STOPPING == ki.mStartedState) {
             // The message indicated result of stopping : clean up keepalive slots.
@@ -679,9 +690,12 @@
                     + " stopped: " + reason);
             ki.mStartedState = KeepaliveInfo.NOT_STARTED;
             cleanupStoppedKeepalive(nai, slot);
+            return true;
         } else {
             Log.wtf(TAG, "Event " + NetworkAgent.EVENT_SOCKET_KEEPALIVE + "," + slot + "," + reason
                     + " for keepalive in wrong state: " + ki.toString());
+            // Although this is an unexpected event, the keepalive is not stopped here.
+            return true;
         }
     }
 
diff --git a/tests/common/Android.bp b/tests/common/Android.bp
index 8e47235..7b5c298 100644
--- a/tests/common/Android.bp
+++ b/tests/common/Android.bp
@@ -31,7 +31,7 @@
 //   (currently, CTS 10, 11, and 12).
 java_defaults {
     name: "ConnectivityTestsLatestSdkDefaults",
-    target_sdk_version: "33",
+    target_sdk_version: "34",
 }
 
 java_library {
diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index 9808137..6c411cf 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -548,8 +548,9 @@
         assertTrue(resolvedService.attributes.containsKey("nullBinaryDataAttr"))
         assertNull(resolvedService.attributes["nullBinaryDataAttr"])
         assertTrue(resolvedService.attributes.containsKey("emptyBinaryDataAttr"))
-        // TODO: change the check to target SDK U when this is what the code implements
-        if (isAtLeastU()) {
+        if (isAtLeastU() || CompatChanges.isChangeEnabled(
+                ConnectivityCompatChanges.ENABLE_PLATFORM_MDNS_BACKEND
+            )) {
             assertArrayEquals(byteArrayOf(), resolvedService.attributes["emptyBinaryDataAttr"])
         } else {
             assertNull(resolvedService.attributes["emptyBinaryDataAttr"])
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index c79c295..522ac8c 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -541,8 +541,7 @@
     private static final int TEST_PACKAGE_UID2 = 321;
     private static final int TEST_PACKAGE_UID3 = 456;
 
-    private static final int PACKET_WAKEUP_MASK = 0xffff0000;
-    private static final int PACKET_WAKEUP_MARK = 0x88880000;
+    private static final int PACKET_WAKEUP_MARK_MASK = 0x80000000;
 
     private static final String ALWAYS_ON_PACKAGE = "com.android.test.alwaysonvpn";
 
@@ -1924,9 +1923,9 @@
         doReturn(0).when(mResources).getInteger(R.integer.config_activelyPreferBadWifi);
         doReturn(true).when(mResources)
                 .getBoolean(R.bool.config_cellular_radio_timesharing_capable);
-        doReturn(PACKET_WAKEUP_MASK).when(mResources).getInteger(
+        doReturn(PACKET_WAKEUP_MARK_MASK).when(mResources).getInteger(
                 R.integer.config_networkWakeupPacketMask);
-        doReturn(PACKET_WAKEUP_MARK).when(mResources).getInteger(
+        doReturn(PACKET_WAKEUP_MARK_MASK).when(mResources).getInteger(
                 R.integer.config_networkWakeupPacketMark);
     }
 
@@ -12931,9 +12930,16 @@
             throws Exception {
         InOrder inOrder = inOrder(mMockNetd, mDestroySocketsWrapper);
         final Set<Integer> exemptUidSet = new ArraySet<>(List.of(exemptUid, Process.VPN_UID));
+        ArgumentCaptor<int[]> exemptUidCaptor = ArgumentCaptor.forClass(int[].class);
 
-        inOrder.verify(mDestroySocketsWrapper).destroyLiveTcpSockets(
-                UidRange.toIntRanges(vpnRanges), exemptUidSet);
+        if (mDeps.isAtLeastU()) {
+            inOrder.verify(mDestroySocketsWrapper).destroyLiveTcpSockets(
+                    UidRange.toIntRanges(vpnRanges), exemptUidSet);
+        } else {
+            inOrder.verify(mMockNetd).socketDestroy(eq(toUidRangeStableParcels(vpnRanges)),
+                    exemptUidCaptor.capture());
+            assertContainsExactly(exemptUidCaptor.getValue(), Process.VPN_UID, exemptUid);
+        }
 
         if (add) {
             inOrder.verify(mMockNetd, times(1)).networkAddUidRangesParcel(
@@ -12945,8 +12951,14 @@
                             toUidRangeStableParcels(vpnRanges), PREFERENCE_ORDER_VPN));
         }
 
-        inOrder.verify(mDestroySocketsWrapper).destroyLiveTcpSockets(
-                UidRange.toIntRanges(vpnRanges), exemptUidSet);
+        if (mDeps.isAtLeastU()) {
+            inOrder.verify(mDestroySocketsWrapper).destroyLiveTcpSockets(
+                    UidRange.toIntRanges(vpnRanges), exemptUidSet);
+        } else {
+            inOrder.verify(mMockNetd).socketDestroy(eq(toUidRangeStableParcels(vpnRanges)),
+                    exemptUidCaptor.capture());
+            assertContainsExactly(exemptUidCaptor.getValue(), Process.VPN_UID, exemptUid);
+        }
     }
 
     @Test
@@ -18224,8 +18236,8 @@
 
         final String expectedPrefix = makeNflogPrefix(WIFI_IFNAME,
                 mWiFiAgent.getNetwork().getNetworkHandle());
-        verify(mMockNetd).wakeupAddInterface(WIFI_IFNAME, expectedPrefix, PACKET_WAKEUP_MARK,
-                PACKET_WAKEUP_MASK);
+        verify(mMockNetd).wakeupAddInterface(WIFI_IFNAME, expectedPrefix, PACKET_WAKEUP_MARK_MASK,
+                PACKET_WAKEUP_MARK_MASK);
     }
 
     @Test
@@ -18238,8 +18250,8 @@
         if (mDeps.isAtLeastU()) {
             final String expectedPrefix = makeNflogPrefix(MOBILE_IFNAME,
                     mCellAgent.getNetwork().getNetworkHandle());
-            verify(mMockNetd).wakeupAddInterface(MOBILE_IFNAME, expectedPrefix, PACKET_WAKEUP_MARK,
-                    PACKET_WAKEUP_MASK);
+            verify(mMockNetd).wakeupAddInterface(MOBILE_IFNAME, expectedPrefix,
+                    PACKET_WAKEUP_MARK_MASK, PACKET_WAKEUP_MARK_MASK);
         } else {
             verify(mMockNetd, never()).wakeupAddInterface(eq(MOBILE_IFNAME), anyString(), anyInt(),
                     anyInt());
diff --git a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
index 9e0435d..608e6d8 100644
--- a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
@@ -20,9 +20,12 @@
 import static android.net.ConnectivityManager.TYPE_MOBILE;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 
+import static com.android.testutils.HandlerUtils.visibleOnHandlerThread;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
@@ -30,17 +33,20 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.longThat;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.ignoreStubs;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.app.AlarmManager;
 import android.content.Context;
 import android.content.res.Resources;
 import android.net.INetd;
 import android.net.ISocketKeepaliveCallback;
-import android.net.KeepalivePacketData;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.MarkMaskParcel;
@@ -48,6 +54,7 @@
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
+import android.net.SocketKeepalive;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Handler;
@@ -63,6 +70,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.connectivity.resources.R;
+import com.android.server.connectivity.AutomaticOnOffKeepaliveTracker.AutomaticOnOffKeepalive;
 import com.android.server.connectivity.KeepaliveTracker.KeepaliveInfo;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
@@ -70,6 +78,7 @@
 
 import libcore.util.HexEncoding;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -82,12 +91,15 @@
 import java.net.Socket;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.util.ArrayList;
+import java.util.List;
 
 @RunWith(DevSdkIgnoreRunner.class)
 @SmallTest
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.TIRAMISU)
 public class AutomaticOnOffKeepaliveTrackerTest {
     private static final String TAG = AutomaticOnOffKeepaliveTrackerTest.class.getSimpleName();
+    private static final int TEST_SLOT = 1;
     private static final int TEST_NETID = 0xA85;
     private static final int TEST_NETID_FWMARK = 0x0A85;
     private static final int OTHER_NETID = 0x1A85;
@@ -95,6 +107,8 @@
     private static final int TIMEOUT_MS = 30_000;
     private static final int MOCK_RESOURCE_ID = 5;
     private static final int TEST_KEEPALIVE_INTERVAL_SEC = 10;
+    private static final int TEST_KEEPALIVE_INVALID_INTERVAL_SEC = 9;
+
     private AutomaticOnOffKeepaliveTracker mAOOKeepaliveTracker;
     private HandlerThread mHandlerThread;
 
@@ -102,6 +116,8 @@
     @Mock AutomaticOnOffKeepaliveTracker.Dependencies mDependencies;
     @Mock Context mCtx;
     @Mock AlarmManager mAlarmManager;
+    @Mock NetworkAgentInfo mNai;
+
     TestKeepaliveTracker mKeepaliveTracker;
     AOOTestHandler mTestHandler;
 
@@ -202,6 +218,37 @@
     private static final byte[] TEST_RESPONSE_BYTES =
             HexEncoding.decode(TEST_RESPONSE_HEX.toCharArray(), false);
 
+    private static class TestKeepaliveInfo {
+        private static List<Socket> sOpenSockets = new ArrayList<>();
+
+        public static void closeAllSockets() throws Exception {
+            for (final Socket socket : sOpenSockets) {
+                socket.close();
+            }
+            sOpenSockets.clear();
+        }
+
+        public final Socket socket;
+        public final Binder binder;
+        public final FileDescriptor fd;
+        public final ISocketKeepaliveCallback socketKeepaliveCallback;
+        public final Network underpinnedNetwork;
+        public final NattKeepalivePacketData kpd;
+
+        TestKeepaliveInfo(NattKeepalivePacketData kpd) throws Exception {
+            this.kpd = kpd;
+            socket = new Socket();
+            socket.bind(null);
+            sOpenSockets.add(socket);
+            fd = socket.getFileDescriptor$();
+
+            binder = new Binder();
+            socketKeepaliveCallback = mock(ISocketKeepaliveCallback.class);
+            doReturn(binder).when(socketKeepaliveCallback).asBinder();
+            underpinnedNetwork = mock(Network.class);
+        }
+    }
+
     private class TestKeepaliveTracker extends KeepaliveTracker {
         private KeepaliveInfo mKi;
 
@@ -231,6 +278,14 @@
     public void setup() throws Exception {
         MockitoAnnotations.initMocks(this);
 
+        mNai.networkCapabilities =
+                new NetworkCapabilities.Builder().addTransportType(TRANSPORT_CELLULAR).build();
+        mNai.networkInfo = new NetworkInfo(TYPE_MOBILE, 0 /* subtype */, "LTE", "LTE");
+        mNai.networkInfo.setDetailedState(
+                NetworkInfo.DetailedState.CONNECTED, "test reason", "test extra info");
+        doReturn(new Network(TEST_NETID)).when(mNai).network();
+        mNai.linkProperties = new LinkProperties();
+
         doReturn(PERMISSION_GRANTED).when(mCtx).checkPermission(any() /* permission */,
                 anyInt() /* pid */, anyInt() /* uid */);
         ConnectivityResources.setResourcesContextForTest(mCtx);
@@ -255,6 +310,11 @@
                 new AutomaticOnOffKeepaliveTracker(mCtx, mTestHandler, mDependencies);
     }
 
+    @After
+    public void teardown() throws Exception {
+        TestKeepaliveInfo.closeAllSockets();
+    }
+
     private final class AOOTestHandler extends Handler {
         public AutomaticOnOffKeepaliveTracker.AutomaticOnOffKeepalive mLastAutoKi = null;
 
@@ -305,45 +365,70 @@
                 () -> assertFalse(mAOOKeepaliveTracker.isAnyTcpSocketConnected(TEST_NETID)));
     }
 
-    @Test
-    public void testAlarm() throws Exception {
+    private void triggerEventKeepalive(int slot, int reason) {
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleEventSocketKeepalive(mNai, slot, reason));
+    }
+
+    private TestKeepaliveInfo doStartNattKeepalive(int intervalSeconds) throws Exception {
         final InetAddress srcAddress = InetAddress.getByAddress(
                 new byte[] { (byte) 192, 0, 0, (byte) 129 });
         final int srcPort = 12345;
-        final InetAddress dstAddress = InetAddress.getByAddress(new byte[] { 8, 8, 8, 8});
+        final InetAddress dstAddress = InetAddress.getByAddress(new byte[] {8, 8, 8, 8});
         final int dstPort = 12345;
 
-        final NetworkAgentInfo nai = mock(NetworkAgentInfo.class);
-        nai.networkCapabilities = new NetworkCapabilities.Builder()
-                .addTransportType(TRANSPORT_CELLULAR).build();
-        nai.networkInfo = new NetworkInfo(TYPE_MOBILE, 0 /* subtype */, "LTE", "LTE");
-        nai.networkInfo.setDetailedState(NetworkInfo.DetailedState.CONNECTED, "test reason",
-                "test extra info");
-        nai.linkProperties = new LinkProperties();
-        nai.linkProperties.addLinkAddress(new LinkAddress(srcAddress, 24));
+        mNai.linkProperties.addLinkAddress(new LinkAddress(srcAddress, 24));
 
-        final Socket socket = new Socket();
-        socket.bind(null);
-        final FileDescriptor fd = socket.getFileDescriptor$();
-        final IBinder binder = new Binder();
-        final ISocketKeepaliveCallback cb = mock(ISocketKeepaliveCallback.class);
-        doReturn(binder).when(cb).asBinder();
-        final Network underpinnedNetwork = mock(Network.class);
-
-        final KeepalivePacketData kpd = new NattKeepalivePacketData(srcAddress, srcPort,
+        final NattKeepalivePacketData kpd = new NattKeepalivePacketData(srcAddress, srcPort,
                 dstAddress, dstPort, new byte[] {1});
-        final KeepaliveInfo ki = mKeepaliveTracker.new KeepaliveInfo(cb, nai, kpd,
-                TEST_KEEPALIVE_INTERVAL_SEC, KeepaliveInfo.TYPE_NATT, fd);
+
+        final TestKeepaliveInfo testInfo = new TestKeepaliveInfo(kpd);
+
+        final KeepaliveInfo ki = mKeepaliveTracker.new KeepaliveInfo(
+                testInfo.socketKeepaliveCallback, mNai, kpd, intervalSeconds,
+                KeepaliveInfo.TYPE_NATT, testInfo.fd);
         mKeepaliveTracker.setReturnedKeepaliveInfo(ki);
 
+        mAOOKeepaliveTracker.startNattKeepalive(mNai, testInfo.fd, intervalSeconds,
+                testInfo.socketKeepaliveCallback, srcAddress.toString(), srcPort,
+                dstAddress.toString(), dstPort, true /* automaticOnOffKeepalives */,
+                testInfo.underpinnedNetwork);
+        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+
+        return testInfo;
+    }
+
+    private TestKeepaliveInfo doStartNattKeepalive() throws Exception {
+        return doStartNattKeepalive(TEST_KEEPALIVE_INTERVAL_SEC);
+    }
+
+    private void doPauseKeepalive(AutomaticOnOffKeepalive autoKi) throws Exception {
+        setupResponseWithoutSocketExisting();
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleMonitorAutomaticKeepalive(autoKi, TEST_NETID));
+    }
+
+    private void doResumeKeepalive(AutomaticOnOffKeepalive autoKi) throws Exception {
+        setupResponseWithSocketExisting();
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleMonitorAutomaticKeepalive(autoKi, TEST_NETID));
+    }
+
+    private void doStopKeepalive(AutomaticOnOffKeepalive autoKi) throws Exception {
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleStopKeepalive(autoKi, SocketKeepalive.SUCCESS));
+    }
+
+    @Test
+    public void testAlarm() throws Exception {
         // Mock elapsed real time to verify the alarm timer.
         final long time = SystemClock.elapsedRealtime();
         doReturn(time).when(mDependencies).getElapsedRealtime();
-
-        mAOOKeepaliveTracker.startNattKeepalive(nai, fd, 10 /* intervalSeconds */, cb,
-                srcAddress.toString(), srcPort, dstAddress.toString(), dstPort,
-                true /* automaticOnOffKeepalives */, underpinnedNetwork);
-        HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
 
         final ArgumentCaptor<AlarmManager.OnAlarmListener> listenerCaptor =
                 ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
@@ -362,9 +447,8 @@
         HandlerUtils.waitForIdle(mTestHandler, TIMEOUT_MS);
 
         assertNotNull(mTestHandler.mLastAutoKi);
-        assertEquals(cb, mTestHandler.mLastAutoKi.getCallback());
-        assertEquals(underpinnedNetwork, mTestHandler.mLastAutoKi.getUnderpinnedNetwork());
-        socket.close();
+        assertEquals(testInfo.socketKeepaliveCallback, mTestHandler.mLastAutoKi.getCallback());
+        assertEquals(testInfo.underpinnedNetwork, mTestHandler.mLastAutoKi.getUnderpinnedNetwork());
     }
 
     private void setupResponseWithSocketExisting() throws Exception {
@@ -391,4 +475,301 @@
         buffer.order(ByteOrder.nativeOrder());
         return buffer;
     }
+
+    private AutomaticOnOffKeepalive getAutoKiForBinder(IBinder binder) {
+        return visibleOnHandlerThread(
+                mTestHandler, () -> mAOOKeepaliveTracker.getKeepaliveForBinder(binder));
+    }
+
+    private void checkAndProcessKeepaliveStart(final NattKeepalivePacketData kpd) throws Exception {
+        checkAndProcessKeepaliveStart(TEST_SLOT, kpd);
+    }
+
+    private void checkAndProcessKeepaliveStart(
+            int slot, final NattKeepalivePacketData kpd) throws Exception {
+        verify(mNai).onStartNattSocketKeepalive(slot, TEST_KEEPALIVE_INTERVAL_SEC, kpd);
+        verify(mNai).onAddNattKeepalivePacketFilter(slot, kpd);
+        triggerEventKeepalive(slot, SocketKeepalive.SUCCESS);
+    }
+
+    private void checkAndProcessKeepaliveStop() throws Exception {
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+    }
+
+    private void checkAndProcessKeepaliveStop(int slot) throws Exception {
+        verify(mNai).onStopSocketKeepalive(slot);
+        verify(mNai).onRemoveKeepalivePacketFilter(slot);
+        triggerEventKeepalive(slot, SocketKeepalive.SUCCESS);
+    }
+
+    @Test
+    public void testStartNattKeepalive_valid() throws Exception {
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+
+        final AutomaticOnOffKeepalive autoKi = getAutoKiForBinder(testInfo.binder);
+        assertNotNull(autoKi);
+        assertEquals(testInfo.socketKeepaliveCallback, autoKi.getCallback());
+
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testStartNattKeepalive_invalidInterval() throws Exception {
+        final TestKeepaliveInfo testInfo =
+                doStartNattKeepalive(TEST_KEEPALIVE_INVALID_INTERVAL_SEC);
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_INTERVAL);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testHandleEventSocketKeepalive_startingFailureHardwareError() throws Exception {
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+
+        verify(mNai)
+                .onStartNattSocketKeepalive(TEST_SLOT, TEST_KEEPALIVE_INTERVAL_SEC, testInfo.kpd);
+        verify(mNai).onAddNattKeepalivePacketFilter(TEST_SLOT, testInfo.kpd);
+        // Network agent returns an error, fails to start the keepalive.
+        triggerEventKeepalive(TEST_SLOT, SocketKeepalive.ERROR_HARDWARE_ERROR);
+
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_HARDWARE_ERROR);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testHandleCheckKeepalivesStillValid_linkPropertiesChanged() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        // Source address is removed from link properties by clearing.
+        mNai.linkProperties.clear();
+
+        // Check for valid keepalives
+        visibleOnHandlerThread(
+                mTestHandler, () -> mAOOKeepaliveTracker.handleCheckKeepalivesStillValid(mNai));
+
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_IP_ADDRESS);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testStopKeepalive() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doStopKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+        verify(testInfo.socketKeepaliveCallback).onStopped();
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testPauseKeepalive() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        // Pausing does not cleanup the autoKi
+        assertNotNull(getAutoKiForBinder(testInfo.binder));
+
+        clearInvocations(mNai);
+        doStopKeepalive(getAutoKiForBinder(testInfo.binder));
+        // The keepalive is already stopped.
+        verify(mNai, never()).onStopSocketKeepalive(TEST_SLOT);
+        verify(mNai, never()).onRemoveKeepalivePacketFilter(TEST_SLOT);
+
+        // Stopping while paused still calls onStopped.
+        verify(testInfo.socketKeepaliveCallback).onStopped();
+        // autoKi is cleaned up.
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testResumeKeepalive() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        clearInvocations(mNai);
+        doResumeKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        assertNotNull(getAutoKiForBinder(testInfo.binder));
+        verify(testInfo.socketKeepaliveCallback).onResumed();
+
+        doStopKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onStopped();
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testResumeKeepalive_invalidSourceAddress() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        mNai.linkProperties.clear();
+
+        clearInvocations(mNai);
+        doResumeKeepalive(getAutoKiForBinder(testInfo.binder));
+        verify(mNai, never()).onStartNattSocketKeepalive(anyInt(), anyInt(), any());
+        verify(mNai, never()).onAddNattKeepalivePacketFilter(anyInt(), any());
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_IP_ADDRESS);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testResumeKeepalive_startingFailureHardwareError() throws Exception {
+        // Successful start of NATT keepalive.
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo.kpd);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        doPauseKeepalive(getAutoKiForBinder(testInfo.binder));
+        checkAndProcessKeepaliveStop();
+        verify(testInfo.socketKeepaliveCallback).onPaused();
+
+        clearInvocations(mNai);
+        doResumeKeepalive(getAutoKiForBinder(testInfo.binder));
+
+        verify(mNai)
+                .onStartNattSocketKeepalive(TEST_SLOT, TEST_KEEPALIVE_INTERVAL_SEC, testInfo.kpd);
+        verify(mNai).onAddNattKeepalivePacketFilter(TEST_SLOT, testInfo.kpd);
+        // Network agent returns error on starting the keepalive.
+        triggerEventKeepalive(TEST_SLOT, SocketKeepalive.ERROR_HARDWARE_ERROR);
+
+        checkAndProcessKeepaliveStop();
+
+        assertNull(getAutoKiForBinder(testInfo.binder));
+        verify(testInfo.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_HARDWARE_ERROR);
+        verifyNoMoreInteractions(ignoreStubs(testInfo.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testStopAllKeepalives() throws Exception {
+        final TestKeepaliveInfo testInfo1 = doStartNattKeepalive();
+        final TestKeepaliveInfo testInfo2 = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(TEST_SLOT, testInfo1.kpd);
+        checkAndProcessKeepaliveStart(TEST_SLOT + 1, testInfo2.kpd);
+
+        verify(testInfo1.socketKeepaliveCallback).onStarted();
+        verify(testInfo2.socketKeepaliveCallback).onStarted();
+
+        // Pause the first keepalive
+        doPauseKeepalive(getAutoKiForBinder(testInfo1.binder));
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+        verify(testInfo1.socketKeepaliveCallback).onPaused();
+
+        visibleOnHandlerThread(
+                mTestHandler,
+                () -> mAOOKeepaliveTracker.handleStopAllKeepalives(
+                        mNai, SocketKeepalive.ERROR_INVALID_NETWORK));
+
+        // Note that checkAndProcessKeepaliveStop is not called since the network agent is assumed
+        // to be disconnected for a handleStopAllKeepalives call.
+        assertNull(getAutoKiForBinder(testInfo1.binder));
+        assertNull(getAutoKiForBinder(testInfo2.binder));
+
+        verify(testInfo1.socketKeepaliveCallback, never()).onStopped();
+        verify(testInfo2.socketKeepaliveCallback, never()).onStopped();
+        verify(testInfo1.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_NETWORK);
+        verify(testInfo2.socketKeepaliveCallback).onError(SocketKeepalive.ERROR_INVALID_NETWORK);
+
+        verifyNoMoreInteractions(ignoreStubs(testInfo1.socketKeepaliveCallback));
+        verifyNoMoreInteractions(ignoreStubs(testInfo2.socketKeepaliveCallback));
+    }
+
+    @Test
+    public void testTwoKeepalives_startAfterPause() throws Exception {
+        final TestKeepaliveInfo testInfo1 = doStartNattKeepalive();
+        checkAndProcessKeepaliveStart(testInfo1.kpd);
+        verify(testInfo1.socketKeepaliveCallback).onStarted();
+        assertNotNull(getAutoKiForBinder(testInfo1.binder));
+
+        final AutomaticOnOffKeepalive autoKi1  = getAutoKiForBinder(testInfo1.binder);
+        doPauseKeepalive(autoKi1);
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+        verify(testInfo1.socketKeepaliveCallback).onPaused();
+        assertNotNull(getAutoKiForBinder(testInfo1.binder));
+
+        clearInvocations(mNai);
+        // Start the second keepalive while the first is paused.
+        final TestKeepaliveInfo testInfo2 = doStartNattKeepalive();
+        // The slot used is TEST_SLOT since it is now a free slot.
+        checkAndProcessKeepaliveStart(TEST_SLOT, testInfo2.kpd);
+        verify(testInfo2.socketKeepaliveCallback).onStarted();
+        assertNotNull(getAutoKiForBinder(testInfo2.binder));
+
+        clearInvocations(mNai);
+        doResumeKeepalive(autoKi1);
+        // The next free slot is TEST_SLOT + 1.
+        checkAndProcessKeepaliveStart(TEST_SLOT + 1, testInfo1.kpd);
+        verify(testInfo1.socketKeepaliveCallback).onResumed();
+
+        clearInvocations(mNai);
+        doStopKeepalive(autoKi1);
+        // TODO: The slot should be consistent with the checkAndProcessKeepaliveStart directly above
+        checkAndProcessKeepaliveStop(TEST_SLOT);
+        // TODO: onStopped should only be called on the first keepalive callback.
+        verify(testInfo1.socketKeepaliveCallback, never()).onStopped();
+        verify(testInfo2.socketKeepaliveCallback).onStopped();
+        assertNull(getAutoKiForBinder(testInfo1.binder));
+
+        clearInvocations(mNai);
+        assertNotNull(getAutoKiForBinder(testInfo2.binder));
+        doStopKeepalive(getAutoKiForBinder(testInfo2.binder));
+        // This slot should be consistent with its corresponding checkAndProcessKeepaliveStart.
+        // TODO: checkAndProcessKeepaliveStop should be called instead but the keepalive is
+        // unexpectedly already stopped above.
+        verify(mNai, never()).onStopSocketKeepalive(TEST_SLOT);
+        verify(mNai, never()).onRemoveKeepalivePacketFilter(TEST_SLOT);
+
+        verify(testInfo2.socketKeepaliveCallback).onStopped();
+        assertNull(getAutoKiForBinder(testInfo2.binder));
+
+        verifyNoMoreInteractions(ignoreStubs(testInfo1.socketKeepaliveCallback));
+        verifyNoMoreInteractions(ignoreStubs(testInfo2.socketKeepaliveCallback));
+    }
 }
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 f15e8ff..da51240 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsServiceTypeClientTests.java
@@ -705,34 +705,8 @@
     }
 
     @Test
-    public void processResponse_notAllowRemoveSearch_shouldNotRemove() throws Exception {
-        final String serviceInstanceName = "service-instance-1";
-        client.startSendAndReceive(
-                mockListenerOne,
-                MdnsSearchOptions.newBuilder().build());
-        Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
-
-        // Process the initial response.
-        client.processResponse(createResponse(
-                serviceInstanceName, "192.168.1.1", 5353, /* subtype= */ "ABCDE",
-                Collections.emptyMap(), TEST_TTL), INTERFACE_INDEX, mockNetwork);
-
-        // Clear the scheduled runnable.
-        currentThreadExecutor.getAndClearLastScheduledRunnable();
-
-        // Simulate the case where the response is after TTL.
-        doReturn(TEST_ELAPSED_REALTIME + TEST_TTL + 1L).when(mockDecoderClock).elapsedRealtime();
-        firstMdnsTask.run();
-
-        // Verify removed callback was not called.
-        verifyServiceRemovedNoCallback(mockListenerOne);
-    }
-
-    @Test
-    @Ignore("MdnsConfigs is not configurable currently.")
-    public void processResponse_allowSearchOptionsToRemoveExpiredService_shouldRemove()
+    public void processResponse_searchOptionsEnableServiceRemoval_shouldRemove()
             throws Exception {
-        //MdnsConfigsFlagsImpl.allowSearchOptionsToRemoveExpiredService.override(true);
         final String serviceInstanceName = "service-instance-1";
         client =
                 new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
@@ -742,7 +716,9 @@
                         return mockPacketWriter;
                     }
                 };
-        client.startSendAndReceive(mockListenerOne, MdnsSearchOptions.getDefaultOptions());
+        MdnsSearchOptions searchOptions = MdnsSearchOptions.newBuilder().setRemoveExpiredService(
+                true).build();
+        client.startSendAndReceive(mockListenerOne, searchOptions);
         Runnable firstMdnsTask = currentThreadExecutor.getAndClearSubmittedRunnable();
 
         // Process the initial response.
@@ -956,15 +932,11 @@
 
         final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
                 new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
-        final List<MdnsRecord> srvTxtQuestions = srvTxtQueryPacket.questions;
 
-        final String[] serviceName = Stream.concat(Stream.of(instanceName),
-                Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
-        assertFalse(srvTxtQuestions.stream().anyMatch(q -> q.getType() == MdnsRecord.TYPE_PTR));
-        assertTrue(srvTxtQuestions.stream().anyMatch(q ->
-                q.getType() == MdnsRecord.TYPE_SRV && Arrays.equals(q.name, serviceName)));
-        assertTrue(srvTxtQuestions.stream().anyMatch(q ->
-                q.getType() == MdnsRecord.TYPE_TXT && Arrays.equals(q.name, serviceName)));
+        final String[] serviceName = getTestServiceName(instanceName);
+        assertFalse(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_PTR));
+        assertTrue(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_SRV, serviceName));
+        assertTrue(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_TXT, serviceName));
 
         // Process a response with SRV+TXT
         final MdnsPacket srvTxtResponse = new MdnsPacket(
@@ -991,11 +963,8 @@
 
         final MdnsPacket addressQueryPacket = MdnsPacket.parse(
                 new MdnsPacketReader(addressQueryCaptor.getValue()));
-        final List<MdnsRecord> addressQueryQuestions = addressQueryPacket.questions;
-        assertTrue(addressQueryQuestions.stream().anyMatch(q ->
-                q.getType() == MdnsRecord.TYPE_A && Arrays.equals(q.name, hostname)));
-        assertTrue(addressQueryQuestions.stream().anyMatch(q ->
-                q.getType() == MdnsRecord.TYPE_AAAA && Arrays.equals(q.name, hostname)));
+        assertTrue(hasQuestion(addressQueryPacket, MdnsRecord.TYPE_A, hostname));
+        assertTrue(hasQuestion(addressQueryPacket, MdnsRecord.TYPE_AAAA, hostname));
 
         // Process a response with address records
         final MdnsPacket addressResponse = new MdnsPacket(
@@ -1028,6 +997,81 @@
     }
 
     @Test
+    public void testRenewTxtSrvInResolve() throws Exception {
+        client = new MdnsServiceTypeClient(SERVICE_TYPE, mockSocketClient, currentThreadExecutor,
+                mockDecoderClock, mockNetwork, mockSharedLog);
+
+        final String instanceName = "service-instance";
+        final String[] hostname = new String[] { "testhost "};
+        final String ipV4Address = "192.0.2.0";
+        final String ipV6Address = "2001:db8::";
+
+        final MdnsSearchOptions resolveOptions = MdnsSearchOptions.newBuilder()
+                .setResolveInstanceName(instanceName).build();
+
+        client.startSendAndReceive(mockListenerOne, resolveOptions);
+        InOrder inOrder = inOrder(mockListenerOne, mockSocketClient);
+
+        // Get the query for SRV/TXT
+        final ArgumentCaptor<DatagramPacket> srvTxtQueryCaptor =
+                ArgumentCaptor.forClass(DatagramPacket.class);
+        currentThreadExecutor.getAndClearLastScheduledRunnable().run();
+        // Send twice for IPv4 and IPv6
+        inOrder.verify(mockSocketClient, times(2)).sendUnicastPacket(srvTxtQueryCaptor.capture(),
+                eq(mockNetwork));
+
+        final MdnsPacket srvTxtQueryPacket = MdnsPacket.parse(
+                new MdnsPacketReader(srvTxtQueryCaptor.getValue()));
+
+        final String[] serviceName = getTestServiceName(instanceName);
+        assertTrue(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_SRV, serviceName));
+        assertTrue(hasQuestion(srvTxtQueryPacket, MdnsRecord.TYPE_TXT, serviceName));
+
+        // Process a response with all records
+        final MdnsPacket srvTxtResponse = new MdnsPacket(
+                0 /* flags */,
+                Collections.emptyList() /* questions */,
+                List.of(
+                        new MdnsServiceRecord(serviceName, TEST_ELAPSED_REALTIME,
+                                true /* cacheFlush */, TEST_TTL, 0 /* servicePriority */,
+                                0 /* serviceWeight */, 1234 /* servicePort */, hostname),
+                        new MdnsTextRecord(serviceName, TEST_ELAPSED_REALTIME,
+                                true /* cacheFlush */, TEST_TTL,
+                                Collections.emptyList() /* entries */),
+                        new MdnsInetAddressRecord(hostname, TEST_ELAPSED_REALTIME,
+                                true /* cacheFlush */, TEST_TTL,
+                                InetAddresses.parseNumericAddress(ipV4Address)),
+                        new MdnsInetAddressRecord(hostname, TEST_ELAPSED_REALTIME,
+                                true /* cacheFlush */, TEST_TTL,
+                                InetAddresses.parseNumericAddress(ipV6Address))),
+                Collections.emptyList() /* authorityRecords */,
+                Collections.emptyList() /* additionalRecords */);
+        client.processResponse(srvTxtResponse, INTERFACE_INDEX, mockNetwork);
+        inOrder.verify(mockListenerOne).onServiceNameDiscovered(any());
+        inOrder.verify(mockListenerOne).onServiceFound(any());
+
+        // Expect no query on the next run
+        currentThreadExecutor.getAndClearLastScheduledRunnable().run();
+        inOrder.verifyNoMoreInteractions();
+
+        // Advance time so 75% of TTL passes and re-execute
+        doReturn(TEST_ELAPSED_REALTIME + (long) (TEST_TTL * 0.75))
+                .when(mockDecoderClock).elapsedRealtime();
+        currentThreadExecutor.getAndClearLastScheduledRunnable().run();
+
+        // Expect a renewal query
+        final ArgumentCaptor<DatagramPacket> renewalQueryCaptor =
+                ArgumentCaptor.forClass(DatagramPacket.class);
+        // Second and later sends are sent as "expect multicast response" queries
+        inOrder.verify(mockSocketClient, times(2)).sendMulticastPacket(renewalQueryCaptor.capture(),
+                eq(mockNetwork));
+        final MdnsPacket renewalPacket = MdnsPacket.parse(
+                new MdnsPacketReader(renewalQueryCaptor.getValue()));
+        assertTrue(hasQuestion(renewalPacket, MdnsRecord.TYPE_SRV, serviceName));
+        assertTrue(hasQuestion(renewalPacket, MdnsRecord.TYPE_TXT, serviceName));
+    }
+
+    @Test
     public void testProcessResponse_ResolveExcludesOtherServices() {
         client = new MdnsServiceTypeClient(
                 SERVICE_TYPE, mockSocketClient, currentThreadExecutor, mockNetwork, mockSharedLog);
@@ -1270,6 +1314,20 @@
         }
     }
 
+    private static String[] getTestServiceName(String instanceName) {
+        return Stream.concat(Stream.of(instanceName),
+                Arrays.stream(SERVICE_TYPE_LABELS)).toArray(String[]::new);
+    }
+
+    private static boolean hasQuestion(MdnsPacket packet, int type) {
+        return hasQuestion(packet, type, null);
+    }
+
+    private static boolean hasQuestion(MdnsPacket packet, int type, @Nullable String[] name) {
+        return packet.questions.stream().anyMatch(q -> q.getType() == type
+                && (name == null || Arrays.equals(q.name, name)));
+    }
+
     // A fake ScheduledExecutorService that keeps tracking the last scheduled Runnable and its delay
     // time.
     private class FakeExecutor extends ScheduledThreadPoolExecutor {
diff --git a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java
index 4b87556..4f56857 100644
--- a/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java
+++ b/tests/unit/java/com/android/server/connectivity/mdns/MdnsSocketProviderTest.java
@@ -31,6 +31,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doCallRealMethod;
@@ -40,7 +41,9 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
 import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.net.LinkAddress;
@@ -49,6 +52,9 @@
 import android.net.NetworkCapabilities;
 import android.net.TetheringManager;
 import android.net.TetheringManager.TetheringEventCallback;
+import android.net.wifi.p2p.WifiP2pGroup;
+import android.net.wifi.p2p.WifiP2pInfo;
+import android.net.wifi.p2p.WifiP2pManager;
 import android.os.Build;
 import android.os.Handler;
 import android.os.HandlerThread;
@@ -88,6 +94,7 @@
     private static final String TAG = MdnsSocketProviderTest.class.getSimpleName();
     private static final String TEST_IFACE_NAME = "test";
     private static final String LOCAL_ONLY_IFACE_NAME = "local_only";
+    private static final String WIFI_P2P_IFACE_NAME = "p2p_wifi";
     private static final String TETHERED_IFACE_NAME = "tethered";
     private static final int TETHERED_IFACE_IDX = 32;
     private static final long DEFAULT_TIMEOUT = 2000L;
@@ -136,11 +143,15 @@
         doReturn(true).when(mTetheredIfaceWrapper).supportsMulticast();
         doReturn(mLocalOnlyIfaceWrapper).when(mDeps)
                 .getNetworkInterfaceByName(LOCAL_ONLY_IFACE_NAME);
+        doReturn(mLocalOnlyIfaceWrapper).when(mDeps)
+                .getNetworkInterfaceByName(WIFI_P2P_IFACE_NAME);
         doReturn(mTetheredIfaceWrapper).when(mDeps).getNetworkInterfaceByName(TETHERED_IFACE_NAME);
         doReturn(mock(MdnsInterfaceSocket.class))
                 .when(mDeps).createMdnsInterfaceSocket(any(), anyInt(), any(), any());
         doReturn(TETHERED_IFACE_IDX).when(mDeps).getNetworkInterfaceIndexByName(
                 TETHERED_IFACE_NAME);
+        doReturn(789).when(mDeps).getNetworkInterfaceIndexByName(
+                WIFI_P2P_IFACE_NAME);
         final HandlerThread thread = new HandlerThread("MdnsSocketProviderTest");
         thread.start();
         mHandler = new Handler(thread.getLooper());
@@ -157,22 +168,41 @@
         mSocketProvider = new MdnsSocketProvider(mContext, thread.getLooper(), mDeps, mLog);
     }
 
+    private void runOnHandler(Runnable r) {
+        mHandler.post(r);
+        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+    }
+
+    private BroadcastReceiver expectWifiP2PChangeBroadcastReceiver() {
+        final ArgumentCaptor<BroadcastReceiver> receiverCaptor =
+                ArgumentCaptor.forClass(BroadcastReceiver.class);
+        verify(mContext, times(1)).registerReceiver(receiverCaptor.capture(),
+                argThat(filter -> filter.hasAction(
+                        WifiP2pManager.WIFI_P2P_CONNECTION_CHANGED_ACTION)),
+                any(), any());
+        final BroadcastReceiver originalReceiver = receiverCaptor.getValue();
+        return new BroadcastReceiver() {
+            @Override
+            public void onReceive(Context context, Intent intent) {
+                runOnHandler(() -> originalReceiver.onReceive(context, intent));
+            }
+        };
+    }
+
     private void startMonitoringSockets() {
         final ArgumentCaptor<NetworkCallback> nwCallbackCaptor =
                 ArgumentCaptor.forClass(NetworkCallback.class);
         final ArgumentCaptor<TetheringEventCallback> teCallbackCaptor =
                 ArgumentCaptor.forClass(TetheringEventCallback.class);
 
-        mHandler.post(mSocketProvider::startMonitoringSockets);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::startMonitoringSockets);
         verify(mCm).registerNetworkCallback(any(), nwCallbackCaptor.capture(), any());
         verify(mTm).registerTetheringEventCallback(any(), teCallbackCaptor.capture());
 
         mNetworkCallback = nwCallbackCaptor.getValue();
         mTetheringEventCallback = teCallbackCaptor.getValue();
 
-        mHandler.post(mSocketProvider::startNetLinkMonitor);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::startNetLinkMonitor);
     }
 
     private static class TestNetlinkMonitor extends SocketNetlinkMonitor {
@@ -281,9 +311,8 @@
         testLp.setInterfaceName(TEST_IFACE_NAME);
         testLp.setLinkAddresses(List.of(LINKADDRV4));
         final NetworkCapabilities testNc = makeCapabilities(transports);
-        mHandler.post(() -> mNetworkCallback.onCapabilitiesChanged(TEST_NETWORK, testNc));
-        mHandler.post(() -> mNetworkCallback.onLinkPropertiesChanged(TEST_NETWORK, testLp));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mNetworkCallback.onCapabilitiesChanged(TEST_NETWORK, testNc));
+        runOnHandler(() -> mNetworkCallback.onLinkPropertiesChanged(TEST_NETWORK, testLp));
     }
 
     @Test
@@ -291,62 +320,53 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback1 = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback1));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback1));
         testCallback1.expectedNoCallback();
 
         postNetworkAvailable(TRANSPORT_WIFI);
         testCallback1.expectedSocketCreatedForNetwork(TEST_NETWORK, List.of(LINKADDRV4));
 
         final TestSocketCallback testCallback2 = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback2));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback2));
         testCallback1.expectedNoCallback();
         testCallback2.expectedSocketCreatedForNetwork(TEST_NETWORK, List.of(LINKADDRV4));
 
         final TestSocketCallback testCallback3 = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(null /* network */, testCallback3));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(null /* network */, testCallback3));
         testCallback1.expectedNoCallback();
         testCallback2.expectedNoCallback();
         testCallback3.expectedSocketCreatedForNetwork(TEST_NETWORK, List.of(LINKADDRV4));
 
-        mHandler.post(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(
+        runOnHandler(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(
                 List.of(LOCAL_ONLY_IFACE_NAME)));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mLocalOnlyIfaceWrapper).getNetworkInterface();
         testCallback1.expectedNoCallback();
         testCallback2.expectedNoCallback();
         testCallback3.expectedSocketCreatedForNetwork(null /* network */, List.of());
 
-        mHandler.post(() -> mTetheringEventCallback.onTetheredInterfacesChanged(
+        runOnHandler(() -> mTetheringEventCallback.onTetheredInterfacesChanged(
                 List.of(TETHERED_IFACE_NAME)));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mTetheredIfaceWrapper).getNetworkInterface();
         testCallback1.expectedNoCallback();
         testCallback2.expectedNoCallback();
         testCallback3.expectedSocketCreatedForNetwork(null /* network */, List.of());
 
-        mHandler.post(() -> mSocketProvider.unrequestSocket(testCallback1));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.unrequestSocket(testCallback1));
         testCallback1.expectedNoCallback();
         testCallback2.expectedNoCallback();
         testCallback3.expectedNoCallback();
 
-        mHandler.post(() -> mNetworkCallback.onLost(TEST_NETWORK));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mNetworkCallback.onLost(TEST_NETWORK));
         testCallback1.expectedNoCallback();
         testCallback2.expectedInterfaceDestroyedForNetwork(TEST_NETWORK);
         testCallback3.expectedInterfaceDestroyedForNetwork(TEST_NETWORK);
 
-        mHandler.post(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(List.of()));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(List.of()));
         testCallback1.expectedNoCallback();
         testCallback2.expectedNoCallback();
         testCallback3.expectedInterfaceDestroyedForNetwork(null /* network */);
 
-        mHandler.post(() -> mSocketProvider.unrequestSocket(testCallback3));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.unrequestSocket(testCallback3));
         testCallback1.expectedNoCallback();
         testCallback2.expectedNoCallback();
         // There was still a tethered interface, but no callback should be sent once unregistered
@@ -376,8 +396,7 @@
     public void testDownstreamNetworkAddressUpdateFromNetlink() {
         startMonitoringSockets();
         final TestSocketCallback testCallbackAll = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(null /* network */, testCallbackAll));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(null /* network */, testCallbackAll));
 
         // Address add message arrived before the interface is created.
         RtNetlinkAddressMessage addIpv4AddrMsg = createNetworkAddressUpdateNetLink(
@@ -385,15 +404,13 @@
                 LINKADDRV4,
                 TETHERED_IFACE_IDX,
                 0 /* flags */);
-        mHandler.post(
+        runOnHandler(
                 () -> mTestSocketNetLinkMonitor.processNetlinkMessage(addIpv4AddrMsg,
                         0 /* whenMs */));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
 
         // Interface is created.
-        mHandler.post(() -> mTetheringEventCallback.onTetheredInterfacesChanged(
+        runOnHandler(() -> mTetheringEventCallback.onTetheredInterfacesChanged(
                 List.of(TETHERED_IFACE_NAME)));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         verify(mTetheredIfaceWrapper).getNetworkInterface();
         testCallbackAll.expectedSocketCreatedForNetwork(null /* network */, List.of(LINKADDRV4));
 
@@ -403,10 +420,9 @@
                 LINKADDRV4,
                 TETHERED_IFACE_IDX,
                 0 /* flags */);
-        mHandler.post(
+        runOnHandler(
                 () -> mTestSocketNetLinkMonitor.processNetlinkMessage(removeIpv4AddrMsg,
                         0 /* whenMs */));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         testCallbackAll.expectedAddressesChangedForNetwork(null /* network */, List.of());
 
         // New address added.
@@ -415,9 +431,8 @@
                 LINKADDRV6,
                 TETHERED_IFACE_IDX,
                 0 /* flags */);
-        mHandler.post(() -> mTestSocketNetLinkMonitor.processNetlinkMessage(addIpv6AddrMsg,
+        runOnHandler(() -> mTestSocketNetLinkMonitor.processNetlinkMessage(addIpv6AddrMsg,
                 0 /* whenMs */));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         testCallbackAll.expectedAddressesChangedForNetwork(null /* network */, List.of(LINKADDRV6));
 
         // Address updated
@@ -426,10 +441,9 @@
                 LINKADDRV6,
                 TETHERED_IFACE_IDX,
                 1 /* flags */);
-        mHandler.post(
+        runOnHandler(
                 () -> mTestSocketNetLinkMonitor.processNetlinkMessage(updateIpv6AddrMsg,
                         0 /* whenMs */));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
         testCallbackAll.expectedAddressesChangedForNetwork(null /* network */,
                 List.of(LINKADDRV6_FLAG_CHANGE));
     }
@@ -439,8 +453,7 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
         testCallback.expectedNoCallback();
 
         postNetworkAvailable(TRANSPORT_WIFI);
@@ -449,8 +462,7 @@
         final LinkProperties newTestLp = new LinkProperties();
         newTestLp.setInterfaceName(TEST_IFACE_NAME);
         newTestLp.setLinkAddresses(List.of(LINKADDRV4, LINKADDRV6));
-        mHandler.post(() -> mNetworkCallback.onLinkPropertiesChanged(TEST_NETWORK, newTestLp));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mNetworkCallback.onLinkPropertiesChanged(TEST_NETWORK, newTestLp));
         testCallback.expectedAddressesChangedForNetwork(
                 TEST_NETWORK, List.of(LINKADDRV4, LINKADDRV6));
     }
@@ -458,8 +470,7 @@
     @Test
     public void testStartAndStopMonitoringSockets() {
         // Stop monitoring sockets before start. Should not unregister any network callback.
-        mHandler.post(mSocketProvider::requestStopWhenInactive);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::requestStopWhenInactive);
         verify(mCm, never()).unregisterNetworkCallback(any(NetworkCallback.class));
         verify(mTm, never()).unregisterTetheringEventCallback(any(TetheringEventCallback.class));
 
@@ -467,39 +478,32 @@
         startMonitoringSockets();
         // Request a socket then unrequest it. Expect no network callback unregistration.
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
         testCallback.expectedNoCallback();
-        mHandler.post(()-> mSocketProvider.unrequestSocket(testCallback));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(()-> mSocketProvider.unrequestSocket(testCallback));
         verify(mCm, never()).unregisterNetworkCallback(any(NetworkCallback.class));
         verify(mTm, never()).unregisterTetheringEventCallback(any(TetheringEventCallback.class));
         // Request stop and it should unregister network callback immediately because there is no
         // socket request.
-        mHandler.post(mSocketProvider::requestStopWhenInactive);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::requestStopWhenInactive);
         verify(mCm, times(1)).unregisterNetworkCallback(any(NetworkCallback.class));
         verify(mTm, times(1)).unregisterTetheringEventCallback(any(TetheringEventCallback.class));
 
         // Start sockets monitoring and request a socket again.
-        mHandler.post(mSocketProvider::startMonitoringSockets);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::startMonitoringSockets);
         verify(mCm, times(2)).registerNetworkCallback(any(), any(NetworkCallback.class), any());
         verify(mTm, times(2)).registerTetheringEventCallback(
                 any(), any(TetheringEventCallback.class));
         final TestSocketCallback testCallback2 = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback2));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback2));
         testCallback2.expectedNoCallback();
         // Try to stop monitoring sockets but should be ignored and wait until all socket are
         // unrequested.
-        mHandler.post(mSocketProvider::requestStopWhenInactive);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::requestStopWhenInactive);
         verify(mCm, times(1)).unregisterNetworkCallback(any(NetworkCallback.class));
         verify(mTm, times(1)).unregisterTetheringEventCallback(any());
         // Unrequest the socket then network callbacks should be unregistered.
-        mHandler.post(()-> mSocketProvider.unrequestSocket(testCallback2));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(()-> mSocketProvider.unrequestSocket(testCallback2));
         verify(mCm, times(2)).unregisterNetworkCallback(any(NetworkCallback.class));
         verify(mTm, times(2)).unregisterTetheringEventCallback(any(TetheringEventCallback.class));
     }
@@ -510,24 +514,20 @@
 
         // Request a socket with null network.
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(null, testCallback));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(null, testCallback));
         testCallback.expectedNoCallback();
 
         // Notify a LinkPropertiesChanged with TEST_NETWORK.
         final LinkProperties testLp = new LinkProperties();
         testLp.setInterfaceName(TEST_IFACE_NAME);
         testLp.setLinkAddresses(List.of(LINKADDRV4));
-        mHandler.post(() -> mNetworkCallback.onLinkPropertiesChanged(TEST_NETWORK, testLp));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mNetworkCallback.onLinkPropertiesChanged(TEST_NETWORK, testLp));
         verify(mTestNetworkIfaceWrapper, times(1)).getNetworkInterface();
         testCallback.expectedSocketCreatedForNetwork(TEST_NETWORK, List.of(LINKADDRV4));
 
         // Try to stop monitoring and unrequest the socket.
-        mHandler.post(mSocketProvider::requestStopWhenInactive);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
-        mHandler.post(()-> mSocketProvider.unrequestSocket(testCallback));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::requestStopWhenInactive);
+        runOnHandler(()-> mSocketProvider.unrequestSocket(testCallback));
         // No callback sent when unregistered
         testCallback.expectedNoCallback();
         verify(mCm, times(1)).unregisterNetworkCallback(any(NetworkCallback.class));
@@ -535,13 +535,11 @@
 
         // Start sockets monitoring and request a socket again. Expected no socket created callback
         // because all saved LinkProperties has been cleared.
-        mHandler.post(mSocketProvider::startMonitoringSockets);
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(mSocketProvider::startMonitoringSockets);
         verify(mCm, times(2)).registerNetworkCallback(any(), any(NetworkCallback.class), any());
         verify(mTm, times(2)).registerTetheringEventCallback(
                 any(), any(TetheringEventCallback.class));
-        mHandler.post(() -> mSocketProvider.requestSocket(null, testCallback));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mSocketProvider.requestSocket(null, testCallback));
         testCallback.expectedNoCallback();
 
         // Notify a LinkPropertiesChanged with another network.
@@ -550,8 +548,7 @@
         final Network otherNetwork = new Network(456);
         otherLp.setInterfaceName("test2");
         otherLp.setLinkAddresses(List.of(otherAddress));
-        mHandler.post(() -> mNetworkCallback.onLinkPropertiesChanged(otherNetwork, otherLp));
-        HandlerUtils.waitForIdle(mHandler, DEFAULT_TIMEOUT);
+        runOnHandler(() -> mNetworkCallback.onLinkPropertiesChanged(otherNetwork, otherLp));
         verify(mTestNetworkIfaceWrapper, times(2)).getNetworkInterface();
         testCallback.expectedSocketCreatedForNetwork(otherNetwork, List.of(otherAddress));
     }
@@ -561,7 +558,7 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
 
         postNetworkAvailable(TRANSPORT_CELLULAR);
         testCallback.expectedNoCallback();
@@ -573,7 +570,7 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
 
         postNetworkAvailable(TRANSPORT_BLUETOOTH);
         testCallback.expectedNoCallback();
@@ -585,7 +582,7 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
 
         postNetworkAvailable(TRANSPORT_BLUETOOTH);
         testCallback.expectedSocketCreatedForNetwork(TEST_NETWORK, List.of(LINKADDRV4));
@@ -597,7 +594,7 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
 
         postNetworkAvailable(TRANSPORT_BLUETOOTH);
         testCallback.expectedNoCallback();
@@ -611,7 +608,7 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
 
         postNetworkAvailable(TRANSPORT_VPN, TRANSPORT_WIFI);
         testCallback.expectedNoCallback();
@@ -623,9 +620,146 @@
         startMonitoringSockets();
 
         final TestSocketCallback testCallback = new TestSocketCallback();
-        mHandler.post(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
+        runOnHandler(() -> mSocketProvider.requestSocket(TEST_NETWORK, testCallback));
 
         postNetworkAvailable(TRANSPORT_WIFI);
         testCallback.expectedSocketCreatedForNetwork(TEST_NETWORK, List.of(LINKADDRV4));
     }
+
+    private Intent buildWifiP2PConnectionChangedIntent(boolean groupFormed) {
+        final Intent intent = new Intent(WifiP2pManager.WIFI_P2P_CONNECTION_CHANGED_ACTION);
+        final WifiP2pInfo formedInfo = new WifiP2pInfo();
+        formedInfo.groupFormed = groupFormed;
+        final WifiP2pGroup group;
+        if (groupFormed) {
+            group = mock(WifiP2pGroup.class);
+            doReturn(WIFI_P2P_IFACE_NAME).when(group).getInterface();
+        } else {
+            group = null;
+        }
+        intent.putExtra(WifiP2pManager.EXTRA_WIFI_P2P_INFO, formedInfo);
+        intent.putExtra(WifiP2pManager.EXTRA_WIFI_P2P_GROUP, group);
+        return intent;
+    }
+
+    @Test
+    public void testWifiP2PInterfaceChange() {
+        final BroadcastReceiver receiver = expectWifiP2PChangeBroadcastReceiver();
+        startMonitoringSockets();
+
+        // Request a socket with null network.
+        final TestSocketCallback testCallback = new TestSocketCallback();
+        runOnHandler(() -> mSocketProvider.requestSocket(null /* network */, testCallback));
+
+        // Wifi p2p is connected and the interface is up. Get a wifi p2p change intent then expect
+        // a socket creation.
+        final Intent formedIntent = buildWifiP2PConnectionChangedIntent(true /* groupFormed */);
+        receiver.onReceive(mContext, formedIntent);
+        verify(mLocalOnlyIfaceWrapper).getNetworkInterface();
+        testCallback.expectedSocketCreatedForNetwork(null /* network */, List.of());
+
+        // Wifi p2p is disconnected. Get a wifi p2p change intent then expect the socket destroy.
+        final Intent unformedIntent = buildWifiP2PConnectionChangedIntent(false /* groupFormed */);
+        receiver.onReceive(mContext, unformedIntent);
+        testCallback.expectedInterfaceDestroyedForNetwork(null /* network */);
+    }
+
+    @Test
+    public void testWifiP2PInterfaceChangeBeforeStartMonitoringSockets() {
+        final BroadcastReceiver receiver = expectWifiP2PChangeBroadcastReceiver();
+
+        // Get a wifi p2p change intent before start monitoring sockets.
+        final Intent formedIntent = buildWifiP2PConnectionChangedIntent(true /* groupFormed */);
+        receiver.onReceive(mContext, formedIntent);
+
+        // Start monitoring sockets and request a socket with null network.
+        startMonitoringSockets();
+        final TestSocketCallback testCallback = new TestSocketCallback();
+        runOnHandler(() -> mSocketProvider.requestSocket(null /* network */, testCallback));
+        verify(mLocalOnlyIfaceWrapper).getNetworkInterface();
+        testCallback.expectedSocketCreatedForNetwork(null /* network */, List.of());
+    }
+
+    @Test
+    public void testWifiP2PInterfaceChangeBeforeGetAllNetworksRequest() {
+        final BroadcastReceiver receiver = expectWifiP2PChangeBroadcastReceiver();
+        startMonitoringSockets();
+
+        // Get a wifi p2p change intent before request socket for all networks.
+        final Intent formedIntent = buildWifiP2PConnectionChangedIntent(true /* groupFormed */);
+        receiver.onReceive(mContext, formedIntent);
+
+        // Request a socket with null network.
+        final TestSocketCallback testCallback = new TestSocketCallback();
+        runOnHandler(() -> mSocketProvider.requestSocket(null /* network */, testCallback));
+        verify(mLocalOnlyIfaceWrapper).getNetworkInterface();
+        testCallback.expectedSocketCreatedForNetwork(null /* network */, List.of());
+    }
+
+    @Test
+    public void testNoDuplicatedSocketCreation() {
+        final BroadcastReceiver receiver = expectWifiP2PChangeBroadcastReceiver();
+        startMonitoringSockets();
+
+        // Request a socket with null network.
+        final TestSocketCallback testCallback = new TestSocketCallback();
+        runOnHandler(() -> mSocketProvider.requestSocket(null, testCallback));
+        testCallback.expectedNoCallback();
+
+        // Receive an interface added change for the wifi p2p interface. Expect a socket creation
+        // callback.
+        runOnHandler(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(
+                List.of(WIFI_P2P_IFACE_NAME)));
+        verify(mLocalOnlyIfaceWrapper, times(1)).getNetworkInterface();
+        testCallback.expectedSocketCreatedForNetwork(null /* network */, List.of());
+
+        // Receive a wifi p2p connected intent. Expect no callback because the socket is created.
+        final Intent formedIntent = buildWifiP2PConnectionChangedIntent(true /* groupFormed */);
+        receiver.onReceive(mContext, formedIntent);
+        testCallback.expectedNoCallback();
+
+        // Request other socket with null network. Should receive socket created callback once.
+        final TestSocketCallback testCallback2 = new TestSocketCallback();
+        runOnHandler(() -> mSocketProvider.requestSocket(null, testCallback2));
+        testCallback2.expectedSocketCreatedForNetwork(null /* network */, List.of());
+        testCallback2.expectedNoCallback();
+
+        // Receive a wifi p2p disconnected intent. Expect a socket destroy callback.
+        final Intent unformedIntent = buildWifiP2PConnectionChangedIntent(false /* groupFormed */);
+        receiver.onReceive(mContext, unformedIntent);
+        testCallback.expectedInterfaceDestroyedForNetwork(null /* network */);
+
+        // Receive an interface removed change for the wifi p2p interface. Expect no callback
+        // because the socket is destroyed.
+        runOnHandler(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(List.of()));
+        testCallback.expectedNoCallback();
+
+        // Receive a wifi p2p connected intent again. Expect a socket creation callback.
+        receiver.onReceive(mContext, formedIntent);
+        verify(mLocalOnlyIfaceWrapper, times(2)).getNetworkInterface();
+        testCallback.expectedSocketCreatedForNetwork(null /* network */, List.of());
+
+        // Receive an interface added change for the wifi p2p interface again. Expect no callback
+        // because the socket is created.
+        runOnHandler(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(
+                List.of(WIFI_P2P_IFACE_NAME)));
+        testCallback.expectedNoCallback();
+    }
+
+    @Test
+    public void testTetherInterfacesChangedBeforeGetAllNetworksRequest() {
+        startMonitoringSockets();
+
+        // Receive an interface added change for the wifi p2p interface. Expect a socket creation
+        // callback.
+        runOnHandler(() -> mTetheringEventCallback.onLocalOnlyInterfacesChanged(
+                List.of(TETHERED_IFACE_NAME)));
+        verify(mTetheredIfaceWrapper, never()).getNetworkInterface();
+
+        // Request a socket with null network.
+        final TestSocketCallback testCallback = new TestSocketCallback();
+        runOnHandler(() -> mSocketProvider.requestSocket(null /* network */, testCallback));
+        verify(mTetheredIfaceWrapper).getNetworkInterface();
+        testCallback.expectedSocketCreatedForNetwork(null /* network */, List.of());
+    }
 }