// SPDX-License-Identifier: Apache-2.0
//
// ldlsd -- LDLS, on SchOS.
//
// docs/plan/ldls.md milestone L3b. `ldls-run` did this on the build machine so
// that loader failures cost a second instead of a boot; this is the same
// loader, the same BTF reader and the same trampolines, in a process on the
// system they were written for.
//
//     ldlsd <module.ko> [vmlinux-btf]
//
// **Neither file ships with SchOS.** Both are the user's: a module they
// obtained, and the BTF of the kernel it was built for. docs/licensing.md §6
// is the whole argument, and `.gitignore` refusing `*.ko` is how it is kept
// mechanically rather than by intention.
//
// What it does *not* yet do is implement any of the symbols the driver asks
// for -- every one is still a stub that reports itself and returns zero, which
// is L3c. So the useful output here is the same as `ldls-run`'s: what the
// driver asks for, in what order, and how far it gets.

#include "types.h"
#include "user.h"
#include "fcntl.h"
#include "posix.h"
#include "stat.h"

#include "ko.h"
#include "stub.h"
#include "btf.h"

// --- small output, since lib/ has no printf ---------------------------------

static void
puts_(const char *s)
{
    int n = 0;
    while (s[n]) n++;
    write(1, (void *)s, n);
}

static void
putn(long long v)
{
    char b[24];
    int i = 0;
    if (v < 0) { write(1, "-", 1); v = -v; }
    if (v == 0) { write(1, "0", 1); return; }
    while (v > 0) { b[i++] = (char)('0' + (v % 10)); v /= 10; }
    while (i-- > 0) write(1, &b[i], 1);
}

static void
puthexn(unsigned long long v, int digits)
{
    static const char h[] = "0123456789abcdef";
    char b[16];
    int i;
    if (digits > 16) digits = 16;
    for (i = digits - 1; i >= 0; i--) { b[i] = h[v & 0xF]; v >>= 4; }
    write(1, b, digits);
}

static void
puthex(unsigned long long v)
{
    puts_("0x");
    puthexn(v, 16);
}

// lib/ has no string.h either, so the two byte helpers this needs are here.
static void
bcopy_(void *d, const void *s, int n)
{
    unsigned char *a = (unsigned char *)d;
    const unsigned char *b = (const unsigned char *)s;
    while (n-- > 0) *a++ = *b++;
}

static int
streq(const char *a, const char *b)
{
    while (*a && *a == *b) { a++; b++; }
    return *a == 0 && *b == 0;
}

// --- allocation -------------------------------------------------------------
//
// mmap and nothing else. lib/ has no malloc, and this needs both large
// buffers (a kernel's BTF is several megabytes) and memory it can execute --
// which on SchOS is every mapping, because the kernel never sets EFER.NXE.
// That is convenient here and is a gap rather than a feature; see the note in
// docs/plan/ldls.md.
//
// The module lands in the low 2 GB automatically: a process's mmap area starts
// at 0x40000000 and bumps upward, which is exactly the window R_X86_64_32S
// relocations need. L1 predicted that and this is where it is collected.
static void *
getmem(unsigned long long n)
{
    void *p = mmap(0, n, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS,
                   -1, 0);
    if (p == 0 || (long)p == -1)
        return 0;
    return p;
}

static unsigned char *
slurp(const char *path, unsigned long long *len)
{
    struct stat st;
    unsigned char *buf;
    int fd, got = 0;

    fd = open((char *)path, O_RDONLY);
    if (fd < 0)
        return 0;
    if (fstat(fd, &st) < 0 || st.size == 0) { close(fd); return 0; }

    buf = getmem(st.size);
    if (!buf) { close(fd); return 0; }

    while (got < (int)st.size) {
        int n = read(fd, (char *)buf + got, (int)st.size - got);
        if (n <= 0)
            break;
        got += n;
    }
    close(fd);
    if (got != (int)st.size)
        return 0;
    *len = st.size;
    return buf;
}

// --- the stubs ---------------------------------------------------------------

static const char **g_names;
static int          g_count;
static int          g_calls;
static unsigned long g_pci_driver;

