#!/usr/bin/env python3
# Copyright 2017 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Checks the number of static initializers in an APK's library."""

import argparse
import os
import re
import subprocess
import sys

from util import build_utils

_DUMP_STATIC_INITIALIZERS_PATH = os.path.join(
    build_utils.DIR_SOURCE_ROOT, 'tools', 'linux', 'dump-static-initializers.py'
)


def _RunReadelf(so_path, options, tool_prefix=''):
    return subprocess.check_output(
        [tool_prefix + 'readobj', '--elf-output-style=GNU']
        + options
        + [so_path],
        encoding='utf-8',
    )


def _DumpStaticInitializers(so_path):
    subprocess.check_call([_DUMP_STATIC_INITIALIZERS_PATH, so_path])


def _ReadInitArray(so_path, tool_prefix):
    stdout = _RunReadelf(so_path, ['-SW'], tool_prefix)
    # Matches: .init_array INIT_ARRAY 000000000516add0 5169dd0 000010 00 WA 0 0 8
    match = re.search(r'\.init_array.*$', stdout, re.MULTILINE)
    if not match:
        raise Exception(
            'Did not find section: .init_array in {}:\n{}'.format(
                so_path, stdout
            )
        )
    size_str = re.split(r'\W+', match.group(0))[5]
    return int(size_str, 16)


def _CountStaticInitializers(so_path, tool_prefix):
    # Find the number of files with at least one static initializer.
    # First determine if we're 32 or 64 bit
    stdout = _RunReadelf(so_path, ['-h'], tool_prefix)
    elf_class_line = re.search('Class:.*$', stdout, re.MULTILINE).group(0)
    elf_class = re.split(r'\W+', elf_class_line)[1]
    if elf_class == 'ELF32':
        word_size = 4
    else:
        word_size = 8

    # Then find the number of files with global static initializers.
    # NOTE: this is very implementation-specific and makes assumptions
    # about how compiler and linker implement global static initializers.
    init_array_size = _ReadInitArray(so_path, tool_prefix)
    assert init_array_size % word_size == 0
    return init_array_size // word_size


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--touch', help='File to touch upon success')
    parser.add_argument(
        '--tool-prefix', required=True, help='Prefix for nm and friends'
    )
    parser.add_argument(
        '--expected-count',
        required=True,
        type=int,
        help='Fail if number of static initializers is not '
        'equal to this value.',
    )
    parser.add_argument(
        '--unstripped-so-path',
        help='Path to the unstripped version of the .so '
        'file if needed for better dumps.',
    )
    parser.add_argument('so_path', help='Path to .so file.')
    args = parser.parse_args()

    si_count = _CountStaticInitializers(args.so_path, args.tool_prefix)
    if si_count != args.expected_count:
        print(
            'Expected {} static initializers, but found {}.'.format(
                args.expected_count, si_count
            )
        )
        if args.expected_count > si_count:
            print('You have removed one or more static initializers. Thanks!')
            print('To fix the build, update the expectation in:')
            print('    //chrome/android/static_initializers.gni')
            print()

        print('Dumping static initializers via dump-static-initializers.py:')
        sys.stdout.flush()
        dump_so_path = args.so_path
        if args.unstripped_so_path:
            dump_so_path = args.unstripped_so_path
        _DumpStaticInitializers(dump_so_path)
        print()
        print('For more information:')
        print(
            '    https://chromium.googlesource.com/chromium/src/+/main/docs/'
            'static_initializers.md'
        )
        sys.exit(1)

    if args.touch:
        with open(args.touch, 'w', encoding='utf-8'):
            pass


if __name__ == '__main__':
    main()