[lldb] Enable locate module callback in GetSharedModule

Author: George Hu

lldb/include/lldb/Core/ModuleList.h

@@ -477,7 +477,8 @@ class ModuleList {
   static Status
   GetSharedModule(const ModuleSpec &module_spec, lldb::ModuleSP &module_sp,
                   llvm::SmallVectorImpl<lldb::ModuleSP> *old_modules,
-                  bool *did_create_ptr, bool always_create = false);
+                  bool *did_create_ptr, bool always_create = false,
+                  bool allow_locate_callback = true);
 
   static bool RemoveSharedModule(lldb::ModuleSP &module_sp);
 

lldb/include/lldb/Core/ModuleSpec.h

@@ -19,6 +19,7 @@
 
 #include "llvm/Support/Chrono.h"
 
+#include <memory>
 #include <mutex>
 #include <vector>
 
@@ -126,11 +127,17 @@ class ModuleSpec {
 
   lldb::DataBufferSP GetData() const { return m_data; }
 
-  Target *GetTargetPtr() { return m_target; }
+  Target *GetTargetPtr() {
+    auto locked = m_target.lock();
+    return locked.get();
+  }
 
-  const Target *GetTargetPtr() const { return m_target; }
+  const Target *GetTargetPtr() const {
+    auto locked = m_target.lock();
+    return locked.get();
+  }
 
-  void SetTarget(Target *target) { m_target = target; }
+  void SetTarget(std::shared_ptr<Target> target) { m_target = target; }
 
   void Clear() {
     m_file.Clear();
@@ -143,7 +150,7 @@ class ModuleSpec {
     m_object_size = 0;
     m_source_mappings.Clear(false);
     m_object_mod_time = llvm::sys::TimePoint<>();
-    m_target = nullptr;
+    m_target.reset();
   }
 
   explicit operator bool() const {
@@ -272,8 +279,9 @@ class ModuleSpec {
   ArchSpec m_arch;
   UUID m_uuid;
   ConstString m_object_name;
-  Target *m_target; // This is set to take advantage of the target's search path
-                    // and platform's locate module callback
+  std::weak_ptr<Target>
+      m_target; // This is set to take advantage of the target's search path
+                // and platform's locate module callback
   uint64_t m_object_offset = 0;
   uint64_t m_object_size = 0;
   llvm::sys::TimePoint<> m_object_mod_time;

lldb/source/Core/DynamicLoader.cpp

@@ -227,7 +227,7 @@ ModuleSP DynamicLoader::LoadBinaryWithUUIDAndAddress(
     }
   }
   ModuleSpec module_spec;
-  module_spec.SetTarget(&target);
+  module_spec.SetTarget(target.shared_from_this());
   module_spec.GetUUID() = uuid;
   FileSpec name_filespec(name);
   if (FileSystem::Instance().Exists(name_filespec))

lldb/source/Core/ModuleList.cpp

@@ -19,6 +19,7 @@
 #include "lldb/Symbol/SymbolContext.h"
 #include "lldb/Symbol/TypeList.h"
 #include "lldb/Symbol/VariableList.h"
+#include "lldb/Target/Platform.h"
 #include "lldb/Target/Target.h"
 #include "lldb/Utility/ArchSpec.h"
 #include "lldb/Utility/ConstString.h"
@@ -1031,7 +1032,8 @@ size_t ModuleList::RemoveOrphanSharedModules(bool mandatory) {
 Status
 ModuleList::GetSharedModule(const ModuleSpec &module_spec, ModuleSP &module_sp,
                             llvm::SmallVectorImpl<lldb::ModuleSP> *old_modules,
-                            bool *did_create_ptr, bool always_create) {
+                            bool *did_create_ptr, bool always_create,
+                            bool allow_locate_callback) {
   ModuleList &shared_module_list = GetSharedModuleList();
   std::lock_guard<std::recursive_mutex> guard(
       shared_module_list.m_modules_mutex);
@@ -1087,6 +1089,24 @@ ModuleList::GetSharedModule(const ModuleSpec &module_spec, ModuleSP &module_sp,
   if (module_sp)
     return error;
 
+  // Try target's platform locate module callback before second attempt
+  if (allow_locate_callback) {
+    ModuleSpec module_spec_copy(module_spec);
+    Target *target = module_spec_copy.GetTargetPtr();
+    if (target && target->IsValid()) {
+      Platform *platform = target->GetPlatform().get();
+      if (platform) {
+        FileSpec symbol_file_spec;
+        platform->CallLocateModuleCallbackIfSet(
+            module_spec, module_sp, symbol_file_spec, did_create_ptr);
+        if (module_sp) {
+          // Success! The callback found a module
+          return error;
+        }
+      }
+    }
+  }
+
   module_sp = std::make_shared<Module>(module_spec);
   // Make sure there are a module and an object file since we can specify a
   // valid file path with an architecture that might not be in that file. By
@@ -1119,7 +1139,7 @@ ModuleList::GetSharedModule(const ModuleSpec &module_spec, ModuleSP &module_sp,
   Target *target = module_spec_copy.GetTargetPtr();
   FileSpecList module_search_paths;
   FileSpecList *module_search_paths_ptr = nullptr;
-  if (target) {
+  if (target && target->IsValid()) {
     module_search_paths = target->GetExecutableSearchPaths();
     module_search_paths_ptr = &module_search_paths;
   }

lldb/source/Plugins/DynamicLoader/Darwin-Kernel/DynamicLoaderDarwinKernel.cpp

@@ -789,7 +789,7 @@ bool DynamicLoaderDarwinKernel::KextImageInfo::LoadImageUsingMemoryModule(
     // Search for the kext on the local filesystem via the UUID
     if (!m_module_sp && m_uuid.IsValid()) {
       ModuleSpec module_spec;
-      module_spec.SetTarget(&target);
+      module_spec.SetTarget(target.shared_from_this());
       module_spec.GetUUID() = m_uuid;
       if (!m_uuid.IsValid())
         module_spec.GetArchitecture() = target.GetArchitecture();

lldb/source/Plugins/DynamicLoader/POSIX-DYLD/DynamicLoaderPOSIXDYLD.cpp

@@ -901,7 +901,7 @@ void DynamicLoaderPOSIXDYLD::ResolveExecutableModule(
   if (module_sp && module_sp->MatchesModuleSpec(module_spec))
     return;
 
-  module_spec.SetTarget(&target);
+  module_spec.SetTarget(target.shared_from_this());
   const auto executable_search_paths(Target::GetDefaultExecutableSearchPaths());
   auto error = platform_sp->ResolveExecutable(module_spec, module_sp);
   if (error.Fail()) {

lldb/source/Plugins/Process/elf-core/ProcessElfCore.cpp

@@ -84,7 +84,7 @@ bool ProcessElfCore::CanDebug(lldb::TargetSP target_sp,
   // For now we are just making sure the file exists for a given module
   if (!m_core_module_sp && FileSystem::Instance().Exists(m_core_file)) {
     ModuleSpec core_module_spec(m_core_file, target_sp->GetArchitecture());
-    core_module_spec.SetTarget(target_sp.get());
+    core_module_spec.SetTarget(target_sp);
     Status error(ModuleList::GetSharedModule(core_module_spec, m_core_module_sp,
                                              nullptr, nullptr));
     if (m_core_module_sp) {

lldb/source/Plugins/Process/mach-core/ProcessMachCore.cpp

@@ -95,7 +95,7 @@ bool ProcessMachCore::CanDebug(lldb::TargetSP target_sp,
     // header but we should still try to use it -
     // ModuleSpecList::FindMatchingModuleSpec enforces a strict arch mach.
     ModuleSpec core_module_spec(m_core_file);
-    core_module_spec.SetTarget(target_sp.get());
+    core_module_spec.SetTarget(target_sp);
     Status error(ModuleList::GetSharedModule(core_module_spec, m_core_module_sp,
                                              nullptr, nullptr));
 

lldb/source/Target/Platform.cpp

@@ -747,9 +747,8 @@ Status Platform::ResolveExecutable(const ModuleSpec &module_spec,
 
   if (resolved_module_spec.GetArchitecture().IsValid() ||
       resolved_module_spec.GetUUID().IsValid()) {
-    Status error =
-        ModuleList::GetSharedModule(resolved_module_spec, exe_module_sp,
-                                    nullptr, nullptr);
+    Status error = ModuleList::GetSharedModule(resolved_module_spec,
+                                               exe_module_sp, nullptr, nullptr);
 
     if (exe_module_sp && exe_module_sp->GetObjectFile())
       return error;
@@ -775,9 +774,8 @@ Status Platform::ResolveExecutable(const ModuleSpec &module_spec,
                                           nullptr, nullptr);
     }
 
-    error =
-        ModuleList::GetSharedModule(resolved_module_spec, exe_module_sp,
-                                    nullptr, nullptr);
+    error = ModuleList::GetSharedModule(resolved_module_spec, exe_module_sp,
+                                        nullptr, nullptr);
     if (error.Success()) {
       if (exe_module_sp && exe_module_sp->GetObjectFile())
         break;
@@ -1678,11 +1676,12 @@ void Platform::CallLocateModuleCallbackIfSet(const ModuleSpec &module_spec,
   cached_module_spec.GetUUID().Clear(); // Clear UUID since it may contain md5
                                         // content hash instead of real UUID.
   cached_module_spec.GetFileSpec() = module_file_spec;
+  cached_module_spec.GetSymbolFileSpec() = symbol_file_spec;
   cached_module_spec.GetPlatformFileSpec() = module_spec.GetFileSpec();
   cached_module_spec.SetObjectOffset(0);
 
   error = ModuleList::GetSharedModule(cached_module_spec, module_sp, nullptr,
-                                      did_create_ptr, false);
+                                      did_create_ptr, false, false);
   if (error.Success() && module_sp) {
     // Succeeded to load the module file.
     LLDB_LOGF(log, "%s: locate module callback succeeded: module=%s symbol=%s",

lldb/source/Target/Target.cpp

@@ -1773,7 +1773,7 @@ bool Target::SetArchitecture(const ArchSpec &arch_spec, bool set_platform,
               arch_spec.GetArchitectureName(),
               arch_spec.GetTriple().getTriple().c_str());
     ModuleSpec module_spec(executable_sp->GetFileSpec(), other);
-    module_spec.SetTarget(this);
+    module_spec.SetTarget(this->shared_from_this());
     Status error = ModuleList::GetSharedModule(module_spec, executable_sp,
                                                nullptr, nullptr);
 
@@ -2344,7 +2344,7 @@ ModuleSP Target::GetOrCreateModule(const ModuleSpec &orig_module_spec,
 
   // Apply any remappings specified in target.object-map:
   ModuleSpec module_spec(orig_module_spec);
-  module_spec.SetTarget(this);
+  module_spec.SetTarget(this->shared_from_this());
   PathMappingList &obj_mapping = GetObjectPathMap();
   if (std::optional<FileSpec> remapped_obj_file =
           obj_mapping.RemapPath(orig_module_spec.GetFileSpec().GetPath(),
@@ -2403,7 +2403,7 @@ ModuleSP Target::GetOrCreateModule(const ModuleSpec &orig_module_spec,
           transformed_spec.GetFileSpec().SetDirectory(transformed_dir);
           transformed_spec.GetFileSpec().SetFilename(
                 module_spec.GetFileSpec().GetFilename());
-          transformed_spec.SetTarget(this);
+          transformed_spec.SetTarget(this->shared_from_this());
           error = ModuleList::GetSharedModule(transformed_spec, module_sp,
                                               &old_modules, &did_create_module);
         }