1#!/usr/bin/python3
2# ELF editor for load align tests.
3# Copyright (C) 2022 Free Software Foundation, Inc.
4# Copyright The GNU Toolchain Authors.
5# This file is part of the GNU C Library.
6#
7# The GNU C Library is free software; you can redistribute it and/or
8# modify it under the terms of the GNU Lesser General Public
9# License as published by the Free Software Foundation; either
10# version 2.1 of the License, or (at your option) any later version.
11#
12# The GNU C Library is distributed in the hope that it will be useful,
13# but WITHOUT ANY WARRANTY; without even the implied warranty of
14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15# Lesser General Public License for more details.
16#
17# You should have received a copy of the GNU Lesser General Public
18# License along with the GNU C Library; if not, see
19# <https://www.gnu.org/licenses/>.
20
21import argparse
22import os
23import sys
24import struct
25
26EI_NIDENT=16
27
28EI_MAG0=0
29ELFMAG0=b'\x7f'
30EI_MAG1=1
31ELFMAG1=b'E'
32EI_MAG2=2
33ELFMAG2=b'L'
34EI_MAG3=3
35ELFMAG3=b'F'
36
37EI_CLASS=4
38ELFCLASSNONE=b'0'
39ELFCLASS32=b'\x01'
40ELFCLASS64=b'\x02'
41
42EI_DATA=5
43ELFDATA2LSB=b'\x01'
44ELFDATA2MSB=b'\x02'
45
46ET_EXEC=2
47ET_DYN=3
48
49PT_LOAD=1
50PT_TLS=7
51
52def elf_types_fmts(e_ident):
53    endian = '<' if e_ident[EI_DATA] == ELFDATA2LSB else '>'
54    addr = 'I' if e_ident[EI_CLASS] == ELFCLASS32 else 'Q'
55    off = 'I' if e_ident[EI_CLASS] == ELFCLASS32 else 'Q'
56    return (endian, addr, off)
57
58class Elf_Ehdr:
59    def __init__(self, e_ident):
60        endian, addr, off = elf_types_fmts(e_ident)
61        self.fmt = '{0}HHI{1}{2}{2}IHHHHHH'.format(endian, addr, off)
62        self.len = struct.calcsize(self.fmt)
63
64    def read(self, f):
65        buf = f.read(self.len)
66        if not buf:
67            error('{}: header too small'.format(f.name))
68        data = struct.unpack(self.fmt, buf)
69        self.e_type = data[0]
70        self.e_machine = data[1]
71        self.e_version = data[2]
72        self.e_entry = data[3]
73        self.e_phoff = data[4]
74        self.e_shoff = data[5]
75        self.e_flags = data[6]
76        self.e_ehsize = data[7]
77        self.e_phentsize= data[8]
78        self.e_phnum = data[9]
79        self.e_shstrndx = data[10]
80
81
82class Elf_Phdr:
83    def __init__(self, e_ident):
84        endian, addr, off = elf_types_fmts(e_ident)
85        self.ei_class = e_ident[EI_CLASS]
86        if self.ei_class == ELFCLASS32:
87            self.fmt = '{0}I{2}{1}{1}IIII'.format(endian, addr, off)
88        else:
89            self.fmt = '{0}II{2}{1}{1}QQQ'.format(endian, addr, off)
90        self.len = struct.calcsize(self.fmt)
91
92    def read(self, f):
93        buf = f.read(self.len)
94        if len(buf) < self.len:
95            error('{}: program header too small'.format(f.name))
96        data = struct.unpack(self.fmt, buf)
97        if self.ei_class == ELFCLASS32:
98            self.p_type = data[0]
99            self.p_offset = data[1]
100            self.p_vaddr = data[2]
101            self.p_paddr = data[3]
102            self.p_filesz = data[4]
103            self.p_memsz = data[5]
104            self.p_flags = data[6]
105            self.p_align = data[7]
106        else:
107            self.p_type = data[0]
108            self.p_flags = data[1]
109            self.p_offset = data[2]
110            self.p_vaddr = data[3]
111            self.p_paddr = data[4]
112            self.p_filesz = data[5]
113            self.p_memsz = data[6]
114            self.p_align = data[7]
115
116    def write(self, f):
117        if self.ei_class == ELFCLASS32:
118            data = struct.pack(self.fmt,
119                               self.p_type,
120                               self.p_offset,
121                               self.p_vaddr,
122                               self.p_paddr,
123                               self.p_filesz,
124                               self.p_memsz,
125                               self.p_flags,
126                               self.p_align)
127        else:
128            data = struct.pack(self.fmt,
129                               self.p_type,
130                               self.p_flags,
131                               self.p_offset,
132                               self.p_vaddr,
133                               self.p_paddr,
134                               self.p_filesz,
135                               self.p_memsz,
136                               self.p_align)
137        f.write(data)
138
139
140def error(msg):
141    print(msg, file=sys.stderr)
142    sys.exit(1)
143
144
145def elf_edit_align(phdr, align):
146    if align == 'half':
147        phdr.p_align = phdr.p_align >> 1
148    else:
149        phdr.p_align = int(align)
150
151def elf_edit_maximize_tls_size(phdr, elfclass):
152    if elfclass == ELFCLASS32:
153        # It is possible that the kernel can allocate half of the
154        # address space, so use something larger.
155        phdr.p_memsz = 0xfff00000
156    else:
157        phdr.p_memsz = 1 << 63
158
159def elf_edit(f, opts):
160    ei_nident_fmt = 'c' * EI_NIDENT
161    ei_nident_len = struct.calcsize(ei_nident_fmt)
162
163    data = f.read(ei_nident_len)
164    if len(data) < ei_nident_len:
165      error('{}: e_nident too small'.format(f.name))
166    e_ident = struct.unpack(ei_nident_fmt, data)
167
168    if e_ident[EI_MAG0] != ELFMAG0 \
169       or e_ident[EI_MAG1] != ELFMAG1 \
170       or e_ident[EI_MAG2] != ELFMAG2 \
171       or e_ident[EI_MAG3] != ELFMAG3:
172      error('{}: bad ELF header'.format(f.name))
173
174    if e_ident[EI_CLASS] != ELFCLASS32 \
175       and e_ident[EI_CLASS] != ELFCLASS64:
176      error('{}: unsupported ELF class: {}'.format(f.name, e_ident[EI_CLASS]))
177
178    if e_ident[EI_DATA] != ELFDATA2LSB \
179       and e_ident[EI_DATA] != ELFDATA2MSB: \
180      error('{}: unsupported ELF data: {}'.format(f.name, e_ident[EI_DATA]))
181
182    ehdr = Elf_Ehdr(e_ident)
183    ehdr.read(f)
184    if ehdr.e_type not in (ET_EXEC, ET_DYN):
185       error('{}: not an executable or shared library'.format(f.name))
186
187    phdr = Elf_Phdr(e_ident)
188    maximize_tls_size_done = False
189    for i in range(0, ehdr.e_phnum):
190        f.seek(ehdr.e_phoff + i * phdr.len)
191        phdr.read(f)
192        if phdr.p_type == PT_LOAD and opts.align is not None:
193            elf_edit_align(phdr, opts.align)
194            f.seek(ehdr.e_phoff + i * phdr.len)
195            phdr.write(f)
196            break
197        if phdr.p_type == PT_TLS and opts.maximize_tls_size:
198            elf_edit_maximize_tls_size(phdr, e_ident[EI_CLASS])
199            f.seek(ehdr.e_phoff + i * phdr.len)
200            phdr.write(f)
201            maximize_tls_size_done = True
202            break
203
204    if opts.maximize_tls_size and not maximize_tls_size_done:
205        error('{}: TLS maximum size was not updated'.format(f.name))
206
207def get_parser():
208    parser = argparse.ArgumentParser(description=__doc__)
209    parser.add_argument('-a', dest='align',
210                        help='How to set the LOAD alignment')
211    parser.add_argument('--maximize-tls-size', action='store_true',
212                        help='Set maximum PT_TLS size')
213    parser.add_argument('output',
214                        help='ELF file to edit')
215    return parser
216
217
218def main(argv):
219    parser = get_parser()
220    opts = parser.parse_args(argv)
221    with open(opts.output, 'r+b') as fout:
222       elf_edit(fout, opts)
223
224
225if __name__ == '__main__':
226    main(sys.argv[1:])
227