Added more functionality to dex file tracking

Additional functions were added for (un)poisoning only certain sections
of a dex file. All functions related to tracking are now in a class
(DexFileTrackingRegistrar) since unpoisoning subsections of an already
poisoned section cause poison reads. Static variables were added
representing specific configurations that were used for testing.

GetCodeItem from art/dex/layout/dex_ir.cc was moved to
art/runtime/dex_file.cc because dex_file_tracking_registrar calls
GetCodeItem and GetCodeItem returns a property of a DexFile.

Bug: 37754950
Test: export ART_DEX_FILE_ACCESS_TRACKING=true && m -j && m -j
SANITIZE_TARGET=address SANITIZE_LITE=true test-art-host

Change-Id: Ie4e5aefb6f80b18fdaae5b5e8aa3bc99d77cbeac
diff --git a/dexlayout/dex_ir.cc b/dexlayout/dex_ir.cc
index cf453b9..62ee445 100644
--- a/dexlayout/dex_ir.cc
+++ b/dexlayout/dex_ir.cc
@@ -57,31 +57,6 @@
                     entry.reg_)));
 }
 
-static uint32_t GetCodeItemSize(const DexFile::CodeItem& disk_code_item) {
-  uintptr_t code_item_start = reinterpret_cast<uintptr_t>(&disk_code_item);
-  uint32_t insns_size = disk_code_item.insns_size_in_code_units_;
-  uint32_t tries_size = disk_code_item.tries_size_;
-  if (tries_size == 0) {
-    uintptr_t insns_end = reinterpret_cast<uintptr_t>(&disk_code_item.insns_[insns_size]);
-    return insns_end - code_item_start;
-  } else {
-    // Get the start of the handler data.
-    const uint8_t* handler_data = DexFile::GetCatchHandlerData(disk_code_item, 0);
-    uint32_t handlers_size = DecodeUnsignedLeb128(&handler_data);
-    // Manually read each handler.
-    for (uint32_t i = 0; i < handlers_size; ++i) {
-      int32_t uleb128_count = DecodeSignedLeb128(&handler_data) * 2;
-      if (uleb128_count <= 0) {
-        uleb128_count = -uleb128_count + 1;
-      }
-      for (int32_t j = 0; j < uleb128_count; ++j) {
-        DecodeUnsignedLeb128(&handler_data);
-      }
-    }
-    return reinterpret_cast<uintptr_t>(handler_data) - code_item_start;
-  }
-}
-
 static uint32_t GetDebugInfoStreamSize(const uint8_t* debug_info_stream) {
   const uint8_t* stream = debug_info_stream;
   DecodeUnsignedLeb128(&stream);  // line_start
@@ -686,7 +661,7 @@
     }
   }
 
-  uint32_t size = GetCodeItemSize(disk_code_item);
+  uint32_t size = DexFile::GetCodeItemSize(disk_code_item);
   CodeItem* code_item = new CodeItem(
       registers_size, ins_size, outs_size, debug_info, insns_size, insns, tries, handler_list);
   code_item->SetSize(size);
diff --git a/runtime/dex_file.cc b/runtime/dex_file.cc
index aace8eb..1301cc2 100644
--- a/runtime/dex_file.cc
+++ b/runtime/dex_file.cc
@@ -678,6 +678,32 @@
   UNREACHABLE();
 }
 
