smgdc/tests/test_vtable.py

#!/usr/bin/python3

# tests for vtable computations

import hashlib
import os
import pathlib
import typing

import pytest
import smgdc.angr
from smgdc.validate import LinuxBinary


@pytest.fixture(scope="session")
def game_bin(request) -> LinuxBinary:
    # this is tested against the TF2 Linux DS server binary (steam app 232250, depot 232256)
    #
    # the author will not provide further support in obtaining this file, nor will the author
    # update the tests for other game versions
    #
    # if you have the file: to enable tests that depend on the binary, create a `.env` file with
    # TF_LINUX_BINARY set to the binary path (use single quotes on Windows to avoid escaping
    # backslashes), then do `just test-extended`
    #
    # TODO: find a libre binary that we can do similarly extensive tests against in its place
    valid_hashes = {
        # game version 8622567, manifest 2590057218013527366
        "a5a1adde851be2c71f8de73830466cd20f8c34d26ade580d110c4f93c6ef1374",
        # game version 8835751, manifest 9155443851166366439?
        # this one requires a constraint file
        "6a14d6460086836b710aaed901af3304f791b25409c84d182c6d7662a9e09401",
    }

    path_string = os.environ.get("TF_LINUX_BINARY")
    if not path_string:
        pytest.skip("No Linux binary given")

    bin_path = pathlib.Path(path_string)
    with bin_path.open("rb") as f:
        h = hashlib.file_digest(f, "sha256")
        if h.hexdigest() not in valid_hashes:
            pytest.skip("Incorrect Linux binary")

    return LinuxBinary(bin_path)


class VTableCheck(typing.NamedTuple):
    linux_index: int
    windows_index: int | None
    symbol_name: str


@pytest.mark.extended
@pytest.mark.parametrize(
    "vtsym_name,vtasserts",
    [
        (
            "_ZTV9CTFPlayer",
            [
                (149, 147, "_ZN20CBaseCombatCharacter8FVisibleERK6VectoriPP11CBaseEntity"),
                (490, None, "_ZN9CTFPlayer16ReapplyProvisionEv"),
                (493, 486, "_ZN9CTFPlayer13GiveNamedItemEPKciPK13CEconItemViewb"),
            ],
        ),
        (
            "_ZTV11CBaseObject",
            [
                (356, 354, "_ZN11CBaseObject6KilledERK15CTakeDamageInfo"),
                (367, 365, "_ZN11CBaseObject17CheckUpgradeOnHitEP9CTFPlayer"),
                (383, 382, "_ZN11CBaseObject13CanBeUpgradedEP9CTFPlayer"),
                (384, 383, "_ZN11CBaseObject14StartUpgradingEv"),
                (413, 381, "_ZNK11CBaseObject13CanBeUpgradedEv"),
            ],
        ),
        (
            "_ZTV13CBaseNPCMaker",
            [
                # on newer binaries the functions are optimized out to 0
                # this also affects CBaseObject
                (2, 1, "_ZN11CBaseEntity13SetRefEHandleERK11CBaseHandle"),
            ],
        ),
        (
            "_ZTV12CTFGameRules",
            [
                # parent vtable spans are not completely in descending order (some of them increase)
                (140, 139, "_ZNK12CTFGameRules15IsHolidayActiveEi"),
            ],
        ),
    ],
)
def test_vtable(game_bin: LinuxBinary, vtsym_name: str, vtasserts: list[VTableCheck]):
    # tests vtable translation
    vtsym = game_bin.angr.loader.find_symbol(vtsym_name)
    assert vtsym, f"Could not find vtable symbol {vtsym_name}"
    orig_vtable, *thunk_vtables = game_bin.vtable_disambiguator.get_vtables_from_address(vtsym)
    win_vtable = game_bin.vtable_disambiguator.get_windows_vtables_from(vtsym)

    for check in map(VTableCheck._make, vtasserts):
        symbol = game_bin.angr.loader.find_symbol(check.symbol_name)
        assert orig_vtable.index(symbol) == check.linux_index

        if check.windows_index is not None:
            assert win_vtable.index(symbol) == check.windows_index
        else:
            assert symbol not in win_vtable