Skip to content

Commit 218fd69

Browse files
maksfbmaksfb
authored
[BOLT] Decouple new segment creation from PHDR rewrite. NFCI (#146111)
Refactor handling of PHDR table rewrite to make modifications easier.
1 parent da01257 commit 218fd69

File tree

5 files changed

+138
-125
lines changed

5 files changed

+138
-125
lines changed

‎bolt/include/bolt/Core/BinaryContext.h

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,14 +73,15 @@ struct SegmentInfo {
7373
uint64_t FileSize; /// Size in file.
7474
uint64_t Alignment; /// Alignment of the segment.
7575
bool IsExecutable; /// Is the executable bit set on the Segment?
76+
bool IsWritable; /// Is the segment writable.
7677

7778
void print(raw_ostream &OS) const {
7879
OS << "SegmentInfo { Address: 0x" << Twine::utohexstr(Address)
7980
<< ", Size: 0x" << Twine::utohexstr(Size) << ", FileOffset: 0x"
8081
<< Twine::utohexstr(FileOffset) << ", FileSize: 0x"
8182
<< Twine::utohexstr(FileSize) << ", Alignment: 0x"
82-
<< Twine::utohexstr(Alignment) << ", " << (IsExecutable ? "x" : " ")
83-
<< "}";
83+
<< Twine::utohexstr(Alignment) << ", " << (IsExecutable ? "x" : "")
84+
<< (IsWritable ? "w" : "") << " }";
8485
};
8586
};
8687

@@ -333,9 +334,14 @@ class BinaryContext {
333334
std::optional<StringRef> Source,
334335
unsigned CUID, unsigned DWARFVersion);
335336

337+
/// Input file segment info
338+
///
336339
/// [start memory address] -> [segment info] mapping.
337340
std::map<uint64_t, SegmentInfo> SegmentMapInfo;
338341

342+
/// Newly created segments.
343+
std::vector<SegmentInfo> NewSegments;
344+
339345
/// Symbols that are expected to be undefined in MCContext during emission.
340346
std::unordered_set<MCSymbol *> UndefinedSymbols;
341347

‎bolt/include/bolt/Rewrite/RewriteInstance.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,9 @@ class RewriteInstance {
300300
return FUNC(ELF64BE); \
301301
}
302302

303+
/// Update loadable segment information based on new sections.
304+
void updateSegmentInfo();
305+
303306
/// Patch ELF book-keeping info.
304307
void patchELFPHDRTable();
305308

‎bolt/lib/Rewrite/RewriteInstance.cpp

Lines changed: 102 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -547,9 +547,14 @@ Error RewriteInstance::discoverStorage() {
547547
NextAvailableOffset = std::max(NextAvailableOffset,
548548
Phdr.p_offset + Phdr.p_filesz);
549549

550-
BC->SegmentMapInfo[Phdr.p_vaddr] = SegmentInfo{
551-
Phdr.p_vaddr, Phdr.p_memsz, Phdr.p_offset,
552-
Phdr.p_filesz, Phdr.p_align, ((Phdr.p_flags & ELF::PF_X) != 0)};
550+
BC->SegmentMapInfo[Phdr.p_vaddr] =
551+
SegmentInfo{Phdr.p_vaddr,
552+
Phdr.p_memsz,
553+
Phdr.p_offset,
554+
Phdr.p_filesz,
555+
Phdr.p_align,
556+
(Phdr.p_flags & ELF::PF_X) != 0,
557+
(Phdr.p_flags & ELF::PF_W) != 0};
553558
if (BC->TheTriple->getArch() == llvm::Triple::x86_64 &&
554559
Phdr.p_vaddr >= BinaryContext::KernelStartX86_64)
555560
BC->IsLinuxKernel = true;
@@ -4182,6 +4187,74 @@ void RewriteInstance::updateOutputValues(const BOLTLinker &Linker) {
41824187
Function->updateOutputValues(Linker);
41834188
}
41844189

4190+
void RewriteInstance::updateSegmentInfo() {
4191+
// NOTE Currently .eh_frame_hdr appends to the last segment, recalculate
4192+
// last segments size based on the NextAvailableAddress variable.
4193+
if (!NewWritableSegmentSize) {
4194+
if (NewTextSegmentAddress)
4195+
NewTextSegmentSize = NextAvailableAddress - NewTextSegmentAddress;
4196+
} else {
4197+
NewWritableSegmentSize = NextAvailableAddress - NewWritableSegmentAddress;
4198+
}
4199+
4200+
if (NewTextSegmentSize) {
4201+
SegmentInfo TextSegment = {NewTextSegmentAddress,
4202+
NewTextSegmentSize,
4203+
NewTextSegmentOffset,
4204+
NewTextSegmentSize,
4205+
BC->PageAlign,
4206+
true,
4207+
false};
4208+
if (!opts::Instrument) {
4209+
BC->NewSegments.push_back(TextSegment);
4210+
} else {
4211+
ErrorOr<BinarySection &> Sec =
4212+
BC->getUniqueSectionByName(".bolt.instr.counters");
4213+
assert(Sec && "expected one and only one `.bolt.instr.counters` section");
4214+
const uint64_t Addr = Sec->getOutputAddress();
4215+
const uint64_t Offset = Sec->getOutputFileOffset();
4216+
const uint64_t Size = Sec->getOutputSize();
4217+
assert(Addr > TextSegment.Address &&
4218+
Addr + Size < TextSegment.Address + TextSegment.Size &&
4219+
"`.bolt.instr.counters` section is expected to be included in the "
4220+
"new text segment");
4221+
4222+
// Set correct size for the previous header since we are breaking the
4223+
// new text segment into three segments.
4224+
uint64_t Delta = Addr - TextSegment.Address;
4225+
TextSegment.Size = Delta;
4226+
TextSegment.FileSize = Delta;
4227+
BC->NewSegments.push_back(TextSegment);
4228+
4229+
// Create RW segment that includes the `.bolt.instr.counters` section.
4230+
SegmentInfo RWSegment = {Addr, Size, Offset, Size, BC->RegularPageSize,
4231+
false, true};
4232+
BC->NewSegments.push_back(RWSegment);
4233+
4234+
// Create RX segment that includes all RX sections from runtime library.
4235+
const uint64_t AddrRX = alignTo(Addr + Size, BC->RegularPageSize);
4236+
const uint64_t OffsetRX = alignTo(Offset + Size, BC->RegularPageSize);
4237+
const uint64_t SizeRX =
4238+
NewTextSegmentSize - (AddrRX - TextSegment.Address);
4239+
SegmentInfo RXSegment = {
4240+
AddrRX, SizeRX, OffsetRX, SizeRX, BC->RegularPageSize, true, false};
4241+
BC->NewSegments.push_back(RXSegment);
4242+
}
4243+
}
4244+
4245+
if (NewWritableSegmentSize) {
4246+
SegmentInfo DataSegmentInfo = {
4247+
NewWritableSegmentAddress,
4248+
NewWritableSegmentSize,
4249+
getFileOffsetForAddress(NewWritableSegmentAddress),
4250+
NewWritableSegmentSize,
4251+
BC->RegularPageSize,
4252+
false,
4253+
true};
4254+
BC->NewSegments.push_back(DataSegmentInfo);
4255+
}
4256+
}
4257+
41854258
void RewriteInstance::patchELFPHDRTable() {
41864259
auto ELF64LEFile = cast<ELF64LEObjectFile>(InputFile);
41874260
const ELFFile<ELF64LE> &Obj = ELF64LEFile->getELFFile();
@@ -4208,107 +4281,36 @@ void RewriteInstance::patchELFPHDRTable() {
42084281
if (opts::Instrument)
42094282
Phnum += 2;
42104283

4211-
// NOTE Currently .eh_frame_hdr appends to the last segment, recalculate
4212-
// last segments size based on the NextAvailableAddress variable.
4213-
if (!NewWritableSegmentSize) {
4214-
if (NewTextSegmentAddress)
4215-
NewTextSegmentSize = NextAvailableAddress - NewTextSegmentAddress;
4216-
} else {
4217-
NewWritableSegmentSize = NextAvailableAddress - NewWritableSegmentAddress;
4218-
}
4219-
4220-
if (!NewTextSegmentSize && !NewWritableSegmentSize) {
4284+
if (BC->NewSegments.empty()) {
42214285
BC->outs() << "BOLT-INFO: not adding new segments\n";
42224286
return;
42234287
}
42244288

42254289
const uint64_t SavedPos = OS.tell();
42264290
OS.seek(PHDRTableOffset);
42274291

4228-
auto createNewPhdrs = [&]() {
4229-
SmallVector<ELF64LEPhdrTy, 3> NewPhdrs;
4230-
ELF64LEPhdrTy NewPhdr;
4231-
NewPhdr.p_type = ELF::PT_LOAD;
4232-
NewPhdr.p_offset = NewTextSegmentOffset;
4233-
NewPhdr.p_vaddr = NewTextSegmentAddress;
4234-
NewPhdr.p_paddr = NewTextSegmentAddress;
4235-
NewPhdr.p_filesz = NewTextSegmentSize;
4236-
NewPhdr.p_memsz = NewTextSegmentSize;
4237-
NewPhdr.p_flags = ELF::PF_X | ELF::PF_R;
4238-
NewPhdr.p_align = BC->PageAlign;
4239-
4240-
if (!opts::Instrument) {
4241-
NewPhdrs.push_back(NewPhdr);
4242-
} else {
4243-
ErrorOr<BinarySection &> Sec =
4244-
BC->getUniqueSectionByName(".bolt.instr.counters");
4245-
assert(Sec && "expected one and only one `.bolt.instr.counters` section");
4246-
const uint64_t Addr = Sec->getOutputAddress();
4247-
const uint64_t Offset = Sec->getOutputFileOffset();
4248-
const uint64_t Size = Sec->getOutputSize();
4249-
assert(Addr > NewPhdr.p_vaddr &&
4250-
Addr + Size < NewPhdr.p_vaddr + NewPhdr.p_memsz &&
4251-
"`.bolt.instr.counters` section is expected to be included in the "
4252-
"new text sgement");
4253-
4254-
// Set correct size for the previous header since we are breaking the
4255-
// new text segment into three segments.
4256-
uint64_t Delta = Addr - NewPhdr.p_vaddr;
4257-
NewPhdr.p_filesz = Delta;
4258-
NewPhdr.p_memsz = Delta;
4259-
NewPhdrs.push_back(NewPhdr);
4260-
4261-
// Create a program header for a RW segment that includes the
4262-
// `.bolt.instr.counters` section only.
4263-
ELF64LEPhdrTy NewPhdrRWSegment;
4264-
NewPhdrRWSegment.p_type = ELF::PT_LOAD;
4265-
NewPhdrRWSegment.p_offset = Offset;
4266-
NewPhdrRWSegment.p_vaddr = Addr;
4267-
NewPhdrRWSegment.p_paddr = Addr;
4268-
NewPhdrRWSegment.p_filesz = Size;
4269-
NewPhdrRWSegment.p_memsz = Size;
4270-
NewPhdrRWSegment.p_flags = ELF::PF_R | ELF::PF_W;
4271-
NewPhdrRWSegment.p_align = BC->RegularPageSize;
4272-
NewPhdrs.push_back(NewPhdrRWSegment);
4273-
4274-
// Create a program header for a RX segment that includes all the RX
4275-
// sections from runtime library.
4276-
ELF64LEPhdrTy NewPhdrRXSegment;
4277-
NewPhdrRXSegment.p_type = ELF::PT_LOAD;
4278-
const uint64_t AddrRX = alignTo(Addr + Size, BC->RegularPageSize);
4279-
const uint64_t OffsetRX = alignTo(Offset + Size, BC->RegularPageSize);
4280-
const uint64_t SizeRX = NewTextSegmentSize - (AddrRX - NewPhdr.p_paddr);
4281-
NewPhdrRXSegment.p_offset = OffsetRX;
4282-
NewPhdrRXSegment.p_vaddr = AddrRX;
4283-
NewPhdrRXSegment.p_paddr = AddrRX;
4284-
NewPhdrRXSegment.p_filesz = SizeRX;
4285-
NewPhdrRXSegment.p_memsz = SizeRX;
4286-
NewPhdrRXSegment.p_flags = ELF::PF_X | ELF::PF_R;
4287-
NewPhdrRXSegment.p_align = BC->RegularPageSize;
4288-
NewPhdrs.push_back(NewPhdrRXSegment);
4289-
}
4290-
4291-
return NewPhdrs;
4292+
auto createPhdr = [](const SegmentInfo &SI) {
4293+
ELF64LEPhdrTy Phdr;
4294+
Phdr.p_type = ELF::PT_LOAD;
4295+
Phdr.p_offset = SI.FileOffset;
4296+
Phdr.p_vaddr = SI.Address;
4297+
Phdr.p_paddr = SI.Address;
4298+
Phdr.p_filesz = SI.FileSize;
4299+
Phdr.p_memsz = SI.Size;
4300+
Phdr.p_flags = ELF::PF_R;
4301+
if (SI.IsExecutable)
4302+
Phdr.p_flags |= ELF::PF_X;
4303+
if (SI.IsWritable)
4304+
Phdr.p_flags |= ELF::PF_W;
4305+
Phdr.p_align = SI.Alignment;
4306+
4307+
return Phdr;
42924308
};
42934309

42944310
auto writeNewSegmentPhdrs = [&]() {
4295-
if (NewTextSegmentSize) {
4296-
SmallVector<ELF64LE::Phdr, 3> NewPhdrs = createNewPhdrs();
4297-
OS.write(reinterpret_cast<const char *>(NewPhdrs.data()),
4298-
sizeof(ELF64LE::Phdr) * NewPhdrs.size());
4299-
}
4300-
4301-
if (NewWritableSegmentSize) {
4302-
ELF64LEPhdrTy NewPhdr;
4303-
NewPhdr.p_type = ELF::PT_LOAD;
4304-
NewPhdr.p_offset = getFileOffsetForAddress(NewWritableSegmentAddress);
4305-
NewPhdr.p_vaddr = NewWritableSegmentAddress;
4306-
NewPhdr.p_paddr = NewWritableSegmentAddress;
4307-
NewPhdr.p_filesz = NewWritableSegmentSize;
4308-
NewPhdr.p_memsz = NewWritableSegmentSize;
4309-
NewPhdr.p_align = BC->RegularPageSize;
4310-
NewPhdr.p_flags = ELF::PF_R | ELF::PF_W;
4311-
OS.write(reinterpret_cast<const char *>(&NewPhdr), sizeof(NewPhdr));
4311+
for (const SegmentInfo &SI : BC->NewSegments) {
4312+
ELF64LEPhdrTy Phdr = createPhdr(SI);
4313+
OS.write(reinterpret_cast<const char *>(&Phdr), sizeof(Phdr));
43124314
}
43134315
};
43144316

@@ -4344,11 +4346,9 @@ void RewriteInstance::patchELFPHDRTable() {
43444346
case ELF::PT_GNU_STACK:
43454347
if (opts::UseGnuStack) {
43464348
// Overwrite the header with the new segment header.
4347-
assert(!opts::Instrument);
4348-
SmallVector<ELF64LE::Phdr, 3> NewPhdrs = createNewPhdrs();
4349-
assert(NewPhdrs.size() == 1 &&
4350-
"expect exactly one program header was created");
4351-
NewPhdr = NewPhdrs[0];
4349+
assert(BC->NewSegments.size() == 1 &&
4350+
"Expected exactly one new segment");
4351+
NewPhdr = createPhdr(BC->NewSegments.front());
43524352
ModdedGnuStack = true;
43534353
}
43544354
break;
@@ -5973,8 +5973,10 @@ void RewriteInstance::rewriteFile() {
59735973
addBATSection();
59745974

59755975
// Patch program header table.
5976-
if (!BC->IsLinuxKernel)
5976+
if (!BC->IsLinuxKernel) {
5977+
updateSegmentInfo();
59775978
patchELFPHDRTable();
5979+
}
59785980

59795981
// Finalize memory image of section string table.
59805982
finalizeSectionStringTable();

‎bolt/unittests/Core/BinaryContext.cpp

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -199,13 +199,13 @@ TEST_P(BinaryContextTester, BaseAddress) {
199199
// Check that base address calculation is correct for a binary with the
200200
// following segment layout:
201201
BC->SegmentMapInfo[0] =
202-
SegmentInfo{0, 0x10e8c2b4, 0, 0x10e8c2b4, 0x1000, true};
203-
BC->SegmentMapInfo[0x10e8d2b4] =
204-
SegmentInfo{0x10e8d2b4, 0x3952faec, 0x10e8c2b4, 0x3952faec, 0x1000, true};
205-
BC->SegmentMapInfo[0x4a3bddc0] =
206-
SegmentInfo{0x4a3bddc0, 0x148e828, 0x4a3bbdc0, 0x148e828, 0x1000, true};
207-
BC->SegmentMapInfo[0x4b84d5e8] =
208-
SegmentInfo{0x4b84d5e8, 0x294f830, 0x4b84a5e8, 0x3d3820, 0x1000, true};
202+
SegmentInfo{0, 0x10e8c2b4, 0, 0x10e8c2b4, 0x1000, true, false};
203+
BC->SegmentMapInfo[0x10e8d2b4] = SegmentInfo{
204+
0x10e8d2b4, 0x3952faec, 0x10e8c2b4, 0x3952faec, 0x1000, true, false};
205+
BC->SegmentMapInfo[0x4a3bddc0] = SegmentInfo{
206+
0x4a3bddc0, 0x148e828, 0x4a3bbdc0, 0x148e828, 0x1000, true, false};
207+
BC->SegmentMapInfo[0x4b84d5e8] = SegmentInfo{
208+
0x4b84d5e8, 0x294f830, 0x4b84a5e8, 0x3d3820, 0x1000, true, false};
209209

210210
std::optional<uint64_t> BaseAddress =
211211
BC->getBaseAddressForMapping(0x7f13f5556000, 0x10e8c000);
@@ -220,13 +220,14 @@ TEST_P(BinaryContextTester, BaseAddress2) {
220220
// Check that base address calculation is correct for a binary if the
221221
// alignment in ELF file are different from pagesize.
222222
// The segment layout is as follows:
223-
BC->SegmentMapInfo[0] = SegmentInfo{0, 0x2177c, 0, 0x2177c, 0x10000, true};
223+
BC->SegmentMapInfo[0] =
224+
SegmentInfo{0, 0x2177c, 0, 0x2177c, 0x10000, true, false};
224225
BC->SegmentMapInfo[0x31860] =
225-
SegmentInfo{0x31860, 0x370, 0x21860, 0x370, 0x10000, true};
226+
SegmentInfo{0x31860, 0x370, 0x21860, 0x370, 0x10000, true, false};
226227
BC->SegmentMapInfo[0x41c20] =
227-
SegmentInfo{0x41c20, 0x1f8, 0x21c20, 0x1f8, 0x10000, true};
228+
SegmentInfo{0x41c20, 0x1f8, 0x21c20, 0x1f8, 0x10000, true, false};
228229
BC->SegmentMapInfo[0x54e18] =
229-
SegmentInfo{0x54e18, 0x51, 0x24e18, 0x51, 0x10000, true};
230+
SegmentInfo{0x54e18, 0x51, 0x24e18, 0x51, 0x10000, true, false};
230231

231232
std::optional<uint64_t> BaseAddress =
232233
BC->getBaseAddressForMapping(0xaaaaea444000, 0x21000);
@@ -242,13 +243,14 @@ TEST_P(BinaryContextTester, BaseAddressSegmentsSmallerThanAlignment) {
242243
// when multiple segments are close together in the ELF file (closer
243244
// than the required alignment in the process space).
244245
// See https://github.com/llvm/llvm-project/issues/109384
245-
BC->SegmentMapInfo[0] = SegmentInfo{0, 0x1d1c, 0, 0x1d1c, 0x10000, false};
246+
BC->SegmentMapInfo[0] =
247+
SegmentInfo{0, 0x1d1c, 0, 0x1d1c, 0x10000, false, false};
246248
BC->SegmentMapInfo[0x11d40] =
247-
SegmentInfo{0x11d40, 0x11e0, 0x1d40, 0x11e0, 0x10000, true};
249+
SegmentInfo{0x11d40, 0x11e0, 0x1d40, 0x11e0, 0x10000, true, false};
248250
BC->SegmentMapInfo[0x22f20] =
249-
SegmentInfo{0x22f20, 0x10e0, 0x2f20, 0x1f0, 0x10000, false};
251+
SegmentInfo{0x22f20, 0x10e0, 0x2f20, 0x1f0, 0x10000, false, false};
250252
BC->SegmentMapInfo[0x33110] =
251-
SegmentInfo{0x33110, 0x89, 0x3110, 0x88, 0x10000, false};
253+
SegmentInfo{0x33110, 0x89, 0x3110, 0x88, 0x10000, false, false};
252254

253255
std::optional<uint64_t> BaseAddress =
254256
BC->getBaseAddressForMapping(0xaaaaaaab1000, 0x1000);

‎bolt/unittests/Core/MemoryMaps.cpp

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -100,10 +100,10 @@ TEST_P(MemoryMapsTester, ParseMultipleSegments) {
100100
"[0xabc2000000(0x8000000) @ 0x31d0000 103:01 1573523 0]: r-xp {1}\n",
101101
Pid, Filename);
102102

103-
BC->SegmentMapInfo[0x11da000] =
104-
SegmentInfo{0x11da000, 0x10da000, 0x11ca000, 0x10da000, 0x10000, true};
105-
BC->SegmentMapInfo[0x31d0000] =
106-
SegmentInfo{0x31d0000, 0x51ac82c, 0x31d0000, 0x3000000, 0x200000, true};
103+
BC->SegmentMapInfo[0x11da000] = SegmentInfo{
104+
0x11da000, 0x10da000, 0x11ca000, 0x10da000, 0x10000, true, false};
105+
BC->SegmentMapInfo[0x31d0000] = SegmentInfo{
106+
0x31d0000, 0x51ac82c, 0x31d0000, 0x3000000, 0x200000, true, false};
107107

108108
DataAggregator DA("");
109109
BC->setFilename(Filename);
@@ -131,12 +131,12 @@ TEST_P(MemoryMapsTester, MultipleSegmentsMismatchedBaseAddress) {
131131
"[0xabc2000000(0x8000000) @ 0x31d0000 103:01 1573523 0]: r-xp {1}\n",
132132
Pid, Filename);
133133

134-
BC->SegmentMapInfo[0x11da000] =
135-
SegmentInfo{0x11da000, 0x10da000, 0x11ca000, 0x10da000, 0x10000, true};
134+
BC->SegmentMapInfo[0x11da000] = SegmentInfo{
135+
0x11da000, 0x10da000, 0x11ca000, 0x10da000, 0x10000, true, false};
136136
// Using '0x31d0fff' FileOffset which triggers a different base address
137137
// for this second text segment.
138-
BC->SegmentMapInfo[0x31d0000] =
139-
SegmentInfo{0x31d0000, 0x51ac82c, 0x31d0fff, 0x3000000, 0x200000, true};
138+
BC->SegmentMapInfo[0x31d0000] = SegmentInfo{
139+
0x31d0000, 0x51ac82c, 0x31d0fff, 0x3000000, 0x200000, true, false};
140140

141141
DataAggregator DA("");
142142
BC->setFilename(Filename);

0 commit comments

Comments
 (0)