msc_disk.c

#include "tusb.h"
#include "grub_binary.h"
#include <string.h>

// ─── Virtual UEFI Boot Drive ───────────────────────────────────────────────────
//
// Layout: MBR + FAT16 partition
// UEFI finds /EFI/BOOT/BOOTX64.EFI (GRUB) and loads it.
// GRUB reads /boot/grub/grub.cfg (generated dynamically from GPIO15).
// GRUB chainloads to the selected OS bootloader.
//
// Absolute LBA layout:
//   LBA    0        : MBR (partition table)
//   LBA    1        : VBR (FAT16 boot sector, partition start)
//   LBA    2 -  18  : FAT copy 1 (17 sectors)
//   LBA   19 -  35  : FAT copy 2 (17 sectors)
//   LBA   36 -  67  : Root directory (32 sectors, 512 entries)
//   LBA   68        : Cluster  2 → /EFI/
//   LBA   69        : Cluster  3 → /EFI/BOOT/
//   LBA   70 - 1093 : Cluster  4…1027 → /EFI/BOOT/BOOTX64.EFI (512 KB)
//   LBA 1094        : Cluster 1028 → /boot/
//   LBA 1095        : Cluster 1029 → /boot/grub/
//   LBA 1096        : Cluster 1030 → /boot/grub/grub.cfg  ← DYNAMIC
//   LBA 1097+       : Free space (zeroes)

// ─── Disk Constants ────────────────────────────────────────────────────────────

#define DISK_BLOCK_SIZE     512u
#define DISK_BLOCK_NUM      4301u   // ~2.15 MB virtual disk

// Partition parameters
#define PART_START_LBA      1u      // Partition starts immediately after MBR
#define PART_SECTORS        4300u   // Sectors in the partition
#define FAT_COUNT           2u      // Number of FAT copies
#define FAT_SECTORS         17u     // Sectors per FAT
#define ROOT_ENTRIES        512u    // Maximum entries in the root directory
#define ROOT_SECTORS        32u     // (512 * 32 bytes) / 512 = 32 sectors

// Data area start (relative to partition start)
#define DATA_REL_START      (1u + FAT_COUNT * FAT_SECTORS + ROOT_SECTORS)  // = 67
// Data area start (absolute LBA)
#define DATA_ABS_START      (PART_START_LBA + DATA_REL_START)              // = 68

// Cluster → absolute LBA (cluster 2 is the first data cluster)
#define CLUSTER_LBA(c)      (DATA_ABS_START + (c) - 2u)

// ─── Cluster Assignments ───────────────────────────────────────────────────────

#define CLUSTER_EFI          2u
#define CLUSTER_EFI_BOOT     3u
#define CLUSTER_GRUB_FIRST   4u

// Number of clusters for the GRUB binary (rounded up to full sectors)
#define GRUB_CLUSTER_COUNT   ((GRUB_BINARY_SIZE + DISK_BLOCK_SIZE - 1u) / DISK_BLOCK_SIZE)
#define CLUSTER_GRUB_LAST    (CLUSTER_GRUB_FIRST + GRUB_CLUSTER_COUNT - 1u)

// Directories and config file follow the GRUB binary
#define CLUSTER_BOOT         (CLUSTER_GRUB_LAST + 1u)
#define CLUSTER_BOOT_GRUB    (CLUSTER_GRUB_LAST + 2u)
#define CLUSTER_CFG          (CLUSTER_GRUB_LAST + 3u)

// ─── GRUB Config Strings ───────────────────────────────────────────────────────

// GPIO15 HIGH (switch open, pull-up active) → Linux
static const char CFG_LINUX[] =
    "set timeout=0\n"
    "search --file --set=root /EFI/systemd/systemd-bootx64.efi\n"
    "chainloader /EFI/systemd/systemd-bootx64.efi\n"
    "boot\n";

// GPIO15 LOW (switch closed) → Windows
static const char CFG_WINDOWS[] =
    "set timeout=0\n"
    "search --file --set=root /EFI/Microsoft/Boot/bootmgfw.efi\n"
    "chainloader /EFI/Microsoft/Boot/bootmgfw.efi\n"
    "boot\n";

// Current switch state
static bool _linux_selected = true;

void msc_set_boot_target(bool linux_selected) {
    _linux_selected = linux_selected;
}

// ─── Helper Functions (ARM-safe, no pointer casting) ──────────────────────────

static void put_u16(uint8_t *p, uint16_t v) {
    p[0] = (uint8_t)(v & 0xFF);
    p[1] = (uint8_t)(v >> 8);
}

