[lldb][AArch64] Add tag packing and repetition memory tag manager

[DavidSpickett]

Summary

PackTags is used by to compress tags to go in the QMemTags packet
and be passed to ptrace when writing memory tags.

The behaviour of RepeatTagsForRange matches that described for QMemTags
in the GDB documentation:
https://sourceware.org/gdb/current/onlinedocs/gdb/General-Query-Packets.html#General-Query-Packets

In addition, unpacking tags with number of tags 0 now means
do not check that number of tags matches the range.
This will be used by lldb-server to unpack tags before repeating
them to fill the requested range.

[omjavaid]

This looks good. I just made some cosmetic comments.

lldb/include/lldb/Target/MemoryTagManager.h

92

Do you think it will be better if we default granules to 0 then granules will become optional?

lldb/unittests/Process/Utility/MemoryTagManagerAArch64MTETest.cpp

61

May be add a comment here as you have done for all the tests above and below this one.

274

May be add a comment explaining this helper.

[DavidSpickett]

• Make granules default to 0 for UnpackTagsData
• Add more comments to tests

(I'll add a method to repeat packed tags in the later patch
that'll use it)

[DavidSpickett]

Rebase, correct the commit title and printf formatters for 32 bit build.

[omjavaid]

Looks good.

lldb/include/lldb/Target/MemoryTagManager.h

public:
   virtual int32_t GetAllocationTagType() const = 0;
 
   // Return the number of bytes a single tag will be packed into during
   // transport. For example an MTE tag is 4 bits but occupies 1 byte during
   // transport.
   virtual size_t GetTagSizeInBytes() const = 0;
 
   // Unpack tags from their stored format (e.g. gdb qMemTags data) into seperate
-  // tags. Checks that each tag is within the expected value range and that the
-  // number of tags found matches the number of granules we originally asked
-  // for.
+  // tags.
+  //
+  // Checks that each tag is within the expected value range and if granules is
+  // set to non-zero, that the number of tags found matches the number of
+  // granules we expected to cover.
+  virtual llvm::Expected<std::vector<lldb::addr_t>>
+  UnpackTagsData(const std::vector<uint8_t> &tags,
+                 size_t granules = 0) const = 0;
+
+  // Pack uncompressed tags into their storage format (e.g. for gdb QMemTags).
+  // Checks that each tag is within the expected value range.
+  // We do not check the number of tags or range they apply to because
+  // it is up to the remote to repeat them as needed.
+  virtual llvm::Expected<std::vector<uint8_t>>
+  PackTags(const std::vector<lldb::addr_t> &tags) const = 0;
+
+  // Take a set of tags and repeat them as much as needed to cover the given
+  // range. We assume that this range has been previously expanded/aligned to
+  // granules. (this method is used by lldb-server to implement QMemTags
+  // packet handling)
+  //
+  // If the range is empty, zero tags are returned.
+  // If the range is not empty and...
+  //   * there are no tags, an error is returned.
+  //   * there are fewer tags than granules, the tags are repeated to fill the
+  //     range.
+  //   * there are more tags than granules, only the tags required to cover
+  //     the range are returned.
+  //
+  // When repeating tags it will not always return a multiple of the original
+  // list. For example if your range is 3 granules and your tags are 1 and 2.
+  // You will get tags 1, 2 and 1 returned. Rather than getting 1, 2, 1, 2,
+  // which would be one too many tags for the range.
+  //
+  // A single tag will just be repeated as you'd expected. Tag 1 over 3 granules
+  // would return 1, 1, 1.
   virtual llvm::Expected<std::vector<lldb::addr_t>>
-  UnpackTagsData(const std::vector<uint8_t> &tags, size_t granules) const = 0;
+  RepeatTagsForRange(const std::vector<lldb::addr_t> &tags,
+                     TagRange range) const = 0;
 
   virtual ~MemoryTagManager() {}
 };
 
 } // namespace lldb_private
 
 #endif // LLDB_TARGET_MEMORYTAGMANAGER_H

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

public:
   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;
+                 size_t granules = 0) const override;
+
+  llvm::Expected<std::vector<uint8_t>>
+  PackTags(const std::vector<lldb::addr_t> &tags) const override;
+
+  llvm::Expected<std::vector<lldb::addr_t>>
+  RepeatTagsForRange(const std::vector<lldb::addr_t> &tags,
+                     TagRange range) const override;
 };
 
 } // namespace lldb_private
 
 #endif // LLDB_SOURCE_PLUGINS_PROCESS_UTILITY_MEMORYTAGMANAGERAARCH64MTE_H

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

MemoryTagManagerAArch64MTE::MakeTaggedRange(
   }
 
   // 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 {
