// SPDX-License-Identifier: Apache-2.0
//
// Loading a relocatable object. See ldls/ko.h.
//
// Written from the ELF-64 gABI and the x86-64 psABI, both published
// specifications -- ../docs/licensing.md §5. Nothing in this file was taken
// from anyone's loader, and nothing in it needs to be: the format is fully
// described and the relocation arithmetic is three lines per type.
//
// Every read from the image is bounds-checked against the length the caller
// gave, and that is not defensive programming for its own sake. The input is a
// file the *user* supplied, from a kernel we did not build, and a loader that
// trusts its input is a loader that can be handed a crash.

#include "ko.h"

// --- small local helpers, so this file needs no libc on either side ---------

static void
ko_memcpy(unsigned char *d, const unsigned char *s, unsigned long long n)
{
  while (n--)
    *d++ = *s++;
}

static void
ko_memset(unsigned char *d, int c, unsigned long long n)
{
  while (n--)
    *d++ = (unsigned char)c;
}

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

static void
ko_fail(struct ko *k, const char *msg, const char *what)
{
  int i = 0;
  while (msg[i] && i < (int)sizeof(k->err) - 1) { k->err[i] = msg[i]; i++; }
  if (what) {
    int j = 0;
    while (what[j] && i < (int)sizeof(k->err) - 1) k->err[i++] = what[j++];
  }
  k->err[i] = 0;
}

// A number, appended to the error. There is no formatter here and an error
// that says "relocation type" without saying which is not worth printing.
static void
ko_fail_num(struct ko *k, const char *msg, unsigned long long v)
{
  char d[24];
  int  n = 0, i = 0;

  while (msg[i] && i < (int)sizeof(k->err) - 1) { k->err[i] = msg[i]; i++; }
  if (v == 0) d[n++] = '0';
  while (v) { d[n++] = (char)('0' + v % 10); v /= 10; }
  while (n-- > 0 && i < (int)sizeof(k->err) - 1)
    k->err[i++] = d[n];
  k->err[i] = 0;
}

// Is [off, off+n) inside the image? Every accessor goes through this.
static int
ko_fits(const struct ko *k, unsigned long long off, unsigned long long n)
{
  if (off > k->len)
    return 0;
  if (n > k->len - off)
    return 0;
  return 1;
}

// A NUL-terminated string at `off` in a table of `size` bytes starting at
// `base` -- checked so that an unterminated string table cannot run away.
static const char *
ko_str_at(const struct ko *k, unsigned long long base, unsigned long long size,
          unsigned int off)
{
  unsigned long long i;

  if (!ko_fits(k, base, size) || off >= size)
    return 0;
  for (i = base + off; i < base + size; i++)
    if (k->img[i] == 0)
      return (const char *)(k->img + base + off);
  return 0;                      // runs to the end without a NUL
}

static const struct ko_shdr *
ko_sh(const struct ko *k, int i)
{
  if (i < 0 || i >= k->nsh)
    return 0;
  return (const struct ko_shdr *)(k->shtab +
                                  (unsigned long long)i * k->eh->e_shentsize);
}

// --- opening -----------------------------------------------------------------

int
ko_sections_needed(const void *image, unsigned long long len)
{
  const struct ko_ehdr *eh = (const struct ko_ehdr *)image;

  if (len < sizeof(struct ko_ehdr))
    return -1;
  if (eh->e_ident[0] != 0x7f || eh->e_ident[1] != 'E' ||
      eh->e_ident[2] != 'L'  || eh->e_ident[3] != 'F')
    return -1;
  if (eh->e_ident[4] != 2 || eh->e_type != KO_ET_REL)
    return -1;
  return eh->e_shnum;
}

int
ko_open(struct ko *k, const void *image, unsigned long long len,
        struct ko_section *sec, int nsec)
{
  const struct ko_ehdr *eh;
  int i;

  ko_memset((unsigned char *)k, 0, sizeof(*k));
  k->img = (const unsigned char *)image;
  k->len = len;
  k->sec = sec;

