Add storage for snapshots in KeySyncTask

Test: adb shell am instrument -w -e package com.android.server.locksettings.recoverablekeystore com.android.frameworks.servicestests/android.support.test.runner.AndroidJUnitRunner
Change-Id: I25a9d6999bec5639cc91532da1b42a8d1f911b79
diff --git a/core/java/android/security/recoverablekeystore/RecoverableKeyStoreLoader.java b/core/java/android/security/recoverablekeystore/RecoverableKeyStoreLoader.java
index d8fb03f..72a138a 100644
--- a/core/java/android/security/recoverablekeystore/RecoverableKeyStoreLoader.java
+++ b/core/java/android/security/recoverablekeystore/RecoverableKeyStoreLoader.java
@@ -45,6 +45,8 @@
     public static final int NO_ERROR = KeyStore.NO_ERROR;
     public static final int SYSTEM_ERROR = KeyStore.SYSTEM_ERROR;
     public static final int UNINITIALIZED_RECOVERY_PUBLIC_KEY = 20;
+    public static final int NO_SNAPSHOT_PENDING_ERROR = 21;
+
     /**
      * Rate limit is enforced to prevent using too many trusted remote devices, since each device
      * can have its own number of user secret guesses allowed.
diff --git a/services/core/java/com/android/server/locksettings/recoverablekeystore/KeySyncTask.java b/services/core/java/com/android/server/locksettings/recoverablekeystore/KeySyncTask.java
index 3434eee..1ee60d7 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/KeySyncTask.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/KeySyncTask.java
@@ -16,14 +16,20 @@
 
 package com.android.server.locksettings.recoverablekeystore;
 
+import static android.security.recoverablekeystore.KeyStoreRecoveryMetadata.TYPE_LOCKSCREEN;
+
 import android.annotation.NonNull;
 import android.content.Context;
+import android.security.recoverablekeystore.KeyDerivationParameters;
+import android.security.recoverablekeystore.KeyEntryRecoveryData;
+import android.security.recoverablekeystore.KeyStoreRecoveryData;
 import android.security.recoverablekeystore.KeyStoreRecoveryMetadata;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.widget.LockPatternUtils;
 import com.android.server.locksettings.recoverablekeystore.storage.RecoverableKeyStoreDb;
+import com.android.server.locksettings.recoverablekeystore.storage.RecoverySnapshotStorage;
 
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -36,6 +42,8 @@
 import java.security.PublicKey;
 import java.security.SecureRandom;
 import java.security.UnrecoverableKeyException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 
 import javax.crypto.KeyGenerator;
@@ -58,28 +66,27 @@
 
     private final RecoverableKeyStoreDb mRecoverableKeyStoreDb;
     private final int mUserId;
-    private final RecoverableSnapshotConsumer mSnapshotConsumer;
     private final int mCredentialType;
     private final String mCredential;
     private final PlatformKeyManager.Factory mPlatformKeyManagerFactory;
     private final VaultKeySupplier mVaultKeySupplier;
+    private final RecoverySnapshotStorage mRecoverySnapshotStorage;
 
     public static KeySyncTask newInstance(
             Context context,
             RecoverableKeyStoreDb recoverableKeyStoreDb,
+            RecoverySnapshotStorage snapshotStorage,
             int userId,
             int credentialType,
             String credential
     ) throws NoSuchAlgorithmException, KeyStoreException, InsecureUserException {
         return new KeySyncTask(
                 recoverableKeyStoreDb,
+                snapshotStorage,
                 userId,
                 credentialType,
                 credential,
                 () -> PlatformKeyManager.getInstance(context, recoverableKeyStoreDb, userId),
-                (salt, recoveryKey, applicationKeys) -> {
-                    // TODO: implement sending intent
-                },
                 () -> {
                     throw new UnsupportedOperationException("Not implemented vault key.");
                 });
@@ -99,19 +106,19 @@
     @VisibleForTesting
     KeySyncTask(
             RecoverableKeyStoreDb recoverableKeyStoreDb,
+            RecoverySnapshotStorage snapshotStorage,
             int userId,
             int credentialType,
             String credential,
             PlatformKeyManager.Factory platformKeyManagerFactory,
-            RecoverableSnapshotConsumer snapshotConsumer,
             VaultKeySupplier vaultKeySupplier) {
         mRecoverableKeyStoreDb = recoverableKeyStoreDb;
         mUserId = userId;
         mCredentialType = credentialType;
         mCredential = credential;
         mPlatformKeyManagerFactory = platformKeyManagerFactory;
-        mSnapshotConsumer = snapshotConsumer;
         mVaultKeySupplier = vaultKeySupplier;
+        mRecoverySnapshotStorage = snapshotStorage;
     }
 
     @Override
@@ -185,7 +192,21 @@
             return;
         }
 
-        mSnapshotConsumer.accept(salt, encryptedRecoveryKey, encryptedApplicationKeys);
+        // TODO: why is the secret sent here? I thought it wasn't sent in the raw at all.
+        KeyStoreRecoveryMetadata metadata = new KeyStoreRecoveryMetadata(
+                /*userSecretType=*/ TYPE_LOCKSCREEN,
+                /*lockScreenUiFormat=*/ mCredentialType,
+                /*keyDerivationParameters=*/ KeyDerivationParameters.createSHA256Parameters(salt),
+                /*secret=*/ new byte[0]);
+        ArrayList<KeyStoreRecoveryMetadata> metadataList = new ArrayList<>();
+        metadataList.add(metadata);
+
+        // TODO: implement snapshot version
+        mRecoverySnapshotStorage.put(mUserId, new KeyStoreRecoveryData(
+                /*snapshotVersion=*/ 1,
+                /*recoveryMetadata=*/ metadataList,
+                /*applicationKeyBlobs=*/ createApplicationKeyEntries(encryptedApplicationKeys),
+                /*encryptedRecoveryKeyblob=*/ encryptedRecoveryKey));
     }
 
     private PublicKey getVaultPublicKey() {
@@ -290,15 +311,16 @@
         return keyGenerator.generateKey();
     }
 
