Add integration test for encrypted KV B&R

Bug: 11386661
Test: RoundTripTest
Change-Id: Ifd18961347d8ebaf00ddc12196f2ee6ec543b457

2 files changed, 120 insertions, 15 deletions

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;
         }
     }