diff --git a/Tethering/common/TetheringLib/api/module-lib-current.txt b/Tethering/common/TetheringLib/api/module-lib-current.txt
index a680590..460c216 100644
--- a/Tethering/common/TetheringLib/api/module-lib-current.txt
+++ b/Tethering/common/TetheringLib/api/module-lib-current.txt
@@ -46,10 +46,5 @@
     method @Deprecated @NonNull public java.util.List<java.lang.String> getTetherableWifiRegexs();
   }
 
-  public static final class TetheringManager.TetheringRequest implements android.os.Parcelable {
-    method @FlaggedApi("com.android.net.flags.tethering_request_with_soft_ap_config") @Nullable public String getPackageName();
-    method @FlaggedApi("com.android.net.flags.tethering_request_with_soft_ap_config") public int getUid();
-  }
-
 }
 
diff --git a/Tethering/common/TetheringLib/src/android/net/TetheringManager.java b/Tethering/common/TetheringLib/src/android/net/TetheringManager.java
index 7c7a4e0..411971d 100644
--- a/Tethering/common/TetheringLib/src/android/net/TetheringManager.java
+++ b/Tethering/common/TetheringLib/src/android/net/TetheringManager.java
@@ -33,7 +33,6 @@
 import android.os.IBinder;
 import android.os.Parcel;
 import android.os.Parcelable;
-import android.os.Process;
 import android.os.RemoteException;
 import android.os.ResultReceiver;
 import android.util.ArrayMap;
@@ -747,7 +746,6 @@
                 mBuilderParcel.exemptFromEntitlementCheck = false;
                 mBuilderParcel.showProvisioningUi = true;
                 mBuilderParcel.connectivityScope = getDefaultConnectivityScope(type);
-                mBuilderParcel.uid = Process.INVALID_UID;
                 mBuilderParcel.softApConfig = null;
             }
 
@@ -922,47 +920,6 @@
         }
 
         /**
-         * Sets the UID of the app that sent this request. This should always be overridden when
-         * receiving TetheringRequest from an external source.
-         * @hide
-         */
-        public void setUid(int uid) {
-            mRequestParcel.uid = uid;
-        }
-
-        /**
-         * Sets the package name of the app that sent this request. This should always be overridden
-         * when receiving a TetheringRequest from an external source.
-         * @hide
-         */
-        public void setPackageName(String packageName) {
-            mRequestParcel.packageName = packageName;
-        }
-
-        /**
-         * Gets the UID of the app that sent this request. This defaults to
-         * {@link Process#INVALID_UID} if unset.
-         * @hide
-         */
-        @FlaggedApi(Flags.FLAG_TETHERING_REQUEST_WITH_SOFT_AP_CONFIG)
-        @SystemApi(client = MODULE_LIBRARIES)
-        public int getUid() {
-            return mRequestParcel.uid;
-        }
-
-        /**
-         * Gets the package name of the app that sent this request. This defaults to {@code null} if
-         * unset.
-         * @hide
-         */
-        @FlaggedApi(Flags.FLAG_TETHERING_REQUEST_WITH_SOFT_AP_CONFIG)
-        @SystemApi(client = MODULE_LIBRARIES)
-        @Nullable
-        public String getPackageName() {
-            return mRequestParcel.packageName;
-        }
-
-        /**
          * Get a TetheringRequestParcel from the configuration
          * @hide
          */
@@ -978,8 +935,6 @@
                     + ", exemptFromEntitlementCheck= " + mRequestParcel.exemptFromEntitlementCheck
                     + ", showProvisioningUi= " + mRequestParcel.showProvisioningUi
                     + ", softApConfig= " + mRequestParcel.softApConfig
-                    + ", uid= " + mRequestParcel.uid
-                    + ", packageName= " + mRequestParcel.packageName
                     + " ]";
         }
 
@@ -995,9 +950,7 @@
                     && parcel.exemptFromEntitlementCheck == otherParcel.exemptFromEntitlementCheck
                     && parcel.showProvisioningUi == otherParcel.showProvisioningUi
                     && parcel.connectivityScope == otherParcel.connectivityScope
-                    && Objects.equals(parcel.softApConfig, otherParcel.softApConfig)
-                    && parcel.uid == otherParcel.uid
-                    && Objects.equals(parcel.packageName, otherParcel.packageName);
+                    && Objects.equals(parcel.softApConfig, otherParcel.softApConfig);
         }
 
         @Override
@@ -1005,8 +958,7 @@
             TetheringRequestParcel parcel = getParcel();
             return Objects.hash(parcel.tetheringType, parcel.localIPv4Address,
                     parcel.staticClientAddress, parcel.exemptFromEntitlementCheck,
-                    parcel.showProvisioningUi, parcel.connectivityScope, parcel.softApConfig,
-                    parcel.uid, parcel.packageName);
+                    parcel.showProvisioningUi, parcel.connectivityScope, parcel.softApConfig);
         }
     }
 
diff --git a/Tethering/common/TetheringLib/src/android/net/TetheringRequestParcel.aidl b/Tethering/common/TetheringLib/src/android/net/TetheringRequestParcel.aidl
index 789d5bb..ea7a353 100644
--- a/Tethering/common/TetheringLib/src/android/net/TetheringRequestParcel.aidl
+++ b/Tethering/common/TetheringLib/src/android/net/TetheringRequestParcel.aidl
@@ -31,6 +31,4 @@
     boolean showProvisioningUi;
     int connectivityScope;
     SoftApConfiguration softApConfig;
-    int uid;
-    String packageName;
 }
diff --git a/Tethering/src/com/android/networkstack/tethering/TetheringService.java b/Tethering/src/com/android/networkstack/tethering/TetheringService.java
index cea7e82..454cbf1 100644
--- a/Tethering/src/com/android/networkstack/tethering/TetheringService.java
+++ b/Tethering/src/com/android/networkstack/tethering/TetheringService.java
@@ -55,7 +55,6 @@
 import androidx.annotation.Nullable;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.net.module.util.PermissionUtils;
 import com.android.networkstack.apishim.SettingsShimImpl;
 import com.android.networkstack.apishim.common.SettingsShim;
 
@@ -139,10 +138,8 @@
                     listener)) {
                 return;
             }
-            TetheringRequest external = new TetheringRequest(request);
-            external.setUid(getBinderCallingUid());
-            external.setPackageName(callerPkg);
-            mTethering.startTethering(external, callerPkg, listener);
+            // TODO(b/216524590): Add UID/packageName of caller to TetheringRequest here
+            mTethering.startTethering(new TetheringRequest(request), callerPkg, listener);
         }
 
         @Override
@@ -241,12 +238,6 @@
                 final String callingAttributionTag, final boolean onlyAllowPrivileged,
                 final IIntResultListener listener) {
             try {
-                if (!checkPackageNameMatchesUid(getBinderCallingUid(), callerPkg)) {
-                    Log.e(TAG, "Package name " + callerPkg + " does not match UID "
-                            + getBinderCallingUid());
-                    listener.onResult(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION);
-                    return true;
-                }
                 if (!hasTetherChangePermission(callerPkg, callingAttributionTag,
                         onlyAllowPrivileged)) {
                     listener.onResult(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION);
@@ -265,12 +256,6 @@
 
         private boolean checkAndNotifyCommonError(final String callerPkg,
                 final String callingAttributionTag, final ResultReceiver receiver) {
-            if (!checkPackageNameMatchesUid(getBinderCallingUid(), callerPkg)) {
-                Log.e(TAG, "Package name " + callerPkg + " does not match UID "
-                        + getBinderCallingUid());
-                receiver.send(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION, null);
-                return true;
-            }
             if (!hasTetherChangePermission(callerPkg, callingAttributionTag,
                     false /* onlyAllowPrivileged */)) {
                 receiver.send(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION, null);
@@ -305,9 +290,9 @@
 
             if (mTethering.isTetherProvisioningRequired()) return false;
 
-            int uid = getBinderCallingUid();
+            int uid = Binder.getCallingUid();
 
-            // If callerPkg's uid is not same as getBinderCallingUid(),
+            // If callerPkg's uid is not same as Binder.getCallingUid(),
             // checkAndNoteWriteSettingsOperation will return false and the operation will be
             // denied.
             return mService.checkAndNoteWriteSettingsOperation(mService, uid, callerPkg,
@@ -320,14 +305,6 @@
             return mService.checkCallingOrSelfPermission(
                     ACCESS_NETWORK_STATE) == PERMISSION_GRANTED;
         }
-
-        private int getBinderCallingUid() {
-            return mService.getBinderCallingUid();
-        }
-
-        private boolean checkPackageNameMatchesUid(final int uid, final String callerPkg) {
-            return mService.checkPackageNameMatchesUid(mService, uid, callerPkg);
-        }
     }
 
     /**
@@ -345,28 +322,6 @@
     }
 
     /**
-     * Check if the package name matches the uid.
-     */
-    @VisibleForTesting
-    boolean checkPackageNameMatchesUid(@NonNull Context context, int uid,
-            @NonNull String callingPackage) {
-        try {
-            PermissionUtils.enforcePackageNameMatchesUid(context, uid, callingPackage);
-        } catch (SecurityException e) {
-            return false;
-        }
-        return true;
-    }
-
-    /**
-     * Wrapper for the Binder calling UID, used for mocks.
-     */
-    @VisibleForTesting
-    int getBinderCallingUid() {
-        return Binder.getCallingUid();
-    }
-
-    /**
      * An injection method for testing.
      */
     @VisibleForTesting
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/MockTetheringService.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/MockTetheringService.java
index 7fcc5f1..3c07580 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/MockTetheringService.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/MockTetheringService.java
@@ -32,8 +32,6 @@
 public class MockTetheringService extends TetheringService {
     private final Tethering mTethering = mock(Tethering.class);
     private final ArrayMap<String, Integer> mMockedPermissions = new ArrayMap<>();
-    private final ArrayMap<String, Integer> mMockedPackageUids = new ArrayMap<>();
-    private int mMockCallingUid;
 
     @Override
     public IBinder onBind(Intent intent) {
@@ -63,17 +61,6 @@
         return super.checkCallingOrSelfPermission(permission);
     }
 
-    @Override
-    boolean checkPackageNameMatchesUid(@NonNull Context context, int uid,
-            @NonNull String callingPackage) {
-        return mMockedPackageUids.getOrDefault(callingPackage, 0) == uid;
-    }
-
-    @Override
-    int getBinderCallingUid() {
-        return mMockCallingUid;
-    }
-
     public Tethering getTethering() {
         return mTethering;
     }
@@ -104,19 +91,5 @@
                 mMockedPermissions.put(permission, granted);
             }
         }
-
-        /**
-         * Mock a package name matching a uid.
-         */
-        public void setPackageNameUid(String packageName, int uid) {
-            mMockedPackageUids.put(packageName, uid);
-        }
-
-        /**
-         * Mock a package name matching a uid.
-         */
-        public void setCallingUid(int uid) {
-            mMockCallingUid = uid;
-        }
     }
 }
diff --git a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java
index 1988311..c0d7ad4 100644
--- a/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java
+++ b/Tethering/tests/unit/src/com/android/networkstack/tethering/TetheringServiceTest.java
@@ -34,7 +34,6 @@
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -80,7 +79,6 @@
 public final class TetheringServiceTest {
     private static final String TEST_IFACE_NAME = "test_wlan0";
     private static final String TEST_CALLER_PKG = "com.android.shell";
-    private static final int TEST_CALLER_UID = 1234;
     private static final String TEST_ATTRIBUTION_TAG = null;
     @Mock private ITetheringEventCallback mITetheringEventCallback;
     @Rule public ServiceTestRule mServiceTestRule;
@@ -130,8 +128,6 @@
         mTetheringConnector = ITetheringConnector.Stub.asInterface(mMockConnector.getIBinder());
         final MockTetheringService service = mMockConnector.getService();
         mTethering = service.getTethering();
-        mMockConnector.setCallingUid(TEST_CALLER_UID);
-        mMockConnector.setPackageNameUid(TEST_CALLER_PKG, TEST_CALLER_UID);
     }
 
     @After
@@ -334,16 +330,6 @@
         });
 
         runAsTetherPrivileged((result) -> {
-            String wrongPackage = "wrong.package";
-            mTetheringConnector.startTethering(request, wrongPackage,
-                    TEST_ATTRIBUTION_TAG, result);
-            verify(mTethering, never()).startTethering(
-                    eq(new TetheringRequest(request)), eq(wrongPackage), eq(result));
-            result.assertResult(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION);
-            verifyNoMoreInteractionsForTethering();
-        });
-
-        runAsTetherPrivileged((result) -> {
             runStartTethering(result, request);
             verifyNoMoreInteractionsForTethering();
         });
@@ -459,13 +445,6 @@
             verifyNoMoreInteractionsForTethering();
         });
 
