// SPDX-License-Identifier: Apache-2.0
//
// ldls-run -- load a module and call its init, on the build machine.
//
// docs/architecture/ldls.md. The loader answered "what does this driver ask
// for"; this answers "and in what order does it ask", which is the other thing
// that cannot be guessed and falls straight out of running the thing.
//
// This is a **development harness, not the deliverable.** LDLS proper runs on
// SchOS in a process, reaching hardware through /dev/pci. Running the same
// loader here first is the arrangement mkschfs already has with the kernel,
// and for the same reason: the failures at this stage are loader failures, and
// one found here costs a second instead of a boot.
//
//     ldls-run 8139too.ko
//
// Nothing touches hardware. Every symbol the module wants resolves to a stub
// that logs and returns zero, so a driver gets as far as its first real
// dependency on the kernel doing something, and then says so.
//
// **No module is committed to this tree.** See docs/architecture/ldls.md rule 1.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <setjmp.h>
#include <stdarg.h>
#include <sys/mman.h>

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

// --- what the stubs report ---------------------------------------------------

static const char **g_names;
static int          g_count;
static int          g_calls;
static int          g_verbose;

// What init_module handed to __pci_register_driver.
//
// In Linux that call ends at the PCI core, which matches the driver's id_table
// against the bus and calls probe(). There is no PCI core here, so LDLS keeps
// the pointer and plays that part itself -- which is the whole of L2.
static unsigned long g_pci_driver;
static int           g_answer;   // give allocators real memory

static int
is_register_driver(const char *n)
{
  return strcmp(n, "__pci_register_driver") == 0 ||
         strcmp(n, "pci_register_driver") == 0;
}

// A handful of symbols are worth actually implementing rather than logging,
// because what they do is print -- and a driver printing its own banner
// through our code is the clearest evidence there is that the load worked.
static int
is_printk(const char *n)
{
  return strcmp(n, "_printk") == 0 || strcmp(n, "printk") == 0;
}

// Linux format strings are printf's with extensions: a leading KERN_SOH level
// byte, and %p followed by a letter meaning "pretty-print this kind of
// pointer". vprintf understands neither, and handing it a %pS would read an
// argument that is not there.
//
// So the level prefix is stripped and anything carrying a pointer extension is
// printed raw. Enough for a banner, honest about the rest, and the real
// implementation belongs with the rest of the shim.
static void
printk_like(const char *fmt, unsigned long *args)
{
  const char *p = fmt;
  const char *s;

  if (!fmt) { printf("      (printk with a null format)\n"); return; }

  if ((unsigned char)p[0] == 0x01 && p[1] >= '0' && p[1] <= '7')
    p += 2;                                     /* KERN_SOH + level */

  for (s = p; *s; s++)
    if (s[0] == '%' && s[1] == 'p' &&
        ((s[2] >= 'a' && s[2] <= 'z') || (s[2] >= 'A' && s[2] <= 'Z'))) {
      printf("      | %s", p);
      if (p[strlen(p) - 1] != '\n') printf("\n");
      return;
    }

  printf("      | ");
  printf(p, args[1], args[2], args[3], args[4], args[5]);
  if (p[0] && p[strlen(p) - 1] != '\n')
    printf("\n");
}

// The smallest set of answers that lets probe() get past its first allocation.
//
// **This is not the shim.** It is six answers chosen to stop the driver
// bailing out at the first call so that the *rest* of its sequence becomes
// visible -- the same reason L1's stubs return zero instead of stopping.
// Implementing these properly, against /dev/pci, is L3.
//
// sizeof(struct net_device) comes from BTF like everything else: a driver
// reaches its private area with netdev_priv(), which is the net_device
// rounded up to NETDEV_ALIGN, so getting it wrong hands the driver a pointer
// into the middle of its own control block.
static unsigned int g_netdev_size;

