Merge "Rename some variables in NsdService"
diff --git a/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt b/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
index 0760e68..0885f4f 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
+++ b/Cronet/tests/cts/src/android/net/http/cts/BidirectionalStreamTest.kt
@@ -27,6 +27,7 @@
 import androidx.test.core.app.ApplicationProvider
 import com.android.testutils.DevSdkIgnoreRule
 import com.android.testutils.DevSdkIgnoreRunner
+import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import org.hamcrest.MatcherAssert
@@ -81,4 +82,113 @@
             "Received byte count must be > 0", info.receivedByteCount, Matchers.greaterThan(0L))
         assertEquals("h2", info.negotiatedProtocol)
     }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getHttpMethod() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val method = "GET"
+
+        builder.setHttpMethod(method)
+        stream = builder.build()
+        assertThat(stream!!.getHttpMethod()).isEqualTo(method)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_hasTrafficStatsTag() {
+        val builder = createBidirectionalStreamBuilder(URL)
+
+        builder.setTrafficStatsTag(10)
+        stream = builder.build()
+        assertThat(stream!!.hasTrafficStatsTag()).isTrue()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getTrafficStatsTag() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val trafficStatsTag = 10
+
+        builder.setTrafficStatsTag(trafficStatsTag)
+        stream = builder.build()
+        assertThat(stream!!.getTrafficStatsTag()).isEqualTo(trafficStatsTag)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_hasTrafficStatsUid() {
+        val builder = createBidirectionalStreamBuilder(URL)
+
+        builder.setTrafficStatsUid(10)
+        stream = builder.build()
+        assertThat(stream!!.hasTrafficStatsUid()).isTrue()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getTrafficStatsUid() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val trafficStatsUid = 10
+
+        builder.setTrafficStatsUid(trafficStatsUid)
+        stream = builder.build()
+        assertThat(stream!!.getTrafficStatsUid()).isEqualTo(trafficStatsUid)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getHeaders_asList() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val expectedHeaders = mapOf(
+          "Authorization" to "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ==",
+          "Max-Forwards" to "10",
+          "X-Client-Data" to "random custom header content").entries.toList()
+
+        for (header in expectedHeaders) {
+            builder.addHeader(header.key, header.value)
+        }
+
+        stream = builder.build()
+        assertThat(stream!!.getHeaders().getAsList()).containsAtLeastElementsIn(expectedHeaders)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getHeaders_asMap() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val expectedHeaders = mapOf(
+          "Authorization" to listOf("Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="),
+          "Max-Forwards" to listOf("10"),
+          "X-Client-Data" to listOf("random custom header content"))
+
+        for (header in expectedHeaders) {
+            builder.addHeader(header.key, header.value.get(0))
+        }
+
+        stream = builder.build()
+        assertThat(stream!!.getHeaders().getAsMap()).containsAtLeastEntriesIn(expectedHeaders)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_getPriority() {
+        val builder = createBidirectionalStreamBuilder(URL)
+        val priority = BidirectionalStream.STREAM_PRIORITY_LOW
+
+        builder.setPriority(priority)
+        stream = builder.build()
+        assertThat(stream!!.getPriority()).isEqualTo(priority)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun testBidirectionalStream_isDelayRequestHeadersUntilFirstFlushEnabled() {
+        val builder = createBidirectionalStreamBuilder(URL)
+
+        builder.setDelayRequestHeadersUntilFirstFlushEnabled(true)
+        stream = builder.build()
+        assertThat(stream!!.isDelayRequestHeadersUntilFirstFlushEnabled()).isTrue()
+    }
+
 }
diff --git a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
index a9add20..9162546 100644
--- a/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
+++ b/Tethering/tests/integration/base/android/net/EthernetTetheringTestBase.java
@@ -119,9 +119,11 @@
             (Inet4Address) parseNumericAddress("8.8.8.8");
     protected static final Inet6Address REMOTE_IP6_ADDR =
             (Inet6Address) parseNumericAddress("2002:db8:1::515:ca");
+    // The IPv6 network address translation of REMOTE_IP4_ADDR if pref64::/n is 64:ff9b::/96.
+    // For more information, see TetheringTester#PREF64_IPV4ONLY_ADDR, which assumes a prefix
+    // of 64:ff9b::/96.
     protected static final Inet6Address REMOTE_NAT64_ADDR =
             (Inet6Address) parseNumericAddress("64:ff9b::808:808");
-    protected static final IpPrefix TEST_NAT64PREFIX = new IpPrefix("64:ff9b::/96");
 
     // LOCAL_PORT is used by public port and private port. Assume port 9876 has not been used yet
     // before the testing that public port and private port are the same in the testing. Note that
@@ -630,7 +632,6 @@
         final LinkProperties lp = new LinkProperties();
         lp.setLinkAddresses(addresses);
         lp.setDnsServers(dnses);
-        lp.setNat64Prefix(TEST_NAT64PREFIX);
 
         return runAsShell(MANAGE_TEST_NETWORKS, () -> initTestNetwork(mContext, lp, TIMEOUT_MS));
     }
diff --git a/Tethering/tests/integration/base/android/net/TetheringTester.java b/Tethering/tests/integration/base/android/net/TetheringTester.java
index 33baf93..3f3768e 100644
--- a/Tethering/tests/integration/base/android/net/TetheringTester.java
+++ b/Tethering/tests/integration/base/android/net/TetheringTester.java
@@ -16,6 +16,8 @@
 
 package android.net;
 
+import static android.net.DnsResolver.CLASS_IN;
+import static android.net.DnsResolver.TYPE_AAAA;
 import static android.net.InetAddresses.parseNumericAddress;
 import static android.system.OsConstants.ICMP_ECHO;
 import static android.system.OsConstants.ICMP_ECHOREPLY;
@@ -28,6 +30,8 @@
 
 import static com.android.net.module.util.DnsPacket.ANSECTION;
 import static com.android.net.module.util.DnsPacket.ARSECTION;
+import static com.android.net.module.util.DnsPacket.DnsHeader;
+import static com.android.net.module.util.DnsPacket.DnsRecord;
 import static com.android.net.module.util.DnsPacket.NSSECTION;
 import static com.android.net.module.util.DnsPacket.QDSECTION;
 import static com.android.net.module.util.HexDump.dumpHexString;
@@ -130,6 +134,27 @@
     // ICMP definition.
     private static final short ICMPECHO_CODE = 0x0;
 
+    // Prefix64 discovery definition. See RFC 7050 section 8.
+    // Note that the AAAA response Pref64::WKAs consisting of Pref64::/n and WKA.
+    // Use 64:ff9b::/96 as Pref64::/n and WKA 192.0.0.17{0|1} here.
+    //
+    // Host                                          DNS64 server
+    //   |                                                |
+    //   |  "AAAA" query for "ipv4only.arpa."             |
+    //   |----------------------------------------------->|
+    //   |                                                |
+    //   |  "AAAA" response with:                         |
+    //   |  "64:ff9b::192.0.0.170"                        |
+    //   |<-----------------------------------------------|
+    //
+    private static final String PREF64_IPV4ONLY_HOSTNAME = "ipv4only.arpa";
+    private static final InetAddress PREF64_IPV4ONLY_ADDR = parseNumericAddress(
+            "64:ff9b::192.0.0.170");
+
+    // DNS header definition.
+    private static final short FLAG = (short) 0x8100;  // qr, ra
+    private static final short TTL = (short) 0;
+
     public static final String DHCP_HOSTNAME = "testhostname";
 
     private final ArrayMap<MacAddress, TetheredDevice> mTetheredDevices;
@@ -490,6 +515,11 @@
             super(data);
         }
 
