Merge "Revert "Change the visibility of Config to be public, so it can be accessed from"" into main
diff --git a/bpf/loader/NetBpfLoad.cpp b/bpf/loader/NetBpfLoad.cpp
index 9a049c7..c2a1d6e 100644
--- a/bpf/loader/NetBpfLoad.cpp
+++ b/bpf/loader/NetBpfLoad.cpp
@@ -1288,6 +1288,8 @@
 
 #define APEX_MOUNT_POINT "/apex/com.android.tethering"
 const char * const platformBpfLoader = "/system/bin/bpfloader";
+const char *const uprobestatsBpfLoader =
+    "/apex/com.android.uprobestats/bin/uprobestatsbpfload";
 
 static int logTetheringApexVersion(void) {
     char * found_blockdev = NULL;
@@ -1657,8 +1659,17 @@
     }
 
     // unreachable before U QPR3
-    ALOGI("done, transferring control to platform bpfloader.");
+    {
+      ALOGI("done, transferring control to uprobestatsbpfload.");
+      const char *args[] = {
+          uprobestatsBpfLoader,
+          NULL,
+      };
+      execve(args[0], (char **)args, envp);
+    }
 
+    ALOGI("unable to execute uprobestatsbpfload, transferring control to "
+          "platform bpfloader.");
     // platform BpfLoader *needs* to run as root
     const char * args[] = { platformBpfLoader, NULL, };
     execve(args[0], (char**)args, envp);
diff --git a/bpf/netd/BpfHandler.cpp b/bpf/netd/BpfHandler.cpp
index 8e4c2c6..50e0329 100644
--- a/bpf/netd/BpfHandler.cpp
+++ b/bpf/netd/BpfHandler.cpp
@@ -120,18 +120,22 @@
     }
 
     if (modules::sdklevel::IsAtLeastV()) {
-        RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_CONNECT4_PROG_PATH,
-                                    cg_fd, BPF_CGROUP_INET4_CONNECT));
-        RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_CONNECT6_PROG_PATH,
-                                    cg_fd, BPF_CGROUP_INET6_CONNECT));
-        RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP4_RECVMSG_PROG_PATH,
-                                    cg_fd, BPF_CGROUP_UDP4_RECVMSG));
-        RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP6_RECVMSG_PROG_PATH,
-                                    cg_fd, BPF_CGROUP_UDP6_RECVMSG));
-        RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP4_SENDMSG_PROG_PATH,
-                                    cg_fd, BPF_CGROUP_UDP4_SENDMSG));
-        RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP6_SENDMSG_PROG_PATH,
-                                    cg_fd, BPF_CGROUP_UDP6_SENDMSG));
+        // V requires 4.19+, so technically this 2nd 'if' is not required, but it
+        // doesn't hurt us to try to support AOSP forks that try to support older kernels.
+        if (bpf::isAtLeastKernelVersion(4, 19, 0)) {
+            RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_CONNECT4_PROG_PATH,
+                                        cg_fd, BPF_CGROUP_INET4_CONNECT));
+            RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_CONNECT6_PROG_PATH,
+                                        cg_fd, BPF_CGROUP_INET6_CONNECT));
+            RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP4_RECVMSG_PROG_PATH,
+                                        cg_fd, BPF_CGROUP_UDP4_RECVMSG));
+            RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP6_RECVMSG_PROG_PATH,
+                                        cg_fd, BPF_CGROUP_UDP6_RECVMSG));
+            RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP4_SENDMSG_PROG_PATH,
+                                        cg_fd, BPF_CGROUP_UDP4_SENDMSG));
+            RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_UDP6_SENDMSG_PROG_PATH,
+                                        cg_fd, BPF_CGROUP_UDP6_SENDMSG));
+        }
 
         if (bpf::isAtLeastKernelVersion(5, 4, 0)) {
             RETURN_IF_NOT_OK(attachProgramToCgroup(CGROUP_GETSOCKOPT_PROG_PATH,
@@ -161,12 +165,16 @@
     }
 
     if (modules::sdklevel::IsAtLeastV()) {
-        if (bpf::queryProgram(cg_fd, BPF_CGROUP_INET4_CONNECT) <= 0) abort();
-        if (bpf::queryProgram(cg_fd, BPF_CGROUP_INET6_CONNECT) <= 0) abort();
-        if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP4_RECVMSG) <= 0) abort();
-        if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP6_RECVMSG) <= 0) abort();
-        if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP4_SENDMSG) <= 0) abort();
-        if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP6_SENDMSG) <= 0) abort();
+        // V requires 4.19+, so technically this 2nd 'if' is not required, but it
+        // doesn't hurt us to try to support AOSP forks that try to support older kernels.
+        if (bpf::isAtLeastKernelVersion(4, 19, 0)) {
+            if (bpf::queryProgram(cg_fd, BPF_CGROUP_INET4_CONNECT) <= 0) abort();
+            if (bpf::queryProgram(cg_fd, BPF_CGROUP_INET6_CONNECT) <= 0) abort();
+            if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP4_RECVMSG) <= 0) abort();
+            if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP6_RECVMSG) <= 0) abort();
+            if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP4_SENDMSG) <= 0) abort();
+            if (bpf::queryProgram(cg_fd, BPF_CGROUP_UDP6_SENDMSG) <= 0) abort();
+        }
 
         if (bpf::isAtLeastKernelVersion(5, 4, 0)) {
             if (bpf::queryProgram(cg_fd, BPF_CGROUP_GETSOCKOPT) <= 0) abort();
diff --git a/bpf/progs/netd.c b/bpf/progs/netd.c
index cbe856d..ed0eed5 100644
--- a/bpf/progs/netd.c
+++ b/bpf/progs/netd.c
@@ -709,32 +709,32 @@
     return block_port(ctx);
 }
 
