Call server based implementation to persist reboot escrow key

Adds the implementation for RebootEscrowProviderServerBased. It will
bind to the resume on reboot service on device; and ask the server
to encrypt/decrypt the reboot escrow key.

Bug: 172780686
Test: atest FrameworksServicesTests:RebootEscrowDataTest \
            FrameworksServicesTests:LockSettingsServiceTests \
            FrameworksServicesTests:RebootEscrowManagerTests \
            FrameworksServicesTests:RebootEscrowProviderServerBasedImplTests;

Change-Id: I6f301314ba43eb106e67d849bfe14e80c607403f

6 files changed, 479 insertions, 13 deletions

diff --git a/services/core/java/com/android/server/locksettings/LockSettingsStorage.java b/services/core/java/com/android/server/locksettings/LockSettingsStorage.java
index 81d07cc11527..c4225eda7d24 100644
--- a/services/core/java/com/android/server/locksettings/LockSettingsStorage.java
+++ b/services/core/java/com/android/server/locksettings/LockSettingsStorage.java
@@ -88,6 +88,7 @@ class LockSettingsStorage {
     private static final String CHILD_PROFILE_LOCK_FILE = "gatekeeper.profile.key";
 
     private static final String REBOOT_ESCROW_FILE = "reboot.escrow.key";
+    private static final String REBOOT_ESCROW_SERVER_BLOB = "reboot.escrow.server.blob.key";
 
     private static final String SYNTHETIC_PASSWORD_DIRECTORY = "spblob/";
 
@@ -317,6 +318,22 @@ class LockSettingsStorage {
         deleteFile(getRebootEscrowFile(userId));
     }
 
+    public void writeRebootEscrowServerBlob(byte[] serverBlob) {
+        writeFile(getRebootEscrowServerBlob(), serverBlob);
+    }
+
+    public byte[] readRebootEscrowServerBlob() {
+        return readFile(getRebootEscrowServerBlob());
+    }
+
+    public boolean hasRebootEscrowServerBlob() {
+        return hasFile(getRebootEscrowServerBlob());
+    }
+
+    public void removeRebootEscrowServerBlob() {
+        deleteFile(getRebootEscrowServerBlob());
+    }
+
     public boolean hasPassword(int userId) {
         return hasFile(getLockPasswordFilename(userId));
     }
@@ -445,6 +462,12 @@ class LockSettingsStorage {
         return getLockCredentialFilePathForUser(userId, REBOOT_ESCROW_FILE);
     }
 
+    @VisibleForTesting
+    String getRebootEscrowServerBlob() {
+        // There is a single copy of server blob for all users.
+        return getLockCredentialFilePathForUser(UserHandle.USER_SYSTEM, REBOOT_ESCROW_SERVER_BLOB);
+    }
+
     private String getLockCredentialFilePathForUser(int userId, String basename) {
         String dataSystemDirectory = Environment.getDataDirectory().getAbsolutePath() +
                         SYSTEM_DIRECTORY;
diff --git a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
index fbec91576ca1..06962d414009 100644
--- a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
+++ b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
@@ -124,26 +124,28 @@ class RebootEscrowManager {
     static class Injector {
         protected Context mContext;
         private final RebootEscrowKeyStoreManager mKeyStoreManager;
-        private final RebootEscrowProviderInterface mRebootEscrowProvider;
+        private final LockSettingsStorage mStorage;
+        private RebootEscrowProviderInterface mRebootEscrowProvider;
 
-        Injector(Context context) {
+        Injector(Context context, LockSettingsStorage storage) {
             mContext = context;
+            mStorage = storage;
             mKeyStoreManager = new RebootEscrowKeyStoreManager();
+        }
 
-            RebootEscrowProviderInterface rebootEscrowProvider = null;
-            // TODO(xunchang) add implementation for server based ror.
+        private RebootEscrowProviderInterface createRebootEscrowProvider() {
+            RebootEscrowProviderInterface rebootEscrowProvider;
             if (DeviceConfig.getBoolean(DeviceConfig.NAMESPACE_OTA,
                     "server_based_ror_enabled", false)) {
-                Slog.e(TAG, "Server based ror isn't implemented yet.");
+                rebootEscrowProvider = new RebootEscrowProviderServerBasedImpl(mContext, mStorage);
             } else {
                 rebootEscrowProvider = new RebootEscrowProviderHalImpl();
             }
 
-            if (rebootEscrowProvider != null && rebootEscrowProvider.hasRebootEscrowSupport()) {
-                mRebootEscrowProvider = rebootEscrowProvider;
-            } else {
-                mRebootEscrowProvider = null;
+            if (rebootEscrowProvider.hasRebootEscrowSupport()) {
+                return rebootEscrowProvider;
             }
+            return null;
         }
 
         public Context getContext() {
@@ -159,6 +161,12 @@ class RebootEscrowManager {
         }
 
         public RebootEscrowProviderInterface getRebootEscrowProvider() {
+            // Initialize for the provider lazily. Because the device_config and service
+            // implementation apps may change when system server is running.
+            if (mRebootEscrowProvider == null) {
+                mRebootEscrowProvider = createRebootEscrowProvider();
+            }
+
             return mRebootEscrowProvider;
         }
 
@@ -177,7 +185,7 @@ class RebootEscrowManager {
     }
 
     RebootEscrowManager(Context context, Callbacks callbacks, LockSettingsStorage storage) {
-        this(new Injector(context), callbacks, storage);
+        this(new Injector(context, storage), callbacks, storage);
     }
 
     @VisibleForTesting
diff --git a/services/core/java/com/android/server/locksettings/RebootEscrowProviderServerBasedImpl.java b/services/core/java/com/android/server/locksettings/RebootEscrowProviderServerBasedImpl.java
new file mode 100644
index 000000000000..d4bdb853358d
--- /dev/null
+++ b/services/core/java/com/android/server/locksettings/RebootEscrowProviderServerBasedImpl.java
@@ -0,0 +1,189 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.locksettings;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.os.RemoteException;
+import android.util.Slog;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.locksettings.ResumeOnRebootServiceProvider.ResumeOnRebootServiceConnection;
+
+import java.io.IOException;
+import java.util.concurrent.TimeoutException;
+
+import javax.crypto.SecretKey;
+
+/**
+ * An implementation of the {@link RebootEscrowProviderInterface} by communicating with server to
+ * encrypt & decrypt the blob.
+ */
+class RebootEscrowProviderServerBasedImpl implements RebootEscrowProviderInterface {
+    private static final String TAG = "RebootEscrowProvider";
+
+    // Timeout for service binding
+    private static final long SERVICE_TIME_OUT_IN_SECONDS = 10;
+
+    /**
+     * Use the default lifetime of 10 minutes. The lifetime covers the following activities:
+     * Server wrap secret -> device reboot -> server unwrap blob.
+     */
+    private static final long SERVER_BLOB_LIFETIME_IN_MILLIS = 600_1000;
+
+    private final LockSettingsStorage mStorage;
+
+    private final Injector mInjector;
+
+    static class Injector {
+        private ResumeOnRebootServiceConnection mServiceConnection = null;
+
+        Injector(Context context) {
+            mServiceConnection = new ResumeOnRebootServiceProvider(context).getServiceConnection();
+            if (mServiceConnection == null) {
+                Slog.e(TAG, "Failed to resolve resume on reboot server service.");
+            }
+        }
+
+        Injector(ResumeOnRebootServiceConnection serviceConnection) {
+            mServiceConnection = serviceConnection;
+        }
+
+        @Nullable
+        private ResumeOnRebootServiceConnection getServiceConnection() {
+            return mServiceConnection;
+        }
+    }
+
+    RebootEscrowProviderServerBasedImpl(Context context, LockSettingsStorage storage) {
+        this(storage, new Injector(context));
+    }
+
+    @VisibleForTesting
+    RebootEscrowProviderServerBasedImpl(LockSettingsStorage storage, Injector injector) {
+        mStorage = storage;
+        mInjector = injector;
+    }
+
+    @Override
+    public boolean hasRebootEscrowSupport() {
+        return mInjector.getServiceConnection() != null;
+    }
+
+    private byte[] unwrapServerBlob(byte[] serverBlob, SecretKey decryptionKey) throws
+            TimeoutException, RemoteException, IOException {
+        ResumeOnRebootServiceConnection serviceConnection = mInjector.getServiceConnection();
+        if (serviceConnection == null) {
+            Slog.w(TAG, "Had reboot escrow data for users, but resume on reboot server"
+                    + " service is unavailable");
+            return null;
+        }
+
+        // Decrypt with k_k from the key store first.
+        byte[] decryptedBlob = AesEncryptionUtil.decrypt(decryptionKey, serverBlob);
+        if (decryptedBlob == null) {
+            Slog.w(TAG, "Decrypted server blob should not be null");
+            return null;
+        }
+
+        // Ask the server connection service to decrypt the inner layer, to get the reboot
+        // escrow key (k_s).
+        serviceConnection.bindToService(SERVICE_TIME_OUT_IN_SECONDS);
+        byte[] escrowKeyBytes = serviceConnection.unwrap(decryptedBlob,
+                SERVICE_TIME_OUT_IN_SECONDS);
+        serviceConnection.unbindService();
+
+        return escrowKeyBytes;
+    }
+
+    @Override
+    public RebootEscrowKey getAndClearRebootEscrowKey(SecretKey decryptionKey) {
+        byte[] serverBlob = mStorage.readRebootEscrowServerBlob();
+        // Delete the server blob in storage.
+        mStorage.removeRebootEscrowServerBlob();
+        if (serverBlob == null) {
+            Slog.w(TAG, "Failed to read reboot escrow server blob from storage");
+            return null;
+        }
+
+        try {
+            byte[] escrowKeyBytes = unwrapServerBlob(serverBlob, decryptionKey);
+            if (escrowKeyBytes == null) {
+                Slog.w(TAG, "Decrypted reboot escrow key bytes should not be null");
+                return null;
+            } else if (escrowKeyBytes.length != 32) {
+                Slog.e(TAG, "Decrypted reboot escrow key has incorrect size "
+                        + escrowKeyBytes.length);
+                return null;
+            }
+
+            return RebootEscrowKey.fromKeyBytes(escrowKeyBytes);
+        } catch (TimeoutException | RemoteException | IOException e) {
+            Slog.w(TAG, "Failed to decrypt the server blob ", e);
+            return null;
+        }
+    }
+
+    @Override
+    public void clearRebootEscrowKey() {
+        mStorage.removeRebootEscrowServerBlob();
+    }
+
+    private byte[] wrapEscrowKey(byte[] escrowKeyBytes, SecretKey encryptionKey) throws
+            TimeoutException, RemoteException, IOException {
+        ResumeOnRebootServiceConnection serviceConnection = mInjector.getServiceConnection();
+        if (serviceConnection == null) {
+            Slog.w(TAG, "Failed to encrypt the reboot escrow key: resume on reboot server"
+                    + " service is unavailable");
+            return null;
+        }
+
+        serviceConnection.bindToService(SERVICE_TIME_OUT_IN_SECONDS);
+        // Ask the server connection service to encrypt the reboot escrow key.
+        byte[] serverEncryptedBlob = serviceConnection.wrapBlob(escrowKeyBytes,
+                SERVER_BLOB_LIFETIME_IN_MILLIS, SERVICE_TIME_OUT_IN_SECONDS);
+        serviceConnection.unbindService();
+
+        if (serverEncryptedBlob == null) {
+            Slog.w(TAG, "Server encrypted reboot escrow key cannot be null");
+            return null;
+        }
+
+        // Additionally wrap the server blob with a local key.
+        return AesEncryptionUtil.encrypt(encryptionKey, serverEncryptedBlob);
+    }
+
+    @Override
+    public boolean storeRebootEscrowKey(RebootEscrowKey escrowKey, SecretKey encryptionKey) {
+        mStorage.removeRebootEscrowServerBlob();
+        try {
+            byte[] wrappedBlob = wrapEscrowKey(escrowKey.getKeyBytes(), encryptionKey);
+            if (wrappedBlob == null) {
+                Slog.w(TAG, "Failed to encrypt the reboot escrow key");
+                return false;
+            }
+            mStorage.writeRebootEscrowServerBlob(wrappedBlob);
+
+            Slog.i(TAG, "Reboot escrow key encrypted and stored.");
+            return true;
+        } catch (TimeoutException | RemoteException | IOException e) {
+            Slog.w(TAG, "Failed to encrypt the reboot escrow key ", e);
+        }
+
+        return false;
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsStorageTestable.java b/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsStorageTestable.java
index 1581d9ac1811..691d174f55f8 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsStorageTestable.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsStorageTestable.java
@@ -82,6 +82,11 @@ public class LockSettingsStorageTestable extends LockSettingsStorage {
     }
 
     @Override
+    String getRebootEscrowServerBlob() {
+        return makeDirs(mStorageDir, super.getRebootEscrowServerBlob()).getAbsolutePath();
+    }
+
+    @Override
     protected File getSyntheticPasswordDirectoryForUser(int userId) {
         return makeDirs(mStorageDir, super.getSyntheticPasswordDirectoryForUser(
                 userId).getAbsolutePath());
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 f74e45b6e59b..a4ba4c86a8fd 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
@@ -26,6 +26,7 @@ import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
@@ -52,6 +53,7 @@ import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.widget.RebootEscrowListener;
+import com.android.server.locksettings.ResumeOnRebootServiceProvider.ResumeOnRebootServiceConnection;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -92,6 +94,7 @@ public class RebootEscrowManagerTests {
     private UserManager mUserManager;
     private RebootEscrowManager.Callbacks mCallbacks;
     private IRebootEscrow mRebootEscrow;
+    private ResumeOnRebootServiceConnection mServiceConnection;
     private RebootEscrowKeyStoreManager mKeyStoreManager;
 
     LockSettingsStorageTestable mStorage;
@@ -108,6 +111,7 @@ public class RebootEscrowManagerTests {
 
     static class MockInjector extends RebootEscrowManager.Injector {
         private final IRebootEscrow mRebootEscrow;
+        private final ResumeOnRebootServiceConnection mServiceConnection;
         private final RebootEscrowProviderInterface mRebootEscrowProvider;
         private final UserManager mUserManager;
         private final MockableRebootEscrowInjected mInjected;
@@ -116,10 +120,11 @@ public class RebootEscrowManagerTests {
         MockInjector(Context context, UserManager userManager,
                 IRebootEscrow rebootEscrow,
                 RebootEscrowKeyStoreManager keyStoreManager,
+                LockSettingsStorageTestable storage,
                 MockableRebootEscrowInjected injected) {
-            super(context);
+            super(context, storage);
             mRebootEscrow = rebootEscrow;
-
+            mServiceConnection = null;
             RebootEscrowProviderHalImpl.Injector halInjector =
                     new RebootEscrowProviderHalImpl.Injector() {
                         @Override
@@ -133,6 +138,22 @@ public class RebootEscrowManagerTests {
             mInjected = injected;
         }
 
+        MockInjector(Context context, UserManager userManager,
+                ResumeOnRebootServiceConnection serviceConnection,
+                RebootEscrowKeyStoreManager keyStoreManager,
+                LockSettingsStorageTestable storage,
+                MockableRebootEscrowInjected injected) {
+            super(context, storage);
+            mServiceConnection = serviceConnection;
+            mRebootEscrow = null;
+            RebootEscrowProviderServerBasedImpl.Injector injector =
+                    new RebootEscrowProviderServerBasedImpl.Injector(serviceConnection);
+            mRebootEscrowProvider = new RebootEscrowProviderServerBasedImpl(storage, injector);
+            mUserManager = userManager;
+            mKeyStoreManager = keyStoreManager;
+            mInjected = injected;
+        }
+
         @Override
         public UserManager getUserManager() {
             return mUserManager;
@@ -165,6 +186,7 @@ public class RebootEscrowManagerTests {
         mUserManager = mock(UserManager.class);
         mCallbacks = mock(RebootEscrowManager.Callbacks.class);
         mRebootEscrow = mock(IRebootEscrow.class);
+        mServiceConnection = mock(ResumeOnRebootServiceConnection.class);
         mKeyStoreManager = mock(RebootEscrowKeyStoreManager.class);
         mAesKey = new SecretKeySpec(TEST_AES_KEY, "AES");
 
@@ -186,7 +208,12 @@ public class RebootEscrowManagerTests {
         when(mCallbacks.isUserSecure(SECURE_SECONDARY_USER_ID)).thenReturn(true);
         mInjected = mock(MockableRebootEscrowInjected.class);
         mService = new RebootEscrowManager(new MockInjector(mContext, mUserManager, mRebootEscrow,
-                mKeyStoreManager, mInjected), mCallbacks, mStorage);
+                mKeyStoreManager, mStorage, mInjected), mCallbacks, mStorage);
+    }
+
+    private void setServerBasedRebootEscrowProvider() throws Exception {
+        mService = new RebootEscrowManager(new MockInjector(mContext, mUserManager,
+                mServiceConnection, mKeyStoreManager, mStorage, mInjected), mCallbacks, mStorage);
     }
 
     @Test
@@ -202,6 +229,19 @@ public class RebootEscrowManagerTests {
     }
 
     @Test
+    public void prepareRebootEscrowServerBased_Success() throws Exception {
+        setServerBasedRebootEscrowProvider();
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+        assertFalse(mStorage.hasRebootEscrowServerBlob());
+    }
+
+    @Test
     public void prepareRebootEscrow_ClearCredentials_Success() throws Exception {
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
         mService.setRebootEscrowListener(mockListener);
@@ -246,6 +286,28 @@ public class RebootEscrowManagerTests {
     }
 
     @Test
+    public void armServiceServerBased_Success() throws Exception {
+        setServerBasedRebootEscrowProvider();
+        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());
+
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertTrue(mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+
+        assertTrue(mStorage.hasRebootEscrow(PRIMARY_USER_ID));
+        assertFalse(mStorage.hasRebootEscrow(NONSECURE_SECONDARY_USER_ID));
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+    }
+
+    @Test
     public void armService_HalFailure_NonFatal() throws Exception {
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
         mService.setRebootEscrowListener(mockListener);
@@ -346,6 +408,40 @@ public class RebootEscrowManagerTests {
     }
 
     @Test
+    public void loadRebootEscrowDataIfAvailable_ServerBased_Success() 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()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        mService.loadRebootEscrowDataIfAvailable();
+        verify(mServiceConnection).unwrap(any(), anyLong());
+        assertTrue(metricsSuccessCaptor.getValue());
+        verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
+    }
+
+    @Test
     public void loadRebootEscrowDataIfAvailable_TooManyBootsInBetween_NoMetrics() throws Exception {
         when(mInjected.getBootCount()).thenReturn(0);
 
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowProviderServerBasedImplTests.java b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowProviderServerBasedImplTests.java
new file mode 100644
index 000000000000..bc1e025dd99f
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowProviderServerBasedImplTests.java
@@ -0,0 +1,145 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.locksettings;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.content.ContextWrapper;
+import android.platform.test.annotations.Presubmit;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.stubbing.Answer;
+
+import java.io.File;
+import java.io.IOException;
+
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+
+@SmallTest
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class RebootEscrowProviderServerBasedImplTests {
+    private SecretKey mKeyStoreEncryptionKey;
+    private RebootEscrowKey mRebootEscrowKey;
+    private ResumeOnRebootServiceProvider.ResumeOnRebootServiceConnection mServiceConnection;
+    private LockSettingsStorageTestable mStorage;
+    private RebootEscrowProviderServerBasedImpl mRebootEscrowProvider;
+    private Answer<byte[]> mFakeEncryption;
+
+    private static final byte[] TEST_AES_KEY = new byte[] {
+            0x48, 0x19, 0x12, 0x54, 0x13, 0x13, 0x52, 0x31,
+            0x44, 0x74, 0x61, 0x54, 0x29, 0x74, 0x37, 0x61,
+            0x70, 0x70, 0x75, 0x25, 0x27, 0x31, 0x49, 0x09,
+            0x26, 0x52, 0x72, 0x63, 0x63, 0x61, 0x78, 0x23,
+    };
+
+    @Before
+    public void setUp() throws Exception {
+        mKeyStoreEncryptionKey = new SecretKeySpec(TEST_AES_KEY, "AES");
+        mRebootEscrowKey = RebootEscrowKey.generate();
+        mServiceConnection = mock(
+                ResumeOnRebootServiceProvider.ResumeOnRebootServiceConnection.class);
+
+        Context context = new ContextWrapper(InstrumentationRegistry.getContext());
+        mStorage = new LockSettingsStorageTestable(context,
+                new File(InstrumentationRegistry.getContext().getFilesDir(), "locksettings"));
+        mRebootEscrowProvider = new RebootEscrowProviderServerBasedImpl(mStorage,
+                new RebootEscrowProviderServerBasedImpl.Injector(mServiceConnection));
+
+        mFakeEncryption = invocation -> {
+            byte[] secret = invocation.getArgument(0);
+            for (int i = 0; i < secret.length; i++) {
+                secret[i] = (byte) (secret[i] ^ 0xf);
+            }
+            return secret;
+        };
+    }
+
+    @Test
+    public void getAndClearRebootEscrowKey_loopback_success() throws Exception {
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong())).thenAnswer(mFakeEncryption);
+        when(mServiceConnection.unwrap(any(), anyLong())).thenAnswer(mFakeEncryption);
+
+        assertTrue(mRebootEscrowProvider.hasRebootEscrowSupport());
+        mRebootEscrowProvider.storeRebootEscrowKey(mRebootEscrowKey, mKeyStoreEncryptionKey);
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+
+        RebootEscrowKey ks = mRebootEscrowProvider.getAndClearRebootEscrowKey(
+                mKeyStoreEncryptionKey);
+        assertThat(ks.getKeyBytes(), is(mRebootEscrowKey.getKeyBytes()));
+        assertFalse(mStorage.hasRebootEscrowServerBlob());
+    }
+
+    @Test
+    public void getAndClearRebootEscrowKey_WrongDecryptionMethod_failure() throws Exception {
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong())).thenAnswer(mFakeEncryption);
+        when(mServiceConnection.unwrap(any(), anyLong())).thenAnswer(
+                invocation -> {
+                    byte[] secret = invocation.getArgument(0);
+                    for (int i = 0; i < secret.length; i++) {
+                        secret[i] = (byte) (secret[i] ^ 0xe);
+                    }
+                    return secret;
+                }
+        );
+
+        assertTrue(mRebootEscrowProvider.hasRebootEscrowSupport());
+        mRebootEscrowProvider.storeRebootEscrowKey(mRebootEscrowKey, mKeyStoreEncryptionKey);
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // Expect to get wrong key bytes
+        RebootEscrowKey ks = mRebootEscrowProvider.getAndClearRebootEscrowKey(
+                mKeyStoreEncryptionKey);
+        assertNotEquals(ks.getKeyBytes(), mRebootEscrowKey.getKeyBytes());
+        assertFalse(mStorage.hasRebootEscrowServerBlob());
+    }
+
+    @Test
+    public void getAndClearRebootEscrowKey_ServiceConnectionException_failure() throws Exception {
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong())).thenAnswer(mFakeEncryption);
+        doThrow(IOException.class).when(mServiceConnection).unwrap(any(), anyLong());
+
+        assertTrue(mRebootEscrowProvider.hasRebootEscrowSupport());
+        mRebootEscrowProvider.storeRebootEscrowKey(mRebootEscrowKey, mKeyStoreEncryptionKey);
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // Expect to get null key bytes when the server service fails to unwrap the blob.
+        RebootEscrowKey ks = mRebootEscrowProvider.getAndClearRebootEscrowKey(
+                mKeyStoreEncryptionKey);
+        assertNull(ks);
+        assertFalse(mStorage.hasRebootEscrowServerBlob());
+    }
+}