[lldb] Fix use-after-free in Playgrounds

The precise compiler invocations change moved the ownership of the
SwiftPersistentExpressionState from SwiftASTContext to
TypeSystemSwiftTypeRef, which makes it shared between all
SwiftASTCOntextForExpression objects. The LLDB name lookup contained a
check to avoid finding a persistent result from the wrong context, but
there was no such check in the REPL name lookup.

This is a problem when a Playground imports a framework that pulls in
a fresh dylib because this forces a new SwiftASTContext to be created.

Duplicating the check avoids the crash. I further made sure to clear
the persistent state when the SwiftASTContext is replaced, and added
several assertions to ensure consistency.

rdar://143923367

Author: adrian-prantl

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

@@ -138,8 +138,7 @@ LLDBNameLookup::LLDBNameLookup(
     SwiftExpressionParser::SILVariableMap &variable_map, SymbolContext &sc,
     ExecutionContextScope &exe_scope)
     : SILDebuggerClient(source_file.getASTContext()),
-      m_log(GetLog(LLDBLog::Expressions)), m_source_file(source_file),
-      m_variable_map(variable_map), m_sc(sc) {
+      m_source_file(source_file), m_variable_map(variable_map), m_sc(sc) {
   source_file.getParentModule()->setDebugClient(this);
 
   if (!m_sc.target_sp)
@@ -177,11 +176,15 @@ void LLDBNameLookup::RegisterTypeAliases(
 
 /// A name lookup class for debugger expr mode.
 class LLDBExprNameLookup : public LLDBNameLookup {
+  const SwiftASTContextForExpressions *m_this_context;
+
 public:
   LLDBExprNameLookup(swift::SourceFile &source_file,
                      SwiftExpressionParser::SILVariableMap &variable_map,
-                     SymbolContext &sc, ExecutionContextScope &exe_scope)
-      : LLDBNameLookup(source_file, variable_map, sc, exe_scope) {}
+                     SymbolContext &sc, ExecutionContextScope &exe_scope,
+                     const SwiftASTContextForExpressions *this_context)
+      : LLDBNameLookup(source_file, variable_map, sc, exe_scope),
+        m_this_context(this_context) {}
 
   bool shouldGlobalize(swift::Identifier Name, swift::DeclKind Kind) override {
     // Extensions have to be globalized, there's no way to mark them
@@ -195,7 +198,7 @@ class LLDBExprNameLookup : public LLDBNameLookup {
       return true;
 
     if (Name.str().starts_with("$")) {
-      LLDB_LOG(m_log,
+      LLDB_LOG(GetLog(LLDBLog::Expressions),
                "[LLDBExprNameLookup::shouldGlobalize] Returning true to "
                "globalizing {0}",
                Name.str());
@@ -223,7 +226,7 @@ class LLDBExprNameLookup : public LLDBNameLookup {
     static unsigned counter = 0;
     unsigned count = counter++;
 
-    LLDB_LOG(m_log,
+    LLDB_LOG(GetLog(LLDBLog::Expressions),
              "[LLDBExprNameLookup::lookupOverrides({0})] Searching for \"{1}\"",
              count, Name.getIdentifier().get());
 
@@ -234,6 +237,8 @@ class LLDBExprNameLookup : public LLDBNameLookup {
                        swift::SourceLoc Loc, bool IsTypeLookup,
                        ResultVector &RV) override {
     LLDB_SCOPED_TIMER();
+    assert(SwiftASTContext::GetSwiftASTContext(&DC->getASTContext()) ==
+           m_this_context);
     static unsigned counter = 0;
     unsigned count = counter++;
 
@@ -242,7 +247,7 @@ class LLDBExprNameLookup : public LLDBNameLookup {
       return false;
 
     LLDB_LOG(
-        m_log,
+        GetLog(LLDBLog::Expressions),
         "[LLDBExprNameLookup::lookupAdditions ({0})] Searching for \"{1}\"",
         count, NameStr);
 
@@ -290,9 +295,10 @@ class LLDBExprNameLookup : public LLDBNameLookup {
                                                    persistent_results);
 
         for (CompilerDecl & decl : persistent_results) {
-          if (decl.GetTypeSystem() !=
-              SwiftASTContext::GetSwiftASTContext(&DC->getASTContext())) {
-            LLDB_LOG(m_log, "ignoring persistent result from other context");
+          if (decl.GetTypeSystem() != m_this_context) {
+            LLDB_LOG(GetLog(LLDBLog::Expressions),
+                     "Ignoring persistent result from other context: {0}",
+                     NameStr);
             continue;
           }
           auto *value_decl =
@@ -368,11 +374,15 @@ class LLDBExprNameLookup : public LLDBNameLookup {
 
 /// A name lookup class for REPL and Playground mode.
 class LLDBREPLNameLookup : public LLDBNameLookup {
+  const SwiftASTContextForExpressions *m_this_context;
+
 public:
   LLDBREPLNameLookup(swift::SourceFile &source_file,
                      SwiftExpressionParser::SILVariableMap &variable_map,
-                     SymbolContext &sc, ExecutionContextScope &exe_scope)
-      : LLDBNameLookup(source_file, variable_map, sc, exe_scope) {}
+                     SymbolContext &sc, ExecutionContextScope &exe_scope,
+                     const SwiftASTContextForExpressions *this_context)
+      : LLDBNameLookup(source_file, variable_map, sc, exe_scope),
+        m_this_context(this_context) {}
 
   bool shouldGlobalize(swift::Identifier Name, swift::DeclKind kind) override {
     return false;
@@ -390,6 +400,9 @@ class LLDBREPLNameLookup : public LLDBNameLookup {
                        swift::SourceLoc Loc, bool IsTypeLookup,
                        ResultVector &RV) override {
     LLDB_SCOPED_TIMER();
+
+    assert(SwiftASTContext::GetSwiftASTContext(&DC->getASTContext()) ==
+           m_this_context);
     static unsigned counter = 0;
     unsigned count = counter++;
 
@@ -398,7 +411,7 @@ class LLDBREPLNameLookup : public LLDBNameLookup {
       return false;
 
     LLDB_LOG(
-        m_log,
+        GetLog(LLDBLog::Expressions),
         "[LLDBREPLNameLookup::lookupAdditions ({0})] Searching for \"{1}\"",
         count, NameStr);
 
@@ -420,6 +433,12 @@ class LLDBREPLNameLookup : public LLDBNameLookup {
 
     // Append the persistent decls that we found to the result vector.
     for (auto result : persistent_decl_results) {
+      if (result.GetTypeSystem() != m_this_context) {
+        LLDB_LOG(GetLog(LLDBLog::Expressions),
+                 "Ignoring persistent result from other context: {0}", NameStr);
+        continue;
+      }
+
       // No import required.
       auto *result_decl =
           static_cast<swift::ValueDecl *>(result.GetOpaqueDecl());
@@ -1375,11 +1394,11 @@ SwiftExpressionParser::ParseAndImport(
 
   LLDBNameLookup *external_lookup;
   if (m_options.GetPlaygroundTransformEnabled() || m_options.GetREPLEnabled()) {
-    external_lookup =
-        new LLDBREPLNameLookup(*source_file, variable_map, m_sc, *m_exe_scope);
+    external_lookup = new LLDBREPLNameLookup(*source_file, variable_map, m_sc,
+                                             *m_exe_scope, &m_swift_ast_ctx);
   } else {
-    external_lookup =
-        new LLDBExprNameLookup(*source_file, variable_map, m_sc, *m_exe_scope);
+    external_lookup = new LLDBExprNameLookup(*source_file, variable_map, m_sc,
+                                             *m_exe_scope, &m_swift_ast_ctx);
   }
 
   // FIXME: This call is here just so that the we keep the
@@ -1758,8 +1777,15 @@ SwiftExpressionParser::Parse(DiagnosticManager &diagnostic_manager,
 
     // Signal that we want to retry the expression exactly once with a
     // fresh SwiftASTContext.
-    if (retry)
+    if (retry) {
+      if (auto *persistent_state =
+              llvm::cast_or_null<SwiftPersistentExpressionState>(
+                  m_sc.target_sp->GetPersistentExpressionStateForLanguage(
+                      lldb::eLanguageTypeSwift)))
+        persistent_state->Clear();
+
       return ParseResult::retry_fresh_context;
+    }
 
     // Unrecoverable error.
     return ParseResult::unrecoverable_error;
@@ -2023,6 +2049,11 @@ SwiftExpressionParser::Parse(DiagnosticManager &diagnostic_manager,
     return ParseResult::unrecoverable_error;
   }
 
+  if (m_swift_ast_ctx.GetASTContext()->hadError()) {
+    DiagnoseSwiftASTContextError();
+    return ParseResult::unrecoverable_error;
+  }
+
   {
     std::lock_guard<std::recursive_mutex> global_context_locker(
         IRExecutionUnit::GetLLVMGlobalContextMutex());

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

@@ -212,6 +212,7 @@ class SwiftExpressionParser : public ExpressionParser {
   /// The container for the IR, to be JIT-compiled or interpreted.
   lldb::IRExecutionUnitSP m_execution_unit_sp;
   /// The AST context to build the expression into.
+  /// A shared pointer of this is held by SwiftUserexpression.
   SwiftASTContextForExpressions &m_swift_ast_ctx;
   /// Used to manage the memory of a potential on-off context.
   //lldb::TypeSystemSP m_typesystem_sp;
@@ -254,7 +255,6 @@ class LLDBNameLookup : public swift::SILDebuggerClient {
       llvm::SmallVectorImpl<swift::TypeAliasDecl *> &type_aliases);
 
 protected:
-  Log *m_log;
   swift::SourceFile &m_source_file;
   SwiftExpressionParser::SILVariableMap &m_variable_map;
   SymbolContext m_sc;

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

@@ -48,6 +48,7 @@ class SwiftPersistentExpressionState : public PersistentExpressionState {
 
     void CopyDeclsTo(SwiftDeclMap &target_map);
     static bool DeclsAreEquivalent(CompilerDecl lhs, CompilerDecl rhs);
+    void Clear() { m_swift_decls.clear(); }
 
   private:
     /// Each decl also stores the context it comes from.
@@ -102,6 +103,8 @@ class SwiftPersistentExpressionState : public PersistentExpressionState {
       const std::vector<CompilerDecl> &excluding_equivalents,
       std::vector<CompilerDecl> &matches);
 
+  void Clear() { m_swift_persistent_decls.Clear(); }
+
 private:
   /// The counter used by GetNextResultName().
   uint32_t m_next_persistent_variable_id;

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

@@ -1089,9 +1089,6 @@ class SwiftASTContextForExpressions : public SwiftASTContext {
   /// These are the names of modules that we have loaded by hand into
   /// the Contexts we make for parsing.
   HandLoadedModuleSet m_hand_loaded_modules;
-
-private:
-  std::unique_ptr<SwiftPersistentExpressionState> m_persistent_state_up;
 };
 
 } // namespace lldb_private

lldb/test/API/lang/swift/playgrounds-repl/dylib/BeforeImport.swift

@@ -0,0 +1 @@
+let comment = "Populate the persistent state with some fields"

lldb/test/API/lang/swift/playgrounds-repl/dylib/Dylib.swift

@@ -0,0 +1,2 @@
+
+public func f() -> String { return "Hello from the Dylib!" }

lldb/test/API/lang/swift/playgrounds-repl/dylib/Import.swift

@@ -0,0 +1,3 @@
+import Dylib
+f()
+let comment = "and back again"

lldb/test/API/lang/swift/playgrounds-repl/dylib/Makefile

@@ -0,0 +1,11 @@
+all: Dylib.framework
+
+Dylib.framework: Dylib.swift
+	"$(MAKE)" -f $(MAKEFILE_RULES) \
+		FRAMEWORK=Dylib \
+		DYLIB_SWIFT_SOURCES=Dylib.swift \
+		DYLIB_NAME=Dylib \
+		SWIFTFLAGS_EXTRAS="-Xcc -DNEW_OPTION_FROM_DYLIB=1" \
+		MAKE_DSYM=YES
+
+include ../Makefile.common

lldb/test/API/lang/swift/playgrounds-repl/dylib/PlaygroundStub.swift

@@ -0,0 +1,18 @@
+// PlaygroundStub.swift
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+// -----------------------------------------------------------------------------
+
+@_silgen_name ("playground_logger_initialize") func builtin_initialize();
+@_silgen_name ("GetOutput") func get_output() -> String;
+
+builtin_initialize();
+
+print(""); // Set breakpoint here

lldb/test/API/lang/swift/playgrounds-repl/dylib/PlaygroundsRuntime.swift

@@ -0,0 +1,92 @@
+// PlaygroundsRuntime.swift
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+// -----------------------------------------------------------------------------
+
+// If you're modifying this file to add or modify function signatures, you
+// should be notifying the maintainer of PlaygroundLogger and also the
+// maintainer of lib/Sema/PlaygroundTransform.cpp.
+
+var PlaygroundOutput = ""
+
+struct SourceRange {
+  let sl : Int
+  let el : Int
+  let sc : Int
+  let ec : Int
+  var text : String {
+    return "[\(sl):\(sc)-\(el):\(ec)]"
+  }
+}
+
+struct ModuleFileIdentifier {
+  let moduleId : Int
+  let fileId : Int
+  var text : String {
+    return "[\(moduleId):\(fileId)]"
+  }
+}
+
+class LogRecord {
+  let text : String
+
+  init(api : String, object : Any, name : String, id : Int, range : SourceRange, moduleFileId : ModuleFileIdentifier) {
+    var object_description : String = ""
+    print(object, terminator: "", to: &object_description)
+    text = moduleFileId.text + " " + range.text + " " + api + "[" + name + "='" + object_description + "']"
+  }
+  init(api : String, object : Any, name : String, range : SourceRange, moduleFileId : ModuleFileIdentifier) {
+    var object_description : String = ""
+    print(object, terminator: "", to: &object_description)
+    text = moduleFileId.text + " " + range.text + " " + api + "[" + name + "='" + object_description + "']"
+  }
+  init(api : String, object: Any, range : SourceRange, moduleFileId : ModuleFileIdentifier) {
+    var object_description : String = ""
+    print(object, terminator: "", to: &object_description)
+    text = moduleFileId.text + " " + range.text + " " + api + "['" + object_description + "']"
+  }
+  init(api: String, range : SourceRange, moduleFileId : ModuleFileIdentifier) {
+    text = moduleFileId.text + " " + range.text + " " + api
+  }
+}
+
+@_silgen_name ("playground_logger_initialize") public func builtin_initialize() {
+}
+
+@_silgen_name ("playground_log_hidden") public func builtin_log_with_id<T>(_ object : T, _ name : String, _ id : Int, _ sl : Int, _ el : Int, _ sc : Int, _ ec: Int, _ moduleId : Int, _ fileId : Int) -> AnyObject? {
+  let moduleFileId = ModuleFileIdentifier(moduleId:moduleId, fileId:fileId)
+  return LogRecord(api:"__builtin_log", object:object, name:name, id:id, range : SourceRange(sl:sl, el:el, sc:sc, ec:ec), moduleFileId:moduleFileId)
+}
+
+@_silgen_name ("playground_log_scope_entry") public func builtin_log_scope_entry(_ sl : Int, _ el : Int, _ sc : Int, _ ec: Int, _ moduleId : Int, _ fileId : Int) -> AnyObject? {
+  let moduleFileId = ModuleFileIdentifier(moduleId:moduleId, fileId:fileId)
+  return LogRecord(api:"__builtin_log_scope_entry", range : SourceRange(sl:sl, el:el, sc:sc, ec:ec), moduleFileId:moduleFileId)
+}
+
+@_silgen_name ("playground_log_scope_exit") public func builtin_log_scope_exit(_ sl : Int, _ el : Int, _ sc : Int, _ ec: Int, _ moduleId : Int, _ fileId : Int) -> AnyObject? {
+  let moduleFileId = ModuleFileIdentifier(moduleId:moduleId, fileId:fileId)
+  return LogRecord(api:"__builtin_log_scope_exit", range : SourceRange(sl:sl, el:el, sc:sc, ec:ec), moduleFileId:moduleFileId)
+}
+
+@_silgen_name ("playground_log_postprint") public func builtin_postPrint(_ sl : Int, _ el : Int, _ sc : Int, _ ec: Int, _ moduleId : Int, _ fileId : Int) -> AnyObject? {
+  let moduleFileId = ModuleFileIdentifier(moduleId:moduleId, fileId:fileId)
+  return LogRecord(api:"__builtin_postPrint", range : SourceRange(sl:sl, el:el, sc:sc, ec:ec), moduleFileId:moduleFileId)
+}
+
+@_silgen_name ("DVTSendPlaygroundLogData") public func builtin_send_data(_ object:AnyObject?) {
+  print((object as! LogRecord).text)
+  PlaygroundOutput.append((object as! LogRecord).text)
+  PlaygroundOutput.append("\n")
+}
+
+@_silgen_name ("GetOutput") public func get_output() -> String {
+  return PlaygroundOutput
+}
+

lldb/test/API/lang/swift/playgrounds-repl/dylib/TestPlaygroundREPLImport.py

@@ -0,0 +1,27 @@
+from __future__ import print_function
+import lldbsuite.test.lldbplaygroundrepl as repl
+from lldbsuite.test.lldbtest import *
+
+class TestPlaygroundREPLImport(repl.PlaygroundREPLTest):
+
+    mydir = repl.PlaygroundREPLTest.compute_mydir(__file__)
+    
+    def do_test(self):
+        """
+        Test importing a library that adds new Clang options.
+        """
+        self.expect('settings set target.swift-framework-search-paths "%s"' %
+                    self.getBuildDir())
+        self.expect('settings set target.use-all-compiler-flags true')
+
+        log = self.getBuildArtifact('types.log')
+        self.expect('log enable lldb types -f ' + log)
+        result, playground_output = self.execute_code('BeforeImport.swift')
+        self.assertIn("persistent", playground_output.GetSummary())
+        result, playground_output = self.execute_code('Import.swift')
+        self.assertIn("Hello from the Dylib", playground_output.GetSummary())
+        self.assertIn("and back again", playground_output.GetSummary())
+
+        # Scan through the types log to make sure the SwiftASTContext was poisoned.
+        self.filecheck('platform shell cat ""%s"' % log, __file__)
+#       CHECK: New Swift image added{{.*}}Versions/A/Dylib{{.*}}ClangImporter needs to be reinitialized

lldb/test/API/lang/swift/playgrounds/Import.swift

@@ -1,2 +1,3 @@
 import Dylib
 f()
+let comment = "and back again"