Separate out encryption in PersistedTabDataStorage

Previously non-encryption and encryption were
both handled in the same file with a flag passed in.

With the new approach they are in separate files and
selecting the right file is driven by a configuration
offering better separation of concerns.

Bug: 1059638
Change-Id: Icdb8ff751295a913123e9b16f36198efcb9785fa
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2106491
Commit-Queue: David Maunder <davidjm@chromium.org>
Reviewed-by: David Trainor <dtrainor@chromium.org>
Reviewed-by: Tommy Nyquist <nyquist@chromium.org>
Cr-Commit-Position: refs/heads/master@{#753264}

chrome/android/chrome_java_sources.gni

@@ -1601,6 +1601,7 @@ chrome_java_sources = [
   "java/src/org/chromium/chrome/browser/tab/TabWebContentsUserData.java",
   "java/src/org/chromium/chrome/browser/tab/TrustedCdn.java",
   "java/src/org/chromium/chrome/browser/tab/state/CriticalPersistedTabData.java",
+  "java/src/org/chromium/chrome/browser/tab/state/EncryptedFilePersistedTabDataStorage.java",
   "java/src/org/chromium/chrome/browser/tab/state/FilePersistedTabDataStorage.java",
   "java/src/org/chromium/chrome/browser/tab/state/PersistedTabData.java",
   "java/src/org/chromium/chrome/browser/tab/state/PersistedTabDataConfiguration.java",

chrome/android/java/src/org/chromium/chrome/browser/tab/DEPS

@@ -19,6 +19,7 @@ include_rules = [
   "+components/browser_ui/styles/android",
   "+components/browser_ui/widget/android",
   "+content/public/android/java/src/org/chromium/content_public",
+  "+chrome/browser/android/crypto/java/src/org/chromium/chrome/browser/crypto/CipherFactory.java",
   "+chrome/android/java/src/org/chromium/chrome/browser/tabmodel/TabbedModeTabPersistencePolicy.java",
   "+chrome/android/public/crypto/java/src/org/chromium/chrome/browser/crypto/CipherFactory.java",
   "+chrome/browser/profiles/android/java/src/org/chromium/chrome/browser/profiles/Profile.java",

chrome/android/java/src/org/chromium/chrome/browser/tab/state/CriticalPersistedTabData.java

@@ -104,8 +104,8 @@ public class CriticalPersistedTabData extends PersistedTabData {
         // from the {@link Tab}.
         if (tabImpl.isInitialized()) {
             TabState.WebContentsState webContentsState = TabState.getWebContentsState(tabImpl);
-            PersistedTabDataConfiguration config =
-                    PersistedTabDataConfiguration.get(CriticalPersistedTabData.class);
+            PersistedTabDataConfiguration config = PersistedTabDataConfiguration.get(
+                    CriticalPersistedTabData.class, tab.isIncognito());
             CriticalPersistedTabData criticalPersistedTabData = new CriticalPersistedTabData(tab,
                     tab.getParentId(), tabImpl.getRootId(), tab.getTimestampMillis(),
                     webContentsState != null

chrome/android/java/src/org/chromium/chrome/browser/tab/state/EncryptedFilePersistedTabDataStorage.java

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package org.chromium.chrome.browser.tab.state;
+
+import org.chromium.base.Callback;
+import org.chromium.base.Log;
+import org.chromium.chrome.browser.crypto.CipherFactory;
+
+import java.util.Locale;
+
+import javax.crypto.BadPaddingException;
+import javax.crypto.Cipher;
+import javax.crypto.IllegalBlockSizeException;
+
+/**
+ * Implements {@link PersistedTabDataStorage} but encrypts and decrypts
+ * as data is stored and retrieved respectively.
+ */
+public class EncryptedFilePersistedTabDataStorage extends FilePersistedTabDataStorage {
+    private static final String TAG = "EFPTDS";
+
+    @Override
+    public void save(int tabId, String dataId, byte[] data) {
+        super.save(tabId, dataId, encrypt(data));
+    }
+
+    @Override
+    public void restore(int tabId, String dataId, Callback<byte[]> callback) {
+        super.restore(tabId, dataId, (data) -> { callback.onResult(decrypt(data)); });
+    }
+
+    private static byte[] decrypt(byte[] data) {
+        try {
+            if (data == null) {
+                return null;
+            }
+            return CipherFactory.getInstance().getCipher(Cipher.DECRYPT_MODE).doFinal(data);
+        } catch (BadPaddingException | IllegalBlockSizeException e) {
+            Log.e(TAG,
+                    String.format(
+                            Locale.ENGLISH, "Error encrypting data. Details: %s", e.getMessage()));
+            return null;
+        }
+    }
+
+    private static byte[] encrypt(byte[] data) {
+        Cipher cipher = CipherFactory.getInstance().getCipher(Cipher.ENCRYPT_MODE);
+        try {
+            return cipher.doFinal(data);
+        } catch (BadPaddingException | IllegalBlockSizeException e) {
+            Log.e(TAG,
+                    String.format(Locale.ENGLISH, "Problem encrypting data. Details: %s",
+                            e.getMessage()));
+            return null;
+        }
+    }
+}

chrome/android/java/src/org/chromium/chrome/browser/tab/state/FilePersistedTabDataStorage.java

@@ -11,8 +11,9 @@ import androidx.annotation.VisibleForTesting;
 import org.chromium.base.Callback;
 import org.chromium.base.Log;
 import org.chromium.base.StreamUtil;
-import org.chromium.chrome.browser.crypto.CipherFactory;
+import org.chromium.base.task.PostTask;
 import org.chromium.chrome.browser.tabmodel.TabbedModeTabPersistencePolicy;
+import org.chromium.content_public.browser.UiThreadTaskTraits;
 
 import java.io.File;
 import java.io.FileNotFoundException;
@@ -20,10 +21,6 @@ import java.io.FileOutputStream;
 import java.io.IOException;
 import java.util.Locale;
 
-import javax.crypto.BadPaddingException;
-import javax.crypto.Cipher;
-import javax.crypto.IllegalBlockSizeException;
-
 /**
  * {@link PersistedTabDataStorage} which uses a file for the storage
  */
@@ -31,24 +28,12 @@ public class FilePersistedTabDataStorage implements PersistedTabDataStorage {
     private static final String TAG = "FilePTDS";
 
     @Override
-    public void save(int tabId, boolean isEncrypted, String dataId, byte[] data) {
+    public void save(int tabId, String dataId, byte[] data) {
         // TODO(crbug.com/1059636) these should be queued and executed on a background thread
         // TODO(crbug.com/1059637) we should introduce a retry mechanisms
-        // TODO(crbug.com/1059638) abstract out encrypt/decrypt
-        if (isEncrypted) {
-            Cipher cipher = CipherFactory.getInstance().getCipher(Cipher.ENCRYPT_MODE);
-            try {
-                data = cipher.doFinal(data);
-            } catch (BadPaddingException | IllegalBlockSizeException e) {
-                Log.e(TAG,
-                        String.format(Locale.ENGLISH, "Problem encrypting data. Details: %s",
-                                e.getMessage()));
-                return;
-            }
-        }
         FileOutputStream outputStream = null;
         try {
-            outputStream = new FileOutputStream(getFile(tabId, isEncrypted, dataId));
+            outputStream = new FileOutputStream(getFile(tabId, dataId));
             outputStream.write(data);
         } catch (FileNotFoundException e) {
             Log.e(TAG,
@@ -68,41 +53,32 @@ public class FilePersistedTabDataStorage implements PersistedTabDataStorage {
     }
 
     @Override
-    public void restore(int tabId, boolean isEncrypted, String dataId, Callback<byte[]> callback) {
-        File file = getFile(tabId, isEncrypted, dataId);
+    public void restore(int tabId, String dataId, Callback<byte[]> callback) {
+        File file = getFile(tabId, dataId);
         try {
             AtomicFile atomicFile = new AtomicFile(file);
             byte[] res = atomicFile.readFully();
-            if (isEncrypted) {
-                Cipher cipher = CipherFactory.getInstance().getCipher(Cipher.DECRYPT_MODE);
-                res = cipher.doFinal(res);
-            }
-            callback.onResult(res);
+            PostTask.postTask(UiThreadTaskTraits.DEFAULT, () -> { callback.onResult(res); });
         } catch (FileNotFoundException e) {
             Log.e(TAG,
                     String.format(Locale.ENGLISH,
                             "FileNotFoundException while attempting to restore "
                                     + " for Tab %d. Details: %s",
                             tabId, e.getMessage()));
-            callback.onResult(null);
+            PostTask.postTask(UiThreadTaskTraits.DEFAULT, () -> { callback.onResult(null); });
         } catch (IOException e) {
             Log.e(TAG,
                     String.format(Locale.ENGLISH,
                             "IOException while attempting to restore for Tab "
                                     + "%d. Details: %s",
                             tabId, e.getMessage()));
-            callback.onResult(null);
-        } catch (BadPaddingException | IllegalBlockSizeException e) {
-            Log.e(TAG,
-                    String.format(
-                            Locale.ENGLISH, "Error encrypting data. Details: %s", e.getMessage()));
-            callback.onResult(null);
+            PostTask.postTask(UiThreadTaskTraits.DEFAULT, () -> { callback.onResult(null); });
         }
     }
 
     @Override
-    public void delete(int tabId, boolean isEncrypted, String dataId) {
-        File file = getFile(tabId, isEncrypted, dataId);
+    public void delete(int tabId, String dataId) {
+        File file = getFile(tabId, dataId);
         if (!file.exists()) {
             return;
         }
@@ -113,9 +89,8 @@ public class FilePersistedTabDataStorage implements PersistedTabDataStorage {
     }
 
     @VisibleForTesting
-    protected File getFile(int tabId, boolean isEncrypted, String dataId) {
-        String encryptedId = isEncrypted ? "Encrypted" : "NotEncrypted";
+    protected File getFile(int tabId, String dataId) {
         return new File(TabbedModeTabPersistencePolicy.getOrCreateTabbedModeStateDirectory(),
-                String.format(Locale.ENGLISH, "%d%s%s", tabId, encryptedId, dataId));
+                String.format(Locale.ENGLISH, "%d%s", tabId, dataId));
     }
 }

chrome/android/java/src/org/chromium/chrome/browser/tab/state/PersistedTabData.java

@@ -74,8 +74,9 @@ public abstract class PersistedTabData implements UserData {
             callback.onResult(persistedTabDataFromTab);
             return;
         }
-        PersistedTabDataConfiguration config = PersistedTabDataConfiguration.get(clazz);
-        config.storage.restore(tab.getId(), tab.isIncognito(), config.id, (data) -> {
+        PersistedTabDataConfiguration config =
+                PersistedTabDataConfiguration.get(clazz, tab.isIncognito());
+        config.storage.restore(tab.getId(), config.id, (data) -> {
             T persistedTabData;
             if (data == null) {
                 persistedTabData = supplier.get();
@@ -116,8 +117,7 @@ public abstract class PersistedTabData implements UserData {
      */
     @VisibleForTesting
     protected void save() {
-        mPersistedTabDataStorage.save(
-                mTab.getId(), mTab.isIncognito(), mPersistedTabDataId, serialize());
+        mPersistedTabDataStorage.save(mTab.getId(), mPersistedTabDataId, serialize());
     }
 
     /**
@@ -138,7 +138,7 @@ public abstract class PersistedTabData implements UserData {
      * @param profile profile
      */
     protected void delete() {
-        mPersistedTabDataStorage.delete(mTab.getId(), mTab.isIncognito(), mPersistedTabDataId);
+        mPersistedTabDataStorage.delete(mTab.getId(), mPersistedTabDataId);
     }
 
     /**

chrome/android/java/src/org/chromium/chrome/browser/tab/state/PersistedTabDataConfiguration.java

@@ -14,14 +14,18 @@ import java.util.Map;
 public enum PersistedTabDataConfiguration {
     // TODO(crbug.com/1059650) investigate should this go in the app code?
     // Also investigate if the storage instance should be shared.
-    CRITICAL_PERSISTED_TAB_DATA("CPTD", new FilePersistedTabDataStorage());
+    CRITICAL_PERSISTED_TAB_DATA("CPTD", new FilePersistedTabDataStorage()),
+    ENCRYPTED_CRITICAL_PERSISTED_TAB_DATA("ECPTD", new EncryptedFilePersistedTabDataStorage());
 
     private static final Map<Class<? extends PersistedTabData>, PersistedTabDataConfiguration>
             sLookup = new HashMap<>();
+    private static final Map<Class<? extends PersistedTabData>, PersistedTabDataConfiguration>
+            sEncryptedLookup = new HashMap<>();
 
     static {
         // TODO(crbug.com/1060187) remove static initializer and initialization lazy
         sLookup.put(CriticalPersistedTabData.class, CRITICAL_PERSISTED_TAB_DATA);
+        sEncryptedLookup.put(CriticalPersistedTabData.class, ENCRYPTED_CRITICAL_PERSISTED_TAB_DATA);
     }
 
     public final String id;
@@ -39,7 +43,11 @@ public enum PersistedTabDataConfiguration {
     /**
      * Acquire {@link PersistedTabDataConfiguration} for a given {@link PersistedTabData} class
      */
-    public static PersistedTabDataConfiguration get(Class<? extends PersistedTabData> clazz) {
+    public static PersistedTabDataConfiguration get(
+            Class<? extends PersistedTabData> clazz, boolean isEncrypted) {
+        if (isEncrypted) {
+            return sEncryptedLookup.get(clazz);
+        }
         return sLookup.get(clazz);
     }
 }

chrome/android/java/src/org/chromium/chrome/browser/tab/state/PersistedTabDataStorage.java

@@ -12,24 +12,21 @@ import org.chromium.base.Callback;
 public interface PersistedTabDataStorage {
     /**
      * @param tabId identifier for the {@link Tab}
-     * @param isEncrypted true if the data should be encrypted (i.e. for incognito Tabs)
      * @param tabDataId unique identifier representing the type {@link PersistedTabData}
      * @param data serialized {@link PersistedTabData}
      */
-    void save(int tabId, boolean isEncrypted, String tabDataId, byte[] data);
+    void save(int tabId, String tabDataId, byte[] data);
 
     /**
      * @param tabId identifier for the {@link Tab}
-     * @param isEncrypted true if the stored data is encrypted (i.e. for incognito Tabs)
      * @param tabDataId unique identifier representing the type of {@link PersistedTabData}
      * @param callback to pass back the seraizliaed {@link PersistedTabData} in
      */
-    void restore(int tabId, boolean isEncrypted, String tabDataId, Callback<byte[]> callback);
+    void restore(int tabId, String tabDataId, Callback<byte[]> callback);
 
     /**
      * @param tabId identifier for the {@link Tab}
-     * @param isEncrypted true if the stored data is encrypted (i.e. for incognito Tabs)
      * @param tabDataId unique identifier representing the type of {@link PersistedTabData}
      */
-    void delete(int tabId, boolean isEncrypted, String tabDataId);
+    void delete(int tabId, String tabDataId);
 }

chrome/android/javatests/src/org/chromium/chrome/browser/tab/state/CriticalPersistedTabDataTest.java

@@ -4,9 +4,6 @@
 
 package org.chromium.chrome.browser.tab.state;
 
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-
 import android.support.test.filters.SmallTest;
 
 import org.junit.Assert;
@@ -14,16 +11,17 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
-import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import org.chromium.base.Callback;
+import org.chromium.base.ThreadUtils;
 import org.chromium.base.test.BaseJUnit4ClassRunner;
 import org.chromium.chrome.browser.tab.MockTab;
 import org.chromium.chrome.browser.tab.Tab;
 import org.chromium.chrome.test.ChromeBrowserTestRule;
 
+import java.util.concurrent.Semaphore;
+
 /**
  * Test relating to {@link CriticalPersistedTabData}
  */
@@ -43,14 +41,7 @@ public class CriticalPersistedTabDataTest {
     private static final int THEME_COLOR = 5;
     private static final int LAUNCH_TYPE_AT_CREATION = 3;
 
-    @Mock
-    private Callback<Boolean> mBooleanCallback;
-
-    @Mock
-    private Callback<CriticalPersistedTabData> mRestoreCallback;
-
-    @Mock
-    private Callback<CriticalPersistedTabData> mDeleteCallback;
+    private CriticalPersistedTabData mCriticalPersistedTabData;
 
     private static Tab mockTab(int id, boolean isEncrypted) {
         return new MockTab(id, isEncrypted);
@@ -63,44 +54,54 @@ public class CriticalPersistedTabDataTest {
 
     @SmallTest
     @Test
-    public void testNonEncryptedSaveRestore() {
+    public void testNonEncryptedSaveRestore() throws InterruptedException {
         testSaveRestoreDelete(false);
     }
 
     @SmallTest
     @Test
-    public void testEncryptedSaveRestoreDelete() {
+    public void testEncryptedSaveRestoreDelete() throws InterruptedException {
         testSaveRestoreDelete(true);
     }
 
-    private void testSaveRestoreDelete(boolean isEncrypted) {
-        PersistedTabDataConfiguration config =
-                PersistedTabDataConfiguration.get(CriticalPersistedTabData.class);
-        CriticalPersistedTabData criticalPersistedTabData =
-                new CriticalPersistedTabData(mockTab(TAB_ID, isEncrypted), PARENT_ID, ROOT_ID,
-                        TIMESTAMP, CONTENT_STATE, CONTENT_STATE_VERSION, OPENER_APP_ID, THEME_COLOR,
-                        LAUNCH_TYPE_AT_CREATION, config.storage, config.id);
-        criticalPersistedTabData.save();
-        CriticalPersistedTabData.from(mockTab(TAB_ID, isEncrypted), mRestoreCallback);
-        ArgumentCaptor<CriticalPersistedTabData> cptdCaptor =
-                ArgumentCaptor.forClass(CriticalPersistedTabData.class);
-        verify(mRestoreCallback, times(1)).onResult(cptdCaptor.capture());
-        CriticalPersistedTabData restored = cptdCaptor.getValue();
-        Assert.assertNotNull(restored);
-        Assert.assertEquals(restored.getParentId(), PARENT_ID);
-        Assert.assertEquals(restored.getRootId(), ROOT_ID);
-        Assert.assertEquals(restored.getTimestampMillis(), TIMESTAMP);
-        Assert.assertEquals(restored.getContentStateVersion(), CONTENT_STATE_VERSION);
-        Assert.assertEquals(restored.getOpenerAppId(), OPENER_APP_ID);
-        Assert.assertEquals(restored.getThemeColor(), THEME_COLOR);
-        Assert.assertEquals(restored.getTabLaunchTypeAtCreation(), LAUNCH_TYPE_AT_CREATION);
-        Assert.assertArrayEquals(restored.getContentStateBytes(), CONTENT_STATE);
-        restored.delete();
-        CriticalPersistedTabData.from(mockTab(TAB_ID, isEncrypted), mDeleteCallback);
-        ArgumentCaptor<CriticalPersistedTabData> cptdDeletedCaptor =
-                ArgumentCaptor.forClass(CriticalPersistedTabData.class);
-        verify(mDeleteCallback, times(1)).onResult(cptdDeletedCaptor.capture());
-        Assert.assertNull(cptdDeletedCaptor.getValue());
+    private void testSaveRestoreDelete(boolean isEncrypted) throws InterruptedException {
+        final Semaphore semaphore = new Semaphore(0);
+        Callback<CriticalPersistedTabData> callback = new Callback<CriticalPersistedTabData>() {
+            @Override
+            public void onResult(CriticalPersistedTabData res) {
+                mCriticalPersistedTabData = res;
+                semaphore.release();
+            }
+        };
+        ThreadUtils.runOnUiThreadBlocking(() -> {
+            PersistedTabDataConfiguration config =
+                    PersistedTabDataConfiguration.get(CriticalPersistedTabData.class, isEncrypted);
+            CriticalPersistedTabData criticalPersistedTabData =
+                    new CriticalPersistedTabData(mockTab(TAB_ID, isEncrypted), PARENT_ID, ROOT_ID,
+                            TIMESTAMP, CONTENT_STATE, CONTENT_STATE_VERSION, OPENER_APP_ID,
+                            THEME_COLOR, LAUNCH_TYPE_AT_CREATION, config.storage, config.id);
+            criticalPersistedTabData.save();
+            CriticalPersistedTabData.from(mockTab(TAB_ID, isEncrypted), callback);
+        });
+        semaphore.acquire();
+        Assert.assertNotNull(mCriticalPersistedTabData);
+        Assert.assertEquals(mCriticalPersistedTabData.getParentId(), PARENT_ID);
+        Assert.assertEquals(mCriticalPersistedTabData.getRootId(), ROOT_ID);
+        Assert.assertEquals(mCriticalPersistedTabData.getTimestampMillis(), TIMESTAMP);
+        Assert.assertEquals(
+                mCriticalPersistedTabData.getContentStateVersion(), CONTENT_STATE_VERSION);
+        Assert.assertEquals(mCriticalPersistedTabData.getOpenerAppId(), OPENER_APP_ID);
+        Assert.assertEquals(mCriticalPersistedTabData.getThemeColor(), THEME_COLOR);
+        Assert.assertEquals(
+                mCriticalPersistedTabData.getTabLaunchTypeAtCreation(), LAUNCH_TYPE_AT_CREATION);
+        Assert.assertArrayEquals(mCriticalPersistedTabData.getContentStateBytes(), CONTENT_STATE);
+
+        ThreadUtils.runOnUiThreadBlocking(() -> {
+            mCriticalPersistedTabData.delete();
+            CriticalPersistedTabData.from(mockTab(TAB_ID, isEncrypted), callback);
+        });
+        semaphore.acquire();
+        Assert.assertNull(mCriticalPersistedTabData);
         // TODO(crbug.com/1060232) test restored.save() after restored.delete()
         // Also cover
         // - Multiple (different) TAB_IDs being stored in the same storage.

chrome/android/javatests/src/org/chromium/chrome/browser/tab/state/FilePersistedTabDataStorageTest.java

@@ -4,9 +4,6 @@
 
 package org.chromium.chrome.browser.tab.state;
 
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-
 import android.support.test.filters.SmallTest;
 
 import org.junit.Assert;
@@ -14,15 +11,16 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import org.chromium.base.Callback;
+import org.chromium.base.ThreadUtils;
 import org.chromium.base.test.BaseJUnit4ClassRunner;
 import org.chromium.chrome.test.ChromeBrowserTestRule;
 
 import java.io.File;
+import java.util.concurrent.Semaphore;
 
 /**
  * Tests relating to  {@link FilePersistedTabDataStorage}
@@ -39,6 +37,8 @@ public class FilePersistedTabDataStorageTest {
     private static final byte[] DATA = {13, 14};
     private static final String DATA_ID = "DataId";
 
+    private byte[] mResult;
+
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
@@ -46,27 +46,41 @@ public class FilePersistedTabDataStorageTest {
 
     @SmallTest
     @Test
-    public void testFilePersistedDataStorageNonEncrypted() {
-        testFilePersistedDataStorage(false);
+    public void testFilePersistedDataStorageNonEncrypted() throws InterruptedException {
+        testFilePersistedDataStorage(new FilePersistedTabDataStorage());
     }
 
     @SmallTest
     @Test
-    public void testFilePersistedDataStorageEncrypted() {
-        testFilePersistedDataStorage(true);
+    public void testFilePersistedDataStorageEncrypted() throws InterruptedException {
+        testFilePersistedDataStorage(new EncryptedFilePersistedTabDataStorage());
     }
 
-    private void testFilePersistedDataStorage(boolean isEncrypted) {
-        FilePersistedTabDataStorage filePersistedTabDataStorage = new FilePersistedTabDataStorage();
-        filePersistedTabDataStorage.save(TAB_ID, isEncrypted, DATA_ID, DATA);
-        ArgumentCaptor<byte[]> byteArrayCaptor = ArgumentCaptor.forClass(byte[].class);
-        filePersistedTabDataStorage.restore(TAB_ID, isEncrypted, DATA_ID, mByteArrayCallback);
-        verify(mByteArrayCallback, times(1)).onResult(byteArrayCaptor.capture());
-        Assert.assertEquals(byteArrayCaptor.getValue().length, 2);
-        Assert.assertArrayEquals(byteArrayCaptor.getValue(), DATA);
-        File file = filePersistedTabDataStorage.getFile(TAB_ID, isEncrypted, DATA_ID);
+    private void testFilePersistedDataStorage(FilePersistedTabDataStorage persistedTabDataStorage)
+            throws InterruptedException {
+        final Semaphore semaphore = new Semaphore(0);
+        Callback<byte[]> callback = new Callback<byte[]>() {
+            @Override
+            public void onResult(byte[] res) {
+                mResult = res;
+                semaphore.release();
+            }
+        };
+        ThreadUtils.runOnUiThreadBlocking(() -> {
+            persistedTabDataStorage.save(TAB_ID, DATA_ID, DATA);
+            persistedTabDataStorage.restore(TAB_ID, DATA_ID, callback);
+        });
+        semaphore.acquire();
+        Assert.assertEquals(mResult.length, 2);
+        Assert.assertArrayEquals(mResult, DATA);
+        File file = persistedTabDataStorage.getFile(TAB_ID, DATA_ID);
         Assert.assertTrue(file.exists());
-        filePersistedTabDataStorage.delete(TAB_ID, isEncrypted, DATA_ID);
+
+        ThreadUtils.runOnUiThreadBlocking(() -> {
+            persistedTabDataStorage.delete(TAB_ID, DATA_ID);
+            semaphore.release();
+        });
+        semaphore.acquire();
         Assert.assertFalse(file.exists());
     }
 }