#!/usr/bin/env fuchsia-vendored-python
# Copyright 2021 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""
Script for checking that all vtables in an elf file are in rodata. This script
provides different "modes" for quick checking of different files/project
layouts.

Usage (append -h for details):

  # Check an individual file.
  $ ./scripts/clang/check-vtable-rodata.py file [filename]

  # Check all vtables in binaries provided by the fuchsia toolchain
  $ ./scripts/clang/check-vtable-rodata.py toolchain [toolchain_dir]

  # Check all vtables in binaries provided by an llvm multilib directory.
  $ ./scripts/clang/check-vtable-rodata.py multilib [multilib_dir]

  # Check all vtables in binaries provided in a fuchsia build directory.
  $ ./scripts/clang/check-vtable-rodata.py fuchsia [fuchsia_dir]

  # Exclude checking specific vtable symbols
  $ ./scripts/clang/check-vtable-rodata.py --exclude 'vtable for ...' file [filename]
"""

import argparse
import os.path
import platform
import subprocess
from pathlib import Path


class RawDescriptionDefaultArgsHelpFormatter(
    argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter
):
    pass


if platform.system() == "Darwin":
    target = "mac-x64"
elif platform.system() == "Linux":
    if platform.machine() == "x86_64":
        target = "linux-x64"
    elif platform.machine() == "aarch64":
        target = "linux-arm64"
    else:
        raise Exception(
            "Clang toolchain not provided for platform machine "
            + platform.machine()
        )

# Point to the readelf provided in the clang toolchain.
DEFAULT_READELF = (
    Path(__file__).parent.parent.parent
    / "prebuilt"
    / "third_party"
    / "clang"
    / target
    / "bin"
    / "llvm-readelf"
)


def is_elf_file(filename):
    """Check if a file is an ELF file."""
    with open(filename, "rb") as f:
        return f.read(4) == b"\x7fELF"


def is_elf_dso(filename, readelf):
    """Check if a file is an ELF shared library."""
    try:
        result = subprocess.check_output(
            [str(readelf), "-hW", filename], text=True
        )
    except subprocess.CalledProcessError:
        return False
    for line in result.split("\n"):
        parts = line.split()
        if not parts:
            continue

        if parts[0] == "Type:":
            type = parts[1]
            if type == "DYN":
                return True

    # Note that the readelf could fail here if the file could not have been parsed
    # by readelf in the first place.
    return False


def find_vtable_symbols(filename, readelf):
    """Collect a map of all symbols and their sections in an elf file."""
    symbols = {}

    # Lines will look like:
    #
    #   1: 0000000000000000     0 FUNC    GLOBAL DEFAULT   UND std::terminate()
    #
    # Which will have 8 parts unless the Name is not provided
    numparts = 8

    result = subprocess.check_output(
        [str(readelf), "-W", "--demangle", "--symbols", filename], text=True
    )
    for line in result.split("\n"):
        parts = line.split(maxsplit=numparts - 1)
        if len(parts) != numparts:
            continue

        symbol = parts[-1]
        section = parts[-2]
        if not section.isdigit():
            continue

        if not symbol.startswith("vtable for "):
            continue

        symbols[symbol.strip()] = int(section)

    return symbols


def parse_section(line):
    """Return a tuple of the section number and name in a line from
    `readelf -W --sections`. If one couldn't be parsed from here, return None.
    """
    numstart = None
    numend = None
    for i, c in enumerate(line):
        if numstart is None:
            if c == "[":
                numstart = i + 1
        elif numend is None:
            # Start parsing the number
            if c == "]":
                numend = i
                break

    if (numstart is None) or (numend is None):
        return None

    try:
        section_num = int(line[numstart:numend].strip())
    except ValueError:
        return None

    name = line[numend + 1 :].split(maxsplit=1)[0]
    if name == "NULL":
        assert section_num == 0, "Expected the null section to be 0"

    return section_num, name


def find_sections(filename, readelf):
    """Collect a map of all sections."""
    sections = []

    result = subprocess.check_output(
        [str(readelf), "-W", "--sections", filename], text=True
    )
    for line in result.split("\n"):
        result = parse_section(line)
        if result is None:
            continue

        section_num, name = result
        assert (
            len(sections) == section_num
        ), f"Expected '{name}' to correspond to section number {len(sections)} but is instead section number {section_num}"
        sections.append(name)

    assert sections, f"Found no sections for '{filename}'"
    return sections


def check_vtables_in_dso(filename, readelf, excludes, output):
    assert isinstance(excludes, set)
    print("Checking", filename, file=output)

    symbols = find_vtable_symbols(filename, readelf)
    print("Found {} vtables".format(len(symbols)), file=output)

    sections = find_sections(filename, readelf)
    print("Found {} sections".format(len(sections)), file=output)

    # Find the rodata section.
    for i, section in enumerate(sections):
        if section == ".rodata":
            rodata = i
            break
    else:
        print("Could not find rodata in ", sections, file=output)
        return

    bad_vtables = []
    for vtable, section in symbols.items():
        if section != rodata and vtable not in excludes:
            bad_vtables.append(vtable)

    assert (
        not bad_vtables
    ), f"Found vtables that aren't in rodata:\n{str(bad_vtables)}"
    print("All vtables are in rodata :)", file=output)


def check_vtables_in_file(filename, readelf, excludes, output):
    if is_elf_file(filename):
        # FIXME: We should also be checking executables.
        check_vtables_in_dso(filename, readelf, excludes, output)
    else:
        print(
            "WARN: Could not recognize ELF file type for",
            filename,
            file=sys.stderr,
        )


def check_vtables_in_dir(dirname, readelf, excludes, output):
    assert os.path.exists(dirname), "{} does not exist".format(dirname)
    for root, _, files in os.walk(dirname):
        for file in files:
            path = os.path.join(root, file)
            check_vtables_in_file(path, readelf, excludes, output)


def check_file_main(args):
    filename = args.filename
    excludes = set(args.exclude)

    if os.path.isfile(filename):
        check_vtables_in_file(filename, args.readelf, excludes, args.output)
    else:
        check_vtables_in_dir(filename, args.readelf, excludes, args.output)

    return 0


def check_toolchain_main(args):
    excludes = set(args.exclude)
    check_vtables_in_dir(
        os.path.join(args.toolchain_dir, "lib", "aarch64-unknown-fuchsia"),
        args.readelf,
        excludes,
        args.output,
    )
    check_vtables_in_dir(
        os.path.join(args.toolchain_dir, "lib", "x86_64-unknown-fuchsia"),
        args.readelf,
        excludes,
        args.output,
    )
    return 0


def check_toolchain_multilib_main(args):
    excludes = set(args.exclude)

    targets = (
        "aarch64-unknown-fuchsia",
        "x86_64-unknown-fuchsia",
    )
    multilibs = (
        "relative-vtables",
        "relative-vtables+noexcept",
        "relative-vtables+asan",
        "relative-vtables+asan+noexcept",
    )

    for target in targets:
        for multilib in multilibs:
            p = Path(args.toolchain_dir) / "lib" / target / multilib
            check_vtables_in_dir(p, args.readelf, excludes, args.output)

    # Check hwasan if exists
    p = (
        Path(args.toolchain_dir)
        / "lib"
        / "aarch64-unknown-fuchsia"
        / "relative-vtables+hwasan"
    )
    if p.exists():
        check_vtables_in_dir(p, args.readelf, excludes, args.output)
    p = (
        Path(args.toolchain_dir)
        / "lib"
        / "aarch64-unknown-fuchsia"
        / "relative-vtables+hwasan+noexcept"
    )
    if p.exists():
        check_vtables_in_dir(p, args.readelf, excludes, args.output)

    return 0


def check_fuchsia_main(args):
    excludes = set(args.exclude)

    # Read files to check from binaries.json.
    with open(os.path.join(args.build, "binaries.json"), "r") as f:
        import json

        binaries = json.load(f)

    files = ["binaries.json"]
    for binary in binaries:
        if binary["os"] != "fuchsia":
            continue

        path = binary["debug"]
        if os.path.exists(path):
            check_vtables_in_file(path, args.readelf, excludes, args.output)
            files.append(path)
        else:
            print(
                f"Skipping {path}. This path does not exist.", file=args.output
            )

    if args.depfile:
        with open(args.depfile, "w") as depfile:
            depfile.write(f"{args.output.name}: {' '.join(files)}\n")

    return 0


def main():
    parser = argparse.ArgumentParser(
        description="Script for checking that all vtables in an ELF file are in .rodata",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--readelf", default=DEFAULT_READELF, help="readelf executable to use"
    )
    parser.add_argument(
        "--exclude",
        action="append",
        default=[],
        help="Do not check if this class is in rodata.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=argparse.FileType("w"),
        default="-",
        help="If provided, write to this output file rather than stdout",
    )

    subparsers = parser.add_subparsers(
        dest="mode", help="Choose which mode this tool should run with."
    )

    file_parser = subparsers.add_parser(
        "file", description="Check vtables in a single ELF file."
    )
    file_parser.add_argument(
        "filename",
        type=os.path.abspath,
        help="Path to ELF file to or directory to recursively search through for ELF files",
    )

    toolchain_parser = subparsers.add_parser(
        "toolchain",
        description="""
