[lldb] Fix use after free on ModuleList::RemoveSharedModuleIfOrphaned (llvm#155331)

This fixes a potential use after free where
ModuleList::RemoveSharedModuleIfOrphaned ->
SharedModuleList::RemoveIfOrphaned -> SharedModuleList::RemoveFromMap
would potentially dereference a freed pointer. This fixes it by not
calling ModuleList::RemoveSharedModuleIfOrphaned at all if the pointer
was just freed.

(cherry picked from commit 397181d)

Author: augusto2112

lldb/include/lldb/Core/ModuleList.h

@@ -477,7 +477,7 @@ class ModuleList {
 
   size_t Remove(ModuleList &module_list);
 
-  bool RemoveIfOrphaned(const Module *module_ptr);
+  bool RemoveIfOrphaned(const lldb::ModuleWP module_ptr);
 
   size_t RemoveOrphans(bool mandatory);
 
@@ -531,7 +531,7 @@ class ModuleList {
 
   static size_t RemoveOrphanSharedModules(bool mandatory);
 
-  static bool RemoveSharedModuleIfOrphaned(const Module *module_ptr);
+  static bool RemoveSharedModuleIfOrphaned(const lldb::ModuleWP module_ptr);
 
   /// Applies 'callback' to each module in this ModuleList.
   /// If 'callback' returns false, iteration terminates.
@@ -575,6 +575,9 @@ class ModuleList {
 
   Notifier *m_notifier = nullptr;
 
+  /// An orphaned module that lives only in the ModuleList has a count of 1.
+  static constexpr long kUseCountModuleListOrphaned = 1;
+
 public:
   typedef LockingAdaptedIterable<std::recursive_mutex, collection>
       ModuleIterable;

lldb/source/Core/ModuleList.cpp

@@ -503,17 +503,20 @@ bool ModuleList::ReplaceModule(const lldb::ModuleSP &old_module_sp,
   return true;
 }
 
-bool ModuleList::RemoveIfOrphaned(const Module *module_ptr) {
-  if (module_ptr) {
+bool ModuleList::RemoveIfOrphaned(const ModuleWP module_wp) {
+  if (auto module_sp = module_wp.lock()) {
     std::lock_guard<std::recursive_mutex> guard(m_modules_mutex);
     collection::iterator pos, end = m_modules.end();
     for (pos = m_modules.begin(); pos != end; ++pos) {
-      if (pos->get() == module_ptr) {
-        if (pos->use_count() == 1) {
+      if (pos->get() == module_sp.get()) {
+        // Since module_sp increases the refcount by 1, the use count should be
+        // the regular use count + 1.
+        constexpr long kUseCountOrphaned = kUseCountModuleListOrphaned + 1;
+        if (pos->use_count() == kUseCountOrphaned) {
           pos = RemoveImpl(pos);
           return true;
-        } else
-          return false;
+        }
+        return false;
       }
     }
   }
@@ -540,7 +543,7 @@ size_t ModuleList::RemoveOrphans(bool mandatory) {
     made_progress = false;
     collection::iterator pos = m_modules.begin();
     while (pos != m_modules.end()) {
-      if (pos->use_count() == 1) {
+      if (pos->use_count() == kUseCountModuleListOrphaned) {
         pos = RemoveImpl(pos);
         ++remove_count;
         // We did make progress.
@@ -1094,7 +1097,7 @@ class SharedModuleList {
     if (!module_sp)
       return false;
     std::lock_guard<std::recursive_mutex> guard(GetMutex());
-    RemoveFromMap(*module_sp.get());
+    RemoveFromMap(module_sp);
     return m_list.Remove(module_sp, use_notifier);
   }
 
@@ -1105,10 +1108,10 @@ class SharedModuleList {
     ReplaceEquivalentInMap(module_sp);
   }
 
-  bool RemoveIfOrphaned(const Module *module_ptr) {
+  bool RemoveIfOrphaned(const ModuleWP module_wp) {
     std::lock_guard<std::recursive_mutex> guard(GetMutex());
-    RemoveFromMap(*module_ptr, /*if_orphaned=*/true);
-    return m_list.RemoveIfOrphaned(module_ptr);
+    RemoveFromMap(module_wp, /*if_orphaned=*/true);
+    return m_list.RemoveIfOrphaned(module_wp);
   }
 
   std::recursive_mutex &GetMutex() const { return m_list.GetMutex(); }
@@ -1148,16 +1151,22 @@ class SharedModuleList {
     m_name_to_modules[name].push_back(module_sp);
   }
 
-  void RemoveFromMap(const Module &module, bool if_orphaned = false) {
-    ConstString name = module.GetFileSpec().GetFilename();
-    if (!m_name_to_modules.contains(name))
-      return;
-    llvm::SmallVectorImpl<ModuleSP> &vec = m_name_to_modules[name];
-    for (auto *it = vec.begin(); it != vec.end(); ++it) {
-      if (it->get() == &module) {
-        if (!if_orphaned || it->use_count() == kUseCountOrphaned) {
-          vec.erase(it);
-          break;
+  void RemoveFromMap(const ModuleWP module_wp, bool if_orphaned = false) {
+    if (auto module_sp = module_wp.lock()) {
+      ConstString name = module_sp->GetFileSpec().GetFilename();
+      if (!m_name_to_modules.contains(name))
+        return;
+      llvm::SmallVectorImpl<ModuleSP> &vec = m_name_to_modules[name];
+      for (auto *it = vec.begin(); it != vec.end(); ++it) {
+        if (it->get() == module_sp.get()) {
+          // Since module_sp increases the refcount by 1, the use count should
+          // be the regular use count + 1.
+          constexpr long kUseCountOrphaned =
+              kUseCountSharedModuleListOrphaned + 1;
+          if (!if_orphaned || it->use_count() == kUseCountOrphaned) {
+            vec.erase(it);
+            break;
+          }
         }
       }
     }
@@ -1195,7 +1204,7 @@ class SharedModuleList {
     // remove_if moves the elements that match the condition to the end of the
     // container, and returns an iterator to the first element that was moved.
     auto *to_remove_start = llvm::remove_if(vec, [](const ModuleSP &module) {
-      return module.use_count() == kUseCountOrphaned;
+      return module.use_count() == kUseCountSharedModuleListOrphaned;
     });
 
     ModuleList to_remove;
@@ -1238,7 +1247,7 @@ class SharedModuleList {
   llvm::DenseMap<ConstString, llvm::SmallVector<ModuleSP, 1>> m_name_to_modules;
 
   /// The use count of a module held only by m_list and m_name_to_modules.
-  static constexpr long kUseCountOrphaned = 2;
+  static constexpr long kUseCountSharedModuleListOrphaned = 2;
 };
 
 struct SharedModuleListInfo {
@@ -1540,8 +1549,8 @@ bool ModuleList::RemoveSharedModule(lldb::ModuleSP &module_sp) {
   return GetSharedModuleList().Remove(module_sp);
 }
 
-bool ModuleList::RemoveSharedModuleIfOrphaned(const Module *module_ptr) {
-  return GetSharedModuleList().RemoveIfOrphaned(module_ptr);
+bool ModuleList::RemoveSharedModuleIfOrphaned(const ModuleWP module_wp) {
+  return GetSharedModuleList().RemoveIfOrphaned(module_wp);
 }
 
 bool ModuleList::LoadScriptingResourcesInTarget(Target *target,

lldb/source/Target/Target.cpp

@@ -2596,9 +2596,9 @@ ModuleSP Target::GetOrCreateModule(const ModuleSpec &orig_module_spec,
           m_images.Append(module_sp, notify);
 
         for (ModuleSP &old_module_sp : replaced_modules) {
-          Module *old_module_ptr = old_module_sp.get();
+          auto old_module_wp = old_module_sp->weak_from_this();
           old_module_sp.reset();
-          ModuleList::RemoveSharedModuleIfOrphaned(old_module_ptr);
+          ModuleList::RemoveSharedModuleIfOrphaned(old_module_wp);
         }
       } else
         module_sp.reset();