unsigned long
ldls_stub(long idx, unsigned long *args)
{
    const char *name = (idx >= 0 && idx < g_count) ? g_names[idx] : "(?)";

    g_calls++;
    if (g_calls <= 40) {
        puts_("  "); putn(g_calls); puts_("  "); puts_(name);
        puts_("("); puthex(args[0]); puts_(", "); puthex(args[1]);
        puts_(")\n");
    } else if (g_calls == 41) {
        puts_("  ... (further calls not printed)\n");
    }

    // The driver's own words, which is the clearest evidence a load worked.
    // Linux format strings carry a leading level byte and %p extensions that
    // nothing here can expand, so the string is printed raw.
    if (streq(name, "_printk") || streq(name, "printk")) {
        const char *f = (const char *)args[0];
        if (f) {
            if ((unsigned char)f[0] == 0x01 && f[1] >= '0' && f[1] <= '7')
                f += 2;
            puts_("      | "); puts_(f);
            {
                int n = 0;
                while (f[n]) n++;
                if (n == 0 || f[n - 1] != '\n') puts_("\n");
            }
        }
    }

    if (streq(name, "__pci_register_driver") && !g_pci_driver)
        g_pci_driver = args[0];

    return 0;
}

struct rctx { struct ko *k; struct stubs *s; };

static int
resolve(void *ctx, const char *name, unsigned long long *out)
{
    struct rctx *r = (struct rctx *)ctx;
    int i;

    for (i = 0; i < r->k->nundef; i++)
        if (streq(r->k->undef[i], name)) {
            *out = stub_addr(r->s, i);
            return 0;
        }
    return -1;
}

static unsigned long long
find_symbol(struct ko *k, const char *want)
{
    int i;
    for (i = 0; i < k->nsym; i++) {
        if (!streq(k->str + k->sym[i].st_name, want))
            continue;
        if (k->sym[i].st_shndx == KO_SHN_UNDEF ||
            k->sym[i].st_shndx >= (unsigned short)k->nsh)
            continue;
        if (!k->sec[k->sym[i].st_shndx].alloc)
            continue;
        return k->sec[k->sym[i].st_shndx].addr + k->sym[i].st_value;
    }
    return 0;
}

// --- being the PCI core -------------------------------------------------------

static unsigned int
lay(struct btf *b, const char *type, const char *member)
{
    unsigned int off, size;
    int id = btf_find(b, type, BTF_KIND_STRUCT);

    if (!id) {
        puts_("ldlsd: this BTF has no struct "); puts_(type); puts_("\n");
        return 0xffffffffu;
    }
    if (!member)
        return btf_type_size(b, id);
    if (btf_member_off(b, id, member, &off, &size) < 0) {
        puts_("ldlsd: "); puts_(type); puts_("."); puts_(member);
        puts_(": "); puts_(btf_err(b)); puts_("\n");
        return 0xffffffffu;
    }
    return off;
}

static unsigned int
rd32(unsigned long base, unsigned int off)
{ unsigned int v; bcopy_(&v, (void *)(base + off), 4); return v; }

static unsigned long
rd64(unsigned long base, unsigned int off)
{ unsigned long v; bcopy_(&v, (void *)(base + off), 8); return v; }

static void
wr16(unsigned long base, unsigned int off, unsigned short v)
{ bcopy_((void *)(base + off), &v, 2); }

static void
wr32(unsigned long base, unsigned int off, unsigned int v)
{ bcopy_((void *)(base + off), &v, 4); }

static void
wr64(unsigned long base, unsigned int off, unsigned long v)
{ bcopy_((void *)(base + off), &v, 8); }

#define LDLS_IORESOURCE_IO  0x00000100
#define LDLS_IORESOURCE_MEM 0x00000200