-DEFINE_NETD_V_BPF_PROG_KVER("connect4/inet4_connect", AID_ROOT, AID_ROOT, inet4_connect, KVER_4_14)
+DEFINE_NETD_V_BPF_PROG_KVER("connect4/inet4_connect", AID_ROOT, AID_ROOT, inet4_connect, KVER_4_19)
 (struct bpf_sock_addr *ctx) {
     return check_localhost(ctx);
 }
 
-DEFINE_NETD_V_BPF_PROG_KVER("connect6/inet6_connect", AID_ROOT, AID_ROOT, inet6_connect, KVER_4_14)
+DEFINE_NETD_V_BPF_PROG_KVER("connect6/inet6_connect", AID_ROOT, AID_ROOT, inet6_connect, KVER_4_19)
 (struct bpf_sock_addr *ctx) {
     return check_localhost(ctx);
 }
 
-DEFINE_NETD_V_BPF_PROG_KVER("recvmsg4/udp4_recvmsg", AID_ROOT, AID_ROOT, udp4_recvmsg, KVER_4_14)
+DEFINE_NETD_V_BPF_PROG_KVER("recvmsg4/udp4_recvmsg", AID_ROOT, AID_ROOT, udp4_recvmsg, KVER_4_19)
 (struct bpf_sock_addr *ctx) {
     return check_localhost(ctx);
 }
 
-DEFINE_NETD_V_BPF_PROG_KVER("recvmsg6/udp6_recvmsg", AID_ROOT, AID_ROOT, udp6_recvmsg, KVER_4_14)
+DEFINE_NETD_V_BPF_PROG_KVER("recvmsg6/udp6_recvmsg", AID_ROOT, AID_ROOT, udp6_recvmsg, KVER_4_19)
 (struct bpf_sock_addr *ctx) {
     return check_localhost(ctx);
 }
 
-DEFINE_NETD_V_BPF_PROG_KVER("sendmsg4/udp4_sendmsg", AID_ROOT, AID_ROOT, udp4_sendmsg, KVER_4_14)
+DEFINE_NETD_V_BPF_PROG_KVER("sendmsg4/udp4_sendmsg", AID_ROOT, AID_ROOT, udp4_sendmsg, KVER_4_19)
 (struct bpf_sock_addr *ctx) {
     return check_localhost(ctx);
 }
 
-DEFINE_NETD_V_BPF_PROG_KVER("sendmsg6/udp6_sendmsg", AID_ROOT, AID_ROOT, udp6_sendmsg, KVER_4_14)
+DEFINE_NETD_V_BPF_PROG_KVER("sendmsg6/udp6_sendmsg", AID_ROOT, AID_ROOT, udp6_sendmsg, KVER_4_19)
 (struct bpf_sock_addr *ctx) {
     return check_localhost(ctx);
 }
