Add PerUidCounter
The class keeps track of the counters under different uid,
fire exception if the counter exceeded the specified maximum
value.
This is reimplemented and generalized based on the one
inside ConnectivityService.
Test: atest NetworkStaticLibTests:com.android.net.moduletests.util.PerUidCounterTest
Test: atest ConnectivityCoverageTests:android.net.connectivity.com.android.net.module.util.PerUidCounterTest
Bug: 229103088
Change-Id: I7dfb16342e3ca4eab45bb40f2e1355981e04b44b
diff --git a/staticlibs/framework/com/android/net/module/util/PerUidCounter.java b/staticlibs/framework/com/android/net/module/util/PerUidCounter.java
new file mode 100644
index 0000000..7e0526d
--- /dev/null
+++ b/staticlibs/framework/com/android/net/module/util/PerUidCounter.java
@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.net.module.util;
+
+import android.util.SparseIntArray;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
+
+/**
+ * Keeps track of the counters under different uid, fire exception if the counter
+ * exceeded the specified maximum value.
+ *
+ * @hide
+ */
+public class PerUidCounter {
+ private final int mMaxCountPerUid;
+
+ // Map from UID to count that UID has filed.
+ @VisibleForTesting
+ @GuardedBy("mUidToCount")
+ final SparseIntArray mUidToCount = new SparseIntArray();
+
+ /**
+ * Constructor
+ *
+ * @param maxCountPerUid the maximum count per uid allowed
+ */
+ public PerUidCounter(final int maxCountPerUid) {
+ if (maxCountPerUid < 0) {
+ throw new IllegalArgumentException("Maximum counter value cannot be negative");
+ }
+ mMaxCountPerUid = maxCountPerUid;
+ }
+
+ /**
+ * Increments the count of the given uid. Throws an exception if the number
+ * of the counter for the uid exceeds the value of maxCounterPerUid which is the value
+ * passed into the constructor. see: {@link #PerUidCounter(int)}.
+ *
+ * @throws IllegalStateException if the number of counter for the uid exceed
+ * the allowed number.
+ *
+ * @param uid the uid that the counter was made under
+ */
+ public void incrementCountOrThrow(final int uid) {
+ synchronized (mUidToCount) {
+ incrementCountOrThrow(uid, 1 /* numToIncrement */);
+ }
+ }
+
+ public void incrementCountOrThrow(final int uid, final int numToIncrement) {
+ if (numToIncrement <= 0) {
+ throw new IllegalArgumentException("Increment count must be positive");
+ }
+ final long newCount = ((long) mUidToCount.get(uid, 0)) + numToIncrement;
+ if (newCount > mMaxCountPerUid) {
+ throw new IllegalStateException("Uid " + uid + " exceeded its allowed limit");
+ }
+ // Since the count cannot be greater than Integer.MAX_VALUE here,
+ // it is safe to cast to int.
+ mUidToCount.put(uid, (int) newCount);
+ }
+
+ /**
+ * Decrements the count of the given uid. Throws an exception if the number
+ * of the counter goes below zero.
+ *
+ * @throws IllegalStateException if the number of counter for the uid goes below
+ * zero.
+ *
+ * @param uid the uid that the count was made under
+ */
+ public void decrementCountOrThrow(final int uid) {
+ synchronized (mUidToCount) {
+ decrementCountOrThrow(uid, 1 /* numToDecrement */);
+ }
+ }
+
+ public void decrementCountOrThrow(final int uid, final int numToDecrement) {
+ if (numToDecrement <= 0) {
+ throw new IllegalArgumentException("Decrement count must be positive");
+ }
+ final int newCount = mUidToCount.get(uid, 0) - numToDecrement;
+ if (newCount < 0) {
+ throw new IllegalStateException("BUG: too small count " + newCount + " for UID " + uid);
+ } else if (newCount == 0) {
+ mUidToCount.delete(uid);
+ } else {
+ mUidToCount.put(uid, newCount);
+ }
+ }
+}
diff --git a/staticlibs/tests/unit/src/com/android/net/module/util/PerUidCounterTest.kt b/staticlibs/tests/unit/src/com/android/net/module/util/PerUidCounterTest.kt
new file mode 100644
index 0000000..c479d81
--- /dev/null
+++ b/staticlibs/tests/unit/src/com/android/net/module/util/PerUidCounterTest.kt
@@ -0,0 +1,114 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.net.module.util
+
+import androidx.test.filters.SmallTest
+import androidx.test.runner.AndroidJUnit4
+import org.junit.Test
+import org.junit.runner.RunWith
+import kotlin.test.assertFailsWith
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class PerUidCounterTest {
+ private val UID_A = 1000
+ private val UID_B = 1001
+
+ @Test
+ fun testCounterMaximum() {
+ assertFailsWith<IllegalArgumentException> {
+ PerUidCounter(-1)
+ }
+
+ val uselessCounter = PerUidCounter(0)
+ assertFailsWith<IllegalStateException> {
+ uselessCounter.incrementCountOrThrow(UID_A)
+ }
+ assertFailsWith<IllegalStateException> {
+ uselessCounter.decrementCountOrThrow(UID_A)
+ }
+
+ val largeMaxCounter = PerUidCounter(Integer.MAX_VALUE)
+ largeMaxCounter.incrementCountOrThrow(UID_A, Integer.MAX_VALUE)
+ assertFailsWith<IllegalStateException> {
+ largeMaxCounter.incrementCountOrThrow(UID_A)
+ }
+ }
+
+ @Test
+ fun testIncrementCountOrThrow() {
+ val counter = PerUidCounter(3)
+
+ // Verify the increment count cannot be zero.
+ assertFailsWith<IllegalArgumentException> {
+ counter.incrementCountOrThrow(UID_A, 0)
+ }
+
+ // Verify the counters work independently.
+ counter.incrementCountOrThrow(UID_A)
+ counter.incrementCountOrThrow(UID_B, 2)
+ counter.incrementCountOrThrow(UID_B)
+ counter.incrementCountOrThrow(UID_A)
+ counter.incrementCountOrThrow(UID_A)
+ assertFailsWith<IllegalStateException> {
+ counter.incrementCountOrThrow(UID_A)
+ }
+ assertFailsWith<IllegalStateException> {
+ counter.incrementCountOrThrow(UID_B)
+ }
+
+ // Verify exception can be triggered again.
+ assertFailsWith<IllegalStateException> {
+ counter.incrementCountOrThrow(UID_A)
+ }
+ assertFailsWith<IllegalStateException> {
+ counter.incrementCountOrThrow(UID_A, 3)
+ }
+ }
+
+ @Test
+ fun testDecrementCountOrThrow() {
+ val counter = PerUidCounter(3)
+
+ // Verify the decrement count cannot be zero.
+ assertFailsWith<IllegalArgumentException> {
+ counter.decrementCountOrThrow(UID_A, 0)
+ }
+
+ // Verify the count cannot go below zero.
+ assertFailsWith<IllegalStateException> {
+ counter.decrementCountOrThrow(UID_A)
+ }
+ assertFailsWith<IllegalStateException> {
+ counter.decrementCountOrThrow(UID_A, 5)
+ }
+ assertFailsWith<IllegalStateException> {
+ counter.decrementCountOrThrow(UID_A, Integer.MAX_VALUE)
+ }
+
+ // Verify the counters work independently.
+ counter.incrementCountOrThrow(UID_A)
+ counter.incrementCountOrThrow(UID_B)
+ assertFailsWith<IllegalStateException> {
+ counter.decrementCountOrThrow(UID_A, 3)
+ }
+ counter.decrementCountOrThrow(UID_A)
+ assertFailsWith<IllegalStateException> {
+ counter.decrementCountOrThrow(UID_A)
+ }
+ }
+}
\ No newline at end of file