#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later
#
# Check the LDLS loader against an object we compiled ourselves.
#
# docs/architecture/ldls.md, "The loader". The thing being tested is ldls/ko.c; the
# input is ldls/test/refmod.c compiled with -c, which is exactly what a `.ko`
# is. **No kernel module is involved and none is in this tree**, which is rule
# 1 applied to the repository as well as to the image.
#
# This is a test rather than a demo, so it asserts. `ldls-inspect` printing
# something plausible is not evidence; the specific things the loader has to
# get right are listed below and each one is checked by name.

import re
import subprocess
import sys


def fail(msg):
    print("ldls-check: FAIL -- %s" % msg)
    sys.exit(1)


def main():
    if len(sys.argv) < 3:
        print("usage: ldls_check.py <ldls-inspect> <object.o>")
        return 2

    inspect, obj = sys.argv[1], sys.argv[2]
    r = subprocess.run([inspect, obj], capture_output=True, text=True)
    out = r.stdout

    if r.returncode != 0:
        fail("ldls-inspect exited %d\n%s%s" % (r.returncode, out, r.stderr))

    checks = 0

    def need(cond, what):
        nonlocal checks
        checks += 1
        if not cond:
            fail("%s\n--- output ---\n%s" % (what, out))
        print("  [ok] %s" % what)

    # It is recognised as a module at all, and its declared licence is read
    # straight out of .modinfo. That is the first question anyone should be
    # able to ask about an object before loading it.
    need("module    refmod" in out, "the module names itself from .modinfo")
    need("license   GPL" in out,
         "the declared licence is read without loading anything")

    # Every kind of section the loader has to place.
    for sec, why in (
        (".text",   "text is laid out"),
        (".data",   "initialised data is laid out"),
        (".bss",    "bss is laid out, and zeroed rather than copied"),
    ):
        need(re.search(r"^\s+%s\b" % re.escape(sec), out, re.M) is not None, why)

    need("(zeroed)" in out, "bss is marked NOBITS and gets zeroed, not copied")

    # The undefined symbols -- L0's entire output, and the work queue for L1.
    m = re.search(r"undefined symbols: (\d+)", out)
    need(m is not None, "the undefined symbols are counted")
    need(int(m.group(1)) >= 4, "it found the object's undefined symbols")

    for sym in ("ldls_test_alloc", "ldls_test_free", "ldls_test_print",
                "ldls_test_counter"):
        need(sym in out, "reports the undefined symbol %s" % sym)

    # And the one that matters most: with a resolver that knows nothing, a real
    # object cannot relocate, and the loader must say which name stopped it
    # rather than quietly producing a broken image.
    need("stops at: unresolved: ldls_test_" in out,
         "relocation stops at a named unresolved symbol")

    need("loads to" in out, "the object reports a load size")

    # The relocations are counted by type, and none of them is one the loader
    # cannot handle. That check matters more than it looks: an unhandled type
    # is invisible until a relocation of that type is reached, which for a real
    # driver can be the last one in the file.
    m = re.search(r"relocations: (\d+)", out)
    need(m is not None and int(m.group(1)) > 0,
         "the relocations are counted by type")
    need("does not handle" not in out,
         "every relocation type in the object is one the loader handles")

    # And with everything resolved, relocation must run to the end. This is a
    # different question from the list of undefined symbols and it is the one
    # that says the loader is finished rather than merely parsing.
    r2 = subprocess.run([inspect, "-q", "-r", obj], capture_output=True,
                        text=True)
    need(r2.returncode == 0 and "relocated in full against stubs" in r2.stdout,
         "with every symbol resolved, relocation runs to the end")

    # A file that is not a relocatable object is refused rather than parsed.
    r3 = subprocess.run([inspect, sys.argv[0]], capture_output=True, text=True)
    need(r3.returncode != 0, "a non-ELF file is refused, not parsed")

    # --- L1: the object is loaded and its init actually runs ----------------
    if len(sys.argv) < 4:
        print("ldls-check: %d checks passed (L0 only -- no ldls-run given)"
              % checks)
        return 0

    run = sys.argv[3]
    r4 = subprocess.run([run, obj], capture_output=True, text=True)
    out = r4.stdout
    need(r4.returncode == 0, "ldls-run loads the object and returns cleanly")
    need("calling init_module" in out, "it finds and calls init_module")
    need("init_module returned 0" in out,
         "init runs to the end and returns what it meant to")

    # The call *order*, which is the whole point of L1: a symbol list is a set
    # and says nothing about sequence.
    seq = re.findall(r"^\s+\d+\s+(\w+)\(", out, re.M)
    need(seq == ["ldls_test_print", "ldls_test_alloc", "ldls_test_free"],
         "the stub calls are reported in the order the code makes them")

    # Arguments survive the trampoline and the thunk. If the register save were
    # wrong this is where it would show, because the first argument is a real
    # pointer into the loaded object rather than a small integer.
    m = re.search(r"ldls_test_print\((0x[0-9a-f]+)", out)
    need(m is not None and int(m.group(1), 16) != 0,
         "arguments reach the stub intact, not zeroed by the thunk")

    # --- L2: the layouts come from a kernel's own BTF -----------------------
    #
    # Conditional, and it says so when it skips. The BTF is the *host's*
    # kernel type data -- not ours, not shipped, and not present on every
    # machine. A check that silently passes when it did not run is worse than
    # one that is honest about being unable to.
    import os
    types = os.path.join(os.path.dirname(inspect), "ldls-types")
    btf = "/sys/kernel/btf/vmlinux"
    if not (os.path.exists(types) and os.path.exists(btf)):
        print("ldls-check: %d checks passed "
              "(L2 skipped -- no %s on this machine)" % (checks, btf))
        return 0

    r5 = subprocess.run([types, btf, "--need"], capture_output=True, text=True)
    need(r5.returncode == 0,
         "every layout the PCI path needs is readable from the kernel's BTF")
    need("all present" in r5.stdout, "none of them is missing")

    # The offsets are not ours to have guessed, so the check is that they come
    # back at all and are self-consistent -- pci_driver.probe after .id_table
    # after .name, which is the order the struct declares them in.
    o = {}
    for line in r5.stdout.splitlines():
        m2 = re.match(r"\s+(\w+)\s+(\w+)\s+offset\s+(\d+)", line)
        if m2:
            o[(m2.group(1), m2.group(2))] = int(m2.group(3))
    need(o.get(("pci_driver", "name")) == 0 and
         0 < o.get(("pci_driver", "id_table"), -1) <
             o.get(("pci_driver", "probe"), -1),
         "the pci_driver offsets come back in declaration order")
    need(o.get(("resource", "start")) == 0 and
         o.get(("resource", "end"), -1) == 8,
         "struct resource reads as the ABI says it is laid out")

    print("ldls-check: %d checks passed" % checks)
    return 0


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