-        runAsTetherPrivileged((none) -> {
-            mTetheringConnector.requestLatestTetheringEntitlementResult(TETHERING_WIFI, result,
-                    true /* showEntitlementUi */, "wrong.package", TEST_ATTRIBUTION_TAG);
-            result.assertResult(TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION);
-            verifyNoMoreInteractions(mTethering);
-        });
-
         runAsWriteSettings((none) -> {
             runRequestLatestTetheringEntitlementResult();
             verify(mTethering).isTetherProvisioningRequired();
diff --git a/bpf/headers/include/bpf_helpers.h b/bpf/headers/include/bpf_helpers.h
index ac5ffda..b994a9f 100644
--- a/bpf/headers/include/bpf_helpers.h
+++ b/bpf/headers/include/bpf_helpers.h
@@ -403,6 +403,8 @@
 static unsigned long long (*bpf_get_smp_processor_id)(void) = (void*) BPF_FUNC_get_smp_processor_id;
 static long (*bpf_get_stackid)(void* ctx, void* map, uint64_t flags) = (void*) BPF_FUNC_get_stackid;
 static long (*bpf_get_current_comm)(void* buf, uint32_t buf_size) = (void*) BPF_FUNC_get_current_comm;
+// bpf_sk_fullsock requires 5.1+ kernel
+static struct bpf_sock* (*bpf_sk_fullsock)(struct bpf_sock* sk) = (void*) BPF_FUNC_sk_fullsock;
 
 // GPL only:
 static int (*bpf_trace_printk)(const char* fmt, int fmt_size, ...) = (void*) BPF_FUNC_trace_printk;
diff --git a/framework-t/src/android/net/INetworkStatsService.aidl b/framework-t/src/android/net/INetworkStatsService.aidl
index 01ac106..b459a13 100644
--- a/framework-t/src/android/net/INetworkStatsService.aidl
+++ b/framework-t/src/android/net/INetworkStatsService.aidl
@@ -21,10 +21,11 @@
 import android.net.Network;
 import android.net.NetworkStateSnapshot;
 import android.net.NetworkStats;
-import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
 import android.net.UnderlyingNetworkInfo;
 import android.net.netstats.IUsageCallback;
+import android.net.netstats.StatsResult;
+import android.net.netstats.TrafficStatsRateLimitCacheConfig;
 import android.net.netstats.provider.INetworkStatsProvider;
 import android.net.netstats.provider.INetworkStatsProviderCallback;
 import android.os.IBinder;
@@ -78,16 +79,13 @@
     void unregisterUsageRequest(in DataUsageRequest request);
 
     /** Get the uid stats information since boot */
-    NetworkStats getTypelessUidStats(int uid);
+    StatsResult getUidStats(int uid);
 
     /** Get the iface stats information since boot */
-    NetworkStats getTypelessIfaceStats(String iface);
+    StatsResult getIfaceStats(String iface);
 
     /** Get the total network stats information since boot */
-    NetworkStats getTypelessTotalStats();
-
-    /** Get the uid stats information (with specified type) since boot */
-    long getUidStats(int uid, int type);
+    StatsResult getTotalStats();
 
     /** Registers a network stats provider */
     INetworkStatsProviderCallback registerNetworkStatsProvider(String tag,
@@ -107,4 +105,7 @@
 
      /** Clear TrafficStats rate-limit caches. */
      void clearTrafficStatsRateLimitCaches();
+
+     /** Get rate-limit cache config. */
+     TrafficStatsRateLimitCacheConfig getRateLimitCacheConfig();
 }
diff --git a/framework-t/src/android/net/TrafficStats.java b/framework-t/src/android/net/TrafficStats.java
index 3b6a69b..1294b3e 100644
--- a/framework-t/src/android/net/TrafficStats.java
+++ b/framework-t/src/android/net/TrafficStats.java
@@ -17,8 +17,13 @@
 package android.net;
 
 import static android.annotation.SystemApi.Client.MODULE_LIBRARIES;
+import static android.net.NetworkStats.UID_ALL;
+import static android.os.Process.SYSTEM_UID;
+
+import static com.android.internal.annotations.VisibleForTesting.Visibility.PRIVATE;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
@@ -29,19 +34,21 @@
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.media.MediaPlayer;
+import android.net.netstats.StatsResult;
 import android.os.Binder;
 import android.os.Build;
 import android.os.RemoteException;
 import android.os.StrictMode;
 import android.util.Log;
 
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.net.DatagramSocket;
 import java.net.Socket;
 import java.net.SocketException;
-import java.util.Iterator;
-import java.util.Objects;
 
 
 /**
@@ -177,10 +184,16 @@
     /** @hide */
     public static final int TAG_SYSTEM_PROBE = 0xFFFFFF42;
 
+    @GuardedBy("TrafficStats.class")
     private static INetworkStatsService sStatsService;
+    @GuardedBy("TrafficStats.class")
+    private static INetworkStatsService sStatsServiceForTest = null;
+    @GuardedBy("TrafficStats.class")
+    private static int sMyUidForTest = UID_ALL;
 
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.P, trackingBug = 130143562)
     private synchronized static INetworkStatsService getStatsService() {
+        if (sStatsServiceForTest != null) return sStatsServiceForTest;
         if (sStatsService == null) {
             throw new IllegalStateException("TrafficStats not initialized, uid="
                     + Binder.getCallingUid());
@@ -188,6 +201,40 @@
         return sStatsService;
     }
 
+    /** @hide */
+    protected static int getMyUid() {
+        synchronized (TrafficStats.class) {
+            if (sMyUidForTest != UID_ALL) {
+                return sMyUidForTest;
+            }
+        }
+        return android.os.Process.myUid();
+    }
+
+    /**
+     * Set the network stats service for testing, or null to reset.
+     *
+     * @hide
+     */
+    @VisibleForTesting(visibility = PRIVATE)
+    public static void setServiceForTest(INetworkStatsService statsService) {
+        synchronized (TrafficStats.class) {
+            sStatsServiceForTest = statsService;
+        }
+    }
+
+    /**
+     * Set myUid for test, or UID_ALL to reset.
+     *
+     * @hide
+     */
+    @VisibleForTesting(visibility = PRIVATE)
+    public static void setMyUidForTest(int myUid) {
+        synchronized (TrafficStats.class) {
+            sMyUidForTest = myUid;
+        }
+    }
+
     /**
      * Snapshot of {@link NetworkStats} when the currently active profiling
      * session started, or {@code null} if no session active.
@@ -450,7 +497,7 @@
      */
     @Deprecated
     public static void setThreadStatsUidSelf() {
-        setThreadStatsUid(android.os.Process.myUid());
+        setThreadStatsUid(getMyUid());
     }
 
     /**
@@ -591,7 +638,7 @@
      * @param operationCount Number of operations to increment count by.
      */
     public static void incrementOperationCount(int tag, int operationCount) {
-        final int uid = android.os.Process.myUid();
+        final int uid = getMyUid();
         try {
             getStatsService().incrementOperationCount(uid, tag, operationCount);
         } catch (RemoteException e) {
@@ -959,17 +1006,20 @@
 
     /** @hide */
     public static long getUidStats(int uid, int type) {
-        if (!isEntryValueTypeValid(type)
-                || android.os.Process.myUid() != uid) {
+        // Perform a quick check on the UID to avoid unnecessary work.
+        // This mirrors a similar check on the service side, but is primarily for
+        // efficiency rather than security, as user-space checks can be bypassed.
+        final int myUid = getMyUid();
+        if (!isEntryValueTypeValid(type) || (myUid != SYSTEM_UID && myUid != uid)) {
             return UNSUPPORTED;
         }
-        final NetworkStats stats;
+        final StatsResult stats;
         try {
-            stats = getStatsService().getTypelessUidStats(uid);
+            stats = getStatsService().getUidStats(uid);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
-        return getValueForTypeFromFirstEntry(stats, type);
+        return getEntryValueForType(stats, type);
     }
 
     /** @hide */
@@ -977,13 +1027,13 @@
         if (!isEntryValueTypeValid(type)) {
             return UNSUPPORTED;
         }
-        final NetworkStats stats;
+        final StatsResult stats;
         try {
-            stats = getStatsService().getTypelessTotalStats();
+            stats = getStatsService().getTotalStats();
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
-        return getValueForTypeFromFirstEntry(stats, type);
+        return getEntryValueForType(stats, type);
     }
 
     /** @hide */
@@ -991,13 +1041,13 @@
         if (!isEntryValueTypeValid(type)) {
             return UNSUPPORTED;
         }
-        final NetworkStats stats;
+        final StatsResult stats;
         try {
-            stats = getStatsService().getTypelessIfaceStats(iface);
+            stats = getStatsService().getIfaceStats(iface);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
-        return getValueForTypeFromFirstEntry(stats, type);
+        return getEntryValueForType(stats, type);
     }
 
     /**
@@ -1094,7 +1144,7 @@
      */
     private static NetworkStats getDataLayerSnapshotForUid(Context context) {
         // TODO: take snapshot locally, since proc file is now visible
-        final int uid = android.os.Process.myUid();
+        final int uid = getMyUid();
         try {
             return getStatsService().getDataLayerSnapshotForUid(uid);
         } catch (RemoteException e) {
@@ -1127,18 +1177,18 @@
     public static final int TYPE_TX_PACKETS = 3;
 
     /** @hide */
-    private static long getEntryValueForType(@NonNull NetworkStats.Entry entry, int type) {
-        Objects.requireNonNull(entry);
+    private static long getEntryValueForType(@Nullable StatsResult stats, int type) {
+        if (stats == null) return UNSUPPORTED;
         if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
         switch (type) {
             case TYPE_RX_BYTES:
-                return entry.getRxBytes();
+                return stats.rxBytes;
             case TYPE_RX_PACKETS:
-                return entry.getRxPackets();
+                return stats.rxPackets;
             case TYPE_TX_BYTES:
-                return entry.getTxBytes();
+                return stats.txBytes;
             case TYPE_TX_PACKETS:
-                return entry.getTxPackets();
+                return stats.txPackets;
             default:
                 throw new IllegalStateException("Bug: Invalid type: "
                         + type + " should not reach here.");
@@ -1157,13 +1207,5 @@
                 return false;
         }
     }
-
-    /** @hide */
-    public static long getValueForTypeFromFirstEntry(@NonNull NetworkStats stats, int type) {
-        Objects.requireNonNull(stats);
-        Iterator<NetworkStats.Entry> iter = stats.iterator();
-        if (!iter.hasNext()) return UNSUPPORTED;
-        return getEntryValueForType(iter.next(), type);
-    }
 }
 
diff --git a/framework-t/src/android/net/netstats/StatsResult.aidl b/framework-t/src/android/net/netstats/StatsResult.aidl
new file mode 100644
index 0000000..3f09566
--- /dev/null
+++ b/framework-t/src/android/net/netstats/StatsResult.aidl
@@ -0,0 +1,31 @@
+/**
+ * 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.netstats;
+
+/**
+ * A lightweight class to pass result of TrafficStats#get{Total|Iface|Uid}Stats.
+ *
+ * @hide
+ */
+@JavaDerive(equals=true, toString=true)
+@JavaOnlyImmutable
+parcelable StatsResult {
+    long rxBytes;
+    long rxPackets;
+    long txBytes;
+    long txPackets;
+}
\ No newline at end of file
diff --git a/framework-t/src/android/net/netstats/TrafficStatsRateLimitCacheConfig.aidl b/framework-t/src/android/net/netstats/TrafficStatsRateLimitCacheConfig.aidl
new file mode 100644
index 0000000..cdf0b7c
--- /dev/null
+++ b/framework-t/src/android/net/netstats/TrafficStatsRateLimitCacheConfig.aidl
@@ -0,0 +1,42 @@
+/**
+ * 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.netstats;
+
+/**
+ * Configuration for the TrafficStats rate limit cache.
+ *
+ * @hide
+ */
+@JavaDerive(equals=true, toString=true)
+@JavaOnlyImmutable
+parcelable TrafficStatsRateLimitCacheConfig {
+
+    /**
+     * Whether the cache is enabled for V+ device or target Sdk V+ apps.
+     */
+    boolean isCacheEnabled;
+
+    /**
+     * The duration for which cache entries are valid, in milliseconds.
+     */
+    int expiryDurationMs;
+
+    /**
+     * The maximum number of entries to store in the cache.
+     */
+    int maxEntries;
+}
diff --git a/nearby/framework/java/android/nearby/NearbyManager.java b/nearby/framework/java/android/nearby/NearbyManager.java
index 150394b..e78f999 100644
--- a/nearby/framework/java/android/nearby/NearbyManager.java
+++ b/nearby/framework/java/android/nearby/NearbyManager.java
@@ -32,7 +32,6 @@
 import android.nearby.aidl.IOffloadCallback;
 import android.os.RemoteException;
 import android.os.SystemProperties;
-import android.provider.Settings;
 import android.util.Log;
 
 import com.android.internal.annotations.GuardedBy;
@@ -129,16 +128,6 @@
     private static final String POWER_OFF_FINDING_SUPPORTED_PROPERTY_PERSIST =
             "persist.bluetooth.finder.supported";
 
-    /**
-     * TODO(b/286137024): Remove this when CTS R5 is rolled out.
-     * Whether allows Fast Pair to scan.
-     *
-     * (0 = disabled, 1 = enabled)
-     *
-     * @hide
-     */
-    public static final String FAST_PAIR_SCAN_ENABLED = "fast_pair_scan_enabled";
-
     @GuardedBy("sScanListeners")
     private static final WeakHashMap<ScanCallback, WeakReference<ScanListenerTransport>>
             sScanListeners = new WeakHashMap<>();
@@ -479,36 +468,6 @@
     }
 
     /**
-     * TODO(b/286137024): Remove this when CTS R5 is rolled out.
-     * Read from {@link Settings} whether Fast Pair scan is enabled.
-     *
-     * @param context the {@link Context} to query the setting
-     * @return whether the Fast Pair is enabled
-     * @hide
-     */
-    public static boolean getFastPairScanEnabled(@NonNull Context context) {
-        final int enabled = Settings.Secure.getInt(
-                context.getContentResolver(), FAST_PAIR_SCAN_ENABLED, 0);
-        return enabled != 0;
-    }
-
-    /**
-     * TODO(b/286137024): Remove this when CTS R5 is rolled out.
-     * Write into {@link Settings} whether Fast Pair scan is enabled
-     *
-     * @param context the {@link Context} to set the setting
-     * @param enable whether the Fast Pair scan should be enabled
-     * @hide
-     */
-    @RequiresPermission(Manifest.permission.WRITE_SECURE_SETTINGS)
-    public static void setFastPairScanEnabled(@NonNull Context context, boolean enable) {
-        Settings.Secure.putInt(
-                context.getContentResolver(), FAST_PAIR_SCAN_ENABLED, enable ? 1 : 0);
-        Log.v(TAG, String.format(
-                "successfully %s Fast Pair scan", enable ? "enables" : "disables"));
-    }
-
-    /**
      * Sets the precomputed EIDs for advertising when the phone is powered off. The Bluetooth
      * controller will store these EIDs in its memory, and will start advertising them in Find My
      * Device network EID frames when powered off, only if the powered off finding mode was
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
index 0b2003f..58defa9 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
@@ -416,13 +416,6 @@
         // recvbuf and src are reused after this returns; ensure references to src are not kept.
         final InetSocketAddress srcCopy = new InetSocketAddress(src.getAddress(), src.getPort());
 
-        if (DBG) {
-            mSharedLog.v("Parsed packet with " + packet.questions.size() + " questions, "
-                    + packet.answers.size() + " answers, "
-                    + packet.authorityRecords.size() + " authority, "
-                    + packet.additionalRecords.size() + " additional from " + srcCopy);
-        }
-
         Map<Integer, Integer> conflictingServices =
                 mRecordRepository.getConflictingServices(packet);
 
@@ -440,7 +433,14 @@
         // answer. One exception is simultaneous probe tiebreaking (rfc6762 8.2), in which case the
         // conflicting service is still probing and won't reply either.
         final MdnsReplyInfo answers = mRecordRepository.getReply(packet, srcCopy);
-
+        // Dump the query packet.
+        if (DBG || answers != null) {
+            mSharedLog.v("Parsed packet with transactionId(" + packet.transactionId + "): "
+                    + packet.questions.size() + " questions, "
+                    + packet.answers.size() + " answers, "
+                    + packet.authorityRecords.size() + " authority, "
+                    + packet.additionalRecords.size() + " additional from " + srcCopy);
+        }
         if (answers == null) return;
         mReplySender.queueReply(answers);
     }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java b/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java
index a89b004..4708cb6 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsReplySender.java
@@ -245,7 +245,7 @@
                 return;
             }
 
-            if (mEnableDebugLog) mSharedLog.v("Sending " + replyInfo);
+            mSharedLog.log("Sending " + replyInfo);
 
             final int flags = 0x8400; // Response, authoritative (rfc6762 18.4)
             final MdnsPacket packet = new MdnsPacket(flags,
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 294a85a..fb712a1 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -52,7 +52,6 @@
 import static android.net.TrafficStats.KB_IN_BYTES;
 import static android.net.TrafficStats.MB_IN_BYTES;
 import static android.net.TrafficStats.UID_TETHERING;
-import static android.net.TrafficStats.getValueForTypeFromFirstEntry;
 import static android.net.connectivity.ConnectivityCompatChanges.ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID_TAG;
@@ -127,6 +126,8 @@
 import android.net.Uri;
 import android.net.netstats.IUsageCallback;
 import android.net.netstats.NetworkStatsDataMigrationUtils;
+import android.net.netstats.StatsResult;
+import android.net.netstats.TrafficStatsRateLimitCacheConfig;
 import android.net.netstats.provider.INetworkStatsProvider;
 import android.net.netstats.provider.INetworkStatsProviderCallback;
 import android.net.netstats.provider.NetworkStatsProvider;
@@ -485,16 +486,23 @@
     @GuardedBy("mStatsLock")
     private long mLatestNetworkStatsUpdatedBroadcastScheduledTime = Long.MIN_VALUE;
 
+    @Nullable
     private final TrafficStatsRateLimitCache mTrafficStatsTotalCache;
+    @Nullable
     private final TrafficStatsRateLimitCache mTrafficStatsIfaceCache;
+    @Nullable
     private final TrafficStatsRateLimitCache mTrafficStatsUidCache;
+    // A feature flag to control whether the client-side rate limit cache should be enabled.
+    static final String TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG =
+            "trafficstats_client_rate_limit_cache_enabled_flag";
     static final String TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG =
             "trafficstats_rate_limit_cache_enabled_flag";
     static final String BROADCAST_NETWORK_STATS_UPDATED_RATE_LIMIT_ENABLED_FLAG =
             "broadcast_network_stats_updated_rate_limit_enabled_flag";
-    private final boolean mAlwaysUseTrafficStatsServiceRateLimitCache;
+    private final boolean mIsTrafficStatsServiceRateLimitCacheEnabled;
     private final int mTrafficStatsRateLimitCacheExpiryDuration;
     private final int mTrafficStatsServiceRateLimitCacheMaxEntries;
+    private final TrafficStatsRateLimitCacheConfig mTrafficStatsRateLimitCacheClientSideConfig;
     private final boolean mBroadcastNetworkStatsUpdatedRateLimitEnabled;
 
 
@@ -688,23 +696,34 @@
             mEventLogger = null;
         }
 
-        mAlwaysUseTrafficStatsServiceRateLimitCache =
-                mDeps.alwaysUseTrafficStatsServiceRateLimitCache(mContext);
+        mTrafficStatsRateLimitCacheClientSideConfig =
+                mDeps.getTrafficStatsRateLimitCacheClientSideConfig(mContext);
+        // If the client side cache feature is enabled, disable the service side
+        // cache unconditionally.
+        mIsTrafficStatsServiceRateLimitCacheEnabled =
+                mDeps.isTrafficStatsServiceRateLimitCacheEnabled(mContext,
+                        mTrafficStatsRateLimitCacheClientSideConfig.isCacheEnabled);
         mBroadcastNetworkStatsUpdatedRateLimitEnabled =
                 mDeps.enabledBroadcastNetworkStatsUpdatedRateLimiting(mContext);
         mTrafficStatsRateLimitCacheExpiryDuration =
                 mDeps.getTrafficStatsRateLimitCacheExpiryDuration();
         mTrafficStatsServiceRateLimitCacheMaxEntries =
                 mDeps.getTrafficStatsServiceRateLimitCacheMaxEntries();
-        mTrafficStatsTotalCache = new TrafficStatsRateLimitCache(mClock,
-                mTrafficStatsRateLimitCacheExpiryDuration,
-                mTrafficStatsServiceRateLimitCacheMaxEntries);
-        mTrafficStatsIfaceCache = new TrafficStatsRateLimitCache(mClock,
-                mTrafficStatsRateLimitCacheExpiryDuration,
-                mTrafficStatsServiceRateLimitCacheMaxEntries);
-        mTrafficStatsUidCache = new TrafficStatsRateLimitCache(mClock,
-                mTrafficStatsRateLimitCacheExpiryDuration,
-                mTrafficStatsServiceRateLimitCacheMaxEntries);
+        if (mIsTrafficStatsServiceRateLimitCacheEnabled) {
+            mTrafficStatsTotalCache = new TrafficStatsRateLimitCache(mClock,
+                    mTrafficStatsRateLimitCacheExpiryDuration,
+                    mTrafficStatsServiceRateLimitCacheMaxEntries);
+            mTrafficStatsIfaceCache = new TrafficStatsRateLimitCache(mClock,
+                    mTrafficStatsRateLimitCacheExpiryDuration,
+                    mTrafficStatsServiceRateLimitCacheMaxEntries);
+            mTrafficStatsUidCache = new TrafficStatsRateLimitCache(mClock,
+                    mTrafficStatsRateLimitCacheExpiryDuration,
+                    mTrafficStatsServiceRateLimitCacheMaxEntries);
+        } else {
+            mTrafficStatsTotalCache = null;
+            mTrafficStatsIfaceCache = null;
+            mTrafficStatsUidCache = null;
+        }
 
         // TODO: Remove bpfNetMaps creation and always start SkDestroyListener
         // Following code is for the experiment to verify the SkDestroyListener refactoring. Based
@@ -964,13 +983,42 @@
         }
 
         /**
-         * Get whether TrafficStats service side rate-limit cache is always applied.
+         * Get client side traffic stats rate-limit cache config.
          *
          * This method should only be called once in the constructor,
          * to ensure that the code does not need to deal with flag values changing at runtime.
          */
-        public boolean alwaysUseTrafficStatsServiceRateLimitCache(@NonNull Context ctx) {
-            return SdkLevel.isAtLeastV() && DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
+        @NonNull
+        public TrafficStatsRateLimitCacheConfig getTrafficStatsRateLimitCacheClientSideConfig(
+                @NonNull Context ctx) {
+            final TrafficStatsRateLimitCacheConfig config =
+                    new TrafficStatsRateLimitCacheConfig.Builder()
+                            .setIsCacheEnabled(DeviceConfigUtils.isTetheringFeatureEnabled(
+                                    ctx, TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG))
+                            .setExpiryDurationMs(getDeviceConfigPropertyInt(
+                                    NAMESPACE_TETHERING, TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME,
+                                    DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS))
+                            .setMaxEntries(getDeviceConfigPropertyInt(
+                                    NAMESPACE_TETHERING,
+                                    TRAFFIC_STATS_SERVICE_CACHE_MAX_ENTRIES_NAME,
+                                    DEFAULT_TRAFFIC_STATS_SERVICE_CACHE_MAX_ENTRIES))
+                            .build();
+            return config;
+        }
+
+        /**
+         * Determines whether the service-side rate-limiting cache is enabled.
+         *
+         * The cache is enabled for devices running Android V+ or apps targeting SDK V+
+         * if the `TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG` feature flag
+         * is enabled and client-side caching is disabled.
+         *
+         * This method should only be called once in the constructor,
+         * to ensure that the code does not need to deal with flag values changing at runtime.
+         */
+        public boolean isTrafficStatsServiceRateLimitCacheEnabled(@NonNull Context ctx,
+                boolean clientCacheEnabled) {
+            return !clientCacheEnabled && DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
                     ctx, TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG);
         }
 
@@ -2133,30 +2181,48 @@
         }
     }
 
-    @Override
-    public long getUidStats(int uid, int type) {
-        return getValueForTypeFromFirstEntry(getTypelessUidStats(uid), type);
+    /**
+     * Determines whether to use the client-side cache for traffic stats rate limiting.
+     *
+     * This is based on the cache enabled feature flag. If enabled, the client-side cache
+     * is used for V+ devices or callers with V+ target sdk.
+     *
+     * @param callingUid The UID of the app making the request.
+     * @return True if the client-side cache should be used, false otherwise.
+     */
+    private boolean useClientSideCache(int callingUid) {
+        return mTrafficStatsRateLimitCacheClientSideConfig.isCacheEnabled && (SdkLevel.isAtLeastV()
+                || mDeps.isChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, callingUid));
     }
 
-    @NonNull
+    /**
+     * Determines whether to use the service-side cache for traffic stats rate limiting.
+     *
+     * This is based on the cache enabled feature flag. If enabled, the service-side cache
+     * is used for V+ devices or callers with V+ target sdk.
+     *
+     * @param callingUid The UID of the app making the request.
+     * @return True if the service-side cache should be used, false otherwise.
+     */
+    private boolean useServiceSideCache(int callingUid) {
+        return mIsTrafficStatsServiceRateLimitCacheEnabled && (SdkLevel.isAtLeastV()
+                || mDeps.isChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, callingUid));
+    }
+
+    @Nullable
     @Override
-    public NetworkStats getTypelessUidStats(int uid) {
-        final NetworkStats stats = new NetworkStats(0, 0);
+    public StatsResult getUidStats(int uid) {
         final int callingUid = Binder.getCallingUid();
         if (callingUid != android.os.Process.SYSTEM_UID && callingUid != uid) {
-            return stats;
+            return null;
         }
         final NetworkStats.Entry entry;
-        if (mAlwaysUseTrafficStatsServiceRateLimitCache
-                || mDeps.isChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, callingUid)) {
+        if (useServiceSideCache(callingUid)) {
             entry = mTrafficStatsUidCache.getOrCompute(IFACE_ALL, uid,
                     () -> mDeps.nativeGetUidStat(uid));
         } else entry = mDeps.nativeGetUidStat(uid);
 
-        if (entry != null) {
-            stats.insertEntry(entry);
-        }
-        return stats;
+        return getStatsResultFromEntryOrNull(entry);
     }
 
     @Nullable
@@ -2173,24 +2239,24 @@
         return entry;
     }
 
-    @NonNull
+    @Nullable
     @Override
-    public NetworkStats getTypelessIfaceStats(@NonNull String iface) {
+    public StatsResult getIfaceStats(@NonNull String iface) {
         Objects.requireNonNull(iface);
 
         final NetworkStats.Entry entry;
-        if (mAlwaysUseTrafficStatsServiceRateLimitCache
-                || mDeps.isChangeEnabled(
-                        ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, Binder.getCallingUid())) {
+        if (useServiceSideCache(Binder.getCallingUid())) {
             entry = mTrafficStatsIfaceCache.getOrCompute(iface, UID_ALL,
                     () -> getIfaceStatsInternal(iface));
         } else entry = getIfaceStatsInternal(iface);
 
-        NetworkStats stats = new NetworkStats(0, 0);
-        if (entry != null) {
-            stats.insertEntry(entry);
-        }
-        return stats;
+        return getStatsResultFromEntryOrNull(entry);
+    }
+
+    @Nullable
+    private StatsResult getStatsResultFromEntryOrNull(@Nullable NetworkStats.Entry entry) {
+        if (entry == null) return null;
+        return new StatsResult(entry.rxBytes, entry.rxPackets, entry.txBytes, entry.txPackets);
     }
 
     @Nullable
@@ -2203,30 +2269,39 @@
         return entry;
     }
 
-    @NonNull
+    @Nullable
     @Override
-    public NetworkStats getTypelessTotalStats() {
+    public StatsResult getTotalStats() {
         final NetworkStats.Entry entry;
-        if (mAlwaysUseTrafficStatsServiceRateLimitCache
-                || mDeps.isChangeEnabled(
-                        ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, Binder.getCallingUid())) {
+        if (useServiceSideCache(Binder.getCallingUid())) {
             entry = mTrafficStatsTotalCache.getOrCompute(
                     IFACE_ALL, UID_ALL, () -> getTotalStatsInternal());
         } else entry = getTotalStatsInternal();
 
-        final NetworkStats stats = new NetworkStats(0, 0);
-        if (entry != null) {
-            stats.insertEntry(entry);
-        }
-        return stats;
+        return getStatsResultFromEntryOrNull(entry);
     }
 
     @Override
     public void clearTrafficStatsRateLimitCaches() {
         PermissionUtils.enforceNetworkStackPermissionOr(mContext, NETWORK_SETTINGS);
-        mTrafficStatsUidCache.clear();
-        mTrafficStatsIfaceCache.clear();
-        mTrafficStatsTotalCache.clear();
+        if (mIsTrafficStatsServiceRateLimitCacheEnabled) {
+            mTrafficStatsUidCache.clear();
+            mTrafficStatsIfaceCache.clear();
+            mTrafficStatsTotalCache.clear();
+        }
+    }
+
+    @Override
+    public TrafficStatsRateLimitCacheConfig getRateLimitCacheConfig() {
+        // Build a per uid config for the client based on the checking result.
+        final TrafficStatsRateLimitCacheConfig config =
+                new TrafficStatsRateLimitCacheConfig.Builder()
+                        .setIsCacheEnabled(useClientSideCache(Binder.getCallingUid()))
+                        .setExpiryDurationMs(
+                                mTrafficStatsRateLimitCacheClientSideConfig.expiryDurationMs)
+                        .setMaxEntries(mTrafficStatsRateLimitCacheClientSideConfig.maxEntries)
+                        .build();
+        return config;
     }
 
     private NetworkStats.Entry getProviderIfaceStats(@Nullable String iface) {
@@ -2996,8 +3071,8 @@
             } catch (IOException e) {
                 pw.println("(failed to dump FastDataInput counters)");
             }
-            pw.print("trafficstats.service.cache.alwaysuse",
-                    mAlwaysUseTrafficStatsServiceRateLimitCache);
+            pw.print("trafficstats.service.cache.isenabled",
+                    mIsTrafficStatsServiceRateLimitCacheEnabled);
             pw.println();
             pw.print(TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME,
                     mTrafficStatsRateLimitCacheExpiryDuration);
@@ -3005,6 +3080,9 @@
             pw.print(TRAFFIC_STATS_SERVICE_CACHE_MAX_ENTRIES_NAME,
                     mTrafficStatsServiceRateLimitCacheMaxEntries);
             pw.println();
+            pw.print("trafficstats.client.cache.config",
+                    mTrafficStatsRateLimitCacheClientSideConfig);
+            pw.println();
 
             pw.decreaseIndent();
 
diff --git a/staticlibs/testutils/app/connectivitychecker/src/com/android/testutils/connectivitypreparer/CarrierConfigSetupTest.kt b/staticlibs/testutils/app/connectivitychecker/src/com/android/testutils/connectivitypreparer/CarrierConfigSetupTest.kt
new file mode 100644
index 0000000..46e511e
--- /dev/null
+++ b/staticlibs/testutils/app/connectivitychecker/src/com/android/testutils/connectivitypreparer/CarrierConfigSetupTest.kt
@@ -0,0 +1,132 @@
+/*
+ * 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.testutils.connectivitypreparer
+
+import android.Manifest.permission.MODIFY_PHONE_STATE
+import android.Manifest.permission.READ_PHONE_STATE
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.content.pm.PackageManager.FEATURE_TELEPHONY_IMS
+import android.content.pm.PackageManager.FEATURE_WIFI
+import android.os.Build
+import android.os.Build.VERSION_CODES.UPSIDE_DOWN_CAKE
+import android.os.ParcelFileDescriptor
+import android.os.PersistableBundle
+import android.telephony.CarrierConfigManager
+import android.telephony.CarrierConfigManager.ACTION_CARRIER_CONFIG_CHANGED
+import android.telephony.SubscriptionManager
+import android.util.Log
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.net.module.util.ArrayTrackRecord
+import com.android.testutils.runAsShell
+import kotlin.test.assertNotNull
+import org.junit.Test
+import org.junit.runner.RunWith
+
+private const val CONFIG_CHANGE_TIMEOUT_MS = 10_000L
+private val TAG = CarrierConfigSetupTest::class.simpleName
+
+@RunWith(AndroidJUnit4::class)
+class CarrierConfigSetupTest {
+    private val context by lazy { InstrumentationRegistry.getInstrumentation().context }
+    private val pm by lazy { context.packageManager }
+    private val carrierConfigManager by lazy {
+        context.getSystemService(CarrierConfigManager::class.java)
+    }
+
+    @Test
+    fun testSetCarrierConfig() {
+        if (!shouldDisableIwlan()) return
+        overrideAllSubscriptions(PersistableBundle().apply {
+            putBoolean(CarrierConfigManager.KEY_CARRIER_WFC_IMS_AVAILABLE_BOOL, false)
+        })
+    }
+
+    @Test
+    fun testClearCarrierConfig() {
+        // set/clear are in different test runs so it is difficult to share state between them.
+        // The conditions to disable IWLAN should not change over time (in particular
+        // force_iwlan_mms is a readonly flag), so just perform the same check again on teardown.
+        // CarrierConfigManager overrides are cleared on reboot by default anyway, so any missed
+        // cleanup should not be too damaging.
+        if (!shouldDisableIwlan()) return
+        overrideAllSubscriptions(null)
+    }
+
+    private class ConfigChangedReceiver : BroadcastReceiver() {
+        val receivedSubIds = ArrayTrackRecord<Int>()
+        override fun onReceive(context: Context, intent: Intent) {
+            if (intent.action != ACTION_CARRIER_CONFIG_CHANGED) return
+            val subIdx = intent.getIntExtra(SubscriptionManager.EXTRA_SUBSCRIPTION_INDEX, -1)
+            // It is possible this is a configuration change for a different setting, so the test
+            // may not wait for long enough. Unfortunately calling CarrierConfigManager to check
+            // if the config was applied does not help because it will always return the override,
+            // even if it was not applied to the subscription yet.
+            // In practice, it is very unlikely that a different broadcast arrives, and then a test
+            // flakes because of the iwlan behavior in the time it takes for the config to be
+            // applied.
+            Log.d(TAG, "Received config change for sub $subIdx")
+            receivedSubIds.add(subIdx)
+        }
+    }
+
+    private fun overrideAllSubscriptions(bundle: PersistableBundle?) {
+        runAsShell(READ_PHONE_STATE, MODIFY_PHONE_STATE) {
+            val receiver = ConfigChangedReceiver()
+            context.registerReceiver(receiver, IntentFilter(ACTION_CARRIER_CONFIG_CHANGED))
+            val subscriptions = context.getSystemService(SubscriptionManager::class.java)
+                .activeSubscriptionInfoList
+            subscriptions?.forEach { subInfo ->
+                Log.d(TAG, "Overriding config for subscription $subInfo")
+                carrierConfigManager.overrideConfig(subInfo.subscriptionId, bundle)
+            }
+            // Don't wait after each update before the next one, but expect all updates to be done
+            // eventually
+            subscriptions?.forEach { subInfo ->
+                assertNotNull(receiver.receivedSubIds.poll(CONFIG_CHANGE_TIMEOUT_MS, pos = 0) {
+                    it == subInfo.subscriptionId
+                }, "Config override broadcast not received for subscription $subInfo")
+            }
+        }
+    }
+
+    private fun shouldDisableIwlan(): Boolean {
+        // IWLAN on U 24Q2 release (U QPR3) causes cell data to reconnect when Wi-Fi is toggled due
+        // to the implementation of the force_iwlan_mms feature, which does not work well with
+        // multinetworking tests. Disable the feature on such builds (b/368477391).
+        // The behavior changed in more recent releases (V) so only U 24Q2 is affected.
+        return pm.hasSystemFeature(FEATURE_TELEPHONY_IMS) && pm.hasSystemFeature(FEATURE_WIFI) &&
+                Build.VERSION.SDK_INT == UPSIDE_DOWN_CAKE &&
+                isForceIwlanMmsEnabled()
+    }
+
+    private fun isForceIwlanMmsEnabled(): Boolean {
+        val uiAutomation = InstrumentationRegistry.getInstrumentation().uiAutomation
+        val flagEnabledRegex = Regex(
+            """telephony/com\.android\.internal\.telephony\.flags\.force_iwlan_mms:""" +
+                    """.*ENABLED \(system\)""")
+        ParcelFileDescriptor.AutoCloseInputStream(
+            uiAutomation.executeShellCommand("printflags")).bufferedReader().use { reader ->
+                return reader.lines().anyMatch {
+                    it.contains(flagEnabledRegex)
+                }
+        }
+    }
+}
\ No newline at end of file
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
index d5e91c2..7b970d3 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
@@ -57,6 +57,7 @@
      * @param enabled The desired state (true for enabled, false for disabled) of the feature flag.
      */
     @Target(AnnotationTarget.FUNCTION)
+    @Repeatable
     @Retention(AnnotationRetention.RUNTIME)
     annotation class FeatureFlag(val name: String, val enabled: Boolean = true)
 
diff --git a/staticlibs/testutils/host/java/com/android/testutils/ConnectivityTestTargetPreparer.kt b/staticlibs/testutils/host/java/com/android/testutils/ConnectivityTestTargetPreparer.kt
index 435fdd8..f6168af 100644
--- a/staticlibs/testutils/host/java/com/android/testutils/ConnectivityTestTargetPreparer.kt
+++ b/staticlibs/testutils/host/java/com/android/testutils/ConnectivityTestTargetPreparer.kt
@@ -28,6 +28,7 @@
 private const val CONNECTIVITY_CHECKER_APK = "ConnectivityTestPreparer.apk"
 private const val CONNECTIVITY_PKG_NAME = "com.android.testutils.connectivitypreparer"
 private const val CONNECTIVITY_CHECK_CLASS = "$CONNECTIVITY_PKG_NAME.ConnectivityCheckTest"
+private const val CARRIER_CONFIG_SETUP_CLASS = "$CONNECTIVITY_PKG_NAME.CarrierConfigSetupTest"
 
 // As per the <instrumentation> defined in the checker manifest
 private const val CONNECTIVITY_CHECK_RUNNER_NAME = "androidx.test.runner.AndroidJUnitRunner"
@@ -84,27 +85,28 @@
         installer.setShouldGrantPermission(true)
         installer.setUp(testInfo)
 
-        val testMethods = mutableListOf<String>()
+        val testMethods = mutableListOf<Pair<String, String>>()
         if (!ignoreWifiCheck) {
-            testMethods.add("testCheckWifiSetup")
+            testMethods.add(CONNECTIVITY_CHECK_CLASS to "testCheckWifiSetup")
         }
         if (!ignoreMobileDataCheck) {
-            testMethods.add("testCheckTelephonySetup")
+            testMethods.add(CARRIER_CONFIG_SETUP_CLASS to "testSetCarrierConfig")
+            testMethods.add(CONNECTIVITY_CHECK_CLASS to "testCheckTelephonySetup")
         }
 
         testMethods.forEach {
-            runTestMethod(testInfo, it)
+            runTestMethod(testInfo, it.first, it.second)
         }
     }
 
-    private fun runTestMethod(testInfo: TestInformation, method: String) {
+    private fun runTestMethod(testInfo: TestInformation, clazz: String, method: String) {
         val runner = DefaultRemoteAndroidTestRunner(
             CONNECTIVITY_PKG_NAME,
             CONNECTIVITY_CHECK_RUNNER_NAME,
             testInfo.device.iDevice
         )
         runner.runOptions = "--no-hidden-api-checks"
-        runner.setMethodName(CONNECTIVITY_CHECK_CLASS, method)
+        runner.setMethodName(clazz, method)
 
         val receiver = CollectingTestListener()
         if (!testInfo.device.runInstrumentationTests(runner, receiver)) {
@@ -187,6 +189,9 @@
 
     override fun tearDown(testInfo: TestInformation, e: Throwable?) {
         if (isTearDownDisabled) return
+        if (!ignoreMobileDataCheck) {
+            runTestMethod(testInfo, CARRIER_CONFIG_SETUP_CLASS, "testClearCarrierConfig")
+        }
         installer.tearDown(testInfo, e)
         setUpdaterNetworkingEnabled(
             testInfo,
diff --git a/tests/cts/net/src/android/net/cts/NetworkStatsBinderTest.java b/tests/cts/net/src/android/net/cts/NetworkStatsBinderTest.java
index 1a48983..10adee0 100644
--- a/tests/cts/net/src/android/net/cts/NetworkStatsBinderTest.java
+++ b/tests/cts/net/src/android/net/cts/NetworkStatsBinderTest.java
@@ -19,28 +19,28 @@
 import static android.os.Process.INVALID_UID;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.net.INetworkStatsService;
 import android.net.TrafficStats;
+import android.net.connectivity.android.net.netstats.StatsResult;
 import android.os.Build;
 import android.os.IBinder;
 import android.os.Process;
 import android.os.RemoteException;
-import android.test.AndroidTestCase;
-import android.util.SparseArray;
 
 import androidx.test.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.util.CollectionUtils;
+import com.android.testutils.ConnectivityModuleTest;
 import com.android.testutils.DevSdkIgnoreRule;
+import com.android.testutils.DevSdkIgnoreRunner;
 
-import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -48,37 +48,20 @@
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
-@RunWith(AndroidJUnit4.class)
+@ConnectivityModuleTest
+@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2) // Mainline NetworkStats starts from T.
+@RunWith(DevSdkIgnoreRunner.class)
 public class NetworkStatsBinderTest {
-    // NOTE: These are shamelessly copied from TrafficStats.
-    private static final int TYPE_RX_BYTES = 0;
-    private static final int TYPE_RX_PACKETS = 1;
-    private static final int TYPE_TX_BYTES = 2;
-    private static final int TYPE_TX_PACKETS = 3;
-
-    @Rule
-    public DevSdkIgnoreRule mIgnoreRule = new DevSdkIgnoreRule(
-            Build.VERSION_CODES.Q /* ignoreClassUpTo */);
-
-    private final SparseArray<Function<Integer, Long>> mUidStatsQueryOpArray = new SparseArray<>();
-
-    @Before
-    public void setUp() throws Exception {
-        mUidStatsQueryOpArray.put(TYPE_RX_BYTES, uid -> TrafficStats.getUidRxBytes(uid));
-        mUidStatsQueryOpArray.put(TYPE_RX_PACKETS, uid -> TrafficStats.getUidRxPackets(uid));
-        mUidStatsQueryOpArray.put(TYPE_TX_BYTES, uid -> TrafficStats.getUidTxBytes(uid));
-        mUidStatsQueryOpArray.put(TYPE_TX_PACKETS, uid -> TrafficStats.getUidTxPackets(uid));
-    }
-
-    private long getUidStatsFromBinder(int uid, int type) throws Exception {
-        Method getServiceMethod = Class.forName("android.os.ServiceManager")
+    @Nullable
+    private StatsResult getUidStatsFromBinder(int uid) throws Exception {
+        final Method getServiceMethod = Class.forName("android.os.ServiceManager")
                 .getDeclaredMethod("getService", new Class[]{String.class});
-        IBinder binder = (IBinder) getServiceMethod.invoke(null, Context.NETWORK_STATS_SERVICE);
-        INetworkStatsService nss = INetworkStatsService.Stub.asInterface(binder);
-        return nss.getUidStats(uid, type);
+        final IBinder binder = (IBinder) getServiceMethod.invoke(
+                null, Context.NETWORK_STATS_SERVICE);
+        final INetworkStatsService nss = INetworkStatsService.Stub.asInterface(binder);
+        return nss.getUidStats(uid);
     }
 
     private int getFirstAppUidThat(@NonNull Predicate<Integer> predicate) {
@@ -108,38 +91,34 @@
         if (notMyUid != INVALID_UID) testUidList.add(notMyUid);
 
         for (final int uid : testUidList) {
-            for (int i = 0; i < mUidStatsQueryOpArray.size(); i++) {
-                final int type = mUidStatsQueryOpArray.keyAt(i);
-                try {
-                    final long uidStatsFromBinder = getUidStatsFromBinder(uid, type);
-                    final long uidTrafficStats = mUidStatsQueryOpArray.get(type).apply(uid);
+            try {
+                final StatsResult uidStatsFromBinder = getUidStatsFromBinder(uid);
 
-                    // Verify that UNSUPPORTED is returned if the uid is not current app uid.
-                    if (uid != myUid) {
-                        assertEquals(uidStatsFromBinder, TrafficStats.UNSUPPORTED);
-                    }
+                if (uid != myUid) {
+                    // Verify that null is returned if the uid is not current app uid.
+                    assertNull(uidStatsFromBinder);
+                } else {
                     // Verify that returned result is the same with the result get from
                     // TrafficStats.
-                    // TODO: If the test is flaky then it should instead assert that the values
-                    //  are approximately similar.
-                    assertEquals("uidStats is not matched for query type " + type
-                                    + ", uid=" + uid + ", myUid=" + myUid, uidTrafficStats,
-                            uidStatsFromBinder);
-                } catch (IllegalAccessException e) {
-                    /* Java language access prevents exploitation. */
-                    return;
-                } catch (InvocationTargetException e) {
-                    /* Underlying method has been changed. */
-                    return;
-                } catch (ClassNotFoundException e) {
-                    /* not vulnerable if hidden API no longer available */
-                    return;
-                } catch (NoSuchMethodException e) {
-                    /* not vulnerable if hidden API no longer available */
-                    return;
-                } catch (RemoteException e) {
-                    return;
+                    assertEquals(uidStatsFromBinder.rxBytes, TrafficStats.getUidRxBytes(uid));
+                    assertEquals(uidStatsFromBinder.rxPackets, TrafficStats.getUidRxPackets(uid));
+                    assertEquals(uidStatsFromBinder.txBytes, TrafficStats.getUidTxBytes(uid));
+                    assertEquals(uidStatsFromBinder.txPackets, TrafficStats.getUidTxPackets(uid));
                 }
+            } catch (IllegalAccessException e) {
+                /* Java language access prevents exploitation. */
+                return;
+            } catch (InvocationTargetException e) {
+                /* Underlying method has been changed. */
+                return;
+            } catch (ClassNotFoundException e) {
+                /* not vulnerable if hidden API no longer available */
+                return;
+            } catch (NoSuchMethodException e) {
+                /* not vulnerable if hidden API no longer available */
+                return;
+            } catch (RemoteException e) {
+                return;
             }
         }
     }
diff --git a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
index ad6fe63..7fc8863 100644
--- a/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
+++ b/tests/cts/net/src/android/net/cts/NsdManagerTest.kt
@@ -326,6 +326,15 @@
         it.port = TEST_PORT
     }
 
+    private fun makePacketReader(network: TestTapNetwork = testNetwork1) = PollPacketReader(
+            Handler(handlerThread.looper),
+            network.iface.fileDescriptor.fileDescriptor,
+            1500 /* maxPacketSize */
+    ).also {
+        it.startAsyncForTest()
+        handlerThread.waitForIdle(TIMEOUT_MS)
+    }
+
     @After
     fun tearDown() {
         runAsShell(MANAGE_TEST_NETWORKS) {
@@ -1298,14 +1307,7 @@
         assumeTrue(TestUtils.shouldTestTApis())
 
         val si = makeTestServiceInfo(testNetwork1.network)
-
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         // Register service on testNetwork1
         val registrationRecord = NsdRegistrationRecord()
@@ -1345,13 +1347,7 @@
                     parseNumericAddress("2001:db8::3"))
         }
 
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         // Register service on testNetwork1
         val registrationRecord = NsdRegistrationRecord()
@@ -1391,13 +1387,7 @@
             hostname = customHostname
         }
 
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         // Register service on testNetwork1
         val registrationRecord = NsdRegistrationRecord()
@@ -1438,13 +1428,7 @@
         val registrationRecord = NsdRegistrationRecord()
         val discoveryRecord = NsdDiscoveryRecord()
         val registeredService = registerService(registrationRecord, si)
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         tryTest {
             assertNotNull(packetReader.pollForAdvertisement(serviceName, serviceType),
@@ -1518,13 +1502,7 @@
         val registrationRecord = NsdRegistrationRecord()
         val discoveryRecord = NsdDiscoveryRecord()
         val registeredService = registerService(registrationRecord, si)
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         tryTest {
             assertNotNull(packetReader.pollForAdvertisement(serviceName, serviceType),
@@ -1587,13 +1565,7 @@
         val registrationRecord = NsdRegistrationRecord()
         val discoveryRecord = NsdDiscoveryRecord()
         val registeredService = registerService(registrationRecord, si)
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         tryTest {
             assertNotNull(packetReader.pollForAdvertisement(serviceName, serviceType),
@@ -1630,13 +1602,7 @@
     fun testDiscoveryWithPtrOnlyResponse_ServiceIsFound() {
         // Register service on testNetwork1
         val discoveryRecord = NsdDiscoveryRecord()
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         nsdManager.discoverServices(
             serviceType,
@@ -1675,9 +1641,12 @@
                 assertEmpty(it.hostAddresses)
                 assertEquals(0, it.attributes.size)
             }
-        } cleanup {
+        } cleanupStep {
             nsdManager.stopServiceDiscovery(discoveryRecord)
             discoveryRecord.expectCallback<DiscoveryStopped>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -1688,79 +1657,77 @@
     fun testResolveWhenServerSendsNoAdditionalRecord() {
         // Resolve service on testNetwork1
         val resolveRecord = NsdResolveRecord()
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
-
+        val packetReader = makePacketReader()
         val si = makeTestServiceInfo(testNetwork1.network)
         nsdManager.resolveService(si, { it.run() }, resolveRecord)
 
-        val serviceFullName = "$serviceName.$serviceType.local"
-        // The query should ask for ANY, since both SRV and TXT are requested. Note legacy
-        // mdnsresponder will ask for SRV and TXT separately, and will not proceed to asking for
-        // address records without an answer for both.
-        val srvTxtQuery = packetReader.pollForQuery(serviceFullName, DnsResolver.TYPE_ANY)
-        assertNotNull(srvTxtQuery)
+        tryTest {
+            val serviceFullName = "$serviceName.$serviceType.local"
+            // The query should ask for ANY, since both SRV and TXT are requested. Note legacy
+            // mdnsresponder will ask for SRV and TXT separately, and will not proceed to asking for
+            // address records without an answer for both.
+            val srvTxtQuery = packetReader.pollForQuery(serviceFullName, DnsResolver.TYPE_ANY)
+            assertNotNull(srvTxtQuery)
 
-        /*
-        Generated with:
-        scapy.raw(scapy.dns_compress(scapy.DNS(rd=0, qr=1, aa=1, qd = None, an =
-            scapy.DNSRRSRV(rrname='NsdTest123456789._nmt123456789._tcp.local',
-                rclass=0x8001, port=31234, target='testhost.local', ttl=120) /
-            scapy.DNSRR(rrname='NsdTest123456789._nmt123456789._tcp.local', type='TXT', ttl=120,
-                rdata='testkey=testvalue')
-        ))).hex()
-         */
-        val srvTxtResponsePayload = HexDump.hexStringToByteArray(
-            "000084000000000200000000104" +
-                "e7364546573743132333435363738390d5f6e6d74313233343536373839045f746370056c6f6" +
-                "3616c0000218001000000780011000000007a020874657374686f7374c030c00c00100001000" +
-                "00078001211746573746b65793d7465737476616c7565"
-        )
-        replaceServiceNameAndTypeWithTestSuffix(srvTxtResponsePayload)
-        packetReader.sendResponse(buildMdnsPacket(srvTxtResponsePayload))
+            /*
+            Generated with:
+            scapy.raw(scapy.dns_compress(scapy.DNS(rd=0, qr=1, aa=1, qd = None, an =
+                scapy.DNSRRSRV(rrname='NsdTest123456789._nmt123456789._tcp.local',
+                    rclass=0x8001, port=31234, target='testhost.local', ttl=120) /
+                scapy.DNSRR(rrname='NsdTest123456789._nmt123456789._tcp.local', type='TXT', ttl=120,
+                    rdata='testkey=testvalue')
+            ))).hex()
+             */
+            val srvTxtResponsePayload = HexDump.hexStringToByteArray(
+                    "000084000000000200000000104" +
+                            "e7364546573743132333435363738390d5f6e6d74313233343536373839045f7463" +
+                            "70056c6f63616c0000218001000000780011000000007a020874657374686f7374c" +
+                            "030c00c0010000100000078001211746573746b65793d7465737476616c7565"
+            )
+            replaceServiceNameAndTypeWithTestSuffix(srvTxtResponsePayload)
+            packetReader.sendResponse(buildMdnsPacket(srvTxtResponsePayload))
 
-        val testHostname = "testhost.local"
-        val addressQuery = packetReader.pollForQuery(
-            testHostname,
-            DnsResolver.TYPE_A,
-            DnsResolver.TYPE_AAAA
-        )
-        assertNotNull(addressQuery)
+            val testHostname = "testhost.local"
+            val addressQuery = packetReader.pollForQuery(
+                    testHostname,
+                    DnsResolver.TYPE_A,
+                    DnsResolver.TYPE_AAAA
+            )
+            assertNotNull(addressQuery)
 
-        /*
-        Generated with:
-        scapy.raw(scapy.dns_compress(scapy.DNS(rd=0, qr=1, aa=1, qd = None, an =
-            scapy.DNSRR(rrname='testhost.local', type='A', ttl=120,
-                rdata='192.0.2.123') /
-            scapy.DNSRR(rrname='testhost.local', type='AAAA', ttl=120,
-                rdata='2001:db8::123')
-        ))).hex()
-         */
-        val addressPayload = HexDump.hexStringToByteArray(
-            "0000840000000002000000000874657374" +
-                "686f7374056c6f63616c0000010001000000780004c000027bc00c001c000100000078001020" +
-                "010db8000000000000000000000123"
-        )
-        packetReader.sendResponse(buildMdnsPacket(addressPayload))
+            /*
+            Generated with:
+            scapy.raw(scapy.dns_compress(scapy.DNS(rd=0, qr=1, aa=1, qd = None, an =
+                scapy.DNSRR(rrname='testhost.local', type='A', ttl=120,
+                    rdata='192.0.2.123') /
+                scapy.DNSRR(rrname='testhost.local', type='AAAA', ttl=120,
+                    rdata='2001:db8::123')
+            ))).hex()
+             */
+            val addressPayload = HexDump.hexStringToByteArray(
+                    "0000840000000002000000000874657374" +
+                            "686f7374056c6f63616c0000010001000000780004c000027bc00c001c000100000" +
+                            "078001020010db8000000000000000000000123"
+            )
+            packetReader.sendResponse(buildMdnsPacket(addressPayload))
 
-        val serviceResolved = resolveRecord.expectCallback<ServiceResolved>()
-        serviceResolved.serviceInfo.let {
-            assertEquals(serviceName, it.serviceName)
-            assertEquals(".$serviceType", it.serviceType)
-            assertEquals(testNetwork1.network, it.network)
-            assertEquals(31234, it.port)
-            assertEquals(1, it.attributes.size)
-            assertArrayEquals("testvalue".encodeToByteArray(), it.attributes["testkey"])
+            val serviceResolved = resolveRecord.expectCallback<ServiceResolved>()
+            serviceResolved.serviceInfo.let {
+                assertEquals(serviceName, it.serviceName)
+                assertEquals(".$serviceType", it.serviceType)
+                assertEquals(testNetwork1.network, it.network)
+                assertEquals(31234, it.port)
+                assertEquals(1, it.attributes.size)
+                assertArrayEquals("testvalue".encodeToByteArray(), it.attributes["testkey"])
+            }
+            assertEquals(
+                    setOf(parseNumericAddress("192.0.2.123"), parseNumericAddress("2001:db8::123")),
+                    serviceResolved.serviceInfo.hostAddresses.toSet()
+            )
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
-        assertEquals(
-                setOf(parseNumericAddress("192.0.2.123"), parseNumericAddress("2001:db8::123")),
-                serviceResolved.serviceInfo.hostAddresses.toSet()
-        )
     }
 
     @Test
@@ -1774,13 +1741,9 @@
         // Register service on testNetwork1
         val registrationRecord = NsdRegistrationRecord()
         var nsResponder: NSResponder? = null
+        val packetReader = makePacketReader()
         tryTest {
             registerService(registrationRecord, si)
-            val packetReader = PollPacketReader(Handler(handlerThread.looper),
-                    testNetwork1.iface.fileDescriptor.fileDescriptor, 1500 /* maxPacketSize */)
-            packetReader.startAsyncForTest()
-
-            handlerThread.waitForIdle(TIMEOUT_MS)
             /*
             Send a "query unicast" query.
             Generated with:
@@ -1805,10 +1768,13 @@
                         pkt.dstAddr == testSrcAddr
             }
             assertNotNull(reply)
-        } cleanup {
+        } cleanupStep {
             nsResponder?.stop()
             nsdManager.unregisterService(registrationRecord)
             registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -1824,13 +1790,9 @@
         // Register service on testNetwork1
         val registrationRecord = NsdRegistrationRecord()
         var nsResponder: NSResponder? = null
+        val packetReader = makePacketReader()
         tryTest {
             registerService(registrationRecord, si)
-            val packetReader = PollPacketReader(Handler(handlerThread.looper),
-                    testNetwork1.iface.fileDescriptor.fileDescriptor, 1500 /* maxPacketSize */)
-            packetReader.startAsyncForTest()
-
-            handlerThread.waitForIdle(TIMEOUT_MS)
             /*
             Send a query with a known answer. Expect to receive a response containing TXT record
             only.
@@ -1895,10 +1857,13 @@
                         pkt.dstAddr == testSrcAddr
             }
             assertNotNull(reply2)
-        } cleanup {
+        } cleanupStep {
             nsResponder?.stop()
             nsdManager.unregisterService(registrationRecord)
             registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -1914,13 +1879,9 @@
         // Register service on testNetwork1
         val registrationRecord = NsdRegistrationRecord()
         var nsResponder: NSResponder? = null
+        val packetReader = makePacketReader()
         tryTest {
             registerService(registrationRecord, si)
-            val packetReader = PollPacketReader(Handler(handlerThread.looper),
-                    testNetwork1.iface.fileDescriptor.fileDescriptor, 1500 /* maxPacketSize */)
-            packetReader.startAsyncForTest()
-
-            handlerThread.waitForIdle(TIMEOUT_MS)
             /*
             Send a query with truncated bit set.
             Generated with:
@@ -1976,10 +1937,13 @@
                         pkt.dstAddr == testSrcAddr
             }
             assertNotNull(reply)
-        } cleanup {
+        } cleanupStep {
             nsResponder?.stop()
             nsdManager.unregisterService(registrationRecord)
             registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -1991,13 +1955,7 @@
 
         // Register service on testNetwork1
         val discoveryRecord = NsdDiscoveryRecord()
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         nsdManager.discoverServices(
             serviceType,
@@ -2043,9 +2001,12 @@
                         pkt.isReplyFor("$serviceType.local", DnsResolver.TYPE_PTR)
             }
             assertNotNull(query)
-        } cleanup {
+        } cleanupStep {
             nsdManager.stopServiceDiscovery(discoveryRecord)
             discoveryRecord.expectCallback<DiscoveryStopped>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -2355,14 +2316,7 @@
             it.port = TEST_PORT
             it.publicKey = publicKey
         }
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
-
+        val packetReader = makePacketReader()
         val registrationRecord = NsdRegistrationRecord()
         val discoveryRecord = NsdDiscoveryRecord()
         tryTest {
@@ -2394,8 +2348,11 @@
             nsdManager.stopServiceDiscovery(discoveryRecord)
 
             discoveryRecord.expectCallback<DiscoveryStopped>()
-        } cleanup {
+        } cleanupStep {
             nsdManager.unregisterService(registrationRecord)
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -2410,14 +2367,7 @@
                     parseNumericAddress("2001:db8::2"))
             it.publicKey = publicKey
         }
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
-
+        val packetReader = makePacketReader()
         val registrationRecord = NsdRegistrationRecord()
         tryTest {
             registerService(registrationRecord, si)
@@ -2439,8 +2389,11 @@
                         it.nsType == DnsResolver.TYPE_A
             }
             assertEquals(3, addressRecords.size)
-        } cleanup {
+        } cleanupStep {
             nsdManager.unregisterService(registrationRecord)
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -2467,14 +2420,7 @@
             it.hostAddresses = listOf()
             it.publicKey = publicKey
         }
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
-
+        val packetReader = makePacketReader()
         val registrationRecord1 = NsdRegistrationRecord()
         val registrationRecord2 = NsdRegistrationRecord()
         tryTest {
@@ -2508,9 +2454,12 @@
             assertTrue(keyRecords.any { it.dName == "$customHostname.local" })
             assertTrue(keyRecords.all { it.ttl == NAME_RECORDS_TTL_MILLIS })
             assertTrue(keyRecords.all { it.rr.contentEquals(publicKey) })
-        } cleanup {
+        } cleanupStep {
             nsdManager.unregisterService(registrationRecord1)
             nsdManager.unregisterService(registrationRecord2)
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
@@ -2582,13 +2531,7 @@
             "test_nsd_avoid_advertising_empty_txt_records",
             "1"
         )
-        val packetReader = PollPacketReader(
-                Handler(handlerThread.looper),
-                testNetwork1.iface.fileDescriptor.fileDescriptor,
-                1500 /* maxPacketSize */
-        )
-        packetReader.startAsyncForTest()
-        handlerThread.waitForIdle(TIMEOUT_MS)
+        val packetReader = makePacketReader()
 
         // Test behavior described in RFC6763 6.1: empty TXT records are not allowed, but TXT
         // records with a zero length string are equivalent.
@@ -2607,12 +2550,85 @@
             assertEquals(1, txtRecords.size)
             // The TXT record should contain as single zero
             assertContentEquals(byteArrayOf(0), txtRecords[0].rr)
-        } cleanup {
+        } cleanupStep {
             nsdManager.unregisterService(registrationRecord)
             registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
         }
     }
 
+    private fun verifyCachedServicesRemoval(isCachedServiceRemoved: Boolean) {
+        val si = makeTestServiceInfo(testNetwork1.network)
+        // Register service on testNetwork1
+        val registrationRecord = NsdRegistrationRecord()
+        registerService(registrationRecord, si)
+        // Register a discovery request.
+        val discoveryRecord = NsdDiscoveryRecord()
+        val packetReader = makePacketReader()
+
+        tryTest {
+            nsdManager.discoverServices(
+                    serviceType,
+                    NsdManager.PROTOCOL_DNS_SD,
+                    testNetwork1.network,
+                    { it.run() },
+                    discoveryRecord
+            )
+
+            discoveryRecord.expectCallback<DiscoveryStarted>()
+            val foundInfo = discoveryRecord.waitForServiceDiscovered(
+                    serviceName, serviceType, testNetwork1.network)
+            assertEquals(testNetwork1.network, foundInfo.network)
+            // Verify that the service is not in the cache (a query is sent).
+            assertNotNull(packetReader.pollForQuery(
+                    "$serviceType.local", DnsResolver.TYPE_PTR, timeoutMs = 0L))
+
+            // Stop discovery to trigger the cached services removal process.
+            nsdManager.stopServiceDiscovery(discoveryRecord)
+            discoveryRecord.expectCallback<DiscoveryStopped>()
+
+            val serviceFullName = "$serviceName.$serviceType.local"
+            if (isCachedServiceRemoved) {
+                Thread.sleep(100L)
+                resolveService(foundInfo)
+                // Verify the resolution query will send because cached services are remove after
+                // exceeding the retention time.
+                assertNotNull(packetReader.pollForQuery(
+                        serviceFullName, DnsResolver.TYPE_ANY, timeoutMs = 0L))
+            } else {
+                resolveService(foundInfo)
+                // Verify the resolution query will not be sent because services are still cached.
+                assertNull(packetReader.pollForQuery(
+                        serviceFullName, DnsResolver.TYPE_ANY, timeoutMs = 0L))
+            }
+        } cleanupStep {
+            nsdManager.unregisterService(registrationRecord)
+            registrationRecord.expectCallback<ServiceUnregistered>()
+        } cleanup {
+            packetReader.handler.post { packetReader.stop() }
+            handlerThread.waitForIdle(TIMEOUT_MS)
+        }
+    }
+
+    @Test
+    fun testRemoveCachedServices() {
+        deviceConfigRule.setConfig(NAMESPACE_TETHERING, "test_nsd_cached_services_removal", "1")
+        verifyCachedServicesRemoval(isCachedServiceRemoved = false)
+    }
+
+    @Test
+    fun testRemoveCachedServices_ShortRetentionTime() {
+        deviceConfigRule.setConfig(NAMESPACE_TETHERING, "test_nsd_cached_services_removal", "1")
+        deviceConfigRule.setConfig(
+                NAMESPACE_TETHERING,
+                "test_nsd_cached_services_retention_time",
+                "1"
+        )
+        verifyCachedServicesRemoval(isCachedServiceRemoved = true)
+    }
+
     private fun hasServiceTypeClientsForNetwork(clients: List<String>, network: Network): Boolean {
         return clients.any { client -> client.substring(
                 client.indexOf("network=") + "network=".length,
diff --git a/tests/cts/tethering/Android.bp b/tests/cts/tethering/Android.bp
index d9bc7f7..83818be 100644
--- a/tests/cts/tethering/Android.bp
+++ b/tests/cts/tethering/Android.bp
@@ -19,10 +19,7 @@
 
 java_defaults {
     name: "CtsTetheringTestDefaults",
-    defaults: [
-        "cts_defaults",
-        "framework-connectivity-test-defaults",
-    ],
+    defaults: ["cts_defaults"],
 
     libs: [
         "android.test.base.stubs.system",
diff --git a/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java b/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java
index a07c9ea..1454d9a 100644
--- a/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java
+++ b/tests/cts/tethering/src/android/tethering/cts/TetheringManagerTest.java
@@ -32,7 +32,6 @@
 import static android.net.TetheringManager.TETHER_ERROR_NO_CHANGE_TETHERING_PERMISSION;
 import static android.net.TetheringManager.TETHER_ERROR_NO_ERROR;
 import static android.net.cts.util.CtsTetheringUtils.isAnyIfaceMatch;
-import static android.os.Process.INVALID_UID;
 
 import static com.android.testutils.TestPermissionUtil.runAsShell;
 
@@ -245,35 +244,24 @@
         assertFalse(tr.isExemptFromEntitlementCheck());
         assertTrue(tr.getShouldShowEntitlementUi());
         assertEquals(softApConfiguration, tr.getSoftApConfiguration());
-        assertEquals(INVALID_UID, tr.getUid());
-        assertNull(tr.getPackageName());
 
         final LinkAddress localAddr = new LinkAddress("192.168.24.5/24");
         final LinkAddress clientAddr = new LinkAddress("192.168.24.100/24");
         final TetheringRequest tr2 = new TetheringRequest.Builder(TETHERING_USB)
                 .setStaticIpv4Addresses(localAddr, clientAddr)
                 .setExemptFromEntitlementCheck(true)
-                .setShouldShowEntitlementUi(false)
-                .build();
-        int uid = 1000;
-        String packageName = "package";
-        tr2.setUid(uid);
-        tr2.setPackageName(packageName);
+                .setShouldShowEntitlementUi(false).build();
 
         assertEquals(localAddr, tr2.getLocalIpv4Address());
         assertEquals(clientAddr, tr2.getClientStaticIpv4Address());
         assertEquals(TETHERING_USB, tr2.getTetheringType());
         assertTrue(tr2.isExemptFromEntitlementCheck());
         assertFalse(tr2.getShouldShowEntitlementUi());
-        assertEquals(uid, tr2.getUid());
-        assertEquals(packageName, tr2.getPackageName());
 
         final TetheringRequest tr3 = new TetheringRequest.Builder(TETHERING_USB)
                 .setStaticIpv4Addresses(localAddr, clientAddr)
                 .setExemptFromEntitlementCheck(true)
                 .setShouldShowEntitlementUi(false).build();
-        tr3.setUid(uid);
-        tr3.setPackageName(packageName);
         assertEquals(tr2, tr3);
     }
 
diff --git a/tests/unit/java/android/net/TrafficStatsTest.kt b/tests/unit/java/android/net/TrafficStatsTest.kt
index c61541e..6e8f3db 100644
--- a/tests/unit/java/android/net/TrafficStatsTest.kt
+++ b/tests/unit/java/android/net/TrafficStatsTest.kt
@@ -1,46 +1,92 @@
 /*
-* 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.
-*/
+ * 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
 
-import android.net.TrafficStats.getValueForTypeFromFirstEntry
-import android.net.TrafficStats.TYPE_RX_BYTES
+import android.net.NetworkStats.UID_ALL
 import android.net.TrafficStats.UNSUPPORTED
+import android.net.netstats.StatsResult
 import android.os.Build
+import android.os.Process.SYSTEM_UID
 import com.android.testutils.DevSdkIgnoreRule
 import com.android.testutils.DevSdkIgnoreRunner
+import org.junit.After
+import org.junit.Assert.assertEquals
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import kotlin.test.assertEquals
-
-private const val TEST_IFACE1 = "test_iface1"
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Mockito.doReturn
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
 
 @RunWith(DevSdkIgnoreRunner::class)
-@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.TIRAMISU)
+@DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2)
 class TrafficStatsTest {
+    private val binder = mock(INetworkStatsService::class.java)
+    private val myUid = android.os.Process.myUid()
+    private val notMyUid = myUid + 1
+    private val mockSystemUidStatsResult = StatsResult(1L, 2L, 3L, 4L)
+    private val mockMyUidStatsResult = StatsResult(5L, 6L, 7L, 8L)
+    private val mockNotMyUidStatsResult = StatsResult(9L, 10L, 11L, 12L)
+    private val unsupportedStatsResult =
+            StatsResult(UNSUPPORTED.toLong(), UNSUPPORTED.toLong(),
+                    UNSUPPORTED.toLong(), UNSUPPORTED.toLong())
 
-    @Test
-    fun testGetValueForTypeFromFirstEntry() {
-        var stats: NetworkStats = NetworkStats(0, 0)
-        // empty stats
-        assertEquals(getValueForTypeFromFirstEntry(stats, TYPE_RX_BYTES), UNSUPPORTED.toLong())
-        // invalid type
-        stats.insertEntry(TEST_IFACE1, 1, 2, 3, 4)
-        assertEquals(getValueForTypeFromFirstEntry(stats, 1000), UNSUPPORTED.toLong())
-        // valid type
-        assertEquals(getValueForTypeFromFirstEntry(stats, TYPE_RX_BYTES), 1)
+    @Before
+    fun setUp() {
+        TrafficStats.setServiceForTest(binder)
+        doReturn(mockSystemUidStatsResult).`when`(binder).getUidStats(SYSTEM_UID)
+        doReturn(mockMyUidStatsResult).`when`(binder).getUidStats(myUid)
+        doReturn(mockNotMyUidStatsResult).`when`(binder).getUidStats(notMyUid)
     }
-}
\ No newline at end of file
+
+    @After
+    fun tearDown() {
+        TrafficStats.setServiceForTest(null)
+        TrafficStats.setMyUidForTest(UID_ALL)
+    }
+
+    private fun assertUidStats(uid: Int, stats: StatsResult) {
+        assertEquals(stats.rxBytes, TrafficStats.getUidRxBytes(uid))
+        assertEquals(stats.rxPackets, TrafficStats.getUidRxPackets(uid))
+        assertEquals(stats.txBytes, TrafficStats.getUidTxBytes(uid))
+        assertEquals(stats.txPackets, TrafficStats.getUidTxPackets(uid))
+    }
+
+    // Verify a normal caller could get a quick UNSUPPORTED result in the TrafficStats
+    // without accessing the service if query stats other than itself.
+    @Test
+    fun testGetUidStats_appCaller() {
+        assertUidStats(SYSTEM_UID, unsupportedStatsResult)
+        assertUidStats(notMyUid, unsupportedStatsResult)
+        verify(binder, never()).getUidStats(anyInt())
+        assertUidStats(myUid, mockMyUidStatsResult)
+    }
+
+    // Verify that callers with SYSTEM_UID can access network
+    // stats for other UIDs. While this behavior is not officially documented
+    // in the API, it exists for compatibility with existing callers that may
+    // rely on it.
+    @Test
+    fun testGetUidStats_systemCaller() {
+        TrafficStats.setMyUidForTest(SYSTEM_UID)
+        assertUidStats(SYSTEM_UID, mockSystemUidStatsResult)
+        assertUidStats(myUid, mockMyUidStatsResult)
+        assertUidStats(notMyUid, mockNotMyUidStatsResult)
+    }
+}
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index ef4c44d..b528480 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -56,7 +56,6 @@
 import static android.net.TrafficStats.MB_IN_BYTES;
 import static android.net.TrafficStats.UID_REMOVED;
 import static android.net.TrafficStats.UID_TETHERING;
-import static android.net.TrafficStats.getValueForTypeFromFirstEntry;
 import static android.net.connectivity.ConnectivityCompatChanges.ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID_TAG;
@@ -79,6 +78,7 @@
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_ATTEMPTS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_FALLBACKS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_SUCCESSES_COUNTER_NAME;
+import static com.android.server.net.NetworkStatsService.TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG;
 import static com.android.server.net.NetworkStatsService.TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
@@ -128,10 +128,12 @@
 import android.net.TestNetworkSpecifier;
 import android.net.TetherStatsParcel;
 import android.net.TetheringManager;
-import android.net.TrafficStats;
 import android.net.UnderlyingNetworkInfo;
+import android.net.netstats.StatsResult;
+import android.net.netstats.TrafficStatsRateLimitCacheConfig;
 import android.net.netstats.provider.INetworkStatsProviderCallback;
 import android.net.wifi.WifiInfo;
+import android.os.Build;
 import android.os.DropBoxManager;
 import android.os.Handler;
 import android.os.HandlerThread;
@@ -209,7 +211,6 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Function;
 
 /**
  * Tests for {@link NetworkStatsService}.
@@ -223,6 +224,8 @@
 // NetworkStatsService is not updatable before T, so tests do not need to be backwards compatible
 @DevSdkIgnoreRule.IgnoreUpTo(SC_V2)
 public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
+    @Rule
+    public final DevSdkIgnoreRule ignoreRule = new DevSdkIgnoreRule();
 
     private static final String TAG = "NetworkStatsServiceTest";
 
@@ -621,8 +624,9 @@
         }
 
         @Override
-        public boolean alwaysUseTrafficStatsServiceRateLimitCache(Context ctx) {
-            return mFeatureFlags.getOrDefault(
+        public boolean isTrafficStatsServiceRateLimitCacheEnabled(Context ctx,
+                boolean isClientCacheEnabled) {
+            return !isClientCacheEnabled && mFeatureFlags.getOrDefault(
                     TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG, false);
         }
 
@@ -643,6 +647,19 @@
         }
 
         @Override
+        public TrafficStatsRateLimitCacheConfig getTrafficStatsRateLimitCacheClientSideConfig(
+                @NonNull Context ctx) {
+            final TrafficStatsRateLimitCacheConfig config =
+                    new TrafficStatsRateLimitCacheConfig.Builder()
+                            .setIsCacheEnabled(mFeatureFlags.getOrDefault(
+                                    TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG, false))
+                            .setExpiryDurationMs(DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS)
+                            .setMaxEntries(DEFAULT_TRAFFIC_STATS_SERVICE_CACHE_MAX_ENTRIES)
+                            .build();
+            return config;
+        }
+
+        @Override
         public boolean isChangeEnabled(long changeId, int uid) {
             return mCompatChanges.getOrDefault(changeId, true);
         }
@@ -2453,11 +2470,49 @@
         assertUidTotal(sTemplateWifi, UID_GREEN, 64L, 3L, 1024L, 8L, 0);
     }
 
+    @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false)
+    @Test
+    public void testGetRateLimitCacheConfig_featureDisabled() {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, false);
+        assertFalse(mService.getRateLimitCacheConfig().isCacheEnabled);
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, true);
+        assertFalse(mService.getRateLimitCacheConfig().isCacheEnabled);
+    }
+
+    @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    @Test
+    public void testGetRateLimitCacheConfig_vOrAbove() {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, false);
+        assertTrue(mService.getRateLimitCacheConfig().isCacheEnabled);
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, true);
+        assertTrue(mService.getRateLimitCacheConfig().isCacheEnabled);
+    }
+
+    @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @DevSdkIgnoreRule.IgnoreAfter(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    @Test
+    public void testGetRateLimitCacheConfig_belowV() {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, false);
+        assertFalse(mService.getRateLimitCacheConfig().isCacheEnabled);
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, true);
+        assertTrue(mService.getRateLimitCacheConfig().isCacheEnabled);
+    }
+
+    @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @FeatureFlag(name = TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @Test
+    public void testTrafficStatsRateLimitCache_clientCacheEnabledDisableServiceCache()
+            throws Exception {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, true);
+        doTestTrafficStatsRateLimitCache(false /* expectCached */);
+    }
+
     @FeatureFlag(name = TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false)
     @Test
     public void testTrafficStatsRateLimitCache_disabledWithCompatChangeEnabled() throws Exception {
         mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, true);
-        doTestTrafficStatsRateLimitCache(true /* expectCached */);
+        doTestTrafficStatsRateLimitCache(false /* expectCached */);
     }
 
     @FeatureFlag(name = TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG)
@@ -2475,8 +2530,19 @@
     }
 
     @FeatureFlag(name = TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @DevSdkIgnoreRule.IgnoreAfter(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
     @Test
-    public void testTrafficStatsRateLimitCache_enabledWithCompatChangeDisabled() throws Exception {
+    public void testTrafficStatsRateLimitCache_enabledWithCompatChangeDisabled_belowV()
+            throws Exception {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, false);
+        doTestTrafficStatsRateLimitCache(false /* expectCached */);
+    }
+
+    @FeatureFlag(name = TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    @Test
+    public void testTrafficStatsRateLimitCache_enabledWithCompatChangeDisabled_vOrAbove()
+            throws Exception {
         mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, false);
         doTestTrafficStatsRateLimitCache(true /* expectCached */);
     }
@@ -2515,22 +2581,18 @@
     // Assert for 3 different API return values respectively.
     private void assertTrafficStatsValues(String iface, int uid, long rxBytes, long rxPackets,
             long txBytes, long txPackets) {
-        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
-                (type) -> getValueForTypeFromFirstEntry(mService.getTypelessTotalStats(), type));
-        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
-                (type) -> getValueForTypeFromFirstEntry(
-                        mService.getTypelessIfaceStats(iface), type)
-        );
-        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
-                (type) -> getValueForTypeFromFirstEntry(mService.getTypelessUidStats(uid), type));
+        assertStatsResultEquals(mService.getTotalStats(), rxBytes, rxPackets, txBytes, txPackets);
+        assertStatsResultEquals(mService.getIfaceStats(iface), rxBytes, rxPackets, txBytes,
+                txPackets);
+        assertStatsResultEquals(mService.getUidStats(uid), rxBytes, rxPackets, txBytes, txPackets);
     }
 
-    private void assertTrafficStatsValuesThat(long rxBytes, long rxPackets, long txBytes,
-            long txPackets, Function<Integer, Long> fetcher) {
-        assertEquals(rxBytes, (long) fetcher.apply(TrafficStats.TYPE_RX_BYTES));
-        assertEquals(rxPackets, (long) fetcher.apply(TrafficStats.TYPE_RX_PACKETS));
-        assertEquals(txBytes, (long) fetcher.apply(TrafficStats.TYPE_TX_BYTES));
-        assertEquals(txPackets, (long) fetcher.apply(TrafficStats.TYPE_TX_PACKETS));
+    private void assertStatsResultEquals(StatsResult stats, long rxBytes, long rxPackets,
+            long txBytes, long txPackets) {
+        assertEquals(rxBytes, stats.rxBytes);
+        assertEquals(rxPackets, stats.rxPackets);
+        assertEquals(txBytes, stats.txBytes);
+        assertEquals(txPackets, stats.txPackets);
     }
 
     private void assertShouldRunComparison(boolean expected, boolean isDebuggable) {
diff --git a/thread/tests/integration/src/android/net/thread/InternetAccessTest.kt b/thread/tests/integration/src/android/net/thread/InternetAccessTest.kt
new file mode 100644
index 0000000..9b787f5
--- /dev/null
+++ b/thread/tests/integration/src/android/net/thread/InternetAccessTest.kt
@@ -0,0 +1,173 @@
+/*
+ * 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
+
+import android.content.Context
+import android.net.DnsResolver.CLASS_IN
+import android.net.DnsResolver.TYPE_A
+import android.net.DnsResolver.TYPE_AAAA
+import android.net.InetAddresses.parseNumericAddress
+import android.net.thread.utils.FullThreadDevice
+import android.net.thread.utils.InfraNetworkDevice
+import android.net.thread.utils.IntegrationTestUtils.DEFAULT_DATASET
+import android.net.thread.utils.IntegrationTestUtils.joinNetworkAndWaitForOmr
+import android.net.thread.utils.IntegrationTestUtils.newPacketReader
+import android.net.thread.utils.IntegrationTestUtils.setUpInfraNetwork
+import android.net.thread.utils.IntegrationTestUtils.startInfraDeviceAndWaitForOnLinkAddr
+import android.net.thread.utils.IntegrationTestUtils.tearDownInfraNetwork
+import android.net.thread.utils.IntegrationTestUtils.waitFor
+import android.net.thread.utils.OtDaemonController
+import android.net.thread.utils.TestDnsServer
+import android.net.thread.utils.ThreadFeatureCheckerRule
+import android.net.thread.utils.ThreadFeatureCheckerRule.RequiresSimulationThreadDevice
+import android.net.thread.utils.ThreadFeatureCheckerRule.RequiresThreadFeature
+import android.net.thread.utils.ThreadNetworkControllerWrapper
+import android.os.Handler
+import android.os.HandlerThread
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.LargeTest
+import com.android.net.module.util.DnsPacket
+import com.android.net.module.util.DnsPacket.ANSECTION
+import com.android.testutils.PollPacketReader
+import com.android.testutils.TestNetworkTracker
+import com.google.common.truth.Truth.assertThat
+import java.net.Inet4Address
+import java.net.InetAddress
+import java.time.Duration
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Integration test cases for Thread Internet Access features. */
+@RunWith(AndroidJUnit4::class)
+@RequiresThreadFeature
+@RequiresSimulationThreadDevice
+@LargeTest
+class InternetAccessTest {
+    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 ANSWER_RECORDS =
+        listOf(
+            DnsPacket.DnsRecord.makeAOrAAAARecord(
+                ANSECTION,
+                "google.com",
+                CLASS_IN,
+                30 /* ttl */,
+                parseNumericAddress("1.2.3.4"),
+            ),
+            DnsPacket.DnsRecord.makeAOrAAAARecord(
+                ANSECTION,
+                "google.com",
+                CLASS_IN,
+                30 /* ttl */,
+                parseNumericAddress("2001::234"),
+            ),
+        )
+
+    @get:Rule val threadRule = ThreadFeatureCheckerRule()
+
+    private val context: Context = ApplicationProvider.getApplicationContext()
+    private val controller = requireNotNull(ThreadNetworkControllerWrapper.newInstance(context))
+    private lateinit var otCtl: OtDaemonController
+    private lateinit var handlerThread: HandlerThread
+    private lateinit var handler: Handler
+    private lateinit var infraNetworkTracker: TestNetworkTracker
+    private lateinit var ftds: ArrayList<FullThreadDevice>
+    private lateinit var infraNetworkReader: PollPacketReader
+    private lateinit var infraDevice: InfraNetworkDevice
+    private lateinit var dnsServer: TestDnsServer
+
+    @Before
+    @Throws(Exception::class)
+    fun setUp() {
+        // TODO: b/323301831 - This is a workaround to avoid unnecessary delay to re-form a network
+        otCtl = OtDaemonController()
+        otCtl.factoryReset()
+
+        handlerThread = HandlerThread(javaClass.simpleName)
+        handlerThread.start()
+        handler = Handler(handlerThread.looper)
+        ftds = ArrayList()
+
+        infraNetworkTracker = setUpInfraNetwork(context, controller)
+        controller.setEnabledAndWait(true)
+        controller.joinAndWait(DEFAULT_DATASET)
+
+        // Creates a infra network device.
+        infraNetworkReader = newPacketReader(infraNetworkTracker.testIface, handler)
+        infraDevice = startInfraDeviceAndWaitForOnLinkAddr(infraNetworkReader)
+
+        // Create a DNS server
+        dnsServer = TestDnsServer(infraNetworkReader, DNS_SERVER_ADDR, ANSWER_RECORDS)
+
+        // Create Ftds
+        for (i in 0 until NUM_FTD) {
+            ftds.add(FullThreadDevice(15 + i /* node ID */))
+        }
+    }
+
+    @After
+    @Throws(Exception::class)
+    fun tearDown() {
+        controller.setTestNetworkAsUpstreamAndWait(null)
+        controller.leaveAndWait()
+        tearDownInfraNetwork(infraNetworkTracker)
+
+        dnsServer.stop()
+
+        handlerThread.quitSafely()
+        handlerThread.join()
+
+        ftds.forEach { it.destroy() }
+        ftds.clear()
+    }
+
+    @Test
+    fun nat64Enabled_threadDeviceResolvesHost_hostIsResolved() {
+        controller.setNat64EnabledAndWait(true)
+        waitFor({ otCtl.hasNat64PrefixInNetdata() }, Duration.ofSeconds(10))
+        val ftd = ftds[0]
+        joinNetworkAndWaitForOmr(ftd, DEFAULT_DATASET)
+        dnsServer.start()
+
+        val ipv4Addresses =
+            ftd.resolveHost("google.com", TYPE_A).map { extractIpv4AddressFromMappedAddress(it) }
+        assertThat(ipv4Addresses).isEqualTo(listOf(parseNumericAddress("1.2.3.4")))
+        val ipv6Addresses = ftd.resolveHost("google.com", TYPE_AAAA)
+        assertThat(ipv6Addresses).isEqualTo(listOf(parseNumericAddress("2001::234")))
+    }
+
+    @Test
+    fun nat64Disabled_threadDeviceResolvesHost_hostIsNotResolved() {
+        controller.setNat64EnabledAndWait(false)
+        val ftd = ftds[0]
+        joinNetworkAndWaitForOmr(ftd, DEFAULT_DATASET)
+        dnsServer.start()
+
+        assertThat(ftd.resolveHost("google.com", TYPE_A)).isEmpty()
+        assertThat(ftd.resolveHost("google.com", TYPE_AAAA)).isEmpty()
+    }
+
+    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 083a841..5b532c7 100644
--- a/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
+++ b/thread/tests/integration/src/android/net/thread/utils/FullThreadDevice.java
@@ -15,6 +15,8 @@
  */
 package android.net.thread.utils;
 
+import static android.net.DnsResolver.TYPE_A;
+import static android.net.DnsResolver.TYPE_AAAA;
 import static android.net.thread.utils.IntegrationTestUtils.SERVICE_DISCOVERY_TIMEOUT;
 import static android.net.thread.utils.IntegrationTestUtils.waitFor;
 
@@ -354,6 +356,31 @@
         executeCommand("dns config " + address);
     }
 
+    /** Resolves the {@code queryType} record of the {@code hostname} via DNS. */
+    public List<InetAddress> resolveHost(String hostname, int queryType) {
+        // CLI output:
+        // DNS response for hostname.com. - fd12::abc1 TTL:50 fd12::abc2 TTL:50 fd12::abc3 TTL:50
+
+        String command;
+        switch (queryType) {
+            case TYPE_A -> command = "resolve4";
+            case TYPE_AAAA -> command = "resolve";
+            default -> throw new IllegalArgumentException("Invalid query type: " + queryType);
+        }
+        final List<InetAddress> addresses = new ArrayList<>();
+        String line;
+        try {
+            line = executeCommand("dns " + command + " " + hostname).get(0);
+        } catch (IllegalStateException e) {
+            return addresses;
+        }
+        final String[] addressTtlPairs = line.split("-")[1].strip().split(" ");
+        for (int i = 0; i < addressTtlPairs.length; i += 2) {
+            addresses.add(InetAddresses.parseNumericAddress(addressTtlPairs[i]));
+        }
+        return addresses;
+    }
+
     /** Returns the first browsed service instance of {@code serviceType}. */
     public NsdServiceInfo browseService(String serviceType) {
         // CLI output:
diff --git a/thread/tests/integration/src/android/net/thread/utils/TestDnsServer.kt b/thread/tests/integration/src/android/net/thread/utils/TestDnsServer.kt
new file mode 100644
index 0000000..c52fc49
--- /dev/null
+++ b/thread/tests/integration/src/android/net/thread/utils/TestDnsServer.kt
@@ -0,0 +1,170 @@
+/*
+ * 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 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.nio.ByteBuffer
+import kotlin.concurrent.thread
+
+/**
+ * A class that simulates a DNS server.
+ *
+ * <p>The server responds to DNS requests with the given {@code answerRecords}.
+ *
+ * @param packetReader the packet reader to poll DNS requests from
+ * @param serverAddress the address of the DNS server
+ * @param answerRecords the records to respond to the DNS requests
+ */
+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 class TestDnsPacket : DnsPacket {
+
+        constructor(buf: ByteArray) : super(buf)
+
+        constructor(
+            header: DnsHeader,
+            qd: List<DnsRecord>,
+            an: List<DnsRecord>,
+        ) : super(header, qd, an) {}
+
+        val header = super.mHeader
+        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>,
+    ): 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 requestDnsHeader = requestDnsPacket.header
+
+        val answerRecords =
+            buildDnsAnswerRecords(requestDnsPacket.records[DnsPacket.QDSECTION], serverAnswers)
+        // TODO: return NXDOMAIN if no answer is found.
+        val responseFlags = 1 shl 15 // QR bit
+        val responseDnsHeader =
+            DnsPacket.DnsHeader(
+                requestDnsHeader.id,
+                responseFlags,
+                requestDnsPacket.records[DnsPacket.QDSECTION].size,
+                answerRecords.size,
+            )
+        val responseDnsPacket =
+            TestDnsPacket(
+                responseDnsHeader,
+                requestDnsPacket.records[DnsPacket.QDSECTION],
+                answerRecords,
+            )
+
+        val buf =
+            PacketBuilder.allocate(
+                false /* hasEther */,
+                IPPROTO_IP,
+                IPPROTO_UDP,
+                responseDnsPacket.bytes.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(responseDnsPacket.bytes)
+
+        return packetBuilder.finalizePacket()
+    }
+
+    private fun buildDnsAnswerRecords(
+        questions: List<DnsPacket.DnsRecord>,
+        serverAnswers: List<DnsPacket.DnsRecord>,
+    ): List<DnsPacket.DnsRecord> {
+        val answers = ArrayList<DnsPacket.DnsRecord>()
+        for (answer in serverAnswers) {
+            if (
+                questions.any {
+                    answer.dName.equals(it.dName, ignoreCase = true) && answer.nsType == it.nsType
+                }
+            ) {
+                answers.add(answer)
+            }
+        }
+        return answers
+    }
+}