Fix use-after-free with precompiled preambles

llvm-svn: 109505

2 files changed, 14 insertions, 3 deletions

diff --git a/clang/include/clang/Frontend/ASTUnit.h b/clang/include/clang/Frontend/ASTUnit.h
index d0202dfe280..10cf3d50ba5 100644
--- a/clang/include/clang/Frontend/ASTUnit.h
+++ b/clang/include/clang/Frontend/ASTUnit.h
@@ -135,6 +135,11 @@ private:
   /// file within the precompiled preamble.
   unsigned PreambleReservedSize;
   
+  /// \brief When non-NULL, this is the buffer used to store the contents of
+  /// the main file when it has been padded for use with the precompiled
+  /// preamble.
+  llvm::MemoryBuffer *SavedMainFileBuffer;
+  
   ASTUnit(const ASTUnit&); // DO NOT IMPLEMENT
   ASTUnit &operator=(const ASTUnit &); // DO NOT IMPLEMENT
   
diff --git a/clang/lib/Frontend/ASTUnit.cpp b/clang/lib/Frontend/ASTUnit.cpp
index 0463db18dc0..d67a6652dac 100644
--- a/clang/lib/Frontend/ASTUnit.cpp
+++ b/clang/lib/Frontend/ASTUnit.cpp
@@ -39,7 +39,7 @@ using namespace clang;
 
 ASTUnit::ASTUnit(bool _MainFileIsAST)
   : CaptureDiagnostics(false), MainFileIsAST(_MainFileIsAST), 
-    ConcurrencyCheckValue(CheckUnlocked) { }
+    ConcurrencyCheckValue(CheckUnlocked), SavedMainFileBuffer(0) { }
 
 ASTUnit::~ASTUnit() {
   ConcurrencyCheckValue = CheckLocked;
@@ -60,6 +60,8 @@ ASTUnit::~ASTUnit() {
          ++FB)
       delete FB->second;
   }
+  
+  delete SavedMainFileBuffer;
 }
 
 void ASTUnit::CleanTemporaryFiles() {
@@ -328,6 +330,9 @@ public:
 /// \returns True if a failure occurred that causes the ASTUnit not to
 /// contain any translation-unit information, false otherwise.
 bool ASTUnit::Parse(llvm::MemoryBuffer *OverrideMainBuffer) {
+  delete SavedMainFileBuffer;
+  SavedMainFileBuffer = 0;
+  
   if (!Invocation.get())
     return true;
   
@@ -395,6 +400,9 @@ bool ASTUnit::Parse(llvm::MemoryBuffer *OverrideMainBuffer) {
                                                     = PreambleEndsAtStartOfLine;
     PreprocessorOpts.ImplicitPCHInclude = PreambleFile.str();
     PreprocessorOpts.DisablePCHValidation = true;
+    
+    // Keep track of the override buffer;
+    SavedMainFileBuffer = OverrideMainBuffer;
   }
   
   llvm::OwningPtr<TopLevelDeclTrackerAction> Act;
@@ -787,7 +795,6 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocation(CompilerInvocation *CI,
   if (!AST->Parse(OverrideMainBuffer))
     return AST.take();
   
-  delete OverrideMainBuffer;
   return 0;
 }
 
@@ -884,6 +891,5 @@ bool ASTUnit::Reparse(RemappedFile *RemappedFiles, unsigned NumRemappedFiles) {
 
   // Parse the sources
   bool Result = Parse(OverrideMainBuffer);  
-  delete OverrideMainBuffer;
   return Result;
 }