Detect rollback and cleanup space allocated for apex

As part of this change, a new enum OTAResult is added to represent state
of OTA.

Test: th
Test: Apply OTA, force rollback, make sure space is cleaned up
Test: Apply OTA, kill update_engine and restart, make sure update_engine
is still in UPDATE_NEED_REBOOT state
Test: Apply OTA, reboot, make sure sure "OTA succeeded
 is printed

Change-Id: I8e1bf20e658c22111c447bf70b2d25d6d10d6083
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index bcae3aa..28c193e 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -19,6 +19,7 @@
 #include <algorithm>
 #include <map>
 #include <memory>
+#include <ostream>
 #include <utility>
 
 #include <android-base/properties.h>
@@ -163,15 +164,37 @@
   return old_boot_id != boot_id;
 }
 
+std::ostream& operator<<(std::ostream& out, OTAResult result) {
+  switch (result) {
+    case OTAResult::NOT_ATTEMPTED:
+      out << "OTAResult::NOT_ATTEMPTED";
+      break;
+    case OTAResult::ROLLED_BACK:
+      out << "OTAResult::ROLLED_BACK";
+      break;
+    case OTAResult::UPDATED_NEED_REBOOT:
+      out << "OTAResult::UPDATED_NEED_REBOOT";
+      break;
+    case OTAResult::OTA_SUCCESSFUL:
+      out << "OTAResult::OTA_SUCCESSFUL";
+      break;
+  }
+  return out;
+}
+
 void UpdateAttempterAndroid::Init() {
   // In case of update_engine restart without a reboot we need to restore the
   // reboot needed state.
   if (UpdateCompletedOnThisBoot()) {
+    LOG(INFO) << "Updated installed but update_engine is restarted without "
+                 "device reboot. Resuming old state.";
     SetStatusAndNotify(UpdateStatus::UPDATED_NEED_REBOOT);
   } else {
+    const auto result = GetOTAUpdateResult();
+    LOG(INFO) << result;
     SetStatusAndNotify(UpdateStatus::IDLE);
     if (DidSystemReboot(prefs_)) {
-      UpdateStateAfterReboot();
+      UpdateStateAfterReboot(result);
     }
 
 #ifdef _UE_SIDELOAD
@@ -383,6 +406,15 @@
     std::vector<ApexInfo> apex_infos_blank;
     apex_handler_android_->AllocateSpace(apex_infos_blank);
   }
+  // Remove the reboot marker so that if the machine is rebooted
+  // after resetting to idle state, it doesn't go back to
+  // UpdateStatus::UPDATED_NEED_REBOOT state.
+  if (!ClearUpdateCompletedMarker()) {
+    return LogAndSetError(error,
+                          FROM_HERE,
+                          "Failed to reset the status because "
+                          "ClearUpdateCompletedMarker() failed");
+  }
 
   switch (status_) {
     case UpdateStatus::IDLE: {
@@ -416,11 +448,6 @@
       if (!boot_control_->GetDynamicPartitionControl()->ResetUpdate(prefs_))
         ret_value = false;
 
-      // Remove the reboot marker so that if the machine is rebooted
-      // after resetting to idle state, it doesn't go back to
-      // UpdateStatus::UPDATED_NEED_REBOOT state.
-      if (!prefs_->Delete(kPrefsUpdateCompletedOnBootId))
-        ret_value = false;
       ClearMetricsPrefs();
 
       if (!ret_value) {
@@ -567,7 +594,9 @@
   switch (code) {
     case ErrorCode::kSuccess:
       // Update succeeded.
-      WriteUpdateCompletedMarker();
+      if (!WriteUpdateCompletedMarker()) {
+        LOG(ERROR) << "Failed to write update completion marker";
+      }
       prefs_->SetInt64(kPrefsDeltaUpdateFailures, 0);
 
       LOG(INFO) << "Update successfully applied, waiting to reboot.";
@@ -705,6 +734,7 @@
   }
 
   if (status_ == UpdateStatus::CLEANUP_PREVIOUS_UPDATE) {
+    ClearUpdateCompletedMarker();
     LOG(INFO) << "Terminating cleanup previous update.";
     SetStatusAndNotify(UpdateStatus::IDLE);
     for (auto observer : daemon_state_->service_observers())
@@ -798,9 +828,20 @@
 }
 
 bool UpdateAttempterAndroid::WriteUpdateCompletedMarker() {
+  LOG(INFO) << "Writing update complete marker.";
   string boot_id;
   TEST_AND_RETURN_FALSE(utils::GetBootId(&boot_id));
-  prefs_->SetString(kPrefsUpdateCompletedOnBootId, boot_id);
+  TEST_AND_RETURN_FALSE(
+      prefs_->SetString(kPrefsUpdateCompletedOnBootId, boot_id));
+  TEST_AND_RETURN_FALSE(
+      prefs_->SetInt64(kPrefsPreviousSlot, boot_control_->GetCurrentSlot()));
+  return true;
+}
+
+bool UpdateAttempterAndroid::ClearUpdateCompletedMarker() {
+  LOG(INFO) << "Clearing update complete marker.";
+  TEST_AND_RETURN_FALSE(prefs_->Delete(kPrefsUpdateCompletedOnBootId));
+  TEST_AND_RETURN_FALSE(prefs_->Delete(kPrefsPreviousSlot));
   return true;
 }
 
@@ -897,15 +938,66 @@
   }
 }
 
-void UpdateAttempterAndroid::UpdateStateAfterReboot() {
+bool UpdateAttempterAndroid::OTARebootSucceeded() const {
+  const auto current_slot = boot_control_->GetCurrentSlot();
+  const string current_version =
+      android::base::GetProperty("ro.build.version.incremental", "");
+  int64_t previous_slot = -1;
+  TEST_AND_RETURN_FALSE(prefs_->GetInt64(kPrefsPreviousSlot, &previous_slot));
+  string previous_version;
+  TEST_AND_RETURN_FALSE(
+      prefs_->GetString(kPrefsPreviousVersion, &previous_version));
+  if (previous_slot != current_slot) {
+    LOG(INFO) << "Detected a slot switch, OTA succeeded, device updated from "
+              << previous_version << " to " << current_version;
+    if (previous_version == current_version) {
+      LOG(INFO) << "Previous version is the same as current version, this is "
+                   "possibly a self-OTA.";
+    }
+    return true;
+  } else {
+    LOG(INFO) << "Slot didn't switch, either the OTA is rolled back, or slot "
+                 "switch never happened, or system not rebooted at all.";
+    if (previous_version != current_version) {
+      LOG(INFO) << "Slot didn't change, but version changed from "
+                << previous_version << " to " << current_version
+                << " device could be flashed.";
+    }
+    return false;
+  }
+}
+
+OTAResult UpdateAttempterAndroid::GetOTAUpdateResult() const {
+  // We only set |kPrefsSystemUpdatedMarker| if slot is actually switched, so
+  // existence of this pref is sufficient indicator. Given that we have to
+  // delete this pref after checking it. This is done in
+  // |DeltaPerformer::ResetUpdateProgress|
+  auto slot_switch_attempted = prefs_->Exists(kPrefsUpdateCompletedOnBootId);
+  auto system_rebooted = DidSystemReboot(prefs_);
+  auto ota_successful = OTARebootSucceeded();
+  if (ota_successful) {
+    return OTAResult::OTA_SUCCESSFUL;
+  }
+  if (slot_switch_attempted) {
+    if (system_rebooted) {
+      // If we attempted slot switch, but still end up on the same slot, we
+      // probably rolled back.
+      return OTAResult::ROLLED_BACK;
+    } else {
+      return OTAResult::UPDATED_NEED_REBOOT;
+    }
+  }
+  return OTAResult::NOT_ATTEMPTED;
+}
+
+void UpdateAttempterAndroid::UpdateStateAfterReboot(const OTAResult result) {
   // Example: [ro.build.version.incremental]: [4292972]
   string current_version =
       android::base::GetProperty("ro.build.version.incremental", "");
   TEST_AND_RETURN(!current_version.empty());
-  const auto current_slot = boot_control_->GetCurrentSlot();
 
-  // |InitAfterReboot()| is only called after system reboot, so record boot id
-  // unconditionally
+  // |UpdateStateAfterReboot()| is only called after system reboot, so record
+  // boot id unconditionally
   string current_boot_id;
   TEST_AND_RETURN(utils::GetBootId(&current_boot_id));
   prefs_->SetString(kPrefsBootId, current_boot_id);
@@ -918,13 +1010,8 @@
     ClearMetricsPrefs();
     return;
   }
-  int64_t previous_slot = -1;
-  prefs_->GetInt64(kPrefsPreviousSlot, &previous_slot);
-  string previous_version;
   // update_engine restarted under the same build and same slot.
-  // TODO(xunchang) identify and report rollback by checking UpdateMarker.
-  if (prefs_->GetString(kPrefsPreviousVersion, &previous_version) &&
-      previous_version == current_version && previous_slot == current_slot) {
+  if (result != OTAResult::OTA_SUCCESSFUL) {
     // Increment the reboot number if |kPrefsNumReboots| exists. That pref is
     // set when we start a new update.
     if (prefs_->Exists(kPrefsNumReboots)) {
@@ -932,6 +1019,17 @@
           metrics_utils::GetPersistedValue(kPrefsNumReboots, prefs_);
       metrics_utils::SetNumReboots(reboot_count + 1, prefs_);
     }
+
+    if (result == OTAResult::ROLLED_BACK) {
+      // This will release all space previously allocated for apex
+      // decompression. If we detect a rollback, we should release space and
+      // return the space to user. Any subsequent attempt to install OTA will
+      // allocate space again anyway.
+      LOG(INFO) << "Detected a rollback, releasing space allocated for apex "
+                   "deompression.";
+      apex_handler_android_->AllocateSpace({});
+      DeltaPerformer::ResetUpdateProgress(prefs_, false);
+    }
     return;
   }
 
@@ -969,6 +1067,7 @@
   }
   metrics_utils::SetUpdateTimestampStart(clock_->GetMonotonicTime(), prefs_);
   metrics_utils::SetUpdateBootTimestampStart(clock_->GetBootTime(), prefs_);
+  ClearUpdateCompletedMarker();
 }
 
 void UpdateAttempterAndroid::ClearMetricsPrefs() {
diff --git a/aosp/update_attempter_android.h b/aosp/update_attempter_android.h
index 7a5a635..3633178 100644
--- a/aosp/update_attempter_android.h
+++ b/aosp/update_attempter_android.h
@@ -45,6 +45,13 @@
 
 namespace chromeos_update_engine {
 
+enum class OTAResult {
+  NOT_ATTEMPTED,
+  ROLLED_BACK,
+  UPDATED_NEED_REBOOT,
+  OTA_SUCCESSFUL,
+};
+
 class UpdateAttempterAndroid
     : public ServiceDelegateAndroidInterface,
       public ActionProcessorDelegate,
@@ -114,9 +121,24 @@
   // CleanupPreviousUpdateActionDelegateInterface
   void OnCleanupProgressUpdate(double progress) override;
 
+  // Check the result of an OTA update. Intended to be called after reboot, this
+  // will use prefs on disk to determine if OTA was installed, or rolledback.
+  [[nodiscard]] OTAResult GetOTAUpdateResult() const;
+  // Intended to be called:
+  // 1. When system rebooted and slot switch is attempted
+  // 2. When a new update is started
+  // 3. When user called |ResetStatus()|
+  bool ClearUpdateCompletedMarker();
+
  private:
   friend class UpdateAttempterAndroidTest;
 
+  // Return |true| only if slot switched successfully after an OTA reboot.
+  // This will return |false| if an downgrade OTA is applied. Because after a
+  // downgrade OTA, we wipe /data, and there's no way for update_engine to
+  // "remember" that a downgrade OTA took place.
+  [[nodiscard]] bool OTARebootSucceeded() const;
+
   // Schedules an event loop callback to start the action processor. This is
   // scheduled asynchronously to unblock the event loop.
   void ScheduleProcessingStart();
@@ -136,10 +158,10 @@
 
   // Writes to the processing completed marker. Does nothing if
   // |update_completed_marker_| is empty.
-  bool WriteUpdateCompletedMarker();
+  [[nodiscard]] bool WriteUpdateCompletedMarker();
 
   // Returns whether an update was completed in the current boot.
-  bool UpdateCompletedOnThisBoot();
+  [[nodiscard]] bool UpdateCompletedOnThisBoot();
 
   // Prefs to use for metrics report
   // |kPrefsPayloadAttemptNumber|: number of update attempts for the current
@@ -171,7 +193,7 @@
   //   |ReportTimeToRebootMetrics|
   // Prefs to update:
   //   |kPrefsBootId|, |kPrefsPreviousVersion|
-  void UpdateStateAfterReboot();
+  void UpdateStateAfterReboot(OTAResult result);
 
   // Prefs to update:
   //   |kPrefsPayloadAttemptNumber|, |kPrefsUpdateTimestampStart|,
diff --git a/aosp/update_attempter_android_unittest.cc b/aosp/update_attempter_android_unittest.cc
index f73df16..969f191 100644
--- a/aosp/update_attempter_android_unittest.cc
+++ b/aosp/update_attempter_android_unittest.cc
@@ -106,6 +106,7 @@
   prefs_.SetString(kPrefsPreviousVersion, "00001");  // Set the fake version
   prefs_.SetInt64(kPrefsPayloadAttemptNumber, 1);
   prefs_.SetInt64(kPrefsSystemUpdatedMarker, 23456);
+  prefs_.SetInt64(kPrefsPreviousSlot, 1);
 
   EXPECT_CALL(*metrics_reporter_,
               ReportAbnormallyTerminatedUpdateAttemptMetrics())