static void
drive_probe(const char *btfpath)
{
    struct btf b;
    unsigned char *bimg;
    unsigned long long blen;
    unsigned int *boff;
    long nt;
    unsigned int drv_probe, drv_idtable, drv_name;
    unsigned int id_size, id_vendor, id_device;
    unsigned int dev_size, dev_vendor, dev_device, dev_devfn, dev_irq, dev_res;
    unsigned int res_size, res_start, res_end, res_flags;
    unsigned long idtab, pdev;
    int nid = 0;

    bimg = slurp(btfpath, &blen);
    if (!bimg) { puts_("ldlsd: cannot read the BTF\n"); return; }

    nt = btf_count(bimg, blen);
    if (nt < 0) { puts_("ldlsd: that file is not BTF\n"); return; }

    boff = getmem(((unsigned long long)nt + 2) * 4);
    if (!boff) { puts_("ldlsd: no memory for the BTF index\n"); return; }
    if (btf_open(&b, bimg, blen, boff, (unsigned int)nt + 2) < 0) {
        puts_("ldlsd: "); puts_(btf_err(&b)); puts_("\n");
        return;
    }
    puts_("  btf     "); putn(nt); puts_(" types\n");

    drv_probe   = lay(&b, "pci_driver", "probe");
    drv_idtable = lay(&b, "pci_driver", "id_table");
    drv_name    = lay(&b, "pci_driver", "name");
    id_size     = lay(&b, "pci_device_id", 0);
    id_vendor   = lay(&b, "pci_device_id", "vendor");
    id_device   = lay(&b, "pci_device_id", "device");
    dev_size    = lay(&b, "pci_dev", 0);
    dev_vendor  = lay(&b, "pci_dev", "vendor");
    dev_device  = lay(&b, "pci_dev", "device");
    dev_devfn   = lay(&b, "pci_dev", "devfn");
    dev_irq     = lay(&b, "pci_dev", "irq");
    dev_res     = lay(&b, "pci_dev", "resource");
    res_size    = lay(&b, "resource", 0);
    res_start   = lay(&b, "resource", "start");
    res_end     = lay(&b, "resource", "end");
    res_flags   = lay(&b, "resource", "flags");

    if (drv_probe == 0xffffffffu || dev_size == 0xffffffffu ||
        res_size == 0xffffffffu) {
        puts_("ldlsd: this kernel's BTF cannot drive probe()\n");
        return;
    }

    puts_("  layout  pci_driver.probe +"); putn(drv_probe);
    puts_("  pci_dev "); putn(dev_size); puts_(" bytes\n");

    {
        const char *dn = (const char *)rd64(g_pci_driver, drv_name);
        if (dn) { puts_("  driver  \""); puts_(dn); puts_("\"\n"); }
    }

    idtab = rd64(g_pci_driver, drv_idtable);
    if (!idtab) { puts_("  it has no id_table\n"); return; }

    puts_("\n  devices it claims:\n");
    for (;;) {
        unsigned long e = idtab + (unsigned long)nid * id_size;
        unsigned int v = rd32(e, id_vendor), d = rd32(e, id_device);
        if (!v && !d) break;
        if (nid < 6) {
            puts_("    "); puthexn(v, 4); puts_(":"); puthexn(d, 4);
            puts_("\n");
        }
        nid++;
        if (nid > 4096) break;
    }
    if (nid > 6) { puts_("    ... and "); putn(nid - 6); puts_(" more\n"); }
    if (!nid) return;

    pdev = (unsigned long)getmem(dev_size);
    if (!pdev) { puts_("ldlsd: no memory for a pci_dev\n"); return; }
    {
        unsigned long r0 = pdev + dev_res;
        unsigned long r1 = r0 + res_size;

        wr16(pdev, dev_vendor, (unsigned short)rd32(idtab, id_vendor));
        wr16(pdev, dev_device, (unsigned short)rd32(idtab, id_device));
        wr32(pdev, dev_devfn, 0x18);
        wr32(pdev, dev_irq, 11);
        wr64(r0, res_start, 0xc000);
        wr64(r0, res_end,   0xc0ff);
        wr64(r0, res_flags, LDLS_IORESOURCE_IO);
        wr64(r1, res_start, 0xfebc0000);
        wr64(r1, res_end,   0xfebc00ff);
        wr64(r1, res_flags, LDLS_IORESOURCE_MEM);
    }

    puts_("\n  calling probe\n\n");
    g_calls = 0;
    {
        int (*probe)(void *, void *) =
            (int (*)(void *, void *))(void *)rd64(g_pci_driver, drv_probe);
        if (!probe) { puts_("  the driver has no probe\n"); return; }
        {
            int r = probe((void *)pdev, (void *)idtab);
            puts_("\n  probe returned "); putn(r);
            puts_(" after "); putn(g_calls); puts_(" stub call(s)\n");
        }
    }
}

