[lldb][AArch64] Refactor memory tag range handling

Previously GetMemoryTagManager checked many things in one:
* architecture supports memory tagging
* process supports memory tagging
* memory range isn't inverted
* memory range is all tagged

Since writing follow up patches for tag writing (in review
at the moment) it has become clear that this gets unwieldy
once we add the features needed for that.

It also implies that the memory tag manager is tied to the
range you used to request it with but it is not. It's a per
process object.

Instead:
* GetMemoryTagManager just checks architecture and process.
* Then the MemoryTagManager can later be asked to check a
  memory range.

This is better because:
* We don't imply that range and manager are tied together.
* A slightly diferent range calculation for tag writing
  doesn't add more code to Process.
* Range checking code can now be unit tested.

Reviewed By: omjavaid

Differential Revision: https://reviews.llvm.org/D105630

Author: DavidSpickett

lldb/include/lldb/Target/MemoryTagManager.h

@@ -9,6 +9,7 @@
 #ifndef LLDB_TARGET_MEMORYTAGMANAGER_H
 #define LLDB_TARGET_MEMORYTAGMANAGER_H
 
+#include "lldb/Target/MemoryRegionInfo.h"
 #include "lldb/Utility/RangeMap.h"
 #include "lldb/lldb-private.h"
 #include "llvm/Support/Error.h"
