• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright (C) 2023 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.bedstead.harrier;
18 
19 import static com.android.bedstead.harrier.UserType.INITIAL_USER;
20 import static com.android.bedstead.harrier.UserType.SECONDARY_USER;
21 import static com.android.bedstead.harrier.UserType.SYSTEM_USER;
22 import static com.android.bedstead.harrier.UserType.WORK_PROFILE;
23 import static com.android.bedstead.harrier.annotations.enterprise.EnsureHasDelegate.DELEGATE_KEY;
24 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_AFFILIATED_PROFILE_OWNER;
25 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_AFFILIATED_PROFILE_OWNER_PROFILE;
26 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_AFFILIATED_PROFILE_OWNER_USER;
27 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_DEVICE_OWNER;
28 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_DPM_ROLE_HOLDER;
29 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_FINANCED_DEVICE_OWNER;
30 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE;
31 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_PARENT_INSTANCE_OF_NON_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE;
32 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_PARENT_INSTANCE_OF_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE;
33 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_PROFILE_OWNER_USER_WITH_NO_DO;
34 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE;
35 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_USER;
36 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIES_IN_BACKGROUND;
37 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIES_TO_AFFILIATED_OTHER_USERS;
38 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIES_TO_OWN_USER;
39 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIES_TO_PARENT;
40 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIES_TO_UNAFFILIATED_CHILD_PROFILES_WITHOUT_INHERITANCE;
41 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIES_TO_UNAFFILIATED_OTHER_USERS;
42 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.CANNOT_BE_APPLIED_BY_ROLE_HOLDER;
43 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.CAN_BE_DELEGATED;
44 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.DO_NOT_APPLY_TO_CANNOT_SET_POLICY_TESTS;
45 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.DO_NOT_APPLY_TO_POLICY_DOES_NOT_APPLY_TESTS;
46 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.INHERITABLE;
47 import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.NO;
48 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_APP_RESTRICTIONS;
49 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_BLOCK_UNINSTALL;
50 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_CERT_INSTALL;
51 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_CERT_SELECTION;
52 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_ENABLE_SYSTEM_APP;
53 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_INSTALL_EXISTING_PACKAGE;
54 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_KEEP_UNINSTALLED_PACKAGES;
55 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_NETWORK_LOGGING;
56 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_PACKAGE_ACCESS;
57 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_PERMISSION_GRANT;
58 import static com.android.bedstead.nene.devicepolicy.CommonDevicePolicy.DELEGATION_SECURITY_LOGGING;
59 import static com.android.bedstead.testapp.TestAppQueryBuilder.queryBuilder;
60 
61 import com.android.bedstead.harrier.annotations.EnsureFeatureFlagEnabled;
62 import com.android.bedstead.harrier.annotations.EnsureTestAppDoesNotHavePermission;
63 import com.android.bedstead.harrier.annotations.EnsureTestAppHasAppOp;
64 import com.android.bedstead.harrier.annotations.EnsureTestAppHasPermission;
65 import com.android.bedstead.harrier.annotations.EnsureTestAppInstalled;
66 import com.android.bedstead.harrier.annotations.FailureMode;
67 import com.android.bedstead.harrier.annotations.RequireFeatureFlagEnabled;
68 import com.android.bedstead.harrier.annotations.enterprise.EnsureHasDelegate;
69 import com.android.bedstead.harrier.annotations.enterprise.EnsureHasDevicePolicyManagerRoleHolder;
70 import com.android.bedstead.harrier.annotations.enterprise.EnsureHasNoDelegate;
71 import com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy;
72 import com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.AppOp;
73 import com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.Permission;
74 import com.android.bedstead.harrier.annotations.meta.ParameterizedAnnotation;
75 import com.android.bedstead.harrier.annotations.parameterized.IncludeNone;
76 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnAffiliatedDeviceOwnerSecondaryUser;
77 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnAffiliatedProfileOwnerSecondaryUser;
78 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnBackgroundDeviceOwnerUser;
79 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnCloneProfileAlongsideManagedProfile;
80 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnCloneProfileAlongsideManagedProfileUsingParentInstance;
81 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnCloneProfileAlongsideOrganizationOwnedProfile;
82 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnCloneProfileAlongsideOrganizationOwnedProfileUsingParentInstance;
83 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnDeviceOwnerUser;
84 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnDevicePolicyManagementRoleHolderProfile;
85 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnDevicePolicyManagementRoleHolderSecondaryUser;
86 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnDevicePolicyManagementRoleHolderUser;
87 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnFinancedDeviceOwnerUser;
88 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnOrganizationOwnedProfileOwner;
89 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnParentOfOrganizationOwnedProfileOwner;
90 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnParentOfOrganizationOwnedProfileOwnerUsingParentInstance;
91 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnParentOfProfileOwnerUsingParentInstance;
92 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnParentOfProfileOwnerWithNoDeviceOwner;
93 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnProfileOwnerPrimaryUser;
94 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnProfileOwnerProfileWithNoDeviceOwner;
95 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnSecondaryUserInDifferentProfileGroupToOrganizationOwnedProfileOwnerProfileUsingParentInstance;
96 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnSecondaryUserInDifferentProfileGroupToProfileOwnerProfile;
97 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnUnaffiliatedDeviceOwnerSecondaryUser;
98 import com.android.bedstead.harrier.annotations.parameterized.IncludeRunOnUnaffiliatedProfileOwnerSecondaryUser;
99 import com.android.queryable.annotations.Query;
100 
101 import com.google.auto.value.AutoAnnotation;
102 import com.google.common.collect.ImmutableMap;
103 import com.google.common.collect.ImmutableSet;
104 
105 import java.lang.annotation.Annotation;
106 import java.util.ArrayList;
107 import java.util.Arrays;
108 import java.util.Comparator;
109 import java.util.HashMap;
110 import java.util.HashSet;
111 import java.util.List;
112 import java.util.Map;
113 import java.util.Set;
114 import java.util.function.BiFunction;
115 import java.util.function.Function;
116 import java.util.stream.Collectors;
117 
118 /**
119  * Utility class for enterprise policy tests.
120  */
121 public final class Policy {
122 
123     private static final String DELEGATE_PACKAGE_NAME = "com.android.Delegate";
124 
125     // Delegate scopes to be used for a "CannotSet" state. All delegate scopes except the ones which
126     // should allow use of the API will be granted
127     private static final ImmutableSet<String> ALL_DELEGATE_SCOPES = ImmutableSet.of(
128             DELEGATION_CERT_INSTALL,
129             DELEGATION_APP_RESTRICTIONS,
130             DELEGATION_BLOCK_UNINSTALL,
131             DELEGATION_PERMISSION_GRANT,
132             DELEGATION_PACKAGE_ACCESS,
133             DELEGATION_ENABLE_SYSTEM_APP,
134             DELEGATION_INSTALL_EXISTING_PACKAGE,
135             DELEGATION_KEEP_UNINSTALLED_PACKAGES,
136             DELEGATION_NETWORK_LOGGING,
137             DELEGATION_CERT_SELECTION,
138             DELEGATION_SECURITY_LOGGING
139     );
140 
141     // This is a map containing all Include* annotations and the flags which lead to them
142     // This is not validated - every state must have a single APPLIED_BY annotation
143     private static final ImmutableMap<Integer, Function<EnterprisePolicy, Set<Annotation>>>
144             STATE_ANNOTATIONS =
145             ImmutableMap.<Integer, Function<EnterprisePolicy, Set<Annotation>>>builder()
146                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnDeviceOwnerUser(), /* roleHolderUser= */ SYSTEM_USER))
147                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_OWN_USER | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnDeviceOwnerUser(), /* isPrimary= */ true))
148                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_OWN_USER | APPLIES_IN_BACKGROUND, singleAnnotation(includeRunOnBackgroundDeviceOwnerUser()))
149                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_OWN_USER | APPLIES_IN_BACKGROUND | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnBackgroundDeviceOwnerUser(), /* isPrimary= */ true))
150 
151                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_UNAFFILIATED_OTHER_USERS, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnNonAffiliatedDeviceOwnerSecondaryUser(), /* roleHolderUser= */ SYSTEM_USER))
152                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_UNAFFILIATED_OTHER_USERS | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnNonAffiliatedDeviceOwnerSecondaryUser(), /* isPrimary= */ true))
153                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_AFFILIATED_OTHER_USERS, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnAffiliatedDeviceOwnerSecondaryUser(), /* roleHolderUser= */ SYSTEM_USER))
154                     .put(APPLIED_BY_DEVICE_OWNER | APPLIES_TO_AFFILIATED_OTHER_USERS | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnAffiliatedDeviceOwnerSecondaryUser(), /* isPrimary= */ true))
155 
156                     .put(APPLIED_BY_DPM_ROLE_HOLDER | APPLIES_TO_OWN_USER, singleAnnotation(includeRunOnDevicePolicyManagementRoleHolderUser()))
157                     .put(APPLIED_BY_DPM_ROLE_HOLDER | APPLIES_TO_UNAFFILIATED_OTHER_USERS, singleAnnotation(includeRunOnDevicePolicyManagementRoleHolderSecondaryUser()))
158                     .put(APPLIED_BY_DPM_ROLE_HOLDER | APPLIES_TO_UNAFFILIATED_CHILD_PROFILES_WITHOUT_INHERITANCE, singleAnnotation(includeRunOnDevicePolicyManagementRoleHolderProfile()))
159 
160                     .put(APPLIED_BY_AFFILIATED_PROFILE_OWNER_USER | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnAffiliatedProfileOwnerSecondaryUser(), /* roleHolderUser= */ SECONDARY_USER))
161                     .put(APPLIED_BY_AFFILIATED_PROFILE_OWNER_USER | APPLIES_TO_OWN_USER | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnAffiliatedProfileOwnerSecondaryUser(), /* isPrimary= */ true))
162                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_USER | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnUnaffiliatedProfileOwnerSecondaryUser(), /* roleHolderUser= */ SECONDARY_USER))
163                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_USER | APPLIES_TO_OWN_USER | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnUnaffiliatedProfileOwnerSecondaryUser(), /* isPrimary= */ true))
164                     .put(APPLIED_BY_PROFILE_OWNER_USER_WITH_NO_DO | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnProfileOwnerPrimaryUser(), /* roleHolderUser= */ SYSTEM_USER))
165                     .put(APPLIED_BY_PROFILE_OWNER_USER_WITH_NO_DO | APPLIES_TO_OWN_USER | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnProfileOwnerPrimaryUser(), /* isPrimary= */ true))
166 
167                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnProfileOwnerProfileWithNoDeviceOwner(), /* roleHolderUser= */ WORK_PROFILE))
168                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE | APPLIES_TO_OWN_USER | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnProfileOwnerProfileWithNoDeviceOwner(), /* isPrimary= */ true))
169                     .put(APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnOrganizationOwnedProfileOwner(), /* roleHolderUser= */ WORK_PROFILE))
170                     .put(APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_OWN_USER | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnOrganizationOwnedProfileOwner(), /* isPrimary= */ true))
171 
172                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE | APPLIES_TO_PARENT, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnParentOfProfileOwnerWithNoDeviceOwner(), /* roleHolderUser= */ WORK_PROFILE))
173                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE | APPLIES_TO_PARENT | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnParentOfProfileOwnerWithNoDeviceOwner(), /* isPrimary= */ true))
174                     .put(APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_PARENT, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnParentOfOrganizationOwnedProfileOwner(), /* roleHolderUser= */ WORK_PROFILE))
175                     .put(APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_PARENT | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnParentOfOrganizationOwnedProfileOwner(), /* isPrimary= */ true))
176 
177                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE | APPLIES_TO_UNAFFILIATED_OTHER_USERS, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnSecondaryUserInDifferentProfileGroupToProfileOwnerProfile(), /* roleHolderUser= */ WORK_PROFILE))
178                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE | APPLIES_TO_UNAFFILIATED_OTHER_USERS | CAN_BE_DELEGATED, generateDelegateAnnotation(includeRunOnSecondaryUserInDifferentProfileGroupToProfileOwnerProfile(), /* isPrimary= */ true))
179 
180                     .put(APPLIED_BY_PARENT_INSTANCE_OF_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_UNAFFILIATED_OTHER_USERS, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnSecondaryUserInDifferentProfileGroupToOrganizationOwnedProfileOwnerProfileUsingParentInstance(), /* roleHolderUser= */ WORK_PROFILE))
181 
182                     // The model here is that APPLIED_BY_PARENT + APPLIES_TO_OWN_USER means it applies to the parent of the DPC - I'm not sure this is the best model (APPLIES_TO_PARENT would also be reasonable)
183                     .put(APPLIED_BY_PARENT_INSTANCE_OF_NON_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE
184                             | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnParentOfProfileOwnerUsingParentInstance(), /* roleHolderUser= */ INITIAL_USER))
185                     .put(APPLIED_BY_PARENT_INSTANCE_OF_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnParentOfOrganizationOwnedProfileOwnerUsingParentInstance(), /* roleHolderUser= */ INITIAL_USER))
186 
187                     .put(APPLIED_BY_FINANCED_DEVICE_OWNER | APPLIES_TO_OWN_USER, generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnFinancedDeviceOwnerUser(), /* roleHolderUser= */ SYSTEM_USER))
188 
189                     .put(APPLIED_BY_PARENT_INSTANCE_OF_NON_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_OWN_USER | INHERITABLE, singleAnnotation(includeRunOnCloneProfileAlongsideManagedProfileUsingParentInstance()))
190                     .put(APPLIED_BY_PARENT_INSTANCE_OF_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_OWN_USER | INHERITABLE, singleAnnotation(includeRunOnCloneProfileAlongsideOrganizationOwnedProfileUsingParentInstance()))
191                     .put(APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE | APPLIES_TO_PARENT | INHERITABLE, singleAnnotation(includeRunOnCloneProfileAlongsideOrganizationOwnedProfile()))
192                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE | APPLIES_TO_PARENT | INHERITABLE, singleAnnotation(includeRunOnCloneProfileAlongsideManagedProfile()))
193                     .build();
194     // This must contain one key for every APPLIED_BY that is being used, and maps to the
195     // "default" for testing that DPC type
196     // in general this will be a state which runs on the same user as the dpc.
197     // The key is the APPLIED_BY annotation and the value is a function which takes the policy and
198     // a boolean indicating if this is the "can set" state (if it is false - it must be the
199     // "cannot set" state). It should return a set of annotations to use.
200     private static final ImmutableMap<Integer, BiFunction<EnterprisePolicy, Boolean, Set<Annotation>>>
201             DPC_STATE_ANNOTATIONS_BASE =
202             ImmutableMap.<Integer, BiFunction<EnterprisePolicy, Boolean, Set<Annotation>>>builder()
203                     .put(APPLIED_BY_DEVICE_OWNER, (flags, canSet) -> hasFlag(flags.dpc(), APPLIED_BY_DEVICE_OWNER | APPLIES_IN_BACKGROUND) ? generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnBackgroundDeviceOwnerUser(), /* roleHolderUser= */ SYSTEM_USER).apply(flags) : generateDevicePolicyManagerRoleHolderAnnotation(includeRunOnDeviceOwnerUser(), /* roleHolderUser= */ SYSTEM_USER).apply(flags))
204                     .put(APPLIED_BY_AFFILIATED_PROFILE_OWNER, devicePolicyManagerRoleHolderIfCanSet(includeRunOnAffiliatedProfileOwnerSecondaryUser(), /* roleHolderUser= */ SECONDARY_USER))
205                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_USER, devicePolicyManagerRoleHolderIfCanSet(includeRunOnProfileOwnerPrimaryUser(), /* roleHolderUser= */ SYSTEM_USER))
206                     .put(APPLIED_BY_PROFILE_OWNER_USER_WITH_NO_DO, devicePolicyManagerRoleHolderIfCanSet(includeRunOnProfileOwnerPrimaryUser(), /* roleHolderUser= */ SYSTEM_USER))
207                     .put(APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE, devicePolicyManagerRoleHolderIfCanSet(includeRunOnOrganizationOwnedProfileOwner(), /* roleHolderUser= */ WORK_PROFILE))
208                     .put(APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE, devicePolicyManagerRoleHolderIfCanSet(includeRunOnProfileOwnerProfileWithNoDeviceOwner(), /* roleHolderUser= */ WORK_PROFILE))
209                     .put(APPLIED_BY_FINANCED_DEVICE_OWNER, devicePolicyManagerRoleHolderIfCanSet(includeRunOnFinancedDeviceOwnerUser(), /* roleHolderUser= */ SYSTEM_USER))
210                     .put(APPLIED_BY_PARENT_INSTANCE_OF_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE, (flags, canSet) -> singleAnnotation(includeRunOnParentOfOrganizationOwnedProfileOwnerUsingParentInstance()).apply(flags))
211                     .put(APPLIED_BY_PARENT_INSTANCE_OF_NON_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE, (flags, canSet) -> singleAnnotation(includeRunOnParentOfProfileOwnerUsingParentInstance()).apply(flags))
212                     .put(APPLIED_BY_DPM_ROLE_HOLDER, (flags, canSet) -> singleAnnotation(includeRunOnDevicePolicyManagementRoleHolderUser()).apply(flags))
213                     .build();
214 
devicePolicyManagerRoleHolderIfCanSet(Annotation annotation, UserType roleHolderUser)215     private static BiFunction<EnterprisePolicy, Boolean, Set<Annotation>> devicePolicyManagerRoleHolderIfCanSet(Annotation annotation, UserType roleHolderUser) {
216         return (flags, canSet) -> {
217             // If the policy already allows the DPM to set it - no need to add the special-cased test
218             if (canSet && (!hasFlag(flags.dpc(), APPLIED_BY_DPM_ROLE_HOLDER))) {
219                 return generateDevicePolicyManagerRoleHolderAnnotation(annotation, roleHolderUser).apply(flags);
220             } else {
221                 return singleAnnotation(annotation).apply(flags);
222             }
223         };
224     }
225     private static final Map<Integer, BiFunction<EnterprisePolicy, Boolean, Set<Annotation>>>
226             DPC_STATE_ANNOTATIONS = DPC_STATE_ANNOTATIONS_BASE.entrySet().stream()
227             .collect(Collectors.toMap(Map.Entry::getKey, Policy::addGeneratedStates));
228     private static final int APPLIED_BY_FLAGS =
229             APPLIED_BY_DEVICE_OWNER | APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_PROFILE
230                     | APPLIED_BY_AFFILIATED_PROFILE_OWNER_PROFILE
231                     | APPLIED_BY_UNAFFILIATED_PROFILE_OWNER_USER
232                     | APPLIED_BY_AFFILIATED_PROFILE_OWNER_USER
233                     | APPLIED_BY_FINANCED_DEVICE_OWNER
234                     | APPLIED_BY_ORGANIZATION_OWNED_PROFILE_OWNER_PROFILE
235                     | APPLIED_BY_PARENT_INSTANCE_OF_NON_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE
236                     | APPLIED_BY_PARENT_INSTANCE_OF_ORGANIZATIONAL_OWNED_PROFILE_OWNER_PROFILE
237                     | APPLIED_BY_DPM_ROLE_HOLDER;
238     private static final Map<Function<EnterprisePolicy, Set<Annotation>>, Set<Integer>>
239             ANNOTATIONS_MAP = calculateAnnotationsMap(STATE_ANNOTATIONS);
240 
241     private Policy() {
242 
243     }
244 
245     @AutoAnnotation
246     public static EnterprisePolicy enterprisePolicy(int[] dpc, Permission[] permissions,
247             AppOp[] appOps, String[] delegatedScopes) {
248         return new AutoAnnotation_Policy_enterprisePolicy(
249                 dpc, permissions, appOps, delegatedScopes);
250     }
251 
252     @AutoAnnotation
253     private static EnsureTestAppInstalled ensureTestAppInstalled(
254             String key, Query query, UserType onUser, boolean isPrimary) {
255         return new AutoAnnotation_Policy_ensureTestAppInstalled(
256                 key, query, onUser, isPrimary);
257     }
258 
259     @AutoAnnotation
260     private static EnsureTestAppHasPermission ensureTestAppHasPermission(
261             String testAppKey, String[] value, FailureMode failureMode) {
262         return new AutoAnnotation_Policy_ensureTestAppHasPermission(testAppKey, value, failureMode);
263     }
264 
265     @AutoAnnotation
266     private static EnsureTestAppDoesNotHavePermission ensureTestAppDoesNotHavePermission(
267             String testAppKey, String[] value, FailureMode failureMode) {
268         return new AutoAnnotation_Policy_ensureTestAppDoesNotHavePermission(
269                 testAppKey, value, failureMode);
270     }
271 
272     @AutoAnnotation
273     private static EnsureTestAppHasAppOp ensureTestAppHasAppOp(String testAppKey, String[] value) {
274         return new AutoAnnotation_Policy_ensureTestAppHasAppOp(testAppKey, value);
275     }
276 
277     @AutoAnnotation
278     private static IncludeNone includeNone() {
279         return new AutoAnnotation_Policy_includeNone();
280     }
281 
282     @AutoAnnotation
283     private static IncludeRunOnDeviceOwnerUser includeRunOnDeviceOwnerUser() {
284         return new AutoAnnotation_Policy_includeRunOnDeviceOwnerUser();
285     }
286 
287     @AutoAnnotation
288     private static IncludeRunOnUnaffiliatedDeviceOwnerSecondaryUser includeRunOnNonAffiliatedDeviceOwnerSecondaryUser() {
289         return new AutoAnnotation_Policy_includeRunOnNonAffiliatedDeviceOwnerSecondaryUser();
290     }
291 
292     @AutoAnnotation
293     private static IncludeRunOnAffiliatedDeviceOwnerSecondaryUser includeRunOnAffiliatedDeviceOwnerSecondaryUser() {
294         return new AutoAnnotation_Policy_includeRunOnAffiliatedDeviceOwnerSecondaryUser();
295     }
296 
297     @AutoAnnotation
298     private static IncludeRunOnAffiliatedProfileOwnerSecondaryUser includeRunOnAffiliatedProfileOwnerSecondaryUser() {
299         return new AutoAnnotation_Policy_includeRunOnAffiliatedProfileOwnerSecondaryUser();
300     }
301 
302     @AutoAnnotation
303     private static IncludeRunOnUnaffiliatedProfileOwnerSecondaryUser includeRunOnUnaffiliatedProfileOwnerSecondaryUser() {
304         return new AutoAnnotation_Policy_includeRunOnUnaffiliatedProfileOwnerSecondaryUser();
305     }
306 
307     @AutoAnnotation
308     private static IncludeRunOnProfileOwnerProfileWithNoDeviceOwner includeRunOnProfileOwnerProfileWithNoDeviceOwner() {
309         return new AutoAnnotation_Policy_includeRunOnProfileOwnerProfileWithNoDeviceOwner();
310     }
311 
312     @AutoAnnotation
313     private static IncludeRunOnSecondaryUserInDifferentProfileGroupToProfileOwnerProfile includeRunOnSecondaryUserInDifferentProfileGroupToProfileOwnerProfile() {
314         return new AutoAnnotation_Policy_includeRunOnSecondaryUserInDifferentProfileGroupToProfileOwnerProfile();
315     }
316 
317     @AutoAnnotation
318     public static IncludeRunOnSecondaryUserInDifferentProfileGroupToOrganizationOwnedProfileOwnerProfileUsingParentInstance includeRunOnSecondaryUserInDifferentProfileGroupToOrganizationOwnedProfileOwnerProfileUsingParentInstance() {
319         return new AutoAnnotation_Policy_includeRunOnSecondaryUserInDifferentProfileGroupToOrganizationOwnedProfileOwnerProfileUsingParentInstance();
320     }
321 
322     @AutoAnnotation
323     private static IncludeRunOnParentOfProfileOwnerWithNoDeviceOwner includeRunOnParentOfProfileOwnerWithNoDeviceOwner() {
324         return new AutoAnnotation_Policy_includeRunOnParentOfProfileOwnerWithNoDeviceOwner();
325     }
326 
327     @AutoAnnotation
328     private static IncludeRunOnOrganizationOwnedProfileOwner includeRunOnOrganizationOwnedProfileOwner() {
329         return new AutoAnnotation_Policy_includeRunOnOrganizationOwnedProfileOwner();
330     }
331 
332     @AutoAnnotation
333     private static IncludeRunOnParentOfOrganizationOwnedProfileOwner includeRunOnParentOfOrganizationOwnedProfileOwner() {
334         return new AutoAnnotation_Policy_includeRunOnParentOfOrganizationOwnedProfileOwner();
335     }
336 
337     @AutoAnnotation
338     private static IncludeRunOnProfileOwnerPrimaryUser includeRunOnProfileOwnerPrimaryUser() {
339         return new AutoAnnotation_Policy_includeRunOnProfileOwnerPrimaryUser();
340     }
341 
342     @AutoAnnotation
343     private static IncludeRunOnBackgroundDeviceOwnerUser includeRunOnBackgroundDeviceOwnerUser() {
344         return new AutoAnnotation_Policy_includeRunOnBackgroundDeviceOwnerUser();
345     }
346 
347     @AutoAnnotation
348     private static EnsureHasDelegate ensureHasDelegate(EnsureHasDelegate.AdminType admin,
349             String[] scopes, boolean isPrimary) {
350         return new AutoAnnotation_Policy_ensureHasDelegate(admin, scopes, isPrimary);
351     }
352 
353     @AutoAnnotation
354     private static EnsureHasDevicePolicyManagerRoleHolder ensureHasDevicePolicyManagerRoleHolder(
355             UserType onUser, boolean isPrimary) {
356         return new AutoAnnotation_Policy_ensureHasDevicePolicyManagerRoleHolder(onUser, isPrimary);
357     }
358 
359     @AutoAnnotation
360     private static IncludeRunOnParentOfProfileOwnerUsingParentInstance includeRunOnParentOfProfileOwnerUsingParentInstance() {
361         return new AutoAnnotation_Policy_includeRunOnParentOfProfileOwnerUsingParentInstance();
362     }
363 
364     @AutoAnnotation
365     private static IncludeRunOnParentOfOrganizationOwnedProfileOwnerUsingParentInstance includeRunOnParentOfOrganizationOwnedProfileOwnerUsingParentInstance() {
366         return new AutoAnnotation_Policy_includeRunOnParentOfOrganizationOwnedProfileOwnerUsingParentInstance();
367     }
368 
369     @AutoAnnotation
370     private static IncludeRunOnFinancedDeviceOwnerUser includeRunOnFinancedDeviceOwnerUser() {
371         return new AutoAnnotation_Policy_includeRunOnFinancedDeviceOwnerUser();
372     }
373 
374     @AutoAnnotation
375     private static IncludeRunOnCloneProfileAlongsideManagedProfileUsingParentInstance includeRunOnCloneProfileAlongsideManagedProfileUsingParentInstance() {
376         return new AutoAnnotation_Policy_includeRunOnCloneProfileAlongsideManagedProfileUsingParentInstance();
377     }
378 
379     @AutoAnnotation
380     private static IncludeRunOnCloneProfileAlongsideOrganizationOwnedProfileUsingParentInstance includeRunOnCloneProfileAlongsideOrganizationOwnedProfileUsingParentInstance() {
381         return new AutoAnnotation_Policy_includeRunOnCloneProfileAlongsideOrganizationOwnedProfileUsingParentInstance();
382     }
383 
384     @AutoAnnotation
385     private static IncludeRunOnCloneProfileAlongsideManagedProfile includeRunOnCloneProfileAlongsideManagedProfile() {
386         return new AutoAnnotation_Policy_includeRunOnCloneProfileAlongsideManagedProfile();
387     }
388 
389     @AutoAnnotation
390     private static IncludeRunOnCloneProfileAlongsideOrganizationOwnedProfile includeRunOnCloneProfileAlongsideOrganizationOwnedProfile() {
391         return new AutoAnnotation_Policy_includeRunOnCloneProfileAlongsideOrganizationOwnedProfile();
392     }
393 
394     @AutoAnnotation
395     private static IncludeRunOnDevicePolicyManagementRoleHolderProfile includeRunOnDevicePolicyManagementRoleHolderProfile() {
396         return new AutoAnnotation_Policy_includeRunOnDevicePolicyManagementRoleHolderProfile();
397     }
398 
399     @AutoAnnotation
400     private static IncludeRunOnDevicePolicyManagementRoleHolderUser includeRunOnDevicePolicyManagementRoleHolderUser() {
401         return new AutoAnnotation_Policy_includeRunOnDevicePolicyManagementRoleHolderUser();
402     }
403 
404     @AutoAnnotation
405     private static IncludeRunOnDevicePolicyManagementRoleHolderSecondaryUser includeRunOnDevicePolicyManagementRoleHolderSecondaryUser() {
406         return new AutoAnnotation_Policy_includeRunOnDevicePolicyManagementRoleHolderSecondaryUser();
407     }
408 
409     @AutoAnnotation
410     private static EnsureFeatureFlagEnabled ensureFeatureFlagEnabled(String namespace, String key) {
411         return new AutoAnnotation_Policy_ensureFeatureFlagEnabled(namespace, key);
412     }
413 
414     @AutoAnnotation
415     private static RequireFeatureFlagEnabled requireFeatureFlagEnabled(String namespace, String key) {
416         return new AutoAnnotation_Policy_requireFeatureFlagEnabled(namespace, key);
417     }
418 
419     private static Function<EnterprisePolicy, Set<Annotation>> singleAnnotation(
420             Annotation annotation) {
421         return (i) -> ImmutableSet.of(annotation);
422     }
423 
424     private static Function<EnterprisePolicy, Set<Annotation>> generateDevicePolicyManagerRoleHolderAnnotation(
425             Annotation annotation, UserType roleHolderUser) {
426         return (policy) -> {
427             if (true) {
428                 // Temporarily disabling enforcement of third-party coexistence
429                 return ImmutableSet.of(annotation);
430             }
431 
432             // If DPM role holder is handled elsewhere - we don't special case it here
433             if (hasFlag(policy.dpc(), CANNOT_BE_APPLIED_BY_ROLE_HOLDER)
434                     || hasFlag(policy.dpc(), APPLIED_BY_DPM_ROLE_HOLDER)) {
435                 return ImmutableSet.of(annotation);
436             }
437             Annotation[] existingAnnotations = annotation.annotationType().getAnnotations();
438             Annotation[] newAnnotations = Arrays.copyOf(existingAnnotations,
439                     existingAnnotations.length + 1);
440             newAnnotations[newAnnotations.length - 1] = ensureHasDevicePolicyManagerRoleHolder(
441                     roleHolderUser, /* isPrimary= */ true);
442             return Set.of(annotation,
443                     new DynamicParameterizedAnnotation(
444                     annotation.annotationType().getSimpleName() + "_DPMRH",
445                     newAnnotations
446             ));
447         };
448     }
449 
450     private static Function<EnterprisePolicy, Set<Annotation>> generateDelegateAnnotation(
451             Annotation annotation, boolean isPrimary) {
452         return (policy) -> {
453             Annotation[] existingAnnotations = filterAnnotations(
454                     annotation.annotationType().getAnnotations(), EnsureHasNoDelegate.class);
455             return Arrays.stream(policy.delegatedScopes())
456                     .map(scope -> {
457                         Annotation[] newAnnotations = Arrays.copyOf(existingAnnotations,
458                                 existingAnnotations.length + 1);
459                         newAnnotations[newAnnotations.length - 1] = ensureHasDelegate(
460                                 EnsureHasDelegate.AdminType.PRIMARY, new String[]{scope},
461                                 isPrimary);
462 
463                         return new DynamicParameterizedAnnotation(
464                                 annotation.annotationType().getSimpleName() + "Delegate_" + scope,
465                                 newAnnotations);
466                     }).collect(Collectors.toSet());
467         };
468     }
469 
470     private static Map<Function<EnterprisePolicy, Set<Annotation>>, Set<Integer>> calculateAnnotationsMap(
471             Map<Integer, Function<EnterprisePolicy, Set<Annotation>>> annotations) {
472         Map<Function<EnterprisePolicy, Set<Annotation>>, Set<Integer>> b = new HashMap<>();
473 
474         for (Map.Entry<Integer, Function<EnterprisePolicy, Set<Annotation>>> i :
475                 annotations.entrySet()) {
476             if (!b.containsKey(i.getValue())) {
477                 b.put(i.getValue(), new HashSet<>());
478             }
479 
480             b.get(i.getValue()).add(i.getKey());
481         }
482 
483         return b;
484     }
485 
486     private static BiFunction<EnterprisePolicy, Boolean, Set<Annotation>> addGeneratedStates(
487             ImmutableMap.Entry<Integer, BiFunction<EnterprisePolicy, Boolean, Set<Annotation>>> entry) {
488         return (policy, canSet) -> {
489             if (hasFlag(policy.dpc(), entry.getKey() | CAN_BE_DELEGATED)) {
490                 Set<Annotation> results = new HashSet<>(entry.getValue().apply(policy, canSet));
491                 results.addAll(results.stream().flatMap(
492                         t -> generateDelegateAnnotation(t, /* isPrimary= */ true).apply(
493                                 policy).stream())
494                         .collect(Collectors.toSet()));
495                 return results;
496             }
497 
498             return entry.getValue().apply(policy, canSet);
499         };
500     }
501 
502     /**
503      * Get parameterized test runs for the given policy.
504      *
505      * <p>These are states which should be run where the policy is able to be applied.
506      */
507     public static List<Annotation> policyAppliesStates(EnterprisePolicy enterprisePolicy) {
508         Set<Annotation> annotations = new HashSet<>();
509 
510         for (Map.Entry<Function<EnterprisePolicy, Set<Annotation>>, Set<Integer>> annotation :
511                 ANNOTATIONS_MAP.entrySet()) {
512             if (policyWillApply(enterprisePolicy.dpc(), annotation.getValue())) {
513                 annotations.addAll(annotation.getKey().apply(enterprisePolicy));
514             }
515         }
516 
517         for (AppOp appOp : enterprisePolicy.appOps()) {
518             // TODO(b/219750042): Currently we only test that app ops apply to the current user
519             Annotation[] withAppOpAnnotations = new Annotation[]{
520                     ensureTestAppInstalled(DELEGATE_KEY, queryBuilder()
521                                     .wherePackageName().isEqualTo(DELEGATE_PACKAGE_NAME)
522                                     .toAnnotation(),
523                             UserType.INSTRUMENTED_USER, /* isPrimary= */ true),
524                     ensureTestAppHasAppOp(DELEGATE_KEY, new String[]{appOp.appliedWith()})
525             };
526             annotations.add(
527                     new DynamicParameterizedAnnotation(
528                             "AppOp_" + appOp.appliedWith(), withAppOpAnnotations));
529         }
530 
531         for (Permission permission : enterprisePolicy.permissions()) {
532             // TODO(b/219750042): Currently we only test that permissions apply to the current user
533             Annotation[] withPermissionAnnotations = new Annotation[]{
534                     ensureTestAppInstalled(DELEGATE_KEY, queryBuilder()
535                                     .wherePackageName().isEqualTo(DELEGATE_PACKAGE_NAME)
536                                     .toAnnotation(),
537                             UserType.INSTRUMENTED_USER, /* isPrimary= */ true),
538                     ensureTestAppHasPermission(DELEGATE_KEY,
539                             new String[]{permission.appliedWith()}, FailureMode.SKIP)
540             };
541             // TODO(281651179): Re-enable
542 //            annotations.add(
543 //                    new DynamicParameterizedAnnotation(
544 //                            "Permission_" + formatPermissionForTestName(permission.appliedWith()), withPermissionAnnotations));
545         }
546 
547         removeShadowingAnnotations(annotations);
548 
549         if (annotations.isEmpty()) {
550             // Don't run the original test unparameterized
551             annotations.add(includeNone());
552         }
553 
554         return new ArrayList<>(annotations);
555     }
556 
557     private static boolean policyWillApply(int[] policyFlags, Set<Integer> annotationFlags) {
558         for (int annotationFlag : annotationFlags) {
559             if (hasFlag(policyFlags, annotationFlag)) {
560                 return true;
561             }
562         }
563         return false;
564     }
565 
566     private static boolean policyWillNotApply(int[] policyFlags, Set<Integer> annotationFlags) {
567         for (int annotationFlag : annotationFlags) {
568             if (hasFlag(annotationFlag,
569                     DO_NOT_APPLY_TO_POLICY_DOES_NOT_APPLY_TESTS, /* nonMatchingFlag= */ NO)) {
570                 return false; // We don't support using this annotation for PolicyDoesNotApply tests
571             }
572 
573             int appliedByFlag = APPLIED_BY_FLAGS & annotationFlag;
574             int otherFlags = annotationFlag ^ appliedByFlag; // remove the appliedByFlag
575             if (hasFlag(policyFlags, /* matchingFlag= */ appliedByFlag, /* nonMatchingFlag= */
576                     otherFlags)) {
577                 return true;
578             }
579         }
580 
581         return false;
582     }
583 
584     /**
585      * Get parameterized test runs for the given policy.
586      *
587      * <p>These are states which should be run where the policy is not able to be applied.
588      */
589     public static List<Annotation> policyDoesNotApplyStates(EnterprisePolicy enterprisePolicy) {
590         Set<Annotation> annotations = new HashSet<>();
591 
592         for (Map.Entry<Function<EnterprisePolicy, Set<Annotation>>, Set<Integer>> annotation :
593                 ANNOTATIONS_MAP.entrySet()) {
594             if (policyWillNotApply(enterprisePolicy.dpc(), annotation.getValue())) {
595                 annotations.addAll(annotation.getKey().apply(enterprisePolicy));
596             }
597         }
598 
599         removeShadowedAnnotations(annotations);
600 
601         if (annotations.isEmpty()) {
602             // Don't run the original test unparameterized
603             annotations.add(includeNone());
604         }
605 
606         return new ArrayList<>(annotations);
607     }
608 
609     /**
610      * Get parameterized test runs where the policy cannot be set for the given policy.
611      */
612     public static List<Annotation> cannotSetPolicyStates(
613             EnterprisePolicy enterprisePolicy,
614             boolean includeDeviceAdminStates, boolean includeNonDeviceAdminStates) {
615 
616         Set<Annotation> annotations = new HashSet<>();
617         if (includeDeviceAdminStates) {
618             int allFlags = 0;
619             for (int p : enterprisePolicy.dpc()) {
620                 allFlags = allFlags | p;
621             }
622 
623             for (Map.Entry<Integer, BiFunction<EnterprisePolicy, Boolean, Set<Annotation>>> appliedByFlag :
624                     DPC_STATE_ANNOTATIONS.entrySet()) {
625 
626                 if ((appliedByFlag.getKey()
627                         & APPLIED_BY_DPM_ROLE_HOLDER) == APPLIED_BY_DPM_ROLE_HOLDER) {
628                     if (!includeNonDeviceAdminStates) {
629                         // Temp fix to avoid random failing tests
630                         continue;
631                     }
632                 }
633 
634                 if (hasFlag(appliedByFlag.getKey(),
635                         DO_NOT_APPLY_TO_CANNOT_SET_POLICY_TESTS, /* nonMatchingFlag= */ NO)) {
636                     continue;
637                 }
638 
639                 if ((appliedByFlag.getKey() & allFlags) == 0) {
640                     annotations.addAll(appliedByFlag.getValue().apply(enterprisePolicy, /* canSet= */ false));
641                 }
642             }
643         }
644 
645         if (includeNonDeviceAdminStates) {
646             Set<String> validScopes = ImmutableSet.copyOf(enterprisePolicy.delegatedScopes());
647             String[] scopes = ALL_DELEGATE_SCOPES.stream()
648                     .filter(i -> !validScopes.contains(i))
649                     .toArray(String[]::new);
650             Annotation[] existingAnnotations = filterAnnotations(
651                     IncludeRunOnDeviceOwnerUser.class.getAnnotations(), EnsureHasNoDelegate.class);
652 
653             String[] validPermissions = Arrays.stream(enterprisePolicy.permissions())
654                     .map(p -> p.appliedWith()).toArray(String[]::new);
655 
656             if (BedsteadJUnit4.isDebug()) {
657                 // Add a non-DPC with no delegate scopes
658                 Annotation[] newAnnotations = Arrays.copyOf(existingAnnotations,
659                         existingAnnotations.length + 2);
660                 newAnnotations[newAnnotations.length - 2] = ensureHasDelegate(
661                         EnsureHasDelegate.AdminType.PRIMARY, new String[]{},
662                         /* isPrimary= */ true);
663                 newAnnotations[newAnnotations.length - 1] = ensureTestAppDoesNotHavePermission(
664                         DELEGATE_KEY, validPermissions, FailureMode.SKIP);
665                 annotations.add(
666                         new DynamicParameterizedAnnotation("DelegateWithNoScopes", newAnnotations));
667 
668                 for (String scope : scopes) {
669                     newAnnotations = Arrays.copyOf(existingAnnotations,
670                             existingAnnotations.length + 1);
671                     newAnnotations[newAnnotations.length - 1] = ensureHasDelegate(
672                             EnsureHasDelegate.AdminType.PRIMARY, new String[]{scope},
673                             /* isPrimary= */ true);
674                     annotations.add(
675                             new DynamicParameterizedAnnotation("DelegateWithScope:" + scope, newAnnotations));
676                 }
677             } else {
678                 Annotation[] newAnnotations = Arrays.copyOf(existingAnnotations,
679                         existingAnnotations.length + 2);
680                 newAnnotations[newAnnotations.length - 2] = ensureHasDelegate(
681                         EnsureHasDelegate.AdminType.PRIMARY, scopes, /* isPrimary= */ true);
682                 newAnnotations[newAnnotations.length - 1] = ensureTestAppDoesNotHavePermission(
683                         DELEGATE_KEY, validPermissions, FailureMode.SKIP);
684                 annotations.add(
685                         new DynamicParameterizedAnnotation("DelegateWithoutValidScope",
686                                 newAnnotations));
687             }
688         }
689 
690         removeShadowedAnnotations(annotations);
691 
692         if (annotations.isEmpty()) {
693             // Don't run the original test unparameterized
694             annotations.add(includeNone());
695         }
696 
697         return new ArrayList<>(annotations);
698     }
699 
700     /** Return {@code annotations} excluding any which are of type
701      * {@code filteredAnnotationClass}. */
702     private static Annotation[] filterAnnotations(Annotation[] annotations,
703             Class<? extends Annotation> filteredAnnotationClass) {
704         return Arrays.stream(annotations).filter(
705                 f -> !f.annotationType().equals(filteredAnnotationClass))
706                 .toArray(Annotation[]::new);
707     }
708 
709     /**
710      * Get state annotations where the policy can be set for the given policy.
711      */
712     public static List<Annotation> canSetPolicyStates(
713             EnterprisePolicy enterprisePolicy, boolean singleTestOnly) {
714         Set<Annotation> annotations = new HashSet<>();
715 
716         int allFlags = 0;
717         for (int p : enterprisePolicy.dpc()) {
718             allFlags = allFlags | p;
719         }
720 
721         for (Map.Entry<Integer, BiFunction<EnterprisePolicy, Boolean, Set<Annotation>>> appliedByFlag :
722                 DPC_STATE_ANNOTATIONS.entrySet()) {
723             if ((appliedByFlag.getKey() & allFlags) == appliedByFlag.getKey()) {
724                 annotations.addAll(appliedByFlag.getValue().apply(enterprisePolicy, /* canSet= */ true));
725             }
726         }
727 
728         for (AppOp appOp : enterprisePolicy.appOps()) {
729             // TODO(b/219750042): Currently we only test that app ops can be set as the primary user
730             Annotation[] withAppOpAnnotations = new Annotation[]{
731                     ensureTestAppInstalled(DELEGATE_KEY,
732                             queryBuilder()
733                                     .wherePackageName().isEqualTo(DELEGATE_PACKAGE_NAME)
734                                     .toAnnotation(), UserType.INSTRUMENTED_USER,
735                             /* isPrimary= */ true),
736                     ensureTestAppHasAppOp(DELEGATE_KEY, new String[]{appOp.appliedWith()})
737             };
738             annotations.add(
739                     new DynamicParameterizedAnnotation(
740                             "AppOp_" + appOp.appliedWith(), withAppOpAnnotations));
741         }
742 
743         for (Permission permission : enterprisePolicy.permissions()) {
744             // TODO(b/219750042): Currently we only test that permissions can be set as the primary user
745             Annotation[] withPermissionAnnotations = new Annotation[]{
746                     ensureTestAppInstalled(DELEGATE_KEY,
747                             queryBuilder()
748                                     .wherePackageName().isEqualTo(DELEGATE_PACKAGE_NAME)
749                                     .toAnnotation(), UserType.INSTRUMENTED_USER,
750                             /* isPrimary= */ true),
751                     ensureTestAppHasPermission(
752                             DELEGATE_KEY, new String[]{permission.appliedWith()}, FailureMode.SKIP),
753             };
754             // TODO(281651179): Re-enable
755 //            annotations.add(
756 //                    new DynamicParameterizedAnnotation(
757 //                            "Permission_" + formatPermissionForTestName(permission.appliedWith()), withPermissionAnnotations));
758         }
759 
760         removeShadowingAnnotations(annotations);
761 
762         if (annotations.isEmpty()) {
763             // Don't run the original test unparameterized
764             annotations.add(includeNone());
765         }
766 
767         List<Annotation> annotationList = new ArrayList<>(annotations);
768 
769         if (singleTestOnly) {
770             // We select one annotation in an arbitrary but deterministic way
771             annotationList.sort(Comparator.comparing(
772                     a -> a instanceof DynamicParameterizedAnnotation
773                             ? "DynamicParameterizedAnnotation" : a.annotationType().getName()));
774 
775             // We don't want a delegate to be the representative test
776             Annotation firstAnnotation = annotationList.stream()
777                     .filter(i -> !(i instanceof DynamicParameterizedAnnotation))
778                     .findFirst().get();
779             annotationList.clear();
780             annotationList.add(firstAnnotation);
781         }
782 
783         return annotationList;
784     }
785 
786     /** Validate flags used by a DPC policy. */
787     public static void validateFlags(String policyName, int[] values) {
788         int usedAppliedByFlags = 0;
789 
790         for (int value : values) {
791             validateFlags(policyName, value);
792             int newUsedAppliedByFlags = usedAppliedByFlags | (value & APPLIED_BY_FLAGS);
793             if (newUsedAppliedByFlags == usedAppliedByFlags) {
794                 throw new IllegalStateException(
795                         "Cannot have more than one policy flag APPLIED by the same component. "
796                                 + "Error in policy " + policyName);
797             }
798             usedAppliedByFlags = newUsedAppliedByFlags;
799         }
800     }
801 
802     private static void validateFlags(String policyName, int value) {
803         int matchingAppliedByFlags = APPLIED_BY_FLAGS & value;
804 
805         if (matchingAppliedByFlags == 0) {
806             throw new IllegalStateException(
807                     "All policy flags must specify 1 APPLIED_BY flag. Policy " + policyName
808                             + " did not.");
809         }
810     }
811 
812     private static boolean hasFlag(int[] values, int matchingFlag) {
813         return hasFlag(values, matchingFlag, /* nonMatchingFlag= */ NO);
814     }
815 
816     private static boolean hasFlag(int[] values, int matchingFlag, int nonMatchingFlag) {
817         for (int value : values) {
818             if (hasFlag(value, matchingFlag, nonMatchingFlag)) {
819                 return true;
820             }
821         }
822         return false;
823     }
824 
825     private static boolean hasFlag(int value, int matchingFlag, int nonMatchingFlag) {
826         if (!((value & matchingFlag) == matchingFlag)) {
827             return false;
828         }
829 
830         if (nonMatchingFlag != NO) {
831             return (value & nonMatchingFlag) != nonMatchingFlag;
832         }
833 
834         return true;
835     }
836 
837     /**
838      * Remove entries from {@code annotations} which are shadowed by another entry
839      * in {@code annotations} (directly or indirectly).
840      */
841     private static void removeShadowedAnnotations(Set<Annotation> annotations) {
842         Set<Class<? extends Annotation>> shadowedAnnotations = new HashSet<>();
843         for (Annotation annotation : annotations) {
844             if (annotation instanceof DynamicParameterizedAnnotation) {
845                 continue; // Doesn't shadow anything
846             }
847 
848             ParameterizedAnnotation parameterizedAnnotation =
849                     annotation.annotationType().getAnnotation(ParameterizedAnnotation.class);
850 
851             if (parameterizedAnnotation == null) {
852                 continue; // Not parameterized
853             }
854 
855             for (Class<? extends Annotation> shadowedAnnotationClass
856                     : parameterizedAnnotation.shadows()) {
857                 addShadowed(shadowedAnnotations, shadowedAnnotationClass);
858             }
859         }
860         annotations.removeIf(a -> shadowedAnnotations.contains(a.annotationType()));
861     }
862 
863     private static void addShadowed(Set<Class<? extends Annotation>> shadowedAnnotations,
864             Class<? extends Annotation> annotationClass) {
865         shadowedAnnotations.add(annotationClass);
866         ParameterizedAnnotation parameterizedAnnotation =
867                 annotationClass.getAnnotation(ParameterizedAnnotation.class);
868 
869         if (parameterizedAnnotation == null) {
870             return;
871         }
872 
873         for (Class<? extends Annotation> shadowedAnnotationClass
874                 : parameterizedAnnotation.shadows()) {
875             addShadowed(shadowedAnnotations, shadowedAnnotationClass);
876         }
877     }
878 
879     // This maps classes to classes which shadow them - we just need to ensure it contains all
880     // annotation classes we encounter
881     private static Map<Class<? extends Annotation>, Set<Class<? extends Annotation>>>
882             sReverseShadowMap = new HashMap<>();
883 
884     /**
885      * Remove entries from {@code annotations} which are shadowing another entry
886      * in {@code annotatipns} (directly or indirectly).
887      */
888     private static void removeShadowingAnnotations(Set<Annotation> annotations) {
889         for (Annotation annotation : annotations) {
890             recordInReverseShadowMap(annotation);
891         }
892 
893         Set<Class<? extends Annotation>> shadowingAnnotations = new HashSet<>();
894 
895         for (Annotation annotation : annotations) {
896             shadowingAnnotations.addAll(
897                     sReverseShadowMap.getOrDefault(annotation.annotationType(), Set.of()));
898         }
899 
900         annotations.removeIf(a -> shadowingAnnotations.contains(a.annotationType()));
901     }
902 
903     private static void recordInReverseShadowMap(Annotation annotation) {
904         if (annotation instanceof DynamicParameterizedAnnotation) {
905             return; // Not shadowed by anything
906         }
907 
908         ParameterizedAnnotation parameterizedAnnotation =
909                 annotation.annotationType().getAnnotation(ParameterizedAnnotation.class);
910 
911         if (parameterizedAnnotation == null) {
912             return; // Not parameterized
913         }
914 
915         if (parameterizedAnnotation.shadows().length == 0) {
916             return; // Doesn't shadow anything
917         }
918 
919         recordShadowedInReverseShadowMap(annotation.annotationType(), parameterizedAnnotation);
920     }
921 
922     private static void recordShadowedInReverseShadowMap(Class<? extends Annotation> annotation,
923             ParameterizedAnnotation parameterizedAnnotation) {
924         for (Class<? extends Annotation> shadowedAnnotation : parameterizedAnnotation.shadows()) {
925             ParameterizedAnnotation shadowedParameterizedAnnotation =
926                     shadowedAnnotation.getAnnotation(ParameterizedAnnotation.class);
927 
928             if (shadowedParameterizedAnnotation == null) {
929                 continue; // Not parameterized
930             }
931 
932             if (!sReverseShadowMap.containsKey(shadowedAnnotation)) {
933                 sReverseShadowMap.put(shadowedAnnotation, new HashSet<>());
934             }
935 
936             sReverseShadowMap.get(shadowedAnnotation).add(annotation);
937 
938             recordShadowedInReverseShadowMap(annotation, shadowedParameterizedAnnotation);
939         }
940     }
941 
942     private static String formatPermissionForTestName(String permission) {
943         if (permission.startsWith("android.permission.")) {
944             return permission.substring(19);
945         }
946         return permission;
947     }
948 }
949