  if (!sec || nsec <= 0)
    { ko_fail(k, "no section table was provided", 0); return -1; }

  if (len < sizeof(struct ko_ehdr))
    { ko_fail(k, "not an ELF object: too short", 0); return -1; }

  eh = (const struct ko_ehdr *)image;
  if (eh->e_ident[0] != 0x7f || eh->e_ident[1] != 'E' ||
      eh->e_ident[2] != 'L'  || eh->e_ident[3] != 'F')
    { ko_fail(k, "not an ELF object: bad magic", 0); return -1; }
  if (eh->e_ident[4] != 2)       // ELFCLASS64
    { ko_fail(k, "not a 64-bit object", 0); return -1; }
  if (eh->e_ident[5] != 1)       // ELFDATA2LSB
    { ko_fail(k, "not little-endian", 0); return -1; }
  if (eh->e_type != KO_ET_REL)
    { ko_fail(k, "not a relocatable object -- a .ko is ET_REL, this is not", 0);
      return -1; }
  if (eh->e_machine != KO_EM_X86_64)
    { ko_fail_num(k, "wrong machine, e_machine=", eh->e_machine); return -1; }
  if (eh->e_shentsize < sizeof(struct ko_shdr))
    { ko_fail(k, "section headers are smaller than ELF64 requires", 0);
      return -1; }

  k->eh  = eh;
  k->nsh = eh->e_shnum;
  if (k->nsh <= 0)
    { ko_fail(k, "no section headers", 0); return -1; }
  if (k->nsh > nsec)
    { ko_fail_num(k, "the caller's section table is too small; it needs ",
                  (unsigned long long)k->nsh); return -1; }
  ko_memset((unsigned char *)sec, 0,
            (unsigned long long)k->nsh * sizeof(struct ko_section));

  if (!ko_fits(k, eh->e_shoff,
               (unsigned long long)k->nsh * eh->e_shentsize))
    { ko_fail(k, "section headers run off the end of the file", 0); return -1; }
  k->shtab = k->img + eh->e_shoff;

  // Section names.
  {
    const struct ko_shdr *sh = ko_sh(k, eh->e_shstrndx);
    if (!sh || sh->sh_type != KO_SHT_STRTAB)
      { ko_fail(k, "no section-name string table", 0); return -1; }
    if (!ko_fits(k, sh->sh_offset, sh->sh_size))
      { ko_fail(k, "section names run off the end of the file", 0); return -1; }
    k->shstr = (const char *)(k->img + sh->sh_offset);
  }

  // Index every section, and find the symbol table on the way past.
  for (i = 0; i < k->nsh; i++) {
    const struct ko_shdr *sh = ko_sh(k, i);
    struct ko_section    *s  = &k->sec[i];

    {
      const struct ko_shdr *ss = ko_sh(k, k->eh->e_shstrndx);
      const char *nm = ko_str_at(k, ss->sh_offset, ss->sh_size, sh->sh_name);
      // An unreadable name is not fatal -- the section is still loadable and
      // still relocatable. It only costs the report a label.
      s->name = nm ? nm : "(unnamed)";
    }
    s->type  = sh->sh_type;
    s->flags = sh->sh_flags;
    s->size  = sh->sh_size;
    s->align = sh->sh_addralign ? sh->sh_addralign : 1;
    s->alloc = (sh->sh_flags & KO_SHF_ALLOC) != 0;
    s->data  = 0;

    if (sh->sh_type != KO_SHT_NOBITS && sh->sh_size) {
      if (!ko_fits(k, sh->sh_offset, sh->sh_size))
        { ko_fail(k, "section runs off the end of the file: ", s->name);
          return -1; }
      s->data = k->img + sh->sh_offset;
    }

    if (sh->sh_type == KO_SHT_SYMTAB) {
      const struct ko_shdr *st;
      if (sh->sh_entsize != sizeof(struct ko_sym))
        { ko_fail(k, "symbol entries are the wrong size", 0); return -1; }
      if (!s->data)
        { ko_fail(k, "empty symbol table", 0); return -1; }
      k->sym  = (const struct ko_sym *)s->data;
      k->nsym = (int)(sh->sh_size / sizeof(struct ko_sym));

      st = ko_sh(k, sh->sh_link);
      if (!st || st->sh_type != KO_SHT_STRTAB)
        { ko_fail(k, "symbol table has no string table", 0); return -1; }
      if (!ko_fits(k, st->sh_offset, st->sh_size))
        { ko_fail(k, "symbol names run off the end of the file", 0);
          return -1; }
      k->str = (const char *)(k->img + st->sh_offset);
    }
  }