Script for checking that all vtables in an ELF file are in .rodata. This will
check all relevant fuchsia-target executables and DSOs in a clang toolchain
(libc++, libunwind, etc). This includes:

- ${toolchain_dir}/lib/aarch64-unknown-fuchsia/
- ${toolchain_dir}/lib/x86_64-unknown-fuchsia/
""",
        formatter_class=RawDescriptionDefaultArgsHelpFormatter,
    )
    toolchain_parser.add_argument(
        "toolchain_dir",
        type=os.path.abspath,
        help="Path to the clang toolchain build directory",
    )

    fuchsia_parser = subparsers.add_parser(
        "fuchsia",
        description="""
Script for checking that all vtables in an ELF file are in .rodata. This will
check all relevant fuchsia-target executables and DSOs in a fuchsia build.
""",
        formatter_class=RawDescriptionDefaultArgsHelpFormatter,
    )
    fuchsia_parser.add_argument(
        "build",
        type=os.path.abspath,
        help="Path to either the GN or ZN fuchsia build directory.",
    )
    fuchsia_parser.add_argument(
        "--depfile", help="If provided, create this depfile."
    )

    multilib_parser = subparsers.add_parser(
        "multilib",
        description="""
Script for checking that all vtables in an ELF file are in .rodata. This will
check all relevant fuchsia-target executables and DSOs in a clang toolchain
(libc++, libunwind, etc). This includes:

