1 /*
2  * Copyright (C) 2020 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.server.locksettings;
18 
19 import static android.content.pm.UserInfo.FLAG_FULL;
20 import static android.content.pm.UserInfo.FLAG_PRIMARY;
21 import static android.content.pm.UserInfo.FLAG_PROFILE;
22 import static android.os.UserHandle.USER_SYSTEM;
23 
24 import static com.android.internal.widget.LockSettingsInternal.ARM_REBOOT_ERROR_ESCROW_NOT_READY;
25 import static com.android.internal.widget.LockSettingsInternal.ARM_REBOOT_ERROR_NONE;
26 import static com.android.internal.widget.LockSettingsInternal.ARM_REBOOT_ERROR_PROVIDER_MISMATCH;
27 import static com.android.internal.widget.LockSettingsInternal.ARM_REBOOT_ERROR_STORE_ESCROW_KEY;
28 
29 import static org.junit.Assert.assertEquals;
30 import static org.junit.Assert.assertFalse;
31 import static org.junit.Assert.assertNotNull;
32 import static org.junit.Assert.assertTrue;
33 import static org.mockito.ArgumentMatchers.any;
34 import static org.mockito.ArgumentMatchers.anyBoolean;
35 import static org.mockito.ArgumentMatchers.anyInt;
36 import static org.mockito.ArgumentMatchers.anyLong;
37 import static org.mockito.ArgumentMatchers.eq;
38 import static org.mockito.Mockito.clearInvocations;
39 import static org.mockito.Mockito.doNothing;
40 import static org.mockito.Mockito.doThrow;
41 import static org.mockito.Mockito.mock;
42 import static org.mockito.Mockito.never;
43 import static org.mockito.Mockito.times;
44 import static org.mockito.Mockito.verify;
45 import static org.mockito.Mockito.verifyNoMoreInteractions;
46 import static org.mockito.Mockito.when;
47 import static org.testng.Assert.assertNull;
48 
49 import android.content.Context;
50 import android.content.ContextWrapper;
51 import android.content.pm.UserInfo;
52 import android.hardware.rebootescrow.IRebootEscrow;
53 import android.net.ConnectivityManager;
54 import android.net.Network;
55 import android.os.Handler;
56 import android.os.HandlerThread;
57 import android.os.RemoteException;
58 import android.os.ServiceSpecificException;
59 import android.os.UserManager;
60 import android.platform.test.annotations.Presubmit;
61 
62 import androidx.test.InstrumentationRegistry;
63 import androidx.test.filters.SmallTest;
64 import androidx.test.runner.AndroidJUnit4;
65 
66 import com.android.internal.widget.RebootEscrowListener;
67 import com.android.server.locksettings.ResumeOnRebootServiceProvider.ResumeOnRebootServiceConnection;
68 
69 import org.junit.Before;
70 import org.junit.Test;
71 import org.junit.runner.RunWith;
72 import org.mockito.ArgumentCaptor;
73 
74 import java.io.File;
75 import java.io.IOException;
76 import java.util.ArrayList;
77 import java.util.concurrent.CountDownLatch;
78 import java.util.concurrent.TimeUnit;
79 import java.util.function.Consumer;
80 
81 import javax.crypto.SecretKey;
82 import javax.crypto.spec.SecretKeySpec;
83 
84 @SmallTest
85 @Presubmit
86 @RunWith(AndroidJUnit4.class)
87 public class RebootEscrowManagerTests {
88     protected static final int PRIMARY_USER_ID = 0;
89     protected static final int WORK_PROFILE_USER_ID = 10;
90     protected static final int NONSECURE_SECONDARY_USER_ID = 20;
91     protected static final int SECURE_SECONDARY_USER_ID = 21;
92     private static final byte FAKE_SP_VERSION = 1;
93     private static final byte[] FAKE_AUTH_TOKEN = new byte[] {
94             0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
95             0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
96             0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
97             0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
98     };
99 
100     // Hex encoding of a randomly generated AES key for test.
101     private static final byte[] TEST_AES_KEY = new byte[] {
102             0x48, 0x19, 0x12, 0x54, 0x13, 0x13, 0x52, 0x31,
103             0x44, 0x74, 0x61, 0x54, 0x29, 0x74, 0x37, 0x61,
104             0x70, 0x70, 0x75, 0x25, 0x27, 0x31, 0x49, 0x09,
105             0x26, 0x52, 0x72, 0x63, 0x63, 0x61, 0x78, 0x23,
106     };
107 
108     private Context mContext;
109     private UserManager mUserManager;
110     private RebootEscrowManager.Callbacks mCallbacks;
111     private IRebootEscrow mRebootEscrow;
112     private ResumeOnRebootServiceConnection mServiceConnection;
113     private RebootEscrowKeyStoreManager mKeyStoreManager;
114 
115     LockSettingsStorageTestable mStorage;
116 
117     private MockableRebootEscrowInjected mInjected;
118     private RebootEscrowManager mService;
119     private SecretKey mAesKey;
120     private MockInjector mMockInjector;
121     private Handler mHandler;
122 
123     public interface MockableRebootEscrowInjected {
getBootCount()124         int getBootCount();
125 
getCurrentTimeMillis()126         long getCurrentTimeMillis();
127 
reportMetric(boolean success, int errorCode, int serviceType, int attemptCount, int escrowDurationInSeconds, int vbmetaDigestStatus, int durationSinceBootComplete)128         void reportMetric(boolean success, int errorCode, int serviceType, int attemptCount,
129                 int escrowDurationInSeconds, int vbmetaDigestStatus, int durationSinceBootComplete);
130     }
131 
132     static class MockInjector extends RebootEscrowManager.Injector {
133         private final IRebootEscrow mRebootEscrow;
134         private final RebootEscrowProviderInterface mDefaultRebootEscrowProvider;
135         private final UserManager mUserManager;
136         private final MockableRebootEscrowInjected mInjected;
137         private final RebootEscrowKeyStoreManager mKeyStoreManager;
138         private boolean mServerBased;
139         private RebootEscrowProviderInterface mRebootEscrowProviderInUse;
140         private ConnectivityManager.NetworkCallback mNetworkCallback;
141         private Consumer<ConnectivityManager.NetworkCallback> mNetworkConsumer;
142         private boolean mWaitForInternet;
143 
MockInjector(Context context, UserManager userManager, IRebootEscrow rebootEscrow, RebootEscrowKeyStoreManager keyStoreManager, LockSettingsStorageTestable storage, MockableRebootEscrowInjected injected)144         MockInjector(Context context, UserManager userManager,
145                 IRebootEscrow rebootEscrow,
146                 RebootEscrowKeyStoreManager keyStoreManager,
147                 LockSettingsStorageTestable storage,
148                 MockableRebootEscrowInjected injected) {
149             super(context, storage);
150             mRebootEscrow = rebootEscrow;
151             mServerBased = false;
152             mWaitForInternet = false;
153             RebootEscrowProviderHalImpl.Injector halInjector =
154                     new RebootEscrowProviderHalImpl.Injector() {
155                         @Override
156                         public IRebootEscrow getRebootEscrow() {
157                             return mRebootEscrow;
158                         }
159                     };
160             mDefaultRebootEscrowProvider = new RebootEscrowProviderHalImpl(halInjector);
161             mUserManager = userManager;
162             mKeyStoreManager = keyStoreManager;
163             mInjected = injected;
164         }
165 
MockInjector(Context context, UserManager userManager, ResumeOnRebootServiceConnection serviceConnection, RebootEscrowKeyStoreManager keyStoreManager, LockSettingsStorageTestable storage, MockableRebootEscrowInjected injected)166         MockInjector(Context context, UserManager userManager,
167                 ResumeOnRebootServiceConnection serviceConnection,
168                 RebootEscrowKeyStoreManager keyStoreManager,
169                 LockSettingsStorageTestable storage,
170                 MockableRebootEscrowInjected injected) {
171             super(context, storage);
172             mRebootEscrow = null;
173             mServerBased = true;
174             mWaitForInternet = false;
175             RebootEscrowProviderServerBasedImpl.Injector injector =
176                     new RebootEscrowProviderServerBasedImpl.Injector(serviceConnection) {
177                         @Override
178                         long getServiceTimeoutInSeconds() {
179                             return 30;
180                         }
181 
182                         @Override
183                         long getServerBlobLifetimeInMillis() {
184                             return 600_000;
185                         }
186                     };
187             mDefaultRebootEscrowProvider = new RebootEscrowProviderServerBasedImpl(
188                     storage, injector);
189             mUserManager = userManager;
190             mKeyStoreManager = keyStoreManager;
191             mInjected = injected;
192         }
193 
194         @Override
post(Handler handler, Runnable runnable)195         void post(Handler handler, Runnable runnable) {
196             runnable.run();
197         }
198 
199         @Override
getUserManager()200         public UserManager getUserManager() {
201             return mUserManager;
202         }
203 
204         @Override
serverBasedResumeOnReboot()205         public boolean serverBasedResumeOnReboot() {
206             return mServerBased;
207         }
208 
209         @Override
waitForInternet()210         public boolean waitForInternet() {
211             return mWaitForInternet;
212         }
213 
setWaitForNetwork(boolean waitForNetworkEnabled)214         public void setWaitForNetwork(boolean waitForNetworkEnabled) {
215             mWaitForInternet = waitForNetworkEnabled;
216         }
217 
218         @Override
isNetworkConnected()219         public boolean isNetworkConnected() {
220             return false;
221         }
222 
223         @Override
requestNetworkWithInternet( ConnectivityManager.NetworkCallback networkCallback)224         public boolean requestNetworkWithInternet(
225                 ConnectivityManager.NetworkCallback networkCallback) {
226             mNetworkCallback = networkCallback;
227             mNetworkConsumer.accept(networkCallback);
228             return true;
229         }
230 
231         @Override
stopRequestingNetwork(ConnectivityManager.NetworkCallback networkCallback)232         public void stopRequestingNetwork(ConnectivityManager.NetworkCallback networkCallback) {
233             mNetworkCallback = null;
234         }
235 
236         @Override
createRebootEscrowProviderIfNeeded()237         public RebootEscrowProviderInterface createRebootEscrowProviderIfNeeded() {
238             mRebootEscrowProviderInUse = mDefaultRebootEscrowProvider;
239             return mRebootEscrowProviderInUse;
240         }
241 
242         @Override
getRebootEscrowProvider()243         public RebootEscrowProviderInterface getRebootEscrowProvider() {
244             return mRebootEscrowProviderInUse;
245         }
246 
247         @Override
clearRebootEscrowProvider()248         public void clearRebootEscrowProvider() {
249             mRebootEscrowProviderInUse = null;
250         }
251 
252         @Override
getKeyStoreManager()253         public RebootEscrowKeyStoreManager getKeyStoreManager() {
254             return mKeyStoreManager;
255         }
256 
257         @Override
getBootCount()258         public int getBootCount() {
259             return mInjected.getBootCount();
260         }
261 
262         @Override
getLoadEscrowDataRetryLimit()263         public int getLoadEscrowDataRetryLimit() {
264             // Try two times
265             return 2;
266         }
267 
268         @Override
getLoadEscrowDataRetryIntervalSeconds()269         public int getLoadEscrowDataRetryIntervalSeconds() {
270             // Retry in 1 seconds
271             return 1;
272         }
273 
274         @Override
getLoadEscrowTimeoutMillis()275         public int getLoadEscrowTimeoutMillis() {
276             // Timeout in 3 seconds.
277             return 3000;
278         }
279 
280         @Override
getVbmetaDigest(boolean other)281         public String getVbmetaDigest(boolean other) {
282             return other ? "" : "fake digest";
283         }
284 
285         @Override
getCurrentTimeMillis()286         public long getCurrentTimeMillis() {
287             return mInjected.getCurrentTimeMillis();
288         }
289 
290         @Override
reportMetric(boolean success, int errorCode, int serviceType, int attemptCount, int escrowDurationInSeconds, int vbmetaDigestStatus, int durationSinceBootComplete)291         public void reportMetric(boolean success, int errorCode, int serviceType, int attemptCount,
292                 int escrowDurationInSeconds, int vbmetaDigestStatus,
293                 int durationSinceBootComplete) {
294 
295             mInjected.reportMetric(success, errorCode, serviceType, attemptCount,
296                     escrowDurationInSeconds, vbmetaDigestStatus, durationSinceBootComplete);
297         }
298     }
299 
300     @Before
setUp_baseServices()301     public void setUp_baseServices() throws Exception {
302         mContext = new ContextWrapper(InstrumentationRegistry.getContext());
303         mUserManager = mock(UserManager.class);
304         mCallbacks = mock(RebootEscrowManager.Callbacks.class);
305         mRebootEscrow = mock(IRebootEscrow.class);
306         mServiceConnection = mock(ResumeOnRebootServiceConnection.class);
307         mKeyStoreManager = mock(RebootEscrowKeyStoreManager.class);
308         mAesKey = new SecretKeySpec(TEST_AES_KEY, "AES");
309 
310         when(mKeyStoreManager.getKeyStoreEncryptionKey()).thenReturn(mAesKey);
311         when(mKeyStoreManager.generateKeyStoreEncryptionKeyIfNeeded()).thenReturn(mAesKey);
312 
313         mStorage = new LockSettingsStorageTestable(mContext,
314                 new File(InstrumentationRegistry.getContext().getFilesDir(), "locksettings"));
315 
316         ArrayList<UserInfo> users = new ArrayList<>();
317         users.add(new UserInfo(PRIMARY_USER_ID, "primary", FLAG_PRIMARY));
318         users.add(new UserInfo(WORK_PROFILE_USER_ID, "work", FLAG_PROFILE));
319         users.add(new UserInfo(NONSECURE_SECONDARY_USER_ID, "non-secure", FLAG_FULL));
320         users.add(new UserInfo(SECURE_SECONDARY_USER_ID, "secure", FLAG_FULL));
321         when(mUserManager.getUsers()).thenReturn(users);
322         when(mCallbacks.isUserSecure(PRIMARY_USER_ID)).thenReturn(true);
323         when(mCallbacks.isUserSecure(WORK_PROFILE_USER_ID)).thenReturn(true);
324         when(mCallbacks.isUserSecure(NONSECURE_SECONDARY_USER_ID)).thenReturn(false);
325         when(mCallbacks.isUserSecure(SECURE_SECONDARY_USER_ID)).thenReturn(true);
326         mInjected = mock(MockableRebootEscrowInjected.class);
327         mMockInjector = new MockInjector(mContext, mUserManager, mRebootEscrow,
328                 mKeyStoreManager, mStorage, mInjected);
329         HandlerThread thread = new HandlerThread("RebootEscrowManagerTest");
330         thread.start();
331         mHandler = new Handler(thread.getLooper());
332         mService = new RebootEscrowManager(mMockInjector, mCallbacks, mStorage, mHandler);
333 
334     }
335 
setServerBasedRebootEscrowProvider()336     private void setServerBasedRebootEscrowProvider() throws Exception {
337         mMockInjector = new MockInjector(mContext, mUserManager, mServiceConnection,
338                 mKeyStoreManager, mStorage, mInjected);
339         mService = new RebootEscrowManager(mMockInjector, mCallbacks, mStorage, mHandler);
340     }
341 
waitForHandler()342     private void waitForHandler() throws InterruptedException {
343         // Wait for handler to complete processing.
344         CountDownLatch latch = new CountDownLatch(1);
345         mHandler.post(latch::countDown);
346         assertTrue(latch.await(5, TimeUnit.SECONDS));
347 
348     }
349 
callToRebootEscrowIfNeededAndWait(int userId)350     private void callToRebootEscrowIfNeededAndWait(int userId) throws InterruptedException {
351         mService.callToRebootEscrowIfNeeded(userId, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
352         waitForHandler();
353     }
354 
355     @Test
prepareRebootEscrow_Success()356     public void prepareRebootEscrow_Success() throws Exception {
357         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
358         mService.setRebootEscrowListener(mockListener);
359         mService.prepareRebootEscrow();
360 
361         clearInvocations(mRebootEscrow);
362         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
363         verify(mockListener).onPreparedForReboot(eq(true));
364         verify(mRebootEscrow, never()).storeKey(any());
365     }
366 
367     @Test
prepareRebootEscrowServerBased_Success()368     public void prepareRebootEscrowServerBased_Success() throws Exception {
369         setServerBasedRebootEscrowProvider();
370         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
371         mService.setRebootEscrowListener(mockListener);
372         mService.prepareRebootEscrow();
373 
374         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
375         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
376         assertFalse(mStorage.hasRebootEscrowServerBlob());
377     }
378 
379     @Test
prepareRebootEscrow_ClearCredentials_Success()380     public void prepareRebootEscrow_ClearCredentials_Success() throws Exception {
381         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
382         mService.setRebootEscrowListener(mockListener);
383         mService.prepareRebootEscrow();
384         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
385         verify(mockListener).onPreparedForReboot(eq(true));
386 
387         clearInvocations(mRebootEscrow);
388         mService.clearRebootEscrow();
389         verify(mockListener).onPreparedForReboot(eq(false));
390         verify(mRebootEscrow).storeKey(eq(new byte[32]));
391     }
392 
393     @Test
clearCredentials_HalFailure_NonFatal()394     public void clearCredentials_HalFailure_NonFatal() throws Exception {
395         doThrow(ServiceSpecificException.class).when(mRebootEscrow).storeKey(any());
396         mService.clearRebootEscrow();
397         verify(mRebootEscrow).storeKey(eq(new byte[32]));
398         assertNull(mMockInjector.getRebootEscrowProvider());
399     }
400 
401     @Test
armService_Success()402     public void armService_Success() throws Exception {
403         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
404         mService.setRebootEscrowListener(mockListener);
405         mService.prepareRebootEscrow();
406 
407         clearInvocations(mRebootEscrow);
408         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
409         verify(mockListener).onPreparedForReboot(eq(true));
410         verify(mRebootEscrow, never()).storeKey(any());
411 
412         assertNull(
413                 mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, null, USER_SYSTEM));
414         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
415         assertNotNull(
416                 mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, null, USER_SYSTEM));
417         verify(mRebootEscrow).storeKey(any());
418         verify(mKeyStoreManager).getKeyStoreEncryptionKey();
419 
420         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
421         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
422     }
423 
424     @Test
armServiceServerBased_Success()425     public void armServiceServerBased_Success() throws Exception {
426         setServerBasedRebootEscrowProvider();
427         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
428         mService.setRebootEscrowListener(mockListener);
429         mService.prepareRebootEscrow();
430 
431         clearInvocations(mServiceConnection);
432         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
433         verify(mockListener).onPreparedForReboot(eq(true));
434         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
435 
436         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
437                 .thenAnswer(invocation -> invocation.getArgument(0));
438         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
439         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
440 
441         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
442         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
443         assertTrue(mStorage.hasRebootEscrowServerBlob());
444     }
445 
446     @Test
armService_HalFailure_NonFatal()447     public void armService_HalFailure_NonFatal() throws Exception {
448         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
449         mService.setRebootEscrowListener(mockListener);
450         mService.prepareRebootEscrow();
451 
452         clearInvocations(mRebootEscrow);
453         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
454         verify(mockListener).onPreparedForReboot(eq(true));
455         verify(mRebootEscrow, never()).storeKey(any());
456 
457         assertNull(
458                 mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, null, USER_SYSTEM));
459         doThrow(ServiceSpecificException.class).when(mRebootEscrow).storeKey(any());
460         assertEquals(ARM_REBOOT_ERROR_STORE_ESCROW_KEY, mService.armRebootEscrowIfNeeded());
461         verify(mRebootEscrow).storeKey(any());
462     }
463 
464     @Test
armService_MultipleUsers_Success()465     public void armService_MultipleUsers_Success() throws Exception {
466         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
467         mService.setRebootEscrowListener(mockListener);
468         mService.prepareRebootEscrow();
469 
470         clearInvocations(mRebootEscrow);
471         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
472         verify(mockListener).onPreparedForReboot(eq(true));
473         callToRebootEscrowIfNeededAndWait(SECURE_SECONDARY_USER_ID);
474         verify(mRebootEscrow, never()).storeKey(any());
475 
476         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
477         assertTrue(mStorage.hasRebootEscrow(SECURE_SECONDARY_USER_ID));
478         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
479 
480         assertNull(
481                 mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, null, USER_SYSTEM));
482         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
483         assertNotNull(
484                 mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, null, USER_SYSTEM));
485         verify(mRebootEscrow, times(1)).storeKey(any());
486 
487         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
488         assertTrue(mStorage.hasRebootEscrow(SECURE_SECONDARY_USER_ID));
489         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
490     }
491 
492     @Test
armService_NoInitialization_Failure()493     public void armService_NoInitialization_Failure() throws Exception {
494         assertEquals(ARM_REBOOT_ERROR_ESCROW_NOT_READY, mService.armRebootEscrowIfNeeded());
495         verifyNoMoreInteractions(mRebootEscrow);
496     }
497 
498     @Test
armService_RebootEscrowServiceException_Failure()499     public void armService_RebootEscrowServiceException_Failure() throws Exception {
500         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
501         mService.setRebootEscrowListener(mockListener);
502         mService.prepareRebootEscrow();
503 
504         clearInvocations(mRebootEscrow);
505         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
506         verify(mockListener).onPreparedForReboot(eq(true));
507         verify(mRebootEscrow, never()).storeKey(any());
508 
509         doThrow(RemoteException.class).when(mRebootEscrow).storeKey(any());
510         assertEquals(ARM_REBOOT_ERROR_STORE_ESCROW_KEY, mService.armRebootEscrowIfNeeded());
511         verify(mRebootEscrow).storeKey(any());
512     }
513 
514     @Test
loadRebootEscrowDataIfAvailable_NothingAvailable_Success()515     public void loadRebootEscrowDataIfAvailable_NothingAvailable_Success() throws Exception {
516         mService.loadRebootEscrowDataIfAvailable(mHandler);
517     }
518 
519     @Test
loadRebootEscrowDataIfAvailable_Success()520     public void loadRebootEscrowDataIfAvailable_Success() throws Exception {
521         when(mInjected.getBootCount()).thenReturn(0);
522 
523         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
524         mService.setRebootEscrowListener(mockListener);
525         mService.prepareRebootEscrow();
526 
527         clearInvocations(mRebootEscrow);
528         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
529         verify(mockListener).onPreparedForReboot(eq(true));
530 
531         verify(mRebootEscrow, never()).storeKey(any());
532 
533         ArgumentCaptor<byte[]> keyByteCaptor = ArgumentCaptor.forClass(byte[].class);
534         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
535         verify(mRebootEscrow).storeKey(keyByteCaptor.capture());
536         verify(mKeyStoreManager).getKeyStoreEncryptionKey();
537 
538         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
539         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
540 
541         // pretend reboot happens here
542 
543         when(mInjected.getBootCount()).thenReturn(1);
544         when(mInjected.getCurrentTimeMillis()).thenReturn(30000L);
545         mStorage.setLong(RebootEscrowManager.REBOOT_ESCROW_KEY_ARMED_TIMESTAMP, 10000L,
546                 USER_SYSTEM);
547         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
548         doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture(),
549                 eq(0) /* error code */, eq(1) /* HAL based */, eq(1) /* attempt count */,
550                 eq(20), eq(0) /* vbmeta status */, anyInt());
551         when(mRebootEscrow.retrieveKey()).thenAnswer(invocation -> keyByteCaptor.getValue());
552 
553         mService.loadRebootEscrowDataIfAvailable(mHandler);
554         verify(mRebootEscrow).retrieveKey();
555         assertTrue(metricsSuccessCaptor.getValue());
556         verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
557         assertEquals(mStorage.getLong(RebootEscrowManager.REBOOT_ESCROW_KEY_ARMED_TIMESTAMP,
558                 -1, USER_SYSTEM), -1);
559     }
560 
561     @Test
loadRebootEscrowDataIfAvailable_ServerBased_Success()562     public void loadRebootEscrowDataIfAvailable_ServerBased_Success() throws Exception {
563         setServerBasedRebootEscrowProvider();
564 
565         when(mInjected.getBootCount()).thenReturn(0);
566         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
567         mService.setRebootEscrowListener(mockListener);
568         mService.prepareRebootEscrow();
569 
570         clearInvocations(mServiceConnection);
571         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
572         verify(mockListener).onPreparedForReboot(eq(true));
573         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
574 
575         // Use x -> x for both wrap & unwrap functions.
576         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
577                 .thenAnswer(invocation -> invocation.getArgument(0));
578         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
579         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
580         assertTrue(mStorage.hasRebootEscrowServerBlob());
581 
582         // pretend reboot happens here
583         when(mInjected.getBootCount()).thenReturn(1);
584         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
585         doNothing()
586                 .when(mInjected)
587                 .reportMetric(
588                         metricsSuccessCaptor.capture(),
589                         eq(0) /* error code */,
590                         eq(2) /* Server based */,
591                         eq(1) /* attempt count */,
592                         anyInt(),
593                         eq(0) /* vbmeta status */,
594                         anyInt());
595 
596         when(mServiceConnection.unwrap(any(), anyLong()))
597                 .thenAnswer(invocation -> invocation.getArgument(0));
598         mService.loadRebootEscrowDataIfAvailable(null);
599         verify(mServiceConnection).unwrap(any(), anyLong());
600         assertTrue(metricsSuccessCaptor.getValue());
601         verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
602     }
603 
604     @Test
loadRebootEscrowDataIfAvailable_ServerBasedRemoteException_Failure()605     public void loadRebootEscrowDataIfAvailable_ServerBasedRemoteException_Failure()
606             throws Exception {
607         setServerBasedRebootEscrowProvider();
608 
609         when(mInjected.getBootCount()).thenReturn(0);
610         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
611         mService.setRebootEscrowListener(mockListener);
612         mService.prepareRebootEscrow();
613 
614         clearInvocations(mServiceConnection);
615         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
616         verify(mockListener).onPreparedForReboot(eq(true));
617         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
618 
619         // Use x -> x for both wrap & unwrap functions.
620         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
621                 .thenAnswer(invocation -> invocation.getArgument(0));
622         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
623         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
624         assertTrue(mStorage.hasRebootEscrowServerBlob());
625 
626         // pretend reboot happens here
627         when(mInjected.getBootCount()).thenReturn(1);
628         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
629         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
630         doNothing()
631                 .when(mInjected)
632                 .reportMetric(
633                         metricsSuccessCaptor.capture(),
634                         metricsErrorCodeCaptor.capture(),
635                         eq(2) /* Server based */,
636                         eq(1) /* attempt count */,
637                         anyInt(),
638                         eq(0) /* vbmeta status */,
639                         anyInt());
640 
641         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(RemoteException.class);
642         mService.loadRebootEscrowDataIfAvailable(null);
643         verify(mServiceConnection).unwrap(any(), anyLong());
644         assertFalse(metricsSuccessCaptor.getValue());
645         assertEquals(
646                 Integer.valueOf(RebootEscrowManager.ERROR_LOAD_ESCROW_KEY),
647                 metricsErrorCodeCaptor.getValue());
648     }
649 
650     @Test
loadRebootEscrowDataIfAvailable_ServerBasedIoError_RetryFailure()651     public void loadRebootEscrowDataIfAvailable_ServerBasedIoError_RetryFailure() throws Exception {
652         setServerBasedRebootEscrowProvider();
653 
654         when(mInjected.getBootCount()).thenReturn(0);
655         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
656         mService.setRebootEscrowListener(mockListener);
657         mService.prepareRebootEscrow();
658 
659         clearInvocations(mServiceConnection);
660         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
661         verify(mockListener).onPreparedForReboot(eq(true));
662         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
663 
664         // Use x -> x for both wrap & unwrap functions.
665         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
666                 .thenAnswer(invocation -> invocation.getArgument(0));
667         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
668         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
669         assertTrue(mStorage.hasRebootEscrowServerBlob());
670 
671         // pretend reboot happens here
672         when(mInjected.getBootCount()).thenReturn(1);
673         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
674         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
675         doNothing()
676                 .when(mInjected)
677                 .reportMetric(
678                         metricsSuccessCaptor.capture(),
679                         metricsErrorCodeCaptor.capture(),
680                         eq(2) /* Server based */,
681                         eq(2) /* attempt count */,
682                         anyInt(),
683                         eq(0) /* vbmeta status */,
684                         anyInt());
685         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(IOException.class);
686 
687         mService.loadRebootEscrowDataIfAvailable(mHandler);
688         // Sleep 5s for the retry to complete
689         Thread.sleep(5 * 1000);
690         assertFalse(metricsSuccessCaptor.getValue());
691         assertEquals(
692                 Integer.valueOf(RebootEscrowManager.ERROR_NO_NETWORK),
693                 metricsErrorCodeCaptor.getValue());
694     }
695 
696     @Test
loadRebootEscrowDataIfAvailable_ServerBased_RetrySuccess()697     public void loadRebootEscrowDataIfAvailable_ServerBased_RetrySuccess() throws Exception {
698         setServerBasedRebootEscrowProvider();
699 
700         when(mInjected.getBootCount()).thenReturn(0);
701         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
702         mService.setRebootEscrowListener(mockListener);
703         mService.prepareRebootEscrow();
704 
705         clearInvocations(mServiceConnection);
706         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
707         verify(mockListener).onPreparedForReboot(eq(true));
708         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
709 
710         // Use x -> x for both wrap & unwrap functions.
711         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
712                 .thenAnswer(invocation -> invocation.getArgument(0));
713         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
714         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
715         assertTrue(mStorage.hasRebootEscrowServerBlob());
716 
717         // pretend reboot happens here
718         when(mInjected.getBootCount()).thenReturn(1);
719         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
720         doNothing()
721                 .when(mInjected)
722                 .reportMetric(
723                         metricsSuccessCaptor.capture(),
724                         anyInt(),
725                         anyInt(),
726                         eq(2) /* attempt count */,
727                         anyInt(),
728                         anyInt(),
729                         anyInt());
730 
731         when(mServiceConnection.unwrap(any(), anyLong()))
732                 .thenThrow(new IOException())
733                 .thenAnswer(invocation -> invocation.getArgument(0));
734 
735         mService.loadRebootEscrowDataIfAvailable(mHandler);
736         // Sleep 5s for the retry to complete
737         Thread.sleep(5 * 1000);
738         verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
739         assertTrue(metricsSuccessCaptor.getValue());
740         verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
741     }
742 
743     @Test
loadRebootEscrowDataIfAvailable_serverBasedWaitForInternet_success()744     public void loadRebootEscrowDataIfAvailable_serverBasedWaitForInternet_success()
745             throws Exception {
746         setServerBasedRebootEscrowProvider();
747         mMockInjector.setWaitForNetwork(true);
748 
749         when(mInjected.getBootCount()).thenReturn(0);
750         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
751         mService.setRebootEscrowListener(mockListener);
752         mService.prepareRebootEscrow();
753 
754         clearInvocations(mServiceConnection);
755         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
756         verify(mockListener).onPreparedForReboot(eq(true));
757         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
758 
759         // Use x -> x for both wrap & unwrap functions.
760         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
761                 .thenAnswer(invocation -> invocation.getArgument(0));
762         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
763         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
764         assertTrue(mStorage.hasRebootEscrowServerBlob());
765 
766         // pretend reboot happens here
767         when(mInjected.getBootCount()).thenReturn(1);
768         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
769         doNothing()
770                 .when(mInjected)
771                 .reportMetric(
772                         metricsSuccessCaptor.capture(),
773                         eq(0) /* error code */,
774                         eq(2) /* Server based */,
775                         eq(1) /* attempt count */,
776                         anyInt(),
777                         eq(0) /* vbmeta status */,
778                         anyInt());
779 
780         // load escrow data
781         when(mServiceConnection.unwrap(any(), anyLong()))
782                 .thenAnswer(invocation -> invocation.getArgument(0));
783         Network mockNetwork = mock(Network.class);
784         mMockInjector.mNetworkConsumer =
785                 (callback) -> {
786                     callback.onAvailable(mockNetwork);
787                 };
788 
789         mService.loadRebootEscrowDataIfAvailable(mHandler);
790         verify(mServiceConnection).unwrap(any(), anyLong());
791         assertTrue(metricsSuccessCaptor.getValue());
792         verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
793         assertNull(mMockInjector.mNetworkCallback);
794     }
795 
796     @Test
loadRebootEscrowDataIfAvailable_serverBasedWaitForInternetRemoteException_Failure()797     public void loadRebootEscrowDataIfAvailable_serverBasedWaitForInternetRemoteException_Failure()
798             throws Exception {
799         setServerBasedRebootEscrowProvider();
800         mMockInjector.setWaitForNetwork(true);
801 
802         when(mInjected.getBootCount()).thenReturn(0);
803         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
804         mService.setRebootEscrowListener(mockListener);
805         mService.prepareRebootEscrow();
806 
807         clearInvocations(mServiceConnection);
808         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
809         verify(mockListener).onPreparedForReboot(eq(true));
810         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
811 
812         // Use x -> x for both wrap & unwrap functions.
813         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
814                 .thenAnswer(invocation -> invocation.getArgument(0));
815         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
816         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
817         assertTrue(mStorage.hasRebootEscrowServerBlob());
818 
819         // pretend reboot happens here
820         when(mInjected.getBootCount()).thenReturn(1);
821         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
822         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
823         doNothing()
824                 .when(mInjected)
825                 .reportMetric(
826                         metricsSuccessCaptor.capture(),
827                         metricsErrorCodeCaptor.capture(),
828                         eq(2) /* Server based */,
829                         eq(1) /* attempt count */,
830                         anyInt(),
831                         eq(0) /* vbmeta status */,
832                         anyInt());
833 
834         // load escrow data
835         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(RemoteException.class);
836         Network mockNetwork = mock(Network.class);
837         mMockInjector.mNetworkConsumer =
838                 (callback) -> {
839                     callback.onAvailable(mockNetwork);
840                 };
841 
842         mService.loadRebootEscrowDataIfAvailable(mHandler);
843         verify(mServiceConnection).unwrap(any(), anyLong());
844         assertFalse(metricsSuccessCaptor.getValue());
845         assertEquals(
846                 Integer.valueOf(RebootEscrowManager.ERROR_LOAD_ESCROW_KEY),
847                 metricsErrorCodeCaptor.getValue());
848         assertNull(mMockInjector.mNetworkCallback);
849     }
850 
851     @Test
loadRebootEscrowDataIfAvailable_waitForInternet_networkUnavailable()852     public void loadRebootEscrowDataIfAvailable_waitForInternet_networkUnavailable()
853             throws Exception {
854         setServerBasedRebootEscrowProvider();
855         mMockInjector.setWaitForNetwork(true);
856 
857         when(mInjected.getBootCount()).thenReturn(0);
858         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
859         mService.setRebootEscrowListener(mockListener);
860         mService.prepareRebootEscrow();
861 
862         clearInvocations(mServiceConnection);
863         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
864         verify(mockListener).onPreparedForReboot(eq(true));
865         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
866 
867         // Use x -> x for both wrap & unwrap functions.
868         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
869                 .thenAnswer(invocation -> invocation.getArgument(0));
870         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
871         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
872         assertTrue(mStorage.hasRebootEscrowServerBlob());
873 
874         // pretend reboot happens here
875         when(mInjected.getBootCount()).thenReturn(1);
876         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
877         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
878         doNothing()
879                 .when(mInjected)
880                 .reportMetric(
881                         metricsSuccessCaptor.capture(),
882                         metricsErrorCodeCaptor.capture(),
883                         eq(2) /* Server based */,
884                         eq(0) /* attempt count */,
885                         anyInt(),
886                         eq(0) /* vbmeta status */,
887                         anyInt());
888 
889         // Network is not available within timeout.
890         mMockInjector.mNetworkConsumer = ConnectivityManager.NetworkCallback::onUnavailable;
891         mService.loadRebootEscrowDataIfAvailable(mHandler);
892         assertFalse(metricsSuccessCaptor.getValue());
893         assertEquals(
894                 Integer.valueOf(RebootEscrowManager.ERROR_NO_NETWORK),
895                 metricsErrorCodeCaptor.getValue());
896         assertNull(mMockInjector.mNetworkCallback);
897     }
898 
899     @Test
loadRebootEscrowDataIfAvailable_waitForInternet_networkLost()900     public void loadRebootEscrowDataIfAvailable_waitForInternet_networkLost() throws Exception {
901         setServerBasedRebootEscrowProvider();
902         mMockInjector.setWaitForNetwork(true);
903 
904         when(mInjected.getBootCount()).thenReturn(0);
905         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
906         mService.setRebootEscrowListener(mockListener);
907         mService.prepareRebootEscrow();
908 
909         clearInvocations(mServiceConnection);
910         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
911         verify(mockListener).onPreparedForReboot(eq(true));
912         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
913 
914         // Use x -> x for both wrap & unwrap functions.
915         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
916                 .thenAnswer(invocation -> invocation.getArgument(0));
917         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
918         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
919         assertTrue(mStorage.hasRebootEscrowServerBlob());
920 
921         // pretend reboot happens here
922         when(mInjected.getBootCount()).thenReturn(1);
923         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
924         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
925         doNothing()
926                 .when(mInjected)
927                 .reportMetric(
928                         metricsSuccessCaptor.capture(),
929                         metricsErrorCodeCaptor.capture(),
930                         eq(2) /* Server based */,
931                         eq(2) /* attempt count */,
932                         anyInt(),
933                         eq(0) /* vbmeta status */,
934                         anyInt());
935 
936         // Network is available, then lost.
937         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(new IOException());
938         Network mockNetwork = mock(Network.class);
939         mMockInjector.mNetworkConsumer =
940                 (callback) -> {
941                     callback.onAvailable(mockNetwork);
942                     callback.onLost(mockNetwork);
943                 };
944         mService.loadRebootEscrowDataIfAvailable(mHandler);
945         // Sleep 5s for the retry to complete
946         Thread.sleep(5 * 1000);
947         assertFalse(metricsSuccessCaptor.getValue());
948         assertEquals(
949                 Integer.valueOf(RebootEscrowManager.ERROR_NO_NETWORK),
950                 metricsErrorCodeCaptor.getValue());
951         assertNull(mMockInjector.mNetworkCallback);
952     }
953 
954     @Test
loadRebootEscrowDataIfAvailable_waitForInternet_networkAvailableWithDelay()955     public void loadRebootEscrowDataIfAvailable_waitForInternet_networkAvailableWithDelay()
956             throws Exception {
957         setServerBasedRebootEscrowProvider();
958         mMockInjector.setWaitForNetwork(true);
959 
960         when(mInjected.getBootCount()).thenReturn(0);
961         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
962         mService.setRebootEscrowListener(mockListener);
963         mService.prepareRebootEscrow();
964 
965         clearInvocations(mServiceConnection);
966         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
967         verify(mockListener).onPreparedForReboot(eq(true));
968         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
969 
970         // Use x -> x for both wrap & unwrap functions.
971         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
972                 .thenAnswer(invocation -> invocation.getArgument(0));
973         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
974         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
975         assertTrue(mStorage.hasRebootEscrowServerBlob());
976 
977         // pretend reboot happens here
978         when(mInjected.getBootCount()).thenReturn(1);
979         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
980         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
981         doNothing()
982                 .when(mInjected)
983                 .reportMetric(
984                         metricsSuccessCaptor.capture(),
985                         metricsErrorCodeCaptor.capture(),
986                         eq(2) /* Server based */,
987                         eq(1) /* attempt count */,
988                         anyInt(),
989                         eq(0) /* vbmeta status */,
990                         anyInt());
991 
992         // load escrow data
993         when(mServiceConnection.unwrap(any(), anyLong()))
994                 .thenAnswer(invocation -> invocation.getArgument(0));
995         // network available after 1 sec
996         Network mockNetwork = mock(Network.class);
997         mMockInjector.mNetworkConsumer =
998                 (callback) -> {
999                     try {
1000                         Thread.sleep(1000);
1001                     } catch (InterruptedException e) {
1002                         throw new RuntimeException(e);
1003                     }
1004                     callback.onAvailable(mockNetwork);
1005                 };
1006         mService.loadRebootEscrowDataIfAvailable(mHandler);
1007         verify(mServiceConnection).unwrap(any(), anyLong());
1008         assertTrue(metricsSuccessCaptor.getValue());
1009         verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
1010         assertNull(mMockInjector.mNetworkCallback);
1011     }
1012 
1013     @Test
loadRebootEscrowDataIfAvailable_waitForInternet_timeoutExhausted()1014     public void loadRebootEscrowDataIfAvailable_waitForInternet_timeoutExhausted()
1015             throws Exception {
1016         setServerBasedRebootEscrowProvider();
1017         mMockInjector.setWaitForNetwork(true);
1018 
1019         when(mInjected.getBootCount()).thenReturn(0);
1020         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1021         mService.setRebootEscrowListener(mockListener);
1022         mService.prepareRebootEscrow();
1023 
1024         clearInvocations(mServiceConnection);
1025         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1026         verify(mockListener).onPreparedForReboot(eq(true));
1027         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
1028 
1029         // Use x -> x for both wrap & unwrap functions.
1030         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
1031                 .thenAnswer(invocation -> invocation.getArgument(0));
1032         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
1033         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
1034         assertTrue(mStorage.hasRebootEscrowServerBlob());
1035 
1036         // pretend reboot happens here
1037         when(mInjected.getBootCount()).thenReturn(1);
1038         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
1039         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
1040         doNothing()
1041                 .when(mInjected)
1042                 .reportMetric(
1043                         metricsSuccessCaptor.capture(),
1044                         metricsErrorCodeCaptor.capture(),
1045                         eq(2) /* Server based */,
1046                         eq(1) /* attempt count */,
1047                         anyInt(),
1048                         eq(0) /* vbmeta status */,
1049                         anyInt());
1050 
1051         // load reboot escrow data
1052         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(IOException.class);
1053         Network mockNetwork = mock(Network.class);
1054         // wait past timeout
1055         mMockInjector.mNetworkConsumer =
1056                 (callback) -> {
1057                     try {
1058                         Thread.sleep(3500);
1059                     } catch (InterruptedException e) {
1060                         throw new RuntimeException(e);
1061                     }
1062                     callback.onAvailable(mockNetwork);
1063                 };
1064         mService.loadRebootEscrowDataIfAvailable(mHandler);
1065         verify(mServiceConnection).unwrap(any(), anyLong());
1066         assertFalse(metricsSuccessCaptor.getValue());
1067         assertEquals(
1068                 Integer.valueOf(RebootEscrowManager.ERROR_TIMEOUT_EXHAUSTED),
1069                 metricsErrorCodeCaptor.getValue());
1070         assertNull(mMockInjector.mNetworkCallback);
1071     }
1072 
1073     @Test
loadRebootEscrowDataIfAvailable_serverBasedWaitForNetwork_retryCountExhausted()1074     public void loadRebootEscrowDataIfAvailable_serverBasedWaitForNetwork_retryCountExhausted()
1075             throws Exception {
1076         setServerBasedRebootEscrowProvider();
1077         mMockInjector.setWaitForNetwork(true);
1078 
1079         when(mInjected.getBootCount()).thenReturn(0);
1080         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1081         mService.setRebootEscrowListener(mockListener);
1082         mService.prepareRebootEscrow();
1083 
1084         clearInvocations(mServiceConnection);
1085         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1086         verify(mockListener).onPreparedForReboot(eq(true));
1087         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
1088 
1089         // Use x -> x for both wrap & unwrap functions.
1090         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
1091                 .thenAnswer(invocation -> invocation.getArgument(0));
1092         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
1093         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
1094         assertTrue(mStorage.hasRebootEscrowServerBlob());
1095 
1096         // pretend reboot happens here
1097         when(mInjected.getBootCount()).thenReturn(1);
1098         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
1099         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
1100         doNothing()
1101                 .when(mInjected)
1102                 .reportMetric(
1103                         metricsSuccessCaptor.capture(),
1104                         metricsErrorCodeCaptor.capture(),
1105                         eq(2) /* Server based */,
1106                         eq(2) /* attempt count */,
1107                         anyInt(),
1108                         eq(0) /* vbmeta status */,
1109                         anyInt());
1110 
1111         when(mServiceConnection.unwrap(any(), anyLong())).thenThrow(new IOException());
1112         Network mockNetwork = mock(Network.class);
1113         mMockInjector.mNetworkConsumer =
1114                 (callback) -> {
1115                     callback.onAvailable(mockNetwork);
1116                 };
1117 
1118         mService.loadRebootEscrowDataIfAvailable(mHandler);
1119         // Sleep 5s for the retry to complete
1120         Thread.sleep(5 * 1000);
1121         verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
1122         assertFalse(metricsSuccessCaptor.getValue());
1123         assertEquals(
1124                 Integer.valueOf(RebootEscrowManager.ERROR_RETRY_COUNT_EXHAUSTED),
1125                 metricsErrorCodeCaptor.getValue());
1126         assertNull(mMockInjector.mNetworkCallback);
1127     }
1128 
1129     @Test
loadRebootEscrowDataIfAvailable_ServerBasedWaitForInternet_RetrySuccess()1130     public void loadRebootEscrowDataIfAvailable_ServerBasedWaitForInternet_RetrySuccess()
1131             throws Exception {
1132         setServerBasedRebootEscrowProvider();
1133         mMockInjector.setWaitForNetwork(true);
1134 
1135         when(mInjected.getBootCount()).thenReturn(0);
1136         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1137         mService.setRebootEscrowListener(mockListener);
1138         mService.prepareRebootEscrow();
1139 
1140         clearInvocations(mServiceConnection);
1141         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1142         verify(mockListener).onPreparedForReboot(eq(true));
1143         verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
1144 
1145         // Use x -> x for both wrap & unwrap functions.
1146         when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
1147                 .thenAnswer(invocation -> invocation.getArgument(0));
1148         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
1149         verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
1150         assertTrue(mStorage.hasRebootEscrowServerBlob());
1151 
1152         // pretend reboot happens here
1153         when(mInjected.getBootCount()).thenReturn(1);
1154         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
1155         doNothing()
1156                 .when(mInjected)
1157                 .reportMetric(
1158                         metricsSuccessCaptor.capture(),
1159                         anyInt(),
1160                         anyInt(),
1161                         eq(2) /* attempt count */,
1162                         anyInt(),
1163                         anyInt(),
1164                         anyInt());
1165 
1166         when(mServiceConnection.unwrap(any(), anyLong()))
1167                 .thenThrow(new IOException())
1168                 .thenAnswer(invocation -> invocation.getArgument(0));
1169         Network mockNetwork = mock(Network.class);
1170         mMockInjector.mNetworkConsumer =
1171                 (callback) -> {
1172                     callback.onAvailable(mockNetwork);
1173                 };
1174 
1175         mService.loadRebootEscrowDataIfAvailable(mHandler);
1176         // Sleep 5s for the retry to complete
1177         Thread.sleep(5 * 1000);
1178         verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
1179         assertTrue(metricsSuccessCaptor.getValue());
1180         verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
1181         assertNull(mMockInjector.mNetworkCallback);
1182     }
1183 
1184     @Test
loadRebootEscrowDataIfAvailable_TooManyBootsInBetween_NoMetrics()1185     public void loadRebootEscrowDataIfAvailable_TooManyBootsInBetween_NoMetrics() throws Exception {
1186         when(mInjected.getBootCount()).thenReturn(0);
1187 
1188         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1189         mService.setRebootEscrowListener(mockListener);
1190         mService.prepareRebootEscrow();
1191 
1192         clearInvocations(mRebootEscrow);
1193         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1194         verify(mockListener).onPreparedForReboot(eq(true));
1195 
1196         verify(mRebootEscrow, never()).storeKey(any());
1197 
1198         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
1199         verify(mRebootEscrow).storeKey(any());
1200 
1201         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
1202         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
1203 
1204         // pretend reboot happens here
1205 
1206         when(mInjected.getBootCount()).thenReturn(10);
1207         when(mRebootEscrow.retrieveKey()).thenReturn(new byte[32]);
1208 
1209         mService.loadRebootEscrowDataIfAvailable(mHandler);
1210         verify(mRebootEscrow).retrieveKey();
1211         verify(mInjected, never()).reportMetric(anyBoolean(), anyInt(), anyInt(), anyInt(),
1212                 anyInt(), anyInt(), anyInt());
1213     }
1214 
1215     @Test
loadRebootEscrowDataIfAvailable_ManualReboot_Failure_NoMetrics()1216     public void loadRebootEscrowDataIfAvailable_ManualReboot_Failure_NoMetrics() throws Exception {
1217         when(mInjected.getBootCount()).thenReturn(0);
1218 
1219         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1220         mService.setRebootEscrowListener(mockListener);
1221         mService.prepareRebootEscrow();
1222 
1223         clearInvocations(mRebootEscrow);
1224         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1225         verify(mockListener).onPreparedForReboot(eq(true));
1226 
1227         verify(mRebootEscrow, never()).storeKey(any());
1228 
1229         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
1230         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
1231 
1232         // pretend reboot happens here
1233 
1234         when(mInjected.getBootCount()).thenReturn(10);
1235         when(mRebootEscrow.retrieveKey()).thenReturn(new byte[32]);
1236 
1237         mService.loadRebootEscrowDataIfAvailable(mHandler);
1238         verify(mInjected, never()).reportMetric(anyBoolean(), anyInt(), anyInt(), anyInt(),
1239                 anyInt(), anyInt(), anyInt());
1240     }
1241 
1242     @Test
loadRebootEscrowDataIfAvailable_OTAFromBeforeArmedStatus_SuccessMetrics()1243     public void loadRebootEscrowDataIfAvailable_OTAFromBeforeArmedStatus_SuccessMetrics()
1244             throws Exception {
1245         when(mInjected.getBootCount()).thenReturn(0);
1246 
1247         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1248         mService.setRebootEscrowListener(mockListener);
1249         mService.prepareRebootEscrow();
1250 
1251         clearInvocations(mRebootEscrow);
1252         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1253         verify(mockListener).onPreparedForReboot(eq(true));
1254 
1255         verify(mRebootEscrow, never()).storeKey(any());
1256 
1257         ArgumentCaptor<byte[]> keyByteCaptor = ArgumentCaptor.forClass(byte[].class);
1258         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
1259         verify(mRebootEscrow).storeKey(keyByteCaptor.capture());
1260 
1261         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
1262         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
1263 
1264         // Delete key to simulate old version that didn't have it.
1265         mStorage.removeKey(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, USER_SYSTEM);
1266 
1267         // pretend reboot happens here
1268 
1269         when(mInjected.getBootCount()).thenReturn(10);
1270         when(mRebootEscrow.retrieveKey()).thenAnswer(invocation -> keyByteCaptor.getValue());
1271 
1272         // Trigger a vbmeta digest mismatch
1273         mStorage.setString(RebootEscrowManager.REBOOT_ESCROW_KEY_VBMETA_DIGEST,
1274                 "non sense value", USER_SYSTEM);
1275         mService.loadRebootEscrowDataIfAvailable(mHandler);
1276         verify(mInjected).reportMetric(eq(true), eq(0) /* error code */, eq(1) /* HAL based */,
1277                 eq(1) /* attempt count */, anyInt(), eq(2) /* vbmeta status */, anyInt());
1278         assertEquals(mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_KEY_VBMETA_DIGEST,
1279                 "", USER_SYSTEM), "");
1280     }
1281 
1282     @Test
loadRebootEscrowDataIfAvailable_RestoreUnsuccessful_Failure()1283     public void loadRebootEscrowDataIfAvailable_RestoreUnsuccessful_Failure() throws Exception {
1284         when(mInjected.getBootCount()).thenReturn(0);
1285 
1286         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1287         mService.setRebootEscrowListener(mockListener);
1288         mService.prepareRebootEscrow();
1289 
1290         clearInvocations(mRebootEscrow);
1291         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1292         verify(mockListener).onPreparedForReboot(eq(true));
1293 
1294         verify(mRebootEscrow, never()).storeKey(any());
1295         assertEquals(ARM_REBOOT_ERROR_NONE, mService.armRebootEscrowIfNeeded());
1296         verify(mRebootEscrow).storeKey(any());
1297 
1298         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
1299         assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
1300 
1301         // pretend reboot happens here.
1302 
1303         when(mInjected.getBootCount()).thenReturn(1);
1304         ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
1305         ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
1306         // Return a null escrow key
1307         doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture(),
1308                 metricsErrorCodeCaptor.capture(), eq(1) /* HAL based */,
1309                 eq(1) /* attempt count */, anyInt(), anyInt(), anyInt());
1310 
1311         when(mRebootEscrow.retrieveKey()).thenAnswer(invocation -> null);
1312         mService.loadRebootEscrowDataIfAvailable(mHandler);
1313         verify(mRebootEscrow).retrieveKey();
1314         assertFalse(metricsSuccessCaptor.getValue());
1315         assertEquals(Integer.valueOf(RebootEscrowManager.ERROR_LOAD_ESCROW_KEY),
1316                 metricsErrorCodeCaptor.getValue());
1317     }
1318 
1319     @Test
armServiceProviderMismatch_Failure()1320     public void armServiceProviderMismatch_Failure() throws Exception {
1321         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
1322         mService.setRebootEscrowListener(mockListener);
1323         mService.prepareRebootEscrow();
1324 
1325         clearInvocations(mRebootEscrow);
1326         callToRebootEscrowIfNeededAndWait(PRIMARY_USER_ID);
1327         verify(mockListener).onPreparedForReboot(eq(true));
1328         assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
1329         verify(mRebootEscrow, never()).storeKey(any());
1330 
1331         assertNull(
1332                 mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, null, USER_SYSTEM));
1333         // Change the provider to server based, expect the reboot to fail
1334         mMockInjector.mServerBased = true;
1335         assertEquals(ARM_REBOOT_ERROR_PROVIDER_MISMATCH, mService.armRebootEscrowIfNeeded());
1336         assertNull(
1337                 mStorage.getString(RebootEscrowManager.REBOOT_ESCROW_ARMED_KEY, null, USER_SYSTEM));
1338         // Verify that the escrow key & data have been cleared.
1339         verify(mRebootEscrow).storeKey(eq(new byte[32]));
1340         assertFalse(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
1341     }
1342 }
1343