+        TestDnsPacket(@NonNull DnsHeader header, @Nullable ArrayList<DnsRecord> qd,
+                @Nullable ArrayList<DnsRecord> an) {
+            super(header, qd, an);
+        }
+
         @Nullable
         public static TestDnsPacket getTestDnsPacket(final ByteBuffer buf) {
             try {
@@ -862,10 +892,85 @@
         return null;
     }
 
+    @NonNull
+    private ByteBuffer buildUdpDnsPrefix64ReplyPacket(int dnsId, @NonNull final Inet6Address srcIp,
+            @NonNull final Inet6Address dstIp, short srcPort, short dstPort) throws Exception {
+        // [1] Build prefix64 DNS message.
+        final ArrayList<DnsRecord> qlist = new ArrayList<>();
+        // Fill QD section.
+        qlist.add(DnsRecord.makeQuestion(PREF64_IPV4ONLY_HOSTNAME, TYPE_AAAA, CLASS_IN));
+        final ArrayList<DnsRecord> alist = new ArrayList<>();
+        // Fill AN sections.
+        alist.add(DnsRecord.makeAOrAAAARecord(ANSECTION, PREF64_IPV4ONLY_HOSTNAME, CLASS_IN, TTL,
+                PREF64_IPV4ONLY_ADDR));
+        final TestDnsPacket dns = new TestDnsPacket(
+                new DnsHeader(dnsId, FLAG, qlist.size(), alist.size()), qlist, alist);
+
+        // [2] Build IPv6 UDP DNS packet.
+        return buildUdpPacket(srcIp, dstIp, srcPort, dstPort, ByteBuffer.wrap(dns.getBytes()));
+    }
+
+    private void maybeReplyUdpDnsPrefix64Discovery(@NonNull byte[] packet) {
+        final ByteBuffer buf = ByteBuffer.wrap(packet);
+
+        // [1] Parse the prefix64 discovery DNS query for hostname ipv4only.arpa.
+        // Parse IPv6 and UDP header.
+        Ipv6Header ipv6Header = null;
+        try {
+            ipv6Header = Struct.parse(Ipv6Header.class, buf);
+            if (ipv6Header == null || ipv6Header.nextHeader != IPPROTO_UDP) return;
+        } catch (Exception e) {
+            // Parsing packet fail means it is not IPv6 UDP packet.
+            return;
+        }
+        final UdpHeader udpHeader = Struct.parse(UdpHeader.class, buf);
+
+        // Parse DNS message.
+        final TestDnsPacket pref64Query = TestDnsPacket.getTestDnsPacket(buf);
+        if (pref64Query == null) return;
+        if (pref64Query.getHeader().isResponse()) return;
+        if (pref64Query.getQDCount() != 1) return;
+        if (pref64Query.getANCount() != 0) return;
+        if (pref64Query.getNSCount() != 0) return;
+        if (pref64Query.getARCount() != 0) return;
+
+        final List<DnsRecord> qdRecordList = pref64Query.getRecordList(QDSECTION);
+        if (qdRecordList.size() != 1) return;
+        if (!qdRecordList.get(0).dName.equals(PREF64_IPV4ONLY_HOSTNAME)) return;
+
+        // [2] Build prefix64 DNS discovery reply from received query.
+        // DNS response transaction id must be copied from DNS query. Used by the requester
+        // to match up replies to outstanding queries. See RFC 1035 section 4.1.1. Also reverse
+        // the source/destination address/port of query packet for building reply packet.
+        final ByteBuffer replyPacket;
+        try {
+            replyPacket = buildUdpDnsPrefix64ReplyPacket(pref64Query.getHeader().getId(),
+                    ipv6Header.dstIp /* srcIp */, ipv6Header.srcIp /* dstIp */,
+                    (short) udpHeader.dstPort /* srcPort */,
+                    (short) udpHeader.srcPort /* dstPort */);
+        } catch (Exception e) {
+            fail("Failed to build prefix64 discovery reply for " + ipv6Header.srcIp + ": " + e);
+            return;
+        }
+
+        Log.d(TAG, "Sending prefix64 discovery reply");
+        try {
+            sendDownloadPacket(replyPacket);
+        } catch (Exception e) {
+            fail("Failed to reply prefix64 discovery for " + ipv6Header.srcIp + ": " + e);
+        }
+    }
+
     private byte[] getUploadPacket(Predicate<byte[]> filter) {
         assertNotNull("Can't deal with upstream interface in local only mode", mUpstreamReader);
 
-        return mUpstreamReader.poll(PACKET_READ_TIMEOUT_MS, filter);
+        byte[] packet;
+        while ((packet = mUpstreamReader.poll(PACKET_READ_TIMEOUT_MS)) != null) {
+            if (filter.test(packet)) return packet;
+
+            maybeReplyUdpDnsPrefix64Discovery(packet);
+        }
+        return null;
     }
 
     private @NonNull byte[] verifyPacketNotNull(String message, @Nullable byte[] packet) {
diff --git a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
index 6ba2033..368860e 100644
--- a/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/AutomaticOnOffKeepaliveTracker.java
@@ -291,6 +291,18 @@
             }
         }
 
