Merge pull request swiftlang#9664 from adrian-prantl/71278998

[lldb] Prevent a bunch of potential use-after-free with SwiftASTContext

Author: adrian-prantl

lldb/source/Plugins/ExpressionParser/Swift/SwiftREPL.cpp

@@ -567,7 +567,7 @@ SwiftASTContextForExpressions *SwiftREPL::getSwiftASTContext() {
   // our own copy of the AST and using this separate AST for completion.
   //----------------------------------------------------------------------
   if (m_swift_ast)
-    return m_swift_ast;
+    return m_swift_ast.get();
 
   auto type_system_or_err =
       m_target.GetScratchTypeSystemForLanguage(eLanguageTypeSwift, false);
@@ -585,11 +585,9 @@ SwiftASTContextForExpressions *SwiftREPL::getSwiftASTContext() {
   if (!sc_list.GetSize())
     return nullptr;
 
-  auto *target_swift_ast =
-      llvm::dyn_cast_or_null<SwiftASTContextForExpressions>(
-          swift_ts->GetSwiftASTContext(sc_list[0]));
-  m_swift_ast = target_swift_ast;
-  return m_swift_ast;
+  m_swift_ast = std::static_pointer_cast<SwiftASTContextForExpressions>(
+      swift_ts->GetSwiftASTContext(sc_list[0]));
+  return m_swift_ast.get();
 }
 
 void SwiftREPL::CompleteCode(const std::string &current_code,

lldb/source/Plugins/ExpressionParser/Swift/SwiftREPL.h

@@ -24,6 +24,8 @@ namespace lldb_private {
 
 class IRExecutionUnit;
 class SwiftASTContextForExpressions;
+typedef std::shared_ptr<SwiftASTContextForExpressions>
+    SwiftASTContextForExpressionsSP;
 
 //----------------------------------------------------------------------
 /// @class SwiftREPL SwiftREPL.h "lldb/Expression/SwiftREPL.h"
@@ -78,7 +80,7 @@ class SwiftREPL : public REPL {
                     CompletionRequest &request) override;
 
 private:
-  SwiftASTContextForExpressions *m_swift_ast = nullptr;
+  SwiftASTContextForExpressionsSP m_swift_ast;
   bool m_completion_module_initialized = false;
 };
 }

lldb/source/Plugins/ExpressionParser/Swift/SwiftUserExpression.cpp

@@ -742,9 +742,11 @@ bool SwiftUserExpression::Parse(DiagnosticManager &diagnostic_manager,
                        target->GetExecutableModule());
   else
     sc = frame->GetSymbolContext(lldb::eSymbolContextFunction);
-  auto *swift_ast_ctx = m_swift_scratch_ctx->GetSwiftASTContext(sc);
-  m_swift_ast_ctx =
-      llvm::dyn_cast_or_null<SwiftASTContextForExpressions>(swift_ast_ctx);
+  auto swift_ast_ctx = m_swift_scratch_ctx->GetSwiftASTContext(sc);
+  if (llvm::dyn_cast_or_null<SwiftASTContextForExpressions>(
+          swift_ast_ctx.get()))
+    m_swift_ast_ctx =
+        std::static_pointer_cast<SwiftASTContextForExpressions>(swift_ast_ctx);
 
   if (!m_swift_ast_ctx)
     return error("could not create a Swift AST context");

lldb/source/Plugins/ExpressionParser/Swift/SwiftUserExpression.h

@@ -190,7 +190,7 @@ class SwiftUserExpression : public LLVMUserExpression {
   };
 
   TypeSystemSwiftTypeRefForExpressionsSP m_swift_scratch_ctx;
-  SwiftASTContextForExpressions *m_swift_ast_ctx;
+  SwiftASTContextForExpressionsSP m_swift_ast_ctx;
   PersistentVariableDelegate m_persistent_variable_delegate;
   std::unique_ptr<SwiftExpressionParser> m_parser;
   std::optional<SwiftLanguageRuntime::GenericSignature> m_generic_signature;

lldb/source/Plugins/InstrumentationRuntime/MainThreadChecker/InstrumentationRuntimeMainThreadChecker.cpp

@@ -105,7 +105,7 @@ static std::string TranslateObjCNameToSwiftName(std::string className,
     sc = &swiftFrame->GetSymbolContext(eSymbolContextFunction);
   if (!sc)
     return "";
-  auto *ctx = ts->GetSwiftASTContext(*sc);
+  auto ctx = ts->GetSwiftASTContext(*sc);
   if (!ctx)
     return "";
   swift::ClangImporter *imp = ctx->GetClangImporter();

lldb/source/Plugins/Language/Swift/SwiftLanguage.cpp

@@ -1355,7 +1355,7 @@ std::unique_ptr<Language::TypeScavenger> SwiftLanguage::GetTypeScavenger() {
               auto &sc =
                   frame_sp->GetSymbolContext(lldb::eSymbolContextFunction);
               if (scratch_ctx)
-                if (SwiftASTContext *ast_ctx =
+                if (SwiftASTContextSP ast_ctx =
                         scratch_ctx->GetSwiftASTContext(sc)) {
                   ConstString cs_input{input};
                   Mangled mangled(cs_input);
@@ -1424,7 +1424,7 @@ std::unique_ptr<Language::TypeScavenger> SwiftLanguage::GetTypeScavenger() {
               const SymbolContext &sc =
                   frame_sp->GetSymbolContext(lldb::eSymbolContextFunction);
               if (scratch_ctx)
-                if (SwiftASTContext *ast_ctx =
+                if (SwiftASTContextSP ast_ctx =
                         scratch_ctx->GetSwiftASTContext(sc)) {
                   auto iter = ast_ctx->GetModuleCache().begin(),
                        end = ast_ctx->GetModuleCache().end();

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.cpp

@@ -1200,7 +1200,7 @@ SwiftLanguageRuntime::GetLanguageSpecificData(SymbolContext sc) {
 
   if (auto *ts = llvm::dyn_cast_or_null<TypeSystemSwiftTypeRef>(
           type_system_or_err->get()))
-    if (auto *swift_ast_ctx = ts->GetSwiftASTContextOrNull(sc))
+    if (auto swift_ast_ctx = ts->GetSwiftASTContextOrNull(sc))
       dict_sp->AddBooleanItem("SwiftExplicitModules",
                               swift_ast_ctx->HasExplicitModules());
 
@@ -1229,7 +1229,7 @@ void SwiftLanguageRuntime::FindFunctionPointersInCall(
   auto scratch_ctx = TypeSystemSwiftTypeRefForExpressions::GetForTarget(target);
   if (!scratch_ctx)
     return;
-  SwiftASTContext *swift_ast = scratch_ctx->GetSwiftASTContext(sc);
+  SwiftASTContextSP swift_ast = scratch_ctx->GetSwiftASTContext(sc);
   if (!swift_ast)
     return;
 

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntimeRemoteAST.cpp

@@ -193,7 +193,7 @@ ConstString SwiftLanguageRuntimeImpl::GetDynamicTypeName_ClassRemoteAST(
       in_value.GetTargetSP());
   if (!ts)
     return {};
-  SwiftASTContext *swift_ast_ctx = ts->GetSwiftASTContext(sc);
+  SwiftASTContextSP swift_ast_ctx = ts->GetSwiftASTContext(sc);
   if (!swift_ast_ctx)
     return {};
 
@@ -232,7 +232,7 @@ SwiftLanguageRuntimeImpl::GetDynamicTypeAndAddress_ExistentialRemoteAST(
   if (!stack_frame_sp)
     return {};
   auto &sc = stack_frame_sp->GetSymbolContext(eSymbolContextFunction);
-  SwiftASTContext *swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
+  SwiftASTContextSP swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
   if (!swift_ast_ctx)
     return {};
 
@@ -285,7 +285,7 @@ CompilerType SwiftLanguageRuntimeImpl::BindGenericTypeParametersRemoteAST(
   if (!scratch_ctx)
     return base_type;
 
-  SwiftASTContext *swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
+  SwiftASTContextSP swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
   if (!swift_ast_ctx)
     return base_type;
   base_type = swift_ast_ctx->ImportType(base_type, error);
@@ -407,7 +407,7 @@ CompilerType SwiftLanguageRuntimeImpl::BindGenericTypeParametersRemoteAST(
 
     target_swift_type = target_swift_type.subst(
         [this, &stack_frame,
-         &swift_ast_ctx](swift::SubstitutableType *type) -> swift::Type {
+         swift_ast_ctx](swift::SubstitutableType *type) -> swift::Type {
           StreamString type_name;
           if (!SwiftLanguageRuntime::GetAbstractTypeName(type_name, type))
             return type;
@@ -460,7 +460,7 @@ CompilerType SwiftLanguageRuntimeImpl::MetadataPromise::FulfillTypePromise(
       *error = Status::FromErrorString("couldn't get Swift scratch context");
     return CompilerType();
   }
-  SwiftASTContext *swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
+  SwiftASTContextSP swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
   if (!swift_ast_ctx) {
     if (error)
       *error = Status::FromErrorString("couldn't get Swift scratch context");
@@ -497,7 +497,7 @@ SwiftLanguageRuntimeImpl::GetMetadataPromise(const SymbolContext &sc,
       for_object.GetTargetSP());
   if (!scratch_ctx)
     return nullptr;
-  SwiftASTContext *swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
+  SwiftASTContextSP swift_ast_ctx = scratch_ctx->GetSwiftASTContext(sc);
   if (!swift_ast_ctx)
     return nullptr;
   if (swift_ast_ctx->HasFatalErrors())

lldb/source/Plugins/TypeSystem/Swift/SwiftASTContext.h

@@ -218,7 +218,7 @@ class SwiftASTContext : public TypeSystemSwift {
 
   bool SupportsLanguage(lldb::LanguageType language) override;
 
-  SwiftASTContext *GetSwiftASTContext(const SymbolContext &sc) const override {
+  SwiftASTContextSP GetSwiftASTContext(const SymbolContext &sc) const override {
     return GetTypeSystemSwiftTypeRef().GetSwiftASTContext(sc);
   }
 

lldb/source/Plugins/TypeSystem/Swift/TypeSystemSwift.h

@@ -29,8 +29,17 @@ class Decl;
 namespace lldb_private {
 class TypeSystemClang;
 class SwiftASTContext;
+class SwiftASTContextForExpressions;
+class TypeSystemSwift;
 class TypeSystemSwiftTypeRef;
 class TypeSystemSwiftTypeRefForExpressions;
+typedef std::shared_ptr<TypeSystemSwift> TypeSystemSwiftSP;
+typedef std::shared_ptr<TypeSystemSwiftTypeRefForExpressions>
+    TypeSystemSwiftTypeRefForExpressionsSP;
+typedef std::shared_ptr<SwiftASTContext> SwiftASTContextSP;
+typedef std::shared_ptr<SwiftASTContextForExpressions>
+    SwiftASTContextForExpressionsSP;
+
 /// The implementation of lldb::Type's m_payload field for TypeSystemSwift.
 class TypePayloadSwift {
   /// Layout: bit 1 ... IsFixedValueBuffer.
@@ -116,7 +125,7 @@ class TypeSystemSwift : public TypeSystem {
 
   const std::string &GetDescription() const { return m_description; }
   static LanguageSet GetSupportedLanguagesForTypes();
-  virtual SwiftASTContext *
+  virtual SwiftASTContextSP
   GetSwiftASTContext(const SymbolContext &sc) const = 0;
   virtual TypeSystemSwiftTypeRef &GetTypeSystemSwiftTypeRef() = 0;
   virtual const TypeSystemSwiftTypeRef &GetTypeSystemSwiftTypeRef() const = 0;
@@ -328,10 +337,6 @@ class TypeSystemSwift : public TypeSystem {
   Module *m_module = nullptr;
 };
 
-typedef std::shared_ptr<TypeSystemSwift> TypeSystemSwiftSP;
-typedef std::shared_ptr<TypeSystemSwiftTypeRefForExpressions>
-    TypeSystemSwiftTypeRefForExpressionsSP;
-
 } // namespace lldb_private
 
 namespace llvm {