+                                           size_t granules /*=0*/) const {
+  // 0 means don't check the number of tags before unpacking
+  if (granules) {
     size_t num_tags = tags.size() / GetTagSizeInBytes();
     if (num_tags != granules) {
       return llvm::createStringError(
           llvm::inconvertibleErrorCode(),
           "Packed tag data size does not match expected number of tags. "
-        "Expected %zu tag(s) for %zu granules, got %zu tag(s).",
+          "Expected %zu tag(s) for %zu granule(s), got %zu tag(s).",
           granules, granules, num_tags);
     }
+  }
 
   // (if bytes per tag was not 1, we would reconstruct them here)
 
   std::vector<lldb::addr_t> unpacked;
   unpacked.reserve(tags.size());
   for (auto it = tags.begin(); it != tags.end(); ++it) {
     // Check all tags are in range
     if (*it > MTE_TAG_MAX) {
       return llvm::createStringError(
           llvm::inconvertibleErrorCode(),
           "Found tag 0x%x which is > max MTE tag value of 0x%x.", *it,
           MTE_TAG_MAX);
     }
     unpacked.push_back(*it);
   }
 
   return unpacked;
 }
+
+llvm::Expected<std::vector<uint8_t>> MemoryTagManagerAArch64MTE::PackTags(
+    const std::vector<lldb::addr_t> &tags) const {
+  std::vector<uint8_t> packed;
+  packed.reserve(tags.size() * GetTagSizeInBytes());
+
+  for (auto tag : tags) {
+    if (tag > MTE_TAG_MAX) {
+      return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                     "Found tag 0x%" PRIx64
+                                     " which is > max MTE tag value of 0x%x.",
+                                     tag, MTE_TAG_MAX);
+    }
+    packed.push_back(static_cast<uint8_t>(tag));
+  }
+
+  return packed;
+}
+
+llvm::Expected<std::vector<lldb::addr_t>>
+MemoryTagManagerAArch64MTE::RepeatTagsForRange(
+    const std::vector<lldb::addr_t> &tags, TagRange range) const {
+  std::vector<lldb::addr_t> new_tags;
+
+  // If the range is not empty
+  if (range.IsValid()) {
+    if (tags.empty()) {
+      return llvm::createStringError(
+          llvm::inconvertibleErrorCode(),
+          "Expected some tags to cover given range, got zero.");
+    }
+
+    // We assume that this range has already been expanded/aligned to granules
+    size_t granules = range.GetByteSize() / GetGranuleSize();
+    new_tags.reserve(granules);
+    for (size_t to_copy = 0; granules > 0; granules -= to_copy) {
+      to_copy = granules > tags.size() ? tags.size() : granules;
+      new_tags.insert(new_tags.end(), tags.begin(), tags.begin() + to_copy);
+    }
+  }
+
+  return new_tags;
+}

lldb/unittests/Process/Utility/MemoryTagManagerAArch64MTETest.cpp

TEST(MemoryTagManagerAArch64MTETest, UnpackTagsData) {
   MemoryTagManagerAArch64MTE manager;
 
   // Error for insufficient tag data
   std::vector<uint8_t> input;
   ASSERT_THAT_EXPECTED(
       manager.UnpackTagsData(input, 2),
       llvm::FailedWithMessage(
           "Packed tag data size does not match expected number of tags. "
-          "Expected 2 tag(s) for 2 granules, got 0 tag(s)."));
+          "Expected 2 tag(s) for 2 granule(s), got 0 tag(s)."));
 
   // This is out of the valid tag range
   input.push_back(0x1f);
   ASSERT_THAT_EXPECTED(
       manager.UnpackTagsData(input, 1),
       llvm::FailedWithMessage(
           "Found tag 0x1f which is > max MTE tag value of 0xf."));
 
   // MTE tags are 1 per byte
   input.pop_back();
   input.push_back(0xe);
   input.push_back(0xf);
 
   std::vector<lldb::addr_t> expected{0xe, 0xf};
 
   llvm::Expected<std::vector<lldb::addr_t>> got =
       manager.UnpackTagsData(input, 2);
   ASSERT_THAT_EXPECTED(got, llvm::Succeeded());
   ASSERT_THAT(expected, testing::ContainerEq(*got));
