Enable rollback lifetime API changes during install

Test: CTS and manual test
Bug: 288288803
Change-Id: Ia244caab1284d3468cbe75d3c6741f88a165f043
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 1a7810a..d067dd1 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3872,6 +3872,7 @@
     method public void setInstallAsInstantApp(boolean);
     method public void setInstallAsVirtualPreload();
     method public void setRequestDowngrade(boolean);
+    method @FlaggedApi("android.content.pm.rollback_lifetime") @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS) public void setRollbackLifetimeMillis(long);
     method @RequiresPermission(android.Manifest.permission.INSTALL_PACKAGES) public void setStaged();
   }
 
diff --git a/core/java/android/content/pm/PackageInstaller.java b/core/java/android/content/pm/PackageInstaller.java
index cd8938d..cbb20e0 100644
--- a/core/java/android/content/pm/PackageInstaller.java
+++ b/core/java/android/content/pm/PackageInstaller.java
@@ -2535,6 +2535,8 @@
         public DataLoaderParams dataLoaderParams;
         /** {@hide} */
         public int rollbackDataPolicy = PackageManager.ROLLBACK_DATA_POLICY_RESTORE;
+        /** @hide */
+        public long rollbackLifetimeMillis = 0;
         /** {@hide} */
         public boolean forceQueryableOverride;
         /** {@hide} */
@@ -2589,6 +2591,7 @@
                 dataLoaderParams = new DataLoaderParams(dataLoaderParamsParcel);
             }
             rollbackDataPolicy = source.readInt();
+            rollbackLifetimeMillis = source.readLong();
             requireUserAction = source.readInt();
             packageSource = source.readInt();
             applicationEnabledSettingPersistent = source.readBoolean();
@@ -2621,6 +2624,7 @@
             ret.requiredInstalledVersionCode = requiredInstalledVersionCode;
             ret.dataLoaderParams = dataLoaderParams;
             ret.rollbackDataPolicy = rollbackDataPolicy;
+            ret.rollbackLifetimeMillis = rollbackLifetimeMillis;
             ret.requireUserAction = requireUserAction;
             ret.packageSource = packageSource;
             ret.applicationEnabledSettingPersistent = applicationEnabledSettingPersistent;
@@ -2902,12 +2906,7 @@
          */
         @SystemApi
         public void setEnableRollback(boolean enable) {
-            if (enable) {
-                installFlags |= PackageManager.INSTALL_ENABLE_ROLLBACK;
-            } else {
-                installFlags &= ~PackageManager.INSTALL_ENABLE_ROLLBACK;
-            }
-            rollbackDataPolicy = PackageManager.ROLLBACK_DATA_POLICY_RESTORE;
+            setEnableRollback(enable, PackageManager.ROLLBACK_DATA_POLICY_RESTORE);
         }
 
         /**
@@ -2931,10 +2930,36 @@
                 installFlags |= PackageManager.INSTALL_ENABLE_ROLLBACK;
             } else {
                 installFlags &= ~PackageManager.INSTALL_ENABLE_ROLLBACK;
+                rollbackLifetimeMillis = 0;
             }
             rollbackDataPolicy = dataPolicy;
         }
 
+        /**
+         * If rollback enabled for this session (via {@link #setEnableRollback}, set time
+         * after which rollback will no longer be possible
+         *
+         * <p>For multi-package installs, this value must be set on the parent session.
+         * Child session rollback lifetime will be ignored.
+         *
+         * @param lifetimeMillis time after which rollback expires
+         * @throws IllegalArgumentException if lifetimeMillis is negative or rollback is not
+         * enabled via setEnableRollback.
+         * @hide
+         */
+        @SystemApi
+        @RequiresPermission(android.Manifest.permission.MANAGE_ROLLBACKS)
+        @FlaggedApi(Flags.FLAG_ROLLBACK_LIFETIME)
+        public void setRollbackLifetimeMillis(@DurationMillisLong long lifetimeMillis) {
+            if (lifetimeMillis < 0) {
+                throw new IllegalArgumentException("rollbackLifetimeMillis can't be negative.");
+            }
+            if ((installFlags & PackageManager.INSTALL_ENABLE_ROLLBACK) == 0) {
+                throw new IllegalArgumentException(
+                        "Can't set rollbackLifetimeMillis when rollback is not enabled");
+            }
+            rollbackLifetimeMillis = lifetimeMillis;
+        }
 
         /**
          * @deprecated use {@link #setRequestDowngrade(boolean)}.
@@ -3295,6 +3320,7 @@
             pw.printPair("requiredInstalledVersionCode", requiredInstalledVersionCode);
             pw.printPair("dataLoaderParams", dataLoaderParams);
             pw.printPair("rollbackDataPolicy", rollbackDataPolicy);
+            pw.printPair("rollbackLifetimeMillis", rollbackLifetimeMillis);
             pw.printPair("applicationEnabledSettingPersistent",
                     applicationEnabledSettingPersistent);
             pw.printHexPair("developmentInstallFlags", developmentInstallFlags);
@@ -3336,6 +3362,7 @@
                 dest.writeParcelable(null, flags);
             }
             dest.writeInt(rollbackDataPolicy);
+            dest.writeLong(rollbackLifetimeMillis);
             dest.writeInt(requireUserAction);
             dest.writeInt(packageSource);
             dest.writeBoolean(applicationEnabledSettingPersistent);
@@ -3529,6 +3556,9 @@
         /** {@hide} */
         public int rollbackDataPolicy;
 
