3 files changed, 64 insertions, 0 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/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 2b4861809..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));
 }