Clean up staged session data on validation failure

Bug: 173132101
Test: atest StagedInstallInternalTest#testStagedInstallationShouldCleanUpOnValidationFailure
Test: atest StagedInstallInternalTest#testStagedInstallationShouldCleanUpOnValidationFailureMultiPackage
Change-Id: Idd6597cd0d2dda34a8a7626b585401eeee39c31f
Merged-In: Idd6597cd0d2dda34a8a7626b585401eeee39c31f
(cherry picked from commit b6dde00d03)

services/core/java/com/android/server/pm/PackageInstallerSession.java

@@ -1597,6 +1597,10 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
         destroyInternal();
         // Dispatch message to remove session from PackageInstallerService.
         dispatchSessionFinished(error, detailMessage, null);
+        // TODO(b/173194203): clean up staged session in destroyInternal() call instead
+        if (isStaged() && stageDir != null) {
+            cleanStageDir();
+        }
     }
 
     private void onStorageUnhealthy() {

tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java

@@ -96,6 +96,19 @@ public class StagedInstallInternalTest {
         assertSessionReady(sessionId);
     }
 
+    @Test
+    public void testStagedInstallationShouldCleanUpOnValidationFailure() throws Exception {
+        InstallUtils.commitExpectingFailure(AssertionError.class, "INSTALL_FAILED_INVALID_APK",
+                Install.single(TestApp.AIncompleteSplit).setStaged());
+    }
+
+    @Test
+    public void testStagedInstallationShouldCleanUpOnValidationFailureMultiPackage()
+            throws Exception {
+        InstallUtils.commitExpectingFailure(AssertionError.class, "INSTALL_FAILED_INVALID_APK",
+                Install.multi(TestApp.AIncompleteSplit, TestApp.B1, TestApp.Apex1).setStaged());
+    }
+
     private static void assertSessionReady(int sessionId) {
         assertSessionState(sessionId,
                 (session) -> assertThat(session.isStagedSessionReady()).isTrue());

tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java

@@ -28,6 +28,7 @@ import android.cts.install.lib.host.InstallUtilsHost;
 import com.android.ddmlib.Log;
 import com.android.tests.rollback.host.AbandonSessionsRule;
 import com.android.tests.util.ModuleTestUtils;
+import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.ProcessInfo;
@@ -39,6 +40,9 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import java.io.File;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
 
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class StagedInstallInternalTest extends BaseHostJUnit4Test {
@@ -182,6 +186,39 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
         assertThat(sessionIds.length).isEqualTo(3);
     }
 
+    @Test
+    public void testStagedInstallationShouldCleanUpOnValidationFailure() throws Exception {
+        List<String> before = getStagingDirectories();
+        runPhase("testStagedInstallationShouldCleanUpOnValidationFailure");
+        List<String> after = getStagingDirectories();
+        assertThat(after).isEqualTo(before);
+    }
+
+    @Test
+    public void testStagedInstallationShouldCleanUpOnValidationFailureMultiPackage()
+            throws Exception {
+        List<String> before = getStagingDirectories();
+        runPhase("testStagedInstallationShouldCleanUpOnValidationFailureMultiPackage");
+        List<String> after = getStagingDirectories();
+        assertThat(after).isEqualTo(before);
+    }
+
+    private List<String> getStagingDirectories() throws DeviceNotAvailableException {
+        String baseDir = "/data/app-staging";
+        try {
+            getDevice().enableAdbRoot();
+            return getDevice().getFileEntry(baseDir).getChildren(false)
+                    .stream().filter(entry -> entry.getName().matches("session_\\d+"))
+                    .map(entry -> entry.getName())
+                    .collect(Collectors.toList());
+        } catch (Exception e) {
+            // Return an empty list if any error
+            return Collections.EMPTY_LIST;
+        } finally {
+            getDevice().disableAdbRoot();
+        }
+    }
+
     private void restartSystemServer() throws Exception {
         // Restart the system server
         long oldStartTime = getDevice().getProcessByName("system_server").getStartTime();