From 8288f17aa33be4cc7eaa95bcd045d50be158012f Mon Sep 17 00:00:00 2001
From: Vedant Kumar <vsk@apple.com>
Date: Fri, 3 Jun 2016 23:12:38 +0000
Subject: [PATCH] [llvm-profdata] Fix use-after-free from discarded
 MemoryBuffer (NFC)

Thanks to Justin Bogner for pointing this out!

Caught by ASAN.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@271748 91177308-0d34-0410-b5e6-96231b3b80d8
---
 tools/llvm-profdata/llvm-profdata.cpp | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/tools/llvm-profdata/llvm-profdata.cpp b/tools/llvm-profdata/llvm-profdata.cpp
index 118a7eba7b7..45f348747db 100644
--- a/tools/llvm-profdata/llvm-profdata.cpp
+++ b/tools/llvm-profdata/llvm-profdata.cpp
@@ -223,16 +223,19 @@ static WeightedFile parseWeightedFile(const StringRef &WeightedFilename) {
   return WeightedFile(FileName, Weight);
 }
 
-static void parseInputFilenamesFile(const StringRef &InputFilenamesFile,
-                                    WeightedFileVector &WFV) {
+static std::unique_ptr<MemoryBuffer>
+parseInputFilenamesFile(const StringRef &InputFilenamesFile,
+                        WeightedFileVector &WFV) {
   if (InputFilenamesFile == "")
-    return;
+    return {};
 
-  auto Buf = MemoryBuffer::getFileOrSTDIN(InputFilenamesFile);
-  if (!Buf)
-    exitWithErrorCode(Buf.getError(), InputFilenamesFile);
+  auto BufOrError = MemoryBuffer::getFileOrSTDIN(InputFilenamesFile);
+  if (!BufOrError)
+    exitWithErrorCode(BufOrError.getError(), InputFilenamesFile);
+
+  std::unique_ptr<MemoryBuffer> Buffer = std::move(*BufOrError);
+  StringRef Data = Buffer->getBuffer();
 
-  StringRef Data = Buf.get()->getBuffer();
   SmallVector<StringRef, 8> Entries;
   Data.split(Entries, '\n', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
   for (const StringRef &FileWeightEntry : Entries) {
@@ -246,6 +249,8 @@ static void parseInputFilenamesFile(const StringRef &InputFilenamesFile,
     else
       WFV.emplace_back(parseWeightedFile(SanitizedEntry));
   }
+
+  return Buffer;
 }
 
 static int merge_main(int argc, const char *argv[]) {
@@ -288,7 +293,7 @@ static int merge_main(int argc, const char *argv[]) {
     WeightedInputs.push_back(WeightedFile(Filename, 1));
   for (StringRef WeightedFilename : WeightedInputFilenames)
     WeightedInputs.push_back(parseWeightedFile(WeightedFilename));
-  parseInputFilenamesFile(InputFilenamesFile, WeightedInputs);
+  auto Buf = parseInputFilenamesFile(InputFilenamesFile, WeightedInputs);
 
   if (WeightedInputs.empty())
     exitWithError("No input files specified. See " +
-- 
2.11.0