Merge changes Ia7ea051f,I38bdfe00 into main

* changes:
  Fix system caller cannot query stats for other uids
  Use device config to enable/disable the rate-limit cache feature
diff --git a/framework-t/src/android/net/TrafficStats.java b/framework-t/src/android/net/TrafficStats.java
index ab0aaa0..1294b3e 100644
--- a/framework-t/src/android/net/TrafficStats.java
+++ b/framework-t/src/android/net/TrafficStats.java
@@ -17,6 +17,10 @@
 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;
@@ -37,6 +41,9 @@
 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;
@@ -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,8 +1006,11 @@
 
     /** @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 StatsResult stats;
@@ -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) {
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 7572f0e..ba706d9 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -485,14 +485,17 @@
     @GuardedBy("mStatsLock")
     private long mLatestNetworkStatsUpdatedBroadcastScheduledTime = Long.MIN_VALUE;
 
+    @Nullable
     private final TrafficStatsRateLimitCache mTrafficStatsTotalCache;
+    @Nullable
     private final TrafficStatsRateLimitCache mTrafficStatsIfaceCache;
+    @Nullable
     private final TrafficStatsRateLimitCache mTrafficStatsUidCache;
     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 boolean mBroadcastNetworkStatsUpdatedRateLimitEnabled;
@@ -688,23 +691,29 @@
             mEventLogger = null;
         }
 
-        mAlwaysUseTrafficStatsServiceRateLimitCache =
-                mDeps.alwaysUseTrafficStatsServiceRateLimitCache(mContext);
+        mIsTrafficStatsServiceRateLimitCacheEnabled =
+                mDeps.isTrafficStatsServiceRateLimitCacheEnabled(mContext);
         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 +973,13 @@
         }
 
         /**
-         * Get whether TrafficStats service side rate-limit cache is always applied.
+         * Whether the service side cache is enabled for V+ device or target Sdk V+ apps.
          *
          * 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(
+        public boolean isTrafficStatsServiceRateLimitCacheEnabled(@NonNull Context ctx) {
+            return DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
                     ctx, TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG);
         }
 
@@ -2133,6 +2142,20 @@
         }
     }
 
+    /**
+     * 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 StatsResult getUidStats(int uid) {
@@ -2141,8 +2164,7 @@
             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);
@@ -2170,9 +2192,7 @@
         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);
@@ -2200,9 +2220,7 @@
     @Override
     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();
@@ -2213,9 +2231,11 @@
     @Override
     public void clearTrafficStatsRateLimitCaches() {
         PermissionUtils.enforceNetworkStackPermissionOr(mContext, NETWORK_SETTINGS);
-        mTrafficStatsUidCache.clear();
-        mTrafficStatsIfaceCache.clear();
-        mTrafficStatsTotalCache.clear();
+        if (mIsTrafficStatsServiceRateLimitCacheEnabled) {
+            mTrafficStatsUidCache.clear();
+            mTrafficStatsIfaceCache.clear();
+            mTrafficStatsTotalCache.clear();
+        }
     }
 
     private NetworkStats.Entry getProviderIfaceStats(@Nullable String iface) {
@@ -2985,8 +3005,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);
diff --git a/tests/unit/java/android/net/TrafficStatsTest.kt b/tests/unit/java/android/net/TrafficStatsTest.kt
new file mode 100644
index 0000000..6e8f3db
--- /dev/null
+++ b/tests/unit/java/android/net/TrafficStatsTest.kt
@@ -0,0 +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.
+ */
+
+package android.net
+
+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 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.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())
+
+    @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)
+    }
+
+    @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 e6f8d21..3a52927 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -131,6 +131,7 @@
 import android.net.netstats.StatsResult;
 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;
@@ -221,6 +222,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";
 
@@ -619,7 +622,7 @@
         }
 
         @Override
-        public boolean alwaysUseTrafficStatsServiceRateLimitCache(Context ctx) {
+        public boolean isTrafficStatsServiceRateLimitCacheEnabled(Context ctx) {
             return mFeatureFlags.getOrDefault(
                     TRAFFICSTATS_SERVICE_RATE_LIMIT_CACHE_ENABLED_FLAG, false);
         }
@@ -2455,7 +2458,7 @@
     @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)
@@ -2473,8 +2476,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 */);
     }