Merge "Correctly count nri uid request counts" am: 5ab962df61 am: b039be7617 am: d2c14296e4

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/1655051

Change-Id: Ibec1f37f2062827cf390eb344dd30447bdabec13
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 809ef41..4b96119 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -321,7 +321,8 @@
     private static final int MAX_NETWORK_REQUESTS_PER_UID = 100;
 
     // The maximum number of network request allowed for system UIDs before an exception is thrown.
-    private static final int MAX_NETWORK_REQUESTS_PER_SYSTEM_UID = 250;
+    @VisibleForTesting
+    static final int MAX_NETWORK_REQUESTS_PER_SYSTEM_UID = 250;
 
     @VisibleForTesting
     protected int mLingerDelayMs;  // Can't be final, or test subclass constructors can't change it.
@@ -338,7 +339,8 @@
     protected final PermissionMonitor mPermissionMonitor;
 
     private final PerUidCounter mNetworkRequestCounter;
-    private final PerUidCounter mSystemNetworkRequestCounter;
+    @VisibleForTesting
+    final PerUidCounter mSystemNetworkRequestCounter;
 
     private volatile boolean mLockdownEnabled;
 
@@ -1060,8 +1062,9 @@
         private final int mMaxCountPerUid;
 
         // Map from UID to number of NetworkRequests that UID has filed.
+        @VisibleForTesting
         @GuardedBy("mUidToNetworkRequestCount")
-        private final SparseIntArray mUidToNetworkRequestCount = new SparseIntArray();
+        final SparseIntArray mUidToNetworkRequestCount = new SparseIntArray();
 
         /**
          * Constructor
@@ -1085,15 +1088,20 @@
          */
         public void incrementCountOrThrow(final int uid) {
             synchronized (mUidToNetworkRequestCount) {
-                final int networkRequests = mUidToNetworkRequestCount.get(uid, 0) + 1;
-                if (networkRequests >= mMaxCountPerUid) {
-                    throw new ServiceSpecificException(
-                            ConnectivityManager.Errors.TOO_MANY_REQUESTS);
-                }
-                mUidToNetworkRequestCount.put(uid, networkRequests);
+                incrementCountOrThrow(uid, 1 /* numToIncrement */);
             }
         }
 
+        private void incrementCountOrThrow(final int uid, final int numToIncrement) {
+            final int newRequestCount =
+                    mUidToNetworkRequestCount.get(uid, 0) + numToIncrement;
+            if (newRequestCount >= mMaxCountPerUid) {
+                throw new ServiceSpecificException(
+                        ConnectivityManager.Errors.TOO_MANY_REQUESTS);
+            }
+            mUidToNetworkRequestCount.put(uid, newRequestCount);
+        }
+
         /**
          * Decrements the request count of the given uid.
          *
@@ -1101,16 +1109,50 @@
          */
         public void decrementCount(final int uid) {
             synchronized (mUidToNetworkRequestCount) {
-                final int requests = mUidToNetworkRequestCount.get(uid, 0);
-                if (requests < 1) {
-                    logwtf("BUG: too small request count " + requests + " for UID " + uid);
-                } else if (requests == 1) {
-                    mUidToNetworkRequestCount.delete(uid);
-                } else {
-                    mUidToNetworkRequestCount.put(uid, requests - 1);
-                }
+                decrementCount(uid, 1 /* numToDecrement */);
             }
         }
