1 /*
2  * Copyright (C) 2021 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 android.devicepolicy.cts;
18 
19 import static com.android.bedstead.remotedpc.RemoteDpc.DPC_COMPONENT_NAME;
20 
21 import static com.google.common.truth.Truth.assertThat;
22 
23 import static org.junit.Assert.assertThrows;
24 
25 import static java.util.Collections.singleton;
26 
27 import android.content.Context;
28 import android.net.Uri;
29 import android.os.Process;
30 import android.security.KeyChain;
31 import android.security.KeyChainException;
32 
33 import com.android.activitycontext.ActivityContext;
34 import com.android.bedstead.harrier.BedsteadJUnit4;
35 import com.android.bedstead.harrier.DeviceState;
36 import com.android.bedstead.harrier.annotations.Postsubmit;
37 import com.android.bedstead.harrier.annotations.enterprise.CanSetPolicyTest;
38 import com.android.bedstead.harrier.annotations.enterprise.PositivePolicyTest;
39 import com.android.bedstead.harrier.policies.KeyManagement;
40 import com.android.bedstead.nene.TestApis;
41 import com.android.compatibility.common.util.BlockingCallback;
42 import com.android.compatibility.common.util.FakeKeys;
43 
44 import org.junit.ClassRule;
45 import org.junit.Ignore;
46 import org.junit.Rule;
47 import org.junit.Test;
48 import org.junit.runner.RunWith;
49 import org.testng.Assert;
50 
51 import java.io.ByteArrayInputStream;
52 import java.io.UnsupportedEncodingException;
53 import java.net.URLEncoder;
54 import java.security.InvalidKeyException;
55 import java.security.KeyFactory;
56 import java.security.NoSuchAlgorithmException;
57 import java.security.PrivateKey;
58 import java.security.Signature;
59 import java.security.cert.Certificate;
60 import java.security.cert.CertificateException;
61 import java.security.cert.CertificateFactory;
62 import java.security.spec.InvalidKeySpecException;
63 import java.security.spec.PKCS8EncodedKeySpec;
64 import java.util.Map;
65 import java.util.concurrent.TimeUnit;
66 
67 /**
68  * Test that a DPC can manage keys and certificate on a device by installing, generating and
69  * removing key pairs via DevicePolicyManager APIs. The instrumented test app can use the installed
70  * keys by requesting access to them and retrieving them via KeyChain APIs.
71  */
72 @RunWith(BedsteadJUnit4.class)
73 public final class KeyManagementTest {
74 
75     @ClassRule
76     @Rule
77     public static final DeviceState sDeviceState = new DeviceState();
78     private static final int KEYCHAIN_CALLBACK_TIMEOUT_SECONDS = 600;
79     private static final String RSA = "RSA";
80     private static final String RSA_ALIAS = "com.android.test.valid-rsa-key-1";
81     private static final PrivateKey PRIVATE_KEY =
82             generatePrivateKey(FakeKeys.FAKE_RSA_1.privateKey, RSA);
83     private static final Certificate CERTIFICATE =
84             generateCertificate(FakeKeys.FAKE_RSA_1.caCertificate);
85     private static final Certificate[] CERTIFICATES = new Certificate[]{CERTIFICATE};
86     private static final String NON_EXISTENT_ALIAS = "KeyManagementTest-nonexistent";
87     private static final Context sContext = TestApis.context().instrumentedContext();
88 
getUri(String alias)89     private static Uri getUri(String alias) {
90         try {
91             return Uri.parse("https://example.org/?alias=" + URLEncoder.encode(alias, "UTF-8"));
92         } catch (UnsupportedEncodingException e) {
93             throw new AssertionError("Unable to parse URI." + e);
94         }
95     }
96 
choosePrivateKeyAlias(KeyChainAliasCallback callback, String alias)97     private static void choosePrivateKeyAlias(KeyChainAliasCallback callback, String alias) {
98         /* Pass the alias as a GET to an imaginary server instead of explicitly asking for it,
99          * to make sure the DPC actually has to do some work to grant the cert.
100          */
101         try {
102             ActivityContext.runWithContext(
103                     (activity) -> KeyChain.choosePrivateKeyAlias(activity, callback, /* keyTypes= */
104                             null, /* issuers= */ null, getUri(alias), /* alias = */ null)
105             );
106         } catch (InterruptedException e) {
107             throw new AssertionError("Unable to choose private key alias." + e);
108         }
109     }
110 
getPrivateKey(Context context, String alias)111     private static PrivateKey getPrivateKey(Context context, String alias) {
112         try {
113             return KeyChain.getPrivateKey(context, alias);
114         } catch (KeyChainException | InterruptedException e) {
115             throw new AssertionError("Failed to get private key." + e);
116         }
117     }
118 
generatePrivateKey(final byte[] key, String type)119     private static PrivateKey generatePrivateKey(final byte[] key, String type) {
120         try {
121             return KeyFactory.getInstance(type).generatePrivate(
122                     new PKCS8EncodedKeySpec(key));
123         } catch (InvalidKeySpecException | NoSuchAlgorithmException e) {
124             throw new AssertionError("Unable to get private key." + e);
125         }
126     }
127 
generateCertificate(byte[] cert)128     private static Certificate generateCertificate(byte[] cert) {
129         try {
130             return CertificateFactory.getInstance("X.509").generateCertificate(
131                     new ByteArrayInputStream(cert));
132         } catch (CertificateException e) {
133             throw new AssertionError("Unable to get certificate." + e);
134         }
135     }
136 
137     @Test
138     @Postsubmit(reason = "new test")
139     @CanSetPolicyTest(policy = KeyManagement.class)
installKeyPair_validRsaKeyPair_success()140     public void installKeyPair_validRsaKeyPair_success() {
141         try {
142             // Install keypair
143             assertThat(sDeviceState.dpc().devicePolicyManager()
144                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATE,
145                             RSA_ALIAS)).isTrue();
146         } finally {
147             // Remove keypair
148             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
149         }
150     }
151 
152     @Test
153     @Postsubmit(reason = "new test")
154     @CanSetPolicyTest(policy = KeyManagement.class, singleTestOnly = true)
installKeyPair_nullPrivateKey_throwException()155     public void installKeyPair_nullPrivateKey_throwException() {
156         assertThrows(NullPointerException.class,
157                 () -> sDeviceState.dpc().devicePolicyManager().installKeyPair(
158                         DPC_COMPONENT_NAME, /* privKey = */ null, CERTIFICATE, RSA_ALIAS));
159     }
160 
161     @Test
162     @Postsubmit(reason = "new test")
163     @CanSetPolicyTest(policy = KeyManagement.class, singleTestOnly = true)
installKeyPair_nullCertificate_throwException()164     public void installKeyPair_nullCertificate_throwException() {
165         assertThrows(NullPointerException.class,
166                 () -> sDeviceState.dpc().devicePolicyManager().installKeyPair(
167                         DPC_COMPONENT_NAME, PRIVATE_KEY, /* cert = */ null, RSA_ALIAS));
168     }
169 
170     @Test
171     @Postsubmit(reason = "new test")
172     @CanSetPolicyTest(policy = KeyManagement.class, singleTestOnly = true)
installKeyPair_nullAdminComponent_throwException()173     public void installKeyPair_nullAdminComponent_throwException() {
174         assertThrows(SecurityException.class,
175                 () -> sDeviceState.dpc().devicePolicyManager().installKeyPair(
176                         /* admin = */ null, PRIVATE_KEY, CERTIFICATE, RSA_ALIAS));
177     }
178 
179     @Test
180     @Ignore("TODO(b/204544463): Enable when the key can be serialized")
181     @Postsubmit(reason = "new test")
182     @CanSetPolicyTest(policy = KeyManagement.class)
installKeyPair_withAutomatedAccess_aliasIsGranted()183     public void installKeyPair_withAutomatedAccess_aliasIsGranted() throws Exception {
184         try {
185             // Install keypair with automated access
186             sDeviceState.dpc().devicePolicyManager().installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY,
187                     CERTIFICATES, RSA_ALIAS, /* requestAccess = */ true);
188 
189             // TODO(b/204544478): Remove the null context
190             assertThat(sDeviceState.dpc().keyChain().getPrivateKey(/* context= */ null, RSA_ALIAS))
191                     .isNotNull();
192         } finally {
193             // Remove keypair
194             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
195         }
196     }
197 
198     @Test
199     @Postsubmit(reason = "new test")
200     @CanSetPolicyTest(policy = KeyManagement.class)
installKeyPair_withoutAutomatedAccess_aliasIsNotGranted()201     public void installKeyPair_withoutAutomatedAccess_aliasIsNotGranted() throws Exception {
202         try {
203             // Install keypair with automated access
204             sDeviceState.dpc().devicePolicyManager().installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY,
205                     CERTIFICATES, RSA_ALIAS, /* requestAccess = */ false);
206 
207             // TODO(b/204544478): Remove the null context
208             assertThat(sDeviceState.dpc().keyChain().getPrivateKey(/* context= */ null, RSA_ALIAS))
209                     .isNull();
210         } finally {
211             // Remove keypair
212             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
213         }
214     }
215 
216     @Test
217     @Postsubmit(reason = "new test")
218     @CanSetPolicyTest(policy = KeyManagement.class)
removeKeyPair_validRsaKeyPair_success()219     public void removeKeyPair_validRsaKeyPair_success() {
220         try {
221             // Install keypair
222             sDeviceState.dpc().devicePolicyManager()
223                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATE, RSA_ALIAS);
224         } finally {
225             // Remove keypair
226             assertThat(sDeviceState.dpc().devicePolicyManager()
227                     .removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS)).isTrue();
228         }
229     }
230 
231     @Test
232     @Postsubmit(reason = "new test")
233     @CanSetPolicyTest(policy = KeyManagement.class)
hasKeyPair_nonExistentAlias_false()234     public void hasKeyPair_nonExistentAlias_false() {
235         assertThat(
236                 sDeviceState.dpc().devicePolicyManager().hasKeyPair(NON_EXISTENT_ALIAS)).isFalse();
237     }
238 
239     @Test
240     @Postsubmit(reason = "new test")
241     @CanSetPolicyTest(policy = KeyManagement.class)
hasKeyPair_installedAlias_true()242     public void hasKeyPair_installedAlias_true() {
243         try {
244             // Install keypair
245             sDeviceState.dpc().devicePolicyManager()
246                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATE, RSA_ALIAS);
247 
248             assertThat(sDeviceState.dpc().devicePolicyManager().hasKeyPair(RSA_ALIAS)).isTrue();
249         } finally {
250             // Remove keypair
251             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
252         }
253     }
254 
255     @Test
256     @Postsubmit(reason = "new test")
257     @CanSetPolicyTest(policy = KeyManagement.class)
hasKeyPair_removedAlias_false()258     public void hasKeyPair_removedAlias_false() {
259         try {
260             // Install keypair
261             sDeviceState.dpc().devicePolicyManager()
262                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATE, RSA_ALIAS);
263             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
264 
265             assertThat(sDeviceState.dpc().devicePolicyManager().hasKeyPair(RSA_ALIAS)).isFalse();
266         } finally {
267             // Remove keypair
268             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
269         }
270     }
271 
272     @Test
273     @Postsubmit(reason = "new test")
274     @PositivePolicyTest(policy = KeyManagement.class)
choosePrivateKeyAlias_aliasIsSelectedByAdmin_returnAlias()275     public void choosePrivateKeyAlias_aliasIsSelectedByAdmin_returnAlias() throws Exception {
276         try {
277             // Install keypair
278             sDeviceState.dpc().devicePolicyManager()
279                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATE, RSA_ALIAS);
280             KeyChainAliasCallback callback = new KeyChainAliasCallback();
281 
282             choosePrivateKeyAlias(callback, RSA_ALIAS);
283 
284             assertThat(callback.await(KEYCHAIN_CALLBACK_TIMEOUT_SECONDS, TimeUnit.SECONDS))
285                     .isEqualTo(RSA_ALIAS);
286         } finally {
287             // Remove keypair
288             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
289         }
290     }
291 
292     @Test
293     @Postsubmit(reason = "new test")
294     @PositivePolicyTest(policy = KeyManagement.class)
choosePrivateKeyAlias_nonUserSelectedAliasIsSelectedByAdmin_returnAlias()295     public void choosePrivateKeyAlias_nonUserSelectedAliasIsSelectedByAdmin_returnAlias()
296             throws Exception {
297         try {
298             // Install keypair which is not user selectable
299             sDeviceState.dpc().devicePolicyManager().installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY,
300                     CERTIFICATES, RSA_ALIAS, /* flags = */ 0);
301             KeyChainAliasCallback callback = new KeyChainAliasCallback();
302 
303             choosePrivateKeyAlias(callback, RSA_ALIAS);
304 
305             assertThat(callback.await(KEYCHAIN_CALLBACK_TIMEOUT_SECONDS, TimeUnit.SECONDS))
306                     .isEqualTo(RSA_ALIAS);
307         } finally {
308             // Remove keypair
309             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
310         }
311     }
312 
313     @Test
314     @Postsubmit(reason = "new test")
315     @PositivePolicyTest(policy = KeyManagement.class)
getPrivateKey_aliasIsGranted_returnPrivateKey()316     public void getPrivateKey_aliasIsGranted_returnPrivateKey() throws Exception {
317         try {
318             // Install keypair
319             sDeviceState.dpc().devicePolicyManager()
320                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATE, RSA_ALIAS);
321             // Grant alias via {@code KeyChain.choosePrivateKeyAlias}
322             KeyChainAliasCallback callback = new KeyChainAliasCallback();
323             choosePrivateKeyAlias(callback, RSA_ALIAS);
324             callback.await(KEYCHAIN_CALLBACK_TIMEOUT_SECONDS, TimeUnit.SECONDS);
325 
326             // Get private key for the granted alias
327             final PrivateKey privateKey =
328                     getPrivateKey(TestApis.context().instrumentedContext(), RSA_ALIAS);
329 
330             assertThat(privateKey).isNotNull();
331             assertThat(privateKey.getAlgorithm()).isEqualTo(RSA);
332 
333         } finally {
334             // Remove keypair
335             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
336         }
337     }
338 
339     @Test
340     @Postsubmit(reason = "new test")
341     @CanSetPolicyTest(policy = KeyManagement.class)
install_wasPreviouslyGrantedOnPreviousInstall_grantDoesNotPersist()342     public void install_wasPreviouslyGrantedOnPreviousInstall_grantDoesNotPersist()
343             throws Exception {
344         try {
345             sDeviceState.dpc().devicePolicyManager()
346                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES, RSA_ALIAS, true);
347             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
348 
349             sDeviceState.dpc().devicePolicyManager()
350                     .installKeyPair(DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES, RSA_ALIAS,
351                             false);
352 
353             assertThat(sDeviceState.dpc().keyChain().getPrivateKey(
354                     TestApis.context().instrumentedContext(), RSA_ALIAS))
355                     .isNull();
356         } finally {
357             // Remove keypair
358             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
359         }
360     }
361 
362     @Test
363     @Postsubmit(reason = "new test")
364     @CanSetPolicyTest(policy = KeyManagement.class, singleTestOnly = true)
getKeyPairGrants_nonExistent_throwsIllegalArgumentException()365     public void getKeyPairGrants_nonExistent_throwsIllegalArgumentException() {
366         Assert.assertThrows(IllegalArgumentException.class,
367                 () -> sDeviceState.dpc().devicePolicyManager()
368                         .getKeyPairGrants(NON_EXISTENT_ALIAS));
369     }
370 
371     @Test
372     @Postsubmit(reason = "new test")
373     @CanSetPolicyTest(policy = KeyManagement.class)
getKeyPairGrants_doesNotIncludeNotGranted()374     public void getKeyPairGrants_doesNotIncludeNotGranted() {
375         try {
376             sDeviceState.dpc().devicePolicyManager().installKeyPair(
377                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
378                     RSA_ALIAS, /* requestAccess= */ false);
379 
380             assertThat(
381                     sDeviceState.dpc().devicePolicyManager().getKeyPairGrants(RSA_ALIAS)).isEmpty();
382         } finally {
383             // Remove keypair
384             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
385         }
386     }
387 
388     @Test
389     @Postsubmit(reason = "new test")
390     @CanSetPolicyTest(policy = KeyManagement.class)
getKeyPairGrants_includesGrantedAtInstall()391     public void getKeyPairGrants_includesGrantedAtInstall() {
392         try {
393             sDeviceState.dpc().devicePolicyManager().installKeyPair(
394                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
395                     RSA_ALIAS, /* requestAccess= */ true);
396 
397             assertThat(sDeviceState.dpc().devicePolicyManager().getKeyPairGrants(RSA_ALIAS))
398                     .isEqualTo(Map.of(sDeviceState.dpc().process().uid(),
399                             singleton(sDeviceState.dpc().componentName().getPackageName())));
400         } finally {
401             // Remove keypair
402             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
403         }
404     }
405 
406     @Test
407     @Postsubmit(reason = "new test")
408     @PositivePolicyTest(policy = KeyManagement.class)
getKeyPairGrants_includesGrantedExplicitly()409     public void getKeyPairGrants_includesGrantedExplicitly() {
410         try {
411             sDeviceState.dpc().devicePolicyManager().installKeyPair(
412                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
413                     RSA_ALIAS, /* requestAccess= */ false);
414             sDeviceState.dpc().devicePolicyManager().grantKeyPairToApp(
415                     DPC_COMPONENT_NAME, RSA_ALIAS,
416                     sContext.getPackageName());
417 
418             assertThat(sDeviceState.dpc().devicePolicyManager().getKeyPairGrants(RSA_ALIAS))
419                     .isEqualTo(Map.of(Process.myUid(),
420                             singleton(sContext.getPackageName())));
421         } finally {
422             // Remove keypair
423             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
424         }
425     }
426 
427     @Test
428     @Postsubmit(reason = "new test")
429     @CanSetPolicyTest(policy = KeyManagement.class)
getKeyPairGrants_doesNotIncludeRevoked()430     public void getKeyPairGrants_doesNotIncludeRevoked() {
431         try {
432             sDeviceState.dpc().devicePolicyManager().installKeyPair(
433                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
434                     RSA_ALIAS, /* requestAccess= */ true);
435             sDeviceState.dpc().devicePolicyManager().revokeKeyPairFromApp(
436                     DPC_COMPONENT_NAME, RSA_ALIAS,
437                     sDeviceState.dpc().componentName().getPackageName());
438 
439             assertThat(
440                     sDeviceState.dpc().devicePolicyManager().getKeyPairGrants(RSA_ALIAS)).isEmpty();
441         } finally {
442             // Remove keypair
443             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
444         }
445     }
446 
447     @Test
448     @Postsubmit(reason = "new test")
449     @CanSetPolicyTest(policy = KeyManagement.class)
isKeyPairGrantedToWifiAuth_default_returnsFalse()450     public void isKeyPairGrantedToWifiAuth_default_returnsFalse() {
451         try {
452             sDeviceState.dpc().devicePolicyManager().installKeyPair(
453                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
454                     RSA_ALIAS, /* requestAccess= */ false);
455 
456             assertThat(
457                     sDeviceState.dpc().devicePolicyManager().isKeyPairGrantedToWifiAuth(RSA_ALIAS))
458                     .isFalse();
459         } finally {
460             // Remove keypair
461             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
462         }
463     }
464 
465     @Test
466     @Postsubmit(reason = "new test")
467     @CanSetPolicyTest(policy = KeyManagement.class)
isKeyPairGrantedToWifiAuth_granted_returnsTrue()468     public void isKeyPairGrantedToWifiAuth_granted_returnsTrue() {
469         try {
470             sDeviceState.dpc().devicePolicyManager().installKeyPair(
471                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
472                     RSA_ALIAS, /* requestAccess= */ false);
473             sDeviceState.dpc().devicePolicyManager().grantKeyPairToWifiAuth(RSA_ALIAS);
474 
475             assertThat(
476                     sDeviceState.dpc().devicePolicyManager().isKeyPairGrantedToWifiAuth(RSA_ALIAS))
477                     .isTrue();
478         } finally {
479             // Remove keypair
480             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
481         }
482     }
483 
484     @Test
485     @Postsubmit(reason = "new test")
486     @CanSetPolicyTest(policy = KeyManagement.class)
isKeyPairGrantedToWifiAuth_revoked_returnsFalse()487     public void isKeyPairGrantedToWifiAuth_revoked_returnsFalse() {
488         try {
489             sDeviceState.dpc().devicePolicyManager().installKeyPair(
490                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
491                     RSA_ALIAS, /* requestAccess= */ false);
492             sDeviceState.dpc().devicePolicyManager().grantKeyPairToWifiAuth(RSA_ALIAS);
493             sDeviceState.dpc().devicePolicyManager().revokeKeyPairFromWifiAuth(RSA_ALIAS);
494 
495             assertThat(
496                     sDeviceState.dpc().devicePolicyManager().isKeyPairGrantedToWifiAuth(RSA_ALIAS))
497                     .isFalse();
498         } finally {
499             // Remove keypair
500             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
501         }
502     }
503 
504     @Test
505     @Postsubmit(reason = "new test")
506     @PositivePolicyTest(policy = KeyManagement.class)
grantKeyPair_keyUsable()507     public void grantKeyPair_keyUsable() throws Exception {
508         try {
509             sDeviceState.dpc().devicePolicyManager().installKeyPair(
510                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
511                     RSA_ALIAS, /* requestAccess= */ false);
512             sDeviceState.dpc().devicePolicyManager().grantKeyPairToApp(
513                     DPC_COMPONENT_NAME, RSA_ALIAS, sContext.getPackageName()
514             );
515 
516             PrivateKey key = KeyChain.getPrivateKey(sContext, RSA_ALIAS);
517 
518             signDataWithKey("SHA256withRSA", key); // Doesn't throw exception
519         } finally {
520             // Remove keypair
521             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
522         }
523     }
524 
525     @Test
526     @Postsubmit(reason = "new test")
527     @PositivePolicyTest(policy = KeyManagement.class)
revokeKeyPairFromApp_keyNotUsable()528     public void revokeKeyPairFromApp_keyNotUsable() throws Exception {
529         try {
530             sDeviceState.dpc().devicePolicyManager().installKeyPair(
531                     DPC_COMPONENT_NAME, PRIVATE_KEY, CERTIFICATES,
532                     RSA_ALIAS, /* requestAccess= */ false);
533             sDeviceState.dpc().devicePolicyManager().grantKeyPairToApp(
534                     DPC_COMPONENT_NAME, RSA_ALIAS, sContext.getPackageName()
535             );
536             // Key is requested from KeyChain prior to revoking the grant.
537             PrivateKey key = KeyChain.getPrivateKey(sContext, RSA_ALIAS);
538 
539             sDeviceState.dpc().devicePolicyManager().revokeKeyPairFromApp(
540                     DPC_COMPONENT_NAME, RSA_ALIAS, sContext.getPackageName());
541 
542             // Key shouldn't be valid after the grant is revoked.
543             Assert.assertThrows(
544                     InvalidKeyException.class, () -> signDataWithKey("SHA256withRSA", key));
545         } finally {
546             // Remove keypair
547             sDeviceState.dpc().devicePolicyManager().removeKeyPair(DPC_COMPONENT_NAME, RSA_ALIAS);
548         }
549     }
550 
signDataWithKey(String algoIdentifier, PrivateKey privateKey)551     private byte[] signDataWithKey(String algoIdentifier, PrivateKey privateKey) throws Exception {
552         byte[] data = "hello".getBytes();
553         Signature sign = Signature.getInstance(algoIdentifier);
554         sign.initSign(privateKey);
555         sign.update(data);
556         return sign.sign();
557     }
558 
559     private static class KeyChainAliasCallback extends BlockingCallback<String> implements
560             android.security.KeyChainAliasCallback {
561 
562         @Override
alias(final String chosenAlias)563         public void alias(final String chosenAlias) {
564             callbackTriggered(chosenAlias);
565         }
566     }
567 }
568