Route EncryptedLocalTransport KV backup/restore through encryption code

Bug: 142227548
Test: Verify the device boots successfully
      Verify EncryptedLocalTransport APK is present
      Verify manual backup/restore using bmgr for LocalTransport and EncryptedLocalTransport
      For LocalTransport (unencrypted) and EncryptedLocalTransport:
        atest CtsBackupTestCases
        atest CtsBackupHostTestCases
        atest GtsBackupTestCases
        atest GtsBackupHostTestCases


Change-Id: Iac3a8a50d7f761442a4b784cfba3a980e900dd7f

6 files changed, 501 insertions, 20 deletions

diff --git a/packages/BackupEncryption/Android.bp b/packages/BackupEncryption/Android.bp
index 342d796de402..68e937c24a89 100644
--- a/packages/BackupEncryption/Android.bp
+++ b/packages/BackupEncryption/Android.bp
@@ -17,8 +17,7 @@
 android_app {
     name: "BackupEncryption",
     srcs: ["src/**/*.java"],
-    libs: ["backup-encryption-protos"],
-    static_libs: ["backuplib"],
+    static_libs: ["backup-encryption-protos", "backuplib"],
     optimize: { enabled: false },
     platform_apis: true,
     certificate: "platform",
diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/EncryptionKeyHelper.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/EncryptionKeyHelper.java
new file mode 100644
index 000000000000..2035b6605559
--- /dev/null
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/EncryptionKeyHelper.java
@@ -0,0 +1,81 @@
+/*
+ * 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;
+
+import android.content.Context;
+import android.security.keystore.recovery.InternalRecoveryServiceException;
+import android.security.keystore.recovery.RecoveryController;
+
+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 java.io.IOException;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.security.SecureRandom;
+import java.security.UnrecoverableKeyException;
+
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
+import javax.crypto.SecretKey;
+
+class EncryptionKeyHelper {
+    private static SecureRandom sSecureRandom = new  SecureRandom();
+
+    private final Context mContext;
+    private final RecoverableKeyStoreSecondaryKeyManager
+            .RecoverableKeyStoreSecondaryKeyManagerProvider
+            mSecondaryKeyManagerProvider;
+
+    EncryptionKeyHelper(Context context) {
+        mContext = context;
+        mSecondaryKeyManagerProvider =
+                () ->
+                        new RecoverableKeyStoreSecondaryKeyManager(
+                                RecoveryController.getInstance(mContext), sSecureRandom);
+    }
+
+    RecoverableKeyStoreSecondaryKeyManager
+            .RecoverableKeyStoreSecondaryKeyManagerProvider getKeyManagerProvider() {
+        return mSecondaryKeyManagerProvider;
+    }
+
+    RecoverableKeyStoreSecondaryKey getActiveSecondaryKey()
+            throws UnrecoverableKeyException, InternalRecoveryServiceException {
+        String keyAlias = CryptoSettings.getInstance(mContext).getActiveSecondaryKeyAlias().get();
+        return mSecondaryKeyManagerProvider.get().get(keyAlias).get();
+    }
+
+    SecretKey getTertiaryKey(
+            String packageName,
+            RecoverableKeyStoreSecondaryKey secondaryKey)
+            throws IllegalBlockSizeException, InvalidAlgorithmParameterException,
+            NoSuchAlgorithmException, IOException, NoSuchPaddingException,
+            InvalidKeyException {
+        TertiaryKeyManager tertiaryKeyManager =
+                new TertiaryKeyManager(
+                        mContext,
+                        sSecureRandom,
+                        TertiaryKeyRotationScheduler.getInstance(mContext),
+                        secondaryKey,
+                        packageName);
+        return tertiaryKeyManager.getKey();
+    }
+}
diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/KeyValueEncrypter.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/KeyValueEncrypter.java
new file mode 100644
index 000000000000..1d841b4a2c8f
--- /dev/null
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/KeyValueEncrypter.java
@@ -0,0 +1,143 @@
+/*
+ * 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;
+
+import static com.android.server.backup.encryption.BackupEncryptionService.TAG;
+
+import android.content.Context;
+import android.os.ParcelFileDescriptor;
+import android.util.Log;
+
+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.protos.nano.WrappedKeyProto;
+import com.android.server.backup.encryption.tasks.EncryptedKvBackupTask;
+import com.android.server.backup.encryption.tasks.EncryptedKvRestoreTask;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.security.SecureRandom;
+import java.util.Map;
+
+public class KeyValueEncrypter {
+    private final Context mContext;
+    private final EncryptionKeyHelper mKeyHelper;
+
+    public KeyValueEncrypter(Context context) {
+        mContext = context;
+        mKeyHelper = new EncryptionKeyHelper(mContext);
+    }
+
+    public void encryptKeyValueData(
+            String packageName, ParcelFileDescriptor inputFd, OutputStream outputStream)
+            throws Exception {
+        EncryptedKvBackupTask.EncryptedKvBackupTaskFactory backupTaskFactory =
+                new EncryptedKvBackupTask.EncryptedKvBackupTaskFactory();
+        EncryptedKvBackupTask backupTask =
+                backupTaskFactory.newInstance(
+                        mContext,
+                        new SecureRandom(),
+                        new FileBackupServer(outputStream),
+                        CryptoSettings.getInstance(mContext),
+                        mKeyHelper.getKeyManagerProvider(),
+                        inputFd,
+                        packageName);
+        backupTask.performBackup(/* incremental */ false);
+    }
+
+    public void decryptKeyValueData(String packageName,
+            InputStream encryptedInputStream, ParcelFileDescriptor outputFd) throws Exception {
+        RecoverableKeyStoreSecondaryKey secondaryKey = mKeyHelper.getActiveSecondaryKey();
+
+        EncryptedKvRestoreTask.EncryptedKvRestoreTaskFactory restoreTaskFactory =
+                new EncryptedKvRestoreTask.EncryptedKvRestoreTaskFactory();
+        EncryptedKvRestoreTask restoreTask =
+                restoreTaskFactory.newInstance(
+                        mContext,
+                        mKeyHelper.getKeyManagerProvider(),
+                        new InputStreamFullRestoreDownloader(encryptedInputStream),
+                        secondaryKey.getAlias(),
+                        KeyWrapUtils.wrap(
+                                secondaryKey.getSecretKey(),
+                                mKeyHelper.getTertiaryKey(packageName, secondaryKey)));
+
+        restoreTask.getRestoreData(outputFd);
+    }
+
+    // TODO(b/142455725): Extract into a commong class.
+    private static class FileBackupServer implements CryptoBackupServer {
+        private static final String EMPTY_DOC_ID = "";
+
+        private final OutputStream mOutputStream;
+
+        FileBackupServer(OutputStream outputStream) {
+            mOutputStream = outputStream;
+        }
+
+        @Override
+        public String uploadIncrementalBackup(
+                String packageName,
+                String oldDocId,
+                byte[] diffScript,
+                WrappedKeyProto.WrappedKey tertiaryKey) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public String uploadNonIncrementalBackup(
+                String packageName, byte[] data, WrappedKeyProto.WrappedKey tertiaryKey) {
+            try {
+                mOutputStream.write(data);
+            } catch (IOException e) {
+                Log.w(TAG, "Failed to write encrypted data to file: ", e);
+            }
+
+            return EMPTY_DOC_ID;
+        }
+
+        @Override
+        public void setActiveSecondaryKeyAlias(
+                String keyAlias, Map<String, WrappedKeyProto.WrappedKey> tertiaryKeys) {
+            // Do nothing.
+        }
+    }
+
+    // TODO(b/142455725): Extract into a commong class.
+    private static class InputStreamFullRestoreDownloader extends FullRestoreDownloader {
+        private final InputStream mInputStream;
+
+        InputStreamFullRestoreDownloader(InputStream inputStream) {
+            mInputStream = inputStream;
+        }
+
+        @Override
+        public int readNextChunk(byte[] buffer) throws IOException {
+            return mInputStream.read(buffer);
+        }
+
+        @Override
+        public void finish(FinishType finishType) {
+            try {
+                mInputStream.close();
+            } catch (IOException e) {
+                Log.w(TAG, "Error while reading restore data");
+            }
+        }
+    }
+}
diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransport.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransport.java
index 1d0224d49be7..c3cb335db89e 100644
--- a/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransport.java
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransport.java
@@ -18,27 +18,58 @@ package com.android.server.backup.encryption.transport;
 
 import static com.android.server.backup.encryption.BackupEncryptionService.TAG;
 