static void put_u32(uint8_t *p, uint32_t v) {
    p[0] = (uint8_t)(v & 0xFF);
    p[1] = (uint8_t)((v >> 8)  & 0xFF);
    p[2] = (uint8_t)((v >> 16) & 0xFF);
    p[3] = (uint8_t)((v >> 24) & 0xFF);
}

// Write a FAT directory entry (32 bytes)
// name11:  11 bytes in FAT 8.3 format (8 char name + 3 char extension, uppercase, space-padded)
// attr:    0x10 = directory, 0x20 = file (archive), 0x08 = volume label
// cluster: first cluster of the file or directory
// size:    file size in bytes (0 for directories)
static void write_dir_entry(uint8_t *e, const char name11[11],
                            uint8_t attr, uint16_t cluster, uint32_t size) {
    memset(e, 0, 32);
    memcpy(e, name11, 11);             // Offset  0: name (11 bytes)
    e[11] = attr;                      // Offset 11: attributes
    // e[12-17]: timestamps (0)
    put_u16(e + 20, 0);                // Offset 20: first cluster high (always 0 for FAT16)
    put_u16(e + 22, 0x0000);           // Offset 22: write time
    put_u16(e + 24, 0x5A21);           // Offset 24: write date (2025-01-01)
    put_u16(e + 26, cluster);          // Offset 26: first cluster low
    put_u32(e + 28, size);             // Offset 28: file size
}

// ─── Sector Generators ─────────────────────────────────────────────────────────

// LBA 0: MBR with a single FAT16 partition
static void sector_mbr(uint8_t *buf) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    // Partition entry (offset 446)
    buf[446] = 0x80;                            // bootable
    buf[447] = 0x20;                            // CHS start (simplified)
    buf[448] = 0x21;
    buf[449] = 0x00;
    buf[450] = 0x0E;                            // type: FAT16 with LBA
    buf[451] = 0xFE;                            // CHS end (simplified)
    buf[452] = 0xFF;
    buf[453] = 0xFF;
    put_u32(buf + 454, PART_START_LBA);         // LBA start
    put_u32(buf + 458, PART_SECTORS);           // LBA size
    buf[510] = 0x55;                            // boot signature
    buf[511] = 0xAA;
}

// LBA 1: VBR – FAT16 boot sector (BPB)
static void sector_vbr(uint8_t *buf) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    buf[0] = 0xEB; buf[1] = 0x58; buf[2] = 0x90;           // JMP + NOP
    memcpy(buf + 3, "MSWIN4.1", 8);                         // OEM name
    put_u16(buf + 11, DISK_BLOCK_SIZE);                     // bytes per sector
    buf[13] = 1;                                            // sectors per cluster
    put_u16(buf + 14, 1);                                   // reserved sectors
    buf[16] = FAT_COUNT;                                    // number of FATs
    put_u16(buf + 17, ROOT_ENTRIES);                        // root directory entries
    put_u16(buf + 19, 0);                                   // total sectors 16-bit (0 = use 32-bit)
    buf[21] = 0xF8;                                         // media type: fixed disk
    put_u16(buf + 22, FAT_SECTORS);                         // sectors per FAT
    put_u16(buf + 24, 63);                                  // sectors per track
    put_u16(buf + 26, 255);                                 // number of heads
    put_u32(buf + 28, PART_START_LBA);                      // hidden sectors
    put_u32(buf + 32, PART_SECTORS);                        // total sectors 32-bit
    buf[36] = 0x80;                                         // drive number
    buf[38] = 0x29;                                         // extended boot signature
    put_u32(buf + 39, 0xB007B007u);                         // volume serial number
    memcpy(buf + 43, "BOOTSWITCH ", 11);                    // volume label
    memcpy(buf + 54, "FAT16   ", 8);                        // filesystem type
    buf[510] = 0x55;
    buf[511] = 0xAA;
}