@@ -57,6 +58,20 @@ class MemoryTagManager {
   // expanded to 2 granules.
   virtual TagRange ExpandToGranule(TagRange range) const = 0;
 
+  // Given a range addr to end_addr, check that:
+  // * end_addr >= addr (when memory tags are removed)
+  // * the granule aligned range is completely covered by tagged memory
+  //   (which may include one or more memory regions)
+  //
+  // If so, return a modified range which will have been expanded
+  // to be granule aligned.
+  //
+  // Tags in the input addresses are ignored and not present
+  // in the returned range.
+  virtual llvm::Expected<TagRange> MakeTaggedRange(
+      lldb::addr_t addr, lldb::addr_t end_addr,
+      const lldb_private::MemoryRegionInfos &memory_regions) const = 0;
+
   // Return the type value to use in GDB protocol qMemTags packets to read
   // allocation tags. This is named "Allocation" specifically because the spec
   // allows for logical tags to be read the same way, though we do not use that.

lldb/include/lldb/Target/Process.h

@@ -1709,30 +1709,19 @@ class Process : public std::enable_shared_from_this<Process>,
   lldb::addr_t CallocateMemory(size_t size, uint32_t permissions,
                                Status &error);
 
-  /// If the address range given is in a memory tagged range and this
-  /// architecture and process supports memory tagging, return a tag
+  /// If this architecture and process supports memory tagging, return a tag
   /// manager that can be used to maniupulate those memory tags.
-  /// Tags present in the addresses given are ignored.
-  ///
-  /// \param[in] addr
-  ///     Start of memory range.
-  ///
-  /// \param[in] end_addr
-  ///     End of the memory range. Where end is one beyond the last byte to be
-  ///     included.
   ///
   /// \return
   ///     Either a valid pointer to a tag manager or an error describing why one
   ///     could not be provided.
-  llvm::Expected<const MemoryTagManager *>
-  GetMemoryTagManager(lldb::addr_t addr, lldb::addr_t end_addr);
+  llvm::Expected<const MemoryTagManager *> GetMemoryTagManager();
 
-  /// Expands the range addr to addr+len to align with granule boundaries and
-  /// then calls DoReadMemoryTags to do the target specific operations.
-  /// Tags are returned unpacked so can be used without conversion.
+  /// Read memory tags for the range addr to addr+len. It is assumed
+  /// that this range has already been granule aligned.
+  /// (see MemoryTagManager::MakeTaggedRange)
   ///
-  /// \param[in] tag_manager
-  ///     The tag manager to get memory tagging information from.
+  /// This calls DoReadMemoryTags to do the target specific operations.
   ///
   /// \param[in] addr
   ///     Start of memory range to read tags for.
@@ -1741,11 +1730,12 @@ class Process : public std::enable_shared_from_this<Process>,
   ///     Length of memory range to read tags for (in bytes).
   ///
   /// \return
-  ///     Either the unpacked tags or an error describing a failure to read
-  ///     or unpack them.
-  llvm::Expected<std::vector<lldb::addr_t>>
-  ReadMemoryTags(const MemoryTagManager *tag_manager, lldb::addr_t addr,
-                 size_t len);
+  ///     If this architecture or process does not support memory tagging,
+  ///     an error saying so.
+  ///     If it does, either the memory tags or an error describing a
+  ///     failure to read or unpack them.
+  llvm::Expected<std::vector<lldb::addr_t>> ReadMemoryTags(lldb::addr_t addr,
+                                                           size_t len);
 
   /// Resolve dynamically loaded indirect functions.
   ///

lldb/source/Commands/CommandObjectMemoryTag.cpp

@@ -68,17 +68,29 @@ class CommandObjectMemoryTagRead : public CommandObjectParsed {
 
     Process *process = m_exe_ctx.GetProcessPtr();
     llvm::Expected<const MemoryTagManager *> tag_manager_or_err =
-        process->GetMemoryTagManager(start_addr, end_addr);
+        process->GetMemoryTagManager();
 
     if (!tag_manager_or_err) {
       result.SetError(Status(tag_manager_or_err.takeError()));
       return false;
     }
 
     const MemoryTagManager *tag_manager = *tag_manager_or_err;
-    ptrdiff_t len = tag_manager->AddressDiff(end_addr, start_addr);
-    llvm::Expected<std::vector<lldb::addr_t>> tags =
-        process->ReadMemoryTags(tag_manager, start_addr, len);
+
+    MemoryRegionInfos memory_regions;
+    // If this fails the list of regions is cleared, so we don't need to read
+    // the return status here.
+    process->GetMemoryRegions(memory_regions);
+    llvm::Expected<MemoryTagManager::TagRange> tagged_range =
+        tag_manager->MakeTaggedRange(start_addr, end_addr, memory_regions);
+
+    if (!tagged_range) {
+      result.SetError(Status(tagged_range.takeError()));
+      return false;
+    }
+
+    llvm::Expected<std::vector<lldb::addr_t>> tags = process->ReadMemoryTags(
+        tagged_range->GetRangeBase(), tagged_range->GetByteSize());
 
     if (!tags) {
       result.SetError(Status(tags.takeError()));
@@ -89,8 +101,7 @@ class CommandObjectMemoryTagRead : public CommandObjectParsed {
                                     tag_manager->GetLogicalTag(start_addr));
     result.AppendMessage("Allocation tags:");
 
-    MemoryTagManager::TagRange initial_range(start_addr, len);
-    addr_t addr = tag_manager->ExpandToGranule(initial_range).GetRangeBase();
+    addr_t addr = tagged_range->GetRangeBase();
     for (auto tag : *tags) {
       addr_t next_addr = addr + tag_manager->GetGranuleSize();
       // Showing tagged adresses here until we have non address bit handling

lldb/source/Plugins/Process/Utility/MemoryTagManagerAArch64MTE.cpp

@@ -66,6 +66,63 @@ MemoryTagManagerAArch64MTE::ExpandToGranule(TagRange range) const {
   return TagRange(new_start, new_len);
 }
 
+llvm::Expected<MemoryTagManager::TagRange>
+MemoryTagManagerAArch64MTE::MakeTaggedRange(
+    lldb::addr_t addr, lldb::addr_t end_addr,
+    const lldb_private::MemoryRegionInfos &memory_regions) const {
+  // First check that the range is not inverted.
+  // We must remove tags here otherwise an address with a higher
+  // tag value will always be > the other.
+  ptrdiff_t len = AddressDiff(end_addr, addr);
+  if (len <= 0) {
+    return llvm::createStringError(
+        llvm::inconvertibleErrorCode(),
+        "End address (0x%" PRIx64
+        ") must be greater than the start address (0x%" PRIx64 ")",
+        end_addr, addr);
+  }
+
+  // Region addresses will not have memory tags. So when searching
+  // we must use an untagged address.
+  MemoryRegionInfo::RangeType tag_range(RemoveNonAddressBits(addr), len);
+  tag_range = ExpandToGranule(tag_range);
+
+  // Make a copy so we can use the original for errors and the final return.
+  MemoryRegionInfo::RangeType remaining_range(tag_range);
+
+  // While there are parts of the range that don't have a matching tagged memory
+  // region
+  while (remaining_range.IsValid()) {
+    // Search for a region that contains the start of the range
+    MemoryRegionInfos::const_iterator region = std::find_if(
+        memory_regions.cbegin(), memory_regions.cend(),
+        [&remaining_range](const MemoryRegionInfo &region) {
+          return region.GetRange().Contains(remaining_range.GetRangeBase());
+        });
+
+    if (region == memory_regions.cend() ||
+        region->GetMemoryTagged() != MemoryRegionInfo::eYes) {
+      // Some part of this range is untagged (or unmapped) so error
+      return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                     "Address range 0x%" PRIx64 ":0x%" PRIx64
+                                     " is not in a memory tagged region",
+                                     tag_range.GetRangeBase(),
+                                     tag_range.GetRangeEnd());
+    }
+
+    // We've found some part of the range so remove that part and continue
+    // searching for the rest. Moving the base "slides" the range so we need to
+    // save/restore the original end. If old_end is less than the new base, the
+    // range will be set to have 0 size and we'll exit the while.
+    lldb::addr_t old_end = remaining_range.GetRangeEnd();
+    remaining_range.SetRangeBase(region->GetRange().GetRangeEnd());
+    remaining_range.SetRangeEnd(old_end);
+  }
+
+  // Every part of the range is contained within a tagged memory region.
+  return tag_range;
+}
+
 llvm::Expected<std::vector<lldb::addr_t>>
 MemoryTagManagerAArch64MTE::UnpackTagsData(const std::vector<uint8_t> &tags,
                                            size_t granules) const {

lldb/source/Plugins/Process/Utility/MemoryTagManagerAArch64MTE.h

@@ -32,6 +32,10 @@ class MemoryTagManagerAArch64MTE : public MemoryTagManager {
 
   TagRange ExpandToGranule(TagRange range) const override;
 
+  llvm::Expected<TagRange> MakeTaggedRange(
+      lldb::addr_t addr, lldb::addr_t end_addr,
+      const lldb_private::MemoryRegionInfos &memory_regions) const override;
+
   llvm::Expected<std::vector<lldb::addr_t>>
   UnpackTagsData(const std::vector<uint8_t> &tags,
                  size_t granules) const override;

lldb/source/Target/Process.cpp

@@ -6067,8 +6067,7 @@ bool Process::CallVoidArgVoidPtrReturn(const Address *address,
   return false;
 }
 
-llvm::Expected<const MemoryTagManager *>
-Process::GetMemoryTagManager(lldb::addr_t addr, lldb::addr_t end_addr) {
+llvm::Expected<const MemoryTagManager *> Process::GetMemoryTagManager() {
   Architecture *arch = GetTarget().GetArchitecturePlugin();
   const MemoryTagManager *tag_manager =
       arch ? arch->GetMemoryTagManager() : nullptr;
@@ -6084,66 +6083,22 @@ Process::GetMemoryTagManager(lldb::addr_t addr, lldb::addr_t end_addr) {
                                    "Process does not support memory tagging");
   }
 
-  ptrdiff_t len = tag_manager->AddressDiff(end_addr, addr);
-  if (len <= 0) {
-    return llvm::createStringError(
-        llvm::inconvertibleErrorCode(),
-        "End address (0x%" PRIx64
-        ") must be greater than the start address (0x%" PRIx64 ")",
-        end_addr, addr);
-  }
-
-  // Region lookup is not address size aware so mask the address
-  MemoryRegionInfo::RangeType tag_range(tag_manager->RemoveNonAddressBits(addr),
-                                        len);
-  tag_range = tag_manager->ExpandToGranule(tag_range);
-
-  // Make a copy so we can use the original range in errors
-  MemoryRegionInfo::RangeType remaining_range(tag_range);
-
-  // While we haven't found a matching memory region for some of the range
-  while (remaining_range.IsValid()) {
-    MemoryRegionInfo region;
-    Status status = GetMemoryRegionInfo(remaining_range.GetRangeBase(), region);
-
-    if (status.Fail() || region.GetMemoryTagged() != MemoryRegionInfo::eYes) {
-      return llvm::createStringError(
-          llvm::inconvertibleErrorCode(),
-          "Address range 0x%lx:0x%lx is not in a memory tagged region",
-          tag_range.GetRangeBase(), tag_range.GetRangeEnd());
-    }
-
-    if (region.GetRange().GetRangeEnd() >= remaining_range.GetRangeEnd()) {
-      // We've found a region for the whole range or the last piece of a range
-      remaining_range.SetByteSize(0);
-    } else {
-      // We've found some part of the range, look for the rest
-      remaining_range.SetRangeBase(region.GetRange().GetRangeEnd());
-    }
-  }
-
   return tag_manager;
 }
 
 llvm::Expected<std::vector<lldb::addr_t>>
-Process::ReadMemoryTags(const MemoryTagManager *tag_manager, lldb::addr_t addr,
-                        size_t len) {
-  if (!tag_manager) {
-    return llvm::createStringError(
-        llvm::inconvertibleErrorCode(),
-        "A memory tag manager is required for reading memory tags.");
-  }
-
-  MemoryTagManager::TagRange range(tag_manager->RemoveNonAddressBits(addr),
-                                   len);
-  range = tag_manager->ExpandToGranule(range);
+Process::ReadMemoryTags(lldb::addr_t addr, size_t len) {
+  llvm::Expected<const MemoryTagManager *> tag_manager_or_err =
+      GetMemoryTagManager();
+  if (!tag_manager_or_err)
+    return tag_manager_or_err.takeError();
 
+  const MemoryTagManager *tag_manager = *tag_manager_or_err;
   llvm::Expected<std::vector<uint8_t>> tag_data =
-      DoReadMemoryTags(range.GetRangeBase(), range.GetByteSize(),
-                       tag_manager->GetAllocationTagType());
+      DoReadMemoryTags(addr, len, tag_manager->GetAllocationTagType());
   if (!tag_data)
     return tag_data.takeError();
 
-  return tag_manager->UnpackTagsData(
-      *tag_data, range.GetByteSize() / tag_manager->GetGranuleSize());
+  return tag_manager->UnpackTagsData(*tag_data,
+                                     len / tag_manager->GetGranuleSize());
 }