  if (!k->sym)
    { ko_fail(k, "no symbol table -- nothing could be resolved", 0);
      return -1; }

  // The undefined symbols. This is L0's whole output and the reason the
  // milestone exists: it is the API surface of this driver, measured rather
  // than guessed at.
  for (i = 0; i < k->nsym; i++) {
    const char *nm;

    if (k->sym[i].st_shndx != KO_SHN_UNDEF)
      continue;
    if (k->sym[i].st_name == 0)
      continue;                  // the index-0 placeholder
    nm = (const char *)(k->str + k->sym[i].st_name);
    if (nm[0] == 0)
      continue;

    if (k->nundef < KO_MAXUNDEF)
      k->undef[k->nundef++] = nm;
    else
      k->undef_lost = 1;
  }

  return 0;
}

const char *ko_err(const struct ko *k) { return k->err[0] ? k->err : "no error"; }

// --- layout ------------------------------------------------------------------

static unsigned long long
ko_round(unsigned long long v, unsigned long long a)
{
  if (a <= 1)
    return v;
  return (v + a - 1) & ~(a - 1);
}

// Where each allocatable section lands, and how much that comes to.
//
// One contiguous block, sections in file order, each at its own alignment.
// Linux's own loader splits by permission so that text can be made read-only;
// that is a refinement worth having later and is not what L0 is measuring.
static unsigned long long
ko_plan(struct ko *k, unsigned long long base, void *mem)
{
  unsigned long long off = 0;
  int i;

  for (i = 0; i < k->nsh; i++) {
    struct ko_section *s = &k->sec[i];
    if (!s->alloc || s->size == 0)
      continue;
    off = ko_round(off, s->align);
    if (mem) {
      s->addr  = base + off;
      s->where = (unsigned char *)mem + off;
    }
    off += s->size;
  }
  return off;
}

unsigned long long
ko_span(const struct ko *k)
{
  return ko_plan((struct ko *)k, 0, 0);
}

int
ko_layout(struct ko *k, void *mem, unsigned long long cap,
          unsigned long long base)
{
  unsigned long long need = ko_plan(k, 0, 0);
  int i;

  if (need > cap)
    { ko_fail_num(k, "the object does not fit; it needs ", need); return -1; }

  k->span = ko_plan(k, base, mem);

  for (i = 0; i < k->nsh; i++) {
    struct ko_section *s = &k->sec[i];
    if (!s->alloc || s->size == 0)
      continue;
    if (s->type == KO_SHT_NOBITS)
      ko_memset(s->where, 0, s->size);     // .bss, and it must be zeroed
    else
      ko_memcpy(s->where, s->data, s->size);
  }

  k->laid_out = 1;
  return 0;
}

// --- relocation --------------------------------------------------------------