+
+        private void decrementCount(final int uid, final int numToDecrement) {
+            final int newRequestCount =
+                    mUidToNetworkRequestCount.get(uid, 0) - numToDecrement;
+            if (newRequestCount < 0) {
+                logwtf("BUG: too small request count " + newRequestCount + " for UID " + uid);
+            } else if (newRequestCount == 0) {
+                mUidToNetworkRequestCount.delete(uid);
+            } else {
+                mUidToNetworkRequestCount.put(uid, newRequestCount);
+            }
+        }
+
+        /**
+         * Used to adjust the request counter for the per-app API flows. Directly adjusting the
+         * counter is not ideal however in the per-app flows, the nris can't be removed until they
+         * are used to create the new nris upon set. Therefore the request count limit can be
+         * artificially hit. This method is used as a workaround for this particular case so that
+         * the request counts are accounted for correctly.
+         * @param uid the uid to adjust counts for
+         * @param numOfNewRequests the new request count to account for
+         * @param r the runnable to execute
+         */
+        public void transact(final int uid, final int numOfNewRequests, @NonNull final Runnable r) {
+            // This should only be used on the handler thread as per all current and foreseen
+            // use-cases. ensureRunningOnConnectivityServiceThread() can't be used because there is
+            // no ref to the outer ConnectivityService.
+            synchronized (mUidToNetworkRequestCount) {
+                final int reqCountOverage = getCallingUidRequestCountOverage(uid, numOfNewRequests);
+                decrementCount(uid, reqCountOverage);
+                r.run();
+                incrementCountOrThrow(uid, reqCountOverage);
+            }
+        }
+
+        private int getCallingUidRequestCountOverage(final int uid, final int numOfNewRequests) {
+            final int newUidRequestCount = mUidToNetworkRequestCount.get(uid, 0)
+                    + numOfNewRequests;
+            return newUidRequestCount >= MAX_NETWORK_REQUESTS_PER_SYSTEM_UID
+                    ? newUidRequestCount - (MAX_NETWORK_REQUESTS_PER_SYSTEM_UID - 1) : 0;
+        }
     }
 
     /**
@@ -9680,9 +9722,13 @@
         validateNetworkCapabilitiesOfProfileNetworkPreference(preference.capabilities);
 
         mProfileNetworkPreferences = mProfileNetworkPreferences.plus(preference);
-        final ArraySet<NetworkRequestInfo> nris =
-                createNrisFromProfileNetworkPreferences(mProfileNetworkPreferences);
-        replaceDefaultNetworkRequestsForPreference(nris);
+        mSystemNetworkRequestCounter.transact(
+                mDeps.getCallingUid(), mProfileNetworkPreferences.preferences.size(),
+                () -> {
+                    final ArraySet<NetworkRequestInfo> nris =
+                            createNrisFromProfileNetworkPreferences(mProfileNetworkPreferences);
+                    replaceDefaultNetworkRequestsForPreference(nris);
+                });
         // Finally, rematch.
         rematchAllNetworksAndRequests();
 
@@ -9768,9 +9814,16 @@
         }
 
         mOemNetworkPreferencesLogs.log("UPDATE INITIATED: " + preference);
-        final ArraySet<NetworkRequestInfo> nris =
-                new OemNetworkRequestFactory().createNrisFromOemNetworkPreferences(preference);
-        replaceDefaultNetworkRequestsForPreference(nris);
+        final int uniquePreferenceCount = new ArraySet<>(
+                preference.getNetworkPreferences().values()).size();
+        mSystemNetworkRequestCounter.transact(
+                mDeps.getCallingUid(), uniquePreferenceCount,
+                () -> {
+                    final ArraySet<NetworkRequestInfo> nris =
+                            new OemNetworkRequestFactory()
+                                    .createNrisFromOemNetworkPreferences(preference);
+                    replaceDefaultNetworkRequestsForPreference(nris);
+                });
         mOemNetworkPreferences = preference;
 
         if (null != listener) {
@@ -9795,10 +9848,14 @@
         final ArraySet<NetworkRequestInfo> perAppCallbackRequestsToUpdate =
                 getPerAppCallbackRequestsToUpdate();
         final ArraySet<NetworkRequestInfo> nrisToRegister = new ArraySet<>(nris);
-        nrisToRegister.addAll(
-                createPerAppCallbackRequestsToRegister(perAppCallbackRequestsToUpdate));
-        handleRemoveNetworkRequests(perAppCallbackRequestsToUpdate);
-        handleRegisterNetworkRequests(nrisToRegister);
+        mSystemNetworkRequestCounter.transact(
+                mDeps.getCallingUid(), perAppCallbackRequestsToUpdate.size(),
+                () -> {
+                    nrisToRegister.addAll(
+                            createPerAppCallbackRequestsToRegister(perAppCallbackRequestsToUpdate));
+                    handleRemoveNetworkRequests(perAppCallbackRequestsToUpdate);
+                    handleRegisterNetworkRequests(nrisToRegister);
+                });
     }
 
     /**
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 7ebed39..23e782f 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -10409,12 +10409,15 @@
         return UidRange.createForUser(UserHandle.of(userId));
     }
 
-    private void mockGetApplicationInfo(@NonNull final String packageName, @NonNull final int uid)
-            throws Exception {
+    private void mockGetApplicationInfo(@NonNull final String packageName, @NonNull final int uid) {
         final ApplicationInfo applicationInfo = new ApplicationInfo();
         applicationInfo.uid = uid;
-        when(mPackageManager.getApplicationInfo(eq(packageName), anyInt()))
-                .thenReturn(applicationInfo);
+        try {
+            when(mPackageManager.getApplicationInfo(eq(packageName), anyInt()))
+                    .thenReturn(applicationInfo);
+        } catch (Exception e) {
+            fail(e.getMessage());
+        }
     }
 
     private void mockGetApplicationInfoThrowsNameNotFound(@NonNull final String packageName)
@@ -10435,8 +10438,7 @@
     }
 
     private OemNetworkPreferences createDefaultOemNetworkPreferences(
-            @OemNetworkPreferences.OemNetworkPreference final int preference)
-            throws Exception {
+            @OemNetworkPreferences.OemNetworkPreference final int preference) {
         // Arrange PackageManager mocks
         mockGetApplicationInfo(TEST_PACKAGE_NAME, TEST_PACKAGE_UID);
 
@@ -10913,11 +10915,13 @@
             mDone.complete(new Object());
         }
 
-        void expectOnComplete() throws Exception {
+        void expectOnComplete() {
             try {
                 mDone.get(TIMEOUT_MS, TimeUnit.MILLISECONDS);
             } catch (TimeoutException e) {
                 fail("Expected onComplete() not received after " + TIMEOUT_MS + " ms");
+            } catch (Exception e) {
+                fail(e.getMessage());
             }
         }
 
@@ -12648,4 +12652,72 @@
                 expected,
                 () -> mCm.registerNetworkCallback(getRequestWithSubIds(), new NetworkCallback()));
     }
+
+    /**
+     * Validate request counts are counted accurately on setProfileNetworkPreference on set/replace.
+     */
+    @Test
+    public void testProfileNetworkPrefCountsRequestsCorrectlyOnSet() throws Exception {
+        final UserHandle testHandle = setupEnterpriseNetwork();
+        testRequestCountLimits(() -> {
+            // Set initially to test the limit prior to having existing requests.
+            final TestOnCompleteListener listener = new TestOnCompleteListener();
+            mCm.setProfileNetworkPreference(testHandle, PROFILE_NETWORK_PREFERENCE_ENTERPRISE,
+                    Runnable::run, listener);
+            listener.expectOnComplete();
+
+            // re-set so as to test the limit as part of replacing existing requests.
+            mCm.setProfileNetworkPreference(testHandle, PROFILE_NETWORK_PREFERENCE_ENTERPRISE,
+                    Runnable::run, listener);
+            listener.expectOnComplete();
+        });
+    }
+
+    /**
+     * Validate request counts are counted accurately on setOemNetworkPreference on set/replace.
+     */
+    @Test
+    public void testSetOemNetworkPreferenceCountsRequestsCorrectlyOnSet() throws Exception {
+        mockHasSystemFeature(PackageManager.FEATURE_AUTOMOTIVE, true);
+        @OemNetworkPreferences.OemNetworkPreference final int networkPref =
+                OEM_NETWORK_PREFERENCE_OEM_PRIVATE_ONLY;
+        testRequestCountLimits(() -> {
+            // Set initially to test the limit prior to having existing requests.
+            final TestOemListenerCallback listener = new TestOemListenerCallback();
+            mService.setOemNetworkPreference(
+                    createDefaultOemNetworkPreferences(networkPref), listener);
+            listener.expectOnComplete();
+
+            // re-set so as to test the limit as part of replacing existing requests.
+            mService.setOemNetworkPreference(
+                    createDefaultOemNetworkPreferences(networkPref), listener);
+            listener.expectOnComplete();
+        });
+    }
+
+    private void testRequestCountLimits(@NonNull final Runnable r) throws Exception {
+        final ArraySet<TestNetworkCallback> callbacks = new ArraySet<>();
+        try {
+            final int requestCount = mService.mSystemNetworkRequestCounter
+                    .mUidToNetworkRequestCount.get(Process.myUid());
+            // The limit is hit when total requests <= limit.
+            final int maxCount =
+                    ConnectivityService.MAX_NETWORK_REQUESTS_PER_SYSTEM_UID - requestCount;
+            // Need permission so registerDefaultNetworkCallback uses mSystemNetworkRequestCounter
+            withPermission(NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK, () -> {
+                for (int i = 1; i < maxCount - 1; i++) {
+                    final TestNetworkCallback cb = new TestNetworkCallback();
+                    mCm.registerDefaultNetworkCallback(cb);
+                    callbacks.add(cb);
+                }
+
+                // Code to run to check if it triggers a max request count limit error.
+                r.run();
+            });
+        } finally {
+            for (final TestNetworkCallback cb : callbacks) {
+                mCm.unregisterNetworkCallback(cb);
+            }
+        }
+    }
 }