// LBA 2-18 / 19-35: FAT table (both copies are identical)
// fat_idx: sector index within the FAT (0..FAT_SECTORS-1)
static void sector_fat(uint8_t *buf, uint32_t fat_idx) {
    memset(buf, 0, DISK_BLOCK_SIZE);

    // Each FAT sector holds 256 entries of 2 bytes each
    uint32_t first_cluster = fat_idx * 256u;

    for (uint32_t i = 0; i < 256u; i++) {
        uint32_t cluster = first_cluster + i;
        uint16_t value;

        if (cluster == 0u) {
            value = 0xFFF8u;  // media descriptor
        } else if (cluster == 1u) {
            value = 0xFFFFu;  // reserved
        } else if (cluster == CLUSTER_EFI) {
            value = 0xFFFFu;  // /EFI/ – single cluster, end of chain
        } else if (cluster == CLUSTER_EFI_BOOT) {
            value = 0xFFFFu;  // /EFI/BOOT/ – single cluster, end of chain
        } else if (cluster >= CLUSTER_GRUB_FIRST && cluster < CLUSTER_GRUB_LAST) {
            value = (uint16_t)(cluster + 1u);  // GRUB chain: points to next cluster
        } else if (cluster == CLUSTER_GRUB_LAST) {
            value = 0xFFFFu;  // end of GRUB binary chain
        } else if (cluster == CLUSTER_BOOT) {
            value = 0xFFFFu;  // /boot/ – single cluster, end of chain
        } else if (cluster == CLUSTER_BOOT_GRUB) {
            value = 0xFFFFu;  // /boot/grub/ – single cluster, end of chain
        } else if (cluster == CLUSTER_CFG) {
            value = 0xFFFFu;  // grub.cfg – single cluster, end of chain
        } else {
            value = 0x0000u;  // free cluster
        }

        buf[i * 2u]     = (uint8_t)(value & 0xFF);
        buf[i * 2u + 1u] = (uint8_t)(value >> 8);
    }
}

// LBA 36-67: root directory
// sector_offset: 0..ROOT_SECTORS-1 (only sector 0 contains entries)
static void sector_root_dir(uint8_t *buf, uint32_t sector_offset) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    if (sector_offset != 0u) return;

    // [0] volume label
    uint8_t *e = buf;
    memcpy(e, "BOOTSWITCH ", 11);
    e[11] = 0x08;

    // [1] /EFI/ directory
    write_dir_entry(buf + 32,  "EFI        ", 0x10, (uint16_t)CLUSTER_EFI,  0);

    // [2] /boot/ directory (contains /boot/grub/grub.cfg)
    write_dir_entry(buf + 64,  "BOOT       ", 0x10, (uint16_t)CLUSTER_BOOT, 0);
}

// Cluster 2: /EFI/ directory
static void sector_efi_dir(uint8_t *buf) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    write_dir_entry(buf,      ".          ", 0x10, (uint16_t)CLUSTER_EFI,      0);
    write_dir_entry(buf + 32, "..         ", 0x10, 0,                          0);
    write_dir_entry(buf + 64, "BOOT       ", 0x10, (uint16_t)CLUSTER_EFI_BOOT, 0);
}

// Cluster 3: /EFI/BOOT/ directory
static void sector_efi_boot_dir(uint8_t *buf) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    write_dir_entry(buf,      ".          ", 0x10, (uint16_t)CLUSTER_EFI_BOOT,  0);
    write_dir_entry(buf + 32, "..         ", 0x10, (uint16_t)CLUSTER_EFI,       0);
    write_dir_entry(buf + 64, "BOOTX64 EFI", 0x20, (uint16_t)CLUSTER_GRUB_FIRST, GRUB_BINARY_SIZE);
}

// Cluster CLUSTER_BOOT: /boot/ directory
static void sector_boot_dir(uint8_t *buf) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    write_dir_entry(buf,      ".          ", 0x10, (uint16_t)CLUSTER_BOOT,      0);
    write_dir_entry(buf + 32, "..         ", 0x10, 0,                           0);
    write_dir_entry(buf + 64, "GRUB       ", 0x10, (uint16_t)CLUSTER_BOOT_GRUB, 0);
}

// Cluster CLUSTER_BOOT_GRUB: /boot/grub/ directory
static void sector_boot_grub_dir(uint8_t *buf) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    const char *cfg = _linux_selected ? CFG_LINUX : CFG_WINDOWS;
    uint32_t cfg_size = (uint32_t)strlen(cfg);

    write_dir_entry(buf,      ".          ", 0x10, (uint16_t)CLUSTER_BOOT_GRUB, 0);         // .
    write_dir_entry(buf + 32, "..         ", 0x10, (uint16_t)CLUSTER_BOOT,      0);         // ..
    write_dir_entry(buf + 64, "GRUB    CFG", 0x20, (uint16_t)CLUSTER_CFG,       cfg_size);  // grub.cfg
}

// Cluster CLUSTER_CFG: /boot/grub/grub.cfg content (dynamic)
static void sector_grub_cfg(uint8_t *buf) {
    memset(buf, 0, DISK_BLOCK_SIZE);
    const char *cfg = _linux_selected ? CFG_LINUX : CFG_WINDOWS;
    uint32_t len = (uint32_t)strlen(cfg);
    if (len > DISK_BLOCK_SIZE) len = DISK_BLOCK_SIZE;
    memcpy(buf, cfg, len);
}

