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

2 files changed, 53 insertions, 12 deletions

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 16484df9f328..f90d936a5f4d 100644
--- a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
+++ b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
@@ -76,6 +76,7 @@ import java.io.FileInputStream;
 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 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
         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 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
             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,
@@ -785,9 +781,24 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
     }
 
     @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()) {
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 6359edf190b0..3d220432cc8e 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,6 @@ 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.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 class PerformUnifiedRestoreTaskTest {
                         return null;
                     }
                 });
-
-        mRestoreTask = new PerformUnifiedRestoreTask(Collections.singletonMap(
-                PACKAGE_NAME, mExcludedkeys));
     }
 
     private void populateTestData() {
@@ -116,6 +114,8 @@ public class PerformUnifiedRestoreTaskTest {
 
     @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.
@@ -123,4 +123,34 @@ public class PerformUnifiedRestoreTaskTest {
         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));
+    }
 }