+        /** @hide */
+        public long rollbackLifetimeMillis;
+
         /** {@hide} */
         public int requireUserAction;
 
@@ -3596,6 +3626,7 @@
             isCommitted = source.readBoolean();
             isPreapprovalRequested = source.readBoolean();
             rollbackDataPolicy = source.readInt();
+            rollbackLifetimeMillis = source.readLong();
             createdMillis = source.readLong();
             requireUserAction = source.readInt();
             installerUid = source.readInt();
@@ -4220,6 +4251,7 @@
             dest.writeBoolean(isCommitted);
             dest.writeBoolean(isPreapprovalRequested);
             dest.writeInt(rollbackDataPolicy);
+            dest.writeLong(rollbackLifetimeMillis);
             dest.writeLong(createdMillis);
             dest.writeInt(requireUserAction);
             dest.writeInt(installerUid);
diff --git a/core/java/android/content/pm/flags.aconfig b/core/java/android/content/pm/flags.aconfig
index 96609ad..9ad66ce 100644
--- a/core/java/android/content/pm/flags.aconfig
+++ b/core/java/android/content/pm/flags.aconfig
@@ -58,3 +58,11 @@
     bug: "295827951"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "rollback_lifetime"
+    namespace: "package_manager_service"
+    description: "Feature flag to enable custom rollback lifetime during install."
+    bug: "299670324"
+    is_fixed_read_only: true
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/pm/PackageInstallerService.java b/services/core/java/com/android/server/pm/PackageInstallerService.java
index b9b5908..305e353 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerService.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerService.java
@@ -44,6 +44,7 @@
 import android.content.IntentSender.SendIntentException;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.ArchivedPackageParcel;
+import android.content.pm.Flags;
 import android.content.pm.IPackageInstaller;
 import android.content.pm.IPackageInstallerCallback;
 import android.content.pm.IPackageInstallerSession;
@@ -745,6 +746,22 @@
             params.installFlags &= ~PackageManager.INSTALL_DISABLE_VERIFICATION;
         }
 