// ─── TinyUSB MSC Callbacks ─────────────────────────────────────────────────────

void tud_msc_inquiry_cb(uint8_t lun, uint8_t vendor_id[8],
                         uint8_t product_id[16], uint8_t product_rev[4]) {
    (void)lun;
    memcpy(vendor_id,   "Pico    ", 8);
    memcpy(product_id,  "BootSwitch      ", 16);
    memcpy(product_rev, "1.0 ", 4);
}

bool tud_msc_test_unit_ready_cb(uint8_t lun) {
    (void)lun;
    return true;
}

void tud_msc_capacity_cb(uint8_t lun, uint32_t *block_count, uint16_t *block_size) {
    (void)lun;
    *block_count = DISK_BLOCK_NUM;
    *block_size  = DISK_BLOCK_SIZE;
}

bool tud_msc_start_stop_cb(uint8_t lun, uint8_t power_condition,
                            bool start, bool load_eject) {
    (void)lun; (void)power_condition; (void)start; (void)load_eject;
    return true;
}

// Read callback: returns the requested virtual sector
int32_t tud_msc_read10_cb(uint8_t lun, uint32_t lba, uint32_t offset,
                            void *buffer, uint32_t bufsize) {
    (void)lun; (void)offset;
    uint8_t *buf = (uint8_t *)buffer;
    memset(buf, 0, bufsize);

    if (lba == 0u) {
        sector_mbr(buf);

    } else if (lba == 1u) {
        sector_vbr(buf);

    } else if (lba >= 2u && lba <= (1u + FAT_SECTORS)) {
        // FAT copy 1: LBA 2..18 → fat_idx 0..16
        sector_fat(buf, lba - 2u);

    } else if (lba >= (2u + FAT_SECTORS) && lba <= (1u + 2u * FAT_SECTORS)) {
        // FAT copy 2: LBA 19..35 → fat_idx 0..16 (identical)
        sector_fat(buf, lba - (2u + FAT_SECTORS));

    } else if (lba >= (2u + 2u * FAT_SECTORS) &&
               lba <  (2u + 2u * FAT_SECTORS + ROOT_SECTORS)) {
        // Root directory: LBA 36..67
        sector_root_dir(buf, lba - (2u + 2u * FAT_SECTORS));

    } else if (lba == CLUSTER_LBA(CLUSTER_EFI)) {
        sector_efi_dir(buf);

    } else if (lba == CLUSTER_LBA(CLUSTER_EFI_BOOT)) {
        sector_efi_boot_dir(buf);

    } else if (lba >= CLUSTER_LBA(CLUSTER_GRUB_FIRST) &&
               lba <= CLUSTER_LBA(CLUSTER_GRUB_LAST)) {
        // Serve GRUB binary from flash
        uint32_t byte_offset = (lba - CLUSTER_LBA(CLUSTER_GRUB_FIRST)) * DISK_BLOCK_SIZE;
        uint32_t remaining   = GRUB_BINARY_SIZE - byte_offset;
        uint32_t copy_size   = (remaining < DISK_BLOCK_SIZE) ? remaining : DISK_BLOCK_SIZE;
        memcpy(buf, grub_binary + byte_offset, copy_size);

    } else if (lba == CLUSTER_LBA(CLUSTER_BOOT)) {
        sector_boot_dir(buf);

    } else if (lba == CLUSTER_LBA(CLUSTER_BOOT_GRUB)) {
        sector_boot_grub_dir(buf);

    } else if (lba == CLUSTER_LBA(CLUSTER_CFG)) {
        sector_grub_cfg(buf);
    }
    // All other LBAs: zeroes (already set by memset above)

    return (int32_t)bufsize;
}

// Write callback: virtual drive is read-only, silently discard all writes
int32_t tud_msc_write10_cb(uint8_t lun, uint32_t lba, uint32_t offset,
                             uint8_t *buffer, uint32_t bufsize) {
    (void)lun; (void)lba; (void)offset; (void)buffer;
    return (int32_t)bufsize;
}

// Reject unknown SCSI commands
int32_t tud_msc_scsi_cb(uint8_t lun, uint8_t const scsi_cmd[16],
                          void *buffer, uint16_t bufsize) {
    (void)lun; (void)scsi_cmd; (void)buffer; (void)bufsize;
    tud_msc_set_sense(lun, SCSI_SENSE_ILLEGAL_REQUEST, 0x20, 0x00);
    return -1;
}