[lld][ELF] Allow data.rel.ro.hot and .data.rel.ro.unlikely to be RELRO (#148920)

https://discourse.llvm.org/t/rfc-profile-guided-static-data-partitioning/83744
proposes to partition a static data section (like `.data.rel.ro`) into
two sections, one grouping the cold ones and the other grouping the
rest.

lld requires all relro sections to be contiguous. To place
`.data.rel.ro.unlikely` in the middle of all relro sections, this change
proposes to add `.data.rel.ro.unlikely` explicitly as RELRO section.

---------

Co-authored-by: Sam Elliott <[email protected]>

Author: mingmingl-llvm

lld/ELF/Config.h

@@ -404,6 +404,7 @@ struct Config {
   bool zIfuncNoplt;
   bool zInitfirst;
   bool zInterpose;
+  bool zKeepDataSectionPrefix;
   bool zKeepTextSectionPrefix;
   bool zLrodataAfterBss;
   bool zNoBtCfi;

lld/ELF/Driver.cpp

@@ -1625,6 +1625,8 @@ static void readConfigs(Ctx &ctx, opt::InputArgList &args) {
   ctx.arg.zIfuncNoplt = hasZOption(args, "ifunc-noplt");
   ctx.arg.zInitfirst = hasZOption(args, "initfirst");
   ctx.arg.zInterpose = hasZOption(args, "interpose");
+  ctx.arg.zKeepDataSectionPrefix = getZFlag(
+      args, "keep-data-section-prefix", "nokeep-data-section-prefix", false);
   ctx.arg.zKeepTextSectionPrefix = getZFlag(
       args, "keep-text-section-prefix", "nokeep-text-section-prefix", false);
   ctx.arg.zLrodataAfterBss =

lld/ELF/Writer.cpp

@@ -553,6 +553,27 @@ template <class ELFT> void Writer<ELFT>::addSectionSymbols() {
   }
 }
 
+// Returns true if the section is a data section that's read only and
+// relocatable per its section name.
+static bool isRelRoDataSection(Ctx &ctx, StringRef secName) {
+  // The section name should start with ".data.rel.ro".
+  if (!secName.consume_front(".data.rel.ro"))
+    return false;
+
+  // If the section name is .data.rel.ro, it is a relocatable read-only data
+  // section.
+  if (secName.empty())
+    return true;
+  //  If  -z keep-data-section-prefix is given, '.data.rel.ro.hot' and
+  //  '.data.rel.ro.unlikely' are considered a split of '.data.rel.ro' based on
+  //  hotness.
+  if (ctx.arg.zKeepDataSectionPrefix) {
+    return secName == ".hot" || secName == ".unlikely";
+  }
+
+  return false;
+}
+
 // Today's loaders have a feature to make segments read-only after
 // processing dynamic relocations to enhance security. PT_GNU_RELRO
 // is defined for that.
@@ -629,7 +650,7 @@ static bool isRelroSection(Ctx &ctx, const OutputSection *sec) {
   // magic section names.
   StringRef s = sec->name;
 
-  bool abiAgnostic = s == ".data.rel.ro" || s == ".bss.rel.ro" ||
+  bool abiAgnostic = isRelRoDataSection(ctx, s) || s == ".bss.rel.ro" ||
                      s == ".ctors" || s == ".dtors" || s == ".jcr" ||
                      s == ".eh_frame" || s == ".fini_array" ||
                      s == ".init_array" || s == ".preinit_array";

lld/test/ELF/keep-text-section-prefix.s

@@ -0,0 +1,89 @@
+# REQUIRES: x86
+
+# RUN: rm -rf %t && split-file %s %t && cd %t
+
+# RUN: llvm-mc -filetype=obj -triple=x86_64 a.s -o a.o
+
+# RUN: ld.lld -z keep-data-section-prefix -T x.lds a.o -o out1
+# RUN: llvm-readelf -l out1 | FileCheck --check-prefixes=SEG,LS %s
+# RUN: llvm-readelf -S out1 | FileCheck %s --check-prefix=CHECK-LS
+
+# RUN: ld.lld -z keep-data-section-prefix a.o -o out2
+# RUN: llvm-readelf -l out2 | FileCheck --check-prefixes=SEG,PRE %s
+# RUN: llvm-readelf -S out2 | FileCheck %s --check-prefix=CHECK-PRE
+
+# RUN: ld.lld a.o -o out3
+# RUN: llvm-readelf -l out3 | FileCheck --check-prefixes=SEG,PRE %s
+# RUN: llvm-readelf -S out3 | FileCheck %s --check-prefix=CHECK-PRE
+
+# RUN: not ld.lld -T x.lds a.o 2>&1 | FileCheck %s
+# CHECK: error: section: .relro_padding is not contiguous with other relro sections
+
+## The first RW PT_LOAD segment has FileSiz 0x126f (0x1000 + 0x200 + 0x60 + 0xf),
+## and its p_offset p_vaddr p_paddr p_filesz should match PT_GNU_RELRO.
+#           Type           Offset   VirtAddr           PhysAddr           FileSiz  MemSiz   Flg Align
+# SEG:      LOAD           0x0001c8 0x00000000002011c8 0x00000000002011c8 0x000001 0x000001 R E 0x1000
+# SEG-NEXT: LOAD           0x0001c9 0x00000000002021c9 0x00000000002021c9 0x00126f 0x001e37 RW  0x1000
+# SEG-NEXT: LOAD           0x001438 0x0000000000204438 0x0000000000204438 0x000001 0x000002 RW  0x1000
+# SEG-NEXT: GNU_RELRO      0x0001c9 0x00000000002021c9 0x00000000002021c9 0x00126f 0x001e37 R   0x1
+# SEG-NEXT: GNU_STACK      0x000000 0x0000000000000000 0x0000000000000000 0x000000 0x000000 RW  0x0
+
+## Input to output mapping per linker script
+##   .data.rel.ro.split -> .data.rel.ro
+##   .data.rel.ro -> .data.rel.ro
+##   .data.rel.ro.hot -> .data.rel.ro.hot
+##   .data.rel.ro.unlikely -> .data.rel.ro.unlikely 
+# LS:      .text
+# LS-NEXT: .data.rel.ro.hot .data.rel.ro .data.rel.ro.unlikely .relro_padding
+# LS-NEXT: .data .bss
+
+#        [Nr] Name                    Type            Address          Off    Size
+# CHECK-LS:      .data.rel.ro.hot        PROGBITS        00000000002021c9 0001c9 00000f
+# CHECK-LS-NEXT: .data.rel.ro            PROGBITS        00000000002021d8 0001d8 000260
+# CHECK-LS-NEXT: .data.rel.ro.unlikely   PROGBITS        0000000000202438 000438 001000
+# CHECK-LS-NEXT: .relro_padding          NOBITS          0000000000203438 001438 000bc8
+# CHECK-LS-NEXT: .data                   PROGBITS        0000000000204438 001438 000001
+# CHECK-LS-NEXT: .bss                    NOBITS          0000000000204439 001439 000001
+
+## Linker script is not provided to map data sections.
+## So all input sections with prefix .data.rel.ro will map to .data.rel.ro in the output.
+# PRE:      .text
+# PRE-NEXT: .data.rel.ro .relro_padding
+# PRE-NEXT: .data .bss
+
+#        [Nr] Name                    Type            Address          Off    Size
+# CHECK-PRE:      .data.rel.ro            PROGBITS        00000000002021c9 0001c9 00126f
+# CHECK-PRE-NEXT: .relro_padding          NOBITS          0000000000203438 001438 000bc8
+# CHECK-PRE-NEXT: .data                   PROGBITS        0000000000204438 001438 000001
+# CHECK-PRE-NEXT: .bss                    NOBITS          0000000000204439 001439 000001
+
+#--- x.lds
+SECTIONS {
+  .data.rel.ro.hot : { *(.data.rel.ro.hot) }
+  .data.rel.ro : { .data.rel.ro }
+  .data.rel.ro.unlikely : { *(.data.rel.ro.unlikely) } 
+} INSERT AFTER .text
+
+
+#--- a.s
+.globl _start
+_start:
+  ret
+
+.section .data.rel.ro.hot, "aw"
+.space 15
+
+.section .data.rel.ro, "aw"
+.space 96
+
+.section .data.rel.ro.split,"aw"
+.space 512
+
+.section .data.rel.ro.unlikely, "aw"
+.space 4096
+
+.section .data, "aw"
+.space 1
+
+.section .bss, "aw"
+.space 1