From fdf7b58a43d7a518c76761ab2d2476e6355bd4d4 Mon Sep 17 00:00:00 2001
From: Tianjie <xunchang@google.com>
Date: Wed, 5 May 2021 17:47:02 -0700
Subject: [PATCH] Add a system API UpdateEngine to switch the slot back and
 forth

In some cases, the OTA update want to flip the slot to source
without cancelling the entire update. The functionality was
achieved by ResetStatus(). However ResetStatus() will clean up
the update progress and free up space to support virtual A/B.

Details in http://go/no-virtual-ab-switch-off

Test: build
Bug: 187321613
Change-Id: I2abcf8289f5fd573d2b55fa7004ab4dcd175252f
---
 core/api/system-current.txt            |  2 ++
 core/java/android/os/UpdateEngine.java | 42 +++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 675347085f095..dca8d2fd879eb 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -7269,8 +7269,10 @@ package android.os {
     method public boolean bind(android.os.UpdateEngineCallback);
     method public void cancel();
     method @WorkerThread public int cleanupAppliedPayload();
+    method public void resetShouldSwitchSlotOnReboot();
     method public void resetStatus();
     method public void resume();
+    method public void setShouldSwitchSlotOnReboot(@NonNull String);
     method public void suspend();
     method public boolean unbind();
     method public boolean verifyPayloadMetadata(String);
diff --git a/core/java/android/os/UpdateEngine.java b/core/java/android/os/UpdateEngine.java
index 5a48242f26eae..3e01c53f04693 100644
--- a/core/java/android/os/UpdateEngine.java
+++ b/core/java/android/os/UpdateEngine.java
@@ -409,10 +409,11 @@ public class UpdateEngine {
 
     /**
      * Resets the bootable flag on the non-current partition and all internal
-     * update_engine state. This can be used after an unwanted payload has been
-     * successfully applied and the device has not yet been rebooted to signal
-     * that we no longer want to boot into that updated system. After this call
-     * completes, update_engine will no longer report
+     * update_engine state. Note this call will clear the entire update
+     * progress. So a subsequent {@link #applyPayload} will apply the update
+     * from scratch.
+     *
+     * <p>After this call completes, update_engine will no longer report
      * {@code UPDATED_NEED_REBOOT}, so your callback can remove any outstanding
      * notification that rebooting into the new system is possible.
      */
@@ -424,6 +425,39 @@ public class UpdateEngine {
         }
     }
 
+    /**
+     * Sets the A/B slot switch for the next boot after applying an ota update. If
+     * {@link #applyPayload} hasn't switched the slot, the updater APP can call
+     * this API to switch the slot and apply the update on next boot.
+     *
+     * @param payloadMetadataFilename the location of the metadata without the
+     * {@code file://} prefix.
+     */
+    public void setShouldSwitchSlotOnReboot(@NonNull String payloadMetadataFilename) {
+        try {
+            mUpdateEngine.setShouldSwitchSlotOnReboot(payloadMetadataFilename);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+   /**
+    * Resets the boot slot to the source/current slot, without cancelling the
+    * update progress. This can be called after the update is installed, and to
+    * prevent the device from accidentally taking the update when it reboots.
+    *
+    * This is useful when users don't want to take the update immediately; or
+    * the updater determines some condition hasn't met, e.g. insufficient space
+    * for boot.
+    */
+    public void resetShouldSwitchSlotOnReboot() {
+        try {
+            mUpdateEngine.resetShouldSwitchSlotOnReboot();
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
     /**
      * Unbinds the last bound callback function.
      */