diff --git a/framework-t/src/android/net/TrafficStats.java b/framework-t/src/android/net/TrafficStats.java
index 1294b3e..5eed3b5 100644
--- a/framework-t/src/android/net/TrafficStats.java
+++ b/framework-t/src/android/net/TrafficStats.java
@@ -289,8 +289,8 @@
 
     private static class SocketTagger extends dalvik.system.SocketTagger {
 
-        // TODO: set to false
-        private static final boolean LOGD = true;
+        // Enable log with `setprop log.tag.TrafficStats DEBUG` and restart the module.
+        private static final boolean LOGD = Log.isLoggable(TAG, Log.DEBUG);
 
         SocketTagger() {
         }
diff --git a/service-t/src/com/android/server/net/TrafficStatsRateLimitCache.java b/service-t/src/com/android/server/net/TrafficStatsRateLimitCache.java
index ca97d07..667aad1 100644
--- a/service-t/src/com/android/server/net/TrafficStatsRateLimitCache.java
+++ b/service-t/src/com/android/server/net/TrafficStatsRateLimitCache.java
@@ -19,9 +19,8 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.net.NetworkStats;
-import android.util.LruCache;
 
-import com.android.internal.annotations.GuardedBy;
+import com.android.net.module.util.LruCacheWithExpiry;
 
 import java.time.Clock;
 import java.util.Objects;
@@ -31,9 +30,8 @@
  * A thread-safe cache for storing and retrieving {@link NetworkStats.Entry} objects,
  * with an adjustable expiry duration to manage data freshness.
  */
-class TrafficStatsRateLimitCache {
-    private final Clock mClock;
-    private final long mExpiryDurationMs;
+class TrafficStatsRateLimitCache extends
+        LruCacheWithExpiry<TrafficStatsRateLimitCache.TrafficStatsCacheKey, NetworkStats.Entry> {
 
     /**
      * Constructs a new {@link TrafficStatsRateLimitCache} with the specified expiry duration.
@@ -43,19 +41,17 @@
      * @param maxSize Maximum number of entries.
      */
     TrafficStatsRateLimitCache(@NonNull Clock clock, long expiryDurationMs, int maxSize) {
-        mClock = clock;
-        mExpiryDurationMs = expiryDurationMs;
-        mMap = new LruCache<>(maxSize);
+        super(clock, expiryDurationMs, maxSize, it -> !it.isEmpty());
     }
 
-    private static class TrafficStatsCacheKey {
+    public static class TrafficStatsCacheKey {
         @Nullable
-        public final String iface;
-        public final int uid;
+        private final String mIface;
+        private final int mUid;
 
         TrafficStatsCacheKey(@Nullable String iface, int uid) {
-            this.iface = iface;
-            this.uid = uid;
+            this.mIface = iface;
+            this.mUid = uid;
         }
 
         @Override
@@ -63,29 +59,15 @@
             if (this == o) return true;
             if (!(o instanceof TrafficStatsCacheKey)) return false;
             TrafficStatsCacheKey that = (TrafficStatsCacheKey) o;
-            return uid == that.uid && Objects.equals(iface, that.iface);
+            return mUid == that.mUid && Objects.equals(mIface, that.mIface);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(iface, uid);
+            return Objects.hash(mIface, mUid);
         }
     }
 
-    private static class TrafficStatsCacheValue {
-        public final long timestamp;
-        @NonNull
-        public final NetworkStats.Entry entry;
-
-        TrafficStatsCacheValue(long timestamp, NetworkStats.Entry entry) {
-            this.timestamp = timestamp;
-            this.entry = entry;
-        }
-    }
-
-    @GuardedBy("mMap")
-    private final LruCache<TrafficStatsCacheKey, TrafficStatsCacheValue> mMap;
-
     /**
      * Retrieves a {@link NetworkStats.Entry} from the cache, associated with the given key.
      *
@@ -95,16 +77,7 @@
      */
     @Nullable
     NetworkStats.Entry get(String iface, int uid) {
-        final TrafficStatsCacheKey key = new TrafficStatsCacheKey(iface, uid);
-        synchronized (mMap) { // Synchronize for thread-safety
-            final TrafficStatsCacheValue value = mMap.get(key);
-            if (value != null && !isExpired(value.timestamp)) {
-                return value.entry;
-            } else {
-                mMap.remove(key); // Remove expired entries
-                return null;
-            }
-        }
+        return super.get(new TrafficStatsCacheKey(iface, uid));
     }
 
     /**
@@ -122,19 +95,7 @@
     @Nullable
     NetworkStats.Entry getOrCompute(String iface, int uid,
             @NonNull Supplier<NetworkStats.Entry> supplier) {
-        synchronized (mMap) {
-            final NetworkStats.Entry cachedValue = get(iface, uid);
-            if (cachedValue != null) {
-                return cachedValue;
-            }
-
-            // Entry not found or expired, compute it
-            final NetworkStats.Entry computedEntry = supplier.get();
-            if (computedEntry != null && !computedEntry.isEmpty()) {
-                put(iface, uid, computedEntry);
-            }
-            return computedEntry;
-        }
+        return super.getOrCompute(new TrafficStatsCacheKey(iface, uid), supplier);
     }
 
     /**
@@ -145,23 +106,7 @@
      * @param entry The {@link NetworkStats.Entry} to store in the cache.
      */
     void put(String iface, int uid, @NonNull final NetworkStats.Entry entry) {
-        Objects.requireNonNull(entry);
-        final TrafficStatsCacheKey key = new TrafficStatsCacheKey(iface, uid);
-        synchronized (mMap) { // Synchronize for thread-safety
-            mMap.put(key, new TrafficStatsCacheValue(mClock.millis(), entry));
-        }
+        super.put(new TrafficStatsCacheKey(iface, uid), entry);
     }
 
-    /**
-     * Clear the cache.
-     */
-    void clear() {
-        synchronized (mMap) {
-            mMap.evictAll();
-        }
-    }
-
-    private boolean isExpired(long timestamp) {
-        return mClock.millis() > timestamp + mExpiryDurationMs;
-    }
 }
diff --git a/staticlibs/framework/com/android/net/module/util/LruCacheWithExpiry.java b/staticlibs/framework/com/android/net/module/util/LruCacheWithExpiry.java
new file mode 100644
index 0000000..80088b9
--- /dev/null
+++ b/staticlibs/framework/com/android/net/module/util/LruCacheWithExpiry.java
@@ -0,0 +1,149 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.net.module.util;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.util.LruCache;
+
+import com.android.internal.annotations.GuardedBy;
+
+import java.time.Clock;
+import java.util.Objects;
+import java.util.function.Predicate;
+import java.util.function.Supplier;
+
+/**
+ * An LRU cache that stores key-value pairs with an expiry time.
+ *
+ * <p>This cache uses an {@link LruCache} to store entries and evicts the least
+ * recently used entries when the cache reaches its maximum capacity. It also
+ * supports an expiry time for each entry, allowing entries to be automatically
+ * removed from the cache after a certain duration.
+ *
+ * @param <K> The type of keys used to identify cached entries.
+ * @param <V> The type of values stored in the cache.
+ *
+ * @hide
+ */
+public class LruCacheWithExpiry<K, V> {
+    private final Clock mClock;
+    private final long mExpiryDurationMs;
+    @GuardedBy("mMap")
+    private final LruCache<K, CacheValue<V>> mMap;
+    private final Predicate<V> mShouldCacheValue;
+
+    /**
+     * Constructs a new {@link LruCacheWithExpiry} with the specified parameters.
+     *
+     * @param clock            The {@link Clock} to use for determining timestamps.
+     * @param expiryDurationMs The expiry duration for cached entries in milliseconds.
+     * @param maxSize          The maximum number of entries to hold in the cache.
+     * @param shouldCacheValue A {@link Predicate} that determines whether a given value should be
+     *                         cached. This can be used to filter out certain values from being
+     *                         stored in the cache.
+     */
+    public LruCacheWithExpiry(@NonNull Clock clock, long expiryDurationMs, int maxSize,
+            Predicate<V> shouldCacheValue) {
+        mClock = clock;
+        mExpiryDurationMs = expiryDurationMs;
+        mMap = new LruCache<>(maxSize);
+        mShouldCacheValue = shouldCacheValue;
+    }
+
+    /**
+     * Retrieves a value from the cache, associated with the given key.
+     *
+     * @param key The key to look up in the cache.
+     * @return The cached value, or {@code null} if not found or expired.
+     */
+    @Nullable
+    public V get(@NonNull K key) {
+        synchronized (mMap) {
+            final CacheValue<V> value = mMap.get(key);
+            if (value != null && !isExpired(value.timestamp)) {
+                return value.entry;
+            } else {
+                mMap.remove(key); // Remove expired entries
+                return null;
+            }
+        }
+    }
+
+    /**
+     * Retrieves a value from the cache, associated with the given key.
+     * If the entry is not found in the cache or has expired, computes it using the provided
+     * {@code supplier} and stores the result in the cache.
+     *
+     * @param key      The key to look up in the cache.
+     * @param supplier The {@link Supplier} to compute the value if not found or expired.
+     * @return The cached or computed value, or {@code null} if the {@code supplier} returns null.
+     */
+    @Nullable
+    public V getOrCompute(@NonNull K key, @NonNull Supplier<V> supplier) {
+        synchronized (mMap) {
+            final V cachedValue = get(key);
+            if (cachedValue != null) {
+                return cachedValue;
+            }
+
+            // Entry not found or expired, compute it
+            final V computedValue = supplier.get();
+            if (computedValue != null && mShouldCacheValue.test(computedValue)) {
+                put(key, computedValue);
+            }
+            return computedValue;
+        }
+    }
+
+    /**
+     * Stores a value in the cache, associated with the given key.
+     *
+     * @param key   The key to associate with the value.
+     * @param value The value to store in the cache.
+     */
+    public void put(@NonNull K key, @NonNull V value) {
+        Objects.requireNonNull(value);
+        synchronized (mMap) {
+            mMap.put(key, new CacheValue<>(mClock.millis(), value));
+        }
+    }
+
+    /**
+     * Clear the cache.
+     */
+    public void clear() {
+        synchronized (mMap) {
+            mMap.evictAll();
+        }
+    }
+
+    private boolean isExpired(long timestamp) {
+        return mClock.millis() > timestamp + mExpiryDurationMs;
+    }
+
+    private static class CacheValue<V> {
+        public final long timestamp;
+        @NonNull
+        public final V entry;
+
+        CacheValue(long timestamp, V entry) {
+            this.timestamp = timestamp;
+            this.entry = entry;
+        }
+    }
+}
diff --git a/tests/cts/net/Android.bp b/tests/cts/net/Android.bp
index a5ad7f2..9e57f69 100644
--- a/tests/cts/net/Android.bp
+++ b/tests/cts/net/Android.bp
@@ -70,6 +70,14 @@
         ":ConnectivityTestPreparer",
         ":CtsCarrierServicePackage",
     ],
+    errorprone: {
+        enabled: true,
+        // Error-prone checking only warns of problems when building. To make the build fail with
+        // these errors, list the specific error-prone problems below.
+        javacflags: [
+            "-Xep:NullablePrimitive:ERROR",
+        ],
+    },
 }
 
 // Networking CTS tests for development and release. These tests always target the platform SDK
