xref: /aosp_15_r20/external/angle/build/chromeos/embed_sections.py (revision 8975f5c5ed3d1c378011245431ada316dfb6f244) 
1#!/usr/bin/env vpython3
2#
3# Copyright 2024 The Chromium Authors
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6'''
7Embed the address and size of elf sections into the pre-defined symbols.
8
9The embedded values are used for performance optimization by mlock(2)ing the
10sections on ChromeOS. See chromeos/ash/components/memory/memory.cc for details.
11'''
12
13import argparse
14import os
15import subprocess
16import sys
17import shutil
18
19llvm_readelf = os.path.join(
20    os.path.dirname(sys.argv[0]), '..', '..', 'third_party', 'llvm-build',
21    'Release+Asserts', 'bin', 'llvm-readelf')
22
23TARGET_SECTIONS = {
24    '.rodata': {
25        'addr': 'kRodataAddr',
26        'size': 'kRodataSize'
27    },
28    '.text.hot': {
29        'addr': 'kTextHotAddr',
30        'size': 'kTextHotSize'
31    },
32}
33
34
35def parse_endianess(objdump_result):
36  for line in objdump_result.splitlines():
37    line = line.strip()
38    if line.startswith('Data:'):
39      if '1' in line:
40        return 'big'
41      if '2' in line:
42        return 'little'
43  raise ValueError('No endian found')
44
45
46def assert_elf_type(objdump_result):
47  for line in objdump_result.splitlines():
48    line = line.strip()
49    if line.startswith('Class:'):
50      if 'ELF64' in line or 'ELF32' in line:
51        return
52      raise ValueError('Class is not ELF64 nor ELF32: ' + line)
53  raise ValueError('No class found')
54
55
56def parse_section_info(objdump_result, section_name):
57  for line in objdump_result.splitlines():
58    row = line.strip().split()
59    if len(row) < 2:
60      continue
61    if row[1] == section_name:
62      # 3: Address, 4: Offset, 5: Size
63      return (int(row[3], base=16), int(row[4], base=16), int(row[5], base=16))
64  return (0, 0, 0)
65
66
67def create_symbol_map(binary_input, objdump_result):
68  (rodata_section_addr, rodata_section_offset,
69   rodata_section_size) = parse_section_info(objdump_result, '.rodata')
70
71  command = [llvm_readelf, '--symbols', binary_input]
72  with subprocess.Popen(command, stdout=subprocess.PIPE, text=True) as process:
73
74    variable_names = [
75        var for maps in TARGET_SECTIONS.values() for var in maps.values()
76    ]
77    result = {}
78    while len(result) < len(variable_names):
79      line = process.stdout.readline()
80      if not line:
81        break
82      for var in variable_names:
83        if var in line:
84          row = line.strip().split()
85          if row[2] == '8':
86            size = 8
87          elif row[2] == '4':
88            size = 4
89          else:
90            raise ValueError('variable size is not 8 or 4: ' + line)
91          addr = int(row[1], base=16)
92          rodata_section_end_addr = rodata_section_addr + rodata_section_size
93          if addr < rodata_section_addr or addr >= rodata_section_end_addr:
94            raise ValueError(var + ' is not in .rodata section')
95          offset = addr - rodata_section_addr + rodata_section_offset
96          result[var] = (offset, size)
97
98  return result
99
100
101def overwrite_variable(file, symbol_map, endianess, section_name, var_type,
102                       value):
103  (var_offset, var_size) = symbol_map[TARGET_SECTIONS[section_name][var_type]]
104  file.seek(var_offset)
105  if file.write(
106      value.to_bytes(length=var_size, byteorder=endianess,
107                     signed=False)) != var_size:
108    raise ValueError('failed to write value to file')
109
110
111def main():
112  argparser = argparse.ArgumentParser(
113      description='embed sections informataion into binary.')
114
115  argparser.add_argument('--binary-input', help='exe file path.')
116  argparser.add_argument('--binary-output', help='embedded file path.')
117  args = argparser.parse_args()
118
119  objdump_result = subprocess.run([llvm_readelf, '-e', args.binary_input],
120                                  stdout=subprocess.PIPE,
121                                  check=True,
122                                  text=True).stdout
123
124  assert_elf_type(objdump_result)
125
126  symbol_map = create_symbol_map(args.binary_input, objdump_result)
127  if len(symbol_map) != len(set(symbol_map.values())):
128    raise ValueError(f'symbol_map overlaps: {symbol_map}')
129
130  endianess = parse_endianess(objdump_result)
131
132  shutil.copyfile(args.binary_input, args.binary_output)
133
134  with open(args.binary_output, 'r+b') as file:
135    for section_name in TARGET_SECTIONS:
136      (addr, _, size) = parse_section_info(objdump_result, section_name)
137      overwrite_variable(file, symbol_map, endianess, section_name, 'addr',
138                         addr)
139      overwrite_variable(file, symbol_map, endianess, section_name, 'size',
140                         size)
141
142  objdump_result_after = subprocess.run(
143      [llvm_readelf, '-e', args.binary_output],
144      stdout=subprocess.PIPE,
145      check=True,
146      text=True).stdout
147  if objdump_result_after != objdump_result:
148    raise ValueError('realelf result has changed')
149
150  return 0
151
152
153if __name__ == '__main__':
154  sys.exit(main())
155