// USB mass storage: bulk-only transport, SCSI transparent command set.
//
// Almost every USB stick, external drive and card reader is this, and the
// name is a fair description of the whole protocol: there are exactly two
// bulk endpoints and everything happens by sending three things over them.
//
//   command   a 31-byte wrapper (CBW) carrying a SCSI command block, OUT
//   data      the payload, IN or OUT, if the command has one
//   status    a 13-byte wrapper (CSW) saying whether it worked, IN
//
// That is the entire transport. The interesting content is SCSI, which is
// what the wrapper carries and what the device actually speaks -- so a USB
// stick is a SCSI disk that happens to be reached over USB, and the commands
// below (INQUIRY, READ CAPACITY, READ(10), WRITE(10)) are the same ones a
// SCSI drive answered in 1990.

#include "core/types.h"
#include "core/defs.h"
#include "core/param.h"
#include "arch/memlayout.h"
#include "arch/mmu.h"
#include "core/spinlock.h"
#include "core/sleeplock.h"
#include "fs/fs.h"
#include "fs/buf.h"
#include "core/dma.h"
#include "dev/usb.h"
#include "fs/blkdev.h"

#define CBW_SIGNATURE  0x43425355       // "USBC"
#define CSW_SIGNATURE  0x53425355       // "USBS"

struct cbw {
  uint32 signature;
  uint32 tag;
  uint32 data_len;
  uint8  flags;                         // 0x80 = device-to-host
  uint8  lun;
  uint8  cb_len;                        // 1..16
  uint8  cb[16];
} __attribute__((packed));

struct csw {
  uint32 signature;
  uint32 tag;
  uint32 residue;
  uint8  status;                        // 0 = passed, 1 = failed, 2 = phase error
} __attribute__((packed));

// SCSI opcodes
#define SCSI_TEST_UNIT_READY  0x00
#define SCSI_REQUEST_SENSE    0x03
#define SCSI_INQUIRY          0x12
#define SCSI_READ_CAPACITY10  0x25
#define SCSI_READ10           0x28
#define SCSI_WRITE10          0x2A
// SYNCHRONIZE CACHE(10): commit the device's write cache to the medium. The
// SCSI spelling of ATA's FLUSH CACHE, and needed for the same reason -- see
// ideflush() and the barrier in schfs_vol_commit.
#define SCSI_SYNC_CACHE10     0x35

#define SECTOR_SIZE 512

// One per stick. A desk with a stick in each port is a reasonable thing to
// have, and the first one to arrive is the one the fixed USBDISK_DEV means.
struct ums {
  struct usb_dev *dev;
  int      present;
  int      gone;                        // unplugged; fail everything
  uint32   nsectors;
  uint32   sector_size;
  uint32   tag;
  void    *xfer;                        // DMA bounce buffer, one block
  struct cbw *cbw;                      // allocated once, not per command
  struct csw *csw;
  struct sleeplock lock;                // one command at a time
  char     name[8];                     // usb0, usb1, ...
  int      blkdev;                      // its number in the registry, or -1
};

#define NUMS 4
static struct ums disks[NUMS];

// What USBDISK_DEV (3) means: the first stick still present. It is the whole
// medium, which is how a stick holding a bare SchFS volume is mounted.
static struct ums *
legacy(void)
{
  for(int i = 0; i < NUMS; i++)
    if(disks[i].present && !disks[i].gone)
      return &disks[i];
  return 0;
}

// SCSI is big-endian, everywhere, for everything. x86 is not.
static void
be32(uint8 *p, uint32 v)
{
  p[0] = v >> 24; p[1] = v >> 16; p[2] = v >> 8; p[3] = v;
}

static uint32
rd_be32(const uint8 *p)
{
  return ((uint32)p[0] << 24) | ((uint32)p[1] << 16) |
         ((uint32)p[2] << 8) | p[3];
}

