#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
#
# 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   Apache-2.0" 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())