// SPDX-License-Identifier: GPL-2.0-or-later
//
// The block device registry, and the partition tables on what is in it.
// See include/blkdev.h.

#include "types.h"
#include "defs.h"
#include "param.h"
#include "spinlock.h"
#include "sleeplock.h"
#include "fs.h"
#include "buf.h"
#include "blkdev.h"

struct blkdev {
  int    used;
  char   name[16];
  const struct blkdev_ops *ops;
  void  *priv;
  uint64 start;           // first sector, in the disk's own sectors
  uint64 nsectors;
  uint32 secsize;
  int    parent;          // the disk's slot for a partition, -1 for a disk
  int    schfs;           // 1: typed SchFS; -1: typed FAT; 0: anything else
};

static struct blkdev devs[NBLKDEV];
static struct spinlock blklock;
static int blkinit;

static struct blkdev *
lookup(uint dev)
{
  if (dev < BLKDEV_FIRST || dev >= BLKDEV_FIRST + NBLKDEV)
    return 0;
  struct blkdev *d = &devs[dev - BLKDEV_FIRST];
  return d->used ? d : 0;
}

static int
add(const char *name, const struct blkdev_ops *ops, void *priv,
    uint64 start, uint64 nsectors, uint32 secsize, int parent, int schfs)
{
  if (!blkinit) {
    initlock(&blklock, "blkdev");
    blkinit = 1;
  }
  acquire(&blklock);
  for (int i = 0; i < NBLKDEV; i++) {
    if (devs[i].used)
      continue;
    struct blkdev *d = &devs[i];
    safestrcpy(d->name, name, sizeof(d->name));
    d->ops = ops;
    d->priv = priv;
    d->start = start;
    d->nsectors = nsectors;
    d->secsize = secsize;
    d->parent = parent;
    d->schfs = schfs;
    __sync_synchronize();
    d->used = 1;
    release(&blklock);
    return BLKDEV_FIRST + i;
  }
  release(&blklock);
  cprintf("blkdev: no slot for %s; it will not be usable\n", name);
  return -1;
}

// --- partition tables -------------------------------------------------------

static uint32 le32(const uchar *p) { return p[0] | p[1] << 8 | p[2] << 16 | (uint32)p[3] << 24; }
static uint64 le64(const uchar *p) { return le32(p) | (uint64)le32(p + 4) << 32; }

// A GUID as GPT stores it: the first three fields little-endian, the last two
// as bytes. Compared against the printed form, so the constant in blkdev.h is
// the one a person types into sgdisk.
static int
guid_is(const uchar *g, const char *text)
{
  static const int order[16] = { 3,2,1,0, 5,4, 7,6, 8,9, 10,11,12,13,14,15 };
  uchar want[16];
  int n = 0;

  for (const char *p = text; *p && n < 32; p++) {
    int v;
    if (*p == '-')
      continue;
    if (*p >= '0' && *p <= '9') v = *p - '0';
    else if (*p >= 'a' && *p <= 'f') v = *p - 'a' + 10;
    else if (*p >= 'A' && *p <= 'F') v = *p - 'A' + 10;
    else return 0;
    if (n & 1) want[n / 2] |= v; else want[n / 2] = v << 4;
    n++;
  }
  for (int i = 0; i < 16; i++)
    if (g[order[i]] != want[i])
      return 0;
  return 1;
}

static int
guid_zero(const uchar *g)
{
  for (int i = 0; i < 16; i++)
    if (g[i])
      return 0;
  return 1;
}

static void
partname(char *out, const char *disk, int n)
{
  int i = 0;
  for (; disk[i] && i < 12; i++)
    out[i] = disk[i];
  out[i++] = 'p';
  if (n >= 10)
    out[i++] = '0' + n / 10;
  out[i++] = '0' + n % 10;
  out[i] = 0;
}

static int
readsec(struct blkdev *d, uint64 lba, uchar *buf)
{
  return d->ops->rw(d->priv, lba, 1, buf, 0);
}

#define ESP_TYPE "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"

