Update recovery snapshot version.

There is exactly one snapshot per userId - recovery agent uid pair.
Version is incremented when
1) User credential is updated
2) User unlockes phone and list of application keys was changes since
last snapshot creation.

Bug: 66499222
Test: adb shell am instrument -w -e package \
com.android.server.locksettings.recoverablekeystore \
com.android.frameworks.servicestests/android.support.test.runner.AndroidJUnitRunner

Change-Id: I6ab98fcbbb05e33958e6def644b40441cb52de6a

11 files changed, 517 insertions, 120 deletions

diff --git a/services/core/java/com/android/server/locksettings/LockSettingsService.java b/services/core/java/com/android/server/locksettings/LockSettingsService.java
index 1e852327914c..cc6350aff2f9 100644
--- a/services/core/java/com/android/server/locksettings/LockSettingsService.java
+++ b/services/core/java/com/android/server/locksettings/LockSettingsService.java
@@ -1284,6 +1284,7 @@ public class LockSettingsService extends ILockSettings.Stub {
             fixateNewestUserKeyAuth(userId);
             synchronizeUnifiedWorkChallengeForProfiles(userId, null);
             notifyActivePasswordMetricsAvailable(null, userId);
+            mRecoverableKeyStoreManager.lockScreenSecretChanged(credentialType, credential, userId);
             return;
         }
         if (credential == null) {
@@ -1333,6 +1334,8 @@ public class LockSettingsService extends ILockSettings.Stub {
                     .verifyChallenge(userId, 0, willStore.hash, credential.getBytes());
             setUserKeyProtection(userId, credential, convertResponse(gkResponse));
             fixateNewestUserKeyAuth(userId);
+            mRecoverableKeyStoreManager.lockScreenSecretChanged(credentialType, credential,
+                userId);
             // Refresh the auth token
             doVerifyCredential(credential, credentialType, true, 0, userId, null /* progressCallback */);
             synchronizeUnifiedWorkChallengeForProfiles(userId, null);
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 ded8a796b5ae..3590e3beaa12 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/KeySyncTask.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/KeySyncTask.java
@@ -69,6 +69,7 @@ public class KeySyncTask implements Runnable {
     private final int mUserId;
     private final int mCredentialType;
     private final String mCredential;
+    private final boolean mCredentialUpdated;
     private final PlatformKeyManager.Factory mPlatformKeyManagerFactory;
     private final RecoverySnapshotStorage mRecoverySnapshotStorage;
     private final RecoverySnapshotListenersStorage mSnapshotListenersStorage;
@@ -80,7 +81,8 @@ public class KeySyncTask implements Runnable {
             RecoverySnapshotListenersStorage recoverySnapshotListenersStorage,
             int userId,
             int credentialType,
-            String credential
+            String credential,
+            boolean credentialUpdated
     ) throws NoSuchAlgorithmException, KeyStoreException, InsecureUserException {
         return new KeySyncTask(
                 recoverableKeyStoreDb,
@@ -89,6 +91,7 @@ public class KeySyncTask implements Runnable {
                 userId,
                 credentialType,
                 credential,
+                credentialUpdated,
                 () -> PlatformKeyManager.getInstance(context, recoverableKeyStoreDb));
     }
 
@@ -99,6 +102,7 @@ public class KeySyncTask implements Runnable {
      * @param userId The uid of the user whose profile has been unlocked.
      * @param credentialType The type of credential - i.e., pattern or password.
      * @param credential The credential, encoded as a {@link String}.
+     * @param credentialUpdated signals weather credentials were updated.
      * @param platformKeyManagerFactory Instantiates a {@link PlatformKeyManager} for the user.
      *     This is a factory to enable unit testing, as otherwise it would be impossible to test
      *     without a screen unlock occurring!
@@ -111,12 +115,14 @@ public class KeySyncTask implements Runnable {
             int userId,
             int credentialType,
             String credential,
+            boolean credentialUpdated,
             PlatformKeyManager.Factory platformKeyManagerFactory) {
         mSnapshotListenersStorage = recoverySnapshotListenersStorage;
         mRecoverableKeyStoreDb = recoverableKeyStoreDb;
         mUserId = userId;
         mCredentialType = credentialType;
         mCredential = credential;
+        mCredentialUpdated = credentialUpdated;
         mPlatformKeyManagerFactory = platformKeyManagerFactory;
         mRecoverySnapshotStorage = snapshotStorage;
     }
@@ -124,29 +130,44 @@ public class KeySyncTask implements Runnable {
     @Override
     public void run() {
         try {
-            syncKeys();
+            // Only one task is active If user unlocks phone many times in a short time interval.
+            synchronized(KeySyncTask.class) {
+                syncKeys();
+            }
         } catch (Exception e) {
             Log.e(TAG, "Unexpected exception thrown during KeySyncTask", e);
         }
     }
 
     private void syncKeys() {
-        if (!isSyncPending()) {
-            Log.d(TAG, "Key sync not needed.");
+        if (mCredentialType == LockPatternUtils.CREDENTIAL_TYPE_NONE) {
+            // Application keys for the user will not be available for sync.
+            Log.w(TAG, "Credentials are not set for user " + mUserId);
             return;
         }
 
-        int recoveryAgentUid = mRecoverableKeyStoreDb.getRecoveryAgentUid(mUserId);
-        if (recoveryAgentUid == -1) {
+        List<Integer> recoveryAgents = mRecoverableKeyStoreDb.getRecoveryAgents(mUserId);
+        for (int uid : recoveryAgents) {
+            syncKeysForAgent(uid);
+        }
+        if (recoveryAgents.isEmpty()) {
             Log.w(TAG, "No recovery agent initialized for user " + mUserId);
+        }
+    }
+
+    private void syncKeysForAgent(int recoveryAgentUid) {
+        if (!shoudCreateSnapshot(recoveryAgentUid)) {
+            Log.d(TAG, "Key sync not needed.");
             return;
         }
+
         if (!mSnapshotListenersStorage.hasListener(recoveryAgentUid)) {
             Log.w(TAG, "No pending intent registered for recovery agent " + recoveryAgentUid);
             return;
         }
 
-        PublicKey publicKey = getVaultPublicKey();
+        PublicKey publicKey = mRecoverableKeyStoreDb.getRecoveryServicePublicKey(mUserId,
+                recoveryAgentUid);
         if (publicKey == null) {
             Log.w(TAG, "Not initialized for KeySync: no public key set. Cancelling task.");
             return;
@@ -163,7 +184,7 @@ public class KeySyncTask implements Runnable {
 
         Map<String, SecretKey> rawKeys;
         try {
-            rawKeys = getKeysToSync();
+            rawKeys = getKeysToSync(recoveryAgentUid);
         } catch (GeneralSecurityException e) {
             Log.e(TAG, "Failed to load recoverable keys for sync", e);
             return;
@@ -196,10 +217,19 @@ public class KeySyncTask implements Runnable {
             return;
         }
 
-        // TODO: where do we get counter_id from here?
+        Long counterId;
+        // counter id is generated exactly once for each credentials value.
+        if (mCredentialUpdated) {
+            counterId = generateAndStoreCounterId(recoveryAgentUid);
+        } else {
+            counterId = mRecoverableKeyStoreDb.getCounterId(mUserId, recoveryAgentUid);
+            if (counterId == null) {
+                counterId = generateAndStoreCounterId(recoveryAgentUid);
+            }
+        }
         byte[] vaultParams = KeySyncUtils.packVaultParams(
                 publicKey,
-                /*counterId=*/ 1,
+                counterId,
                 TRUSTED_HARDWARE_MAX_ATTEMPTS,
                 deviceId);
 
@@ -217,7 +247,7 @@ public class KeySyncTask implements Runnable {
             Log.e(TAG,"Could not encrypt with recovery key", e);
             return;
         }
-
+        // TODO: store raw data in RecoveryServiceMetadataEntry and generate Parcelables later
         KeyStoreRecoveryMetadata metadata = new KeyStoreRecoveryMetadata(
                 /*userSecretType=*/ TYPE_LOCKSCREEN,
                 /*lockScreenUiFormat=*/ mCredentialType,
@@ -226,40 +256,62 @@ public class KeySyncTask implements Runnable {
         ArrayList<KeyStoreRecoveryMetadata> metadataList = new ArrayList<>();
         metadataList.add(metadata);
 
-        // TODO: implement snapshot version
-        mRecoverySnapshotStorage.put(mUserId, new KeyStoreRecoveryData(
-                /*snapshotVersion=*/ 1,
+        int snapshotVersion = incrementSnapshotVersion(recoveryAgentUid);
+
+        // If application keys are not updated, snapshot will not be created on next unlock.
+        mRecoverableKeyStoreDb.setShouldCreateSnapshot(mUserId, recoveryAgentUid, false);
+
+        mRecoverySnapshotStorage.put(recoveryAgentUid, new KeyStoreRecoveryData(
+                snapshotVersion,
                 /*recoveryMetadata=*/ metadataList,
                 /*applicationKeyBlobs=*/ createApplicationKeyEntries(encryptedApplicationKeys),
                 /*encryptedRecoveryKeyblob=*/ encryptedRecoveryKey));
+
         mSnapshotListenersStorage.recoverySnapshotAvailable(recoveryAgentUid);
     }
 
-    private PublicKey getVaultPublicKey() {
-        return mRecoverableKeyStoreDb.getRecoveryServicePublicKey(mUserId);
+    @VisibleForTesting
+    int incrementSnapshotVersion(int recoveryAgentUid) {
+        Long snapshotVersion = mRecoverableKeyStoreDb.getSnapshotVersion(mUserId, recoveryAgentUid);
+        snapshotVersion = snapshotVersion == null ? 1 : snapshotVersion + 1;
+        mRecoverableKeyStoreDb.setSnapshotVersion(mUserId, recoveryAgentUid, snapshotVersion);
+
+        return snapshotVersion.intValue();
+    }
+
+    private long generateAndStoreCounterId(int recoveryAgentUid) {
+        long counter = new SecureRandom().nextLong();
+        mRecoverableKeyStoreDb.setCounterId(mUserId, recoveryAgentUid, counter);
+        return counter;
     }
 
     /**
      * Returns all of the recoverable keys for the user.
      */
-    private Map<String, SecretKey> getKeysToSync()
+    private Map<String, SecretKey> getKeysToSync(int recoveryAgentUid)
             throws InsecureUserException, KeyStoreException, UnrecoverableKeyException,
             NoSuchAlgorithmException, NoSuchPaddingException, BadPlatformKeyException {
         PlatformKeyManager platformKeyManager = mPlatformKeyManagerFactory.newInstance();
         PlatformDecryptionKey decryptKey = platformKeyManager.getDecryptKey(mUserId);
         Map<String, WrappedKey> wrappedKeys = mRecoverableKeyStoreDb.getAllKeys(
-                mUserId, decryptKey.getGenerationId());
+                mUserId, recoveryAgentUid, decryptKey.getGenerationId());
         return WrappedKey.unwrapKeys(decryptKey, wrappedKeys);
     }
 
     /**
      * Returns {@code true} if a sync is pending.
+     * @param recoveryAgentUid uid of the recovery agent.
      */
-    private boolean isSyncPending() {
-        // TODO: implement properly. For now just always syncing if the user has any recoverable
-        // keys. We need to keep track of when the store's state actually changes.
-        return !mRecoverableKeyStoreDb.getAllKeys(
-                mUserId, mRecoverableKeyStoreDb.getPlatformKeyGenerationId(mUserId)).isEmpty();
+    private boolean shoudCreateSnapshot(int recoveryAgentUid) {
+        if (mCredentialUpdated) {
+            // Sync credential if at least one snapshot was created.
+            if (mRecoverableKeyStoreDb.getSnapshotVersion(mUserId, recoveryAgentUid) != null) {
+                mRecoverableKeyStoreDb.setShouldCreateSnapshot(mUserId, recoveryAgentUid, true);
+                return true;
+            }
+        }
+
+        return mRecoverableKeyStoreDb.getShouldCreateSnapshot(mUserId, recoveryAgentUid);
     }
 
     /**
diff --git a/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyGenerator.java b/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyGenerator.java
index 8c23d9b4f8e1..2fe3f4e943b3 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyGenerator.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyGenerator.java
@@ -99,6 +99,7 @@ public class RecoverableKeyGenerator {
                             Locale.US, "Failed writing (%d, %s) to database.", uid, alias));
         }
 
+        mDatabase.setShouldCreateSnapshot(userId, uid, true);
         return key.getEncoded();
     }
 }
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 6b8e9ed5fb41..a6f7766e3cf3 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManager.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/RecoverableKeyStoreManager.java
@@ -174,8 +174,8 @@ public class RecoverableKeyStoreManager {
     public @NonNull KeyStoreRecoveryData getRecoveryData(@NonNull byte[] account)
             throws RemoteException {
         checkRecoverKeyStorePermission();
-
-        KeyStoreRecoveryData snapshot = mSnapshotStorage.get(UserHandle.getCallingUserId());
+        int uid = Binder.getCallingUid();
+        KeyStoreRecoveryData snapshot = mSnapshotStorage.get(uid);
         if (snapshot == null) {
             throw new ServiceSpecificException(RecoverableKeyStoreLoader.ERROR_NO_SNAPSHOT_PENDING);
         }
@@ -433,7 +433,13 @@ public class RecoverableKeyStoreManager {
     }
 
     public void removeKey(@NonNull String alias) throws RemoteException {
-        mDatabase.removeKey(Binder.getCallingUid(), alias);
+        int uid = Binder.getCallingUid();
+        int userId = UserHandle.getCallingUserId();
+
+        boolean wasRemoved = mDatabase.removeKey(uid, alias);
+        if (wasRemoved) {
+            mDatabase.setShouldCreateSnapshot(userId, uid, true);
+        }
     }
 
     private byte[] decryptRecoveryKey(
@@ -505,7 +511,8 @@ public class RecoverableKeyStoreManager {
                     mListenersStorage,
                     userId,
                     storedHashType,
-                    credential));
+                    credential,
+                    /*credentialUpdated=*/ false));
         } catch (NoSuchAlgorithmException e) {
             Log.wtf(TAG, "Should never happen - algorithm unavailable for KeySync", e);
         } catch (KeyStoreException e) {
@@ -515,12 +522,35 @@ public class RecoverableKeyStoreManager {
         }
     }
 
-    /** This function can only be used inside LockSettingsService. */
+    /**
+     * This function can only be used inside LockSettingsService.
+     * @param storedHashType from {@code CredentialHash}
+     * @param credential - unencrypted String
+     * @param userId for the user whose lock screen credentials were changed.
+     * @hide
+     */
     public void lockScreenSecretChanged(
-            @KeyStoreRecoveryMetadata.LockScreenUiFormat int type,
+            int storedHashType,
             @Nullable String credential,
             int userId) {
-        throw new UnsupportedOperationException();
+        // So as not to block the critical path unlocking the phone, defer to another thread.
+        try {
+            mExecutorService.execute(KeySyncTask.newInstance(
+                    mContext,
+                    mDatabase,
+                    mSnapshotStorage,
+                    mListenersStorage,
+                    userId,
+                    storedHashType,
+                    credential,
+                    /*credentialUpdated=*/ true));
+        } catch (NoSuchAlgorithmException e) {
+            Log.wtf(TAG, "Should never happen - algorithm unavailable for KeySync", e);
+        } catch (KeyStoreException e) {
+            Log.e(TAG, "Key store error encountered during recoverable key sync", e);
+        } catch (InsecureUserException e) {
+            Log.wtf(TAG, "Impossible - insecure user, but user just entered lock screen", e);
+        }
     }
 
     private void checkRecoverKeyStorePermission() {
diff --git a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDb.java b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDb.java
index 33249e9de8d3..c6f3ede149ef 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDb.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDb.java
@@ -35,8 +35,10 @@ import java.security.NoSuchAlgorithmException;
 import java.security.PublicKey;
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.X509EncodedKeySpec;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.StringJoiner;
@@ -220,16 +222,19 @@ public class RecoverableKeyStoreDb {
     }
 
     /**
-     * Returns all keys for the given {@code userId} and {@code platformKeyGenerationId}.
+     * Returns all keys for the given {@code userId} {@code recoveryAgentUid}
+     * and {@code platformKeyGenerationId}.
      *
      * @param userId User id of the profile to which all the keys are associated.
+     * @param recoveryAgentUid Uid of the recovery agent which will perform the sync
      * @param platformKeyGenerationId The generation ID of the platform key that wrapped these keys.
      *     (i.e., this should be the most recent generation ID, as older platform keys are not
      *     usable.)
      *
      * @hide
      */
-    public Map<String, WrappedKey> getAllKeys(int userId, int platformKeyGenerationId) {
+    public Map<String, WrappedKey> getAllKeys(int userId, int recoveryAgentUid,
+            int platformKeyGenerationId) {
         SQLiteDatabase db = mKeyStoreDbHelper.getReadableDatabase();
         String[] projection = {
                 KeysEntry._ID,
@@ -239,9 +244,13 @@ public class RecoverableKeyStoreDb {
                 KeysEntry.COLUMN_NAME_RECOVERY_STATUS};
         String selection =
                 KeysEntry.COLUMN_NAME_USER_ID + " = ? AND "
+                + KeysEntry.COLUMN_NAME_UID + " = ? AND "
                 + KeysEntry.COLUMN_NAME_GENERATION_ID + " = ?";
         String[] selectionArguments = {
-                Integer.toString(userId), Integer.toString(platformKeyGenerationId) };
+                Integer.toString(userId),
+                Integer.toString(recoveryAgentUid),
+                Integer.toString(platformKeyGenerationId)
+            };
 
         try (
             Cursor cursor = db.query(
@@ -341,9 +350,12 @@ public class RecoverableKeyStoreDb {
     }
 
     /**
-     * Returns the uid of the recovery agent for the given user, or -1 if none is set.
+     * Returns the list of recovery agents initialized for given {@code userId}
+     * @param userId The userId of the profile the application is running under.
+     * @return The list of recovery agents
+     * @hide
      */
-    public int getRecoveryAgentUid(int userId) {
+    public @NonNull List<Integer> getRecoveryAgents(int userId) {
         SQLiteDatabase db = mKeyStoreDbHelper.getReadableDatabase();
 
         String[] projection = { RecoveryServiceMetadataEntry.COLUMN_NAME_UID };
@@ -361,19 +373,20 @@ public class RecoverableKeyStoreDb {
                     /*orderBy=*/ null)
         ) {
             int count = cursor.getCount();
-            if (count == 0) {
-                return -1;
+            ArrayList<Integer> result = new ArrayList<>(count);
+            while (cursor.moveToNext()) {
+                int uid = cursor.getInt(
+                        cursor.getColumnIndexOrThrow(RecoveryServiceMetadataEntry.COLUMN_NAME_UID));
+                result.add(uid);
             }
-            cursor.moveToFirst();
-            return cursor.getInt(
-                    cursor.getColumnIndexOrThrow(RecoveryServiceMetadataEntry.COLUMN_NAME_UID));
+            return result;
         }
     }
 
     /**
      * Returns the public key of the recovery service.
      *
-     * @param userId The uid of the profile the application is running under.
+     * @param userId The userId of the profile the application is running under.
      * @param uid The uid of the application who initializes the local recovery components.
      *
      * @hide
@@ -438,7 +451,7 @@ public class RecoverableKeyStoreDb {
      * If no secret types are set, recovery snapshot will not be created.
      * See {@code KeyStoreRecoveryMetadata}
      *
-     * @param userId The uid of the profile the application is running under.
+     * @param userId The userId of the profile the application is running under.
      * @param uid The uid of the application.
      * @param secretTypes list of secret types
      * @return The primary key of the updated row, or -1 if failed.
@@ -463,7 +476,7 @@ public class RecoverableKeyStoreDb {
     /**
      * Returns the list of secret types used for end-to-end encryption.
      *
-     * @param userId The uid of the profile the application is running under.
+     * @param userId The userId of the profile the application is running under.
      * @param uid The uid of the application who initialized the local recovery components.
      * @return Secret types or empty array, if types were not set.
      *
@@ -529,7 +542,7 @@ public class RecoverableKeyStoreDb {
     /**
      * Returns the first (and only?) public key for {@code userId}.
      *
-     * @param userId The uid of the profile whose keys are to be synced.
+     * @param userId The userId of the profile whose keys are to be synced.
      * @return The public key, or null if none exists.
      */
     @Nullable
@@ -569,10 +582,40 @@ public class RecoverableKeyStoreDb {
     }
 
     /**
+     * Updates the counterId
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application.
+     * @param counterId The counterId.
+     * @return The primary key of the inserted row, or -1 if failed.
+     *
+     * @hide
+     */
+    public long setCounterId(int userId, int uid, long counterId) {
+        return setLong(userId, uid,
+                RecoveryServiceMetadataEntry.COLUMN_NAME_COUNTER_ID, counterId);
+    }
+
+    /**
+     * Returns the counter id.
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application who initialized the local recovery components.
+     * @return The counter id
+     *
+     * @hide
+     */
+    @Nullable
+    public Long getCounterId(int userId, int uid) {
+        return getLong(userId, uid, RecoveryServiceMetadataEntry.COLUMN_NAME_COUNTER_ID);
+    }
+
+
+    /**
      * Updates the server parameters given by the application initializing the local recovery
      * components.
      *
-     * @param userId The uid of the profile the application is running under.
+     * @param userId The userId of the profile the application is running under.
      * @param uid The uid of the application.
      * @param serverParameters The server parameters.
      * @return The primary key of the inserted row, or -1 if failed.
@@ -580,24 +623,15 @@ public class RecoverableKeyStoreDb {
      * @hide
      */
     public long setServerParameters(int userId, int uid, long serverParameters) {
-        SQLiteDatabase db = mKeyStoreDbHelper.getWritableDatabase();
-        ContentValues values = new ContentValues();
-        values.put(RecoveryServiceMetadataEntry.COLUMN_NAME_SERVER_PARAMETERS, serverParameters);
-        String selection =
-                RecoveryServiceMetadataEntry.COLUMN_NAME_USER_ID + " = ? AND "
-                        + RecoveryServiceMetadataEntry.COLUMN_NAME_UID + " = ?";
-        String[] selectionArguments = {Integer.toString(userId), Integer.toString(uid)};
-
-        ensureRecoveryServiceMetadataEntryExists(userId, uid);
-        return db.update(
-                RecoveryServiceMetadataEntry.TABLE_NAME, values, selection, selectionArguments);
+        return setLong(userId, uid,
+                RecoveryServiceMetadataEntry.COLUMN_NAME_SERVER_PARAMETERS, serverParameters);
     }
 
     /**
      * Returns the server paramters that was previously set by the application who initialized the
      * local recovery service components.
      *
-     * @param userId The uid of the profile the application is running under.
+     * @param userId The userId of the profile the application is running under.
      * @param uid The uid of the application who initialized the local recovery components.
      * @return The server parameters that were previously set, or null if there's none.
      *
@@ -605,27 +639,103 @@ public class RecoverableKeyStoreDb {
      */
     @Nullable
     public Long getServerParameters(int userId, int uid) {
+        return getLong(userId, uid, RecoveryServiceMetadataEntry.COLUMN_NAME_SERVER_PARAMETERS);
+
+    }
+
+    /**
+     * Updates the snapshot version.
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application.
+     * @param snapshotVersion The snapshot version
+     * @return The primary key of the inserted row, or -1 if failed.
+     *
+     * @hide
+     */
+    public long setSnapshotVersion(int userId, int uid, long snapshotVersion) {
+        return setLong(userId, uid,
+                RecoveryServiceMetadataEntry.COLUMN_NAME_SNAPSHOT_VERSION, snapshotVersion);
+    }
+
+    /**
+     * Returns the snapshot version
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application who initialized the local recovery components.
+     * @return The server parameters that were previously set, or null if there's none.
+     *
+     * @hide
+     */
+    @Nullable
+    public Long getSnapshotVersion(int userId, int uid) {
+        return getLong(userId, uid,
+            RecoveryServiceMetadataEntry.COLUMN_NAME_SNAPSHOT_VERSION);
+    }
+
+    /**
+     * Updates the snapshot version.
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application.
+     * @param pending The server parameters.
+     * @return The primary key of the inserted row, or -1 if failed.
+     *
+     * @hide
+     */
+    public long setShouldCreateSnapshot(int userId, int uid, boolean pending) {
+        return setLong(userId, uid,
+                RecoveryServiceMetadataEntry.COLUMN_NAME_SHOULD_CREATE_SNAPSHOT, pending ? 1 : 0);
+    }
+
+    /**
+     * Returns {@code true} if new snapshot should be created.
+     * Returns {@code false} if the flag was never set.
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application who initialized the local recovery components.
+     * @return snapshot outdated flag.
+     *
+     * @hide
+     */
+    public boolean getShouldCreateSnapshot(int userId, int uid) {
+        Long res = getLong(userId, uid,
+                RecoveryServiceMetadataEntry.COLUMN_NAME_SHOULD_CREATE_SNAPSHOT);
+        return res != null && res != 0L;
+    }
+
+    /**
+     * Returns given long value from the database.
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application who initialized the local recovery components.
+     * @param key from {@code RecoveryServiceMetadataEntry}
+     * @return The value that were previously set, or null if there's none.
+     *
+     * @hide
+     */
+    private Long getLong(int userId, int uid, String key) {
         SQLiteDatabase db = mKeyStoreDbHelper.getReadableDatabase();
 
         String[] projection = {
                 RecoveryServiceMetadataEntry._ID,
                 RecoveryServiceMetadataEntry.COLUMN_NAME_USER_ID,
                 RecoveryServiceMetadataEntry.COLUMN_NAME_UID,
-                RecoveryServiceMetadataEntry.COLUMN_NAME_SERVER_PARAMETERS};
+                key};
         String selection =
                 RecoveryServiceMetadataEntry.COLUMN_NAME_USER_ID + " = ? AND "
                         + RecoveryServiceMetadataEntry.COLUMN_NAME_UID + " = ?";
         String[] selectionArguments = {Integer.toString(userId), Integer.toString(uid)};
 
         try (
-                Cursor cursor = db.query(
-                        RecoveryServiceMetadataEntry.TABLE_NAME,
-                        projection,
-                        selection,
-                        selectionArguments,
-                        /*groupBy=*/ null,
-                        /*having=*/ null,
-                        /*orderBy=*/ null)
+            Cursor cursor = db.query(
+                    RecoveryServiceMetadataEntry.TABLE_NAME,
+                    projection,
+                    selection,
+                    selectionArguments,
+                    /*groupBy=*/ null,
+                    /*having=*/ null,
+                    /*orderBy=*/ null)
         ) {
             int count = cursor.getCount();
             if (count == 0) {
@@ -634,13 +744,12 @@ public class RecoverableKeyStoreDb {
             if (count > 1) {
                 Log.wtf(TAG,
                         String.format(Locale.US,
-                                "%d deviceId entries found for userId=%d uid=%d. "
+                                "%d entries found for userId=%d uid=%d. "
                                         + "Should only ever be 0 or 1.", count, userId, uid));
                 return null;
             }
             cursor.moveToFirst();
-            int idx = cursor.getColumnIndexOrThrow(
-                    RecoveryServiceMetadataEntry.COLUMN_NAME_SERVER_PARAMETERS);
+            int idx = cursor.getColumnIndexOrThrow(key);
             if (cursor.isNull(idx)) {
                 return null;
             } else {
@@ -650,6 +759,32 @@ public class RecoverableKeyStoreDb {
     }
 
     /**
+     * Sets a long value in the database.
+     *
+     * @param userId The userId of the profile the application is running under.
+     * @param uid The uid of the application who initialized the local recovery components.
+     * @param key defined in {@code RecoveryServiceMetadataEntry}
+     * @param value new value.
+     * @return The primary key of the inserted row, or -1 if failed.
+     *
+     * @hide
+     */
+
+    private long setLong(int userId, int uid, String key, long value) {
+        SQLiteDatabase db = mKeyStoreDbHelper.getWritableDatabase();
+        ContentValues values = new ContentValues();
+        values.put(key, value);
+        String selection =
+                RecoveryServiceMetadataEntry.COLUMN_NAME_USER_ID + " = ? AND "
+                        + RecoveryServiceMetadataEntry.COLUMN_NAME_UID + " = ?";
+        String[] selectionArguments = {Integer.toString(userId), Integer.toString(uid)};
+
+        ensureRecoveryServiceMetadataEntryExists(userId, uid);
+        return db.update(
+                RecoveryServiceMetadataEntry.TABLE_NAME, values, selection, selectionArguments);
+    }
+
+    /**
      * Creates an empty row in the recovery service metadata table if such a row doesn't exist for
      * the given userId and uid, so db.update will succeed.
      */
diff --git a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbContract.java b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbContract.java
index 8f773ddd4598..597ae4c2d1d0 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbContract.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbContract.java
@@ -104,6 +104,15 @@ class RecoverableKeyStoreDbContract {
         static final String COLUMN_NAME_UID = "uid";
 
         /**
+         * Version of the latest recovery snapshot
+         */
+        static final String COLUMN_NAME_SNAPSHOT_VERSION = "snapshot_version";
+        /**
+         * Flag to generate new snapshot.
+         */
+        static final String COLUMN_NAME_SHOULD_CREATE_SNAPSHOT = "should_create_snapshot";
+
+        /**
          * The public key of the recovery service.
          */
         static final String COLUMN_NAME_PUBLIC_KEY = "public_key";
@@ -114,6 +123,11 @@ class RecoverableKeyStoreDbContract {
         static final String COLUMN_NAME_SECRET_TYPES = "secret_types";
 
         /**
+         * Locally generated random number.
+         */
+        static final String COLUMN_NAME_COUNTER_ID = "counter_id";
+
+        /**
          * The server parameters of the recovery service.
          */
         static final String COLUMN_NAME_SERVER_PARAMETERS = "server_parameters";
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 5b07f3e28e83..6eb47ee44f24 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
@@ -51,13 +51,16 @@ class RecoverableKeyStoreDbHelper extends SQLiteOpenHelper {
                     + UserMetadataEntry.COLUMN_NAME_USER_ID + " INTEGER UNIQUE,"
                     + UserMetadataEntry.COLUMN_NAME_PLATFORM_KEY_GENERATION_ID + " INTEGER)";
 
-    private static final String SQL_CREATE_RECOVERY_SERVICE_PUBLIC_KEY_ENTRY =
+    private static final String SQL_CREATE_RECOVERY_SERVICE_METADATA_ENTRY =
             "CREATE TABLE " + RecoveryServiceMetadataEntry.TABLE_NAME + " ("
                     + RecoveryServiceMetadataEntry._ID + " INTEGER PRIMARY KEY,"
                     + RecoveryServiceMetadataEntry.COLUMN_NAME_USER_ID + " INTEGER,"
                     + RecoveryServiceMetadataEntry.COLUMN_NAME_UID + " INTEGER,"
+                    + RecoveryServiceMetadataEntry.COLUMN_NAME_SNAPSHOT_VERSION + " INTEGER,"
+                    + RecoveryServiceMetadataEntry.COLUMN_NAME_SHOULD_CREATE_SNAPSHOT + " INTEGER,"
                     + RecoveryServiceMetadataEntry.COLUMN_NAME_PUBLIC_KEY + " BLOB,"
                     + RecoveryServiceMetadataEntry.COLUMN_NAME_SECRET_TYPES + " TEXT,"
+                    + RecoveryServiceMetadataEntry.COLUMN_NAME_COUNTER_ID + " INTEGER,"
                     + RecoveryServiceMetadataEntry.COLUMN_NAME_SERVER_PARAMETERS + " INTEGER,"
                     + "UNIQUE("
                     + RecoveryServiceMetadataEntry.COLUMN_NAME_USER_ID  + ","
@@ -69,7 +72,7 @@ class RecoverableKeyStoreDbHelper extends SQLiteOpenHelper {
     private static final String SQL_DELETE_USER_METADATA_ENTRY =
             "DROP TABLE IF EXISTS " + UserMetadataEntry.TABLE_NAME;
 
-    private static final String SQL_DELETE_RECOVERY_SERVICE_PUBLIC_KEY_ENTRY =
+    private static final String SQL_DELETE_RECOVERY_SERVICE_METADATA_ENTRY =
             "DROP TABLE IF EXISTS " + RecoveryServiceMetadataEntry.TABLE_NAME;
 
     RecoverableKeyStoreDbHelper(Context context) {
@@ -80,14 +83,14 @@ class RecoverableKeyStoreDbHelper extends SQLiteOpenHelper {
     public void onCreate(SQLiteDatabase db) {
         db.execSQL(SQL_CREATE_KEYS_ENTRY);
         db.execSQL(SQL_CREATE_USER_METADATA_ENTRY);
-        db.execSQL(SQL_CREATE_RECOVERY_SERVICE_PUBLIC_KEY_ENTRY);
+        db.execSQL(SQL_CREATE_RECOVERY_SERVICE_METADATA_ENTRY);
     }
 
     @Override
     public void onUpgrade(SQLiteDatabase db, int oldVersion, int newVersion) {
         db.execSQL(SQL_DELETE_KEYS_ENTRY);
         db.execSQL(SQL_DELETE_USER_METADATA_ENTRY);
-        db.execSQL(SQL_DELETE_RECOVERY_SERVICE_PUBLIC_KEY_ENTRY);
+        db.execSQL(SQL_DELETE_RECOVERY_SERVICE_METADATA_ENTRY);
         onCreate(db);
     }
 }
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
index d1a1629d101e..011b374bcbdd 100644
--- a/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorage.java
+++ b/services/core/java/com/android/server/locksettings/recoverablekeystore/storage/RecoverySnapshotStorage.java
@@ -27,34 +27,34 @@ import com.android.internal.annotations.GuardedBy;
  *
  * <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.
+ * latest snapshot for each recovery agent.
  *
  * <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<>();
+    private final SparseArray<KeyStoreRecoveryData> mSnapshotByUid = new SparseArray<>();
 
     /**
-     * Sets the latest {@code snapshot} for the user {@code userId}.
+     * Sets the latest {@code snapshot} for the recovery agent {@code uid}.
      */
-    public synchronized void put(int userId, KeyStoreRecoveryData snapshot) {
-        mSnapshotByUserId.put(userId, snapshot);
+    public synchronized void put(int uid, KeyStoreRecoveryData snapshot) {
+        mSnapshotByUid.put(uid, snapshot);
     }
 
     /**
-     * Returns the latest snapshot for user {@code userId}, or null if none exists.
+     * Returns the latest snapshot for the recovery agent {@code uid}, or null if none exists.
      */
     @Nullable
-    public synchronized KeyStoreRecoveryData get(int userId) {
-        return mSnapshotByUserId.get(userId);
+    public synchronized KeyStoreRecoveryData get(int uid) {
+        return mSnapshotByUid.get(uid);
     }
 
     /**
-     * Removes any (if any) snapshot associated with user {@code userId}.
+     * Removes any (if any) snapshot associated with recovery agent {@code uid}.
      */
-    public synchronized void remove(int userId) {
-        mSnapshotByUserId.remove(userId);
+    public synchronized void remove(int uid) {
+        mSnapshotByUid.remove(uid);
     }
 }
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 763ee07afd3b..1895e15454e4 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
@@ -23,11 +23,14 @@ import static android.security.recoverablekeystore.KeyStoreRecoveryMetadata.TYPE
 import static com.android.internal.widget.LockPatternUtils.CREDENTIAL_TYPE_PASSWORD;
 import static com.android.internal.widget.LockPatternUtils.CREDENTIAL_TYPE_PATTERN;
 
+import static com.google.common.truth.Truth.assertThat;
+
 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.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -70,8 +73,8 @@ public class KeySyncTaskTest {
     private static final String WRAPPING_KEY_ALIAS = "KeySyncTaskTest/WrappingKey";
     private static final String DATABASE_FILE_NAME = "recoverablekeystore.db";
     private static final int TEST_USER_ID = 1000;
-    private static final int TEST_APP_UID = 10009;
-    private static final int TEST_RECOVERY_AGENT_UID = 90873;
+    private static final int TEST_RECOVERY_AGENT_UID = 10009;
+    private static final int TEST_RECOVERY_AGENT_UID2 = 10010;
     private static final long TEST_DEVICE_ID = 13295035643L;
     private static final String TEST_APP_KEY_ALIAS = "rcleaver";
     private static final int TEST_GENERATION_ID = 2;
@@ -110,6 +113,7 @@ public class KeySyncTaskTest {
                 TEST_USER_ID,
                 TEST_CREDENTIAL_TYPE,
                 TEST_CREDENTIAL,
+                /*credentialUpdated=*/ false,
                 () -> mPlatformKeyManager);
 
         mWrappingKey = generateAndroidKeyStoreKey();
@@ -198,11 +202,16 @@ public class KeySyncTaskTest {
 
     @Test
     public void run_doesNotSendAnythingIfNoKeysToSync() throws Exception {
-        // TODO: proper test here, once we have proper implementation for checking that keys need
-        // to be synced.
         mKeySyncTask.run();
 
-        assertNull(mRecoverySnapshotStorage.get(TEST_USER_ID));
+        assertNull(mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID));
+    }
+
+    @Test
+    public void run_doesNotSendAnythingIfSnapshotIsUpToDate() throws Exception {
+        mKeySyncTask.run();
+
+        assertNull(mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID));
     }
 
     @Test
@@ -211,14 +220,14 @@ public class KeySyncTaskTest {
         mRecoverableKeyStoreDb.setPlatformKeyGenerationId(TEST_USER_ID, TEST_GENERATION_ID);
         mRecoverableKeyStoreDb.insertKey(
                 TEST_USER_ID,
-                TEST_APP_UID,
+                TEST_RECOVERY_AGENT_UID,
                 TEST_APP_KEY_ALIAS,
                 WrappedKey.fromSecretKey(mEncryptKey, applicationKey));
         when(mSnapshotListenersStorage.hasListener(TEST_RECOVERY_AGENT_UID)).thenReturn(true);
 
         mKeySyncTask.run();
 
-        assertNull(mRecoverySnapshotStorage.get(TEST_USER_ID));
+        assertNull(mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID));
     }
 
     @Test
@@ -229,7 +238,7 @@ public class KeySyncTaskTest {
         mRecoverableKeyStoreDb.setPlatformKeyGenerationId(TEST_USER_ID, TEST_GENERATION_ID);
         mRecoverableKeyStoreDb.insertKey(
                 TEST_USER_ID,
-                TEST_APP_UID,
+                TEST_RECOVERY_AGENT_UID,
                 TEST_APP_KEY_ALIAS,
                 WrappedKey.fromSecretKey(mEncryptKey, applicationKey));
         mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
@@ -237,7 +246,7 @@ public class KeySyncTaskTest {
 
         mKeySyncTask.run();
 
-        assertNull(mRecoverySnapshotStorage.get(TEST_USER_ID));
+        assertNull(mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID));
     }
 
     @Test
@@ -246,7 +255,7 @@ public class KeySyncTaskTest {
         mRecoverableKeyStoreDb.setPlatformKeyGenerationId(TEST_USER_ID, TEST_GENERATION_ID);
         mRecoverableKeyStoreDb.insertKey(
                 TEST_USER_ID,
-                TEST_APP_UID,
+                TEST_RECOVERY_AGENT_UID,
                 TEST_APP_KEY_ALIAS,
                 WrappedKey.fromSecretKey(mEncryptKey, applicationKey));
         mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
@@ -255,51 +264,120 @@ public class KeySyncTaskTest {
 
         mKeySyncTask.run();
 
-        assertNull(mRecoverySnapshotStorage.get(TEST_USER_ID));
+        assertNull(mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID));
     }
 
     @Test
     public void run_sendsEncryptedKeysIfAvailableToSync() throws Exception {
-        SecretKey applicationKey = generateKey();
-        mRecoverableKeyStoreDb.setServerParameters(
-                TEST_USER_ID, TEST_RECOVERY_AGENT_UID, TEST_DEVICE_ID);
-        mRecoverableKeyStoreDb.setPlatformKeyGenerationId(TEST_USER_ID, TEST_GENERATION_ID);
-        mRecoverableKeyStoreDb.insertKey(
-                TEST_USER_ID,
-                TEST_APP_UID,
-                TEST_APP_KEY_ALIAS,
-                WrappedKey.fromSecretKey(mEncryptKey, applicationKey));
+
         mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
                 TEST_USER_ID, TEST_RECOVERY_AGENT_UID, mKeyPair.getPublic());
         when(mSnapshotListenersStorage.hasListener(TEST_RECOVERY_AGENT_UID)).thenReturn(true);
-
+        SecretKey applicationKey =
+                addApplicationKey(TEST_USER_ID, TEST_RECOVERY_AGENT_UID, TEST_APP_KEY_ALIAS);
         mKeySyncTask.run();
 
-        KeyStoreRecoveryData recoveryData = mRecoverySnapshotStorage.get(TEST_USER_ID);
+        KeyStoreRecoveryData recoveryData = mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID);
         KeyDerivationParameters keyDerivationParameters =
                 recoveryData.getRecoveryMetadata().get(0).getKeyDerivationParameters();
-        assertEquals(KeyDerivationParameters.ALGORITHM_SHA256,
-                keyDerivationParameters.getAlgorithm());
+        assertThat(keyDerivationParameters.getAlgorithm()).isEqualTo(
+                KeyDerivationParameters.ALGORITHM_SHA256);
         verify(mSnapshotListenersStorage).recoverySnapshotAvailable(TEST_RECOVERY_AGENT_UID);
         byte[] lockScreenHash = KeySyncTask.hashCredentials(
                 keyDerivationParameters.getSalt(),
                 TEST_CREDENTIAL);
-        // TODO: what should counter_id be here?
+        Long counterId = mRecoverableKeyStoreDb.getCounterId(TEST_USER_ID, TEST_RECOVERY_AGENT_UID);
+        assertThat(counterId).isNotNull();
         byte[] recoveryKey = decryptThmEncryptedKey(
                 lockScreenHash,
                 recoveryData.getEncryptedRecoveryKeyBlob(),
                 /*vaultParams=*/ KeySyncUtils.packVaultParams(
                         mKeyPair.getPublic(),
-                        /*counterId=*/ 1,
+                        counterId,
                         /*maxAttempts=*/ 10,
                         TEST_DEVICE_ID));
         List<KeyEntryRecoveryData> applicationKeys = recoveryData.getApplicationKeyBlobs();
-        assertEquals(1, applicationKeys.size());
+        assertThat(applicationKeys).hasSize(1);
         KeyEntryRecoveryData keyData = applicationKeys.get(0);
         assertEquals(TEST_APP_KEY_ALIAS, keyData.getAlias());
+        assertThat(keyData.getAlias()).isEqualTo(keyData.getAlias());
         byte[] appKey = KeySyncUtils.decryptApplicationKey(
                 recoveryKey, keyData.getEncryptedKeyMaterial());
-        assertArrayEquals(applicationKey.getEncoded(), appKey);
+        assertThat(appKey).isEqualTo(applicationKey.getEncoded());
+    }
+
+    @Test
+    public void run_setsCorrectSnapshotVersion() throws Exception {
+
+        mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
+                TEST_USER_ID, TEST_RECOVERY_AGENT_UID, mKeyPair.getPublic());
+        when(mSnapshotListenersStorage.hasListener(TEST_RECOVERY_AGENT_UID)).thenReturn(true);
+        SecretKey applicationKey =
+                addApplicationKey(TEST_USER_ID, TEST_RECOVERY_AGENT_UID, TEST_APP_KEY_ALIAS);
+
+        mKeySyncTask.run();
+
+        KeyStoreRecoveryData recoveryData = mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID);
+        assertThat(recoveryData.getSnapshotVersion()).isEqualTo(1); // default value;
+        mRecoverableKeyStoreDb.setShouldCreateSnapshot(TEST_USER_ID, TEST_RECOVERY_AGENT_UID, true);
+
+        mKeySyncTask.run();
+
+        recoveryData = mRecoverySnapshotStorage.get(TEST_RECOVERY_AGENT_UID);
+        assertThat(recoveryData.getSnapshotVersion()).isEqualTo(2); // Updated
+    }
+
+    @Test
+    public void run_sendsEncryptedKeysWithTwoRegisteredAgents() throws Exception {
+
+        mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
+                TEST_USER_ID, TEST_RECOVERY_AGENT_UID, mKeyPair.getPublic());
+        mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
+                TEST_USER_ID, TEST_RECOVERY_AGENT_UID2, mKeyPair.getPublic());
+        when(mSnapshotListenersStorage.hasListener(TEST_RECOVERY_AGENT_UID)).thenReturn(true);
+        when(mSnapshotListenersStorage.hasListener(TEST_RECOVERY_AGENT_UID2)).thenReturn(true);
+        addApplicationKey(TEST_USER_ID, TEST_RECOVERY_AGENT_UID, TEST_APP_KEY_ALIAS);
+        addApplicationKey(TEST_USER_ID, TEST_RECOVERY_AGENT_UID2, TEST_APP_KEY_ALIAS);
+        mKeySyncTask.run();
+
+        verify(mSnapshotListenersStorage).recoverySnapshotAvailable(TEST_RECOVERY_AGENT_UID);
+        verify(mSnapshotListenersStorage).recoverySnapshotAvailable(TEST_RECOVERY_AGENT_UID2);
+    }
+
+    @Test
+    public void run_doesnSendKeyToNonregisteredAgent() throws Exception {
+
+        mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
+                TEST_USER_ID, TEST_RECOVERY_AGENT_UID, mKeyPair.getPublic());
+        mRecoverableKeyStoreDb.setRecoveryServicePublicKey(
+                TEST_USER_ID, TEST_RECOVERY_AGENT_UID2, mKeyPair.getPublic());
+        when(mSnapshotListenersStorage.hasListener(TEST_RECOVERY_AGENT_UID)).thenReturn(true);
+        when(mSnapshotListenersStorage.hasListener(TEST_RECOVERY_AGENT_UID2)).thenReturn(false);
+        addApplicationKey(TEST_USER_ID, TEST_RECOVERY_AGENT_UID, TEST_APP_KEY_ALIAS);
+        addApplicationKey(TEST_USER_ID, TEST_RECOVERY_AGENT_UID2, TEST_APP_KEY_ALIAS);
+        mKeySyncTask.run();
+
+        verify(mSnapshotListenersStorage).recoverySnapshotAvailable(TEST_RECOVERY_AGENT_UID);
+        verify(mSnapshotListenersStorage, never()).
+                recoverySnapshotAvailable(TEST_RECOVERY_AGENT_UID2);
+    }
+
+    private SecretKey addApplicationKey(int userId, int recoveryAgentUid, String alias)
+            throws Exception{
+        SecretKey applicationKey = generateKey();
+        mRecoverableKeyStoreDb.setServerParameters(
+                userId, recoveryAgentUid, TEST_DEVICE_ID);
+        mRecoverableKeyStoreDb.setPlatformKeyGenerationId(userId, TEST_GENERATION_ID);
+
+        // Newly added key is not synced.
+        mRecoverableKeyStoreDb.setShouldCreateSnapshot(userId, recoveryAgentUid, true);
+
+        mRecoverableKeyStoreDb.insertKey(
+                userId,
+                recoveryAgentUid,
+                alias,
+                WrappedKey.fromSecretKey(mEncryptKey, applicationKey));
+        return applicationKey;
     }
 
     private byte[] decryptThmEncryptedKey(
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 01050667a2c0..c5bce1863ea7 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
@@ -180,10 +180,13 @@ public class RecoverableKeyStoreManagerTest {
     @Test
     public void generateAndStoreKey_storesTheKey() throws Exception {
         int uid = Binder.getCallingUid();
+        int userId = UserHandle.getCallingUserId();
 
         mRecoverableKeyStoreManager.generateAndStoreKey(TEST_ALIAS);
 
         assertThat(mRecoverableKeyStoreDb.getKey(uid, TEST_ALIAS)).isNotNull();
+
+        assertThat(mRecoverableKeyStoreDb.getShouldCreateSnapshot(userId, uid)).isTrue();
     }
 
     @Test
@@ -203,6 +206,20 @@ public class RecoverableKeyStoreManagerTest {
     }
 
     @Test
+    public void removeKey_UpdatesShouldCreateSnapshot() throws Exception {
+        int uid = Binder.getCallingUid();
+        int userId = UserHandle.getCallingUserId();
+
+        mRecoverableKeyStoreManager.generateAndStoreKey(TEST_ALIAS);
+        // Pretend that key was synced
+        mRecoverableKeyStoreDb.setShouldCreateSnapshot(userId, uid, false);
+
+        mRecoverableKeyStoreManager.removeKey(TEST_ALIAS);
+
+        assertThat(mRecoverableKeyStoreDb.getShouldCreateSnapshot(userId, uid)).isTrue();
+    }
+
+    @Test
     public void startRecoverySession_checksPermissionFirst() throws Exception {
         mRecoverableKeyStoreManager.startRecoverySession(
                 TEST_SESSION_ID,
@@ -253,7 +270,8 @@ public class RecoverableKeyStoreManagerTest {
                     ImmutableList.of());
             fail("should have thrown");
         } catch (ServiceSpecificException e) {
-            assertEquals("Only a single KeyStoreRecoveryMetadata is supported", e.getMessage());
+            assertThat(e.getMessage()).startsWith(
+                    "Only a single KeyStoreRecoveryMetadata is supported");
         }
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbTest.java b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbTest.java
index c7b338c6fa49..b8080ab07b4a 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbTest.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/recoverablekeystore/storage/RecoverableKeyStoreDbTest.java
@@ -40,6 +40,7 @@ import java.nio.charset.StandardCharsets;
 import java.security.KeyPairGenerator;
 import java.security.PublicKey;
 import java.security.spec.ECGenParameterSpec;
+import java.util.List;
 import java.util.Map;
 
 @SmallTest
@@ -167,7 +168,7 @@ public class RecoverableKeyStoreDbTest {
         WrappedKey wrappedKey = new WrappedKey(nonce, keyMaterial, generationId);
         mRecoverableKeyStoreDb.insertKey(userId, uid, alias, wrappedKey);
 
-        Map<String, WrappedKey> keys = mRecoverableKeyStoreDb.getAllKeys(userId, generationId);
+        Map<String, WrappedKey> keys = mRecoverableKeyStoreDb.getAllKeys(userId, uid, generationId);
 
         assertEquals(1, keys.size());
         assertTrue(keys.containsKey(alias));
@@ -189,7 +190,7 @@ public class RecoverableKeyStoreDbTest {
                 userId, uid, /*alias=*/ "test", wrappedKey);
 
         Map<String, WrappedKey> keys = mRecoverableKeyStoreDb.getAllKeys(
-                userId, /*generationId=*/ 7);
+                userId, uid, /*generationId=*/ 7);
 
         assertTrue(keys.isEmpty());
     }
@@ -204,7 +205,7 @@ public class RecoverableKeyStoreDbTest {
                 /*userId=*/ 1, uid, /*alias=*/ "test", wrappedKey);
 
         Map<String, WrappedKey> keys = mRecoverableKeyStoreDb.getAllKeys(
-                /*userId=*/ 2, generationId);
+                /*userId=*/ 2, uid, generationId);
 
         assertTrue(keys.isEmpty());
     }
@@ -344,17 +345,33 @@ public class RecoverableKeyStoreDbTest {
     }
 
     @Test
-    public void getRecoveryAgentUid_returnsUidIfSet() throws Exception {
+    public void getRecoveryAgents_returnsUidIfSet() throws Exception {
         int userId = 12;
         int uid = 190992;
         mRecoverableKeyStoreDb.setRecoveryServicePublicKey(userId, uid, genRandomPublicKey());
 
-        assertThat(mRecoverableKeyStoreDb.getRecoveryAgentUid(userId)).isEqualTo(uid);
+        assertThat(mRecoverableKeyStoreDb.getRecoveryAgents(userId)).contains(uid);
     }
 
     @Test
-    public void getRecoveryAgentUid_returnsMinusOneForNonexistentAgent() throws Exception {
-        assertThat(mRecoverableKeyStoreDb.getRecoveryAgentUid(12)).isEqualTo(-1);
+    public void getRecoveryAgents_returnsEmptyListIfThereAreNoAgents() throws Exception {
+        int userId = 12;
+        assertThat(mRecoverableKeyStoreDb.getRecoveryAgents(userId)).isEmpty();
+        assertThat(mRecoverableKeyStoreDb.getRecoveryAgents(userId)).isNotNull();
+    }
+
+    @Test
+    public void getRecoveryAgents_withTwoAgents() throws Exception {
+        int userId = 12;
+        int uid1 = 190992;
+        int uid2 = 190993;
+        mRecoverableKeyStoreDb.setRecoveryServicePublicKey(userId, uid1, genRandomPublicKey());
+        mRecoverableKeyStoreDb.setRecoveryServicePublicKey(userId, uid2, genRandomPublicKey());
+        List<Integer> agents = mRecoverableKeyStoreDb.getRecoveryAgents(userId);
+
+        assertThat(agents).hasSize(2);
+        assertThat(agents).contains(uid1);
+        assertThat(agents).contains(uid2);
     }
 
     public void setRecoverySecretTypes_emptyDefaultValue() throws Exception {
@@ -495,6 +512,52 @@ public class RecoverableKeyStoreDbTest {
     }
 
     @Test
+    public void setCounterId_defaultValueAndTwoUpdates() throws Exception {
+        int userId = 12;
+        int uid = 10009;
+        long value1 = 111L;
+        long value2 = 222L;
+        assertThat(mRecoverableKeyStoreDb.getCounterId(userId, uid)).isNull();
+
+        mRecoverableKeyStoreDb.setCounterId(userId, uid, value1);
+        assertThat(mRecoverableKeyStoreDb.getCounterId(userId, uid)).isEqualTo(
+                value1);
+
+        mRecoverableKeyStoreDb.setCounterId(userId, uid, value2);
+        assertThat(mRecoverableKeyStoreDb.getCounterId(userId, uid)).isEqualTo(
+                value2);
+    }
+
+    @Test
+    public void setSnapshotVersion_defaultValueAndTwoUpdates() throws Exception {
+        int userId = 12;
+        int uid = 10009;
+        long value1 = 111L;
+        long value2 = 222L;
+        assertThat(mRecoverableKeyStoreDb.getSnapshotVersion(userId, uid)).isNull();
+        mRecoverableKeyStoreDb.setSnapshotVersion(userId, uid, value1);
+        assertThat(mRecoverableKeyStoreDb.getSnapshotVersion(userId, uid)).isEqualTo(
+                value1);
+        mRecoverableKeyStoreDb.setSnapshotVersion(userId, uid, value2);
+        assertThat(mRecoverableKeyStoreDb.getSnapshotVersion(userId, uid)).isEqualTo(
+                value2);
+    }
+
+    @Test
+    public void setShouldCreateSnapshot_defaultValueAndTwoUpdates() throws Exception {
+        int userId = 12;
+        int uid = 10009;
+        boolean value1 = true;
+        boolean value2 = false;
+        assertThat(mRecoverableKeyStoreDb.getShouldCreateSnapshot(userId, uid)).isEqualTo(false);
+        mRecoverableKeyStoreDb.setShouldCreateSnapshot(userId, uid, value1);
+        assertThat(mRecoverableKeyStoreDb.getShouldCreateSnapshot(userId, uid)).isEqualTo(value1);
+        mRecoverableKeyStoreDb.setShouldCreateSnapshot(userId, uid, value2);
+        assertThat(mRecoverableKeyStoreDb.getShouldCreateSnapshot(userId, uid)).isEqualTo(
+                value2);
+    }
+
+    @Test
     public void setRecoveryServiceMetadataEntry_allowsAUserToHaveTwoUids() throws Exception {
         int userId = 12;
         int uid1 = 10009;