+        if (Flags.rollbackLifetime()) {
+            if (params.rollbackLifetimeMillis > 0) {
+                if ((params.installFlags & PackageManager.INSTALL_ENABLE_ROLLBACK) == 0) {
+                    throw new IllegalArgumentException(
+                            "Can't set rollbackLifetimeMillis when rollback is not enabled");
+                }
+                if (mContext.checkCallingOrSelfPermission(Manifest.permission.MANAGE_ROLLBACKS)
+                        != PackageManager.PERMISSION_GRANTED) {
+                    throw new SecurityException(
+                            "Setting rollback lifetime requires the MANAGE_ROLLBACKS permission");
+                }
+            } else if (params.rollbackLifetimeMillis < 0) {
+                throw new IllegalArgumentException("rollbackLifetimeMillis can't be negative.");
+            }
+        }
+
         boolean isApex = (params.installFlags & PackageManager.INSTALL_APEX) != 0;
         if (isApex) {
             if (mContext.checkCallingOrSelfPermission(Manifest.permission.INSTALL_PACKAGE_UPDATES)
diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index 5dc7dab..1be28ca 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -1282,6 +1282,7 @@
             info.whitelistedRestrictedPermissions = params.whitelistedRestrictedPermissions;
             info.autoRevokePermissionsMode = params.autoRevokePermissionsMode;
             info.installFlags = params.installFlags;
+            info.rollbackLifetimeMillis = params.rollbackLifetimeMillis;
             info.isMultiPackage = params.isMultiPackage;
             info.isStaged = params.isStaged;
             info.rollbackDataPolicy = params.rollbackDataPolicy;
diff --git a/services/core/java/com/android/server/rollback/Rollback.java b/services/core/java/com/android/server/rollback/Rollback.java
index 85d93f4..a5b90f1 100644
--- a/services/core/java/com/android/server/rollback/Rollback.java
+++ b/services/core/java/com/android/server/rollback/Rollback.java
@@ -174,6 +174,11 @@
     @Nullable private final String mInstallerPackageName;
 
     /**
+     * Time after which rollback expires.
+     */
+    private long mRollbackLifetimeMillis = 0;
+
+    /**
      * Session ids for all packages in the install. For multi-package sessions, this is the list
      * of child session ids. For normal sessions, this list is a single element with the normal
      * session id.
@@ -286,6 +291,24 @@
     }
 
     /**
+     * Sets rollback lifetime in milliseconds, for purposes of expiring rollback data.
+     */
+    @WorkerThread
+    void setRollbackLifetimeMillis(long lifetimeMillis) {
+        assertInWorkerThread();
+        mRollbackLifetimeMillis = lifetimeMillis;
+    }
+
+    /**
+     * Returns rollback lifetime in milliseconds, for purposes of expiring rollback data.
+     */
+    @WorkerThread
+    long getRollbackLifetimeMillis() {
+        assertInWorkerThread();
+        return mRollbackLifetimeMillis;
+    }
+
+    /**
      * Returns the session ID associated with this rollback, or {@code -1} if unknown.
      */
     @AnyThread
@@ -930,6 +953,7 @@
         ipw.println("-state: " + getStateAsString());
         ipw.println("-stateDescription: " + mStateDescription);
         ipw.println("-timestamp: " + getTimestamp());
+        ipw.println("-rollbackLifetimeMillis: " + getRollbackLifetimeMillis());
         ipw.println("-isStaged: " + isStaged());
         ipw.println("-originalSessionId: " + getOriginalSessionId());
         ipw.println("-packages:");
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 720c773..8d93408 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -28,6 +28,7 @@
 import android.content.IntentFilter;
 import android.content.IntentSender;
 import android.content.pm.ApplicationInfo;
+import android.content.pm.Flags;
 import android.content.pm.ModuleInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageInstaller;
@@ -702,6 +703,15 @@
     // Schedules future expiration as appropriate.
     @WorkerThread
     private void runExpiration() {
+        if (Flags.rollbackLifetime()) {
+            runExpirationCustomRollbackLifetime();
+        } else {
+            runExpirationDefaultRollbackLifetime();
+        }
+    }
+
+    @WorkerThread
+    private void runExpirationDefaultRollbackLifetime() {
         getHandler().removeCallbacks(mRunExpiration);
         assertInWorkerThread();
         Instant now = Instant.now();
@@ -729,6 +739,44 @@
         }
     }
 
+    @WorkerThread
+    private void runExpirationCustomRollbackLifetime() {
+        getHandler().removeCallbacks(mRunExpiration);
+        assertInWorkerThread();
+        Instant now = Instant.now();
+        long minDelay = 0;
+        Iterator<Rollback> iter = mRollbacks.iterator();
+        while (iter.hasNext()) {
+            Rollback rollback = iter.next();
+            if (!rollback.isAvailable() && !rollback.isCommitted()) {
+                continue;
+            }
+            long rollbackLifetimeMillis = rollback.getRollbackLifetimeMillis();
+            if (rollbackLifetimeMillis <= 0) {
+                rollbackLifetimeMillis = mRollbackLifetimeDurationInMillis;
+            }
+
+            Instant rollbackExpiryTimestamp = rollback.getTimestamp()
+                    .plusMillis(rollbackLifetimeMillis);
+            if (!now.isBefore(rollbackExpiryTimestamp)) {
+                Slog.i(TAG, "runExpiration id=" + rollback.info.getRollbackId());
+                iter.remove();
+                deleteRollback(rollback, "Expired by timeout");
+                continue;
+            }
+
+            long delay = now.until(
+                    rollbackExpiryTimestamp, ChronoUnit.MILLIS);
+            if (minDelay == 0 || delay < minDelay) {
+                minDelay = delay;
+            }
+        }
+
+        if (minDelay != 0) {
+            getHandler().postDelayed(mRunExpiration, minDelay);
+        }
+    }
+
     @AnyThread
     private Handler getHandler() {
         return mHandler;
@@ -1277,6 +1325,7 @@
         }
 
         final Rollback rollback;
+
         if (parentSession.isStaged()) {
             rollback = mRollbackStore.createStagedRollback(rollbackId, parentSessionId, userId,
                     installerPackageName, packageSessionIds, getExtensionVersions());
@@ -1285,6 +1334,11 @@
                     installerPackageName, packageSessionIds, getExtensionVersions());
         }
 