+uint32_t DexFile::GetCodeItemSize(const DexFile::CodeItem& code_item) {
+  uintptr_t code_item_start = reinterpret_cast<uintptr_t>(&code_item);
+  uint32_t insns_size = code_item.insns_size_in_code_units_;
+  uint32_t tries_size = code_item.tries_size_;
+  const uint8_t* handler_data = GetCatchHandlerData(code_item, 0);
+
+  if (tries_size == 0 || handler_data == nullptr) {
+    uintptr_t insns_end = reinterpret_cast<uintptr_t>(&code_item.insns_[insns_size]);
+    return insns_end - code_item_start;
+  } else {
+    // Get the start of the handler data.
+    uint32_t handlers_size = DecodeUnsignedLeb128(&handler_data);
+    // Manually read each handler.
+    for (uint32_t i = 0; i < handlers_size; ++i) {
+      int32_t uleb128_count = DecodeSignedLeb128(&handler_data) * 2;
+      if (uleb128_count <= 0) {
+        uleb128_count = -uleb128_count + 1;
+      }
+      for (int32_t j = 0; j < uleb128_count; ++j) {
+        DecodeUnsignedLeb128(&handler_data);
+      }
+    }
+    return reinterpret_cast<uintptr_t>(handler_data) - code_item_start;
+  }
+}
+
 const DexFile::FieldId* DexFile::FindFieldId(const DexFile::TypeId& declaring_klass,
                                              const DexFile::StringId& name,
                                              const DexFile::TypeId& type) const {
diff --git a/runtime/dex_file.h b/runtime/dex_file.h
index 3de78ed..3249edb 100644
--- a/runtime/dex_file.h
+++ b/runtime/dex_file.h
@@ -636,6 +636,8 @@
   uint32_t FindCodeItemOffset(const DexFile::ClassDef& class_def,
                               uint32_t dex_method_idx) const;
 
+  static uint32_t GetCodeItemSize(const DexFile::CodeItem& disk_code_item);
+
   // Returns the declaring class descriptor string of a field id.
   const char* GetFieldDeclaringClassDescriptor(const FieldId& field_id) const {
     const DexFile::TypeId& type_id = GetTypeId(field_id.class_idx_);
diff --git a/runtime/dex_file_tracking_registrar.cc b/runtime/dex_file_tracking_registrar.cc
index cfbca3d..f41a50b 100644
--- a/runtime/dex_file_tracking_registrar.cc
+++ b/runtime/dex_file_tracking_registrar.cc
@@ -16,6 +16,9 @@
 
 #include "dex_file_tracking_registrar.h"
 
+#include <deque>
+#include <tuple>
+
 // For dex tracking through poisoning. Note: Requires forcing sanitization. This is the reason for
 // the ifdefs and early include.
 #ifdef ART_DEX_FILE_ACCESS_TRACKING
@@ -26,6 +29,7 @@
 #include "base/memory_tool.h"
 
 #include "base/logging.h"
+#include "dex_file-inl.h"
 
 namespace art {
 namespace dex {
@@ -39,11 +43,161 @@
     false;
 #endif
 
-void RegisterDexFile(const DexFile* const dex_file) {
-  if (kDexFileAccessTracking && dex_file != nullptr) {
-    LOG(ERROR) << dex_file->GetLocation() + " @ " << std::hex
-               << reinterpret_cast<uintptr_t>(dex_file->Begin());
-    MEMORY_TOOL_MAKE_NOACCESS(dex_file->Begin(), dex_file->Size());
+// The following are configurations of poisoning certain sections of a Dex File.
+// More will be added
+enum DexTrackingType {
+  // Poisons all of a Dex File when set.
+  kWholeDexTracking,
+  // Poisons all Code Items of a Dex File when set.
+  kCodeItemTracking,
+  // Poisons all subsections of a Code Item, except the Insns bytecode array
+  // section, when set for all Code Items in a Dex File.
+  kCodeItemNonInsnsTracking,
+  // Poisons all subsections of a Code Item, except the Insns bytecode array
+  // section, when set for all Code Items in a Dex File.
+  // Additionally unpoisons the entire Code Item when method is a class
+  // initializer.
+  kCodeItemNonInsnsNoClinitTracking,
+  // Poisons based on a custom tracking system which can be specified in
+  // SetDexSections
+  kCustomTracking,
+};
+
+// Intended for local changes only.
+// Represents the current configuration being run.
+static constexpr DexTrackingType kCurrentTrackingSystem = kWholeDexTracking;
+
+// Intended for local changes only.
+void DexFileTrackingRegistrar::SetDexSections() {
+  if (kDexFileAccessTracking || dex_file_ != nullptr) {
+    switch (kCurrentTrackingSystem) {
+      case kWholeDexTracking:
+        SetDexFileRegistration(true);
+        break;
+      case kCodeItemTracking:
+        SetAllCodeItemRegistration(true);
+        break;
+      case kCodeItemNonInsnsTracking:
+        SetAllCodeItemRegistration(true);
+        SetAllInsnsRegistration(false);
+        break;
+      case kCodeItemNonInsnsNoClinitTracking:
+        SetAllCodeItemRegistration(true);
+        SetAllInsnsRegistration(false);
+        SetCodeItemRegistration("<clinit>", false);
+        break;
+      case kCustomTracking:
+        // TODO: Add/remove additional calls here to (un)poison sections of
+        // dex_file_
+        break;
+    }
+  }
+}
+
+void RegisterDexFile(const DexFile* dex_file) {
+  DexFileTrackingRegistrar dex_tracking_registrar(dex_file);
+  dex_tracking_registrar.SetDexSections();
+  dex_tracking_registrar.SetCurrentRanges();
+}
+
+inline void SetRegistrationRange(const void* begin, size_t size, bool should_poison) {
+  if (should_poison) {
+    MEMORY_TOOL_MAKE_NOACCESS(begin, size);
+  } else {
+    // Note: MEMORY_TOOL_MAKE_UNDEFINED has the same functionality with Address
+    // Sanitizer. The difference has not been tested with Valgrind
+    MEMORY_TOOL_MAKE_DEFINED(begin, size);
+  }
+}
+
+void DexFileTrackingRegistrar::SetCurrentRanges() {
+  // This also empties range_values_ to avoid redundant (un)poisoning upon
+  // subsequent calls.
+  while (!range_values_.empty()) {
+    const std::tuple<const void*, size_t, bool>& current_range = range_values_.front();
+    SetRegistrationRange(std::get<0>(current_range),
+                         std::get<1>(current_range),
+                         std::get<2>(current_range));
+    range_values_.pop_front();
+  }
+}
+
+void DexFileTrackingRegistrar::SetDexFileRegistration(bool should_poison) {
+  const void* dex_file_begin = reinterpret_cast<const void*>(dex_file_->Begin());
+  size_t dex_file_size = dex_file_->Size();
+  range_values_.push_back(std::make_tuple(dex_file_begin, dex_file_size, should_poison));
+}
+
+void DexFileTrackingRegistrar::SetAllCodeItemRegistration(bool should_poison) {
+  for (size_t classdef_ctr = 0; classdef_ctr < dex_file_->NumClassDefs(); ++classdef_ctr) {
+    const DexFile::ClassDef& cd = dex_file_->GetClassDef(classdef_ctr);
+    const uint8_t* class_data = dex_file_->GetClassData(cd);
+    if (class_data != nullptr) {
+      ClassDataItemIterator cdit(*dex_file_, class_data);
+      // Skipping Fields
+      while (cdit.HasNextStaticField() || cdit.HasNextInstanceField()) {
+        cdit.Next();
+      }
+      while (cdit.HasNextDirectMethod()) {
+        const DexFile::CodeItem* code_item = cdit.GetMethodCodeItem();
+        if (code_item != nullptr) {
+          const void* code_item_begin = reinterpret_cast<const void*>(code_item);
+          size_t code_item_size = DexFile::GetCodeItemSize(*code_item);
+          range_values_.push_back(std::make_tuple(code_item_begin, code_item_size, should_poison));
+        }
+        cdit.Next();
+      }
+    }
+  }
+}
+
+void DexFileTrackingRegistrar::SetAllInsnsRegistration(bool should_poison) {
+  for (size_t classdef_ctr = 0; classdef_ctr < dex_file_->NumClassDefs(); ++classdef_ctr) {
+    const DexFile::ClassDef& cd = dex_file_->GetClassDef(classdef_ctr);
+    const uint8_t* class_data = dex_file_->GetClassData(cd);
+    if (class_data != nullptr) {
+      ClassDataItemIterator cdit(*dex_file_, class_data);
+      // Skipping Fields
+      while (cdit.HasNextStaticField() || cdit.HasNextInstanceField()) {
+        cdit.Next();
+      }
+      while (cdit.HasNextDirectMethod()) {
+        const DexFile::CodeItem* code_item = cdit.GetMethodCodeItem();
+        if (code_item != nullptr) {
+          const void* insns_begin = reinterpret_cast<const void*>(&code_item->insns_);
+          // Member insns_size_in_code_units_ is in 2-byte units
+          size_t insns_size = code_item->insns_size_in_code_units_ * 2;
+          range_values_.push_back(std::make_tuple(insns_begin, insns_size, should_poison));
+        }
+        cdit.Next();
+      }
+    }
+  }
+}
+
+void DexFileTrackingRegistrar::SetCodeItemRegistration(const char* class_name, bool should_poison) {
+  for (size_t classdef_ctr = 0; classdef_ctr < dex_file_->NumClassDefs(); ++classdef_ctr) {
+    const DexFile::ClassDef& cd = dex_file_->GetClassDef(classdef_ctr);
+    const uint8_t* class_data = dex_file_->GetClassData(cd);
+    if (class_data != nullptr) {
+      ClassDataItemIterator cdit(*dex_file_, class_data);
+      // Skipping Fields
+      while (cdit.HasNextStaticField() || cdit.HasNextInstanceField()) {
+        cdit.Next();
+      }
+      while (cdit.HasNextDirectMethod()) {
+        const DexFile::MethodId& methodid_item = dex_file_->GetMethodId(cdit.GetMemberIndex());
+        const char * methodid_name = dex_file_->GetMethodName(methodid_item);
+        const DexFile::CodeItem* code_item = cdit.GetMethodCodeItem();
+        if (code_item != nullptr && strcmp(methodid_name, class_name) == 0) {
+          const void* code_item_begin = reinterpret_cast<const void*>(code_item);
+          size_t code_item_size = DexFile::GetCodeItemSize(*code_item);
+          range_values_.push_back(
+              std::make_tuple(code_item_begin, code_item_size, should_poison));
+        }
+        cdit.Next();
+      }
+    }
   }
 }
 
diff --git a/runtime/dex_file_tracking_registrar.h b/runtime/dex_file_tracking_registrar.h
index 7d5d78d..b0fa275 100644
--- a/runtime/dex_file_tracking_registrar.h
+++ b/runtime/dex_file_tracking_registrar.h
@@ -17,13 +17,53 @@
 #ifndef ART_RUNTIME_DEX_FILE_TRACKING_REGISTRAR_H_
 #define ART_RUNTIME_DEX_FILE_TRACKING_REGISTRAR_H_
 
+#include <deque>
+#include <tuple>
+
 #include "dex_file.h"
 
 namespace art {
 namespace dex {
 namespace tracking {
 
-void RegisterDexFile(const DexFile* const dex_file);
+// Class for (un)poisoning various sections of Dex Files
+//
+// This class provides the means to log accesses only of sections whose
+// accesses are needed. All accesses are displayed as stack traces in
+// logcat.
+class DexFileTrackingRegistrar {
+ public:
+  explicit DexFileTrackingRegistrar(const DexFile* const dex_file)
+      : dex_file_(dex_file) {
+  }
+
+  // This function is where the functions below it are called to actually
+  // poison sections.
+  void SetDexSections();
+
+  // Uses data contained inside range_values_ to poison memory through the
+  // memory tool.
+  void SetCurrentRanges();
+
+ private:
+  void SetDexFileRegistration(bool should_poison);
+
+  // Set of functions concerning Code Items of dex_file_
+  void SetAllCodeItemRegistration(bool should_poison);
+  // Sets the insns_ section of all code items.
+  void SetAllInsnsRegistration(bool should_poison);
+  // This function finds the code item of a class based on class name.
+  void SetCodeItemRegistration(const char* class_name, bool should_poison);
+
+  // Contains tuples of all ranges of memory that need to be explicitly
+  // (un)poisoned by the memory tool.
+  std::deque<std::tuple<const void *, size_t, bool>> range_values_;
+
+  const DexFile* const dex_file_;
+};
+
+// This function is meant to called externally to use DexfileTrackingRegistrar
+void RegisterDexFile(const DexFile* dex_file);
 
 }  // namespace tracking
 }  // namespace dex