+import android.app.backup.BackupTransport;
+import android.app.backup.RestoreDescription;
+import android.content.Context;
+import android.content.pm.PackageInfo;
+import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.backup.IBackupTransport;
+import com.android.server.backup.encryption.KeyValueEncrypter;
 import com.android.server.backup.transport.DelegatingTransport;
 import com.android.server.backup.transport.TransportClient;
 
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.concurrent.atomic.AtomicReference;
+
 /**
  * This is an implementation of {@link IBackupTransport} that encrypts (or decrypts) the data when
  * sending it (or receiving it) from the {@link IBackupTransport} returned by {@link
  * TransportClient.connect(String)}.
  */
 public class IntermediateEncryptingTransport extends DelegatingTransport {
+    private static final String BACKUP_TEMP_DIR = "backup";
+    private static final String RESTORE_TEMP_DIR = "restore";
+
     private final TransportClient mTransportClient;
     private final Object mConnectLock = new Object();
+    private final Context mContext;
     private volatile IBackupTransport mRealTransport;
+    private AtomicReference<String> mNextRestorePackage = new AtomicReference<>();
+    private final KeyValueEncrypter mKeyValueEncrypter;
+    private final boolean mShouldEncrypt;
+
+    IntermediateEncryptingTransport(
+            TransportClient transportClient, Context context, boolean shouldEncrypt) {
+        this(transportClient, context, new KeyValueEncrypter(context), shouldEncrypt);
+    }
 
     @VisibleForTesting
-    IntermediateEncryptingTransport(TransportClient transportClient) {
+    IntermediateEncryptingTransport(
+            TransportClient transportClient, Context context, KeyValueEncrypter keyValueEncrypter,
+            boolean shouldEncrypt) {
         mTransportClient = transportClient;
+        mContext = context;
+        mKeyValueEncrypter = keyValueEncrypter;
+        mShouldEncrypt = shouldEncrypt;
     }
 
     @Override
@@ -46,9 +77,116 @@ public class IntermediateEncryptingTransport extends DelegatingTransport {
         if (mRealTransport == null) {
             connect();
         }
+        Log.d(TAG, "real transport = " + mRealTransport.name());
         return mRealTransport;
     }
 
+    @Override
+    public int performBackup(PackageInfo packageInfo, ParcelFileDescriptor inFd, int flags)
+            throws RemoteException {
+        if (!mShouldEncrypt) {
+            return super.performBackup(packageInfo, inFd, flags);
+        }
+
+        File encryptedStorageFile = getBackupTempStorage(packageInfo.packageName);
+        if (encryptedStorageFile == null) {
+            return BackupTransport.TRANSPORT_ERROR;
+        }
+
+        // Encrypt the backup data and write it into a temp file.
+        try (OutputStream encryptedOutput = new FileOutputStream(encryptedStorageFile)) {
+            mKeyValueEncrypter.encryptKeyValueData(packageInfo.packageName, inFd,
+                    encryptedOutput);
+        } catch (Throwable e) {
+            Log.e(TAG, "Failed to encrypt backup data: ", e);
+            return BackupTransport.TRANSPORT_ERROR;
+        }
+
+        // Pass the temp file to the real transport for backup.
+        try (FileInputStream encryptedInput = new FileInputStream(encryptedStorageFile)) {
+            return super.performBackup(
+                    packageInfo, ParcelFileDescriptor.dup(encryptedInput.getFD()), flags);
+        } catch (IOException e) {
+            Log.e(TAG, "Failed to read encrypted data from temp storage: ", e);
+            return BackupTransport.TRANSPORT_ERROR;
+        }
+    }
+
+    @Override
+    public int getRestoreData(ParcelFileDescriptor outFd) throws RemoteException {
+        if (!mShouldEncrypt) {
+            return super.getRestoreData(outFd);
+        }
+
+        String nextRestorePackage = mNextRestorePackage.get();
+        if (nextRestorePackage == null) {
+            Log.e(TAG, "No next restore package set");
+            return BackupTransport.TRANSPORT_ERROR;
+        }
+
+        File encryptedStorageFile = getRestoreTempStorage(nextRestorePackage);
+        if (encryptedStorageFile == null) {
+            return BackupTransport.TRANSPORT_ERROR;
+        }
+
+        // Get encrypted restore data from the real transport and write it into a temp file.
+        try (FileOutputStream outputStream = new FileOutputStream(encryptedStorageFile)) {
+            int status = super.getRestoreData(ParcelFileDescriptor.dup(outputStream.getFD()));
+            if (status != BackupTransport.TRANSPORT_OK) {
+                Log.e(TAG, "Failed to read restore data from transport, status = " + status);
+                return status;
+            }
+        } catch (IOException e) {
+            Log.e(TAG, "Failed to write encrypted data to temp storage: ", e);
+            return BackupTransport.TRANSPORT_ERROR;
+        }
+
+        // Decrypt the data and write it into the fd given by the real transport.
+        try (InputStream inputStream = new FileInputStream(encryptedStorageFile)) {
+            mKeyValueEncrypter.decryptKeyValueData(nextRestorePackage, inputStream, outFd);
+            encryptedStorageFile.delete();
+        } catch (Exception e) {
+            Log.e(TAG, "Failed to decrypt restored data: ", e);
+            return BackupTransport.TRANSPORT_ERROR;
+        }
+
+        return BackupTransport.TRANSPORT_OK;
+    }
+
+    @Override
+    public RestoreDescription nextRestorePackage() throws RemoteException {
+        if (!mShouldEncrypt) {
+            return super.nextRestorePackage();
+        }
+
+        RestoreDescription restoreDescription = super.nextRestorePackage();
+        mNextRestorePackage.set(restoreDescription.getPackageName());
+
+        return restoreDescription;
+    }
+
+    @VisibleForTesting
+    protected File getBackupTempStorage(String packageName) {
+        return getTempStorage(packageName, BACKUP_TEMP_DIR);
+    }
+
+    @VisibleForTesting
+    protected File getRestoreTempStorage(String packageName) {
+        return getTempStorage(packageName, RESTORE_TEMP_DIR);
+    }
+
+    private File getTempStorage(String packageName, String operationType) {
+        File encryptedDir = new File(mContext.getFilesDir(), operationType);
+        encryptedDir.mkdir();
+        File encryptedFile = new File(encryptedDir, packageName);
+        try {
+            encryptedFile.createNewFile();
+        } catch (IOException e) {
+            Log.e(TAG, "Failed to create temp file for encrypted data: ", e);
+        }
+        return encryptedFile;
+    }
+
     private void connect() throws RemoteException {
         Log.i(TAG, "connecting " + mTransportClient);
         synchronized (mConnectLock) {
@@ -65,4 +203,9 @@ public class IntermediateEncryptingTransport extends DelegatingTransport {
     TransportClient getClient() {
         return mTransportClient;
     }
+
+    @VisibleForTesting
+    void setNextRestorePackage(String nextRestorePackage) {
+        mNextRestorePackage.set(nextRestorePackage);
+    }
 }
diff --git a/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportManager.java b/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportManager.java
index 6e6d571aa3c7..7c4082c2a54d 100644
--- a/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportManager.java
+++ b/packages/BackupEncryption/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportManager.java
@@ -26,20 +26,20 @@ import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.backup.IBackupTransport;
+import com.android.internal.widget.LockPatternUtils;
 import com.android.server.backup.transport.TransportClientManager;
 import com.android.server.backup.transport.TransportStats;
 
 import java.util.HashMap;
 import java.util.Map;
 
-/**
- * Handles creation and cleanup of {@link IntermediateEncryptingTransport} instances.
- */
+/** Handles creation and cleanup of {@link IntermediateEncryptingTransport} instances. */
 public class IntermediateEncryptingTransportManager {
     private static final String CALLER = "IntermediateEncryptingTransportManager";
     private final TransportClientManager mTransportClientManager;
     private final Object mTransportsLock = new Object();
     private final Map<ComponentName, IntermediateEncryptingTransport> mTransports = new HashMap<>();
+    private Context mContext;
 
     @VisibleForTesting
     IntermediateEncryptingTransportManager(TransportClientManager transportClientManager) {
@@ -48,6 +48,7 @@ public class IntermediateEncryptingTransportManager {
 
     public IntermediateEncryptingTransportManager(Context context) {
         this(new TransportClientManager(UserHandle.myUserId(), context, new TransportStats()));
+        mContext = context;
     }
 
     /**
@@ -55,31 +56,42 @@ public class IntermediateEncryptingTransportManager {
      * provide a {@link IntermediateEncryptingTransport} which is an implementation of {@link
      * IBackupTransport} that encrypts (or decrypts) the data when sending it (or receiving it) from
      * the real {@link IBackupTransport}.
+     *
      * @param intent {@link Intent} created with a call to {@link
-     * TransportClientManager.getEncryptingTransportIntent(ComponentName)}.
+     *     TransportClientManager.getEncryptingTransportIntent(ComponentName)}.
      * @return
      */
     public IntermediateEncryptingTransport get(Intent intent) {
         Intent transportIntent = TransportClientManager.getRealTransportIntent(intent);
         Log.i(TAG, "get: intent:" + intent + " transportIntent:" + transportIntent);
         synchronized (mTransportsLock) {
-            return mTransports.computeIfAbsent(transportIntent.getComponent(),
-                    c -> create(transportIntent));
+            return mTransports.computeIfAbsent(
+                    transportIntent.getComponent(), c -> create(transportIntent));
         }
     }
 
-    /**
-     * Create an instance of {@link IntermediateEncryptingTransport}.
-     */
+    /** Create an instance of {@link IntermediateEncryptingTransport}. */
     private IntermediateEncryptingTransport create(Intent realTransportIntent) {
         Log.d(TAG, "create: intent:" + realTransportIntent);
-        return new IntermediateEncryptingTransport(mTransportClientManager.getTransportClient(
-                realTransportIntent.getComponent(), realTransportIntent.getExtras(), CALLER));
+
+        LockPatternUtils patternUtils = new LockPatternUtils(mContext);
+        boolean shouldEncrypt =
+                realTransportIntent.getComponent().getClassName().contains("EncryptedLocalTransport")
+                        && (patternUtils.isLockPatternEnabled(UserHandle.myUserId())
+                                || patternUtils.isLockPasswordEnabled(UserHandle.myUserId()));
+
+        return new IntermediateEncryptingTransport(
+                mTransportClientManager.getTransportClient(
+                        realTransportIntent.getComponent(),
+                        realTransportIntent.getExtras(),
+                        CALLER),
+                mContext,
+                shouldEncrypt);
     }
 
     /**
-     * Cleanup the {@link IntermediateEncryptingTransport} which was created by a call to
-     * {@link #get(Intent)} with this {@link Intent}.
+     * Cleanup the {@link IntermediateEncryptingTransport} which was created by a call to {@link
+     * #get(Intent)} with this {@link Intent}.
      */
     public void cleanup(Intent intent) {
         Intent transportIntent = TransportClientManager.getRealTransportIntent(intent);
diff --git a/packages/BackupEncryption/test/unittest/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportTest.java b/packages/BackupEncryption/test/unittest/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportTest.java
index cc4b0ab1bb36..a85b2e4498a6 100644
--- a/packages/BackupEncryption/test/unittest/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportTest.java
+++ b/packages/BackupEncryption/test/unittest/src/com/android/server/backup/encryption/transport/IntermediateEncryptingTransportTest.java
@@ -18,43 +18,71 @@ package com.android.server.backup.encryption.transport;
 
 import static junit.framework.Assert.assertEquals;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
+import android.app.backup.BackupTransport;
+import android.content.Context;
+import android.content.pm.PackageInfo;
+import android.os.ParcelFileDescriptor;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.backup.IBackupTransport;
+import com.android.server.backup.encryption.KeyValueEncrypter;
 import com.android.server.backup.transport.TransportClient;
 
 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 java.io.File;
+
 @Presubmit
 @RunWith(AndroidJUnit4.class)
 public class IntermediateEncryptingTransportTest {
+    private static final String TEST_PACKAGE_NAME = "test_package";
+
+    private IntermediateEncryptingTransport mIntermediateEncryptingTransport;
+    private final PackageInfo mTestPackage = new PackageInfo();
+
     @Mock private IBackupTransport mRealTransport;
     @Mock private TransportClient mTransportClient;
+    @Mock private ParcelFileDescriptor mParcelFileDescriptor;
+    @Mock private KeyValueEncrypter mKeyValueEncrypter;
+    @Mock private Context mContext;
 
-    private IntermediateEncryptingTransport mIntermediateEncryptingTransport;
+    @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
+
+    private File mTempFile;
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mIntermediateEncryptingTransport = new IntermediateEncryptingTransport(mTransportClient);
+
+        mIntermediateEncryptingTransport =
+                new IntermediateEncryptingTransport(
+                        mTransportClient, mContext, mKeyValueEncrypter, true);
+        mTestPackage.packageName = TEST_PACKAGE_NAME;
+        mTempFile = mTemporaryFolder.newFile();
+
+        when(mTransportClient.connect(anyString())).thenReturn(mRealTransport);
+        when(mRealTransport.getRestoreData(any())).thenReturn(BackupTransport.TRANSPORT_OK);
     }
 
     @Test
     public void testGetDelegate_callsConnect() throws Exception {
-        when(mTransportClient.connect(anyString())).thenReturn(mRealTransport);
-
         IBackupTransport ret = mIntermediateEncryptingTransport.getDelegate();
 
         assertEquals(mRealTransport, ret);
@@ -74,4 +102,79 @@ public class IntermediateEncryptingTransportTest {
         verify(mTransportClient, times(1)).connect(anyString());
         verifyNoMoreInteractions(mTransportClient);
     }
+
+    @Test
+    public void testPerformBackup_shouldEncryptTrue_encryptsDataAndPassesToDelegate()
+            throws Exception {
+        mIntermediateEncryptingTransport =
+                new TestIntermediateTransport(mTransportClient, mContext, mKeyValueEncrypter, true);
+
+        mIntermediateEncryptingTransport.performBackup(mTestPackage, mParcelFileDescriptor, 0);
+
+        verify(mKeyValueEncrypter, times(1))
+                .encryptKeyValueData(eq(TEST_PACKAGE_NAME), eq(mParcelFileDescriptor), any());
+        verify(mRealTransport, times(1)).performBackup(eq(mTestPackage), any(), eq(0));
+    }
+
+    @Test
+    public void testPerformBackup_shouldEncryptFalse_doesntEncryptDataAndPassedToDelegate()
+            throws Exception {
+        mIntermediateEncryptingTransport =
+                new TestIntermediateTransport(
+                        mTransportClient, mContext, mKeyValueEncrypter, false);
+
+        mIntermediateEncryptingTransport.performBackup(mTestPackage, mParcelFileDescriptor, 0);
+
+        verifyZeroInteractions(mKeyValueEncrypter);
+        verify(mRealTransport, times(1))
+                .performBackup(eq(mTestPackage), eq(mParcelFileDescriptor), eq(0));
+    }
+
+    @Test
+    public void testGetRestoreData_shouldEncryptTrue_decryptsDataAndPassesToDelegate()
+            throws Exception {
+        mIntermediateEncryptingTransport =
+                new TestIntermediateTransport(mTransportClient, mContext, mKeyValueEncrypter, true);
+        mIntermediateEncryptingTransport.setNextRestorePackage(TEST_PACKAGE_NAME);
+
+        mIntermediateEncryptingTransport.getRestoreData(mParcelFileDescriptor);
+
+        verify(mKeyValueEncrypter, times(1))
+                .decryptKeyValueData(eq(TEST_PACKAGE_NAME), any(), eq(mParcelFileDescriptor));
+        verify(mRealTransport, times(1)).getRestoreData(any());
+    }
+
+    @Test
+    public void testGetRestoreData_shouldEncryptFalse_doesntDecryptDataAndPassesToDelegate()
+            throws Exception {
+        mIntermediateEncryptingTransport =
+                new TestIntermediateTransport(
+                        mTransportClient, mContext, mKeyValueEncrypter, false);
+        mIntermediateEncryptingTransport.setNextRestorePackage(TEST_PACKAGE_NAME);
+
+        mIntermediateEncryptingTransport.getRestoreData(mParcelFileDescriptor);
+
+        verifyZeroInteractions(mKeyValueEncrypter);
+        verify(mRealTransport, times(1)).getRestoreData(eq(mParcelFileDescriptor));
+    }
+
+    private final class TestIntermediateTransport extends IntermediateEncryptingTransport {
+        TestIntermediateTransport(
+                TransportClient transportClient,
+                Context context,
+                KeyValueEncrypter keyValueEncrypter,
+                boolean shouldEncrypt) {
+            super(transportClient, context, keyValueEncrypter, shouldEncrypt);
+        }
+
+        @Override
+        protected File getBackupTempStorage(String packageName) {
+            return mTempFile;
+        }
+
+        @Override
+        protected File getRestoreTempStorage(String packageName) {
+            return mTempFile;
+        }
+    }
 }