Script for checking that all vtables in an ELF file are in .rodata. This will
check all relevant fuchsia-target executables and DSOs in the relative-vtables
multilib variant of a clang toolchain (libc++, libunwind, etc). This includes:

- ${toolchain_dir}/lib/x86_64-unknown-fuchsia/relative-vtables/
- ${toolchain_dir}/lib/x86_64-unknown-fuchsia/relative-vtables+noexcept/
- ${toolchain_dir}/lib/x86_64-unknown-fuchsia/relative-vtables+asan/
- ${toolchain_dir}/lib/x86_64-unknown-fuchsia/relative-vtables+asan+noexcept/
- ${toolchain_dir}/lib/aarch64-unknown-fuchsia/relative-vtables/
- ${toolchain_dir}/lib/aarch64-unknown-fuchsia/relative-vtables+noexcept/
- ${toolchain_dir}/lib/aarch64-unknown-fuchsia/relative-vtables+asan/
- ${toolchain_dir}/lib/aarch64-unknown-fuchsia/relative-vtables+asan+noexcept/
- ${toolchain_dir}/lib/aarch64-unknown-fuchsia/relative-vtables+hwasan/
- ${toolchain_dir}/lib/aarch64-unknown-fuchsia/relative-vtables+hwasan+noexcept/
""",
        formatter_class=RawDescriptionDefaultArgsHelpFormatter,
    )
    multilib_parser.add_argument(
        "toolchain_dir",
        type=os.path.abspath,
        help="Path to the clang toolchain build directory",
    )

    file_parser.set_defaults(func=check_file_main)
    toolchain_parser.set_defaults(func=check_toolchain_main)
    multilib_parser.set_defaults(func=check_toolchain_multilib_main)
    fuchsia_parser.set_defaults(func=check_fuchsia_main)

    args = parser.parse_args()
    return args.func(args)


if __name__ == "__main__":
    import sys

    sys.exit(main())