diff --git a/thread/tests/integration/src/android/net/thread/InternetAccessTest.kt b/thread/tests/integration/src/android/net/thread/InternetAccessTest.kt
index 9b787f5..f3a09a0 100644
--- a/thread/tests/integration/src/android/net/thread/InternetAccessTest.kt
+++ b/thread/tests/integration/src/android/net/thread/InternetAccessTest.kt
@@ -32,6 +32,7 @@
 import android.net.thread.utils.IntegrationTestUtils.waitFor
 import android.net.thread.utils.OtDaemonController
 import android.net.thread.utils.TestDnsServer
+import android.net.thread.utils.TestUdpEchoServer
 import android.net.thread.utils.ThreadFeatureCheckerRule
 import android.net.thread.utils.ThreadFeatureCheckerRule.RequiresSimulationThreadDevice
 import android.net.thread.utils.ThreadFeatureCheckerRule.RequiresThreadFeature
@@ -48,6 +49,7 @@
 import com.google.common.truth.Truth.assertThat
 import java.net.Inet4Address
 import java.net.InetAddress
+import java.net.InetSocketAddress
 import java.time.Duration
 import org.junit.After
 import org.junit.Before
@@ -64,6 +66,7 @@
     private val TAG = BorderRoutingTest::class.java.simpleName
     private val NUM_FTD = 1
     private val DNS_SERVER_ADDR = parseNumericAddress("8.8.8.8") as Inet4Address
