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.safetycenter.testing
18 
19 import android.app.UiAutomation
20 import android.app.UiAutomation.ALL_PERMISSIONS
21 import androidx.annotation.GuardedBy
22 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
23 import com.android.compatibility.common.util.SystemUtil
24 import com.google.common.collect.HashMultiset
25 import com.google.common.collect.Multiset
26 
27 /** A class to facilitate working with System Shell permissions. */
28 object ShellPermissions {
29 
30     private val lock = Object()
31     @GuardedBy("lock") private val nestedPermissions: Multiset<String> = HashMultiset.create()
32 
33     /**
34      * Behaves the same way as [SystemUtil.callWithShellPermissionIdentity], but allows nesting
35      * calls to it by merging the adopted [permissions] within each [block].
36      *
37      * Note that [SystemUtil.callWithShellPermissionIdentity] should NOT be used together with this
38      * method.
39      */
callWithShellPermissionIdentitynull40     fun <T> callWithShellPermissionIdentity(vararg permissions: String, block: () -> T): T {
41         val uiAutomation = getInstrumentation().getUiAutomation()
42         val permissionsToAddForThisBlock =
43             if (permissions.isEmpty()) {
44                 ALL_PERMISSIONS
45             } else {
46                 permissions.toSet()
47             }
48         synchronized(lock) {
49             permissionsToAddForThisBlock.forEach { nestedPermissions.add(it) }
50             uiAutomation.adoptShellPermissionIdentityFor(nestedPermissions.elementSet())
51         }
52         try {
53             return block()
54         } finally {
55             synchronized(lock) {
56                 permissionsToAddForThisBlock.forEach { nestedPermissions.remove(it) }
57                 if (nestedPermissions.isEmpty()) {
58                     uiAutomation.dropShellPermissionIdentity()
59                 } else {
60                     uiAutomation.adoptShellPermissionIdentityFor(nestedPermissions.elementSet())
61                 }
62             }
63         }
64     }
65 
UiAutomationnull66     private fun UiAutomation.adoptShellPermissionIdentityFor(permissionsToAdopt: Set<String>) {
67         if (permissionsToAdopt.containsAll(ALL_PERMISSIONS)) {
68             adoptShellPermissionIdentity()
69         } else {
70             adoptShellPermissionIdentity(*permissionsToAdopt.toTypedArray())
71         }
72     }
73 }
74