-    /**
-     * TODO: just replace with the Intent call. I'm not sure exactly what this looks like, hence
-     * this interface, so I can test in the meantime.
-     */
-    public interface RecoverableSnapshotConsumer {
-        void accept(
-                byte[] salt,
-                byte[] encryptedRecoveryKey,
-                Map<String, byte[]> encryptedApplicationKeys);
+    private static List<KeyEntryRecoveryData> createApplicationKeyEntries(
+            Map<String, byte[]> encryptedApplicationKeys) {
+        ArrayList<KeyEntryRecoveryData> keyEntries = new ArrayList<>();
+        for (String alias : encryptedApplicationKeys.keySet()) {
+            keyEntries.add(
+                    new KeyEntryRecoveryData(
+                            alias.getBytes(StandardCharsets.UTF_8),
+                            encryptedApplicationKeys.get(alias)));
+        }
+        return keyEntries;
     }
 
     /**
diff --git a/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManager.java b/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManager.java
index 426dc8a..047bf1f 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManager.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManager.java
@@ -34,6 +34,7 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.locksettings.recoverablekeystore.storage.RecoverableKeyStoreDb;
 import com.android.server.locksettings.recoverablekeystore.storage.RecoverySessionStorage;
+import com.android.server.locksettings.recoverablekeystore.storage.RecoverySnapshotStorage;
 
 import java.nio.charset.StandardCharsets;
 import java.security.InvalidKeyException;
@@ -75,6 +76,7 @@
     private final ExecutorService mExecutorService;
     private final ListenersStorage mListenersStorage;
     private final RecoverableKeyGenerator mRecoverableKeyGenerator;
+    private final RecoverySnapshotStorage mSnapshotStorage;
 
     /**
      * Returns a new or existing instance.
@@ -89,7 +91,8 @@
                     db,
                     new RecoverySessionStorage(),
                     Executors.newSingleThreadExecutor(),
-                    ListenersStorage.getInstance());
+                    ListenersStorage.getInstance(),
+                    new RecoverySnapshotStorage());
         }
         return mInstance;
     }
@@ -100,12 +103,14 @@
             RecoverableKeyStoreDb recoverableKeyStoreDb,
             RecoverySessionStorage recoverySessionStorage,
             ExecutorService executorService,
-            ListenersStorage listenersStorage) {
+            ListenersStorage listenersStorage,
+            RecoverySnapshotStorage snapshotStorage) {
         mContext = context;
         mDatabase = recoverableKeyStoreDb;
         mRecoverySessionStorage = recoverySessionStorage;
         mExecutorService = executorService;
         mListenersStorage = listenersStorage;
+        mSnapshotStorage = snapshotStorage;
         try {
             mRecoverableKeyGenerator = RecoverableKeyGenerator.newInstance(mDatabase);
         } catch (NoSuchAlgorithmException e) {
@@ -143,24 +148,12 @@
     public @NonNull KeyStoreRecoveryData getRecoveryData(@NonNull byte[] account, int userId)
             throws RemoteException {
         checkRecoverKeyStorePermission();
-        final int callingUid = Binder.getCallingUid(); // Recovery agent uid.
-        final int callingUserId = UserHandle.getCallingUserId();
-        final long callingIdentiy = Binder.clearCallingIdentity();
-        try {
-            // TODO: Return the latest snapshot for the calling recovery agent.
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentiy);
-        }
 
-        // KeyStoreRecoveryData without application keys and empty recovery blob.
-        KeyStoreRecoveryData recoveryData =
-                new KeyStoreRecoveryData(
-                        /*snapshotVersion=*/ 1,
-                        new ArrayList<KeyStoreRecoveryMetadata>(),
-                        new ArrayList<KeyEntryRecoveryData>(),
-                        /*encryptedRecoveryKeyBlob=*/ new byte[] {});
-        throw new ServiceSpecificException(
-                RecoverableKeyStoreLoader.UNINITIALIZED_RECOVERY_PUBLIC_KEY);
+        KeyStoreRecoveryData snapshot = mSnapshotStorage.get(UserHandle.getCallingUserId());
+        if (snapshot == null) {
+            throw new ServiceSpecificException(RecoverableKeyStoreLoader.NO_SNAPSHOT_PENDING_ERROR);
+        }
+        return snapshot;
     }
 
     public void setSnapshotCreatedPendingIntent(@Nullable PendingIntent intent, int userId)
