Add a unittest to cover the retry of load escrow data

The load escrow data is async. Add a unit test to cover the retry of
temporarily failure.

Bug: 172780686
Test: atest FrameworksServicesTests:RebootEscrowManagerTests
Change-Id: I63b264a5f9b68039868fcf962f4654c8d587985d

2 files changed, 67 insertions, 5 deletions

diff --git a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
index 7e00fd69a148..364aa2cb41ae 100644
--- a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
+++ b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
@@ -193,6 +193,17 @@ class RebootEscrowManager {
                     0);
         }
 
+        public int getLoadEscrowDataRetryLimit() {
+            return DeviceConfig.getInt(DeviceConfig.NAMESPACE_OTA,
+                    "load_escrow_data_retry_count", DEFAULT_LOAD_ESCROW_DATA_RETRY_COUNT);
+        }
+
+        public int getLoadEscrowDataRetryIntervalSeconds() {
+            return DeviceConfig.getInt(DeviceConfig.NAMESPACE_OTA,
+                    "load_escrow_data_retry_interval_seconds",
+                    DEFAULT_LOAD_ESCROW_DATA_RETRY_INTERVAL_SECONDS);
+        }
+
         public void reportMetric(boolean success) {
             // TODO(b/179105110) design error code; and report the true value for other fields.
             FrameworkStatsLog.write(FrameworkStatsLog.REBOOT_ESCROW_RECOVERY_REPORTED, 0, 1, 1,
@@ -251,11 +262,8 @@ class RebootEscrowManager {
             List<UserInfo> users, List<UserInfo> rebootEscrowUsers) {
         Objects.requireNonNull(retryHandler);
 
-        final int retryLimit = DeviceConfig.getInt(DeviceConfig.NAMESPACE_OTA,
-                "load_escrow_data_retry_count", DEFAULT_LOAD_ESCROW_DATA_RETRY_COUNT);
-        final int retryIntervalInSeconds = DeviceConfig.getInt(DeviceConfig.NAMESPACE_OTA,
-                "load_escrow_data_retry_interval_seconds",
-                DEFAULT_LOAD_ESCROW_DATA_RETRY_INTERVAL_SECONDS);
+        final int retryLimit = mInjector.getLoadEscrowDataRetryLimit();
+        final int retryIntervalInSeconds = mInjector.getLoadEscrowDataRetryIntervalSeconds();
 
         if (attemptNumber < retryLimit) {
             Slog.i(TAG, "Scheduling loadRebootEscrowData retry number: " + attemptNumber);
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
index a896f1b0d60f..b51f4df43259 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
@@ -44,6 +44,7 @@ import android.content.ContextWrapper;
 import android.content.pm.UserInfo;
 import android.hardware.rebootescrow.IRebootEscrow;
 import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
 import android.os.UserManager;
@@ -62,6 +63,7 @@ import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 
 import java.io.File;
+import java.io.IOException;
 import java.util.ArrayList;
 
 import javax.crypto.SecretKey;
@@ -181,6 +183,18 @@ public class RebootEscrowManagerTests {
         }
 
         @Override
+        public int getLoadEscrowDataRetryLimit() {
+            // Try two times
+            return 2;
+        }
+
+        @Override
+        public int getLoadEscrowDataRetryIntervalSeconds() {
+            // Retry in 1 seconds
+            return 1;
+        }
+
+        @Override
         public void reportMetric(boolean success) {
             mInjected.reportMetric(success);
         }
@@ -448,6 +462,46 @@ public class RebootEscrowManagerTests {
     }
 
     @Test
+    public void loadRebootEscrowDataIfAvailable_ServerBased_RetrySuccess() throws Exception {
+        setServerBasedRebootEscrowProvider();
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertTrue(mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture());
+
+        when(mServiceConnection.unwrap(any(), anyLong()))
+                .thenThrow(new IOException())
+                .thenAnswer(invocation -> invocation.getArgument(0));
+
+        HandlerThread thread = new HandlerThread("RebootEscrowManagerTest");
+        thread.start();
+        mService.loadRebootEscrowDataIfAvailable(new Handler(thread.getLooper()));
+        // Sleep 5s for the retry to complete
+        Thread.sleep(5 * 1000);
+        verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
+        assertTrue(metricsSuccessCaptor.getValue());
+        verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
+    }
+
+    @Test
     public void loadRebootEscrowDataIfAvailable_TooManyBootsInBetween_NoMetrics() throws Exception {
         when(mInjected.getBootCount()).thenReturn(0);