static unsigned long
answer(const char *name, unsigned long *args)
{
  if (strcmp(name, "alloc_etherdev_mqs") == 0) {
    unsigned long n = (g_netdev_size + 31u) & ~31u;
    return (unsigned long)calloc(1, n + args[0] + 64);
  }
  if (strcmp(name, "pci_iomap") == 0 ||
      strcmp(name, "pci_ioremap_bar") == 0 ||
      strcmp(name, "ioremap") == 0)
    return (unsigned long)calloc(1, 4096);
  if (strcmp(name, "__kmem_cache_create_args") == 0 ||
      strcmp(name, "kmem_cache_create") == 0)
    return (unsigned long)calloc(1, 256);
  if (strcmp(name, "dma_alloc_attrs") == 0)
    return (unsigned long)calloc(1, args[1] ? args[1] : 4096);
  // Success is zero for most of the kernel's int-returning calls, which is
  // already what a stub returns -- so nothing is listed for them here.
  return 0;
}

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

  g_calls++;
  printf("  %3d  %s(", g_calls, name);
  printf("%#lx, %#lx, %#lx", args[0], args[1], args[2]);
  printf(")\n");

  if (is_printk(name))
    printk_like((const char *)args[0], args);

  // netdev_info(dev, fmt, ...) and dev_info(dev, fmt, ...) put the format in
  // the *second* argument. Worth handling for the same reason as _printk: the
  // line a driver prints when it has found its hardware is the one that says
  // whether it really did.
  if (strcmp(name, "netdev_info") == 0 || strcmp(name, "netdev_err") == 0 ||
      strcmp(name, "_dev_info") == 0   || strcmp(name, "_dev_err") == 0)
    printk_like((const char *)args[1], args + 1);

  if (is_register_driver(name) && !g_pci_driver)
    g_pci_driver = args[0];

  if (g_verbose)
    printf("       args: %#lx %#lx %#lx %#lx %#lx %#lx\n",
           args[0], args[1], args[2], args[3], args[4], args[5]);

  // Zero, so the driver carries on and shows the next call. Stopping at the
  // first tells you one name; continuing tells you the sequence, which is what
  // this milestone is for.
  return g_answer ? answer(name, args) : 0;
}

// --- surviving the module -----------------------------------------------------

static sigjmp_buf g_jump;
static volatile int g_signal;

static void
on_fault(int sig)
{
  g_signal = sig;
  siglongjmp(g_jump, 1);
}

// --- playing the PCI core -----------------------------------------------------
//
// Every offset here is read from the kernel's own BTF rather than written
// down. That is the point of doing it this way: LDLS never claims to know
// where `pci_driver.probe` is, it asks the kernel that the driver was built
// for. See ldls/btf.h.

struct pcilayout {
  unsigned int drv_probe, drv_idtable, drv_name;
  unsigned int id_size, id_vendor, id_device, id_class, id_class_mask;
  unsigned int dev_size, dev_vendor, dev_device, dev_devfn, dev_irq,
               dev_resource;
  unsigned int res_size, res_start, res_end, res_flags;
};

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

  if (!id) {
    fprintf(stderr, "ldls-run: this BTF has no struct %s\n", type);
    return -1;
  }
  if (!member) { *out = btf_type_size(b, id); return 0; }
  if (btf_member_off(b, id, member, &off, &size) < 0) {
    fprintf(stderr, "ldls-run: %s.%s: %s\n", type, member, btf_err(b));
    return -1;
  }
  *out = off;
  return 0;
}

static int
pci_layout(struct btf *b, struct pcilayout *L)
{
  return lay(b, "pci_driver", "probe",      &L->drv_probe)   ||
         lay(b, "pci_driver", "id_table",   &L->drv_idtable) ||
         lay(b, "pci_driver", "name",       &L->drv_name)    ||
         lay(b, "pci_device_id", 0,         &L->id_size)     ||
         lay(b, "pci_device_id", "vendor",  &L->id_vendor)   ||
         lay(b, "pci_device_id", "device",  &L->id_device)   ||
         lay(b, "pci_device_id", "class",   &L->id_class)    ||
         lay(b, "pci_device_id", "class_mask", &L->id_class_mask) ||
         lay(b, "pci_dev", 0,               &L->dev_size)    ||
         lay(b, "pci_dev", "vendor",        &L->dev_vendor)  ||
         lay(b, "pci_dev", "device",        &L->dev_device)  ||
         lay(b, "pci_dev", "devfn",         &L->dev_devfn)   ||
         lay(b, "pci_dev", "irq",           &L->dev_irq)     ||
         lay(b, "pci_dev", "resource",      &L->dev_resource)||
         lay(b, "resource", 0,              &L->res_size)    ||
         lay(b, "resource", "start",        &L->res_start)   ||
         lay(b, "resource", "end",          &L->res_end)     ||
         lay(b, "resource", "flags",        &L->res_flags)   ? -1 : 0;
}

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

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

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

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

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