@@ -480,7 +473,7 @@
         // So as not to block the critical path unlocking the phone, defer to another thread.
         try {
             mExecutorService.execute(KeySyncTask.newInstance(
-                    mContext, mDatabase, userId, storedHashType, credential));
+                    mContext, mDatabase, mSnapshotStorage, userId, storedHashType, credential));
         } catch (NoSuchAlgorithmException e) {
             Log.wtf(TAG, "Should never happen - algorithm unavailable for KeySync", e);
         } catch (KeyStoreException e) {
diff --git a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbHelper.java b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbHelper.java
index 80fa20a..c87812d 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbHelper.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbHelper.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2017 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.recoverablekeystore.storage;
 
 import android.content.Context;
diff --git a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorage.java b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorage.java
new file mode 100644
index 0000000..d1a1629
--- /dev/null
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorage.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2017 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.recoverablekeystore.storage;
+
+import android.annotation.Nullable;
+import android.security.recoverablekeystore.KeyStoreRecoveryData;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.GuardedBy;
+
+/**
+ * In-memory storage for recovery snapshots.
+ *
+ * <p>Recovery snapshots are generated after a successful screen unlock. They are only generated if
+ * the recoverable keystore has been mutated since the previous snapshot. This class stores only the
+ * latest snapshot for each user.
+ *
+ * <p>This class is thread-safe. It is used both on the service thread and the
+ * {@link com.android.server.locksettings.recoverablekeystore.KeySyncTask} thread.
+ */
+public class RecoverySnapshotStorage {
+    @GuardedBy("this")
+    private final SparseArray<KeyStoreRecoveryData> mSnapshotByUserId = new SparseArray<>();
+
+    /**
+     * Sets the latest {@code snapshot} for the user {@code userId}.
+     */
+    public synchronized void put(int userId, KeyStoreRecoveryData snapshot) {
+        mSnapshotByUserId.put(userId, snapshot);
+    }
+
+    /**
+     * Returns the latest snapshot for user {@code userId}, or null if none exists.
+     */
+    @Nullable
+    public synchronized KeyStoreRecoveryData get(int userId) {
+        return mSnapshotByUserId.get(userId);
+    }
+
+    /**
+     * Removes any (if any) snapshot associated with user {@code userId}.
+     */
+    public synchronized void remove(int userId) {
+        mSnapshotByUserId.remove(userId);
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/KeySyncTaskTest.java b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/KeySyncTaskTest.java
index b8000c4..393333b 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/KeySyncTaskTest.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/KeySyncTaskTest.java
@@ -26,27 +26,28 @@
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.security.keystore.AndroidKeyStoreSecretKey;
 import android.security.keystore.KeyGenParameterSpec;
 import android.security.keystore.KeyProperties;
+import android.security.recoverablekeystore.KeyDerivationParameters;
+import android.security.recoverablekeystore.KeyEntryRecoveryData;
+import android.security.recoverablekeystore.KeyStoreRecoveryData;
 import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
 
 import com.android.server.locksettings.recoverablekeystore.storage.RecoverableKeyStoreDb;
+import com.android.server.locksettings.recoverablekeystore.storage.RecoverySnapshotStorage;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -54,7 +55,7 @@
 import java.nio.charset.StandardCharsets;
 import java.security.KeyPair;
 import java.util.Arrays;
-import java.util.Map;
+import java.util.List;
 import java.util.Random;
 
 import javax.crypto.KeyGenerator;
@@ -76,13 +77,9 @@
     private static final byte[] THM_ENCRYPTED_RECOVERY_KEY_HEADER =
             "V1 THM_encrypted_recovery_key".getBytes(StandardCharsets.UTF_8);
 
-    @Mock private KeySyncTask.RecoverableSnapshotConsumer mRecoverableSnapshotConsumer;
     @Mock private PlatformKeyManager mPlatformKeyManager;
 
-    @Captor private ArgumentCaptor<byte[]> mSaltCaptor;
-    @Captor private ArgumentCaptor<byte[]> mEncryptedRecoveryKeyCaptor;
-    @Captor private ArgumentCaptor<Map<String, byte[]>> mEncryptedApplicationKeysCaptor;
-
+    private RecoverySnapshotStorage mRecoverySnapshotStorage;
     private RecoverableKeyStoreDb mRecoverableKeyStoreDb;
     private File mDatabaseFile;
     private KeyPair mKeyPair;
@@ -100,13 +97,15 @@
         mRecoverableKeyStoreDb = RecoverableKeyStoreDb.newInstance(context);
         mKeyPair = SecureBox.genKeyPair();
 
+        mRecoverySnapshotStorage = new RecoverySnapshotStorage();
+
         mKeySyncTask = new KeySyncTask(
                 mRecoverableKeyStoreDb,
+                mRecoverySnapshotStorage,
                 TEST_USER_ID,
                 TEST_CREDENTIAL_TYPE,
                 TEST_CREDENTIAL,
                 () -> mPlatformKeyManager,
-                mRecoverableSnapshotConsumer,
                 () -> mKeyPair.getPublic());
 
         mWrappingKey = generateAndroidKeyStoreKey();
@@ -199,7 +198,7 @@
         // to be synced.
         mKeySyncTask.run();
 
-        verifyZeroInteractions(mRecoverableSnapshotConsumer);
+        assertNull(mRecoverySnapshotStorage.get(TEST_USER_ID));
     }
 
     @Test
@@ -214,21 +213,25 @@
 
         mKeySyncTask.run();
 
-        verify(mRecoverableSnapshotConsumer).accept(
-                mSaltCaptor.capture(),
-                mEncryptedRecoveryKeyCaptor.capture(),
-                mEncryptedApplicationKeysCaptor.capture());
+        KeyStoreRecoveryData recoveryData = mRecoverySnapshotStorage.get(TEST_USER_ID);
+        KeyDerivationParameters keyDerivationParameters =
+                recoveryData.getRecoveryMetadata().get(0).getKeyDerivationParameters();
+        assertEquals(KeyDerivationParameters.ALGORITHM_SHA256,
+                keyDerivationParameters.getAlgorithm());
         byte[] lockScreenHash = KeySyncTask.hashCredentials(
-                mSaltCaptor.getValue(), TEST_CREDENTIAL);
+                keyDerivationParameters.getSalt(),
+                TEST_CREDENTIAL);
         // TODO: what should vault params be here?
         byte[] recoveryKey = decryptThmEncryptedKey(
                 lockScreenHash,
-                mEncryptedRecoveryKeyCaptor.getValue(),
+                recoveryData.getEncryptedRecoveryKeyBlob(),
                 /*vaultParams=*/ new byte[0]);
-        Map<String, byte[]> applicationKeys = mEncryptedApplicationKeysCaptor.getValue();
+        List<KeyEntryRecoveryData> applicationKeys = recoveryData.getApplicationKeyBlobs();
         assertEquals(1, applicationKeys.size());
+        KeyEntryRecoveryData keyData = applicationKeys.get(0);
+        assertArrayEquals(TEST_APP_KEY_ALIAS.getBytes(StandardCharsets.UTF_8), keyData.getAlias());
         byte[] appKey = KeySyncUtils.decryptApplicationKey(
-                recoveryKey, applicationKeys.get(TEST_APP_KEY_ALIAS));
+                recoveryKey, keyData.getEncryptedKeyMaterial());
         assertArrayEquals(applicationKey.getEncoded(), appKey);
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManagerTest.java b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManagerTest.java
index 419d7b4..437744d 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManagerTest.java
@@ -48,6 +48,7 @@
 
 import com.android.server.locksettings.recoverablekeystore.storage.RecoverableKeyStoreDb;
 import com.android.server.locksettings.recoverablekeystore.storage.RecoverySessionStorage;
