[veridex] Resolve all type_id/method_id/field_id in app dex files.

Change-Id: I6e28d70b6d3807944ccebc0f429c1b9b1d4b1eeb
diff --git a/tools/veridex/resolver.cc b/tools/veridex/resolver.cc
index c0705e5..8297821 100644
--- a/tools/veridex/resolver.cc
+++ b/tools/veridex/resolver.cc
@@ -55,4 +55,245 @@
   }
 }
 
+static bool HasSameNameAndSignature(const DexFile& dex_file,
+                                    const DexFile::MethodId& method_id,
+                                    const char* method_name,
+                                    const Signature& signature) {
+  return strcmp(method_name, dex_file.GetMethodName(method_id)) == 0 &&
+      dex_file.GetMethodSignature(method_id) == signature;
+}
+
+static bool HasSameNameAndType(const DexFile& dex_file,
+                               const DexFile::FieldId& field_id,
+                               const char* field_name,
+                               const char* field_type) {
+  return strcmp(field_name, dex_file.GetFieldName(field_id)) == 0 &&
+      strcmp(field_type, dex_file.GetFieldTypeDescriptor(field_id)) == 0;
+}
+
+VeriClass* VeridexResolver::GetVeriClass(dex::TypeIndex index) {
+  CHECK_LT(index.index_, dex_file_.NumTypeIds());
+  // Lookup in our local cache.
+  VeriClass* cls = &type_infos_[index.index_];
+  if (cls->IsUninitialized()) {
+    // Class is defined in another dex file. Lookup in the global cache.
+    std::string name(dex_file_.StringByTypeIdx(index));
+    auto existing = type_map_.find(name);
+    if (existing == type_map_.end()) {
+      // Class hasn't been defined, so check if it's an array class.
+      size_t last_array = name.find_last_of('[');
+      if (last_array == std::string::npos) {
+        // There is no such class.
+        return nullptr;
+      } else {
+        // Class is an array class. Check if its most enclosed component type (which is not
+        // an array class) has been defined.
+        std::string klass_name = name.substr(last_array + 1);
+        existing = type_map_.find(klass_name);
+        if (existing == type_map_.end()) {
+          // There is no such class, so there is no such array.
+          return nullptr;
+        } else {
+          // Create the type, and cache it locally and globally.
+          type_infos_[index.index_] = VeriClass(
+              existing->second->GetKind(), last_array + 1, existing->second->GetClassDef());
+          cls = &(type_infos_[index.index_]);
+          type_map_[name] = cls;
+        }
+      }
+    } else {
+      // Cache the found class.
+      cls = existing->second;
+      type_infos_[index.index_] = *cls;
+    }
+  }
+  return cls;
+}
+
+VeridexResolver* VeridexResolver::GetResolverOf(const VeriClass& kls) const {
+  auto resolver_it = dex_resolvers_.lower_bound(reinterpret_cast<uintptr_t>(kls.GetClassDef()));
+  --resolver_it;
+
+  // Check the class def pointer is indeed in the mapped dex file range.
+  const DexFile& dex_file = resolver_it->second->dex_file_;
+  CHECK_LT(reinterpret_cast<uintptr_t>(dex_file.Begin()),
+           reinterpret_cast<uintptr_t>(kls.GetClassDef()));
+  CHECK_GT(reinterpret_cast<uintptr_t>(dex_file.Begin()) + dex_file.Size(),
+           reinterpret_cast<uintptr_t>(kls.GetClassDef()));
+  return resolver_it->second;
+}
+
+VeriMethod VeridexResolver::LookupMethodIn(const VeriClass& kls,
+                                           const char* method_name,
+                                           const Signature& method_signature) {
+  if (kls.IsPrimitive()) {
+    // Primitive classes don't have methods.
+    return nullptr;
+  }
+  if (kls.IsArray()) {
+    // Array classes don't have methods, but inherit the ones in j.l.Object.
+    return LookupMethodIn(*VeriClass::object_, method_name, method_signature);
+  }
+  // Get the resolver where `kls` is from.
+  VeridexResolver* resolver = GetResolverOf(kls);
+
+  // Look at methods declared in `kls`.
+  const DexFile& other_dex_file = resolver->dex_file_;
+  const uint8_t* class_data = other_dex_file.GetClassData(*kls.GetClassDef());
+  if (class_data != nullptr) {
+    ClassDataItemIterator it(other_dex_file, class_data);
+    it.SkipAllFields();
+    for (; it.HasNextMethod(); it.Next()) {
+      const DexFile::MethodId& other_method_id = other_dex_file.GetMethodId(it.GetMemberIndex());
+      if (HasSameNameAndSignature(other_dex_file,
+                                  other_method_id,
+                                  method_name,
+                                  method_signature)) {
+        return it.DataPointer();
+      }
+    }
+  }
+
+  // Look at methods in `kls`'s super class hierarchy.
+  if (kls.GetClassDef()->superclass_idx_.IsValid()) {
+    VeriClass* super = resolver->GetVeriClass(kls.GetClassDef()->superclass_idx_);
+    if (super != nullptr) {
+      VeriMethod super_method = resolver->LookupMethodIn(*super, method_name, method_signature);
+      if (super_method != nullptr) {
+        return super_method;
+      }
+    }
+  }
+
+  // Look at methods in `kls`'s interface hierarchy.
+  const DexFile::TypeList* interfaces = other_dex_file.GetInterfacesList(*kls.GetClassDef());
+  if (interfaces != nullptr) {
+    for (size_t i = 0; i < interfaces->Size(); i++) {
+      dex::TypeIndex idx = interfaces->GetTypeItem(i).type_idx_;
+      VeriClass* itf = resolver->GetVeriClass(idx);
+      if (itf != nullptr) {
+        VeriMethod itf_method = resolver->LookupMethodIn(*itf, method_name, method_signature);
+        if (itf_method != nullptr) {
+          return itf_method;
+        }
+      }
+    }
+  }
+  return nullptr;
+}
+
+VeriField VeridexResolver::LookupFieldIn(const VeriClass& kls,
+                                         const char* field_name,
+                                         const char* field_type) {
+  if (kls.IsPrimitive()) {
+    // Primitive classes don't have fields.
+    return nullptr;
+  }
+  if (kls.IsArray()) {
+    // Array classes don't have fields.
+    return nullptr;
+  }
+  // Get the resolver where `kls` is from.
+  VeridexResolver* resolver = GetResolverOf(kls);
+
+  // Look at fields declared in `kls`.
+  const DexFile& other_dex_file = resolver->dex_file_;
+  const uint8_t* class_data = other_dex_file.GetClassData(*kls.GetClassDef());
+  if (class_data != nullptr) {
+    ClassDataItemIterator it(other_dex_file, class_data);
+    for (; it.HasNextStaticField() || it.HasNextInstanceField(); it.Next()) {
+      const DexFile::FieldId& other_field_id = other_dex_file.GetFieldId(it.GetMemberIndex());
+      if (HasSameNameAndType(other_dex_file,
+                             other_field_id,
+                             field_name,
+                             field_type)) {
+        return it.DataPointer();
+      }
+    }
+  }
+
+  // Look at fields in `kls`'s interface hierarchy.
+  const DexFile::TypeList* interfaces = other_dex_file.GetInterfacesList(*kls.GetClassDef());
+  if (interfaces != nullptr) {
+    for (size_t i = 0; i < interfaces->Size(); i++) {
+      dex::TypeIndex idx = interfaces->GetTypeItem(i).type_idx_;
+      VeriClass* itf = resolver->GetVeriClass(idx);
+      if (itf != nullptr) {
+        VeriField itf_field = resolver->LookupFieldIn(*itf, field_name, field_type);
+        if (itf_field != nullptr) {
+          return itf_field;
+        }
+      }
+    }
+  }
+
+  // Look at fields in `kls`'s super class hierarchy.
+  if (kls.GetClassDef()->superclass_idx_.IsValid()) {
+    VeriClass* super = resolver->GetVeriClass(kls.GetClassDef()->superclass_idx_);
+    if (super != nullptr) {
+      VeriField super_field = resolver->LookupFieldIn(*super, field_name, field_type);
+      if (super_field != nullptr) {
+        return super_field;
+      }
+    }
+  }
+  return nullptr;
+}
+
+VeriMethod VeridexResolver::GetMethod(uint32_t method_index) {
+  VeriMethod method_info = method_infos_[method_index];
+  if (method_info == nullptr) {
+    // Method is defined in another dex file.
+    const DexFile::MethodId& method_id = dex_file_.GetMethodId(method_index);
+    VeriClass* kls = GetVeriClass(method_id.class_idx_);
+    if (kls == nullptr) {
+      return nullptr;
+    }
+    // Class found, now lookup the method in it.
+    method_info = LookupMethodIn(*kls,
+                                 dex_file_.GetMethodName(method_id),
+                                 dex_file_.GetMethodSignature(method_id));
+    method_infos_[method_index] = method_info;
+  }
+  return method_info;
+}
+
+VeriField VeridexResolver::GetField(uint32_t field_index) {
+  VeriField field_info = field_infos_[field_index];
+  if (field_info == nullptr) {
+    // Field is defined in another dex file.
+    const DexFile::FieldId& field_id = dex_file_.GetFieldId(field_index);
+    VeriClass* kls = GetVeriClass(field_id.class_idx_);
+    if (kls == nullptr) {
+      return nullptr;
+    }
+    // Class found, now lookup the field in it.
+    field_info = LookupFieldIn(*kls,
+                               dex_file_.GetFieldName(field_id),
+                               dex_file_.GetFieldTypeDescriptor(field_id));
+    field_infos_[field_index] = field_info;
+  }
+  return field_info;
+}
+
+void VeridexResolver::ResolveAll() {
+  for (uint32_t i = 0; i < dex_file_.NumTypeIds(); ++i) {
+    if (GetVeriClass(dex::TypeIndex(i)) == nullptr) {
+      LOG(WARNING) << "Unresolved " << dex_file_.PrettyType(dex::TypeIndex(i));
+    }
+  }
+
+  for (uint32_t i = 0; i < dex_file_.NumMethodIds(); ++i) {
+    if (GetMethod(i) == nullptr) {
+      LOG(WARNING) << "Unresolved: " << dex_file_.PrettyMethod(i);
+    }
+  }
+
+  for (uint32_t i = 0; i < dex_file_.NumFieldIds(); ++i) {
+    if (GetField(i) == nullptr) {
+      LOG(WARNING) << "Unresolved: " << dex_file_.PrettyField(i);
+    }
+  }
+}
+
 }  // namespace art
