[lldb] Add "memory tag write" --end-addr option

The default mode of "memory tag write" is to calculate the
range from the start address and the number of tags given.
(just like "memory write" does)

(lldb) memory tag write mte_buf 1 2
(lldb) memory tag read mte_buf mte_buf+48
Logical tag: 0x0
Allocation tags:
[0xfffff7ff9000, 0xfffff7ff9010): 0x1
[0xfffff7ff9010, 0xfffff7ff9020): 0x2
[0xfffff7ff9020, 0xfffff7ff9030): 0x0

This new option allows you to set an end address and have
the tags repeat until that point.

(lldb) memory tag write mte_buf 1 2 --end-addr mte_buf+64
(lldb) memory tag read mte_buf mte_buf+80
Logical tag: 0x0
Allocation tags:
[0xfffff7ff9000, 0xfffff7ff9010): 0x1
[0xfffff7ff9010, 0xfffff7ff9020): 0x2
[0xfffff7ff9020, 0xfffff7ff9030): 0x1
[0xfffff7ff9030, 0xfffff7ff9040): 0x2
[0xfffff7ff9040, 0xfffff7ff9050): 0x0

This is implemented using the QMemTags packet previously
added. We skip validating the number of tags in lldb and send
them on to lldb-server, which repeats them as needed.

Apart from the number of tags, all the other client side checks
remain. Tag values, memory range must be tagged, etc.

Reviewed By: omjavaid

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

lldb/source/Commands/CommandObjectMemoryTag.cpp

@@ -7,8 +7,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "CommandObjectMemoryTag.h"
+#include "lldb/Host/OptionParser.h"
 #include "lldb/Interpreter/CommandReturnObject.h"
 #include "lldb/Interpreter/OptionArgParser.h"
+#include "lldb/Interpreter/OptionGroupFormat.h"
+#include "lldb/Interpreter/OptionValueString.h"
 #include "lldb/Target/Process.h"
 
 using namespace lldb;
