drgn/tests/elfwriter.py
Omar Sandoval 104a14781d tests: test compressed debug sections
Test both .zdebug_* sections and SHF_COMPRESSED .debug_* sections.

Signed-off-by: Omar Sandoval <osandov@osandov.com>
2023-06-20 13:45:04 -07:00

279 lines
8.9 KiB
Python

# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later
import struct
from typing import List, NamedTuple, Optional, Sequence
import zlib
from tests.elf import ET, PT, SHF, SHN, SHT, STB, STT, STV
class ElfSection:
def __init__(
self,
data: bytes = b"",
name: Optional[str] = None,
sh_type: Optional[SHT] = None,
p_type: Optional[PT] = None,
vaddr: int = 0,
paddr: int = 0,
memsz: Optional[int] = None,
p_align: int = 0,
sh_link: int = 0,
sh_info: int = 0,
sh_entsize: int = 0,
compressed=False,
):
self.data = data
self.name = name
self.sh_type = sh_type
self.sh_flags = SHF.COMPRESSED if compressed else 0
self.p_type = p_type
self.vaddr = vaddr
self.paddr = paddr
self.memsz = memsz
self.p_align = p_align
self.sh_link = sh_link
self.sh_info = sh_info
self.sh_entsize = sh_entsize
assert (self.name is not None) or (self.p_type is not None)
assert (self.name is None) == (self.sh_type is None)
assert self.p_type is None or not compressed
class ElfSymbol(NamedTuple):
name: str
value: int
size: int
type: STT
binding: STB
shindex: Optional[int] = None
visibility: STV = STV.DEFAULT
def st_info(self) -> int:
return (self.binding << 4) + (self.type & 0xF)
def _create_symtab(
sections: List[ElfSection],
symbols: Sequence[ElfSymbol],
little_endian: bool,
bits: int,
):
assert not any(section.name in (".symtab", ".strtab") for section in sections)
endian = "<" if little_endian else ">"
if bits == 64:
symbol_struct = struct.Struct(endian + "IBBHQQ")
def symbol_fields(sym: ElfSymbol):
return (
sym.st_info(),
sym.visibility,
SHN.UNDEF if sym.shindex is None else sym.shindex,
sym.value,
sym.size,
)
else:
symbol_struct = struct.Struct(endian + "IIIBBH")
def symbol_fields(sym: ElfSymbol):
return (
sym.value,
sym.size,
sym.st_info(),
sym.visibility,
SHN.UNDEF if sym.shindex is None else sym.shindex,
)
symtab_data = bytearray((len(symbols) + 1) * symbol_struct.size)
strtab_data = bytearray(1)
sh_info = 1
for i, sym in enumerate(symbols, 1):
symbol_struct.pack_into(
symtab_data, i * symbol_struct.size, len(strtab_data), *symbol_fields(sym)
)
strtab_data.extend(sym.name.encode())
strtab_data.append(0)
if sym.binding == STB.LOCAL:
assert sh_info == i, "local symbol after non-local symbol"
sh_info = i + 1
sections.append(
ElfSection(
name=".symtab",
sh_type=SHT.SYMTAB,
data=symtab_data,
sh_link=sum((1 for section in sections if section.name is not None), 2),
sh_info=sh_info,
sh_entsize=symbol_struct.size,
)
)
sections.append(ElfSection(name=".strtab", sh_type=SHT.STRTAB, data=strtab_data))
def create_elf_file(
type: ET,
sections: Sequence[ElfSection],
symbols: Sequence[ElfSymbol] = (),
little_endian: bool = True,
bits: int = 64,
):
endian = "<" if little_endian else ">"
if bits == 64:
ehdr_struct = struct.Struct(endian + "16BHHIQQQIHHHHHH")
shdr_struct = struct.Struct(endian + "IIQQQQIIQQ")
phdr_struct = struct.Struct(endian + "IIQQQQQQ")
chdr_struct = struct.Struct(endian + "IIQQ")
e_machine = 62 if little_endian else 43 # EM_X86_64 or EM_SPARCV9
else:
assert bits == 32
ehdr_struct = struct.Struct(endian + "16BHHIIIIIHHHHHH")
shdr_struct = struct.Struct(endian + "10I")
phdr_struct = struct.Struct(endian + "8I")
chdr_struct = struct.Struct(endian + "III")
e_machine = 3 if little_endian else 8 # EM_386 or EM_MIPS
sections = list(sections)
if symbols:
_create_symtab(sections, symbols, little_endian=little_endian, bits=bits)
shnum = 0
phnum = 0
shstrtab = bytearray(1)
for section in sections:
if section.name is not None:
shstrtab.extend(section.name.encode())
shstrtab.append(0)
shnum += 1
if section.p_type is not None:
phnum += 1
if shnum > 0:
shnum += 2 # One for the SHT_NULL section, one for .shstrtab.
shstrtab.extend(b".shstrtab\0")
sections = list(sections)
sections.append(ElfSection(name=".shstrtab", sh_type=SHT.STRTAB, data=shstrtab))
shdr_offset = ehdr_struct.size
phdr_offset = shdr_offset + shdr_struct.size * shnum
headers_size = phdr_offset + phdr_struct.size * phnum
buf = bytearray(headers_size)
ehdr_struct.pack_into(
buf,
0,
0x7F, # ELFMAG0
ord("E"), # ELFMAG1
ord("L"), # ELFMAG2
ord("F"), # ELFMAG3
2 if bits == 64 else 1, # EI_CLASS = ELFCLASS64 or ELFCLASS32
1 if little_endian else 2, # EI_DATA = ELFDATA2LSB or ELFDATA2MSB
1, # EI_VERSION = EV_CURRENT
0, # EI_OSABI = ELFOSABI_NONE
0, # EI_ABIVERSION
0,
0,
0,
0,
0,
0,
0, # EI_PAD
type, # e_type
e_machine,
1, # e_version = EV_CURRENT
0, # e_entry
phdr_offset if phnum else 0, # e_phoff
shdr_offset if shnum else 0, # e_shoff
0, # e_flags
ehdr_struct.size, # e_ehsize
phdr_struct.size, # e_phentsize
phnum, # e_phnum
shdr_struct.size, # e_shentsize
shnum, # e_shnum
shnum - 1 if shnum else 0, # e_shstrndx
)
shdr_offset += shdr_struct.size
for section in sections:
ch_addralign = 1 if section.p_type is None else bits // 8
memsz = len(section.data) if section.memsz is None else section.memsz
if section.sh_flags & SHF.COMPRESSED:
sh_addralign = bits // 8
compressed_data = zlib.compress(section.data)
sh_size = chdr_struct.size + len(compressed_data)
else:
sh_addralign = ch_addralign
sh_size = memsz
if section.p_align:
padding = section.vaddr % section.p_align - len(buf) % section.p_align
buf.extend(bytes(padding))
if section.name is not None:
shdr_struct.pack_into(
buf,
shdr_offset,
shstrtab.index(section.name.encode()), # sh_name
section.sh_type, # sh_type
section.sh_flags, # sh_flags
section.vaddr, # sh_addr
len(buf), # sh_offset
sh_size, # sh_size
section.sh_link, # sh_link
section.sh_info, # sh_info
sh_addralign, # sh_addralign
section.sh_entsize, # sh_entsize
)
shdr_offset += shdr_struct.size
if section.p_type is not None:
flags = 7 # PF_R | PF_W | PF_X
if bits == 64:
phdr_struct.pack_into(
buf,
phdr_offset,
section.p_type, # p_type
flags, # p_flags
len(buf), # p_offset
section.vaddr, # p_vaddr
section.paddr, # p_paddr
len(section.data), # p_filesz
memsz, # p_memsz
section.p_align, # p_align
)
else:
phdr_struct.pack_into(
buf,
phdr_offset,
section.p_type, # p_type
len(buf), # p_offset
section.vaddr, # p_vaddr
section.paddr, # p_paddr
len(section.data), # p_filesz
memsz, # p_memsz
flags, # p_flags
section.p_align, # p_align
)
phdr_offset += phdr_struct.size
if section.sh_flags & SHF.COMPRESSED:
ELFCOMPRESS_ZLIB = 1
if bits == 64:
buf.extend(
chdr_struct.pack(
ELFCOMPRESS_ZLIB, # ch_type
0, # ch_reserved
memsz, # ch_size
ch_addralign, # ch_addralign
)
)
else:
buf.extend(
chdr_struct.pack(
ELFCOMPRESS_ZLIB, # ch_type
memsz, # ch_size
ch_addralign, # ch_addralign
)
)
buf.extend(compressed_data)
else:
buf.extend(section.data)
return buf