// The address a symbol will have once loaded.
static int
ko_value(struct ko *k, unsigned int idx, ko_resolver r, void *ctx,
         unsigned long long *out)
{
  const struct ko_sym *sy;
  const char *nm;

  if ((int)idx >= k->nsym)
    { ko_fail_num(k, "relocation names symbol ", idx); return -1; }
  sy = &k->sym[idx];

  if (sy->st_shndx == KO_SHN_ABS) {
    *out = sy->st_value;
    return 0;
  }

  if (sy->st_shndx == KO_SHN_UNDEF) {
    nm = (const char *)(k->str + sy->st_name);
    if (!r || r(ctx, nm, out) != 0)
      { ko_fail(k, "unresolved: ", nm); return -1; }
    return 0;
  }

  if (sy->st_shndx == KO_SHN_COMMON)
    { ko_fail(k, "common symbols are not supported: ",
              (const char *)(k->str + sy->st_name)); return -1; }

  if (sy->st_shndx >= k->nsh)
    { ko_fail_num(k, "symbol points at section ", sy->st_shndx); return -1; }

  // A symbol defined in a section that was not laid out has no address. That
  // happens for debug sections, and a relocation against one means the object
  // is asking for something the loader deliberately did not load.
  if (!k->sec[sy->st_shndx].alloc)
    { ko_fail(k, "relocation against a section that is not loaded: ",
              k->sec[sy->st_shndx].name); return -1; }

  *out = k->sec[sy->st_shndx].addr + sy->st_value;
  return 0;
}

void
ko_reltypes(struct ko *k)
{
  int i;

  for (i = 0; i < 64; i++)
    k->reltype[i] = 0;
  k->reltype_other = 0;
  k->nrel = 0;

  for (i = 0; i < k->nsh; i++) {
    const struct ko_shdr *sh = ko_sh(k, i);
    unsigned long long n, j;

    if (!sh || sh->sh_type != KO_SHT_RELA)
      continue;
    if (sh->sh_entsize != sizeof(struct ko_rela) || !k->sec[i].data)
      continue;
    if (sh->sh_info >= (unsigned int)k->nsh || !k->sec[sh->sh_info].alloc)
      continue;                  // relocations for sections we do not load

    n = sh->sh_size / sizeof(struct ko_rela);
    for (j = 0; j < n; j++) {
      const struct ko_rela *rl = (const struct ko_rela *)(k->sec[i].data) + j;
      unsigned int t = KO_R_TYPE(rl->r_info);
      k->nrel++;
      if (t < 64)
        k->reltype[t]++;
      else
        k->reltype_other++;
    }
  }
}

int
ko_relocate(struct ko *k, ko_resolver r, void *ctx)
{
  int i;

  if (!k->laid_out)
    { ko_fail(k, "relocate before layout", 0); return -1; }

  for (i = 0; i < k->nsh; i++) {
    const struct ko_shdr *sh = ko_sh(k, i);
    struct ko_section *tgt;
    unsigned long long n, j;

    if (!sh || sh->sh_type != KO_SHT_RELA)
      continue;
    if (sh->sh_entsize != sizeof(struct ko_rela))
      { ko_fail(k, "relocation entries are the wrong size", 0); return -1; }
    if (sh->sh_info >= (unsigned int)k->nsh)
      { ko_fail_num(k, "relocations target section ", sh->sh_info); return -1; }

    tgt = &k->sec[sh->sh_info];
    if (!tgt->alloc || tgt->size == 0)
      continue;                  // relocations for .debug_* and friends

    n = sh->sh_size / sizeof(struct ko_rela);
    for (j = 0; j < n; j++) {
      const struct ko_rela *rl =
          (const struct ko_rela *)(k->sec[i].data) + j;
      unsigned int  type = KO_R_TYPE(rl->r_info);
      unsigned long long S, P;
      unsigned char *loc;
      long long      X;

      if (rl->r_offset + 4 > tgt->size &&
          !(type == KO_R_X86_64_64 || type == KO_R_X86_64_PC64))
        { ko_fail(k, "relocation points outside ", tgt->name); return -1; }
      if ((type == KO_R_X86_64_64 || type == KO_R_X86_64_PC64) &&
          rl->r_offset + 8 > tgt->size)
        { ko_fail(k, "relocation points outside ", tgt->name); return -1; }

      if (type == KO_R_X86_64_NONE)
        continue;

      if (ko_value(k, KO_R_SYM(rl->r_info), r, ctx, &S) < 0)
        return -1;

      loc = tgt->where + rl->r_offset;
      P   = tgt->addr  + rl->r_offset;
      X   = rl->r_addend;

      switch (type) {
      case KO_R_X86_64_64:
        *(unsigned long long *)loc = S + (unsigned long long)X;
        break;
      case KO_R_X86_64_PC64:
        *(unsigned long long *)loc = S + (unsigned long long)X - P;
        break;
      case KO_R_X86_64_PC32:
      case KO_R_X86_64_PLT32: {
        // PLT32 is treated as PC32. There is no procedure linkage table here:
        // every call is resolved at load time, so the "call through the PLT"
        // a compiler emitted is just a call to where the symbol ended up.
        long long v = (long long)S + X - (long long)P;
        if (v < -0x80000000LL || v > 0x7fffffffLL)
          { ko_fail(k, "a 32-bit PC-relative relocation does not reach; the "
                       "object was loaded too far away", 0); return -1; }
        *(unsigned int *)loc = (unsigned int)(int)v;
        break;
      }
      case KO_R_X86_64_32: {
        unsigned long long v = S + (unsigned long long)X;
        if (v > 0xffffffffULL)
          { ko_fail(k, "a 32-bit absolute relocation does not fit", 0);
            return -1; }
        *(unsigned int *)loc = (unsigned int)v;
        break;
      }
      case KO_R_X86_64_32S: {
        long long v = (long long)S + X;
        if (v < -0x80000000LL || v > 0x7fffffffLL)
          { ko_fail(k, "a signed 32-bit relocation does not fit", 0);
            return -1; }
        *(unsigned int *)loc = (unsigned int)(int)v;
        break;
      }
      default:
        // Named, not ignored. A relocation quietly skipped is a wrong address
        // discovered a long way from here.
        ko_fail_num(k, "unhandled relocation type ", type);
        return -1;
      }
    }
  }

  return 0;
}