+
+  // Error for too much tag data
+  ASSERT_THAT_EXPECTED(
+      manager.UnpackTagsData(input, 1),
+      llvm::FailedWithMessage(
+          "Packed tag data size does not match expected number of tags. "
+          "Expected 1 tag(s) for 1 granule(s), got 2 tag(s)."));
+
+  // By default, we don't check number of tags
+  llvm::Expected<std::vector<lldb::addr_t>> got_zero =
+      manager.UnpackTagsData(input);
+  ASSERT_THAT_EXPECTED(got_zero, llvm::Succeeded());
+  ASSERT_THAT(expected, testing::ContainerEq(*got));
+
+  // Which is the same as granules=0
+  got_zero = manager.UnpackTagsData(input, 0);
+  ASSERT_THAT_EXPECTED(got_zero, llvm::Succeeded());
+  ASSERT_THAT(expected, testing::ContainerEq(*got));
+}
+
+TEST(MemoryTagManagerAArch64MTETest, PackTags) {
+  MemoryTagManagerAArch64MTE manager;
+
+  // Error for tag out of range
+  llvm::Expected<std::vector<uint8_t>> invalid_tag_err =
+      manager.PackTags({0x10});
+  ASSERT_THAT_EXPECTED(
+      invalid_tag_err,
+      llvm::FailedWithMessage(
+          "Found tag 0x10 which is > max MTE tag value of 0xf."));
+
+  // 0xf here is the max tag value that we can pack
+  std::vector<lldb::addr_t> tags{0, 1, 0xf};
+  std::vector<uint8_t> expected{0, 1, 0xf};
+  llvm::Expected<std::vector<uint8_t>> packed = manager.PackTags(tags);
+  ASSERT_THAT_EXPECTED(packed, llvm::Succeeded());
+  ASSERT_THAT(expected, testing::ContainerEq(*packed));
 }
 
 TEST(MemoryTagManagerAArch64MTETest, GetLogicalTag) {
   MemoryTagManagerAArch64MTE manager;
 
   // Set surrounding bits to check shift is correct
   ASSERT_EQ((lldb::addr_t)0, manager.GetLogicalTag(0xe0e00000ffffffff));
   // Max tag value
TEST(MemoryTagManagerAArch64MTETest, AddressDiff) {
   // Result is signed
   ASSERT_EQ(10, manager.AddressDiff(10, 0));
   ASSERT_EQ(-10, manager.AddressDiff(0, 10));
   // Anything in the top byte is ignored
   ASSERT_EQ(0, manager.AddressDiff(0x2211222233334444, 0x3311222233334444));
   ASSERT_EQ(-32, manager.AddressDiff(0x5511222233334400, 0x4411222233334420));
   ASSERT_EQ(65, manager.AddressDiff(0x9911222233334441, 0x6611222233334400));
 }
+
+// Helper to check that repeating "tags" over "range" gives you
+// "expected_tags".
+static void
+test_repeating_tags(const std::vector<lldb::addr_t> &tags,
+                    MemoryTagManagerAArch64MTE::TagRange range,
+                    const std::vector<lldb::addr_t> &expected_tags) {
+  MemoryTagManagerAArch64MTE manager;
+  llvm::Expected<std::vector<lldb::addr_t>> tags_or_err =
+      manager.RepeatTagsForRange(tags, range);
+  ASSERT_THAT_EXPECTED(tags_or_err, llvm::Succeeded());
+  ASSERT_THAT(expected_tags, testing::ContainerEq(*tags_or_err));
+}
+
+TEST(MemoryTagManagerAArch64MTETest, RepeatTagsForRange) {
+  MemoryTagManagerAArch64MTE manager;
+
+  // Must have some tags if your range is not empty
+  llvm::Expected<std::vector<lldb::addr_t>> no_tags_err =
+      manager.RepeatTagsForRange({},
+                                 MemoryTagManagerAArch64MTE::TagRange{0, 16});
+  ASSERT_THAT_EXPECTED(
+      no_tags_err, llvm::FailedWithMessage(
+                       "Expected some tags to cover given range, got zero."));
+
+  // If the range is empty, you get no tags back
+  test_repeating_tags({1, 2, 3}, MemoryTagManagerAArch64MTE::TagRange{0, 0},
+                      {});
+  // And you don't need tags for an empty range
+  test_repeating_tags({}, MemoryTagManagerAArch64MTE::TagRange{0, 0}, {});
+
+  // A single tag will just be multiplied as many times as needed
+  test_repeating_tags({5}, MemoryTagManagerAArch64MTE::TagRange{0, 16}, {5});
+  test_repeating_tags({6}, MemoryTagManagerAArch64MTE::TagRange{0, 32}, {6, 6});
+
+  // If you've got as many tags as granules, it's a roundtrip
+  test_repeating_tags({7, 8}, MemoryTagManagerAArch64MTE::TagRange{0, 32},
+                      {7, 8});
+
+  // If you've got fewer tags than granules, they repeat. Exactly or partially
+  // as needed.
+  test_repeating_tags({7, 8}, MemoryTagManagerAArch64MTE::TagRange{0, 64},
+                      {7, 8, 7, 8});
+  test_repeating_tags({7, 8}, MemoryTagManagerAArch64MTE::TagRange{0, 48},
+                      {7, 8, 7});
+
+  // If you've got more tags than granules you get back only those needed
+  test_repeating_tags({1, 2, 3, 4}, MemoryTagManagerAArch64MTE::TagRange{0, 32},
+                      {1, 2});
+}