// --- main ----------------------------------------------------------------------

int
main(int argc, char *argv[])
{
    struct ko k;
    struct ko_section *sec;
    struct stubs stubs;
    struct rctx rc;
    unsigned char *img, *mem, *code;
    unsigned long long len, span, entry;
    int nsec, i;

    if (argc < 2) {
        puts_("usage: ldlsd <module.ko> [vmlinux-btf]\n");
        puts_("  neither file ships with SchOS; both are yours.\n");
        exit();
    }

    img = slurp(argv[1], &len);
    if (!img) { puts_("ldlsd: cannot read "); puts_(argv[1]); puts_("\n"); exit(); }

    nsec = ko_sections_needed(img, len);
    if (nsec < 0) {
        puts_("ldlsd: that is not an ELF64 relocatable object\n");
        exit();
    }
    sec = getmem((unsigned long long)(nsec + 1) * sizeof(struct ko_section));
    if (!sec) { puts_("ldlsd: out of memory\n"); exit(); }

    if (ko_open(&k, img, len, sec, nsec) < 0) {
        puts_("ldlsd: "); puts_(ko_err(&k)); puts_("\n");
        exit();
    }

    puts_(argv[1]); puts_("\n");
    puts_("  module  ");
    puts_(ko_name(&k) ? ko_name(&k) : "(not a kernel module)"); puts_("\n");
    puts_("  license ");
    puts_(ko_modinfo(&k, "license") ? ko_modinfo(&k, "license")
                                    : "(none declared)"); puts_("\n");
    puts_("  wants   "); putn(k.nundef); puts_(" symbols\n");

    span = ko_span(&k);
    mem  = getmem(span ? span : 4096);
    code = getmem((unsigned long long)(k.nundef + 1) * STUB_SIZE);
    if (!mem || !code) { puts_("ldlsd: out of memory\n"); exit(); }

    stub_build(&stubs, code, k.nundef);
    g_names = k.undef;
    g_count = k.nundef;
    for (i = 0; i < k.nundef; i++)
        if (stub_is_bare(k.undef[i]))
            stub_make_bare(code + (unsigned long long)i * STUB_SIZE);

    if (ko_layout(&k, mem, span, (unsigned long long)mem) < 0) {
        puts_("ldlsd: layout: "); puts_(ko_err(&k)); puts_("\n");
        exit();
    }

    rc.k = &k; rc.s = &stubs;
    if (ko_relocate(&k, resolve, &rc) < 0) {
        puts_("ldlsd: relocate: "); puts_(ko_err(&k)); puts_("\n");
        exit();
    }
    // The stack canary is fetched through %gs, which belongs to the kernel
    // here. See ko_defuse_gs().
    {
        int g = ko_defuse_gs(&k, "__ref_stack_chk_guard");
        if (g < 0) g = 0;
        if (g == 0) g = ko_defuse_gs(&k, "__stack_chk_guard");
        if (g > 0) { puts_("  defused "); putn(g); puts_(" %gs stack-guard load(s)\n"); }
    }

    puts_("  loaded  "); putn(span); puts_(" bytes at "); puthex((unsigned long long)mem);
    puts_("\n");

    entry = find_symbol(&k, "init_module");
    if (!entry) { puts_("  no init_module -- nothing to call\n"); exit(); }

    puts_("\n  calling init_module\n\n");
    {
        int (*init)(void) = (int (*)(void))(void *)entry;
        int r = init();
        puts_("\n  init_module returned "); putn(r);
        puts_(" after "); putn(g_calls); puts_(" stub call(s)\n");
    }

    if (argc > 2 && g_pci_driver)
        drive_probe(argv[2]);
    else if (argc > 2)
        puts_("\n  no pci_driver was registered -- nothing to probe\n");

    exit();
}