// --- the %gs stack guard ------------------------------------------------------

// Rewrite one `mov %gs:disp32(%rip), %reg` into `xor %reg, %reg` plus padding.
//
// The encoding is fixed and short: 65 (the GS prefix), a REX byte, 8b (MOV
// r64, r/m64), a ModRM of mod=00 reg=DEST rm=101 (RIP-relative), then the
// displacement. Eight bytes, of which the last four are what the relocation
// was going to fill in.
//
// The replacement has to clear the *same* register, so the destination is
// decoded rather than assumed: REX.R is the high bit and ModRM.reg the low
// three. `xor %rN, %rN` is three bytes, and the rest is padded with nops.
static int
defuse_one(unsigned char *insn)
{
  unsigned int rex, modrm, reg;

  if (insn[0] != 0x65)                 return -1;   /* not GS-prefixed   */
  if ((insn[1] & 0xF0) != 0x40)        return -1;   /* not a REX byte    */
  // Two forms, and both have to go. The prologue *loads* the guard with
  // MOV (8b); the epilogue *compares* it with SUB (2b), followed by a jne to
  // __stack_chk_fail. Replacing either with `xor %reg,%reg` is right: the
  // first makes the saved canary zero, and the second sets the register to
  // zero *and* sets ZF, so the jne falls through exactly as a passing check
  // does.
  if (insn[2] != 0x8b && insn[2] != 0x2b) return -1;
  rex   = insn[1];
  modrm = insn[3];
  if ((modrm & 0xC7) != 0x05)          return -1;   /* not mod=00 rm=101 */

  reg = ((rex & 0x4) << 1) | ((modrm >> 3) & 0x7);

  insn[0] = (unsigned char)(0x48 | ((reg >= 8) ? 0x05 : 0x00));  /* REX.W(+R+B) */
  insn[1] = 0x31;                                                /* XOR r/m64,r64 */
  insn[2] = (unsigned char)(0xC0 | ((reg & 7) << 3) | (reg & 7));
  insn[3] = 0x90;
  insn[4] = 0x90; insn[5] = 0x90; insn[6] = 0x90; insn[7] = 0x90;
  return 0;
}