+        /**
+         * Construct a new AutomaticOnOffKeepalive from existing AutomaticOnOffKeepalive with a
+         * new KeepaliveInfo.
+         */
+        public AutomaticOnOffKeepalive withKeepaliveInfo(KeepaliveTracker.KeepaliveInfo ki)
+                throws InvalidSocketException {
+            return new AutomaticOnOffKeepalive(
+                    ki,
+                    mAutomaticOnOffState != STATE_ALWAYS_ON /* autoOnOff */,
+                    mUnderpinnedNetwork);
+        }
+
         @Override
         public String toString() {
             return "AutomaticOnOffKeepalive [ "
@@ -470,13 +482,26 @@
      * The message is expected to contain a KeepaliveTracker.KeepaliveInfo.
      */
     public void handleStartKeepalive(Message message) {
-        final AutomaticOnOffKeepalive autoKi = (AutomaticOnOffKeepalive) message.obj;
-        final int error = mKeepaliveTracker.handleStartKeepalive(autoKi.mKi);
+        final AutomaticOnOffKeepalive target = (AutomaticOnOffKeepalive) message.obj;
+        final Pair<Integer, KeepaliveTracker.KeepaliveInfo> res =
+                mKeepaliveTracker.handleStartKeepalive(target.mKi);
+        final int error = res.first;
         if (error != SUCCESS) {
-            mEventLog.log("Failed to start keepalive " + autoKi.mCallback + " on "
-                    + autoKi.getNetwork() + " with error " + error);
+            mEventLog.log("Failed to start keepalive " + target.mCallback + " on "
+                    + target.getNetwork() + " with error " + error);
             return;
         }
+        // Generate a new auto ki with the started keepalive info.
+        final AutomaticOnOffKeepalive autoKi;
+        try {
+            autoKi = target.withKeepaliveInfo(res.second);
+            // Close the duplicated fd.
+            target.close();
+        } catch (InvalidSocketException e) {
+            Log.wtf(TAG, "Fail to create AutomaticOnOffKeepalive", e);
+            return;
+        }
+
         mEventLog.log("Start keepalive " + autoKi.mCallback + " on " + autoKi.getNetwork());
         mKeepaliveStatsTracker.onStartKeepalive(
                 autoKi.getNetwork(),
@@ -506,14 +531,19 @@
      * @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);
+        final Pair<Integer, KeepaliveTracker.KeepaliveInfo> res =
+                mKeepaliveTracker.handleStartKeepalive(ki);
+        final KeepaliveTracker.KeepaliveInfo startedKi = res.second;
+        final int error = res.first;
         if (error != SUCCESS) {
-            mEventLog.log("Failed to resume keepalive " + ki.mCallback + " on " + ki.mNai
-                    + " with error " + error);
+            mEventLog.log("Failed to resume keepalive " + startedKi.mCallback + " on "
+                    + startedKi.mNai + " with error " + error);
             return error;
         }
-        mKeepaliveStatsTracker.onResumeKeepalive(ki.getNai().network(), ki.getSlot());
-        mEventLog.log("Resumed successfully keepalive " + ki.mCallback + " on " + ki.mNai);
+
+        mKeepaliveStatsTracker.onResumeKeepalive(startedKi.getNai().network(), startedKi.getSlot());
+        mEventLog.log("Resumed successfully keepalive " + startedKi.mCallback
+                + " on " + startedKi.mNai);
 
         return SUCCESS;
     }
diff --git a/service/src/com/android/server/connectivity/KeepaliveTracker.java b/service/src/com/android/server/connectivity/KeepaliveTracker.java
index 76e97e2..125c269 100644
--- a/service/src/com/android/server/connectivity/KeepaliveTracker.java
+++ b/service/src/com/android/server/connectivity/KeepaliveTracker.java
@@ -54,6 +54,7 @@
 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;
@@ -62,6 +63,8 @@
 import com.android.net.module.util.IpUtils;
 
 import java.io.FileDescriptor;
+import java.net.Inet4Address;
+import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
@@ -292,11 +295,15 @@
 
         private int checkSourceAddress() {
             // Check that we have the source address.
-            for (InetAddress address : mNai.linkProperties.getAddresses()) {
+            for (InetAddress address : mNai.linkProperties.getAllAddresses()) {
                 if (address.equals(mPacket.getSrcAddress())) {
                     return SUCCESS;
                 }
             }
+            // Or the address is the clat source address.
+            if (mPacket.getSrcAddress().equals(mNai.getClatv6SrcAddress())) {
+                return SUCCESS;
+            }
             return ERROR_INVALID_IP_ADDRESS;
         }
 
@@ -479,6 +486,15 @@
             return new KeepaliveInfo(mCallback, mNai, mPacket, mPid, mUid, mInterval, mType,
                     fd, mSlot, true /* resumed */);
         }
+
+        /**
+         * Construct a new KeepaliveInfo from existing KeepaliveInfo with a new KeepalivePacketData.
+         */
+        public KeepaliveInfo withPacketData(@NonNull KeepalivePacketData packet)
+                throws InvalidSocketException {
+            return new KeepaliveInfo(mCallback, mNai, packet, mPid, mUid, mInterval, mType,
+                    mFd, mSlot, mResumed);
+        }
     }
 
     void notifyErrorCallback(ISocketKeepaliveCallback cb, int error) {
@@ -512,15 +528,47 @@
      * 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.
+     * @return Pair of (SUCCESS if the keepalive is successfully starting and the error reason
+     *         otherwise, the started KeepaliveInfo object)
      */
-    public int handleStartKeepalive(KeepaliveInfo ki) {
-        NetworkAgentInfo nai = ki.getNai();
+    public Pair<Integer, KeepaliveInfo> handleStartKeepalive(KeepaliveInfo ki) {
+        final KeepaliveInfo newKi;
+        try {
+            newKi = handleUpdateKeepaliveForClat(ki);
+        } catch (InvalidSocketException | InvalidPacketException e) {
+            Log.e(TAG, "Fail to construct keepalive packet");
+            notifyErrorCallback(ki.mCallback, ERROR_INVALID_IP_ADDRESS);
+            // Fail to create new keepalive packet for clat. Return the original keepalive info.
+            return new Pair<>(ERROR_INVALID_IP_ADDRESS, ki);
+        }
+
+        final NetworkAgentInfo nai = newKi.getNai();
         // If this was a paused keepalive, then reuse the same slot that was kept for it. Otherwise,
         // use the first free slot for this network agent.
-        final int slot = NO_KEEPALIVE != ki.mSlot ? ki.mSlot : findFirstFreeSlot(nai);
-        mKeepalives.get(nai).put(slot, ki);
-        return ki.start(slot);
+        final int slot = NO_KEEPALIVE != newKi.mSlot ? newKi.mSlot : findFirstFreeSlot(nai);
+        mKeepalives.get(nai).put(slot, newKi);
+
+        return new Pair<>(newKi.start(slot), newKi);
+    }
+
+    private KeepaliveInfo handleUpdateKeepaliveForClat(KeepaliveInfo ki)
+            throws InvalidSocketException, InvalidPacketException {
+        // Only try to translate address if the packet source address is the clat's source address.
+        if (!ki.mPacket.getSrcAddress().equals(ki.getNai().getClatv4SrcAddress())) return ki;
+
+        final InetAddress dstAddr = ki.mPacket.getDstAddress();
+        // Do not perform translation for a v6 dst address.
+        if (!(dstAddr instanceof Inet4Address)) return ki;
+
+        final Inet6Address address = ki.getNai().translateV4toClatV6((Inet4Address) dstAddr);
+
+        if (address == null) return ki;
+
+        final int srcPort = ki.mPacket.getSrcPort();
+        final KeepaliveInfo newInfo = ki.withPacketData(NattKeepalivePacketData.nattKeepalivePacket(
+                ki.getNai().getClatv6SrcAddress(), srcPort, address, NATT_PORT));
+        Log.d(TAG, "Src is clat v4 address. Convert from " + ki + " to " + newInfo);
+        return newInfo;
     }
 
     public void handleStopAllKeepalives(NetworkAgentInfo nai, int reason) {
diff --git a/service/src/com/android/server/connectivity/Nat464Xlat.java b/service/src/com/android/server/connectivity/Nat464Xlat.java
index b315235..f9e07fd 100644
--- a/service/src/com/android/server/connectivity/Nat464Xlat.java
+++ b/service/src/com/android/server/connectivity/Nat464Xlat.java
@@ -567,7 +567,7 @@
         try {
             return (Inet6Address) Inet6Address.getByAddress(v6Addr);
         } catch (UnknownHostException e) {
-            Log.e(TAG, "getByAddress should never throw for a numeric address");
+            Log.wtf(TAG, "getByAddress should never throw for a numeric address", e);
             return null;
         }
     }
@@ -583,6 +583,21 @@
         return mIPv6Address;
     }
 
+    /**
+     * Get the generated v4 address of clat.
+     */
+    @Nullable
+    public Inet4Address getClatv4SrcAddress() {
+        // Variables in Nat464Xlat should only be accessed from handler thread.
+        ensureRunningOnHandlerThread();
+        if (!isStarted()) return null;
+
+        final LinkAddress v4Addr = getLinkAddress(mIface);
+        if (v4Addr == null) return null;
+
+        return (Inet4Address) v4Addr.getAddress();
+    }
+
     private void ensureRunningOnHandlerThread() {
         if (mNetwork.handler().getLooper().getThread() != Thread.currentThread()) {
             throw new IllegalStateException(
diff --git a/service/src/com/android/server/connectivity/NetworkAgentInfo.java b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
index 08c1455..845c04c 100644
--- a/service/src/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/service/src/com/android/server/connectivity/NetworkAgentInfo.java
@@ -1043,6 +1043,14 @@
     }
 
     /**
+     * Get the generated v4 address of clat.
+     */
+    @Nullable
+    public Inet4Address getClatv4SrcAddress() {
+        return clatd.getClatv4SrcAddress();
+    }
+
+    /**
      * Translate the input v4 address to v6 clat address.
      */
     @Nullable
diff --git a/tests/unit/java/android/net/nsd/NsdServiceInfoTest.java b/tests/common/java/android/net/nsd/NsdServiceInfoTest.java
similarity index 88%
rename from tests/unit/java/android/net/nsd/NsdServiceInfoTest.java
rename to tests/common/java/android/net/nsd/NsdServiceInfoTest.java
index 9ce0693..ffe0e91 100644
--- a/tests/unit/java/android/net/nsd/NsdServiceInfoTest.java
+++ b/tests/common/java/android/net/nsd/NsdServiceInfoTest.java
@@ -26,10 +26,10 @@
 import android.os.Build;
 import android.os.Bundle;
 import android.os.Parcel;
-import android.os.StrictMode;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.testutils.ConnectivityModuleTest;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
 
@@ -37,7 +37,6 @@
 import org.junit.runner.RunWith;
 
 import java.net.InetAddress;
-import java.net.UnknownHostException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -45,22 +44,11 @@
 @RunWith(DevSdkIgnoreRunner.class)
 @SmallTest
 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
+@ConnectivityModuleTest
 public class NsdServiceInfoTest {
 
     private static final InetAddress IPV4_ADDRESS = InetAddresses.parseNumericAddress("192.0.2.1");
     private static final InetAddress IPV6_ADDRESS = InetAddresses.parseNumericAddress("2001:db8::");
-    public final static InetAddress LOCALHOST;
-    static {
-        // Because test.
-        StrictMode.ThreadPolicy policy = new StrictMode.ThreadPolicy.Builder().permitAll().build();
-        StrictMode.setThreadPolicy(policy);
-
-        InetAddress _host = null;
-        try {
-            _host = InetAddress.getLocalHost();
-        } catch (UnknownHostException e) { }
-        LOCALHOST = _host;
-    }
 
     @Test
     public void testLimits() throws Exception {
@@ -89,10 +77,10 @@
         // Single key + value length too long.
         exceptionThrown = false;
         try {
-            String longValue = "loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
-                    "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
-                    "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo" +
-                    "ooooooooooooooooooooooooooooong";  // 248 characters.
+            String longValue = "loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
+                    + "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
+                    + "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
+                    + "ooooooooooooooooooooooooooooong";  // 248 characters.
             info.setAttribute("longcat", longValue);  // Key + value == 255 characters.
         } catch (IllegalArgumentException e) {
             exceptionThrown = true;
@@ -127,7 +115,6 @@
         fullInfo.setServiceName("kitten");
         fullInfo.setServiceType("_kitten._tcp");
         fullInfo.setPort(4242);
-        fullInfo.setHost(LOCALHOST);
         fullInfo.setHostAddresses(List.of(IPV4_ADDRESS));
         fullInfo.setNetwork(new Network(123));
         fullInfo.setInterfaceIndex(456);
@@ -143,8 +130,7 @@
         attributedInfo.setServiceName("kitten");
         attributedInfo.setServiceType("_kitten._tcp");
         attributedInfo.setPort(4242);
-        attributedInfo.setHost(LOCALHOST);
-        fullInfo.setHostAddresses(List.of(IPV6_ADDRESS, IPV4_ADDRESS));
+        attributedInfo.setHostAddresses(List.of(IPV6_ADDRESS, IPV4_ADDRESS));
         attributedInfo.setAttribute("color", "pink");
         attributedInfo.setAttribute("sound", (new String("にゃあ")).getBytes("UTF-8"));
         attributedInfo.setAttribute("adorable", (String) null);
diff --git a/tests/cts/net/src/android/net/cts/RateLimitTest.java b/tests/cts/net/src/android/net/cts/RateLimitTest.java
index 36b98fc..5c93738 100644
--- a/tests/cts/net/src/android/net/cts/RateLimitTest.java
+++ b/tests/cts/net/src/android/net/cts/RateLimitTest.java
@@ -36,6 +36,7 @@
 import android.icu.text.MessageFormat;
 import android.net.ConnectivityManager;
 import android.net.ConnectivitySettingsManager;
+import android.net.ConnectivityThread;
 import android.net.InetAddresses;
 import android.net.IpPrefix;
 import android.net.LinkAddress;
@@ -189,7 +190,19 @@
             // whatever happens, don't leave the device in rate limited state.
             ConnectivitySettingsManager.setIngressRateLimitInBytesPerSecond(mContext, -1);
         }
-        if (mSocket != null) mSocket.close();
+        if (mSocket == null) {
+            // HACK(b/272147742): dump ConnectivityThread if test initialization failed.
+            final StackTraceElement[] elements = ConnectivityThread.get().getStackTrace();
+            final StringBuilder sb = new StringBuilder();
+            // Skip first element as it includes the invocation of getStackTrace()
+            for (int i = 1; i < elements.length; i++) {
+                sb.append(elements[i]);
+                sb.append("\n");
+            }
+            Log.e(TAG, sb.toString());
+        } else {
+            mSocket.close();
+        }
         if (mNetworkAgent != null) mNetworkAgent.unregister();
         if (mTunInterface != null) mTunInterface.getFileDescriptor().close();
         if (mCm != null) mCm.unregisterNetworkCallback(mNetworkCallback);
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 9406edd..bbde9b4 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -6857,17 +6857,19 @@
 
     @Test
     public void testPacketKeepalives() throws Exception {
-        InetAddress myIPv4 = InetAddress.getByName("192.0.2.129");
+        final LinkAddress v4Addr = new LinkAddress("192.0.2.129/24");
+        final InetAddress myIPv4 = v4Addr.getAddress();
         InetAddress notMyIPv4 = InetAddress.getByName("192.0.2.35");
         InetAddress myIPv6 = InetAddress.getByName("2001:db8::1");
         InetAddress dstIPv4 = InetAddress.getByName("8.8.8.8");
         InetAddress dstIPv6 = InetAddress.getByName("2001:4860:4860::8888");
-
+        doReturn(getClatInterfaceConfigParcel(v4Addr)).when(mMockNetd)
+                .interfaceGetCfg(CLAT_MOBILE_IFNAME);
         final int validKaInterval = 15;
         final int invalidKaInterval = 9;
 
         LinkProperties lp = new LinkProperties();
-        lp.setInterfaceName("wlan12");
+        lp.setInterfaceName(MOBILE_IFNAME);
         lp.addLinkAddress(new LinkAddress(myIPv6, 64));
         lp.addLinkAddress(new LinkAddress(myIPv4, 25));
         lp.addRoute(new RouteInfo(InetAddress.getByName("fe80::1234")));
diff --git a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
index 9e604e3..eeffbe1 100644
--- a/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
+++ b/tests/unit/java/com/android/server/connectivity/AutomaticOnOffKeepaliveTrackerTest.java
@@ -52,6 +52,7 @@
 import android.content.res.Resources;
 import android.net.INetd;
 import android.net.ISocketKeepaliveCallback;
+import android.net.InetAddresses;
 import android.net.KeepalivePacketData;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
@@ -116,7 +117,8 @@
     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 static final byte[] V4_SRC_ADDR = new byte[] { (byte) 192, 0, 0, (byte) 129 };
+    private static final String TEST_V4_IFACE = "v4-testIface";
     private AutomaticOnOffKeepaliveTracker mAOOKeepaliveTracker;
     private HandlerThread mHandlerThread;
 
@@ -327,6 +329,8 @@
                 NetworkInfo.DetailedState.CONNECTED, "test reason", "test extra info");
         doReturn(new Network(TEST_NETID)).when(mNai).network();
         mNai.linkProperties = new LinkProperties();
+        doReturn(null).when(mNai).translateV4toClatV6(any());
+        doReturn(null).when(mNai).getClatv6SrcAddress();
 
         doReturn(PERMISSION_GRANTED).when(mCtx).checkPermission(any() /* permission */,
                 anyInt() /* pid */, anyInt() /* uid */);
@@ -429,8 +433,7 @@
     }
 
     private TestKeepaliveInfo doStartNattKeepalive(int intervalSeconds) throws Exception {
-        final InetAddress srcAddress = InetAddress.getByAddress(
-                new byte[] { (byte) 192, 0, 0, (byte) 129 });
+        final InetAddress srcAddress = InetAddress.getByAddress(V4_SRC_ADDR);
         final int srcPort = 12345;
         final InetAddress dstAddress = InetAddress.getByAddress(new byte[] {8, 8, 8, 8});
         final int dstPort = 12345;
@@ -610,6 +613,42 @@
     }
 
     @Test
+    public void testStartNattKeepalive_addressTranslationOnClat() throws Exception {
+        final InetAddress v6AddrSrc = InetAddresses.parseNumericAddress("2001:db8::1");
+        final InetAddress v6AddrDst = InetAddresses.parseNumericAddress("2001:db8::2");
+        doReturn(v6AddrDst).when(mNai).translateV4toClatV6(any());
+        doReturn(v6AddrSrc).when(mNai).getClatv6SrcAddress();
+        doReturn(InetAddress.getByAddress(V4_SRC_ADDR)).when(mNai).getClatv4SrcAddress();
+        // Setup nai to add clat address
+        final LinkProperties stacked = new LinkProperties();
+        stacked.setInterfaceName(TEST_V4_IFACE);
+        mNai.linkProperties.addStackedLink(stacked);
+
+        final TestKeepaliveInfo testInfo = doStartNattKeepalive();
+        final ArgumentCaptor<NattKeepalivePacketData> kpdCaptor =
+                ArgumentCaptor.forClass(NattKeepalivePacketData.class);
+        verify(mNai).onStartNattSocketKeepalive(
+                eq(TEST_SLOT), eq(TEST_KEEPALIVE_INTERVAL_SEC), kpdCaptor.capture());
+        final NattKeepalivePacketData kpd = kpdCaptor.getValue();
+        // Verify the addresses are updated to v6 when clat is started.
+        assertEquals(v6AddrSrc, kpd.getSrcAddress());
+        assertEquals(v6AddrDst, kpd.getDstAddress());
+
+        triggerEventKeepalive(TEST_SLOT, SocketKeepalive.SUCCESS);
+        verify(testInfo.socketKeepaliveCallback).onStarted();
+
+        // Remove clat address should stop the keepalive.
+        doReturn(null).when(mNai).getClatv6SrcAddress();
+        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 testHandleEventSocketKeepalive_startingFailureHardwareError() throws Exception {
         final TestKeepaliveInfo testInfo = doStartNattKeepalive();