// The GUID partition table: the header in sector 1, and its entries wherever
// it says. Only the primary copy is read -- a disk whose primary is damaged
// gets reported as having no partitions rather than half of some.
//
// The EFI system partition is registered like any other, as FAT: the root
// probe passes over it (see blkdev_is_schfs), and kernel/fs/fat mounts it.
static int
scan_gpt(int disk, uchar *sec)
{
  struct blkdev *d = &devs[disk];
  int found = 0;

  if (readsec(d, 1, sec) < 0 || memcmp(sec, "EFI PART", 8) != 0)
    return -1;

  uint64 elba  = le64(sec + 72);
  uint32 count = le32(sec + 80);
  uint32 esize = le32(sec + 84);
  if (esize < 128 || esize > d->secsize || (esize & (esize - 1)) ||
      count > 256 || elba < 2 || elba >= d->nsectors) {
    cprintf("blkdev: %s: GPT header makes no sense; ignoring it\n", d->name);
    return -1;
  }

  uint32 per = d->secsize / esize;
  for (uint32 i = 0; i < count; i++) {
    if (i % per == 0 && readsec(d, elba + i / per, sec) < 0)
      break;
    uchar *e = sec + (i % per) * esize;
    if (guid_zero(e))
      continue;
    uint64 first = le64(e + 32), last = le64(e + 40);
    if (first == 0 || last < first || last >= d->nsectors)
      continue;
    char name[16];
    partname(name, d->name, i + 1);
    int schfs = guid_is(e, SCHFS_GPT_TYPE) ? 1 : guid_is(e, ESP_TYPE) ? -1 : 0;
    if (add(name, d->ops, d->priv, first, last - first + 1, d->secsize,
            disk, schfs) >= 0) {
      kinfo("blkdev: %s: %d MB%s\n", name,
            (int)(((last - first + 1) * d->secsize) >> 20),
            schfs > 0 ? ", SchFS" : schfs < 0 ? ", EFI system" : "");
      found++;
    }
  }
  return found;
}

// The PC partition table, for disks that have one instead. Primary entries
// only; a protective entry (0xEE) means the disk is GPT and was read above.
//
// A FAT volume with no partition table -- a "superfloppy", which is how a
// good many USB sticks leave the factory -- also ends in 55 AA, and its boot
// code would read as four nonsense entries. Linux tells them apart by the
// status byte, which in a real entry is 0x00 or 0x80 and in boot code is
// whatever the code happens to be.
static int
fat_type(uint32 type)
{
  switch (type) {
  case 0x01: case 0x04: case 0x06: case 0x0B: case 0x0C: case 0x0E: case 0xEF:
    return 1;
  }
  return 0;
}

static void
scan_mbr(int disk, uchar *sec)
{
  struct blkdev *d = &devs[disk];

  if (readsec(d, 0, sec) < 0 || sec[510] != 0x55 || sec[511] != 0xAA)
    return;
  for (int i = 0; i < 4; i++) {
    uchar *e = sec + 446 + i * 16;
    uint32 type = e[4], first = le32(e + 8), n = le32(e + 12);
    // Empty, protective, and extended (which would need its own chain read).
    if ((e[0] != 0x00 && e[0] != 0x80) ||
        type == 0 || type == 0xEE || type == 0x05 || type == 0x0F ||
        n == 0 || first + (uint64)n > d->nsectors)
      continue;
    char name[16];
    partname(name, d->name, i + 1);
    if (add(name, d->ops, d->priv, first, n, d->secsize, disk,
            fat_type(type) ? -1 : 0) >= 0)
      kinfo("blkdev: %s: %d MB, MBR type %x\n", name,
            (int)(((uint64)n * d->secsize) >> 20), type);
  }
}

int
blkdev_register(const char *name, const struct blkdev_ops *ops, void *priv,
                uint64 nsectors, uint32 secsize)
{
  if (secsize == 0 || secsize > BSIZE || BSIZE % secsize)
    return -1;
  int dev = add(name, ops, priv, 0, nsectors, secsize, -1, 0);
  if (dev < 0)
    return -1;

  uchar *sec = (uchar *)kalloc();
  if (sec) {
    if (scan_gpt(dev - BLKDEV_FIRST, sec) < 0)
      scan_mbr(dev - BLKDEV_FIRST, sec);
    kfree((char *)sec);
  }
  return dev;
}