int
ko_defuse_gs(struct ko *k, const char *sym)
{
  int i, n = 0;

  if (!k->laid_out)
    { ko_fail(k, "defuse before layout", 0); return -1; }

  // Two passes: check everything first, then patch. A loader that rewrites
  // half the instructions and then refuses has left the object in a state
  // nobody can reason about.
  for (int pass = 0; pass < 2; pass++) {
    n = 0;
    for (i = 0; i < k->nsh; i++) {
      const struct ko_shdr *sh = ko_sh(k, i);
      struct ko_section *tgt;
      unsigned long long cnt, j;

      if (!sh || sh->sh_type != KO_SHT_RELA || !k->sec[i].data)
        continue;
      if (sh->sh_info >= (unsigned int)k->nsh)
        continue;
      tgt = &k->sec[sh->sh_info];
      if (!tgt->alloc || tgt->size == 0)
        continue;

      cnt = sh->sh_size / sizeof(struct ko_rela);
      for (j = 0; j < cnt; j++) {
        const struct ko_rela *rl =
            (const struct ko_rela *)(k->sec[i].data) + j;
        unsigned int si = KO_R_SYM(rl->r_info);
        const char *nm;

        if ((int)si >= k->nsym)
          continue;
        nm = k->str + k->sym[si].st_name;
        if (!ko_streq(nm, sym))
          continue;
        if (rl->r_offset < 4 || rl->r_offset + 4 > tgt->size)
          { ko_fail(k, "a guard relocation is at an impossible offset", 0);
            return -1; }

        {
          unsigned char *insn = tgt->where + rl->r_offset - 4;
          if (pass == 0) {
            // Check only.
            if (insn[0] != 0x65 || (insn[1] & 0xF0) != 0x40 ||
                (insn[2] != 0x8b && insn[2] != 0x2b) ||
                (insn[3] & 0xC7) != 0x05)
              { ko_fail(k, "unrecognised instruction at a guard reference: ",
                        sym); return -1; }
          } else if (defuse_one(insn) < 0) {
            ko_fail(k, "failed to rewrite a guard reference", 0);
            return -1;
          }
          n++;
        }
      }
    }
  }
  return n;
}

// --- what the object says about itself ---------------------------------------

static const struct ko_section *
ko_find(const struct ko *k, const char *name)
{
  int i;
  for (i = 0; i < k->nsh; i++)
    if (ko_streq(k->sec[i].name, name))
      return &k->sec[i];
  return 0;
}

// `.modinfo` is a run of NUL-terminated "key=value" strings. Reading it needs
// nothing loaded and no symbols resolved, which is why the licence a module
// declares can be printed before deciding whether to load it at all.
const char *
ko_modinfo(const struct ko *k, const char *key)
{
  const struct ko_section *s = ko_find(k, ".modinfo");
  unsigned long long i = 0;

  if (!s || !s->data)
    return 0;

  while (i < s->size) {
    const char *e = (const char *)(s->data + i);
    unsigned long long n = 0;
    const char *p = e;
    const char *q = key;

    while (i + n < s->size && e[n])
      n++;
    if (n == 0) { i++; continue; }

    while (*q && *p == *q) { p++; q++; }
    if (*q == 0 && *p == '=')
      return p + 1;

    i += n + 1;
  }
  return 0;
}

const char *
ko_name(const struct ko *k)
{
  int i;

  // The module's `struct module` lives in `.gnu.linkonce.this_module` and its
  // name is at a fixed offset -- but that offset is a property of the kernel
  // version, which is exactly the kind of thing docs/plan/ldls.md rule 4 says
  // not to assume. The symbol table has the same information portably: the
  // object defines `__this_module`, and `.modinfo` carries the rest.
  for (i = 0; i < k->nsym; i++) {
    const char *nm = (const char *)(k->str + k->sym[i].st_name);
    if (ko_streq(nm, "__this_module"))
      return ko_modinfo(k, "name") ? ko_modinfo(k, "name") : "(unnamed)";
  }
  return 0;
}