Merge changes Ifd189613,Id603dabc,Id9cd0a57

* changes:
  Add integration test for encrypted KV B&R
  Import EncryptedKvRestoreTask
  Import EncryptedKvBackupTask

9 files changed, 1147 insertions, 15 deletions

diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/EncryptedKvBackupTask.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/EncryptedKvBackupTask.java
new file mode 100644
index 000000000000..619438c7f6fe
--- /dev/null
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/EncryptedKvBackupTask.java
@@ -0,0 +1,244 @@
+/*
+ * Copyright (C) 2019 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.backup.encryption.tasks;
+
+import android.annotation.Nullable;
+import android.app.backup.BackupDataInput;
+import android.content.Context;
+import android.os.ParcelFileDescriptor;
+import android.security.keystore.recovery.InternalRecoveryServiceException;
+import android.security.keystore.recovery.LockScreenRequiredException;
+import android.util.Pair;
+import android.util.Slog;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.backup.encryption.CryptoSettings;
+import com.android.server.backup.encryption.chunking.ProtoStore;
+import com.android.server.backup.encryption.client.CryptoBackupServer;
+import com.android.server.backup.encryption.keys.RecoverableKeyStoreSecondaryKey;
+import com.android.server.backup.encryption.keys.RecoverableKeyStoreSecondaryKeyManager;
+import com.android.server.backup.encryption.keys.TertiaryKeyManager;
+import com.android.server.backup.encryption.keys.TertiaryKeyRotationScheduler;
+import com.android.server.backup.encryption.protos.nano.ChunksMetadataProto;
+import com.android.server.backup.encryption.protos.nano.KeyValueListingProto;
+
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.security.GeneralSecurityException;
+import java.security.InvalidKeyException;
+import java.security.SecureRandom;
+import java.security.UnrecoverableKeyException;
+import java.util.Optional;
+
+// TODO(b/141975695): Create a base class for EncryptedKvBackupTask and EncryptedFullBackupTask.
+/** Performs encrypted key value backup, handling rotating the tertiary key as necessary. */
+public class EncryptedKvBackupTask {
+    private static final String TAG = "EncryptedKvBackupTask";
+
+    private final TertiaryKeyManager mTertiaryKeyManager;
+    private final RecoverableKeyStoreSecondaryKey mSecondaryKey;
+    private final ProtoStore<KeyValueListingProto.KeyValueListing> mKeyValueListingStore;
+    private final ProtoStore<ChunksMetadataProto.ChunkListing> mChunkListingStore;
+    private final KvBackupEncrypter mKvBackupEncrypter;
+    private final EncryptedBackupTask mEncryptedBackupTask;
+    private final String mPackageName;
+
+    /** Constructs new instances of {@link EncryptedKvBackupTask}. */
+    public static class EncryptedKvBackupTaskFactory {
+        /**
+         * Creates a new instance.
+         *
+         * <p>Either initializes encrypted backup or loads an existing secondary key as necessary.
+         *
+         * @param cryptoSettings to load secondary key state from
+         * @param fileDescriptor to read the backup data from
+         */
+        public EncryptedKvBackupTask newInstance(
+                Context context,
+                SecureRandom secureRandom,
+                CryptoBackupServer cryptoBackupServer,
+                CryptoSettings cryptoSettings,
+                RecoverableKeyStoreSecondaryKeyManager
+                                .RecoverableKeyStoreSecondaryKeyManagerProvider
+                        recoverableSecondaryKeyManagerProvider,
+                ParcelFileDescriptor fileDescriptor,
+                String packageName)
+                throws IOException, UnrecoverableKeyException, LockScreenRequiredException,
+                        InternalRecoveryServiceException, InvalidKeyException {
+            RecoverableKeyStoreSecondaryKey secondaryKey =
+                    new InitializeRecoverableSecondaryKeyTask(
+                                    context,
+                                    cryptoSettings,
+                                    recoverableSecondaryKeyManagerProvider.get(),
+                                    cryptoBackupServer)
+                            .run();
+            KvBackupEncrypter backupEncrypter =
+                    new KvBackupEncrypter(new BackupDataInput(fileDescriptor.getFileDescriptor()));
+            TertiaryKeyManager tertiaryKeyManager =
+                    new TertiaryKeyManager(
+                            context,
+                            secureRandom,
+                            TertiaryKeyRotationScheduler.getInstance(context),
+                            secondaryKey,
+                            packageName);
+
+            return new EncryptedKvBackupTask(
+                    tertiaryKeyManager,
+                    ProtoStore.createKeyValueListingStore(context),
+                    secondaryKey,
+                    ProtoStore.createChunkListingStore(context),
+                    backupEncrypter,
+                    new EncryptedBackupTask(
+                            cryptoBackupServer, secureRandom, packageName, backupEncrypter),
+                    packageName);
+        }
+    }
+
+    @VisibleForTesting
+    EncryptedKvBackupTask(
+            TertiaryKeyManager tertiaryKeyManager,
+            ProtoStore<KeyValueListingProto.KeyValueListing> keyValueListingStore,
+            RecoverableKeyStoreSecondaryKey secondaryKey,
+            ProtoStore<ChunksMetadataProto.ChunkListing> chunkListingStore,
+            KvBackupEncrypter kvBackupEncrypter,
+            EncryptedBackupTask encryptedBackupTask,
+            String packageName) {
+        mTertiaryKeyManager = tertiaryKeyManager;
+        mSecondaryKey = secondaryKey;
+        mKeyValueListingStore = keyValueListingStore;
+        mChunkListingStore = chunkListingStore;
+        mKvBackupEncrypter = kvBackupEncrypter;
+        mEncryptedBackupTask = encryptedBackupTask;
+        mPackageName = packageName;
+    }
+
+    /**
+     * Reads backup data from the file descriptor provided in the construtor, encrypts it and
+     * uploads it to the server.
+     *
+     * <p>The {@code incremental} flag indicates if the backup data provided is incremental or a
+     * complete set. Incremental backup is not possible if no previous crypto state exists, or the
+     * tertiary key must be rotated in the next backup. If the caller requests incremental backup
+     * but it is not possible, then the backup will not start and this method will throw {@link
+     * NonIncrementalBackupRequiredException}.
+     *
+     * <p>TODO(b/70704456): Update return code to indicate that we require non-incremental backup.
+     *
+     * @param incremental {@code true} if the data provided is a diff from the previous backup,
+     *     {@code false} if it is a complete set
+     * @throws NonIncrementalBackupRequiredException if the caller provides an incremental backup but the task
+     *     requires non-incremental backup
+     */
+    public void performBackup(boolean incremental)
+            throws GeneralSecurityException, IOException, NoSuchMethodException,
+            InstantiationException, IllegalAccessException, InvocationTargetException,
+            NonIncrementalBackupRequiredException {
+        if (mTertiaryKeyManager.wasKeyRotated()) {
+            Slog.d(TAG, "Tertiary key is new so clearing package state.");
+            deleteListings(mPackageName);
+        }
+
+        Optional<Pair<KeyValueListingProto.KeyValueListing, ChunksMetadataProto.ChunkListing>>
+                oldListings = getListingsAndEnsureConsistency(mPackageName);
+
+        if (oldListings.isPresent() && !incremental) {
+            Slog.d(
+                    TAG,
+                    "Non-incremental backup requested but incremental state existed, clearing it");
+            deleteListings(mPackageName);
+            oldListings = Optional.empty();
+        }
+
+        if (!oldListings.isPresent() && incremental) {
+            // If we don't have any state then we require a non-incremental backup, but this backup
+            // is incremental.
+            throw new NonIncrementalBackupRequiredException();
+        }
+
+        if (oldListings.isPresent()) {
+            mKvBackupEncrypter.setOldKeyValueListing(oldListings.get().first);
+        }
+
+        ChunksMetadataProto.ChunkListing newChunkListing;
+        if (oldListings.isPresent()) {
+            Slog.v(TAG, "Old listings existed, performing incremental backup");
+            newChunkListing =
+                    mEncryptedBackupTask.performIncrementalBackup(
+                            mTertiaryKeyManager.getKey(),
+                            mTertiaryKeyManager.getWrappedKey(),
+                            oldListings.get().second);
+        } else {
+            Slog.v(TAG, "Old listings did not exist, performing non-incremental backup");
+            // kv backups don't use this salt because they don't involve content-defined chunking.
+            byte[] fingerprintMixerSalt = null;
+            newChunkListing =
+                    mEncryptedBackupTask.performNonIncrementalBackup(
+                            mTertiaryKeyManager.getKey(),
+                            mTertiaryKeyManager.getWrappedKey(),
+                            fingerprintMixerSalt);
+        }
+
+        Slog.v(TAG, "Backup and upload succeeded, saving new listings");
+        saveListings(mPackageName, mKvBackupEncrypter.getNewKeyValueListing(), newChunkListing);
+    }
+
+    private Optional<Pair<KeyValueListingProto.KeyValueListing, ChunksMetadataProto.ChunkListing>>
+            getListingsAndEnsureConsistency(String packageName)
+                    throws IOException, InvocationTargetException, NoSuchMethodException,
+                            InstantiationException, IllegalAccessException {
+        Optional<KeyValueListingProto.KeyValueListing> keyValueListing =
+                mKeyValueListingStore.loadProto(packageName);
+        Optional<ChunksMetadataProto.ChunkListing> chunkListing =
+                mChunkListingStore.loadProto(packageName);
+
+        // Normally either both protos exist or neither exist, but we correct this just in case.
+        boolean bothPresent = keyValueListing.isPresent() && chunkListing.isPresent();
+        if (!bothPresent) {
+            Slog.d(
+                    TAG,
+                    "Both listing were not present, clearing state, key value="
+                            + keyValueListing.isPresent()
+                            + ", chunk="
+                            + chunkListing.isPresent());
+            deleteListings(packageName);
+            return Optional.empty();
+        }
+
+        return Optional.of(Pair.create(keyValueListing.get(), chunkListing.get()));
+    }
+
+    private void saveListings(
+            String packageName,
+            KeyValueListingProto.KeyValueListing keyValueListing,
+            ChunksMetadataProto.ChunkListing chunkListing) {
+        try {
+            mKeyValueListingStore.saveProto(packageName, keyValueListing);
+            mChunkListingStore.saveProto(packageName, chunkListing);
+        } catch (IOException e) {
+            // If a problem occurred while saving either listing then they may be inconsistent, so
+            // delete
+            // both.
+            Slog.w(TAG, "Unable to save listings, deleting both for consistency", e);
+            deleteListings(packageName);
+        }
+    }
+
+    private void deleteListings(String packageName) {
+        mKeyValueListingStore.deleteProto(packageName);
+        mChunkListingStore.deleteProto(packageName);
+    }
+}
diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/EncryptedKvRestoreTask.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/EncryptedKvRestoreTask.java
new file mode 100644
index 000000000000..12b44590ebe6
--- /dev/null
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/EncryptedKvRestoreTask.java
@@ -0,0 +1,139 @@
+/*
+ * Copyright (C) 2019 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.backup.encryption.tasks;
+
+import static com.android.internal.util.Preconditions.checkArgument;
+
+import android.app.backup.BackupDataOutput;
+import android.content.Context;
+import android.os.ParcelFileDescriptor;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.backup.encryption.FullRestoreDownloader;
+import com.android.server.backup.encryption.chunking.ChunkHasher;
+import com.android.server.backup.encryption.keys.RecoverableKeyStoreSecondaryKeyManager;
+import com.android.server.backup.encryption.keys.RestoreKeyFetcher;
+import com.android.server.backup.encryption.kv.DecryptedChunkKvOutput;
+import com.android.server.backup.encryption.protos.nano.KeyValuePairProto;
+import com.android.server.backup.encryption.protos.nano.WrappedKeyProto;
+
+import java.io.File;
+import java.io.IOException;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
+import java.security.KeyException;
+import java.security.NoSuchAlgorithmException;
+
+import javax.crypto.BadPaddingException;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
+import javax.crypto.SecretKey;
+import javax.crypto.ShortBufferException;
+
+/**
+ * Performs a key value restore by downloading the backup set, decrypting it and writing it to the
+ * file provided by backup manager.
+ */
+public class EncryptedKvRestoreTask {
+    private static final String ENCRYPTED_FILE_NAME = "encrypted_kv";
+
+    private final File mTemporaryFolder;
+    private final ChunkHasher mChunkHasher;
+    private final FullRestoreToFileTask mFullRestoreToFileTask;
+    private final BackupFileDecryptorTask mBackupFileDecryptorTask;
+
+    /** Constructs new instances of the task. */
+    public static class EncryptedKvRestoreTaskFactory {
+        /**
+         * Constructs a new instance.
+         *
+         * <p>Fetches the appropriate secondary key and uses this to unwrap the tertiary key. Stores
+         * temporary files in {@link Context#getFilesDir()}.
+         */
+        public EncryptedKvRestoreTask newInstance(
+                Context context,
+                RecoverableKeyStoreSecondaryKeyManager
+                                .RecoverableKeyStoreSecondaryKeyManagerProvider
+                        recoverableSecondaryKeyManagerProvider,
+                FullRestoreDownloader fullRestoreDownloader,
+                String secondaryKeyAlias,
+                WrappedKeyProto.WrappedKey wrappedTertiaryKey)
+                throws EncryptedRestoreException, NoSuchAlgorithmException, NoSuchPaddingException,
+                        KeyException, InvalidAlgorithmParameterException {
+            SecretKey tertiaryKey =
+                    RestoreKeyFetcher.unwrapTertiaryKey(
+                            recoverableSecondaryKeyManagerProvider,
+                            secondaryKeyAlias,
+                            wrappedTertiaryKey);
+
+            return new EncryptedKvRestoreTask(
+                    context.getFilesDir(),
+                    new ChunkHasher(tertiaryKey),
+                    new FullRestoreToFileTask(fullRestoreDownloader),
+                    new BackupFileDecryptorTask(tertiaryKey));
+        }
+    }
+
+    @VisibleForTesting
+    EncryptedKvRestoreTask(
+            File temporaryFolder,
+            ChunkHasher chunkHasher,
+            FullRestoreToFileTask fullRestoreToFileTask,
+            BackupFileDecryptorTask backupFileDecryptorTask) {
+        checkArgument(
+                temporaryFolder.isDirectory(), "Temporary folder must be an existing directory");
+
+        mTemporaryFolder = temporaryFolder;
+        mChunkHasher = chunkHasher;
+        mFullRestoreToFileTask = fullRestoreToFileTask;
+        mBackupFileDecryptorTask = backupFileDecryptorTask;
+    }
+
+    /**
+     * Runs the restore, writing the pairs in lexicographical order to the given file descriptor.
+     *
+     * <p>This will block for the duration of the restore.
+     *
+     * @throws EncryptedRestoreException if there is a problem decrypting or verifying the backup
+     */
+    public void getRestoreData(ParcelFileDescriptor output)
+            throws IOException, EncryptedRestoreException, BadPaddingException,
+                    InvalidAlgorithmParameterException, NoSuchAlgorithmException,
+                    IllegalBlockSizeException, ShortBufferException, InvalidKeyException {
+        File encryptedFile = new File(mTemporaryFolder, ENCRYPTED_FILE_NAME);
+        try {
+            downloadDecryptAndWriteBackup(encryptedFile, output);
+        } finally {
+            encryptedFile.delete();
+        }
+    }
+
+    private void downloadDecryptAndWriteBackup(File encryptedFile, ParcelFileDescriptor output)
+            throws EncryptedRestoreException, IOException, BadPaddingException, InvalidKeyException,
+                    NoSuchAlgorithmException, IllegalBlockSizeException, ShortBufferException,
+                    InvalidAlgorithmParameterException {
+        mFullRestoreToFileTask.restoreToFile(encryptedFile);
+        DecryptedChunkKvOutput decryptedChunkKvOutput = new DecryptedChunkKvOutput(mChunkHasher);
+        mBackupFileDecryptorTask.decryptFile(encryptedFile, decryptedChunkKvOutput);
+
+        BackupDataOutput backupDataOutput = new BackupDataOutput(output.getFileDescriptor());
+        for (KeyValuePairProto.KeyValuePair pair : decryptedChunkKvOutput.getPairs()) {
+            backupDataOutput.writeEntityHeader(pair.key, pair.value.length);
+            backupDataOutput.writeEntityData(pair.value, pair.value.length);
+        }
+    }
+}
diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/NonIncrementalBackupRequiredException.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/NonIncrementalBackupRequiredException.java
new file mode 100644
index 000000000000..a3eda7d1270f
--- /dev/null
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/NonIncrementalBackupRequiredException.java
@@ -0,0 +1,25 @@
+/*
+ * Copyright (C) 2019 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.backup.encryption.tasks;
+
+// TODO(141840878): Update documentation.
+/**
+ * Exception thrown when the framework provides an incremental backup but the transport requires a
+ * non-incremental backup.
+ */
+public class NonIncrementalBackupRequiredException extends Exception {}
diff --git a/packages/BackupEncryption/test/robolectric-integration/Android.bp b/packages/BackupEncryption/test/robolectric-integration/Android.bp
index f696278ab967..67365df4b28f 100644
--- a/packages/BackupEncryption/test/robolectric-integration/Android.bp
+++ b/packages/BackupEncryption/test/robolectric-integration/Android.bp
@@ -23,6 +23,7 @@ android_robolectric_test {
         "platform-test-annotations",
         "testng",
         "truth-prebuilt",
+        "BackupEncryptionRoboTests",
     ],
     static_libs: [
         "androidx.test.core",
diff --git a/packages/BackupEncryption/test/robolectric-integration/src/com/android/server/backup/encryption/RoundTripTest.java b/packages/BackupEncryption/test/robolectric-integration/src/com/android/server/backup/encryption/RoundTripTest.java
index 8ec68fdf822d..a432d91828cf 100644
--- a/packages/BackupEncryption/test/robolectric-integration/src/com/android/server/backup/encryption/RoundTripTest.java
+++ b/packages/BackupEncryption/test/robolectric-integration/src/com/android/server/backup/encryption/RoundTripTest.java
@@ -19,21 +19,35 @@ package com.android.server.backup.encryption;
 import static com.google.common.truth.Truth.assertThat;
 
 import android.content.Context;
+import android.os.ParcelFileDescriptor;
+import android.security.keystore.recovery.InternalRecoveryServiceException;
+import android.security.keystore.recovery.RecoveryController;
 
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.server.backup.encryption.client.CryptoBackupServer;
+import com.android.server.backup.encryption.keys.KeyWrapUtils;
 import com.android.server.backup.encryption.keys.RecoverableKeyStoreSecondaryKey;
+import com.android.server.backup.encryption.keys.RecoverableKeyStoreSecondaryKeyManager;
 import com.android.server.backup.encryption.keys.TertiaryKeyManager;
 import com.android.server.backup.encryption.keys.TertiaryKeyRotationScheduler;
 import com.android.server.backup.encryption.protos.nano.WrappedKeyProto;
 import com.android.server.backup.encryption.tasks.EncryptedFullBackupTask;
 import com.android.server.backup.encryption.tasks.EncryptedFullRestoreTask;
+import com.android.server.backup.encryption.tasks.EncryptedKvBackupTask;
+import com.android.server.backup.encryption.tasks.EncryptedKvRestoreTask;
+import com.android.server.testing.shadows.DataEntity;
+import com.android.server.testing.shadows.ShadowBackupDataInput;
+import com.android.server.testing.shadows.ShadowBackupDataOutput;
+import com.android.server.testing.shadows.ShadowRecoveryController;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -42,15 +56,29 @@ import java.security.InvalidAlgorithmParameterException;
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
+import java.util.Optional;
 import java.util.Map;
+import java.util.Set;
 
 import javax.crypto.IllegalBlockSizeException;
 import javax.crypto.KeyGenerator;
 import javax.crypto.NoSuchPaddingException;
 import javax.crypto.SecretKey;
 
+@Config(
+        shadows = {
+            ShadowBackupDataInput.class,
+            ShadowBackupDataOutput.class,
+            ShadowRecoveryController.class
+        })
 @RunWith(RobolectricTestRunner.class)
 public class RoundTripTest {
+    private static final DataEntity[] KEY_VALUE_DATA = {
+        new DataEntity("test_key_1", "test_value_1"),
+        new DataEntity("test_key_2", "test_value_2"),
+        new DataEntity("test_key_3", "test_value_3")
+    };
+
     /** Amount of data we want to round trip in this test */
     private static final int TEST_DATA_SIZE = 1024 * 1024; // 1MB
 
@@ -59,6 +87,7 @@ public class RoundTripTest {
 
     /** Key parameters used for the secondary encryption key */
     private static final String KEY_ALGORITHM = "AES";
+
     private static final int KEY_SIZE_BITS = 256;
 
     /** Package name for our test package */
@@ -77,25 +106,82 @@ public class RoundTripTest {
     private RecoverableKeyStoreSecondaryKey mSecondaryKey;
 
     /** Source of random material which is considered non-predictable in its' generation */
-    private SecureRandom mSecureRandom = new SecureRandom();
+    private final SecureRandom mSecureRandom = new SecureRandom();
+
+    private RecoverableKeyStoreSecondaryKeyManager.RecoverableKeyStoreSecondaryKeyManagerProvider
+            mSecondaryKeyManagerProvider;
+    private DummyServer mDummyServer;
+    private RecoveryController mRecoveryController;
+
+    @Mock private ParcelFileDescriptor mParcelFileDescriptor;
 
     @Before
-    public void setUp() throws NoSuchAlgorithmException {
+    public void setUp() throws NoSuchAlgorithmException, InternalRecoveryServiceException {
+        MockitoAnnotations.initMocks(this);
+
+        ShadowBackupDataInput.reset();
+        ShadowBackupDataOutput.reset();
+
         mContext = ApplicationProvider.getApplicationContext();
         mSecondaryKey = new RecoverableKeyStoreSecondaryKey(TEST_KEY_ALIAS, generateAesKey());
+        mDummyServer = new DummyServer();
+        mSecondaryKeyManagerProvider =
+                () ->
+                        new RecoverableKeyStoreSecondaryKeyManager(
+                                RecoveryController.getInstance(mContext), mSecureRandom);
+
         fillBuffer(mOriginalData);
     }
 
     @Test
-    public void testRoundTrip() throws Exception {
-        byte[] backupData = performBackup(mOriginalData);
+    public void testFull_nonIncrementalBackupAndRestoreAreSuccessful() throws Exception {
+        byte[] backupData = performFullBackup(mOriginalData);
         assertThat(backupData).isNotEqualTo(mOriginalData);
-        byte[] restoredData = performRestore(backupData);
+        byte[] restoredData = performFullRestore(backupData);
         assertThat(restoredData).isEqualTo(mOriginalData);
     }
 
-    /** Perform a backup and return the backed-up representation of the data */
-    private byte[] performBackup(byte[] backupData) throws Exception {
+    @Test
+    public void testKeyValue_nonIncrementalBackupAndRestoreAreSuccessful() throws Exception {
+        byte[] backupData = performNonIncrementalKeyValueBackup(KEY_VALUE_DATA);
+
+        // Get the secondary key used to do backup.
+        Optional<RecoverableKeyStoreSecondaryKey> secondaryKey =
+                mSecondaryKeyManagerProvider.get().get(mDummyServer.mSecondaryKeyAlias);
+        assertThat(secondaryKey.isPresent()).isTrue();
+
+        Set<DataEntity> restoredData = performKeyValueRestore(backupData, secondaryKey.get());
+
+        assertThat(restoredData).containsExactly(KEY_VALUE_DATA).inOrder();
+    }
+
+    /** Perform a key/value backup and return the backed-up representation of the data */
+    private byte[] performNonIncrementalKeyValueBackup(DataEntity[] backupData)
+            throws Exception {
+        // Populate test key/value data.
+        for (DataEntity entity : backupData) {
+            ShadowBackupDataInput.addEntity(entity);
+        }
+
+        EncryptedKvBackupTask.EncryptedKvBackupTaskFactory backupTaskFactory =
+                new EncryptedKvBackupTask.EncryptedKvBackupTaskFactory();
+        EncryptedKvBackupTask backupTask =
+                backupTaskFactory.newInstance(
+                        mContext,
+                        mSecureRandom,
+                        mDummyServer,
+                        CryptoSettings.getInstance(mContext),
+                        mSecondaryKeyManagerProvider,
+                        mParcelFileDescriptor,
+                        TEST_PACKAGE_NAME);
+
+        backupTask.performBackup(/* incremental */ false);
+
+        return mDummyServer.mStoredData;
+    }
+
+    /** Perform a full backup and return the backed-up representation of the data */
+    private byte[] performFullBackup(byte[] backupData) throws Exception {
         DummyServer dummyServer = new DummyServer();
         EncryptedFullBackupTask backupTask =
                 EncryptedFullBackupTask.newInstance(
@@ -109,8 +195,24 @@ public class RoundTripTest {
         return dummyServer.mStoredData;
     }
 
-    /** Perform a restore and resturn the bytes obtained from the restore process */
-    private byte[] performRestore(byte[] backupData)
+    private Set<DataEntity> performKeyValueRestore(
+            byte[] backupData, RecoverableKeyStoreSecondaryKey secondaryKey) throws Exception {
+        EncryptedKvRestoreTask.EncryptedKvRestoreTaskFactory restoreTaskFactory =
+                new EncryptedKvRestoreTask.EncryptedKvRestoreTaskFactory();
+        EncryptedKvRestoreTask restoreTask =
+                restoreTaskFactory.newInstance(
+                        mContext,
+                        mSecondaryKeyManagerProvider,
+                        new FakeFullRestoreDownloader(backupData),
+                        secondaryKey.getAlias(),
+                        KeyWrapUtils.wrap(
+                                secondaryKey.getSecretKey(), getTertiaryKey(secondaryKey)));
+        restoreTask.getRestoreData(mParcelFileDescriptor);
+        return ShadowBackupDataOutput.getEntities();
+    }
+
+    /** Perform a full restore and return the bytes obtained from the restore process */
+    private byte[] performFullRestore(byte[] backupData)
             throws IOException, NoSuchAlgorithmException, NoSuchPaddingException,
                     InvalidAlgorithmParameterException, InvalidKeyException,
                     IllegalBlockSizeException {
@@ -118,7 +220,9 @@ public class RoundTripTest {
 
         EncryptedFullRestoreTask restoreTask =
                 EncryptedFullRestoreTask.newInstance(
-                        mContext, new FakeFullRestoreDownloader(backupData), getTertiaryKey());
+                        mContext,
+                        new FakeFullRestoreDownloader(backupData),
+                        getTertiaryKey(mSecondaryKey));
 
         byte[] buffer = new byte[READ_BUFFER_SIZE];
         int bytesRead = restoreTask.readNextChunk(buffer);
@@ -131,7 +235,7 @@ public class RoundTripTest {
     }
 
     /** Get the tertiary key for our test package from the key manager */
-    private SecretKey getTertiaryKey()
+    private SecretKey getTertiaryKey(RecoverableKeyStoreSecondaryKey secondaryKey)
             throws IllegalBlockSizeException, InvalidAlgorithmParameterException,
                     NoSuchAlgorithmException, IOException, NoSuchPaddingException,
                     InvalidKeyException {
@@ -140,7 +244,7 @@ public class RoundTripTest {
                         mContext,
                         mSecureRandom,
                         TertiaryKeyRotationScheduler.getInstance(mContext),
-                        mSecondaryKey,
+                        secondaryKey,
                         TEST_PACKAGE_NAME);
         return tertiaryKeyManager.getKey();
     }
@@ -162,13 +266,13 @@ public class RoundTripTest {
     }
 
     /**
-     * Dummy backup data endpoint. This stores the data so we can use it
-     * in subsequent test steps.
+     * Dummy backup data endpoint. This stores the data so we can use it in subsequent test steps.
      */
     private static class DummyServer implements CryptoBackupServer {
         private static final String DUMMY_DOC_ID = "DummyDoc";
 
         byte[] mStoredData = null;
+        String mSecondaryKeyAlias;
 
         @Override
         public String uploadIncrementalBackup(
@@ -190,7 +294,7 @@ public class RoundTripTest {
         @Override
         public void setActiveSecondaryKeyAlias(
                 String keyAlias, Map<String, WrappedKeyProto.WrappedKey> tertiaryKeys) {
-            throw new RuntimeException("Not Implemented");
+            mSecondaryKeyAlias = keyAlias;
         }
     }
 
diff --git a/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/EncryptedKvBackupTaskTest.java b/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/EncryptedKvBackupTaskTest.java
new file mode 100644
index 000000000000..fa4fef50ac1a
--- /dev/null
+++ b/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/EncryptedKvBackupTaskTest.java
@@ -0,0 +1,356 @@
+/*
+ * Copyright (C) 2019 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.backup.encryption.tasks;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+
+import android.app.Application;
+import android.util.Pair;
+
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.server.backup.encryption.chunk.ChunkHash;
+import com.android.server.backup.encryption.chunking.ProtoStore;
+import com.android.server.backup.encryption.keys.RecoverableKeyStoreSecondaryKey;
+import com.android.server.backup.encryption.keys.TertiaryKeyManager;
+import com.android.server.backup.encryption.kv.KeyValueListingBuilder;
+import com.android.server.backup.encryption.protos.nano.ChunksMetadataProto;
+import com.android.server.backup.encryption.protos.nano.KeyValueListingProto;
+import com.android.server.backup.encryption.protos.nano.WrappedKeyProto;
+import com.android.server.backup.testing.CryptoTestUtils;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Map.Entry;
+
+import javax.crypto.SecretKey;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+
+@RunWith(RobolectricTestRunner.class)
+public class EncryptedKvBackupTaskTest {
+    private static final boolean INCREMENTAL = true;
+    private static final boolean NON_INCREMENTAL = false;
+
+    private static final String TEST_PACKAGE_1 = "com.example.app1";
+    private static final String TEST_KEY_1 = "key_1";
+    private static final String TEST_KEY_2 = "key_2";
+    private static final ChunkHash TEST_HASH_1 =
+            new ChunkHash(Arrays.copyOf(new byte[] {1}, ChunkHash.HASH_LENGTH_BYTES));
+    private static final ChunkHash TEST_HASH_2 =
+            new ChunkHash(Arrays.copyOf(new byte[] {2}, ChunkHash.HASH_LENGTH_BYTES));
+    private static final int TEST_LENGTH_1 = 200;
+    private static final int TEST_LENGTH_2 = 300;
+
+    @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
+
+    @Captor private ArgumentCaptor<ChunksMetadataProto.ChunkListing> mChunkListingCaptor;
+
+    @Mock private TertiaryKeyManager mTertiaryKeyManager;
+    @Mock private RecoverableKeyStoreSecondaryKey mSecondaryKey;
+    @Mock private ProtoStore<KeyValueListingProto.KeyValueListing> mKeyValueListingStore;
+    @Mock private ProtoStore<ChunksMetadataProto.ChunkListing> mChunkListingStore;
+    @Mock private KvBackupEncrypter mKvBackupEncrypter;
+    @Mock private EncryptedBackupTask mEncryptedBackupTask;
+    @Mock private SecretKey mTertiaryKey;
+
+    private WrappedKeyProto.WrappedKey mWrappedTertiaryKey;
+    private KeyValueListingProto.KeyValueListing mNewKeyValueListing;
+    private ChunksMetadataProto.ChunkListing mNewChunkListing;
+    private EncryptedKvBackupTask mTask;
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        Application application = ApplicationProvider.getApplicationContext();
+        mKeyValueListingStore = ProtoStore.createKeyValueListingStore(application);
+        mChunkListingStore = ProtoStore.createChunkListingStore(application);
+
+        mWrappedTertiaryKey = new WrappedKeyProto.WrappedKey();
+
+        when(mTertiaryKeyManager.wasKeyRotated()).thenReturn(false);
+        when(mTertiaryKeyManager.getKey()).thenReturn(mTertiaryKey);
+        when(mTertiaryKeyManager.getWrappedKey()).thenReturn(mWrappedTertiaryKey);
+
+        mNewKeyValueListing =
+                createKeyValueListing(
+                        CryptoTestUtils.mapOf(
+                                new Pair<>(TEST_KEY_1, TEST_HASH_1),
+                                new Pair<>(TEST_KEY_2, TEST_HASH_2)));
+        mNewChunkListing =
+                createChunkListing(
+                        CryptoTestUtils.mapOf(
+                                new Pair<>(TEST_HASH_1, TEST_LENGTH_1),
+                                new Pair<>(TEST_HASH_2, TEST_LENGTH_2)));
+        when(mKvBackupEncrypter.getNewKeyValueListing()).thenReturn(mNewKeyValueListing);
+        when(mEncryptedBackupTask.performIncrementalBackup(
+                        eq(mTertiaryKey), eq(mWrappedTertiaryKey), any()))
+                .thenReturn(mNewChunkListing);
+        when(mEncryptedBackupTask.performNonIncrementalBackup(
+                        eq(mTertiaryKey), eq(mWrappedTertiaryKey), any()))
+                .thenReturn(mNewChunkListing);
+
+        mTask =
+                new EncryptedKvBackupTask(
+                        mTertiaryKeyManager,
+                        mKeyValueListingStore,
+                        mSecondaryKey,
+                        mChunkListingStore,
+                        mKvBackupEncrypter,
+                        mEncryptedBackupTask,
+                        TEST_PACKAGE_1);
+    }
+
+    @Test
+    public void testPerformBackup_rotationRequired_deletesListings() throws Exception {
+        mKeyValueListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createKeyValueListing(CryptoTestUtils.mapOf(new Pair<>(TEST_KEY_1, TEST_HASH_1))));
+        mChunkListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createChunkListing(CryptoTestUtils.mapOf(new Pair<>(TEST_HASH_1, TEST_LENGTH_1))));
+
+        when(mTertiaryKeyManager.wasKeyRotated()).thenReturn(true);
+        // Throw an IOException so it aborts before saving the new listings.
+        when(mEncryptedBackupTask.performNonIncrementalBackup(any(), any(), any()))
+                .thenThrow(IOException.class);
+
+        assertThrows(IOException.class, () -> mTask.performBackup(NON_INCREMENTAL));
+
+        assertFalse(mKeyValueListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+        assertFalse(mChunkListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+    }
+
+    @Test
+    public void testPerformBackup_rotationRequiredButIncremental_throws() throws Exception {
+        mKeyValueListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createKeyValueListing(CryptoTestUtils.mapOf(new Pair<>(TEST_KEY_1, TEST_HASH_1))));
+        mChunkListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createChunkListing(CryptoTestUtils.mapOf(new Pair<>(TEST_HASH_1, TEST_LENGTH_1))));
+
+        when(mTertiaryKeyManager.wasKeyRotated()).thenReturn(true);
+
+        assertThrows(NonIncrementalBackupRequiredException.class,
+                () -> mTask.performBackup(INCREMENTAL));
+    }
+
+    @Test
+    public void testPerformBackup_rotationRequiredAndNonIncremental_performsNonIncrementalBackup()
+            throws Exception {
+        mKeyValueListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createKeyValueListing(CryptoTestUtils.mapOf(new Pair<>(TEST_KEY_1, TEST_HASH_1))));
+        mChunkListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createChunkListing(CryptoTestUtils.mapOf(new Pair<>(TEST_HASH_1, TEST_LENGTH_1))));
+
+        when(mTertiaryKeyManager.wasKeyRotated()).thenReturn(true);
+
+        mTask.performBackup(NON_INCREMENTAL);
+
+        verify(mEncryptedBackupTask)
+                .performNonIncrementalBackup(eq(mTertiaryKey), eq(mWrappedTertiaryKey), any());
+    }
+
+    @Test
+    public void testPerformBackup_existingStateButNonIncremental_deletesListings() throws Exception {
+        mKeyValueListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createKeyValueListing(CryptoTestUtils.mapOf(new Pair<>(TEST_KEY_1, TEST_HASH_1))));
+        mChunkListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createChunkListing(CryptoTestUtils.mapOf(new Pair<>(TEST_HASH_1, TEST_LENGTH_1))));
+
+        // Throw an IOException so it aborts before saving the new listings.
+        when(mEncryptedBackupTask.performNonIncrementalBackup(any(), any(), any()))
+                .thenThrow(IOException.class);
+
+        assertThrows(IOException.class, () -> mTask.performBackup(NON_INCREMENTAL));
+
+        assertFalse(mKeyValueListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+        assertFalse(mChunkListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+    }
+
+    @Test
+    public void testPerformBackup_keyValueListingMissing_deletesChunkListingAndPerformsNonIncremental()
+            throws Exception {
+        mChunkListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createChunkListing(CryptoTestUtils.mapOf(new Pair<>(TEST_HASH_1, TEST_LENGTH_1))));
+
+        // Throw an IOException so it aborts before saving the new listings.
+        when(mEncryptedBackupTask.performNonIncrementalBackup(any(), any(), any()))
+                .thenThrow(IOException.class);
+
+        assertThrows(IOException.class, () -> mTask.performBackup(NON_INCREMENTAL));
+
+        verify(mEncryptedBackupTask).performNonIncrementalBackup(any(), any(), any());
+        assertFalse(mKeyValueListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+        assertFalse(mChunkListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+    }
+
+    @Test
+    public void testPerformBackup_chunkListingMissing_deletesKeyValueListingAndPerformsNonIncremental()
+            throws Exception {
+        mKeyValueListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createKeyValueListing(CryptoTestUtils.mapOf(new Pair<>(TEST_KEY_1, TEST_HASH_1))));
+
+        // Throw an IOException so it aborts before saving the new listings.
+        when(mEncryptedBackupTask.performNonIncrementalBackup(any(), any(), any()))
+                .thenThrow(IOException.class);
+
+        assertThrows(IOException.class, () -> mTask.performBackup(NON_INCREMENTAL));
+
+        verify(mEncryptedBackupTask).performNonIncrementalBackup(any(), any(), any());
+        assertFalse(mKeyValueListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+        assertFalse(mChunkListingStore.loadProto(TEST_PACKAGE_1).isPresent());
+    }
+
+    @Test
+    public void testPerformBackup_existingStateAndIncremental_performsIncrementalBackup()
+            throws Exception {
+        mKeyValueListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createKeyValueListing(CryptoTestUtils.mapOf(new Pair<>(TEST_KEY_1, TEST_HASH_1))));
+        ChunksMetadataProto.ChunkListing oldChunkListing =
+                createChunkListing(CryptoTestUtils.mapOf(new Pair<>(TEST_HASH_1, TEST_LENGTH_1)));
+        mChunkListingStore.saveProto(TEST_PACKAGE_1, oldChunkListing);
+
+        mTask.performBackup(INCREMENTAL);
+
+        verify(mEncryptedBackupTask)
+                .performIncrementalBackup(
+                        eq(mTertiaryKey), eq(mWrappedTertiaryKey), mChunkListingCaptor.capture());
+        assertChunkListingsEqual(mChunkListingCaptor.getValue(), oldChunkListing);
+    }
+
+    @Test
+    public void testPerformBackup_noExistingStateAndNonIncremental_performsNonIncrementalBackup()
+            throws Exception {
+        mTask.performBackup(NON_INCREMENTAL);
+
+        verify(mEncryptedBackupTask)
+                .performNonIncrementalBackup(eq(mTertiaryKey), eq(mWrappedTertiaryKey), eq(null));
+    }
+
+    @Test
+    public void testPerformBackup_incremental_savesNewListings() throws Exception {
+        mKeyValueListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createKeyValueListing(CryptoTestUtils.mapOf(new Pair<>(TEST_KEY_1, TEST_HASH_1))));
+        mChunkListingStore.saveProto(
+                TEST_PACKAGE_1,
+                createChunkListing(CryptoTestUtils.mapOf(new Pair<>(TEST_HASH_1, TEST_LENGTH_1))));
+
+        mTask.performBackup(INCREMENTAL);
+
+        KeyValueListingProto.KeyValueListing actualKeyValueListing =
+                mKeyValueListingStore.loadProto(TEST_PACKAGE_1).get();
+        ChunksMetadataProto.ChunkListing actualChunkListing =
+                mChunkListingStore.loadProto(TEST_PACKAGE_1).get();
+        assertKeyValueListingsEqual(actualKeyValueListing, mNewKeyValueListing);
+        assertChunkListingsEqual(actualChunkListing, mNewChunkListing);
+    }
+
+    @Test
+    public void testPerformBackup_nonIncremental_savesNewListings() throws Exception {
+        mTask.performBackup(NON_INCREMENTAL);
+
+        KeyValueListingProto.KeyValueListing actualKeyValueListing =
+                mKeyValueListingStore.loadProto(TEST_PACKAGE_1).get();
+        ChunksMetadataProto.ChunkListing actualChunkListing =
+                mChunkListingStore.loadProto(TEST_PACKAGE_1).get();
+        assertKeyValueListingsEqual(actualKeyValueListing, mNewKeyValueListing);
+        assertChunkListingsEqual(actualChunkListing, mNewChunkListing);
+    }
+
+    private static KeyValueListingProto.KeyValueListing createKeyValueListing(
+            Map<String, ChunkHash> pairs) {
+        return new KeyValueListingBuilder().addAll(pairs).build();
+    }
+
+    private static ChunksMetadataProto.ChunkListing createChunkListing(
+            Map<ChunkHash, Integer> chunks) {
+        ChunksMetadataProto.Chunk[] listingChunks = new ChunksMetadataProto.Chunk[chunks.size()];
+        int chunksAdded = 0;
+        for (Entry<ChunkHash, Integer> entry : chunks.entrySet()) {
+            listingChunks[chunksAdded] = CryptoTestUtils.newChunk(entry.getKey(), entry.getValue());
+            chunksAdded++;
+        }
+        return CryptoTestUtils.newChunkListingWithoutDocId(
+                /* fingerprintSalt */ new byte[0],
+                ChunksMetadataProto.AES_256_GCM,
+                ChunksMetadataProto.CHUNK_ORDERING_TYPE_UNSPECIFIED,
+                listingChunks);
+    }
+
+    private static void assertKeyValueListingsEqual(
+            KeyValueListingProto.KeyValueListing actual,
+            KeyValueListingProto.KeyValueListing expected) {
+        KeyValueListingProto.KeyValueEntry[] actualEntries = actual.entries;
+        KeyValueListingProto.KeyValueEntry[] expectedEntries = expected.entries;
+        assertThat(actualEntries.length).isEqualTo(expectedEntries.length);
+        for (int i = 0; i < actualEntries.length; i++) {
+            assertWithMessage("entry " + i)
+                    .that(actualEntries[i].key)
+                    .isEqualTo(expectedEntries[i].key);
+            assertWithMessage("entry " + i)
+                    .that(actualEntries[i].hash)
+                    .isEqualTo(expectedEntries[i].hash);
+        }
+    }
+
+    private static void assertChunkListingsEqual(
+            ChunksMetadataProto.ChunkListing actual, ChunksMetadataProto.ChunkListing expected) {
+        ChunksMetadataProto.Chunk[] actualChunks = actual.chunks;
+        ChunksMetadataProto.Chunk[] expectedChunks = expected.chunks;
+        assertThat(actualChunks.length).isEqualTo(expectedChunks.length);
+        for (int i = 0; i < actualChunks.length; i++) {
+            assertWithMessage("chunk " + i)
+                    .that(actualChunks[i].hash)
+                    .isEqualTo(expectedChunks[i].hash);
+            assertWithMessage("chunk " + i)
+                    .that(actualChunks[i].length)
+                    .isEqualTo(expectedChunks[i].length);
+        }
+        assertThat(actual.cipherType).isEqualTo(expected.cipherType);
+        assertThat(actual.documentId)
+                .isEqualTo(expected.documentId == null ? "" : expected.documentId);
+    }
+}
diff --git a/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/EncryptedKvRestoreTaskTest.java b/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/EncryptedKvRestoreTaskTest.java
new file mode 100644
index 000000000000..6666d95d9a2d
--- /dev/null
+++ b/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/EncryptedKvRestoreTaskTest.java
@@ -0,0 +1,185 @@
+/*
+ * Copyright (C) 2019 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.backup.encryption.tasks;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertThrows;
+
+import android.os.ParcelFileDescriptor;
+
+import com.android.server.backup.encryption.chunk.ChunkHash;
+import com.android.server.backup.encryption.chunking.ChunkHasher;
+import com.android.server.backup.testing.CryptoTestUtils;
+import com.android.server.testing.shadows.DataEntity;
+import com.android.server.testing.shadows.ShadowBackupDataOutput;
+
+import com.google.protobuf.nano.MessageNano;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
+
+@Config(shadows = {ShadowBackupDataOutput.class})
+@RunWith(RobolectricTestRunner.class)
+public class EncryptedKvRestoreTaskTest {
+    private static final String TEST_KEY_1 = "test_key_1";
+    private static final String TEST_KEY_2 = "test_key_2";
+    private static final String TEST_KEY_3 = "test_key_3";
+    private static final byte[] TEST_VALUE_1 = {1, 2, 3};
+    private static final byte[] TEST_VALUE_2 = {4, 5, 6};
+    private static final byte[] TEST_VALUE_3 = {20, 25, 30, 35};
+
+    @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+    private File temporaryDirectory;
+
+    @Mock private ParcelFileDescriptor mParcelFileDescriptor;
+    @Mock private ChunkHasher mChunkHasher;
+    @Mock private FullRestoreToFileTask mFullRestoreToFileTask;
+    @Mock private BackupFileDecryptorTask mBackupFileDecryptorTask;
+
+    private EncryptedKvRestoreTask task;
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        when(mChunkHasher.computeHash(any()))
+                .thenAnswer(invocation -> fakeHash(invocation.getArgument(0)));
+        doAnswer(invocation -> writeTestPairsToFile(invocation.getArgument(0)))
+                .when(mFullRestoreToFileTask)
+                .restoreToFile(any());
+        doAnswer(
+                        invocation ->
+                                readPairsFromFile(
+                                        invocation.getArgument(0), invocation.getArgument(1)))
+                .when(mBackupFileDecryptorTask)
+                .decryptFile(any(), any());
+
+        temporaryDirectory = temporaryFolder.newFolder();
+        task =
+                new EncryptedKvRestoreTask(
+                        temporaryDirectory,
+                        mChunkHasher,
+                        mFullRestoreToFileTask,
+                        mBackupFileDecryptorTask);
+    }
+
+    @Test
+    public void testGetRestoreData_writesPairsToOutputInOrder() throws Exception {
+        task.getRestoreData(mParcelFileDescriptor);
+
+        assertThat(ShadowBackupDataOutput.getEntities())
+                .containsExactly(
+                        new DataEntity(TEST_KEY_1, TEST_VALUE_1),
+                        new DataEntity(TEST_KEY_2, TEST_VALUE_2),
+                        new DataEntity(TEST_KEY_3, TEST_VALUE_3))
+                .inOrder();
+    }
+
+    @Test
+    public void testGetRestoreData_exceptionDuringDecryption_throws() throws Exception {
+        doThrow(IOException.class).when(mBackupFileDecryptorTask).decryptFile(any(), any());
+        assertThrows(IOException.class, () -> task.getRestoreData(mParcelFileDescriptor));
+    }
+
+    @Test
+    public void testGetRestoreData_exceptionDuringDownload_throws() throws Exception {
+        doThrow(IOException.class).when(mFullRestoreToFileTask).restoreToFile(any());
+        assertThrows(IOException.class, () -> task.getRestoreData(mParcelFileDescriptor));
+    }
+
+    @Test
+    public void testGetRestoreData_exceptionDuringDecryption_deletesTemporaryFiles() throws Exception {
+        doThrow(InvalidKeyException.class).when(mBackupFileDecryptorTask).decryptFile(any(), any());
+        assertThrows(InvalidKeyException.class, () -> task.getRestoreData(mParcelFileDescriptor));
+        assertThat(temporaryDirectory.listFiles()).isEmpty();
+    }
+
+    @Test
+    public void testGetRestoreData_exceptionDuringDownload_deletesTemporaryFiles() throws Exception {
+        doThrow(IOException.class).when(mFullRestoreToFileTask).restoreToFile(any());
+        assertThrows(IOException.class, () -> task.getRestoreData(mParcelFileDescriptor));
+        assertThat(temporaryDirectory.listFiles()).isEmpty();
+    }
+
+    private static Void writeTestPairsToFile(File file) throws IOException {
+        // Write the pairs out of order to check the task sorts them.
+        Set<byte[]> pairs =
+                new HashSet<>(
+                        Arrays.asList(
+                                createPair(TEST_KEY_1, TEST_VALUE_1),
+                                createPair(TEST_KEY_3, TEST_VALUE_3),
+                                createPair(TEST_KEY_2, TEST_VALUE_2)));
+
+        try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(file))) {
+            oos.writeObject(pairs);
+        }
+        return null;
+    }
+
+    private static Void readPairsFromFile(File file, DecryptedChunkOutput decryptedChunkOutput)
+            throws IOException, ClassNotFoundException, InvalidKeyException,
+                    NoSuchAlgorithmException {
+        try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(file));
+                DecryptedChunkOutput output = decryptedChunkOutput.open()) {
+            Set<byte[]> pairs = readPairs(ois);
+            for (byte[] pair : pairs) {
+                output.processChunk(pair, pair.length);
+            }
+        }
+
+        return null;
+    }
+
+    private static byte[] createPair(String key, byte[] value) {
+        return MessageNano.toByteArray(CryptoTestUtils.newPair(key, value));
+    }
+
+    @SuppressWarnings("unchecked") // deserialization.
+    private static Set<byte[]> readPairs(ObjectInputStream ois)
+            throws IOException, ClassNotFoundException {
+        return (Set<byte[]>) ois.readObject();
+    }
+
+    private static ChunkHash fakeHash(byte[] data) {
+        return new ChunkHash(Arrays.copyOf(data, ChunkHash.HASH_LENGTH_BYTES));
+    }
+}
diff --git a/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/CryptoTestUtils.java b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/CryptoTestUtils.java
index 5dfd5ee8ad53..b0c02ba637e0 100644
--- a/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/CryptoTestUtils.java
+++ b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/CryptoTestUtils.java
@@ -16,6 +16,8 @@
 
 package com.android.server.backup.testing;
 
+import android.util.Pair;
+
 import com.android.server.backup.encryption.chunk.ChunkHash;
 import com.android.server.backup.encryption.protos.nano.ChunksMetadataProto;
 import com.android.server.backup.encryption.protos.nano.KeyValuePairProto;
@@ -23,6 +25,8 @@ import com.android.server.backup.encryption.protos.nano.KeyValuePairProto;
 import java.nio.charset.Charset;
 import java.security.NoSuchAlgorithmException;
 import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Random;
 
 import javax.crypto.KeyGenerator;
@@ -162,4 +166,12 @@ public class CryptoTestUtils {
         clone.checksum = Arrays.copyOf(original.checksum, original.checksum.length);
         return clone;
     }
+
+    public static <K, V> Map<K, V> mapOf(Pair<K, V>... pairs) {
+        Map<K, V> map = new HashMap<>();
+        for (Pair<K, V> pair : pairs) {
+            map.put(pair.first, pair.second);
+        }
+        return map;
+    }
 }
diff --git a/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/ShadowBackupDataOutput.java b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/ShadowBackupDataOutput.java
new file mode 100644
index 000000000000..2302e555fb44
--- /dev/null
+++ b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/ShadowBackupDataOutput.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright (C) 2019 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.testing.shadows;
+
+import android.app.backup.BackupDataOutput;
+
+import java.io.FileDescriptor;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
+
+import org.junit.Assert;
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+
+/** Shadow for BackupDataOutput. */
+@Implements(BackupDataOutput.class)
+public class ShadowBackupDataOutput {
+    private static final List<DataEntity> ENTRIES = new ArrayList<>();
+
+    private String mCurrentKey;
+    private int mDataSize;
+
+    public static void reset() {
+        ENTRIES.clear();
+    }
+
+    public static Set<DataEntity> getEntities() {
+        return new LinkedHashSet<>(ENTRIES);
+    }
+
+    public void __constructor__(FileDescriptor fd) {}
+
+    public void __constructor__(FileDescriptor fd, long quota) {}
+
+    public void __constructor__(FileDescriptor fd, long quota, int transportFlags) {}
+
+    @Implementation
+    public int writeEntityHeader(String key, int size) {
+        mCurrentKey = key;
+        mDataSize = size;
+        return 0;
+    }
+
+    @Implementation
+    public int writeEntityData(byte[] data, int size) {
+        Assert.assertEquals("ShadowBackupDataOutput expects size = mDataSize", size, mDataSize);
+        ENTRIES.add(new DataEntity(mCurrentKey, data, mDataSize));
+        return 0;
+    }
+}