Get blocked restore keys directly from UserBMS

Bug: 145126096
Test: atest KeyValueRestoreExclusionHostSideTest
      atest PerformUnifiedRestoreHostSideTest
      atest UserBackupPreferencesTest

Currently PerformUnifiedRestoreTask gets the list of blocked restore
keys at construction. However, at that point the list might not be fully
constructed yet. We should get the keys through the getter avaialble in
UserBMS when we need them.

Change-Id: I62ad34138ba7a893e66d6af05d2e242c9c964a44

9 files changed, 67 insertions, 94 deletions

diff --git a/services/backup/java/com/android/server/backup/UserBackupManagerService.java b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
index e3d2dcc8141f..6247a635233a 100644
--- a/services/backup/java/com/android/server/backup/UserBackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
@@ -1080,12 +1080,8 @@ public class UserBackupManagerService {
         }
     }
 
-    public Map<String, Set<String>> getExcludedRestoreKeys(String... packages) {
-        return mBackupPreferences.getExcludedRestoreKeysForPackages(packages);
-    }
-
-    public Map<String, Set<String>> getAllExcludedRestoreKeys() {
-        return mBackupPreferences.getAllExcludedRestoreKeys();
+    public Set<String> getExcludedRestoreKeys(String packageName) {
+        return mBackupPreferences.getExcludedRestoreKeysForPackage(packageName);
     }
 
     /** Used for generating random salts or passwords. */
@@ -3356,8 +3352,7 @@ public class UserBackupManagerService {
                                 restoreSet,
                                 packageName,
                                 token,
-                                listener,
-                                getExcludedRestoreKeys(packageName));
+                                listener);
                 mBackupHandler.sendMessage(msg);
             } catch (Exception e) {
                 // Calling into the transport broke; back off and proceed with the installation.
diff --git a/services/backup/java/com/android/server/backup/UserBackupPreferences.java b/services/backup/java/com/android/server/backup/UserBackupPreferences.java
index 41b9719d273b..bb8bf52187c5 100644
--- a/services/backup/java/com/android/server/backup/UserBackupPreferences.java
+++ b/services/backup/java/com/android/server/backup/UserBackupPreferences.java
@@ -48,16 +48,7 @@ public class UserBackupPreferences {
         mEditor.commit();
     }
 
-    Map<String, Set<String>> getExcludedRestoreKeysForPackages(String... packages) {
-        Map<String, Set<String>> excludedKeys = new HashMap<>();
-        for (String packageName : packages) {
-            excludedKeys.put(packageName,
-                    mPreferences.getStringSet(packageName, Collections.emptySet()));
-        }
-        return excludedKeys;
-    }
-
-    Map<String, Set<String>> getAllExcludedRestoreKeys() {
-        return (Map<String, Set<String>>) mPreferences.getAll();
+    Set<String> getExcludedRestoreKeysForPackage(String packageName) {
+        return mPreferences.getStringSet(packageName, Collections.emptySet());
     }
 }
diff --git a/services/backup/java/com/android/server/backup/internal/BackupHandler.java b/services/backup/java/com/android/server/backup/internal/BackupHandler.java
index 05396f36b364..87a8e4982529 100644
--- a/services/backup/java/com/android/server/backup/internal/BackupHandler.java
+++ b/services/backup/java/com/android/server/backup/internal/BackupHandler.java
@@ -299,8 +299,7 @@ public class BackupHandler extends Handler {
                                 params.pmToken,
                                 params.isSystemRestore,
                                 params.filterSet,
-                                params.listener,
-                                params.excludedKeys);
+                                params.listener);
 
                 synchronized (backupManagerService.getPendingRestores()) {
                     if (backupManagerService.isRestoreInProgress()) {
diff --git a/services/backup/java/com/android/server/backup/params/RestoreParams.java b/services/backup/java/com/android/server/backup/params/RestoreParams.java
index 09b7e3535e2e..a6fea6cc75a0 100644
--- a/services/backup/java/com/android/server/backup/params/RestoreParams.java
+++ b/services/backup/java/com/android/server/backup/params/RestoreParams.java
@@ -37,7 +37,6 @@ public class RestoreParams {
     public final boolean isSystemRestore;
     @Nullable public final String[] filterSet;
     public final OnTaskFinishedListener listener;
-    public final Map<String, Set<String>> excludedKeys;
 
     /**
      * No kill after restore.
@@ -48,8 +47,7 @@ public class RestoreParams {
             IBackupManagerMonitor monitor,
             long token,
             PackageInfo packageInfo,
-            OnTaskFinishedListener listener,
-            Map<String, Set<String>> excludedKeys) {
+            OnTaskFinishedListener listener) {
         return new RestoreParams(
                 transportClient,
                 observer,
@@ -59,8 +57,7 @@ public class RestoreParams {
                 /* pmToken */ 0,
                 /* isSystemRestore */ false,
                 /* filterSet */ null,
-                listener,
-                excludedKeys);
+                listener);
     }
 
     /**
@@ -73,8 +70,7 @@ public class RestoreParams {
             long token,
             String packageName,
             int pmToken,
-            OnTaskFinishedListener listener,
-            Map<String, Set<String>> excludedKeys) {
+            OnTaskFinishedListener listener) {
         String[] filterSet = {packageName};
         return new RestoreParams(
                 transportClient,
@@ -85,8 +81,7 @@ public class RestoreParams {
                 pmToken,
                 /* isSystemRestore */ false,
                 filterSet,
