Stage 'android' package in PerformUnifiedRestoreTask

Bug: 145126096
Test: atest KeyValueRestoreExclusionHostSideTest
      atest PerformUnifiedRestoreTaskTest

In a KV restore after getting data from the transport, we save it into a
stage file. Then we go through the keys and do filtering: skip the keys
that should be excluded and extract the widget data into a separate
file. The rest of the data is wirtten into the file where the app's
backup agent will read it from.

However, this process is skipped for 'android' package. It was done as
an optimization before the ability to exclude keys from restore was
introduced: as 'android' backup data doesn't contain any widget info.
However, now we need to process 'android' package as well because it can
contain keys to be excluded.

Change-Id: I612f8cc9c6903c9bd257762360dadb81ed12d106

services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java

@@ -76,6 +76,7 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -697,11 +698,7 @@ private void initiateOneRestore(PackageInfo app, long appVersionCode) {
         mStageName = new File(backupManagerService.getDataDir(), packageName + ".stage");
         mNewStateName = new File(mStateDir, packageName + ".new");
 
-        // don't stage the 'android' package where the wallpaper data lives.  this is
-        // an optimization: we know there's no widget data hosted/published by that
-        // package, and this way we avoid doing a spurious copy of MB-sized wallpaper
-        // data following the download.
-        boolean staging = !packageName.equals(PLATFORM_PACKAGE_NAME);
+        boolean staging = shouldStageBackupData(packageName);
         ParcelFileDescriptor stage;
         File downloadFile = (staging) ? mStageName : mBackupDataName;
         boolean startedAgentRestore = false;
@@ -768,8 +765,7 @@ private void initiateOneRestore(PackageInfo app, long appVersionCode) {
             startedAgentRestore = true;
             mAgent.doRestoreWithExcludedKeys(mBackupData, appVersionCode, mNewState,
                     mEphemeralOpToken, backupManagerService.getBackupManagerBinder(),
-                    mExcludedKeys.containsKey(packageName)
-                            ? new ArrayList<>(mExcludedKeys.get(packageName)) : null);
+                    new ArrayList<>(getExcludedKeysForPackage(packageName)));
         } catch (Exception e) {
             Slog.e(TAG, "Unable to call app for restore: " + packageName, e);
             EventLog.writeEvent(EventLogTags.RESTORE_AGENT_FAILURE,
@@ -784,10 +780,25 @@ private void initiateOneRestore(PackageInfo app, long appVersionCode) {
         }
     }
 
+    @VisibleForTesting
+    boolean shouldStageBackupData(String packageName) {
+        // Backup data is staged for 2 reasons:
+        // 1. We might need to exclude keys from the data before passing it to the agent
+        // 2. Widget metadata needs to be separated from the rest to be handled separately
+        // But 'android' package doesn't contain widget metadata so we want to skip staging for it
+        // when there are no keys to be excluded either.
+        return !packageName.equals(PLATFORM_PACKAGE_NAME) ||
+                !getExcludedKeysForPackage(PLATFORM_PACKAGE_NAME).isEmpty();
+    }
+
+    private Set<String> getExcludedKeysForPackage(String packageName) {
+        return mExcludedKeys.getOrDefault(packageName, Collections.emptySet());
+    }
+
     @VisibleForTesting
     void filterExcludedKeys(String packageName, BackupDataInput in, BackupDataOutput out)
             throws Exception {
-        Set<String> excludedKeysForPackage = mExcludedKeys.get(packageName);
+        Set<String> excludedKeysForPackage = getExcludedKeysForPackage(packageName);
 
         byte[] buffer = new byte[8192]; // will grow when needed
         while (in.readNextHeader()) {

services/tests/servicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java

@@ -20,7 +20,6 @@
 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.Mockito.when;
 
@@ -55,6 +54,8 @@ public class PerformUnifiedRestoreTaskTest {
     private static final String INCLUDED_KEY = "included_key";
     private static final String EXCLUDED_KEY_1 = "excluded_key_1";
     private static final String EXCLUDED_KEY_2 = "excluded_key_2";
+    private static final String SYSTEM_PACKAGE_NAME = "android";
+    private static final String NON_SYSTEM_PACKAGE_NAME = "package";
 
     @Mock private BackupDataInput mBackupDataInput;
     @Mock private BackupDataOutput mBackupDataOutput;
@@ -98,9 +99,6 @@ public Void answer(InvocationOnMock invocation) throws Throwable {
                         return null;
                     }
                 });
-
-        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
-                PACKAGE_NAME, mExcludedkeys));
     }
 
     private void populateTestData() {
@@ -116,11 +114,43 @@ private void populateTestData() {
 
     @Test
     public void testFilterExcludedKeys() throws Exception {
+        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
+                PACKAGE_NAME, mExcludedkeys));
         mRestoreTask.filterExcludedKeys(PACKAGE_NAME, mBackupDataInput, mBackupDataOutput);
 
         // Verify only the correct were written into BackupDataOutput object.
         Set<String> allowedBackupKeys = new HashSet<>(mBackupData.keySet());
         allowedBackupKeys.removeAll(mExcludedkeys);
         assertEquals(allowedBackupKeys, mBackupDataDump);
     }
+
+    @Test
+    public void testStageBackupData_stageForNonSystemPackageWithKeysToExclude() {
+        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
+                PACKAGE_NAME, mExcludedkeys));
+
+        assertTrue(mRestoreTask.shouldStageBackupData(NON_SYSTEM_PACKAGE_NAME));
+    }
+
+    @Test
+    public void testStageBackupData_stageForNonSystemPackageWithNoKeysToExclude() {
+        mRestoreTask = new PerformUnifiedRestoreTask(Collections.emptyMap());
+
+        assertTrue(mRestoreTask.shouldStageBackupData(NON_SYSTEM_PACKAGE_NAME));
+    }
+
+    @Test
+    public void testStageBackupData_doNotStageForSystemPackageWithNoKeysToExclude() {
+        mRestoreTask = new PerformUnifiedRestoreTask(Collections.emptyMap());
+
+        assertFalse(mRestoreTask.shouldStageBackupData(SYSTEM_PACKAGE_NAME));
+    }
+
+    @Test
+    public void testStageBackupData_stageForSystemPackageWithKeysToExclude() {
+        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
+                PACKAGE_NAME, mExcludedkeys));
+
+        assertTrue(mRestoreTask.shouldStageBackupData(NON_SYSTEM_PACKAGE_NAME));
+    }
 }