1 /*
2  * Copyright (C) 2022 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 com.android.net.module.util;
18 
19 import android.util.SparseIntArray;
20 
21 import com.android.internal.annotations.GuardedBy;
22 import com.android.internal.annotations.VisibleForTesting;
23 
24 /**
25  * Keeps track of the counters under different uid, fire exception if the counter
26  * exceeded the specified maximum value.
27  *
28  * @hide
29  */
30 public class PerUidCounter {
31     private final int mMaxCountPerUid;
32 
33     // Map from UID to count that UID has filed.
34     @VisibleForTesting
35     @GuardedBy("this")
36     final SparseIntArray mUidToCount = new SparseIntArray();
37 
38     /**
39      * Constructor
40      *
41      * @param maxCountPerUid the maximum count per uid allowed
42      */
PerUidCounter(final int maxCountPerUid)43     public PerUidCounter(final int maxCountPerUid) {
44         if (maxCountPerUid <= 0) {
45             throw new IllegalArgumentException("Maximum counter value must be positive");
46         }
47         mMaxCountPerUid = maxCountPerUid;
48     }
49 
50     /**
51      * Increments the count of the given uid.  Throws an exception if the number
52      * of the counter for the uid exceeds the value of maxCounterPerUid which is the value
53      * passed into the constructor. see: {@link #PerUidCounter(int)}.
54      *
55      * @throws IllegalStateException if the number of counter for the uid exceed
56      *         the allowed number.
57      *
58      * @param uid the uid that the counter was made under
59      */
incrementCountOrThrow(final int uid)60     public synchronized void incrementCountOrThrow(final int uid) {
61         final long newCount = ((long) mUidToCount.get(uid, 0)) + 1;
62         if (newCount > mMaxCountPerUid) {
63             throw new IllegalStateException("Uid " + uid + " exceeded its allowed limit");
64         }
65         // Since the count cannot be greater than Integer.MAX_VALUE here since mMaxCountPerUid
66         // is an integer, it is safe to cast to int.
67         mUidToCount.put(uid, (int) newCount);
68     }
69 
70     /**
71      * Decrements the count of the given uid. Throws an exception if the number
72      * of the counter goes below zero.
73      *
74      * @throws IllegalStateException if the number of counter for the uid goes below
75      *         zero.
76      *
77      * @param uid the uid that the count was made under
78      */
decrementCountOrThrow(final int uid)79     public synchronized void decrementCountOrThrow(final int uid) {
80         final int newCount = mUidToCount.get(uid, 0) - 1;
81         if (newCount < 0) {
82             throw new IllegalStateException("BUG: too small count " + newCount + " for UID " + uid);
83         } else if (newCount == 0) {
84             mUidToCount.delete(uid);
85         } else {
86             mUidToCount.put(uid, newCount);
87         }
88     }
89 
90     @VisibleForTesting
get(int uid)91     public synchronized int get(int uid) {
92         return mUidToCount.get(uid, 0);
93     }
94 }
95