+    private val UDP_ECHO_SERVER_ADDRESS = InetSocketAddress(parseNumericAddress("1.2.3.4"), 12345)
     private val ANSWER_RECORDS =
         listOf(
             DnsPacket.DnsRecord.makeAOrAAAARecord(
@@ -94,6 +97,7 @@
     private lateinit var infraNetworkReader: PollPacketReader
     private lateinit var infraDevice: InfraNetworkDevice
     private lateinit var dnsServer: TestDnsServer
+    private lateinit var udpEchoServer: TestUdpEchoServer
 
     @Before
     @Throws(Exception::class)
@@ -111,13 +115,16 @@
         controller.setEnabledAndWait(true)
         controller.joinAndWait(DEFAULT_DATASET)
 
-        // Creates a infra network device.
+        // Create an infra network device.
         infraNetworkReader = newPacketReader(infraNetworkTracker.testIface, handler)
         infraDevice = startInfraDeviceAndWaitForOnLinkAddr(infraNetworkReader)
 
         // Create a DNS server
         dnsServer = TestDnsServer(infraNetworkReader, DNS_SERVER_ADDR, ANSWER_RECORDS)
 
+        // Create a UDP echo server
+        udpEchoServer = TestUdpEchoServer(infraNetworkReader, UDP_ECHO_SERVER_ADDRESS)
+
         // Create Ftds
         for (i in 0 until NUM_FTD) {
             ftds.add(FullThreadDevice(15 + i /* node ID */))
@@ -132,6 +139,7 @@
         tearDownInfraNetwork(infraNetworkTracker)
 
         dnsServer.stop()
+        udpEchoServer.stop()
 
         handlerThread.quitSafely()
         handlerThread.join()
@@ -166,6 +174,20 @@
         assertThat(ftd.resolveHost("google.com", TYPE_AAAA)).isEmpty()
     }
 
+    @Test
+    fun nat64Enabled_threadDeviceSendsUdpToEchoServer_replyIsReceived() {
+        controller.setNat64EnabledAndWait(true)
+        waitFor({ otCtl.hasNat64PrefixInNetdata() }, Duration.ofSeconds(10))
+        val ftd = ftds[0]
+        joinNetworkAndWaitForOmr(ftd, DEFAULT_DATASET)
+        udpEchoServer.start()
+
+        ftd.udpOpen()
+        ftd.udpSend("Hello,Thread", UDP_ECHO_SERVER_ADDRESS.address, UDP_ECHO_SERVER_ADDRESS.port)
+        val reply = ftd.udpReceive()
+        assertThat(reply).isEqualTo("Hello,Thread")
+    }
+
     private fun extractIpv4AddressFromMappedAddress(address: InetAddress): Inet4Address {
         return InetAddress.getByAddress(address.address.slice(12 until 16).toByteArray())
             as Inet4Address
diff --git a/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java b/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
index 5b532c7..209eed6 100644
--- a/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
+++ b/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
@@ -234,8 +234,8 @@
         return matcher.group(4);
     }
 
-    /** Sends a UDP message to given IPv6 address and port. */
-    public void udpSend(String message, Inet6Address serverAddr, int serverPort) {
+    /** Sends a UDP message to given IP address and port. */
+    public void udpSend(String message, InetAddress serverAddr, int serverPort) {
         executeCommand("udp send %s %d %s", serverAddr.getHostAddress(), serverPort, message);
     }
 
diff --git a/thread/tests/integration/src/android/net/thread/utils/TestDnsServer.kt b/thread/tests/integration/src/android/net/thread/utils/TestDnsServer.kt
index c52fc49..f97c0f2 100644
--- a/thread/tests/integration/src/android/net/thread/utils/TestDnsServer.kt
+++ b/thread/tests/integration/src/android/net/thread/utils/TestDnsServer.kt
@@ -16,18 +16,16 @@
 
 package android.net.thread.utils
 
-import android.net.thread.utils.IntegrationTestUtils.pollForPacket
 import android.system.OsConstants.IPPROTO_IP
 import android.system.OsConstants.IPPROTO_UDP
 import com.android.net.module.util.DnsPacket
 import com.android.net.module.util.PacketBuilder
-import com.android.net.module.util.Struct
 import com.android.net.module.util.structs.Ipv4Header
 import com.android.net.module.util.structs.UdpHeader
 import com.android.testutils.PollPacketReader
 import java.net.InetAddress
+import java.net.InetSocketAddress
 import java.nio.ByteBuffer
-import kotlin.concurrent.thread
 
 /**
  * A class that simulates a DNS server.
@@ -41,11 +39,12 @@
 class TestDnsServer(
     private val packetReader: PollPacketReader,
     private val serverAddress: InetAddress,
-    private val answerRecords: List<DnsPacket.DnsRecord>,
-) {
-    private val TAG = TestDnsServer::class.java.simpleName
-    private val DNS_UDP_PORT = 53
-    private var workerThread: Thread? = null
+    private val serverAnswers: List<DnsPacket.DnsRecord>,
+) : TestUdpServer(packetReader, InetSocketAddress(serverAddress, DNS_UDP_PORT)) {
+    companion object {
+        private val TAG = TestDnsServer::class.java.simpleName
+        private const val DNS_UDP_PORT = 53
+    }
 
     private class TestDnsPacket : DnsPacket {
 
@@ -61,49 +60,12 @@
         val records = super.mRecords
     }
 
-    /**
-     * Starts the DNS server to respond to DNS requests.
-     *
-     * <p> The server polls the DNS requests from the {@code packetReader} and responds with the
-     * {@code answerRecords}. The server will automatically stop when it fails to poll a DNS request
-     * within the timeout (3000 ms, as defined in IntegrationTestUtils).
-     */
-    fun start() {
-        workerThread = thread {
-            var requestPacket: ByteArray
-            while (true) {
-                requestPacket = pollForDnsPacket() ?: break
-                val buf = ByteBuffer.wrap(requestPacket)
-                packetReader.sendResponse(buildDnsResponse(buf, answerRecords))
-            }
-        }
-    }
-
-    /** Stops the DNS server. */
-    fun stop() {
-        workerThread?.join()
-    }
-
-    private fun pollForDnsPacket(): ByteArray? {
-        val filter =
-            fun(packet: ByteArray): Boolean {
-                val buf = ByteBuffer.wrap(packet)
-                val ipv4Header = Struct.parse(Ipv4Header::class.java, buf) ?: return false
-                val udpHeader = Struct.parse(UdpHeader::class.java, buf) ?: return false
-                return ipv4Header.dstIp == serverAddress && udpHeader.dstPort == DNS_UDP_PORT
-            }
-        return pollForPacket(packetReader, filter)
-    }
-
-    private fun buildDnsResponse(
-        requestPacket: ByteBuffer,
-        serverAnswers: List<DnsPacket.DnsRecord>,
+    override fun buildResponse(
+        requestIpv4Header: Ipv4Header,
+        requestUdpHeader: UdpHeader,
+        requestUdpPayload: ByteArray,
     ): ByteBuffer? {
-        val requestIpv4Header = Struct.parse(Ipv4Header::class.java, requestPacket) ?: return null
-        val requestUdpHeader = Struct.parse(UdpHeader::class.java, requestPacket) ?: return null
-        val remainingRequestPacket = ByteArray(requestPacket.remaining())
-        requestPacket.get(remainingRequestPacket)
-        val requestDnsPacket = TestDnsPacket(remainingRequestPacket)
+        val requestDnsPacket = TestDnsPacket(requestUdpPayload)
         val requestDnsHeader = requestDnsPacket.header
 
         val answerRecords =
diff --git a/thread/tests/integration/src/android/net/thread/utils/TestUdpEchoServer.kt b/thread/tests/integration/src/android/net/thread/utils/TestUdpEchoServer.kt
new file mode 100644
index 0000000..9fcd6a4
--- /dev/null
+++ b/thread/tests/integration/src/android/net/thread/utils/TestUdpEchoServer.kt
@@ -0,0 +1,74 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.thread.utils
+
+import android.system.OsConstants.IPPROTO_IP
+import android.system.OsConstants.IPPROTO_UDP
+import com.android.net.module.util.PacketBuilder
+import com.android.net.module.util.structs.Ipv4Header
+import com.android.net.module.util.structs.UdpHeader
+import com.android.testutils.PollPacketReader
+import java.net.InetSocketAddress
+import java.nio.ByteBuffer
+
+/**
+ * A class that simulates a UDP echo server that replies to incoming UDP message with the same
+ * payload.
+ *
+ * @param packetReader the packet reader to poll UDP requests from
+ * @param serverAddress the address and port of the UDP server
+ */
+class TestUdpEchoServer(
+    private val packetReader: PollPacketReader,
+    private val serverAddress: InetSocketAddress,
+) : TestUdpServer(packetReader, serverAddress) {
+    companion object {
+        private val TAG = TestUdpEchoServer::class.java.simpleName
+    }
+
+    override fun buildResponse(
+        requestIpv4Header: Ipv4Header,
+        requestUdpHeader: UdpHeader,
+        requestUdpPayload: ByteArray,
+    ): ByteBuffer? {
+        val buf =
+            PacketBuilder.allocate(
+                false /* hasEther */,
+                IPPROTO_IP,
+                IPPROTO_UDP,
+                requestUdpPayload.size,
+            )
+
+        val packetBuilder = PacketBuilder(buf)
+        packetBuilder.writeIpv4Header(
+            requestIpv4Header.tos,
+            requestIpv4Header.id,
+            requestIpv4Header.flagsAndFragmentOffset,
+            0x40 /* ttl */,
+            IPPROTO_UDP.toByte(),
+            requestIpv4Header.dstIp, /* srcIp */
+            requestIpv4Header.srcIp, /* dstIp */
+        )
+        packetBuilder.writeUdpHeader(
+            requestUdpHeader.dstPort.toShort() /* srcPort */,
+            requestUdpHeader.srcPort.toShort(), /* dstPort */
+        )
+        buf.put(requestUdpPayload)
+
+        return packetBuilder.finalizePacket()
+    }
+}
diff --git a/thread/tests/integration/src/android/net/thread/utils/TestUdpServer.kt b/thread/tests/integration/src/android/net/thread/utils/TestUdpServer.kt
new file mode 100644
index 0000000..fb0942e
--- /dev/null
+++ b/thread/tests/integration/src/android/net/thread/utils/TestUdpServer.kt
@@ -0,0 +1,98 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.thread.utils
+
+import android.net.thread.utils.IntegrationTestUtils.pollForPacket
+import com.android.net.module.util.Struct
+import com.android.net.module.util.structs.Ipv4Header
+import com.android.net.module.util.structs.UdpHeader
+import com.android.testutils.PollPacketReader
+import java.net.InetSocketAddress
+import java.nio.ByteBuffer
+import kotlin.concurrent.thread
+
+/**
+ * A class that simulates a UDP server that replies to incoming UDP messages.
+ *
+ * @param packetReader the packet reader to poll UDP requests from
+ * @param serverAddress the address and port of the UDP server
+ */
+abstract class TestUdpServer(
+    private val packetReader: PollPacketReader,
+    private val serverAddress: InetSocketAddress,
+) {
+    private val TAG = TestUdpServer::class.java.simpleName
+    private var workerThread: Thread? = null
+
+    /**
+     * Starts the UDP server to respond to UDP messages.
+     *
+     * <p> The server polls the UDP messages from the {@code packetReader} and responds with a
+     * message built by {@code buildResponse}. The server will automatically stop when it fails to
+     * poll a UDP request within the timeout (3000 ms, as defined in IntegrationTestUtils).
+     */
+    fun start() {
+        workerThread = thread {
+            var requestPacket: ByteArray
+            while (true) {
+                requestPacket = pollForUdpPacket() ?: break
+                val buf = ByteBuffer.wrap(requestPacket)
+                packetReader.sendResponse(buildResponse(buf) ?: break)
+            }
+        }
+    }
+
+    /** Stops the UDP server. */
+    fun stop() {
+        workerThread?.join()
+    }
+
+    /**
+     * Builds the UDP response for the given UDP request.
+     *
+     * @param ipv4Header the IPv4 header of the UDP request
+     * @param udpHeader the UDP header of the UDP request
+     * @param udpPayload the payload of the UDP request
+     * @return the UDP response
+     */
+    abstract fun buildResponse(
+        requestIpv4Header: Ipv4Header,
+        requestUdpHeader: UdpHeader,
+        requestUdpPayload: ByteArray,
+    ): ByteBuffer?
+
+    private fun pollForUdpPacket(): ByteArray? {
+        val filter =
+            fun(packet: ByteArray): Boolean {
+                val buf = ByteBuffer.wrap(packet)
+                val ipv4Header = Struct.parse(Ipv4Header::class.java, buf) ?: return false
+                val udpHeader = Struct.parse(UdpHeader::class.java, buf) ?: return false
+                return ipv4Header.dstIp == serverAddress.address &&
+                    udpHeader.dstPort == serverAddress.port
+            }
+        return pollForPacket(packetReader, filter)
+    }
+
+    private fun buildResponse(requestPacket: ByteBuffer): ByteBuffer? {
+        val requestIpv4Header = Struct.parse(Ipv4Header::class.java, requestPacket) ?: return null
+        val requestUdpHeader = Struct.parse(UdpHeader::class.java, requestPacket) ?: return null
+        val remainingRequestPacket = ByteArray(requestPacket.remaining())
+        requestPacket.get(remainingRequestPacket)
+
+        return buildResponse(requestIpv4Header, requestUdpHeader, remainingRequestPacket)
+    }
+}