#!/usr/bin/env vpython3
#
# Copyright 2024 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
'''
Embed the address and size of elf sections into the pre-defined symbols.

The embedded values are used for performance optimization by mlock(2)ing the
sections on ChromeOS. See chromeos/ash/components/memory/memory.cc for details.
'''

import argparse
import os
import subprocess
import sys
import shutil

llvm_readelf = os.path.join(
  os.path.dirname(sys.argv[0]),
  '..',
  '..',
  'third_party',
  'llvm-build',
  'Release+Asserts',
  'bin',
  'llvm-readelf',
)

TARGET_SECTIONS = {
  '.rodata': {'addr': 'kRodataAddr', 'size': 'kRodataSize'},
  '.text.hot': {'addr': 'kTextHotAddr', 'size': 'kTextHotSize'},
}


def parse_endianess(objdump_result):
  for line in objdump_result.splitlines():
    line = line.strip()
    if line.startswith('Data:'):
      if '1' in line:
        return 'big'
      if '2' in line:
        return 'little'
  raise ValueError('No endian found')


def assert_elf_type(objdump_result):
  for line in objdump_result.splitlines():
    line = line.strip()
    if line.startswith('Class:'):
      if 'ELF64' in line or 'ELF32' in line:
        return
      raise ValueError('Class is not ELF64 nor ELF32: ' + line)
  raise ValueError('No class found')


def parse_section_info(objdump_result, section_name):
  for line in objdump_result.splitlines():
    row = line.strip().split()
    if len(row) < 2:
      continue
    if row[1] == section_name:
      # 3: Address, 4: Offset, 5: Size
      return (int(row[3], base=16), int(row[4], base=16), int(row[5], base=16))
  return (0, 0, 0)


def create_symbol_map(binary_input, objdump_result):
  (rodata_section_addr, rodata_section_offset, rodata_section_size) = (
    parse_section_info(objdump_result, '.rodata')
  )

  command = [llvm_readelf, '--symbols', binary_input]
  with subprocess.Popen(command, stdout=subprocess.PIPE, text=True) as process:
    variable_names = [
      var for maps in TARGET_SECTIONS.values() for var in maps.values()
    ]
    result = {}
    while len(result) < len(variable_names):
      line = process.stdout.readline()
      if not line:
        break
      for var in variable_names:
        if var in line:
          row = line.strip().split()
          if row[2] == '8':
            size = 8
          elif row[2] == '4':
            size = 4
          else:
            raise ValueError('variable size is not 8 or 4: ' + line)
          addr = int(row[1], base=16)
          rodata_section_end_addr = rodata_section_addr + rodata_section_size
          if addr < rodata_section_addr or addr >= rodata_section_end_addr:
            raise ValueError(var + ' is not in .rodata section')
          offset = addr - rodata_section_addr + rodata_section_offset
          result[var] = (offset, size)

  return result


def overwrite_variable(
  file, symbol_map, endianess, section_name, var_type, value
):
  (var_offset, var_size) = symbol_map[TARGET_SECTIONS[section_name][var_type]]
  file.seek(var_offset)
  if (
    file.write(
      value.to_bytes(length=var_size, byteorder=endianess, signed=False)
    )
    != var_size
  ):
    raise ValueError('failed to write value to file')


def main():
  argparser = argparse.ArgumentParser(
    description='embed sections informataion into binary.'
  )

  argparser.add_argument('--binary-input', help='exe file path.')
  argparser.add_argument('--binary-output', help='embedded file path.')
  args = argparser.parse_args()

  objdump_result = subprocess.run(
    [llvm_readelf, '-e', args.binary_input],
    stdout=subprocess.PIPE,
    check=True,
    text=True,
  ).stdout

  assert_elf_type(objdump_result)

  symbol_map = create_symbol_map(args.binary_input, objdump_result)
  if len(symbol_map) != len(set(symbol_map.values())):
    raise ValueError(f'symbol_map overlaps: {symbol_map}')

  endianess = parse_endianess(objdump_result)

  shutil.copyfile(args.binary_input, args.binary_output)

  with open(args.binary_output, 'r+b') as file:
    for section_name in TARGET_SECTIONS:
      (addr, _, size) = parse_section_info(objdump_result, section_name)
      overwrite_variable(
        file, symbol_map, endianess, section_name, 'addr', addr
      )
      overwrite_variable(
        file, symbol_map, endianess, section_name, 'size', size
      )

  objdump_result_after = subprocess.run(
    [llvm_readelf, '-e', args.binary_output],
    stdout=subprocess.PIPE,
    check=True,
    text=True,
  ).stdout
  if objdump_result_after != objdump_result:
    raise ValueError('realelf result has changed')

  return 0


if __name__ == '__main__':
  sys.exit(main())