1 /* <lambda>null2 * Copyright (C) 2024 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package android.net 18 19 import android.net.TrafficStats.UNSUPPORTED 20 import android.net.netstats.StatsResult 21 import android.net.netstats.TrafficStatsRateLimitCacheConfig 22 import android.os.Build 23 import com.android.server.net.NetworkStatsService.TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG 24 import com.android.testutils.DevSdkIgnoreRule 25 import com.android.testutils.DevSdkIgnoreRunner 26 import com.android.testutils.com.android.testutils.SetFeatureFlagsRule 27 import com.android.testutils.com.android.testutils.SetFeatureFlagsRule.FeatureFlag 28 import org.junit.After 29 import org.junit.Assert.assertEquals 30 import org.junit.Before 31 import org.junit.Rule 32 import org.junit.Test 33 import org.junit.runner.RunWith 34 import org.mockito.Mockito.clearInvocations 35 import org.mockito.Mockito.doReturn 36 import org.mockito.Mockito.mock 37 import org.mockito.Mockito.times 38 import org.mockito.Mockito.verify 39 import java.util.HashMap 40 import java.util.function.LongSupplier 41 42 const val TEST_EXPIRY_DURATION_MS = 1000 43 const val TEST_IFACE = "wlan0" 44 45 @RunWith(DevSdkIgnoreRunner::class) 46 @DevSdkIgnoreRule.IgnoreUpTo(Build.VERSION_CODES.S_V2) 47 class TrafficStatsTest { 48 private val binder = mock(INetworkStatsService::class.java) 49 private val myUid = android.os.Process.myUid() 50 private val mockMyUidStatsResult = StatsResult(5L, 6L, 7L, 8L) 51 private val mockIfaceStatsResult = StatsResult(7L, 3L, 10L, 21L) 52 private val mockTotalStatsResult = StatsResult(8L, 1L, 5L, 2L) 53 private val secondUidStatsResult = StatsResult(3L, 7L, 10L, 5L) 54 private val secondIfaceStatsResult = StatsResult(9L, 8L, 7L, 6L) 55 private val secondTotalStatsResult = StatsResult(4L, 3L, 2L, 1L) 56 private val emptyStatsResult = StatsResult(0L, 0L, 0L, 0L) 57 private val unsupportedStatsResult = 58 StatsResult(UNSUPPORTED.toLong(), UNSUPPORTED.toLong(), 59 UNSUPPORTED.toLong(), UNSUPPORTED.toLong()) 60 61 private val cacheDisabledConfig = TrafficStatsRateLimitCacheConfig.Builder() 62 .setIsCacheEnabled(false) 63 .setExpiryDurationMs(0) 64 .setMaxEntries(0) 65 .build() 66 private val cacheEnabledConfig = TrafficStatsRateLimitCacheConfig.Builder() 67 .setIsCacheEnabled(true) 68 .setExpiryDurationMs(TEST_EXPIRY_DURATION_MS) 69 .setMaxEntries(100) 70 .build() 71 private val mTestTimeSupplier = TestTimeSupplier() 72 73 private val featureFlags = HashMap<String, Boolean>() 74 75 // This will set feature flags from @FeatureFlag annotations 76 // into the map before setUp() runs. 77 @get:Rule 78 val setFeatureFlagsRule = SetFeatureFlagsRule( 79 { name, enabled -> featureFlags.put(name, enabled == true) }, 80 { name -> featureFlags.getOrDefault(name, false) } 81 ) 82 83 class TestTimeSupplier : LongSupplier { 84 private var currentTimeMillis = 0L 85 86 override fun getAsLong() = currentTimeMillis 87 88 fun advanceTime(millis: Int) { 89 currentTimeMillis += millis 90 } 91 } 92 93 @Before 94 fun setUp() { 95 TrafficStats.setServiceForTest(binder) 96 TrafficStats.setTimeSupplierForTest(mTestTimeSupplier) 97 mockStats(mockMyUidStatsResult, mockIfaceStatsResult, mockTotalStatsResult) 98 if (featureFlags.getOrDefault(TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG, false)) { 99 doReturn(cacheEnabledConfig).`when`(binder).getRateLimitCacheConfig() 100 } else { 101 doReturn(cacheDisabledConfig).`when`(binder).getRateLimitCacheConfig() 102 } 103 TrafficStats.reinitRateLimitCacheForTest() 104 } 105 106 @After 107 fun tearDown() { 108 TrafficStats.setServiceForTest(null) 109 TrafficStats.setTimeSupplierForTest(null) 110 TrafficStats.reinitRateLimitCacheForTest() 111 } 112 113 private fun assertUidStats(uid: Int, stats: StatsResult) { 114 assertEquals(stats.rxBytes, TrafficStats.getUidRxBytes(uid)) 115 assertEquals(stats.rxPackets, TrafficStats.getUidRxPackets(uid)) 116 assertEquals(stats.txBytes, TrafficStats.getUidTxBytes(uid)) 117 assertEquals(stats.txPackets, TrafficStats.getUidTxPackets(uid)) 118 } 119 120 private fun assertIfaceStats(iface: String, stats: StatsResult) { 121 assertEquals(stats.rxBytes, TrafficStats.getRxBytes(iface)) 122 assertEquals(stats.rxPackets, TrafficStats.getRxPackets(iface)) 123 assertEquals(stats.txBytes, TrafficStats.getTxBytes(iface)) 124 assertEquals(stats.txPackets, TrafficStats.getTxPackets(iface)) 125 } 126 127 private fun assertTotalStats(stats: StatsResult) { 128 assertEquals(stats.rxBytes, TrafficStats.getTotalRxBytes()) 129 assertEquals(stats.rxPackets, TrafficStats.getTotalRxPackets()) 130 assertEquals(stats.txBytes, TrafficStats.getTotalTxBytes()) 131 assertEquals(stats.txPackets, TrafficStats.getTotalTxPackets()) 132 } 133 134 private fun mockStats(uidStats: StatsResult?, ifaceStats: StatsResult?, 135 totalStats: StatsResult?) { 136 doReturn(uidStats).`when`(binder).getUidStats(myUid) 137 doReturn(ifaceStats).`when`(binder).getIfaceStats(TEST_IFACE) 138 doReturn(totalStats).`when`(binder).getTotalStats() 139 } 140 141 private fun assertStats(uidStats: StatsResult, ifaceStats: StatsResult, 142 totalStats: StatsResult) { 143 assertUidStats(myUid, uidStats) 144 assertIfaceStats(TEST_IFACE, ifaceStats) 145 assertTotalStats(totalStats) 146 } 147 148 private fun assertStatsFetchInvocations(wantedInvocations: Int) { 149 verify(binder, times(wantedInvocations)).getUidStats(myUid) 150 verify(binder, times(wantedInvocations)).getIfaceStats(TEST_IFACE) 151 verify(binder, times(wantedInvocations)).getTotalStats() 152 } 153 154 @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG) 155 @Test 156 fun testRateLimitCacheExpiry_cacheEnabled() { 157 // Initial fetch, verify binder calls. 158 assertStats(mockMyUidStatsResult, mockIfaceStatsResult, mockTotalStatsResult) 159 assertStatsFetchInvocations(1) 160 161 // Advance time within expiry, verify cached values used. 162 clearInvocations(binder) 163 mockStats(secondUidStatsResult, secondIfaceStatsResult, secondTotalStatsResult) 164 mTestTimeSupplier.advanceTime(1) 165 assertStats(mockMyUidStatsResult, mockIfaceStatsResult, mockTotalStatsResult) 166 assertStatsFetchInvocations(0) 167 168 // Advance time to expire cache, verify new values fetched. 169 clearInvocations(binder) 170 mTestTimeSupplier.advanceTime(TEST_EXPIRY_DURATION_MS) 171 assertStats(secondUidStatsResult, secondIfaceStatsResult, secondTotalStatsResult) 172 assertStatsFetchInvocations(1) 173 } 174 175 @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false) 176 @Test 177 fun testRateLimitCacheExpiry_cacheDisabled() { 178 // Initial fetch, verify binder calls. 179 assertStats(mockMyUidStatsResult, mockIfaceStatsResult, mockTotalStatsResult) 180 assertStatsFetchInvocations(4) 181 182 // Advance time within expiry, verify new values fetched. 183 clearInvocations(binder) 184 mockStats(secondUidStatsResult, secondIfaceStatsResult, secondTotalStatsResult) 185 mTestTimeSupplier.advanceTime(1) 186 assertStats(secondUidStatsResult, secondIfaceStatsResult, secondTotalStatsResult) 187 assertStatsFetchInvocations(4) 188 } 189 190 @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG) 191 @Test 192 fun testInvalidStatsNotCached_cacheEnabled() { 193 doTestInvalidStatsNotCached() 194 } 195 196 @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false) 197 @Test 198 fun testInvalidStatsNotCached_cacheDisabled() { 199 doTestInvalidStatsNotCached() 200 } 201 202 private fun doTestInvalidStatsNotCached() { 203 // Mock null stats, this usually happens when the query is not valid, 204 // e.g. query uid stats of other application. 205 mockStats(null, null, null) 206 assertStats(unsupportedStatsResult, unsupportedStatsResult, unsupportedStatsResult) 207 assertStatsFetchInvocations(4) 208 209 // Verify null stats is not cached, and mock empty stats. This usually 210 // happens when queries with non-existent interface names. 211 clearInvocations(binder) 212 mockStats(emptyStatsResult, emptyStatsResult, emptyStatsResult) 213 assertStats(emptyStatsResult, emptyStatsResult, emptyStatsResult) 214 assertStatsFetchInvocations(4) 215 216 // Verify empty result is also not cached. 217 clearInvocations(binder) 218 assertStats(emptyStatsResult, emptyStatsResult, emptyStatsResult) 219 assertStatsFetchInvocations(4) 220 } 221 222 @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG) 223 @Test 224 fun testClearRateLimitCaches_cacheEnabled() { 225 doTestClearRateLimitCaches(true) 226 } 227 228 @FeatureFlag(name = TRAFFICSTATS_CLIENT_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false) 229 @Test 230 fun testClearRateLimitCaches_cacheDisabled() { 231 doTestClearRateLimitCaches(false) 232 } 233 234 private fun doTestClearRateLimitCaches(cacheEnabled: Boolean) { 235 // Initial fetch, verify binder calls. 236 assertStats(mockMyUidStatsResult, mockIfaceStatsResult, mockTotalStatsResult) 237 assertStatsFetchInvocations(if (cacheEnabled) 1 else 4) 238 239 // Verify cached values are used. 240 clearInvocations(binder) 241 assertStats(mockMyUidStatsResult, mockIfaceStatsResult, mockTotalStatsResult) 242 assertStatsFetchInvocations(if (cacheEnabled) 0 else 4) 243 244 // Clear caches, verify fetching from the service. 245 clearInvocations(binder) 246 TrafficStats.clearRateLimitCaches() 247 mockStats(secondUidStatsResult, secondIfaceStatsResult, secondTotalStatsResult) 248 assertStats(secondUidStatsResult, secondIfaceStatsResult, secondTotalStatsResult) 249 assertStatsFetchInvocations(if (cacheEnabled) 1 else 4) 250 } 251 } 252