[BOLT][instr] Avoid WX segment (llvm#128982)

BOLT instrumented binary today has a readable (R), writeable (W) and also
executable (X) segment, which Android system won't load due to its WX
attribute. Such RWX segment was produced because BOLT has a two step linking,
first for everything in the updated or rewritten input binary and next for
runtime library. Each linking will layout sections in the order of RX sections
followed by RO sections and then followed by RW sections. So we could end up
having a RW section `.bolt.instr.counters` surrounded by a number of RO and RX
sections, and a new text segment was then formed by including all RX sections
which includes the RW section in the middle, and hence the RWX segment. One
way to fix this is to separate the RW `.bolt.instr.counters` section into its
own segment by a). assigning the starting addresses for section
`.bolt.instr.counters` and its following section with regular page aligned
addresses and b). creating two extra program headers accordingly.

Author: yozhu

bolt/include/bolt/Rewrite/RewriteInstance.h

@@ -505,6 +505,9 @@ class RewriteInstance {
   /// Number of local symbols in newly written symbol table.
   uint64_t NumLocalSymbols{0};
 
+  /// Flag indicating runtime library linking just started.
+  bool StartLinkingRuntimeLib{false};
+
   /// Information on special Procedure Linkage Table sections. There are
   /// multiple variants generated by different linkers.
   struct PLTSectionInfo {

bolt/lib/Passes/Instrumentation.cpp

@@ -604,7 +604,7 @@ Error Instrumentation::runOnFunctions(BinaryContext &BC) {
                                                  /*IsText=*/false,
                                                  /*IsAllocatable=*/true);
   BC.registerOrUpdateSection(".bolt.instr.counters", ELF::SHT_PROGBITS, Flags,
-                             nullptr, 0, 1);
+                             nullptr, 0, BC.RegularPageSize);
 
   BC.registerOrUpdateNoteSection(".bolt.instr.tables", nullptr, 0,
                                  /*Alignment=*/1,

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -628,6 +628,11 @@ Error RewriteInstance::discoverStorage() {
     unsigned Phnum = Obj.getHeader().e_phnum;
     Phnum += 3;
 
+    // Reserve two more pheaders to avoid having writeable and executable
+    // segment in instrumented binary.
+    if (opts::Instrument)
+      Phnum += 2;
+
     NextAvailableAddress += Phnum * sizeof(ELF64LEPhdrTy);
     NextAvailableOffset += Phnum * sizeof(ELF64LEPhdrTy);
   }
@@ -2083,6 +2088,13 @@ void RewriteInstance::adjustCommandLineOptions() {
     opts::HotText = false;
   }
 
+  if (opts::Instrument && opts::UseGnuStack) {
+    BC->errs() << "BOLT-ERROR: cannot avoid having writeable and executable "
+                  "segment in instrumented binary if program headers will be "
+                  "updated in place\n";
+    exit(1);
+  }
+
   if (opts::HotText && opts::HotTextMoveSections.getNumOccurrences() == 0) {
     opts::HotTextMoveSections.addValue(".stub");
     opts::HotTextMoveSections.addValue(".mover");
@@ -3612,11 +3624,13 @@ void RewriteInstance::emitAndLink() {
         static_cast<MCObjectStreamer &>(*Streamer).getAssembler());
   }
 
-  if (RuntimeLibrary *RtLibrary = BC->getRuntimeLibrary())
+  if (RuntimeLibrary *RtLibrary = BC->getRuntimeLibrary()) {
+    StartLinkingRuntimeLib = true;
     RtLibrary->link(*BC, ToolPath, *Linker, [this](auto MapSection) {
       // Map newly registered sections.
       this->mapAllocatableSections(MapSection);
     });
+  }
 
   // Once the code is emitted, we can rename function sections to actual
   // output sections and de-register sections used for emission.
@@ -4011,12 +4025,17 @@ void RewriteInstance::mapAllocatableSections(
         Section.setOutputFileOffset(Section.getInputFileOffset());
         MapSection(Section, Section.getAddress());
       } else {
-        NextAvailableAddress =
-            alignTo(NextAvailableAddress, Section.getAlignment());
+        uint64_t Alignment = Section.getAlignment();
+        if (opts::Instrument && StartLinkingRuntimeLib) {
+          Alignment = BC->RegularPageSize;
+          StartLinkingRuntimeLib = false;
+        }
+        NextAvailableAddress = alignTo(NextAvailableAddress, Alignment);
+
         LLVM_DEBUG({
-          dbgs() << "BOLT: mapping section " << Section.getName() << " (0x"
-                 << Twine::utohexstr(Section.getAllocAddress()) << ") to 0x"
-                 << Twine::utohexstr(NextAvailableAddress) << ":0x"
+          dbgs() << "BOLT-DEBUG: mapping section " << Section.getName()
+                 << " (0x" << Twine::utohexstr(Section.getAllocAddress())
+                 << ") to 0x" << Twine::utohexstr(NextAvailableAddress) << ":0x"
                  << Twine::utohexstr(NextAvailableAddress +
                                      Section.getOutputSize())
                  << '\n';
@@ -4079,6 +4098,9 @@ void RewriteInstance::patchELFPHDRTable() {
     }
   }
 
+  if (opts::Instrument)
+    Phnum += 2;
+
   // NOTE Currently .eh_frame_hdr appends to the last segment, recalculate
   // last segments size based on the NextAvailableAddress variable.
   if (!NewWritableSegmentSize) {
@@ -4093,7 +4115,8 @@ void RewriteInstance::patchELFPHDRTable() {
   const uint64_t SavedPos = OS.tell();
   OS.seek(PHDRTableOffset);
 
-  auto createNewTextPhdr = [&]() {
+  auto createNewPhdrs = [&]() {
+    SmallVector<ELF64LEPhdrTy, 3> NewPhdrs;
     ELF64LEPhdrTy NewPhdr;
     NewPhdr.p_type = ELF::PT_LOAD;
     if (PHDRTableAddress) {
@@ -4108,20 +4131,67 @@ void RewriteInstance::patchELFPHDRTable() {
     NewPhdr.p_filesz = NewTextSegmentSize;
     NewPhdr.p_memsz = NewTextSegmentSize;
     NewPhdr.p_flags = ELF::PF_X | ELF::PF_R;
-    if (opts::Instrument) {
-      // FIXME: Currently instrumentation is experimental and the runtime data
-      // is emitted with code, thus everything needs to be writable.
-      NewPhdr.p_flags |= ELF::PF_W;
-    }
     NewPhdr.p_align = BC->PageAlign;
 
-    return NewPhdr;
+    if (!opts::Instrument) {
+      NewPhdrs.push_back(NewPhdr);
+    } else {
+      ErrorOr<BinarySection &> Sec =
+          BC->getUniqueSectionByName(".bolt.instr.counters");
+      assert(Sec && "expected one and only one `.bolt.instr.counters` section");
+      const uint64_t Addr = Sec->getOutputAddress();
+      const uint64_t Offset = Sec->getOutputFileOffset();
+      const uint64_t Size = Sec->getOutputSize();
+      assert(Addr > NewPhdr.p_vaddr &&
+             Addr + Size < NewPhdr.p_vaddr + NewPhdr.p_memsz &&
+             "`.bolt.instr.counters` section is expected to be included in the "
+             "new text sgement");
+
+      // Set correct size for the previous header since we are breaking the
+      // new text segment into three segments.
+      uint64_t Delta = Addr - NewPhdr.p_vaddr;
+      NewPhdr.p_filesz = Delta;
+      NewPhdr.p_memsz = Delta;
+      NewPhdrs.push_back(NewPhdr);
+
+      // Create a program header for a RW segment that includes the
+      // `.bolt.instr.counters` section only.
+      ELF64LEPhdrTy NewPhdrRWSegment;
+      NewPhdrRWSegment.p_type = ELF::PT_LOAD;
+      NewPhdrRWSegment.p_offset = Offset;
+      NewPhdrRWSegment.p_vaddr = Addr;
+      NewPhdrRWSegment.p_paddr = Addr;
+      NewPhdrRWSegment.p_filesz = Size;
+      NewPhdrRWSegment.p_memsz = Size;
+      NewPhdrRWSegment.p_flags = ELF::PF_R | ELF::PF_W;
+      NewPhdrRWSegment.p_align = BC->RegularPageSize;
+      NewPhdrs.push_back(NewPhdrRWSegment);
+
+      // Create a program header for a RX segment that includes all the RX
+      // sections from runtime library.
+      ELF64LEPhdrTy NewPhdrRXSegment;
+      NewPhdrRXSegment.p_type = ELF::PT_LOAD;
+      const uint64_t AddrRX = alignTo(Addr + Size, BC->RegularPageSize);
+      const uint64_t OffsetRX = alignTo(Offset + Size, BC->RegularPageSize);
+      const uint64_t SizeRX = NewTextSegmentSize - (AddrRX - NewPhdr.p_paddr);
+      NewPhdrRXSegment.p_offset = OffsetRX;
+      NewPhdrRXSegment.p_vaddr = AddrRX;
+      NewPhdrRXSegment.p_paddr = AddrRX;
+      NewPhdrRXSegment.p_filesz = SizeRX;
+      NewPhdrRXSegment.p_memsz = SizeRX;
+      NewPhdrRXSegment.p_flags = ELF::PF_X | ELF::PF_R;
+      NewPhdrRXSegment.p_align = BC->RegularPageSize;
+      NewPhdrs.push_back(NewPhdrRXSegment);
+    }
+
+    return NewPhdrs;
   };
 
   auto writeNewSegmentPhdrs = [&]() {
     if (PHDRTableAddress || NewTextSegmentSize) {
-      ELF64LE::Phdr NewPhdr = createNewTextPhdr();
-      OS.write(reinterpret_cast<const char *>(&NewPhdr), sizeof(NewPhdr));
+      SmallVector<ELF64LE::Phdr, 3> NewPhdrs = createNewPhdrs();
+      OS.write(reinterpret_cast<const char *>(NewPhdrs.data()),
+               sizeof(ELF64LE::Phdr) * NewPhdrs.size());
     }
 
     if (NewWritableSegmentSize) {
@@ -4169,8 +4239,12 @@ void RewriteInstance::patchELFPHDRTable() {
     }
     case ELF::PT_GNU_STACK:
       if (opts::UseGnuStack) {
-        // Overwrite the header with the new text segment header.
-        NewPhdr = createNewTextPhdr();
+        // Overwrite the header with the new segment header.
+        assert(!opts::Instrument);
+        SmallVector<ELF64LE::Phdr, 3> NewPhdrs = createNewPhdrs();
+        assert(NewPhdrs.size() == 1 &&
+               "expect exactly one program header was created");
+        NewPhdr = NewPhdrs[0];
         ModdedGnuStack = true;
       }
       break;

bolt/test/avoid-wx-segment.c

@@ -0,0 +1,15 @@
+// Test bolt instrumentation won't generate a binary with any segment that
+// is writable and executable. Basically we want to put `.bolt.instr.counters`
+// section into its own segment, separated from its surrounding RX sections.
+
+// REQUIRES: system-linux
+
+void foo() {}
+void bar() { foo(); }
+
+// RUN: %clang %cflags -c %s -o %t.o
+// RUN: ld.lld -q -o %t.so %t.o -shared --init=foo --fini=foo
+// RUN: llvm-bolt --instrument %t.so -o %tt.so
+// RUN: llvm-readelf -l %tt.so | FileCheck %s
+// CHECK-NOT: RWE
+// CHECK: {{[0-9]*}} .bolt.instr.counters {{$}}