+        if (Flags.rollbackLifetime()) {
+            rollback.setRollbackLifetimeMillis(parentSession.rollbackLifetimeMillis);
+        }
+
+
         mRollbacks.add(rollback);
         return rollback;
     }
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index 8068c6f..0af137f 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -19,6 +19,7 @@
 import static com.android.server.rollback.Rollback.rollbackStateFromString;
 
 import android.annotation.NonNull;
+import android.content.pm.Flags;
 import android.content.pm.PackageManager;
 import android.content.pm.VersionedPackage;
 import android.content.rollback.PackageRollbackInfo;
@@ -312,6 +313,9 @@
             JSONObject dataJson = new JSONObject();
             dataJson.put("info", rollbackInfoToJson(rollback.info));
             dataJson.put("timestamp", rollback.getTimestamp().toString());
+            if (Flags.rollbackLifetime()) {
+                dataJson.put("rollbackLifetimeMillis", rollback.getRollbackLifetimeMillis());
+            }
             dataJson.put("originalSessionId", rollback.getOriginalSessionId());
             dataJson.put("state", rollback.getStateAsString());
             dataJson.put("stateDescription", rollback.getStateDescription());
@@ -375,7 +379,7 @@
     @VisibleForTesting
     static Rollback rollbackFromJson(JSONObject dataJson, File backupDir)
             throws JSONException, ParseException {
-        return new Rollback(
+        Rollback rollback = new Rollback(
                 rollbackInfoFromJson(dataJson.getJSONObject("info")),
                 backupDir,
                 Instant.parse(dataJson.getString("timestamp")),
@@ -388,6 +392,10 @@
                 dataJson.optInt("userId", UserHandle.SYSTEM.getIdentifier()),
                 dataJson.optString("installerPackageName", ""),
                 extensionVersionsFromJson(dataJson.optJSONArray("extensionVersions")));
+        if (Flags.rollbackLifetime()) {
+            rollback.setRollbackLifetimeMillis(dataJson.optLong("rollbackLifetimeMillis"));
+        }
+        return rollback;
     }
 
     private static JSONObject toJson(VersionedPackage pkg) throws JSONException {