-                listener,
-                excludedKeys);
+                listener);
     }
 
     /**
@@ -97,8 +92,7 @@ public class RestoreParams {
             IRestoreObserver observer,
             IBackupManagerMonitor monitor,
             long token,
-            OnTaskFinishedListener listener,
-            Map<String, Set<String>> excludedKeys) {
+            OnTaskFinishedListener listener) {
         return new RestoreParams(
                 transportClient,
                 observer,
@@ -108,8 +102,7 @@ public class RestoreParams {
                 /* pmToken */ 0,
                 /* isSystemRestore */ true,
                 /* filterSet */ null,
-                listener,
-                excludedKeys);
+                listener);
     }
 
     /**
@@ -122,8 +115,7 @@ public class RestoreParams {
             long token,
             String[] filterSet,
             boolean isSystemRestore,
-            OnTaskFinishedListener listener,
-            Map<String, Set<String>> excludedKeys) {
+            OnTaskFinishedListener listener) {
         return new RestoreParams(
                 transportClient,
                 observer,
@@ -133,8 +125,7 @@ public class RestoreParams {
                 /* pmToken */ 0,
                 isSystemRestore,
                 filterSet,
-                listener,
-                excludedKeys);
+                listener);
     }
 
     private RestoreParams(
@@ -146,8 +137,7 @@ public class RestoreParams {
             int pmToken,
             boolean isSystemRestore,
             @Nullable String[] filterSet,
-            OnTaskFinishedListener listener,
-            Map<String, Set<String>> excludedKeys) {
+            OnTaskFinishedListener listener) {
         this.transportClient = transportClient;
         this.observer = observer;
         this.monitor = monitor;
@@ -157,6 +147,5 @@ public class RestoreParams {
         this.isSystemRestore = isSystemRestore;
         this.filterSet = filterSet;
         this.listener = listener;
-        this.excludedKeys = excludedKeys;
     }
 }
diff --git a/services/backup/java/com/android/server/backup/restore/ActiveRestoreSession.java b/services/backup/java/com/android/server/backup/restore/ActiveRestoreSession.java
index c0f76c39e04f..5a57cdc39402 100644
--- a/services/backup/java/com/android/server/backup/restore/ActiveRestoreSession.java
+++ b/services/backup/java/com/android/server/backup/restore/ActiveRestoreSession.java
@@ -178,8 +178,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
                                                 observer,
                                                 monitor,
                                                 token,
-                                                listener,
-                                                mBackupManagerService.getAllExcludedRestoreKeys()),
+                                                listener),
                                 "RestoreSession.restoreAll()");
                     } finally {
                         Binder.restoreCallingIdentity(oldId);
@@ -272,8 +271,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
                                                 token,
                                                 packages,
                                                 /* isSystemRestore */ packages.length > 1,
-                                                listener,
-                                                mBackupManagerService.getExcludedRestoreKeys(packages)),
+                                                listener),
                                 "RestoreSession.restorePackages(" + packages.length + " packages)");
                     } finally {
                         Binder.restoreCallingIdentity(oldId);
@@ -365,8 +363,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
                                     monitor,
                                     token,
                                     app,
-                                    listener,
-                                    mBackupManagerService.getExcludedRestoreKeys(app.packageName)),
+                                    listener),
                     "RestoreSession.restorePackage(" + packageName + ")");
         } finally {
             Binder.restoreCallingIdentity(oldId);
diff --git a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
index f90d936a5f4d..3c37f737f8be 100644
--- a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
+++ b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
@@ -155,8 +155,6 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
     // When finished call listener
     private final OnTaskFinishedListener mListener;
 
-    private final Map<String, Set<String>> mExcludedKeys;
-
     // Key/value: bookkeeping about staged data and files for agent access
     private File mBackupDataName;
     private File mStageName;
@@ -168,14 +166,14 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
     private final BackupAgentTimeoutParameters mAgentTimeoutParameters;
 
     @VisibleForTesting
-    PerformUnifiedRestoreTask(Map<String, Set<String>> excludedKeys) {
-        mExcludedKeys = excludedKeys;
+    PerformUnifiedRestoreTask(UserBackupManagerService backupManagerService) {
         mListener = null;
         mAgentTimeoutParameters = null;
         mTransportClient = null;
         mTransportManager = null;
         mEphemeralOpToken = 0;
         mUserId = 0;
+        this.backupManagerService = backupManagerService;
     }
 
     // This task can assume that the wakelock is properly held for it and doesn't have to worry
@@ -190,8 +188,7 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
             int pmToken,
             boolean isFullSystemRestore,
             @Nullable String[] filterSet,
-            OnTaskFinishedListener listener,
-            Map<String, Set<String>> excludedKeys) {
+            OnTaskFinishedListener listener) {
         this.backupManagerService = backupManagerService;
         mUserId = backupManagerService.getUserId();
         mTransportManager = backupManagerService.getTransportManager();
@@ -213,8 +210,6 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
                 backupManagerService.getAgentTimeoutParameters(),
                 "Timeout parameters cannot be null");
 
-        mExcludedKeys = excludedKeys;
-
         if (targetPackage != null) {
             // Single package restore
             mAcceptSet = new ArrayList<>();
@@ -791,8 +786,9 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
                 !getExcludedKeysForPackage(PLATFORM_PACKAGE_NAME).isEmpty();
     }
 
-    private Set<String> getExcludedKeysForPackage(String packageName) {
-        return mExcludedKeys.getOrDefault(packageName, Collections.emptySet());
+    @VisibleForTesting
+    Set<String> getExcludedKeysForPackage(String packageName) {
+        return backupManagerService.getExcludedRestoreKeys(packageName);
     }
 
     @VisibleForTesting
diff --git a/services/robotests/src/com/android/server/testing/shadows/ShadowPerformUnifiedRestoreTask.java b/services/robotests/src/com/android/server/testing/shadows/ShadowPerformUnifiedRestoreTask.java
index 223a98b6c8f6..8daef5fad032 100644
--- a/services/robotests/src/com/android/server/testing/shadows/ShadowPerformUnifiedRestoreTask.java
+++ b/services/robotests/src/com/android/server/testing/shadows/ShadowPerformUnifiedRestoreTask.java
@@ -67,8 +67,7 @@ public class ShadowPerformUnifiedRestoreTask {
             int pmToken,
             boolean isFullSystemRestore,
             @Nullable String[] filterSet,
-            OnTaskFinishedListener listener,
-            Map<String, Set<String>> excludedKeys) {
+            OnTaskFinishedListener listener) {
         mBackupManagerService = backupManagerService;
         mPackage = targetPackage;
         mIsFullSystemRestore = isFullSystemRestore;
diff --git a/services/tests/servicestests/src/com/android/server/backup/UserBackupPreferencesTest.java b/services/tests/servicestests/src/com/android/server/backup/UserBackupPreferencesTest.java
index d6efe35723db..5800acabe8a5 100644
--- a/services/tests/servicestests/src/com/android/server/backup/UserBackupPreferencesTest.java
+++ b/services/tests/servicestests/src/com/android/server/backup/UserBackupPreferencesTest.java
@@ -16,8 +16,6 @@
 
 package com.android.server.backup;
 
-import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
 
 import androidx.test.InstrumentationRegistry;
@@ -33,18 +31,15 @@ import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
 
 import java.util.Arrays;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 
 @Presubmit
 @RunWith(AndroidJUnit4.class)
 public class UserBackupPreferencesTest {
     private static final String EXCLUDED_PACKAGE_1 = "package1";
-    private static final String EXCLUDED_PACKAGE_2 = "package2";
     private static final List<String> EXCLUDED_KEYS_1 = Arrays.asList("key1", "key2");
-    private static final List<String> EXCLUDED_KEYS_2 = Arrays.asList("key1");
+    private static final List<String> EXCLUDED_KEYS_2 = Arrays.asList("key3");
 
     @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
 
@@ -60,27 +55,13 @@ public class UserBackupPreferencesTest {
     }
 
     @Test
-    public void testGetExcludedKeysForPackages_returnsExcludedKeys() {
+    public void testGetExcludedKeysForPackage_returnsExcludedKeys() {
         mExcludedRestoreKeysStorage.addExcludedKeys(EXCLUDED_PACKAGE_1, EXCLUDED_KEYS_1);
-        mExcludedRestoreKeysStorage.addExcludedKeys(EXCLUDED_PACKAGE_2, EXCLUDED_KEYS_2);
+        mExcludedRestoreKeysStorage.addExcludedKeys(EXCLUDED_PACKAGE_1, EXCLUDED_KEYS_2);
 
-        Map<String, Set<String>> excludedKeys =
-                mExcludedRestoreKeysStorage.getExcludedRestoreKeysForPackages(EXCLUDED_PACKAGE_1);
-        assertTrue(excludedKeys.containsKey(EXCLUDED_PACKAGE_1));
-        assertFalse(excludedKeys.containsKey(EXCLUDED_PACKAGE_2));
-        assertEquals(new HashSet<>(EXCLUDED_KEYS_1), excludedKeys.get(EXCLUDED_PACKAGE_1));
-    }
-
-    @Test
-    public void testGetExcludedKeysForPackages_withEmpty_list_returnsAllExcludedKeys() {
-        mExcludedRestoreKeysStorage.addExcludedKeys(EXCLUDED_PACKAGE_1, EXCLUDED_KEYS_1);
-        mExcludedRestoreKeysStorage.addExcludedKeys(EXCLUDED_PACKAGE_2, EXCLUDED_KEYS_2);
-
-        Map<String, Set<String>> excludedKeys =
-                mExcludedRestoreKeysStorage.getAllExcludedRestoreKeys();
-        assertTrue(excludedKeys.containsKey(EXCLUDED_PACKAGE_1));
-        assertTrue(excludedKeys.containsKey(EXCLUDED_PACKAGE_2));
-        assertEquals(new HashSet<>(EXCLUDED_KEYS_1), excludedKeys.get(EXCLUDED_PACKAGE_1));
-        assertEquals(new HashSet<>(EXCLUDED_KEYS_2), excludedKeys.get(EXCLUDED_PACKAGE_2));
+        Set<String> excludedKeys =
+                mExcludedRestoreKeysStorage.getExcludedRestoreKeysForPackage(EXCLUDED_PACKAGE_1);
+        assertTrue(excludedKeys.containsAll(EXCLUDED_KEYS_1));
+        assertTrue(excludedKeys.containsAll(EXCLUDED_KEYS_2));
     }
 }
diff --git a/services/tests/servicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java b/services/tests/servicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
index 3d220432cc8e..017c93975286 100644
--- a/services/tests/servicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
+++ b/services/tests/servicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
@@ -20,7 +20,9 @@ import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
 import androidx.test.InstrumentationRegistry;
@@ -30,6 +32,8 @@ import android.app.backup.BackupDataInput;
 import android.app.backup.BackupDataOutput;
 import android.platform.test.annotations.Presubmit;
 
+import com.android.server.backup.UserBackupManagerService;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -40,6 +44,7 @@ import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
 import java.util.ArrayDeque;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -59,6 +64,7 @@ public class PerformUnifiedRestoreTaskTest {
 
     @Mock private BackupDataInput mBackupDataInput;
     @Mock private BackupDataOutput mBackupDataOutput;
+    @Mock private UserBackupManagerService mBackupManagerService;
 
     private Set<String> mExcludedkeys = new HashSet<>();
     private Map<String, String> mBackupData = new HashMap<>();
@@ -99,6 +105,8 @@ public class PerformUnifiedRestoreTaskTest {
                         return null;
                     }
                 });
+
+        mRestoreTask = new PerformUnifiedRestoreTask(mBackupManagerService);
     }
 
     private void populateTestData() {
@@ -114,8 +122,9 @@ public class PerformUnifiedRestoreTaskTest {
 
     @Test
     public void testFilterExcludedKeys() throws Exception {
-        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
-                PACKAGE_NAME, mExcludedkeys));
+        when(mBackupManagerService.getExcludedRestoreKeys(eq(PACKAGE_NAME))).thenReturn(
+                mExcludedkeys);
+
         mRestoreTask.filterExcludedKeys(PACKAGE_NAME, mBackupDataInput, mBackupDataOutput);
 
         // Verify only the correct were written into BackupDataOutput object.
@@ -125,32 +134,49 @@ public class PerformUnifiedRestoreTaskTest {
     }
 
     @Test
+    public void testGetExcludedKeysForPackage_alwaysReturnsLatestKeys() {
+        Set<String> firstExcludedKeys = new HashSet<>(Collections.singletonList(EXCLUDED_KEY_1));
+        when(mBackupManagerService.getExcludedRestoreKeys(eq(PACKAGE_NAME))).thenReturn(
+                firstExcludedKeys);
+        assertEquals(firstExcludedKeys, mRestoreTask.getExcludedKeysForPackage(PACKAGE_NAME));
+
+
+        Set<String> secondExcludedKeys = new HashSet<>(Arrays.asList(EXCLUDED_KEY_1,
+                EXCLUDED_KEY_2));
+        when(mBackupManagerService.getExcludedRestoreKeys(eq(PACKAGE_NAME))).thenReturn(
+                secondExcludedKeys);
+        assertEquals(secondExcludedKeys, mRestoreTask.getExcludedKeysForPackage(PACKAGE_NAME));
+    }
+
+    @Test
     public void testStageBackupData_stageForNonSystemPackageWithKeysToExclude() {
-        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
-                PACKAGE_NAME, mExcludedkeys));
+        when(mBackupManagerService.getExcludedRestoreKeys(eq(NON_SYSTEM_PACKAGE_NAME))).thenReturn(
+                mExcludedkeys);
 
         assertTrue(mRestoreTask.shouldStageBackupData(NON_SYSTEM_PACKAGE_NAME));
     }
 
     @Test
     public void testStageBackupData_stageForNonSystemPackageWithNoKeysToExclude() {
-        mRestoreTask = new PerformUnifiedRestoreTask(Collections.emptyMap());
+        when(mBackupManagerService.getExcludedRestoreKeys(any())).thenReturn(
+                Collections.emptySet());
 
         assertTrue(mRestoreTask.shouldStageBackupData(NON_SYSTEM_PACKAGE_NAME));
     }
 
     @Test
     public void testStageBackupData_doNotStageForSystemPackageWithNoKeysToExclude() {
-        mRestoreTask = new PerformUnifiedRestoreTask(Collections.emptyMap());
+        when(mBackupManagerService.getExcludedRestoreKeys(any())).thenReturn(
+                Collections.emptySet());
 
         assertFalse(mRestoreTask.shouldStageBackupData(SYSTEM_PACKAGE_NAME));
     }
 
     @Test
     public void testStageBackupData_stageForSystemPackageWithKeysToExclude() {
-        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
-                PACKAGE_NAME, mExcludedkeys));
+        when(mBackupManagerService.getExcludedRestoreKeys(eq(SYSTEM_PACKAGE_NAME))).thenReturn(
+                mExcludedkeys);
 
-        assertTrue(mRestoreTask.shouldStageBackupData(NON_SYSTEM_PACKAGE_NAME));
+        assertTrue(mRestoreTask.shouldStageBackupData(SYSTEM_PACKAGE_NAME));
     }
 }