Add permission control methods override

Bug: 325984919
Test: atest FrameworksNetTests
Change-Id: Ib57a4a936538295a6df10ae7e334bc969ea455d1
diff --git a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
index e401434..ea9ba75 100644
--- a/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
+++ b/tests/unit/java/com/android/server/connectivityservice/base/CSTest.kt
@@ -47,8 +47,10 @@
 import android.os.Bundle
 import android.os.Handler
 import android.os.HandlerThread
+import android.os.Process
 import android.os.UserHandle
 import android.os.UserManager
+import android.permission.PermissionManager.PermissionResult
 import android.telephony.TelephonyManager
 import android.testing.TestableContext
 import android.util.ArraySet
@@ -300,13 +302,64 @@
         val pacProxyManager = mock<PacProxyManager>()
         val networkPolicyManager = mock<NetworkPolicyManager>()
 
+        // Map of permission name -> PermissionManager.Permission_{GRANTED|DENIED} constant
+        // For permissions granted across the board, the key is only the permission name.
+        // For permissions only granted to a combination of uid/pid, the key
+        // is "<permission name>,<pid>,<uid>". PID+UID permissions have priority over generic ones.
+        private val mMockedPermissions: HashMap<String, Int> = HashMap()
         override fun getPackageManager() = this@CSTest.packageManager
         override fun getContentResolver() = this@CSTest.contentResolver
 
-        // TODO : buff up the capabilities of this permission scheme to allow checking for
-        // permission rejections
-        override fun checkPermission(permission: String, pid: Int, uid: Int) = PERMISSION_GRANTED
-        override fun checkCallingOrSelfPermission(permission: String) = PERMISSION_GRANTED
+        // If the permission result does not set in the mMockedPermissions, it will be
+        // considered as PERMISSION_GRANTED as existing design to prevent breaking other tests.
+        override fun checkPermission(permission: String, pid: Int, uid: Int) =
+            checkMockedPermission(permission, pid, uid, PERMISSION_GRANTED)
+
+        override fun enforceCallingOrSelfPermission(permission: String, message: String?) {
+            // If the permission result does not set in the mMockedPermissions, it will be
+            // considered as PERMISSION_GRANTED as existing design to prevent breaking other tests.
+            val granted = checkMockedPermission(permission, Process.myPid(), Process.myUid(),
+                PERMISSION_GRANTED)
+            if (!granted.equals(PERMISSION_GRANTED)) {
+                throw SecurityException("[Test] permission denied: " + permission)
+            }
+        }
+
+        // If the permission result does not set in the mMockedPermissions, it will be
+        // considered as PERMISSION_GRANTED as existing design to prevent breaking other tests.
+        override fun checkCallingOrSelfPermission(permission: String) =
+            checkMockedPermission(permission, Process.myPid(), Process.myUid(), PERMISSION_GRANTED)
+
+        private fun checkMockedPermission(permission: String, pid: Int, uid: Int, default: Int):
+                Int {
+            val processSpecificKey = "$permission,$pid,$uid"
+            return mMockedPermissions[processSpecificKey]
+                    ?: mMockedPermissions[permission] ?: default
+        }
+
+        /**
+         * Mock checks for the specified permission, and have them behave as per `granted` or
+         * `denied`.
+         *
+         * This will apply to all calls no matter what the checked UID and PID are.
+         *
+         * @param granted One of {@link PackageManager#PermissionResult}.
+         */
+        fun setPermission(permission: String, @PermissionResult granted: Int) {
+            mMockedPermissions.put(permission, granted)
+        }
+
+        /**
+         * Mock checks for the specified permission, and have them behave as per `granted` or
+         * `denied`.
+         *
+         * This will only apply to the passed UID and PID.
+         *
+         * @param granted One of {@link PackageManager#PermissionResult}.
+         */
+        fun setPermission(permission: String, pid: Int, uid: Int, @PermissionResult granted: Int) {
+            mMockedPermissions.put("$permission,$pid,$uid", granted)
+        }
 
         // Necessary for MultinetworkPolicyTracker, which tries to register a receiver for
         // all users. The test can't do that since it doesn't hold INTERACT_ACROSS_USERS.