Merge "Import KvBackupEncrypter"

5 files changed, 673 insertions, 1 deletions

diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/KvBackupEncrypter.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/KvBackupEncrypter.java
new file mode 100644
index 000000000000..d20cd4c07f88
--- /dev/null
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/tasks/KvBackupEncrypter.java
@@ -0,0 +1,179 @@
+/*
+ * 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.checkState;
+
+import android.annotation.Nullable;
+import android.app.backup.BackupDataInput;
+
+import com.android.server.backup.encryption.chunk.ChunkHash;
+import com.android.server.backup.encryption.chunking.ChunkEncryptor;
+import com.android.server.backup.encryption.chunking.ChunkHasher;
+import com.android.server.backup.encryption.chunking.EncryptedChunk;
+import com.android.server.backup.encryption.kv.KeyValueListingBuilder;
+import com.android.server.backup.encryption.protos.nano.KeyValueListingProto;
+import com.android.server.backup.encryption.protos.nano.KeyValuePairProto;
+
+import java.io.IOException;
+import java.security.GeneralSecurityException;
+import java.security.InvalidKeyException;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.security.SecureRandom;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.SecretKey;
+
+/**
+ * Reads key value backup data from an input, converts each pair into a chunk and encrypts the
+ * chunks.
+ *
+ * <p>The caller should pass in the key value listing from the previous backup, if there is one.
+ * This class emits chunks for both existing and new pairs, using the provided listing to
+ * determine the hashes of pairs that already exist. During the backup it computes the new listing,
+ * which the caller should store on disk and pass in at the start of the next backup.
+ *
+ * <p>Also computes the message digest, which is {@code SHA-256(chunk hashes sorted
+ * lexicographically)}.
+ */
+public class KvBackupEncrypter implements BackupEncrypter {
+    private final BackupDataInput mBackupDataInput;
+
+    private KeyValueListingProto.KeyValueListing mOldKeyValueListing;
+    @Nullable private KeyValueListingBuilder mNewKeyValueListing;
+
+    /**
+     * Constructs a new instance which reads data from the given input.
+     *
+     * <p>By default this performs non-incremental backup, call {@link #setOldKeyValueListing} to
+     * perform incremental backup.
+     */
+    public KvBackupEncrypter(BackupDataInput backupDataInput) {
+        mBackupDataInput = backupDataInput;
+        mOldKeyValueListing = KeyValueListingBuilder.emptyListing();
+    }
+
+    /** Sets the old listing to perform incremental backup against. */
+    public void setOldKeyValueListing(KeyValueListingProto.KeyValueListing oldKeyValueListing) {
+        mOldKeyValueListing = oldKeyValueListing;
+    }
+
+    @Override
+    public Result backup(
+            SecretKey secretKey,
+            @Nullable byte[] unusedFingerprintMixerSalt,
+            Set<ChunkHash> unusedExistingChunks)
+            throws IOException, GeneralSecurityException {
+        ChunkHasher chunkHasher = new ChunkHasher(secretKey);
+        ChunkEncryptor chunkEncryptor = new ChunkEncryptor(secretKey, new SecureRandom());
+        mNewKeyValueListing = new KeyValueListingBuilder();
+        List<ChunkHash> allChunks = new ArrayList<>();
+        List<EncryptedChunk> newChunks = new ArrayList<>();
+
+        Map<String, ChunkHash> existingChunksToReuse = buildPairMap(mOldKeyValueListing);
+
+        while (mBackupDataInput.readNextHeader()) {
+            String key = mBackupDataInput.getKey();
+            Optional<byte[]> value = readEntireValue(mBackupDataInput);
+
+            // As this pair exists in the new backup, we don't need to add it from the previous
+            // backup.
+            existingChunksToReuse.remove(key);
+
+            // If the value is not present then this key has been deleted.
+            if (value.isPresent()) {
+                EncryptedChunk newChunk =
+                        createEncryptedChunk(chunkHasher, chunkEncryptor, key, value.get());
+                allChunks.add(newChunk.key());
+                newChunks.add(newChunk);
+                mNewKeyValueListing.addPair(key, newChunk.key());
+            }
+        }
+
+        allChunks.addAll(existingChunksToReuse.values());
+
+        mNewKeyValueListing.addAll(existingChunksToReuse);
+
+        return new Result(allChunks, newChunks, createMessageDigest(allChunks));
+    }
+
+    /**
+     * Returns a listing containing the pairs in the new backup.
+     *
+     * <p>You must call {@link #backup} first.
+     */
+    public KeyValueListingProto.KeyValueListing getNewKeyValueListing() {
+        checkState(mNewKeyValueListing != null, "Must call backup() first");
+        return mNewKeyValueListing.build();
+    }
+
+    private static Map<String, ChunkHash> buildPairMap(
+            KeyValueListingProto.KeyValueListing listing) {
+        Map<String, ChunkHash> map = new HashMap<>();
+        for (KeyValueListingProto.KeyValueEntry entry : listing.entries) {
+            map.put(entry.key, new ChunkHash(entry.hash));
+        }
+        return map;
+    }
+
+    private EncryptedChunk createEncryptedChunk(
+            ChunkHasher chunkHasher, ChunkEncryptor chunkEncryptor, String key, byte[] value)
+            throws InvalidKeyException, IllegalBlockSizeException {
+        KeyValuePairProto.KeyValuePair pair = new KeyValuePairProto.KeyValuePair();
+        pair.key = key;
+        pair.value = Arrays.copyOf(value, value.length);
+
+        byte[] plaintext = KeyValuePairProto.KeyValuePair.toByteArray(pair);
+        return chunkEncryptor.encrypt(chunkHasher.computeHash(plaintext), plaintext);
+    }
+
+    private static byte[] createMessageDigest(List<ChunkHash> allChunks)
+            throws NoSuchAlgorithmException {
+        MessageDigest messageDigest =
+                MessageDigest.getInstance(BackupEncrypter.MESSAGE_DIGEST_ALGORITHM);
+        // TODO:b/141531271 Extract sorted chunks code to utility class
+        List<ChunkHash> sortedChunks = new ArrayList<>(allChunks);
+        Collections.sort(sortedChunks);
+        for (ChunkHash hash : sortedChunks) {
+            messageDigest.update(hash.getHash());
+        }
+        return messageDigest.digest();
+    }
+
+    private static Optional<byte[]> readEntireValue(BackupDataInput input) throws IOException {
+        // A negative data size indicates that this key should be deleted.
+        if (input.getDataSize() < 0) {
+            return Optional.empty();
+        }
+
+        byte[] value = new byte[input.getDataSize()];
+        int bytesRead = 0;
+        while (bytesRead < value.length) {
+            bytesRead += input.readEntityData(value, bytesRead, value.length - bytesRead);
+        }
+        return Optional.of(value);
+    }
+}
diff --git a/packages/BackupEncryption/test/robolectric/Android.bp b/packages/BackupEncryption/test/robolectric/Android.bp
index 4e42ce7366f0..2a36dcf0baba 100644
--- a/packages/BackupEncryption/test/robolectric/Android.bp
+++ b/packages/BackupEncryption/test/robolectric/Android.bp
@@ -16,7 +16,7 @@ android_robolectric_test {
     name: "BackupEncryptionRoboTests",
     srcs: [
         "src/**/*.java",
-        ":FrameworksServicesRoboShadows",
+//        ":FrameworksServicesRoboShadows",
     ],
     java_resource_dirs: ["config"],
     libs: [
diff --git a/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/KvBackupEncrypterTest.java b/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/KvBackupEncrypterTest.java
new file mode 100644
index 000000000000..ccfbfa4b25e9
--- /dev/null
+++ b/packages/BackupEncryption/test/robolectric/src/com/android/server/backup/encryption/tasks/KvBackupEncrypterTest.java
@@ -0,0 +1,287 @@
+/*
+ * 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.testng.Assert.assertThrows;
+
+import android.app.backup.BackupDataInput;
+import android.platform.test.annotations.Presubmit;
+import android.util.Pair;
+
+import com.android.server.backup.encryption.chunk.ChunkHash;
+import com.android.server.backup.encryption.chunking.ChunkHasher;
+import com.android.server.backup.encryption.chunking.EncryptedChunk;
+import com.android.server.backup.encryption.kv.KeyValueListingBuilder;
+import com.android.server.backup.encryption.protos.nano.KeyValueListingProto.KeyValueListing;
+import com.android.server.backup.encryption.protos.nano.KeyValuePairProto.KeyValuePair;
+import com.android.server.backup.encryption.tasks.BackupEncrypter.Result;
+import com.android.server.testing.shadows.DataEntity;
+import com.android.server.testing.shadows.ShadowBackupDataInput;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Ordering;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
+
+import java.security.MessageDigest;
+import java.util.Arrays;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import javax.crypto.Cipher;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.GCMParameterSpec;
+import javax.crypto.spec.SecretKeySpec;
+
+@RunWith(RobolectricTestRunner.class)
+@Presubmit
+@Config(shadows = {ShadowBackupDataInput.class})
+public class KvBackupEncrypterTest {
+    private static final String KEY_ALGORITHM = "AES";
+    private static final String CIPHER_ALGORITHM = "AES/GCM/NoPadding";
+    private static final int GCM_TAG_LENGTH_BYTES = 16;
+
+    private static final byte[] TEST_TERTIARY_KEY = Arrays.copyOf(new byte[0], 256 / Byte.SIZE);
+    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 = {10, 11, 12};
+    private static final byte[] TEST_VALUE_2 = {13, 14, 15};
+    private static final byte[] TEST_VALUE_2B = {13, 14, 15, 16};
+    private static final byte[] TEST_VALUE_3 = {16, 17, 18};
+
+    private SecretKey mSecretKey;
+    private ChunkHasher mChunkHasher;
+
+    @Before
+    public void setUp() {
+        mSecretKey = new SecretKeySpec(TEST_TERTIARY_KEY, KEY_ALGORITHM);
+        mChunkHasher = new ChunkHasher(mSecretKey);
+
+        ShadowBackupDataInput.reset();
+    }
+
+    private KvBackupEncrypter createEncrypter(KeyValueListing keyValueListing) {
+        KvBackupEncrypter encrypter = new KvBackupEncrypter(new BackupDataInput(null));
+        encrypter.setOldKeyValueListing(keyValueListing);
+        return encrypter;
+    }
+
+    @Test
+    public void backup_noExistingBackup_encryptsAllPairs() throws Exception {
+        ShadowBackupDataInput.addEntity(TEST_KEY_1, TEST_VALUE_1);
+        ShadowBackupDataInput.addEntity(TEST_KEY_2, TEST_VALUE_2);
+
+        KeyValueListing emptyKeyValueListing = new KeyValueListingBuilder().build();
+        ImmutableSet<ChunkHash> emptyExistingChunks = ImmutableSet.of();
+        KvBackupEncrypter encrypter = createEncrypter(emptyKeyValueListing);
+
+        Result result =
+                encrypter.backup(
+                        mSecretKey, /*unusedFingerprintMixerSalt=*/ null, emptyExistingChunks);
+
+        assertThat(result.getAllChunks()).hasSize(2);
+        EncryptedChunk chunk1 = result.getNewChunks().get(0);
+        EncryptedChunk chunk2 = result.getNewChunks().get(1);
+        assertThat(chunk1.key()).isEqualTo(getChunkHash(TEST_KEY_1, TEST_VALUE_1));
+        KeyValuePair pair1 = decryptChunk(chunk1);
+        assertThat(pair1.key).isEqualTo(TEST_KEY_1);
+        assertThat(pair1.value).isEqualTo(TEST_VALUE_1);
+        assertThat(chunk2.key()).isEqualTo(getChunkHash(TEST_KEY_2, TEST_VALUE_2));
+        KeyValuePair pair2 = decryptChunk(chunk2);
+        assertThat(pair2.key).isEqualTo(TEST_KEY_2);
+        assertThat(pair2.value).isEqualTo(TEST_VALUE_2);
+    }
+
+    @Test
+    public void backup_existingBackup_encryptsNewAndUpdatedPairs() throws Exception {
+        Pair<KeyValueListing, Set<ChunkHash>> initialResult = runInitialBackupOfPairs1And2();
+
+        // Update key 2 and add the new key 3.
+        ShadowBackupDataInput.reset();
+        ShadowBackupDataInput.addEntity(TEST_KEY_2, TEST_VALUE_2B);
+        ShadowBackupDataInput.addEntity(TEST_KEY_3, TEST_VALUE_3);
+
+        KvBackupEncrypter encrypter = createEncrypter(initialResult.first);
+        BackupEncrypter.Result secondResult =
+                encrypter.backup(
+                        mSecretKey, /*unusedFingerprintMixerSalt=*/ null, initialResult.second);
+
+        assertThat(secondResult.getAllChunks()).hasSize(3);
+        assertThat(secondResult.getNewChunks()).hasSize(2);
+        EncryptedChunk newChunk2 = secondResult.getNewChunks().get(0);
+        EncryptedChunk newChunk3 = secondResult.getNewChunks().get(1);
+        assertThat(newChunk2.key()).isEqualTo(getChunkHash(TEST_KEY_2, TEST_VALUE_2B));
+        assertThat(decryptChunk(newChunk2).value).isEqualTo(TEST_VALUE_2B);
+        assertThat(newChunk3.key()).isEqualTo(getChunkHash(TEST_KEY_3, TEST_VALUE_3));
+        assertThat(decryptChunk(newChunk3).value).isEqualTo(TEST_VALUE_3);
+    }
+
+    @Test
+    public void backup_allChunksContainsHashesOfAllChunks() throws Exception {
+        Pair<KeyValueListing, Set<ChunkHash>> initialResult = runInitialBackupOfPairs1And2();
+
+        ShadowBackupDataInput.reset();
+        ShadowBackupDataInput.addEntity(TEST_KEY_3, TEST_VALUE_3);
+
+        KvBackupEncrypter encrypter = createEncrypter(initialResult.first);
+        BackupEncrypter.Result secondResult =
+                encrypter.backup(
+                        mSecretKey, /*unusedFingerprintMixerSalt=*/ null, initialResult.second);
+
+        assertThat(secondResult.getAllChunks())
+                .containsExactly(
+                        getChunkHash(TEST_KEY_1, TEST_VALUE_1),
+                        getChunkHash(TEST_KEY_2, TEST_VALUE_2),
+                        getChunkHash(TEST_KEY_3, TEST_VALUE_3));
+    }
+
+    @Test
+    public void backup_negativeSize_deletesKeyFromExistingBackup() throws Exception {
+        Pair<KeyValueListing, Set<ChunkHash>> initialResult = runInitialBackupOfPairs1And2();
+
+        ShadowBackupDataInput.reset();
+        ShadowBackupDataInput.addEntity(new DataEntity(TEST_KEY_2));
+
+        KvBackupEncrypter encrypter = createEncrypter(initialResult.first);
+        Result secondResult =
+                encrypter.backup(
+                        mSecretKey, /*unusedFingerprintMixerSalt=*/ null, initialResult.second);
+
+        assertThat(secondResult.getAllChunks())
+                .containsExactly(getChunkHash(TEST_KEY_1, TEST_VALUE_1));
+        assertThat(secondResult.getNewChunks()).isEmpty();
+    }
+
+    @Test
+    public void backup_returnsMessageDigestOverChunkHashes() throws Exception {
+        Pair<KeyValueListing, Set<ChunkHash>> initialResult = runInitialBackupOfPairs1And2();
+
+        ShadowBackupDataInput.reset();
+        ShadowBackupDataInput.addEntity(TEST_KEY_3, TEST_VALUE_3);
+
+        KvBackupEncrypter encrypter = createEncrypter(initialResult.first);
+        Result secondResult =
+                encrypter.backup(
+                        mSecretKey, /*unusedFingerprintMixerSalt=*/ null, initialResult.second);
+
+        MessageDigest messageDigest =
+                MessageDigest.getInstance(BackupEncrypter.MESSAGE_DIGEST_ALGORITHM);
+        ImmutableList<ChunkHash> sortedHashes =
+                Ordering.natural()
+                        .immutableSortedCopy(
+                                ImmutableList.of(
+                                        getChunkHash(TEST_KEY_1, TEST_VALUE_1),
+                                        getChunkHash(TEST_KEY_2, TEST_VALUE_2),
+                                        getChunkHash(TEST_KEY_3, TEST_VALUE_3)));
+        messageDigest.update(sortedHashes.get(0).getHash());
+        messageDigest.update(sortedHashes.get(1).getHash());
+        messageDigest.update(sortedHashes.get(2).getHash());
+        assertThat(secondResult.getDigest()).isEqualTo(messageDigest.digest());
+    }
+
+    @Test
+    public void getNewKeyValueListing_noExistingBackup_returnsCorrectListing() throws Exception {
+        KeyValueListing keyValueListing = runInitialBackupOfPairs1And2().first;
+
+        assertThat(keyValueListing.entries.length).isEqualTo(2);
+        assertThat(keyValueListing.entries[0].key).isEqualTo(TEST_KEY_1);
+        assertThat(keyValueListing.entries[0].hash)
+                .isEqualTo(getChunkHash(TEST_KEY_1, TEST_VALUE_1).getHash());
+        assertThat(keyValueListing.entries[1].key).isEqualTo(TEST_KEY_2);
+        assertThat(keyValueListing.entries[1].hash)
+                .isEqualTo(getChunkHash(TEST_KEY_2, TEST_VALUE_2).getHash());
+    }
+
+    @Test
+    public void getNewKeyValueListing_existingBackup_returnsCorrectListing() throws Exception {
+        Pair<KeyValueListing, Set<ChunkHash>> initialResult = runInitialBackupOfPairs1And2();
+
+        ShadowBackupDataInput.reset();
+        ShadowBackupDataInput.addEntity(TEST_KEY_2, TEST_VALUE_2B);
+        ShadowBackupDataInput.addEntity(TEST_KEY_3, TEST_VALUE_3);
+
+        KvBackupEncrypter encrypter = createEncrypter(initialResult.first);
+        encrypter.backup(mSecretKey, /*unusedFingerprintMixerSalt=*/ null, initialResult.second);
+
+        ImmutableMap<String, ChunkHash> keyValueListing =
+                listingToMap(encrypter.getNewKeyValueListing());
+        assertThat(keyValueListing).hasSize(3);
+        assertThat(keyValueListing)
+                .containsEntry(TEST_KEY_1, getChunkHash(TEST_KEY_1, TEST_VALUE_1));
+        assertThat(keyValueListing)
+                .containsEntry(TEST_KEY_2, getChunkHash(TEST_KEY_2, TEST_VALUE_2B));
+        assertThat(keyValueListing)
+                .containsEntry(TEST_KEY_3, getChunkHash(TEST_KEY_3, TEST_VALUE_3));
+    }
+
+    @Test
+    public void getNewKeyValueChunkListing_beforeBackup_throws() throws Exception {
+        KvBackupEncrypter encrypter = createEncrypter(new KeyValueListing());
+        assertThrows(IllegalStateException.class, encrypter::getNewKeyValueListing);
+    }
+
+    private ImmutableMap<String, ChunkHash> listingToMap(KeyValueListing listing) {
+        // We can't use the ImmutableMap collector directly because it isn't supported in Android
+        // guava.
+        return ImmutableMap.copyOf(
+                Arrays.stream(listing.entries)
+                        .collect(
+                                Collectors.toMap(
+                                        entry -> entry.key, entry -> new ChunkHash(entry.hash))));
+    }
+
+    private Pair<KeyValueListing, Set<ChunkHash>> runInitialBackupOfPairs1And2() throws Exception {
+        ShadowBackupDataInput.addEntity(TEST_KEY_1, TEST_VALUE_1);
+        ShadowBackupDataInput.addEntity(TEST_KEY_2, TEST_VALUE_2);
+
+        KeyValueListing initialKeyValueListing = new KeyValueListingBuilder().build();
+        ImmutableSet<ChunkHash> initialExistingChunks = ImmutableSet.of();
+        KvBackupEncrypter encrypter = createEncrypter(initialKeyValueListing);
+        Result firstResult =
+                encrypter.backup(
+                        mSecretKey, /*unusedFingerprintMixerSalt=*/ null, initialExistingChunks);
+
+        return Pair.create(
+                encrypter.getNewKeyValueListing(), ImmutableSet.copyOf(firstResult.getAllChunks()));
+    }
+
+    private ChunkHash getChunkHash(String key, byte[] value) throws Exception {
+        KeyValuePair pair = new KeyValuePair();
+        pair.key = key;
+        pair.value = Arrays.copyOf(value, value.length);
+        return mChunkHasher.computeHash(KeyValuePair.toByteArray(pair));
+    }
+
+    private KeyValuePair decryptChunk(EncryptedChunk encryptedChunk) throws Exception {
+        Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
+        cipher.init(
+                Cipher.DECRYPT_MODE,
+                mSecretKey,
+                new GCMParameterSpec(GCM_TAG_LENGTH_BYTES * Byte.SIZE, encryptedChunk.nonce()));
+        byte[] decryptedBytes = cipher.doFinal(encryptedChunk.encryptedBytes());
+        return KeyValuePair.parseFrom(decryptedBytes);
+    }
+}
diff --git a/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/DataEntity.java b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/DataEntity.java
new file mode 100644
index 000000000000..6d3b5e9f1d7b
--- /dev/null
+++ b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/DataEntity.java
@@ -0,0 +1,100 @@
+/*
+ * 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 static com.google.common.base.Preconditions.checkNotNull;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+
+/**
+ * Represents a key value pair in {@link ShadowBackupDataInput} and {@link ShadowBackupDataOutput}.
+ */
+public class DataEntity {
+    public final String mKey;
+    public final byte[] mValue;
+    public final int mSize;
+
+    /**
+     * Constructs a pair with a string value. The value will be converted to a byte array in {@link
+     * StandardCharsets#UTF_8}.
+     */
+    public DataEntity(String key, String value) {
+        this.mKey = checkNotNull(key);
+        this.mValue = value.getBytes(StandardCharsets.UTF_8);
+        mSize = this.mValue.length;
+    }
+
+    /**
+     * Constructs a new entity with the given key but a negative size. This represents a deleted
+     * pair.
+     */
+    public DataEntity(String key) {
+        this.mKey = checkNotNull(key);
+        mSize = -1;
+        mValue = null;
+    }
+
+    /** Constructs a new entity where the size of the value is the entire array. */
+    public DataEntity(String key, byte[] value) {
+        this(key, value, value.length);
+    }
+
+    /**
+     * Constructs a new entity.
+     *
+     * @param key the key of the pair
+     * @param data the value to associate with the key
+     * @param size the length of the value in bytes
+     */
+    public DataEntity(String key, byte[] data, int size) {
+        this.mKey = checkNotNull(key);
+        this.mSize = size;
+        mValue = new byte[size];
+        for (int i = 0; i < size; i++) {
+            mValue[i] = data[i];
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        DataEntity that = (DataEntity) o;
+
+        if (mSize != that.mSize) {
+            return false;
+        }
+        if (!mKey.equals(that.mKey)) {
+            return false;
+        }
+        return Arrays.equals(mValue, that.mValue);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = mKey.hashCode();
+        result = 31 * result + Arrays.hashCode(mValue);
+        result = 31 * result + mSize;
+        return result;
+    }
+}
diff --git a/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/ShadowBackupDataInput.java b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/ShadowBackupDataInput.java
new file mode 100644
index 000000000000..7ac6ec40508d
--- /dev/null
+++ b/packages/BackupEncryption/test/robolectric/src/com/android/server/testing/shadows/ShadowBackupDataInput.java
@@ -0,0 +1,106 @@
+/*
+ * 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 static com.google.common.base.Preconditions.checkState;
+
+import android.annotation.Nullable;
+import android.app.backup.BackupDataInput;
+
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+
+import java.io.ByteArrayInputStream;
+import java.io.FileDescriptor;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+/** Shadow for BackupDataInput. */
+@Implements(BackupDataInput.class)
+public class ShadowBackupDataInput {
+    private static final List<DataEntity> ENTITIES = new ArrayList<>();
+    @Nullable private static IOException sReadNextHeaderException;
+
+    @Nullable private ByteArrayInputStream mCurrentEntityInputStream;
+    private int mCurrentEntity = -1;
+
+    /** Resets the shadow, clearing any entities or exception. */
+    public static void reset() {
+        ENTITIES.clear();
+        sReadNextHeaderException = null;
+    }
+
+    /** Sets the exception which the input will throw for any call to {@link #readNextHeader}. */
+    public static void setReadNextHeaderException(@Nullable IOException readNextHeaderException) {
+        ShadowBackupDataInput.sReadNextHeaderException = readNextHeaderException;
+    }
+
+    /** Adds the given entity to the input. */
+    public static void addEntity(DataEntity e) {
+        ENTITIES.add(e);
+    }
+
+    /** Adds an entity to the input with the given key and value. */
+    public static void addEntity(String key, byte[] value) {
+        ENTITIES.add(new DataEntity(key, value, value.length));
+    }
+
+    public void __constructor__(FileDescriptor fd) {}
+
+    @Implementation
+    public boolean readNextHeader() throws IOException {
+        if (sReadNextHeaderException != null) {
+            throw sReadNextHeaderException;
+        }
+
+        mCurrentEntity++;
+
+        if (mCurrentEntity >= ENTITIES.size()) {
+            return false;
+        }
+
+        byte[] value = ENTITIES.get(mCurrentEntity).mValue;
+        if (value == null) {
+            mCurrentEntityInputStream = new ByteArrayInputStream(new byte[0]);
+        } else {
+            mCurrentEntityInputStream = new ByteArrayInputStream(value);
+        }
+        return true;
+    }
+
+    @Implementation
+    public String getKey() {
+        return ENTITIES.get(mCurrentEntity).mKey;
+    }
+
+    @Implementation
+    public int getDataSize() {
+        return ENTITIES.get(mCurrentEntity).mSize;
+    }
+
+    @Implementation
+    public void skipEntityData() {
+        // Do nothing.
+    }
+
+    @Implementation
+    public int readEntityData(byte[] data, int offset, int size) {
+        checkState(mCurrentEntityInputStream != null, "Must call readNextHeader() first");
+        return mCurrentEntityInputStream.read(data, offset, size);
+    }
+}