@@ -120,23 +123,64 @@ protected:
 
 class CommandObjectMemoryTagWrite : public CommandObjectParsed {
 public:
+  class OptionGroupTagWrite : public OptionGroup {
+  public:
+    OptionGroupTagWrite() : OptionGroup(), m_end_addr(LLDB_INVALID_ADDRESS) {}
+
+    ~OptionGroupTagWrite() override = default;
+
+    llvm::ArrayRef<OptionDefinition> GetDefinitions() override {
+      return llvm::makeArrayRef(g_memory_tag_write_options);
+    }
+
+    Status SetOptionValue(uint32_t option_idx, llvm::StringRef option_value,
+                          ExecutionContext *execution_context) override {
+      Status status;
+      const int short_option =
+          g_memory_tag_write_options[option_idx].short_option;
+
+      switch (short_option) {
+      case 'e':
+        m_end_addr = OptionArgParser::ToAddress(execution_context, option_value,
+                                                LLDB_INVALID_ADDRESS, &status);
+        break;
+      default:
+        llvm_unreachable("Unimplemented option");
+      }
+
+      return status;
+    }
+
+    void OptionParsingStarting(ExecutionContext *execution_context) override {
+      m_end_addr = LLDB_INVALID_ADDRESS;
+    }
+
+    lldb::addr_t m_end_addr;
+  };
+
   CommandObjectMemoryTagWrite(CommandInterpreter &interpreter)
       : CommandObjectParsed(interpreter, "tag",
                             "Write memory tags starting from the granule that "
                             "contains the given address.",
                             nullptr,
                             eCommandRequiresTarget | eCommandRequiresProcess |
-                                eCommandProcessMustBePaused) {
+                                eCommandProcessMustBePaused),
+        m_option_group(), m_tag_write_options() {
     // Address
     m_arguments.push_back(
         CommandArgumentEntry{CommandArgumentData(eArgTypeAddressOrExpression)});
     // One or more tag values
     m_arguments.push_back(CommandArgumentEntry{
         CommandArgumentData(eArgTypeValue, eArgRepeatPlus)});
+
+    m_option_group.Append(&m_tag_write_options);
+    m_option_group.Finalize();
   }
 
   ~CommandObjectMemoryTagWrite() override = default;
 
+  Options *GetOptions() override { return &m_option_group; }
+
 protected:
   bool DoExecute(Args &command, CommandReturnObject &result) override {
     if (command.GetArgumentCount() < 2) {
@@ -196,14 +240,24 @@ protected:
         tag_manager->ExpandToGranule(MemoryTagManager::TagRange(start_addr, 1))
             .GetRangeBase();
 
+    lldb::addr_t end_addr = 0;
+    // When you have an end address you want to align the range like tag read
+    // does. Meaning, align the start down (which we've done) and align the end
+    // up.
+    if (m_tag_write_options.m_end_addr != LLDB_INVALID_ADDRESS)
+      end_addr = m_tag_write_options.m_end_addr;
+    else
+      // Without an end address assume number of tags matches number of granules
+      // to write to
+      end_addr =
+          aligned_start_addr + (tags.size() * tag_manager->GetGranuleSize());
+
     // Now we've aligned the start address so if we ask for another range
     // using the number of tags N, we'll get back a range that is also N
     // granules in size.
     llvm::Expected<MemoryTagManager::TagRange> tagged_range =
-        tag_manager->MakeTaggedRange(
-            aligned_start_addr,
-            aligned_start_addr + (tags.size() * tag_manager->GetGranuleSize()),
-            memory_regions);
+        tag_manager->MakeTaggedRange(aligned_start_addr, end_addr,
+                                     memory_regions);
 
     if (!tagged_range) {
       result.SetError(Status(tagged_range.takeError()));
@@ -221,6 +275,9 @@ protected:
     result.SetStatus(eReturnStatusSuccessFinishResult);
     return true;
   }
+
+  OptionGroupOptions m_option_group;
+  OptionGroupTagWrite m_tag_write_options;
 };
 
 CommandObjectMemoryTag::CommandObjectMemoryTag(CommandInterpreter &interpreter)

lldb/source/Commands/Options.td

@@ -504,6 +504,14 @@ let Command = "memory write" in {
     Desc<"Start writing bytes from an offset within the input file.">;
 }
 
+let Command = "memory tag write" in {
+  def memory_write_end_addr : Option<"end-addr", "e">, Group<1>,
+  Arg<"AddressOrExpression">, Desc<
+    "Set tags for start address to end-addr, repeating tags as needed"
+    " to cover the range. (instead of calculating the range from the"
+    " number of tags given)">;
+}
+
 let Command = "register read" in {
   def register_read_alternate : Option<"alternate", "A">,
     Desc<"Display register names using the alternate register name if there "

lldb/test/API/linux/aarch64/mte_tag_access/TestAArch64LinuxMTEMemoryTagAccess.py

@@ -216,3 +216,59 @@ class AArch64LinuxMTEMemoryTagAccessTestCase(TestBase):
         self.expect("memory tag write mte_buf 99",
                 patterns=["error: Found tag 0x63 which is > max MTE tag value of 0xf."],
                 error=True)
+
+        # You can provide an end address and have lldb repeat the tags as needed
+        # The range is checked in the same way it is for "memory tag read"
+        self.expect("memory tag write mte_buf 9 -e",
+                patterns=["error: last option requires an argument"],
+                error=True)
+        self.expect("memory tag write mte_buf 9 -e food",
+                patterns=["error: address expression \"food\" evaluation failed"],
+                error=True)
+        self.expect("memory tag write mte_buf_2 9 --end-addr mte_buf_2",
+                patterns=["error: End address \(0x[A-Fa-f0-9]+\) must be "
+                          "greater than the start address \(0x[A-Fa-f0-9]+\)"],
+                error=True)
+        self.expect("memory tag write mte_buf_2 9 --end-addr mte_buf_2-16",
+                patterns=["error: End address \(0x[A-Fa-f0-9]+\) must be "
+                          "greater than the start address \(0x[A-Fa-f0-9]+\)"],
+                error=True)
+        self.expect("memory tag write mte_buf_2 9 --end-addr mte_buf_2+page_size+16",
+                patterns=["error: Address range 0x[0-9A-fa-f]+00:0x[0-9A-Fa-f]+10 "
+                          "is not in a memory tagged region"],
+                error=True)
+
+        # Tags are repeated across the range
+        # For these we'll read one extra to make sure we don't over write
+        self.expect("memory tag write mte_buf_2 4 5 --end-addr mte_buf_2+48")
+        self.expect("memory tag read mte_buf_2 mte_buf_2+64",
+                patterns=["Logical tag: 0x0\n"
+                          "Allocation tags:\n"
+                          "\[0x[0-9A-Fa-f]+00, 0x[0-9A-Fa-f]+10\): 0x4\n"
+                          "\[0x[0-9A-Fa-f]+10, 0x[0-9A-Fa-f]+20\): 0x5\n"
+                          "\[0x[0-9A-Fa-f]+20, 0x[0-9A-Fa-f]+30\): 0x4\n"
+                          "\[0x[0-9A-Fa-f]+30, 0x[0-9A-Fa-f]+40\): 0x0$"])
+
+        # Since this aligns like tag read does, the start is aligned down and the end up.
+        # Meaning that start/end tells you the start/end granule that will be written.
+        # This matters particularly if either are misaligned.
+
+        # Here start moves down so the final range is mte_buf_2 -> mte_buf_2+32
+        self.expect("memory tag write mte_buf_2+8 6 -end-addr mte_buf_2+32")
+        self.expect("memory tag read mte_buf_2 mte_buf_2+48",
+                patterns=["Logical tag: 0x0\n"
+                          "Allocation tags:\n"
+                          "\[0x[0-9A-Fa-f]+00, 0x[0-9A-Fa-f]+10\): 0x6\n"
+                          "\[0x[0-9A-Fa-f]+10, 0x[0-9A-Fa-f]+20\): 0x6\n"
+                          "\[0x[0-9A-Fa-f]+20, 0x[0-9A-Fa-f]+30\): 0x4$"])
+
+        # If we do the same with a misaligned end, it also moves but upward.
+        # The intial range is 2 granules but the final range is mte_buf_2 -> mte_buf_2+48
+        self.expect("memory tag write mte_buf_2+8 3 -end-addr mte_buf_2+32+8")
+        self.expect("memory tag read mte_buf_2 mte_buf_2+64",
+                patterns=["Logical tag: 0x0\n"
+                          "Allocation tags:\n"
+                          "\[0x[0-9A-Fa-f]+00, 0x[0-9A-Fa-f]+10\): 0x3\n"
+                          "\[0x[0-9A-Fa-f]+10, 0x[0-9A-Fa-f]+20\): 0x3\n"
+                          "\[0x[0-9A-Fa-f]+20, 0x[0-9A-Fa-f]+30\): 0x3\n"
+                          "\[0x[0-9A-Fa-f]+30, 0x[0-9A-Fa-f]+40\): 0x0$"])