// One complete command-data-status exchange.
//
// `buf`, when there is data, must be DMA memory -- it is handed to the
// controller as a physical address.
static int
bot_command(struct ums *u, const uint8 *cdb, int cdb_len, void *buf, int len,
            int in)
{
  struct usb_dev *dev = u->dev;
  struct cbw *cbw = u->cbw;
  struct csw *csw = u->csw;
  uint32 tag = ++u->tag;
  int i, data_ok = 1;

  memset(cbw, 0, sizeof(*cbw));
  cbw->signature = CBW_SIGNATURE;
  cbw->tag       = tag;
  cbw->data_len  = len;
  cbw->flags     = in ? 0x80 : 0x00;
  cbw->lun       = 0;
  cbw->cb_len    = cdb_len;
  for(i = 0; i < cdb_len && i < 16; i++)
    cbw->cb[i] = cdb[i];

  if(dev->ops->bulk(dev, dev->bulk_out.bEndpointAddress, cbw,
                    sizeof(*cbw)) != (int)sizeof(*cbw)){
    cprintf("usb-storage: the device would not take the command\n");
    return -1;
  }

  if(len > 0){
    int addr = in ? dev->bulk_in.bEndpointAddress
                  : dev->bulk_out.bEndpointAddress;
    if(dev->ops->bulk(dev, addr, buf, len) < 0)
      data_ok = 0;
  }

  // The status wrapper is not optional and not advisory: a device that has
  // sent its data is still waiting to be asked how it went.
  //
  // It is read even when the data phase failed, which is the whole point --
  // returning early there would leave the device still holding a status it
  // has not been asked for, so the *next* command's data phase would receive
  // this command's status instead. One failed read would turn into every
  // subsequent read failing, which is a far worse fault than the one that
  // started it, and one that looks nothing like its cause.
  memset(csw, 0, sizeof(*csw));
  if(dev->ops->bulk(dev, dev->bulk_in.bEndpointAddress, csw,
                    sizeof(*csw)) != (int)sizeof(*csw)){
    cprintf("usb-storage: no status returned for command 0x%x\n", cdb[0]);
    return -1;
  }

  if(csw->signature != CSW_SIGNATURE || csw->tag != tag){
    cprintf("usb-storage: status wrapper does not match the command\n");
    return -1;
  }
  if(csw->status != 0){
    cprintf("usb-storage: command 0x%x failed (status %d)\n",
            cdb[0], csw->status);
    return -1;
  }
  if(!data_ok)
    return -1;

  // Residue is what was *not* transferred, and a device is allowed to report
  // more than was asked for only if it is broken -- clamp rather than return
  // a negative length that a caller would read as an error code.
  if(csw->residue > (uint32)len)
    return -1;
  return len - (int)csw->residue;
}

static int
scsi_inquiry(struct ums *u, char *vendor, char *product)
{
  uint8 cdb[6] = { SCSI_INQUIRY, 0, 0, 0, 36, 0 };
  uint8 *r = u->xfer;
  int i;

  if(bot_command(u, cdb, sizeof(cdb), r, 36, 1) < 0)
    return -1;

  // Bytes 8-15 are the vendor, 16-31 the product, space padded and not
  // NUL-terminated -- the only fixed-width strings left in computing.
  for(i = 0; i < 8; i++)
    vendor[i] = r[8 + i];
  vendor[8] = 0;
  for(i = 0; i < 16; i++)
    product[i] = r[16 + i];
  product[16] = 0;
  return 0;
}