+import com.android.server.locksettings.recoverablekeystore.storage.RecoverySnapshotStorage;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -114,6 +115,7 @@
     private File mDatabaseFile;
     private RecoverableKeyStoreManager mRecoverableKeyStoreManager;
     private RecoverySessionStorage mRecoverySessionStorage;
+    private RecoverySnapshotStorage mRecoverySnapshotStorage;
 
     @Before
     public void setUp() {
@@ -134,7 +136,8 @@
                 mRecoverableKeyStoreDb,
                 mRecoverySessionStorage,
                 Executors.newSingleThreadExecutor(),
-                mMockListenersStorage);
+                mMockListenersStorage,
+                mRecoverySnapshotStorage);
     }
 
     @After
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorageTest.java b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorageTest.java
new file mode 100644
index 0000000..2759e39
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorageTest.java
@@ -0,0 +1,53 @@
+package com.android.server.locksettings.recoverablekeystore.storage;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+import android.security.recoverablekeystore.KeyStoreRecoveryData;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.ArrayList;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class RecoverySnapshotStorageTest {
+
+    private final RecoverySnapshotStorage mRecoverySnapshotStorage = new RecoverySnapshotStorage();
+
+    @Test
+    public void get_isNullForNonExistentSnapshot() {
+        assertNull(mRecoverySnapshotStorage.get(1000));
+    }
+
+    @Test
+    public void get_returnsSetSnapshot() {
+        int userId = 1000;
+        KeyStoreRecoveryData recoveryData = new KeyStoreRecoveryData(
+                /*snapshotVersion=*/ 1,
+                new ArrayList<>(),
+                new ArrayList<>(),
+                new byte[0]);
+        mRecoverySnapshotStorage.put(userId, recoveryData);
+
+        assertEquals(recoveryData, mRecoverySnapshotStorage.get(userId));
+    }
+
+    @Test
+    public void remove_removesSnapshots() {
+        int userId = 1000;
+        KeyStoreRecoveryData recoveryData = new KeyStoreRecoveryData(
+                /*snapshotVersion=*/ 1,
+                new ArrayList<>(),
+                new ArrayList<>(),
+                new byte[0]);
+        mRecoverySnapshotStorage.put(userId, recoveryData);
+
+        mRecoverySnapshotStorage.remove(userId);
+
+        assertNull(mRecoverySnapshotStorage.get(1000));
+    }
+}