diff --git a/tools/veridex/resolver.h b/tools/veridex/resolver.h
index 4e0c5b3..ae94dad 100644
--- a/tools/veridex/resolver.h
+++ b/tools/veridex/resolver.h
@@ -22,20 +22,59 @@
 
 namespace art {
 
+class VeridexResolver;
+
+/**
+ * Map from the start of a dex file (ie DexFile::Begin()), to
+ * its corresponding resolver.
+ */
+using DexResolverMap = std::map<uintptr_t, VeridexResolver*>;
+
 class VeridexResolver {
  public:
-  VeridexResolver(const DexFile& dex_file, TypeMap& type_map)
+  VeridexResolver(const DexFile& dex_file,
+                  const DexResolverMap& dex_resolvers,
+                  TypeMap& type_map)
       : dex_file_(dex_file),
         type_map_(type_map),
+        dex_resolvers_(dex_resolvers),
         type_infos_(dex_file.NumTypeIds(), VeriClass()),
         method_infos_(dex_file.NumMethodIds(), nullptr),
         field_infos_(dex_file.NumFieldIds(), nullptr) {}
 
+  // Run on the defined classes of that dex file and populate our
+  // local type cache.
   void Run();
 
+  // Return the class declared at `index`.
+  VeriClass* GetVeriClass(dex::TypeIndex index);
+
+  // Return the method declared at `method_index`.
+  VeriMethod GetMethod(uint32_t method_index);
+
+  // Return the field declared at `field_index`.
+  VeriField GetField(uint32_t field_index);
+
+  // Do a JLS lookup in `kls` to find a method.
+  VeriMethod LookupMethodIn(const VeriClass& kls,
+                            const char* method_name,
+                            const Signature& method_signature);
+
+  // Do a JLS lookup in `kls` to find a field.
+  VeriField LookupFieldIn(const VeriClass& kls,
+                          const char* field_name,
+                          const char* field_type);
+
+  // Resolve all type_id/method_id/field_id.
+  void ResolveAll();
+
  private:
+  // Return the resolver where `kls` is from.
+  VeridexResolver* GetResolverOf(const VeriClass& kls) const;
+
   const DexFile& dex_file_;
   TypeMap& type_map_;
+  const DexResolverMap& dex_resolvers_;
   std::vector<VeriClass> type_infos_;
   std::vector<VeriMethod> method_infos_;
   std::vector<VeriField> field_infos_;
diff --git a/tools/veridex/veridex.cc b/tools/veridex/veridex.cc
index 0370a03..9287211 100644
--- a/tools/veridex/veridex.cc
+++ b/tools/veridex/veridex.cc
@@ -26,6 +26,28 @@
 
 namespace art {
 
+static VeriClass z_(Primitive::Type::kPrimBoolean, 0, nullptr);
+static VeriClass b_(Primitive::Type::kPrimByte, 0, nullptr);
+static VeriClass c_(Primitive::Type::kPrimChar, 0, nullptr);
+static VeriClass s_(Primitive::Type::kPrimShort, 0, nullptr);
+static VeriClass i_(Primitive::Type::kPrimInt, 0, nullptr);
+static VeriClass f_(Primitive::Type::kPrimFloat, 0, nullptr);
+static VeriClass d_(Primitive::Type::kPrimDouble, 0, nullptr);
+static VeriClass j_(Primitive::Type::kPrimLong, 0, nullptr);
+static VeriClass v_(Primitive::Type::kPrimVoid, 0, nullptr);
+
+VeriClass* VeriClass::boolean_ = &z_;
+VeriClass* VeriClass::byte_ = &b_;
+VeriClass* VeriClass::char_ = &c_;
+VeriClass* VeriClass::short_ = &s_;
+VeriClass* VeriClass::integer_ = &i_;
+VeriClass* VeriClass::float_ = &f_;
+VeriClass* VeriClass::double_ = &d_;
+VeriClass* VeriClass::long_ = &j_;
+VeriClass* VeriClass::void_ = &v_;
+// Will be set after boot classpath has been resolved.
+VeriClass* VeriClass::object_ = nullptr;
+
 struct VeridexOptions {
   const char* dex_file = nullptr;
   const char* core_stubs = nullptr;
@@ -114,14 +136,35 @@
 
     // Resolve classes/methods/fields defined in each dex file.
 
-    // Cache of types we've seen. This is used in case of duplicate classes.
+    // Cache of types we've seen, for quick class name lookups.
     TypeMap type_map;
+    // Add internally defined primitives.
+    type_map["Z"] = VeriClass::boolean_;
+    type_map["B"] = VeriClass::byte_;
+    type_map["S"] = VeriClass::short_;
+    type_map["C"] = VeriClass::char_;
+    type_map["I"] = VeriClass::integer_;
+    type_map["F"] = VeriClass::float_;
+    type_map["D"] = VeriClass::double_;
+    type_map["J"] = VeriClass::long_;
+    type_map["V"] = VeriClass::void_;
 
-    std::vector<VeridexResolver> boot_resolvers;
-    Resolve(boot_dex_files, type_map, &boot_resolvers);
+    // Cache of resolvers, to easily query address in memory to VeridexResolver.
+    DexResolverMap resolver_map;
 
-    std::vector<VeridexResolver> app_resolvers;
-    Resolve(app_dex_files, type_map, &app_resolvers);
+    std::vector<std::unique_ptr<VeridexResolver>> boot_resolvers;
+    Resolve(boot_dex_files, resolver_map, type_map, &boot_resolvers);
+
+    // Now that boot classpath has been resolved, fill j.l.Object.
+    VeriClass::object_ = type_map["Ljava/lang/Object;"];
+
+    std::vector<std::unique_ptr<VeridexResolver>> app_resolvers;
+    Resolve(app_dex_files, resolver_map, type_map, &app_resolvers);
+
+    // Resolve all type_id/method_id/field_id of app dex files.
+    for (const std::unique_ptr<VeridexResolver>& resolver : app_resolvers) {
+      resolver->ResolveAll();
+    }
 
     return 0;
   }
@@ -159,14 +202,18 @@
   }
 
   static void Resolve(const std::vector<std::unique_ptr<const DexFile>>& dex_files,
+                      DexResolverMap& resolver_map,
                       TypeMap& type_map,
-                      std::vector<VeridexResolver>* resolvers) {
+                      std::vector<std::unique_ptr<VeridexResolver>>* resolvers) {
     for (const std::unique_ptr<const DexFile>& dex_file : dex_files) {
-      resolvers->push_back(VeridexResolver(*dex_file.get(), type_map));
+      VeridexResolver* resolver =
+          new VeridexResolver(*dex_file.get(), resolver_map, type_map);
+      resolvers->emplace_back(resolver);
+      resolver_map[reinterpret_cast<uintptr_t>(dex_file->Begin())] = resolver;
     }
 
-    for (VeridexResolver& resolver : *resolvers) {
-      resolver.Run();
+    for (const std::unique_ptr<VeridexResolver>& resolver : *resolvers) {
+      resolver->Run();
     }
   }
 };
diff --git a/tools/veridex/veridex.h b/tools/veridex/veridex.h
index bbff254..0c928ab 100644
--- a/tools/veridex/veridex.h
+++ b/tools/veridex/veridex.h
@@ -47,6 +47,21 @@
     return dimensions_ != 0;
   }
 
+  Primitive::Type GetKind() const { return kind_; }
+  uint8_t GetDimensions() const { return dimensions_; }
+  const DexFile::ClassDef* GetClassDef() const { return class_def_; }
+
+  static VeriClass* object_;
+  static VeriClass* boolean_;
+  static VeriClass* byte_;
+  static VeriClass* char_;
+  static VeriClass* short_;
+  static VeriClass* integer_;
+  static VeriClass* float_;
+  static VeriClass* double_;
+  static VeriClass* long_;
+  static VeriClass* void_;
+
  private:
   Primitive::Type kind_;
   uint8_t dimensions_;
@@ -65,6 +80,9 @@
  */
 using VeriMethod = const uint8_t*;
 
+/**
+ * Map from name to VeriClass to quickly lookup classes.
+ */
 using TypeMap = std::map<std::string, VeriClass*>;
 
 }  // namespace art