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

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.

Reviewed By: Amir

Differential Revision: https://reviews.llvm.org/D148427

1 files changed, 3 insertions, 1 deletions

diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 57e41d5b5724..4c4b3fa47c62 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -4097,7 +4097,6 @@ void RewriteInstance::mapCodeSections(RuntimeDyld &RTDyld) {
       FF.setImageAddress(0);
       FF.setImageSize(0);
       FF.setFileOffset(0);
-      BC->deregisterSection(*ColdSection);
     } else {
       FF.setAddress(NextAvailableAddress);
       FF.setImageAddress(ColdSection->getAllocAddress());
@@ -4112,6 +4111,9 @@ void RewriteInstance::mapCodeSections(RuntimeDyld &RTDyld) {
             FF.getImageAddress(), FF.getAddress(), FF.getImageSize()));
     RTDyld.reassignSectionAddress(ColdSection->getSectionID(), FF.getAddress());
 
+    if (TooLarge)
+      BC->deregisterSection(*ColdSection);
+
     NextAvailableAddress += FF.getImageSize();
   }