4 files changed, 82 insertions, 10 deletions
diff --git a/otautil/include/otautil/rangeset.h b/otautil/include/otautil/rangeset.h
index e91d02ca6..a18c30e29 100644
--- a/otautil/include/otautil/rangeset.h
+++ b/otautil/include/otautil/rangeset.h
@@ -18,6 +18,7 @@
 
 #include <stddef.h>
 
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -49,6 +50,12 @@ class RangeSet {
   // bounds. For example, "3,5" contains blocks 3 and 4. So "3,5" and "5,7" are not overlapped.
   bool Overlaps(const RangeSet& other) const;
 
+  // Returns a subset of ranges starting from |start_index| with respect to the original range. The
+  // output range will have |num_of_blocks| blocks in size. Returns std::nullopt if the input is
+  // invalid. e.g. RangeSet({{0, 5}, {10, 15}}).GetSubRanges(1, 5) returns
+  // RangeSet({{1, 5}, {10, 11}}).
+  std::optional<RangeSet> GetSubRanges(size_t start_index, size_t num_of_blocks) const;
+
   // Returns a vector of RangeSets that contain the same set of blocks represented by the current
   // RangeSet. The RangeSets in the vector contain similar number of blocks, with a maximum delta
   // of 1-block between any two of them. For example, 14 blocks would be split into 4 + 4 + 3 + 3,
diff --git a/otautil/include/otautil/sysutil.h b/otautil/include/otautil/sysutil.h
index 692a99e9d..48e9011e5 100644
--- a/otautil/include/otautil/sysutil.h
+++ b/otautil/include/otautil/sysutil.h
@@ -14,12 +14,12 @@
  * limitations under the License.
  */
 
-#ifndef _OTAUTIL_SYSUTIL
-#define _OTAUTIL_SYSUTIL
+#pragma once
 
 #include <sys/types.h>
 
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include "rangeset.h"
@@ -101,13 +101,14 @@ class MemMapping {
   std::vector<MappedRange> ranges_;
 };
 
-// Wrapper function to trigger a reboot, by additionally handling quiescent reboot mode. The
-// command should start with "reboot," (e.g. "reboot,bootloader" or "reboot,").
-bool reboot(const std::string& command);
+// Reboots the device into the specified target, by additionally handling quiescent reboot mode.
+// 'target' can be an empty string, which indicates booting into Android.
+bool Reboot(std::string_view target);
+
+// Triggers a shutdown.
+bool Shutdown();
 
 // Returns a null-terminated char* array, where the elements point to the C-strings in the given
 // vector, plus an additional nullptr at the end. This is a helper function that facilitates
 // calling C functions (such as getopt(3)) that expect an array of C-strings.
 std::vector<char*> StringVectorToNullTerminatedArray(const std::vector<std::string>& args);
-
-#endif  // _OTAUTIL_SYSUTIL
diff --git a/otautil/rangeset.cpp b/otautil/rangeset.cpp
index 5ab8e08fe..8ee99dd7a 100644
--- a/otautil/rangeset.cpp
+++ b/otautil/rangeset.cpp
@@ -184,6 +184,58 @@ bool RangeSet::Overlaps(const RangeSet& other) const {
   return false;
 }
 
+std::optional<RangeSet> RangeSet::GetSubRanges(size_t start_index, size_t num_of_blocks) const {
+  size_t end_index = start_index + num_of_blocks;  // The index of final block to read plus one
+  if (start_index > end_index || end_index > blocks_) {
+    LOG(ERROR) << "Failed to get the sub ranges for start_index " << start_index
+               << " num_of_blocks " << num_of_blocks
+               << " total number of blocks the range contains is " << blocks_;
+    return std::nullopt;
+  }
+
+  if (num_of_blocks == 0) {
+    LOG(WARNING) << "num_of_blocks is zero when calling GetSubRanges()";
+    return RangeSet();
+  }
+
+  RangeSet result;
+  size_t current_index = 0;
+  for (const auto& [range_start, range_end] : ranges_) {
+    CHECK_LT(range_start, range_end);
+    size_t blocks_in_range = range_end - range_start;
+    // Linear search to skip the ranges until we reach start_block.
+    if (current_index + blocks_in_range <= start_index) {
+      current_index += blocks_in_range;
+      continue;
+    }
+
+    size_t trimmed_range_start = range_start;
+    // We have found the first block range to read, trim the heading blocks.
+    if (current_index < start_index) {
+      trimmed_range_start += start_index - current_index;
+    }
+    // Trim the trailing blocks if the last range has more blocks than desired; also return the
+    // result.
+    if (current_index + blocks_in_range >= end_index) {
+      size_t trimmed_range_end = range_end - (current_index + blocks_in_range - end_index);
+      if (!result.PushBack({ trimmed_range_start, trimmed_range_end })) {
+        return std::nullopt;
+      }
+
+      return result;
+    }
+
+    if (!result.PushBack({ trimmed_range_start, range_end })) {
+      return std::nullopt;
+    }
+    current_index += blocks_in_range;
+  }
+
+  LOG(ERROR) << "Failed to construct byte ranges to read, start_block: " << start_index
+             << ", num_of_blocks: " << num_of_blocks << " total number of blocks: " << blocks_;
+  return std::nullopt;
+}
+
 // Ranges in the the set should be mutually exclusive; and they're sorted by the start block.
 SortedRangeSet::SortedRangeSet(std::vector<Range>&& pairs) : RangeSet(std::move(pairs)) {
   std::sort(ranges_.begin(), ranges_.end());
diff --git a/otautil/sysutil.cpp b/otautil/sysutil.cpp
index 8366fa0ac..420db4cac 100644
--- a/otautil/sysutil.cpp
+++ b/otautil/sysutil.cpp
@@ -94,6 +94,11 @@ BlockMapData BlockMapData::ParseBlockMapFile(const std::string& block_map_path)
     remaining_blocks -= range_blocks;
   }
 
+  if (remaining_blocks != 0) {
+    LOG(ERROR) << "Invalid ranges: remaining blocks " << remaining_blocks;
+    return {};
+  }
+
   return BlockMapData(block_dev, file_size, blksize, std::move(ranges));
 }
 
@@ -214,14 +219,21 @@ MemMapping::~MemMapping() {
   ranges_.clear();
 }
 
-bool reboot(const std::string& command) {
-  std::string cmd = command;
-  if (android::base::GetBoolProperty("ro.boot.quiescent", false)) {
+bool Reboot(std::string_view target) {
+  std::string cmd = "reboot," + std::string(target);
+  // Honor the quiescent mode if applicable.
+  if (target != "bootloader" && target != "fastboot" &&
+      android::base::GetBoolProperty("ro.boot.quiescent", false)) {
     cmd += ",quiescent";
   }
   return android::base::SetProperty(ANDROID_RB_PROPERTY, cmd);
 }
 
+bool Shutdown() {
+  // "shutdown" doesn't need a "reason" arg nor a comma.
+  return android::base::SetProperty(ANDROID_RB_PROPERTY, "shutdown");
+}
+
 std::vector<char*> StringVectorToNullTerminatedArray(const std::vector<std::string>& args) {
   std::vector<char*> result(args.size());
   std::transform(args.cbegin(), args.cend(), result.begin(),