[BOLT][instr] Avoid WX segment

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 {{$}}