[BOLT] Fix use-after-free in RewriteInstance::mapCodeSections

[jobnoorman]

Summary

When a cold function is too large, its section gets deregistered.
However, the section is still dereferenced later to get its RuntimeDyld
ID. This patch moves the deregistration to after the last dereference.

Note that this came up in D147544 and I haven't found a way to actually
trigger this bug (i.e., I'm not sure how to create a cold function
that's considered "too large"). I tried to resolve the issue without
affecting BOLT's behavior but there might be better ways to solve it
(e.g., not deregistering, not mapping the deregistered section in
RuntimeDyld?).

[Amir]

Good catch!

[Amir]

Please remove "[NFC]" from title as bugfix is a functional change

bolt/lib/Rewrite/RewriteInstance.cpp

for (auto &BFI : BC->getBinaryFunctions()) {
     // Cold fragments are aligned at 16 bytes.
     NextAvailableAddress = alignTo(NextAvailableAddress, 16);
     if (TooLarge) {
       // The corresponding FDE will refer to address 0.
       FF.setAddress(0);
       FF.setImageAddress(0);
       FF.setImageSize(0);
       FF.setFileOffset(0);
-      BC->deregisterSection(*ColdSection);
     } else {
       FF.setAddress(NextAvailableAddress);
       FF.setImageAddress(ColdSection->getAllocAddress());
       FF.setImageSize(ColdSection->getOutputSize());
       FF.setFileOffset(getFileOffsetForAddress(NextAvailableAddress));
       ColdSection->setOutputAddress(FF.getAddress());
     }
 
     LLVM_DEBUG(
         dbgs() << formatv(
             "BOLT: mapping cold fragment {0:x+} to {1:x+} with size {2:x+}\n",
             FF.getImageAddress(), FF.getAddress(), FF.getImageSize()));
     RTDyld.reassignSectionAddress(ColdSection->getSectionID(), FF.getAddress());
 
+    if (TooLarge)
+      BC->deregisterSection(*ColdSection);
+
     NextAvailableAddress += FF.getImageSize();
   }
 
   // Add the new text section aggregating all existing code sections.
   // This is pseudo-section that serves a purpose of creating a corresponding
   // entry in section header table.
   int64_t NewTextSectionSize =
       NextAvailableAddress - NewTextSectionStartAddress;