Add BTI support in the Metadata feature

src/aarch64/simulator-aarch64.cc

@@ -31,6 +31,8 @@
 #include <errno.h>
 #include <unistd.h>
 
+#include <sys/mman.h>
+
 #include <cmath>
 #include <cstring>
 #include <limits>
@@ -14310,6 +14312,8 @@ void Simulator::DoRestoreCPUFeatures(const Instruction* instr) {
   saved_cpu_features_.pop_back();
 }
 
+
+// BTI and MTE support in Simulator is not tested on aarch64.
 void* Simulator::Mmap(
     void* address, size_t length, int prot, int flags, int fd, off_t offset) {
   // The underlying system `mmap` in the simulated environment doesn't recognize
@@ -14321,6 +14325,10 @@ void* Simulator::Mmap(
   uint64_t address2 = reinterpret_cast<uint64_t>(
       mmap(address, length, prot, flags, fd, offset));
 
+  if (intenal_prot & PROT_BTI) {
+    SetGuardedPages(address2, length);
+  }
+
   if (intenal_prot & PROT_MTE) {
     // The returning address of `mmap` isn't tagged.
     int tag = static_cast<int>(GenerateRandomTag());
@@ -14333,6 +14341,10 @@ void* Simulator::Mmap(
 
 
 int Simulator::Munmap(void* address, size_t length, int prot) {
+  if (prot & PROT_BTI) {
+    ClearBTIGuard(address, length);
+  }
+
   if (prot & PROT_MTE) {
     // Untag the address since `munmap` doesn't recognize the memory tagging
     // managed by the Simulator.

src/aarch64/simulator-aarch64.h

@@ -29,6 +29,7 @@
 
 #include <memory>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "../globals-vixl.h"
@@ -260,9 +261,54 @@ class MetaDataDepot {
 
   size_t GetTotalCountMTE() { return metadata_mte_.size(); }
 
+  // BTI
+
+  // All addresses within a page share the same key.
+  uint64_t GenerateBTIKey(uint64_t address) const {
+    return AddressUntag(address) >> kPageSizeLog2;
+  }
+
+  bool InGuardedPage(uintptr_t address) const {
+    uint64_t key = GenerateBTIKey(address);
+    auto m = metadata_bti_.find(key);
+    // Assume the default page attribute is NOT guarded.
+    return m != metadata_bti_.end();
+  }
+
+  void ClearBTIGuard(uintptr_t address, size_t length) {
+    VIXL_ASSERT(IsAligned(address, kPageSize));
+    VIXL_ASSERT(length % kPageSize == 0);
+    for (size_t offset = 0; offset < length; offset += kPageSize) {
+      uint64_t key = GenerateBTIKey(address + offset);
+      if (metadata_bti_.erase(key)) {
+        std::stringstream sstream;
+        sstream << std::hex << "BTI WARNING : The page at 0x" << address
+                << " is already unguarded.\n";
+      }
+    }
+  }
+
+  void SetBTIGuard(uintptr_t address, size_t length) {
+    VIXL_ASSERT(IsAligned(address, kPageSize));
+    VIXL_ASSERT(length % kPageSize == 0);
+    for (size_t offset = 0; offset < length; offset += kPageSize) {
+      uint64_t key = GenerateBTIKey(address + offset);
+      auto m = metadata_bti_.find(key);
+      if (m == metadata_bti_.end()) {
+        metadata_bti_.insert(key);
+      } else {
+        std::stringstream sstream;
+        sstream << std::hex << "BTI WARNING : The page at 0x" << address
+                << " is already guarded.\n";
+      }
+    }
+  }
+
  private:
   // Tag recording of each allocated memory in the tag-granule.
   std::unordered_map<uint64_t, class MetaDataMTE> metadata_mte_;
+  // Recording of BTI-guarded pages.
+  std::unordered_set<uint64_t> metadata_bti_;
 };
 
 
@@ -1261,9 +1307,25 @@ class Simulator : public DecoderVisitor {
   // Helper function to determine BType for branches.
   BType GetBTypeFromInstruction(const Instruction* instr) const;
 
-  bool PcIsInGuardedPage() const { return guard_pages_; }
+  // Query if the page including the specified address is BTI-guarded. The
+  // address is the current PC by default.
+  bool PcIsInGuardedPage(const void* address = nullptr) const {
+    if (guard_pages_) return true;
+    if (address == nullptr) {
+      address = ReadPc();
+    }
+    return meta_data_.InGuardedPage(reinterpret_cast<uintptr_t>(address));
+  }
+
+  // Set the protection to all the pages.
   void SetGuardedPages(bool guard_pages) { guard_pages_ = guard_pages; }
 
+  // Set the protection to the specified consecutive pages.
+  template <typename T>
+  void SetGuardedPages(const T address, size_t length = kPageSize) {
+    meta_data_.SetBTIGuard(reinterpret_cast<uintptr_t>(address), length);
+  }
+
   void ExecuteInstruction() {
     // The program counter should always be aligned.
     VIXL_ASSERT(IsWordAligned(pc_));
@@ -3000,6 +3062,11 @@ class Simulator : public DecoderVisitor {
     return meta_data_.GetMTETag(address);
   }
 
+  template <typename T>
+  void ClearBTIGuard(T address, size_t length) {
+    meta_data_.ClearBTIGuard(reinterpret_cast<uintptr_t>(address), length);
+  }
+
   // Generate a random address tag, and any tags specified in the input are
   // excluded from the selection.
   uint64_t GenerateRandomTag(uint16_t exclude = 0);

test/aarch64/test-api-aarch64.cc

@@ -24,6 +24,8 @@
 // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+#include <sys/mman.h>
+
 #include <cstdio>
 #include <cstring>
 #include <string>
@@ -1765,5 +1767,34 @@ TEST(sim_stack) {
 }
 #endif
 
+#ifdef VIXL_INCLUDE_SIMULATOR_AARCH64
+TEST(bti_set_and_clear) {
+  Decoder simulator_decoder;
+  Simulator simulator(&simulator_decoder);
+  simulator.SetCPUFeatures(CPUFeatures::kBTI);
+
+  // Allocate pages and set up BTI protection.
+  int data_size = 10 * kPageSize;
+  byte* data_start =
+      reinterpret_cast<byte*>(simulator.Mmap(NULL,
+                                             data_size,
+                                             PROT_READ | PROT_WRITE | PROT_BTI,
+                                             MAP_PRIVATE | MAP_ANONYMOUS,
+                                             -1,
+                                             0));
+
+  for (int offset = 0; offset < data_size; offset += kPageSize) {
+    VIXL_CHECK(simulator.PcIsInGuardedPage(data_start + offset));
+  }
+
+  // Unmap above allocated pages and clear BTI protection.
+  simulator.Munmap(data_start, data_size, PROT_BTI);
+
+  for (int offset = 0; offset < data_size; offset += kPageSize) {
+    VIXL_CHECK(!simulator.PcIsInGuardedPage(data_start + offset));
+  }
+}
+#endif
+
 }  // namespace aarch64
 }  // namespace vixl