static int
scsi_read_capacity(struct ums *u)
{
  uint8 cdb[10] = { SCSI_READ_CAPACITY10, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
  uint8 *r = u->xfer;

  if(bot_command(u, cdb, sizeof(cdb), r, 8, 1) < 0)
    return -1;

  // The first word is the address of the *last* sector, not the count.
  u->nsectors    = rd_be32(r) + 1;
  u->sector_size = rd_be32(r + 4);
  return 0;
}

static int
scsi_rw(struct ums *u, uint32 lba, int count, void *buf, int write)
{
  uint8 cdb[10];

  memset(cdb, 0, sizeof(cdb));
  cdb[0] = write ? SCSI_WRITE10 : SCSI_READ10;
  be32(&cdb[2], lba);
  cdb[7] = (count >> 8) & 0xFF;
  cdb[8] = count & 0xFF;

  return bot_command(u, cdb, sizeof(cdb), buf, count * SECTOR_SIZE, !write);
}

// ---------------------------------------------------------------------------
// Block device
// ---------------------------------------------------------------------------

int
usb_storage_present(void)
{
  return legacy() != 0;
}

// Sync one buffer-cache block with the stick. Same contract as iderw() and
// ramdiskrw(): on a read fill b->data and set B_VALID, on a write push it out
// and clear B_DIRTY.
void
usb_storage_rw(struct buf *b)
{
  int per_block = BSIZE / SECTOR_SIZE;
  uint32 lba = b->blockno * per_block;
  int write = (b->flags & B_DIRTY) != 0;
  int i, n = -1;
  struct ums *u = legacy();

  if(!holdingsleep(&b->lock))
    panic("usb_storage_rw: buf not locked");
  if(u == 0)
    panic("usb_storage_rw: no USB disk");
  if(lba + per_block > u->nsectors)
    panic("usb_storage_rw: block out of range");

  // One command at a time. Bulk-only transport has no tagging beyond the CBW
  // tag and no queueing: a second command sent while the first is mid-phase
  // desynchronises the endpoint and every command after it fails.
  acquiresleep(&u->lock);

  // Removable media fails transiently in ways a fixed disk does not, so a
  // couple of retries are worth having before giving up.
  for(i = 0; i < 3; i++){
    if(write)
      memmove(u->xfer, b->data, BSIZE);
    n = scsi_rw(u, lba, per_block, u->xfer, write);
    if(n == BSIZE)
      break;
    cprintf("usb-storage: %s of block %d failed, retrying\n",
            write ? "write" : "read", b->blockno);
  }

  releasesleep(&u->lock);

  // A read that failed must not return quietly. bread() hands the buffer
  // straight back to its caller, which has no way to ask whether it worked --
  // so leaving B_VALID clear would deliver whatever the buffer happened to
  // contain as though it were the block on disk. Corrupting a filesystem
  // silently is worse than stopping.
  if(n != BSIZE)
    panic("usb_storage_rw: transfer failed");

  if(!write)
    memmove(b->data, u->xfer, BSIZE);
  b->flags |= B_VALID;
  b->flags &= ~B_DIRTY;
}

// Make everything written so far durable. Returns 0, or -1 if the stick would
// not confirm it.
//
// A whole-device sync with no LBA range, which is what the callers want: the
// filesystem's barrier is "everything before this point", not "these blocks".
//
// Cheap sticks routinely reject this command outright. That is reported rather
// than ignored, because the caller is about to write an uberblock on the
// strength of it -- and a barrier that silently does nothing is worse than no
// barrier at all, since it makes an unsafe volume look safe.
static int
ums_flush(struct ums *u)
{
  uint8 cdb[10];
  int r;

  memset(cdb, 0, sizeof(cdb));
  cdb[0] = SCSI_SYNC_CACHE10;   // LBA 0, count 0: the entire medium

  acquiresleep(&u->lock);
  r = u->gone ? -1 : bot_command(u, cdb, sizeof(cdb), 0, 0, 0);
  releasesleep(&u->lock);

  if(r < 0){
    cprintf("usb-storage: SYNCHRONIZE CACHE failed; writes are not durable\n");
    return -1;
  }
  return 0;
}

int
usb_storage_flush(void)
{
  struct ums *u = legacy();
  return u ? ums_flush(u) : 0;
}

// Read block 1 through the ordinary buffer cache and report what is there.
//
// This is the only test that proves the whole path rather than a piece of it:
// bread() -> bdevrw() -> usb_storage_rw() -> bulk-only transport -> SCSI
// READ(10) -> an interrupt -> a sleeping process woken with data. Everything
// before it -- INQUIRY answering, a plausible capacity -- can be true of a
// device that cannot actually read a sector.
//
// Reached with `echo s > /proc/sysrq`.
void
usb_storage_dump(void)
{
  struct buf *b;
  struct superblock *sb;
  struct ums *u = legacy();

  if(u == 0){
    cprintf("usb-storage: no disk\n");
    return;
  }

  kinfo("usb-storage: %d sectors of %d bytes\n", u->nsectors,
          u->sector_size);

  b = bread(USBDISK_DEV, 1);
  sb = (struct superblock*)b->data;
  kinfo("usb-storage: block 1 magic %x size %d nblocks %d ninodes %d\n",
          sb->magic, sb->size, sb->nblocks, sb->ninodes);
  brelse(b);
}

// The same disk, as the block device registry sees it: whole, as "usb0", and
// with its partitions read and registered behind it. A stick written with
// tools/posix/mkdisk.sh -- or by any installer -- keeps SchFS in a partition
// next to its ESP, and the fixed USBDISK_DEV, which is the whole medium, can
// never see it there.
static int
ums_blk_rw(void *priv, uint64 lba, uint32 count, void *buf, int write)
{
  struct ums *u = priv;
  int n;
  if(count * SECTOR_SIZE > BSIZE || lba + count > u->nsectors)
    return -1;
  acquiresleep(&u->lock);
  if(u->gone){
    releasesleep(&u->lock);
    return -1;
  }
  if(write)
    memmove(u->xfer, buf, count * SECTOR_SIZE);
  n = scsi_rw(u, (uint32)lba, count, u->xfer, write);
  if(n == (int)(count * SECTOR_SIZE) && !write)
    memmove(buf, u->xfer, count * SECTOR_SIZE);
  releasesleep(&u->lock);
  return n == (int)(count * SECTOR_SIZE) ? 0 : -1;
}

static int
ums_blk_flush(void *priv)
{
  return ums_flush(priv);
}

static const struct blkdev_ops ums_blk_ops = {
  .rw    = ums_blk_rw,
  .flush = ums_blk_flush,
};

// ---------------------------------------------------------------------------

static int
storage_probe(struct usb_dev *dev)
{
  char vendor[9], product[17];
  struct ums *u = 0;
  int i;

  for(i = 0; i < NUMS; i++)
    if(!disks[i].present && disks[i].dev == 0){ u = &disks[i]; break; }
  if(u == 0){
    cprintf("usb-storage: %d sticks already; this one is not used\n", NUMS);
    return -1;
  }
  if(!dev->has_bulk_in || !dev->has_bulk_out){
    cprintf("usb-storage: needs a bulk endpoint in each direction\n");
    return -1;
  }

  initsleeplock(&u->lock, "usbms");
  u->dev = dev;
  u->gone = 0;
  u->blkdev = -1;
  // The name is the slot's, not a count: pull usb0 out and put it back and
  // it comes back as usb0.
  safestrcpy(u->name, "usb0", sizeof(u->name));
  u->name[3] = '0' + (u - disks);

  // One block, used for every transfer. bread() hands us b->data, which lives
  // in the buffer cache rather than in DMA memory, so everything bounces
  // through here. The wrappers are allocated once too: they are used on every
  // single I/O, and churning the DMA extent list twice per block read would
  // fragment an arena that has no business fragmenting.
  u->xfer = dma_alloc(BSIZE, 4096);
  u->cbw  = dma_alloc(sizeof(struct cbw), 64);
  u->csw  = dma_alloc(sizeof(struct csw), 64);
  if(u->xfer == 0 || u->cbw == 0 || u->csw == 0)
    goto fail;
  dev->drvdata = u;

  // Deliberately OUT first, which is the order that used to break.
  //
  // The slot context's Context Entries field must cover the highest DCI in
  // use, and xhci_open_ep originally wrote whichever endpoint it was adding.
  // Bulk OUT here has the higher DCI, so opening it first and then IN would
  // *lower* the field and silently invalidate the endpoint that was already
  // configured. Opening IN first hid that. Now that xhci_open_ep tracks the
  // maximum itself the order genuinely does not matter -- and doing it in the
  // formerly-broken order is what keeps that true.
  if(dev->ops->open_ep(dev, &dev->bulk_out) < 0 ||
     dev->ops->open_ep(dev, &dev->bulk_in) < 0)
    goto fail;

  // Deliberately no GET MAX LUN. It is answered with a STALL by
  // single-LUN devices, and recovering from a STALL on endpoint 0 needs a
  // Reset Endpoint command this driver does not implement yet. LUN 0 is the
  // only one used below, so asking would gain nothing and could wedge us.

  // A device that has just been configured is usually "not ready" for the
  // first command or two while it spins up or reads its own media.
  for(i = 0; i < 5; i++){
    uint8 cdb[6] = { SCSI_TEST_UNIT_READY, 0, 0, 0, 0, 0 };
    if(bot_command(u, cdb, sizeof(cdb), 0, 0, 0) >= 0)
      break;
    microdelay(100000);
  }

  if(scsi_inquiry(u, vendor, product) < 0)
    goto fail;
  if(scsi_read_capacity(u) < 0)
    goto fail;

  if(u->sector_size != SECTOR_SIZE){
    cprintf("usb-storage: %d-byte sectors are not supported\n",
            u->sector_size);
    goto fail;
  }
  if(u->nsectors == 0){
    cprintf("usb-storage: device reports no capacity\n");
    goto fail;
  }

  u->present = 1;
  kinfo("usb-storage: %s: %s %s, %d sectors (%d MB)\n", u->name,
          vendor, product, u->nsectors,
          (int)((uint64)u->nsectors * SECTOR_SIZE / (1024 * 1024)));
  u->blkdev = blkdev_register(u->name, &ums_blk_ops, u, u->nsectors,
                              SECTOR_SIZE);
  return 0;

fail:
  // Give the memory back and forget the device. Leaving u->dev set would let
  // a later block read reach a device we decided we could not drive.
  if(u->xfer) dma_free(u->xfer, BSIZE);
  if(u->cbw)  dma_free(u->cbw, sizeof(struct cbw));
  if(u->csw)  dma_free(u->csw, sizeof(struct csw));
  u->xfer = 0;
  u->cbw  = 0;
  u->csw  = 0;
  u->dev  = 0;
  dev->drvdata = 0;
  return -1;
}

// Pulled out. The controller has already failed whatever transfer was in
// flight, so the lock comes free; from then on every request fails at once,
// and the registry forgets the disk and its partitions. Mounted from, it
// fails the filesystem's next read -- which panics, by the rule in
// blkdev_rw(), because a filesystem that carried on would be guessing.
static void
storage_disconnect(struct usb_dev *dev)
{
  struct ums *u = dev->drvdata;

  if(u == 0)
    return;
  acquiresleep(&u->lock);
  u->gone = 1;
  releasesleep(&u->lock);

  kinfo("usb-storage: %s unplugged\n", u->name);
  if(u->blkdev >= 0)
    blkdev_unregister(u->blkdev);
  // USBDISK_DEV may have been this stick; nothing it cached is true now.
  binval(USBDISK_DEV);

  dma_free(u->xfer, BSIZE);
  dma_free(u->cbw, sizeof(struct cbw));
  dma_free(u->csw, sizeof(struct csw));
  u->xfer = 0;
  u->cbw = 0;
  u->csw = 0;
  u->present = 0;
  u->blkdev = -1;
  u->dev = 0;
  dev->drvdata = 0;
}

const struct usb_driver usb_storage_driver = {
  .name     = "usb-storage",
  .class    = USB_CLASS_MASS_STORAGE,
  .subclass = USB_MSC_SUBCLASS_SCSI,
  .protocol = USB_MSC_PROTO_BOT,
  .devclass = DEVCLASS_STORAGE,
  .probe    = storage_probe,
  .disconnect = storage_disconnect,
};