Fix double-delete when failing to load llvm bitcode (#2169)

We added an exception when the bitcode reader reports an error, however the code higher up in the call stack isn't fully exception-safe and it turns out that two unique_ptrs over the same MemoryBuffer are alive at the same time, hence the double-delete and, more often than not, crash.

lib/Bitcode/Reader/BitcodeReader.cpp

@@ -222,7 +222,7 @@ public:
   std::error_code error(BitcodeError E);
   std::error_code error(const Twine &Message);
 
-  BitcodeReader(MemoryBuffer *Buffer, LLVMContext &Context,
+  BitcodeReader(std::unique_ptr<MemoryBuffer> &&Buffer, LLVMContext &Context, // HLSL Change: unique_ptr
                 DiagnosticHandlerFunction DiagnosticHandler);
   BitcodeReader(LLVMContext &Context,
                 DiagnosticHandlerFunction DiagnosticHandler);
@@ -435,11 +435,11 @@ static void ReportWarning(DiagnosticHandlerFunction F, const char *Msg) {
 }
 // HLSL Change Ends
 
-BitcodeReader::BitcodeReader(MemoryBuffer *Buffer, LLVMContext &Context,
+BitcodeReader::BitcodeReader(std::unique_ptr<MemoryBuffer> &&Buffer, LLVMContext &Context, // HLSL Change: unique_ptr
                              DiagnosticHandlerFunction DiagnosticHandler)
     : Context(Context),
       DiagnosticHandler(getDiagHandler(DiagnosticHandler, Context)),
-      Buffer(Buffer), ValueList(Context), MDValueList(Context) {}
+      Buffer(std::move(Buffer)), ValueList(Context), MDValueList(Context) {} // HLSL Change: std::move
 
 BitcodeReader::BitcodeReader(LLVMContext &Context,
                              DiagnosticHandlerFunction DiagnosticHandler)
@@ -4733,29 +4733,27 @@ const std::error_category &llvm::BitcodeErrorCategory() {
 
 static ErrorOr<std::unique_ptr<Module>>
 getBitcodeModuleImpl(std::unique_ptr<DataStreamer> Streamer, StringRef Name,
-                     BitcodeReader *R, LLVMContext &Context,
+                     std::unique_ptr<BitcodeReader> RPtr, LLVMContext &Context, // HLSL Change: unique_ptr
                      bool MaterializeAll, bool ShouldLazyLoadMetadata) {
   std::unique_ptr<Module> M = make_unique<Module>(Name, Context);
+  // HLSL Change Begin: Transfer ownership of R to M, but keep a raw pointer
+  BitcodeReader* R = RPtr.release();
   M->setMaterializer(R);
-
-  auto cleanupOnError = [&](std::error_code EC) {
-    R->releaseBuffer(); // Never take ownership on error.
-    return EC;
-  };
+  // HLSL Change End
 
   // Delay parsing Metadata if ShouldLazyLoadMetadata is true.
   if (std::error_code EC = R->parseBitcodeInto(std::move(Streamer), M.get(),
                                                ShouldLazyLoadMetadata))
-    return cleanupOnError(EC);
+    return EC; // HLSL Change: Correct memory management of BitcodeReader.buffer
 
   if (MaterializeAll) {
     // Read in the entire module, and destroy the BitcodeReader.
     if (std::error_code EC = M->materializeAllPermanently())
-      return cleanupOnError(EC);
+      return EC; // HLSL Change: Correct memory management of BitcodeReader.buffer
   } else {
     // Resolve forward references from blockaddresses.
     if (std::error_code EC = R->materializeForwardReferencedFunctions())
-      return cleanupOnError(EC);
+      return EC; // HLSL Change: Correct memory management of BitcodeReader.buffer
   }
   return std::move(M);
 }
@@ -4773,16 +4771,20 @@ getLazyBitcodeModuleImpl(std::unique_ptr<MemoryBuffer> &&Buffer,
                          LLVMContext &Context, bool MaterializeAll,
                          DiagnosticHandlerFunction DiagnosticHandler,
                          bool ShouldLazyLoadMetadata = false) {
-  BitcodeReader *R =
-      new BitcodeReader(Buffer.get(), Context, DiagnosticHandler);
+  // HLSL Change Begin: Proper memory management with unique_ptr
+  // Get the buffer identifier before we transfer the ownership to the bitcode reader,
+  // this is ugly but safe as long as it keeps the buffer, and hence identifier string, alive.
+  const char* BufferIdentifier = Buffer->getBufferIdentifier();
+  std::unique_ptr<BitcodeReader> R = llvm::make_unique<BitcodeReader>(
+    std::move(Buffer), Context, DiagnosticHandler);
 
   ErrorOr<std::unique_ptr<Module>> Ret =
-      getBitcodeModuleImpl(nullptr, Buffer->getBufferIdentifier(), R, Context,
+      getBitcodeModuleImpl(nullptr, BufferIdentifier, std::move(R), Context,
                            MaterializeAll, ShouldLazyLoadMetadata);
+  // HLSL Change End
   if (!Ret)
     return Ret;
 
-  Buffer.release(); // The BitcodeReader owns it now.
   return Ret;
 }
 
@@ -4797,9 +4799,9 @@ ErrorOr<std::unique_ptr<Module>> llvm::getStreamedBitcodeModule(
     StringRef Name, std::unique_ptr<DataStreamer> Streamer,
     LLVMContext &Context, DiagnosticHandlerFunction DiagnosticHandler) {
   std::unique_ptr<Module> M = make_unique<Module>(Name, Context);
-  BitcodeReader *R = new BitcodeReader(Context, DiagnosticHandler);
+  std::unique_ptr<BitcodeReader> R = llvm::make_unique<BitcodeReader>(Context, DiagnosticHandler); // HLSL Change: unique_ptr
 
-  return getBitcodeModuleImpl(std::move(Streamer), Name, R, Context, false,
+  return getBitcodeModuleImpl(std::move(Streamer), Name, std::move(R), Context, false, // HLSL Change: std::move
                               false);
 }
 
@@ -4838,7 +4840,7 @@ std::string
 llvm::getBitcodeTargetTriple(MemoryBufferRef Buffer, LLVMContext &Context,
                              DiagnosticHandlerFunction DiagnosticHandler) {
   std::unique_ptr<MemoryBuffer> Buf = MemoryBuffer::getMemBuffer(Buffer, false);
-  auto R = llvm::make_unique<BitcodeReader>(Buf.release(), Context,
+  auto R = llvm::make_unique<BitcodeReader>(std::move(Buf), Context, // HLSL Change: std::move
                                             DiagnosticHandler);
   ErrorOr<std::string> Triple = R->parseTriple();
   if (Triple.getError())

tools/clang/unittests/HLSL/ValidationTest.cpp

@@ -258,6 +258,8 @@ public:
 
   TEST_METHOD(SpaceOnlyRegisterFail)
 
+  TEST_METHOD(WhenDisassembleInvalidBlobThenFail)
+
   dxc::DxcDllSupport m_dllSupport;
   VersionSupportInfo m_ver;
 
@@ -3057,6 +3059,23 @@ TEST_F(ValidationTest, SpaceOnlyRegisterFail) {
   TestCheck(L"..\\CodeGenHLSL\\space-only-register.hlsl");
 }
 
+// Regression test for a double-delete when failing to parse bitcode.
+TEST_F(ValidationTest, WhenDisassembleInvalidBlobThenFail) {
+  if (!m_dllSupport.IsEnabled()) {
+    VERIFY_SUCCEEDED(m_dllSupport.Initialize());
+  }
+
+  CComPtr<IDxcBlobEncoding> pInvalidBitcode;
+  Utf8ToBlob(m_dllSupport, "This text is certainly not bitcode", &pInvalidBitcode);
+
+  CComPtr<IDxcCompiler> pCompiler;
+  VERIFY_SUCCEEDED(m_dllSupport.CreateInstance(CLSID_DxcCompiler, &pCompiler));
+
+  CComPtr<IDxcBlobEncoding> pDisassembly;
+  VERIFY_FAILED(pCompiler->Disassemble(pInvalidBitcode, &pDisassembly));
+}
+
+
 TEST_F(ValidationTest, GSMainMissingAttributeFail) {
   TestCheck(L"..\\CodeGenHLSL\\attributes-gs-no-inout-main.hlsl");
 }