void
blkdev_unregister(int dev)
{
  struct blkdev *d = lookup(dev);
  if (!d)
    return;
  int disk = dev - BLKDEV_FIRST;

  acquire(&blklock);
  for (int i = 0; i < NBLKDEV; i++) {
    if (devs[i].used && (i == disk || devs[i].parent == disk)) {
      devs[i].used = 0;
      kinfo("blkdev: %s removed\n", devs[i].name);
      release(&blklock);
      binval(BLKDEV_FIRST + i);
      acquire(&blklock);
    }
  }
  release(&blklock);
}

// --- I/O ---------------------------------------------------------------------

int
blkdev_present(uint dev)
{
  return lookup(dev) != 0;
}

// Same contract as iderw(): fill b->data and set B_VALID on a read, write it
// out and clear B_DIRTY on a write. A failed transfer panics, for the reason
// usb_storage_rw() gives: bread() has no way to report one, and a buffer that
// silently holds the wrong bytes is how a filesystem is corrupted.
void
blkdev_rw(struct buf *b)
{
  struct blkdev *d = lookup(b->dev);
  int write = (b->flags & B_DIRTY) != 0;

  if (!holdingsleep(&b->lock))
    panic("blkdev_rw: buf not locked");
  if (!d)
    panic("blkdev_rw: no such device");

  uint32 per = BSIZE / d->secsize;
  uint64 lba = b->blockno * per;
  if (lba + per > d->nsectors)
    panic("blkdev_rw: block out of range");

  if (d->ops->rw(d->priv, d->start + lba, per, b->data, write) < 0) {
    cprintf("blkdev: %s: %s of block %d failed\n", d->name,
            write ? "write" : "read", (int)b->blockno);
    panic("blkdev_rw: transfer failed");
  }
  b->flags |= B_VALID;
  b->flags &= ~B_DIRTY;
}

int
blkdev_xfer(uint dev, uint64 lba, uint32 count, void *buf, int write)
{
  struct blkdev *d = lookup(dev);

  if (!d || count == 0 || (uint64)count * d->secsize > BSIZE ||
      lba + count > d->nsectors || lba + count < lba)
    return -1;
  return d->ops->rw(d->priv, d->start + lba, count, buf, write) < 0 ? -1 : 0;
}

uint64
blkdev_sectors(uint dev, uint32 *secsize)
{
  struct blkdev *d = lookup(dev);
  if (!d)
    return 0;
  if (secsize)
    *secsize = d->secsize;
  return d->nsectors;
}

int
blkdev_flush(uint dev)
{
  struct blkdev *d = lookup(dev);
  if (!d || !d->ops->flush)
    return 0;
  return d->ops->flush(d->priv);
}

int
blkdev_nth(int i)
{
  for (int s = 0; s < NBLKDEV; s++)
    if (devs[s].used && i-- == 0)
      return BLKDEV_FIRST + s;
  return -1;
}

const char *
blkdev_name(uint dev)
{
  struct blkdev *d = lookup(dev);
  return d ? d->name : "?";
}

int
blkdev_is_schfs(uint dev)
{
  struct blkdev *d = lookup(dev);
  return d ? d->schfs : 0;
}

// --- /proc/partitions --------------------------------------------------------

static int
put(char *buf, int cap, int off, const char *s)
{
  while (*s && off < cap)
    buf[off++] = *s++;
  return off;
}

static int
putu(char *buf, int cap, int off, uint64 v)
{
  char t[24];
  int n = 0;
  do { t[n++] = '0' + v % 10; v /= 10; } while (v);
  while (n && off < cap)
    buf[off++] = t[--n];
  return off;
}

int
blkdev_info(char *buf, int cap)
{
  int off = put(buf, cap, 0, "dev  name          MB  type\n");
  for (int s = 0; s < NBLKDEV; s++) {
    struct blkdev *d = &devs[s];
    if (!d->used)
      continue;
    off = putu(buf, cap, off, BLKDEV_FIRST + s);
    off = put(buf, cap, off, "    ");
    off = put(buf, cap, off, d->name);
    for (int pad = 12 - strlen(d->name); pad > 0; pad--)
      off = put(buf, cap, off, " ");
    off = putu(buf, cap, off, (d->nsectors * d->secsize) >> 20);
    off = put(buf, cap, off, d->parent < 0 ? "  disk" :
                             d->schfs > 0 ? "  schfs" :
                             d->schfs < 0 ? "  fat" : "  part");
    off = put(buf, cap, off, "\n");
  }
  return off;
}