// The PCI resource flag bits. These are not layout and BTF does not carry
// them -- they are #defines. They are ABI between the kernel and every driver
// ever written against it, which makes them facts in the sense of
// docs/architecture/ldls.md rule 3, and there are two of them.
#define LDLS_IORESOURCE_IO  0x00000100
#define LDLS_IORESOURCE_MEM 0x00000200

// --- loading ------------------------------------------------------------------

static unsigned char *
slurp(const char *path, unsigned long long *len)
{
  FILE *f = fopen(path, "rb");
  unsigned char *buf;
  long n;

  if (!f) return 0;
  if (fseek(f, 0, SEEK_END) != 0) { fclose(f); return 0; }
  n = ftell(f);
  if (n <= 0) { fclose(f); return 0; }
  rewind(f);
  buf = malloc((size_t)n);
  if (!buf) { fclose(f); return 0; }
  if (fread(buf, 1, (size_t)n, f) != (size_t)n) { free(buf); fclose(f); return 0; }
  fclose(f);
  *len = (unsigned long long)n;
  return buf;
}

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 (strcmp(r->k->undef[i], name) == 0) {
      *out = stub_addr(r->s, i);
      return 0;
    }
  return -1;                     // cannot happen: the list came from the object
}

// The module's entry point. `module_init(x)` makes `init_module` an alias for
// x, so it is an ordinary defined global and needs no struct layout to find --
// which matters, because the struct that *would* hold it is one of the things
// rule 4 says not to assume.
static unsigned long long
find_symbol(struct ko *k, const char *want)
{
  int i;
  for (i = 0; i < k->nsym; i++) {
    const char *nm = k->str + k->sym[i].st_name;
    if (strcmp(nm, want) != 0)
      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;
}

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, rv = 0;
  const char *btfpath = 0;

  if (argc > 1 && strcmp(argv[1], "-v") == 0) { g_verbose = 1; argv++; argc--; }
  if (argc > 1 && strcmp(argv[1], "-a") == 0) { g_answer = 1; argv++; argc--; }
  if (argc > 2 && strcmp(argv[1], "--btf") == 0) {
    btfpath = argv[2]; argv += 2; argc -= 2;
  }
  if (argc < 2) {
    fprintf(stderr,
            "usage: ldls-run [-v] [-a] [--btf <btf>] <object.ko>\n");
    return 2;
  }

  img = slurp(argv[1], &len);
  if (!img) { fprintf(stderr, "ldls-run: cannot read %s\n", argv[1]); return 1; }

  nsec = ko_sections_needed(img, len);
  if (nsec < 0) {
    fprintf(stderr, "ldls-run: %s: not an ELF64 relocatable object\n", argv[1]);
    return 1;
  }
  sec = calloc((size_t)(nsec ? nsec : 1), sizeof(struct ko_section));
  if (!sec) { fprintf(stderr, "ldls-run: out of memory\n"); return 1; }

  if (ko_open(&k, img, len, sec, nsec) < 0) {
    fprintf(stderr, "ldls-run: %s\n", ko_err(&k));
    return 1;
  }

  printf("%s\n", argv[1]);
  printf("  module  %s\n", ko_name(&k) ? ko_name(&k) : "(not a module)");
  printf("  license %s\n", ko_modinfo(&k, "license") ?
                           ko_modinfo(&k, "license") : "(none declared)");
  printf("  wants   %d symbols\n", k.nundef);

  // Executable memory for both the module and the stubs. A module's .text has
  // to be executed and a trampoline has to be jumped to, so neither can live
  // in ordinary malloc'd memory on a machine that enforces W^X.
  //
  // **And it has to be low.** A module is compiled with -mcmodel=kernel, which
  // lets the compiler emit R_X86_64_32S -- a 32-bit signed absolute address --
  // for anything it believes lives in the top 2 GB of the address space. Load
  // the object at an ordinary mmap address around 0x7f... and those
  // relocations simply do not fit, which is what the loader said the first
  // time this was run.
  //
  // Linux solves it by giving modules their own region within 2 GB of the
  // kernel text. Here MAP_32BIT does the same job from the other end: the
  // first 2 GB, where a positive signed 32-bit value reaches. On SchOS this
  // costs nothing, because a process's mmap area starts at 0x40000000 and is
  // already inside that window -- but it is a real constraint on where the
  // LDLS host may put a module, and it is not one you would guess.
  //
  // The stubs are mapped low for the same reason: a 32S relocation can name an
  // undefined symbol just as easily as a section.
  span = ko_span(&k);
  mem = mmap(0, (size_t)(span ? span : 4096), PROT_READ | PROT_WRITE | PROT_EXEC,
             MAP_PRIVATE | MAP_ANONYMOUS | MAP_32BIT, -1, 0);
  code = mmap(0, (size_t)(k.nundef + 1) * STUB_SIZE,
              PROT_READ | PROT_WRITE | PROT_EXEC,
              MAP_PRIVATE | MAP_ANONYMOUS | MAP_32BIT, -1, 0);
  if (mem == MAP_FAILED || code == MAP_FAILED) {
    fprintf(stderr, "ldls-run: cannot get executable memory\n");
    return 1;
  }

  stub_build(&stubs, code, k.nundef);
  g_names = k.undef;
  g_count = k.nundef;

  // The ones that must not be logged, because they are called at the top of
  // every function and have register rules of their own. See stub_is_bare().
  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) {
    fprintf(stderr, "ldls-run: layout: %s\n", ko_err(&k));
    return 1;
  }

  rc.k = &k; rc.s = &stubs;
  if (ko_relocate(&k, resolve, &rc) < 0) {
    fprintf(stderr, "ldls-run: relocate: %s\n", ko_err(&k));
    return 1;
  }
  {
    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)
      printf("  defused %d %%gs stack-guard reference(s)\n", g);
  }

  printf("  loaded  %llu bytes at %p, %d stubs at %p\n\n",
         span, (void *)mem, k.nundef, (void *)code);

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

  printf("  calling init_module at %#llx\n\n", entry);

  // A module that dereferences the zero our stubs return will fault, and that
  // is a *result* rather than a crash: it says exactly how far the driver got
  // before it needed the kernel to have done something real.
  signal(SIGSEGV, on_fault);
  signal(SIGBUS,  on_fault);
  signal(SIGILL,  on_fault);
  signal(SIGFPE,  on_fault);

  if (sigsetjmp(g_jump, 1) == 0) {
    int (*init)(void) = (int (*)(void))(void *)entry;
    int r = init();
    printf("\n  init_module returned %d after %d stub call(s)\n", r, g_calls);
  } else {
    printf("\n  stopped by signal %d after %d stub call(s)\n", g_signal, g_calls);
    printf("  that is a result, not a crash: the driver got that far before it\n"
           "  needed one of those stubs to have done something real\n");
    goto done;
  }

  // --- L2: be the PCI core -------------------------------------------------
  if (!btfpath)
    goto done;
  if (!g_pci_driver) {
    printf("\n  no pci_driver was registered -- nothing to probe\n");
    goto done;
  }

  {
    struct btf b;
    struct pcilayout L;
    unsigned char *bimg;
    unsigned long long blen;
    unsigned int *boff;
    long nt;
    unsigned long idtab, pdev;
    int nid = 0;

    bimg = slurp(btfpath, &blen);
    if (!bimg) { fprintf(stderr, "ldls-run: cannot read %s\n", btfpath); goto done; }
    nt = btf_count(bimg, blen);
    if (nt < 0) { fprintf(stderr, "ldls-run: %s is not BTF\n", btfpath); goto done; }
    boff = calloc((size_t)nt + 2, sizeof(unsigned int));
    if (!boff || btf_open(&b, bimg, blen, boff, (unsigned int)nt + 2) < 0) {
      fprintf(stderr, "ldls-run: %s\n", btf_err(&b));
      goto done;
    }
    if (pci_layout(&b, &L) < 0)
      goto done;
    {
      int nd = btf_find(&b, "net_device", BTF_KIND_STRUCT);
      g_netdev_size = nd ? btf_type_size(&b, nd) : 4096;
    }

    printf("\n  pci_driver at %#lx, from BTF: probe +%u, id_table +%u\n",
           g_pci_driver, L.drv_probe, L.drv_idtable);
    {
      const char *dn = (const char *)rd64(g_pci_driver, L.drv_name);
      if (dn) printf("  driver name \"%s\"\n", dn);
    }

    idtab = rd64(g_pci_driver, L.drv_idtable);
    if (!idtab) { printf("  it has no id_table\n"); goto done; }

    printf("\n  devices it claims:\n");
    for (;;) {
      unsigned long e = idtab + (unsigned long)nid * L.id_size;
      unsigned int v = rd32(e, L.id_vendor), d = rd32(e, L.id_device);
      unsigned int c = rd32(e, L.id_class);
      if (!v && !d && !c) break;
      if (nid < 8)
        printf("    %04x:%04x  class %06x/%06x\n", v, d, c,
               rd32(e, L.id_class_mask));
      nid++;
      if (nid > 4096) break;
    }
    if (nid > 8) printf("    ... and %d more\n", nid - 8);
    if (!nid) { printf("    (none)\n"); goto done; }

    // A device to hand it. Zeroed, then the few fields a driver reads before
    // it has called anything: what it is, where it lives, and its BARs.
    //
    // The BAR values are invented. The point of L2 is that probe() is reached
    // with a pci_dev whose *shape* is right; making the values real is L3's
    // job and belongs to /dev/pci rather than to a fake.
    pdev = (unsigned long)calloc(1, L.dev_size);
    if (!pdev) goto done;
    {
      unsigned long e = idtab;              /* the first entry it claims */
      unsigned long r0 = pdev + L.dev_resource;
      unsigned long r1 = r0 + L.res_size;

      wr16(pdev, L.dev_vendor, (unsigned short)rd32(e, L.id_vendor));
      wr16(pdev, L.dev_device, (unsigned short)rd32(e, L.id_device));
      wr32(pdev, L.dev_devfn, 0x18);        /* 03.0 */
      wr32(pdev, L.dev_irq, 11);

      wr64(r0, L.res_start, 0xc000);
      wr64(r0, L.res_end,   0xc0ff);
      wr64(r0, L.res_flags, LDLS_IORESOURCE_IO);
      wr64(r1, L.res_start, 0xfebc0000);
      wr64(r1, L.res_end,   0xfebc00ff);
      wr64(r1, L.res_flags, LDLS_IORESOURCE_MEM);

      printf("\n  calling probe(pci_dev %04x:%04x at %#lx, id %#lx)\n\n",
             rd32(e, L.id_vendor), rd32(e, L.id_device), pdev, idtab);

      g_calls = 0;
      if (sigsetjmp(g_jump, 1) == 0) {
        int (*probe)(void *, void *) =
            (int (*)(void *, void *))(void *)rd64(g_pci_driver, L.drv_probe);
        if (!probe) { printf("  the driver has no probe\n"); goto done; }
        {
          int r = probe((void *)pdev, (void *)idtab);
          printf("\n  probe returned %d after %d stub call(s)\n", r, g_calls);
        }
      } else {
        printf("\n  probe stopped by signal %d after %d stub call(s)\n",
               g_signal, g_calls);
        printf("  which is how far a driver gets before one of those stubs\n"
               "  has to do something real. That list is L3's work queue.\n");
      }
    }
  }

done:
  ;

  free(sec);
  free(img);
  return rv;
}