luajitos

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs

ata.c (21265B)


      1 /* ATA PIO Driver for LuajitOS
      2  * Provides basic ATA/IDE disk access using PIO mode
      3  */
      4 
      5 #include "ata.h"
      6 #include <string.h>
      7 #include <stdlib.h>
      8 #include <stdio.h>
      9 
     10 /* External terminal function for debug output */
     11 extern void terminal_writestring(const char* str);
     12 
     13 /* Detected drives */
     14 static uint32_t drives_canary_before = 0xDEADBEEF;
     15 static ata_drive_info_t drives[4];  /* Primary master/slave, Secondary master/slave */
     16 static uint32_t drives_canary_after = 0xCAFEBABE;
     17 static int num_drives = 0;
     18 
     19 /* I/O port access functions */
     20 static inline void outb(uint16_t port, uint8_t val) {
     21     __asm__ volatile ("outb %0, %1" : : "a"(val), "Nd"(port));
     22 }
     23 
     24 static inline uint8_t inb(uint16_t port) {
     25     uint8_t ret;
     26     __asm__ volatile ("inb %1, %0" : "=a"(ret) : "Nd"(port));
     27     return ret;
     28 }
     29 
     30 static inline void outw(uint16_t port, uint16_t val) {
     31     __asm__ volatile ("outw %0, %1" : : "a"(val), "Nd"(port));
     32 }
     33 
     34 static inline uint16_t inw(uint16_t port) {
     35     uint16_t ret;
     36     __asm__ volatile ("inw %1, %0" : "=a"(ret) : "Nd"(port));
     37     return ret;
     38 }
     39 
     40 /* Small delay (400ns) by reading alternate status register 4 times */
     41 static void ata_delay(uint16_t ctrl_port) {
     42     inb(ctrl_port);
     43     inb(ctrl_port);
     44     inb(ctrl_port);
     45     inb(ctrl_port);
     46 }
     47 
     48 /* Wait for BSY flag to clear */
     49 static int ata_wait_bsy(uint16_t status_port) {
     50     int timeout = ATA_TIMEOUT;
     51     while ((inb(status_port) & ATA_SR_BSY) && timeout > 0) {
     52         timeout--;
     53     }
     54     return (timeout > 0) ? ATA_OK : ATA_ERR_TIMEOUT;
     55 }
     56 
     57 /* Wait for DRQ flag to set (data ready) */
     58 static int ata_wait_drq(uint16_t status_port) {
     59     int timeout = ATA_TIMEOUT;
     60     uint8_t status;
     61 
     62     while (timeout > 0) {
     63         status = inb(status_port);
     64 
     65         if (status & ATA_SR_ERR) {
     66             return ATA_ERR_IO;
     67         }
     68         if (status & ATA_SR_DF) {
     69             return ATA_ERR_IO;
     70         }
     71         if (status & ATA_SR_DRQ) {
     72             return ATA_OK;
     73         }
     74 
     75         timeout--;
     76     }
     77 
     78     return ATA_ERR_TIMEOUT;
     79 }
     80 
     81 /* Wait for drive to be ready */
     82 static int ata_wait_ready(uint16_t status_port) {
     83     int timeout = ATA_TIMEOUT;
     84     uint8_t status;
     85 
     86     while (timeout > 0) {
     87         status = inb(status_port);
     88 
     89         if (status & ATA_SR_ERR) {
     90             return ATA_ERR_IO;
     91         }
     92         if (!(status & ATA_SR_BSY) && (status & ATA_SR_DRDY)) {
     93             return ATA_OK;
     94         }
     95 
     96         timeout--;
     97     }
     98 
     99     return ATA_ERR_TIMEOUT;
    100 }
    101 
    102 /* Get base I/O port for bus */
    103 static uint16_t ata_get_base_port(uint8_t bus) {
    104     return (bus == 0) ? ATA_PRIMARY_DATA : ATA_SECONDARY_DATA;
    105 }
    106 
    107 /* Get control port for bus */
    108 static uint16_t ata_get_ctrl_port(uint8_t bus) {
    109     return (bus == 0) ? ATA_PRIMARY_CTRL : ATA_SECONDARY_CTRL;
    110 }
    111 
    112 /* Select drive */
    113 static void ata_select_drive(uint16_t base_port, uint8_t drive, uint32_t lba) {
    114     /* Drive/Head register format for LBA mode:
    115      * Bit 7: 1 (always set)
    116      * Bit 6: 1 for LBA mode, 0 for CHS
    117      * Bit 5: 1 (always set)
    118      * Bit 4: 0 for master, 1 for slave
    119      * Bits 3-0: LBA bits 27-24
    120      */
    121     uint8_t drive_byte = 0xE0 | ((drive & 1) << 4) | ((lba >> 24) & 0x0F);
    122     outb(base_port + 6, drive_byte);  /* Drive/Head register is at base + 6 */
    123 }
    124 
    125 /* Identify a drive */
    126 static int ata_identify_drive(uint8_t bus, uint8_t drive, ata_drive_info_t* info) {
    127     uint16_t base_port = ata_get_base_port(bus);
    128     uint16_t ctrl_port = ata_get_ctrl_port(bus);
    129     uint16_t identify_data[256];
    130     int i;
    131 
    132     memset(info, 0, sizeof(ata_drive_info_t));
    133     info->bus = bus;
    134     info->drive = drive;
    135 
    136     /* Select drive (no LBA bits needed for IDENTIFY) */
    137     outb(base_port + 6, 0xA0 | ((drive & 1) << 4));
    138     ata_delay(ctrl_port);
    139 
    140     /* Clear sector count, LBA registers */
    141     outb(base_port + 2, 0);  /* Sector count */
    142     outb(base_port + 3, 0);  /* LBA lo */
    143     outb(base_port + 4, 0);  /* LBA mid */
    144     outb(base_port + 5, 0);  /* LBA hi */
    145 
    146     /* Send IDENTIFY command */
    147     outb(base_port + 7, ATA_CMD_IDENTIFY);
    148     ata_delay(ctrl_port);
    149 
    150     /* Check if drive exists */
    151     uint8_t status = inb(base_port + 7);
    152     if (status == 0) {
    153         /* No drive on this bus/slot */
    154         return ATA_ERR_NO_DRIVE;
    155     }
    156 
    157     /* Wait for BSY to clear */
    158     if (ata_wait_bsy(base_port + 7) != ATA_OK) {
    159         return ATA_ERR_TIMEOUT;
    160     }
    161 
    162     /* Check for ATAPI or other non-ATA device */
    163     uint8_t lba_mid = inb(base_port + 4);
    164     uint8_t lba_hi = inb(base_port + 5);
    165     if (lba_mid != 0 || lba_hi != 0) {
    166         /* Not an ATA device (might be ATAPI) */
    167         info->is_ata = 0;
    168         return ATA_ERR_NO_DRIVE;  /* We only support ATA */
    169     }
    170 
    171     /* Wait for DRQ or ERR */
    172     int result = ata_wait_drq(base_port + 7);
    173     if (result != ATA_OK) {
    174         return result;
    175     }
    176 
    177     /* Read 256 words of identification data */
    178     for (i = 0; i < 256; i++) {
    179         identify_data[i] = inw(base_port);
    180     }
    181 
    182     /* Parse identification data */
    183     info->present = 1;
    184     info->is_ata = 1;
    185 
    186     /* Word 60-61: Total sectors in 28-bit LBA */
    187     info->sectors = identify_data[60] | ((uint32_t)identify_data[61] << 16);
    188 
    189     {
    190         char dbg[100];
    191         snprintf(dbg, sizeof(dbg), "[ATA] identify: w60=%d w61=%d sectors=%d\n",
    192                  (int)identify_data[60], (int)identify_data[61], (int)info->sectors);
    193         terminal_writestring(dbg);
    194     }
    195 
    196     /* Word 83: Command set supported (bit 10 = 48-bit LBA) */
    197     info->lba48_supported = (identify_data[83] & (1 << 10)) ? 1 : 0;
    198 
    199     /* Words 100-103: Total sectors in 48-bit LBA */
    200     if (info->lba48_supported) {
    201         info->sectors_48 = (uint64_t)identify_data[100] |
    202                           ((uint64_t)identify_data[101] << 16) |
    203                           ((uint64_t)identify_data[102] << 32) |
    204                           ((uint64_t)identify_data[103] << 48);
    205     } else {
    206         info->sectors_48 = info->sectors;
    207     }
    208 
    209     /* Words 27-46: Model string (40 chars, swapped byte order) */
    210     for (i = 0; i < 20; i++) {
    211         info->model[i*2] = (identify_data[27 + i] >> 8) & 0xFF;
    212         info->model[i*2 + 1] = identify_data[27 + i] & 0xFF;
    213     }
    214     info->model[40] = '\0';
    215 
    216     /* Trim trailing spaces */
    217     for (i = 39; i >= 0 && info->model[i] == ' '; i--) {
    218         info->model[i] = '\0';
    219     }
    220 
    221     /* Words 10-19: Serial number (20 chars, swapped byte order) */
    222     for (i = 0; i < 10; i++) {
    223         info->serial[i*2] = (identify_data[10 + i] >> 8) & 0xFF;
    224         info->serial[i*2 + 1] = identify_data[10 + i] & 0xFF;
    225     }
    226     info->serial[20] = '\0';
    227 
    228     /* Trim trailing spaces */
    229     for (i = 19; i >= 0 && info->serial[i] == ' '; i--) {
    230         info->serial[i] = '\0';
    231     }
    232 
    233     /* Words 23-26: Firmware revision (8 chars, swapped byte order) */
    234     for (i = 0; i < 4; i++) {
    235         info->firmware[i*2] = (identify_data[23 + i] >> 8) & 0xFF;
    236         info->firmware[i*2 + 1] = identify_data[23 + i] & 0xFF;
    237     }
    238     info->firmware[8] = '\0';
    239 
    240     /* Trim trailing spaces */
    241     for (i = 7; i >= 0 && info->firmware[i] == ' '; i--) {
    242         info->firmware[i] = '\0';
    243     }
    244 
    245     return ATA_OK;
    246 }
    247 
    248 /* Initialize ATA driver */
    249 int ata_init(void) {
    250     int i;
    251 
    252     terminal_writestring("[ATA] Initializing ATA/IDE driver...\n");
    253 
    254     {
    255         char dbg[80];
    256         /* Cast sizeof to int to avoid snprintf issues */
    257         int sz = (int)sizeof(drives);
    258         snprintf(dbg, sizeof(dbg), "[ATA] drives array size=%d bytes\n", sz);
    259         terminal_writestring(dbg);
    260     }
    261 
    262     /* Reset both buses */
    263     outb(ATA_PRIMARY_CTRL, 0x04);   /* Set SRST bit */
    264     ata_delay(ATA_PRIMARY_CTRL);
    265     outb(ATA_PRIMARY_CTRL, 0x00);   /* Clear SRST bit */
    266     ata_delay(ATA_PRIMARY_CTRL);
    267 
    268     outb(ATA_SECONDARY_CTRL, 0x04);
    269     ata_delay(ATA_SECONDARY_CTRL);
    270     outb(ATA_SECONDARY_CTRL, 0x00);
    271     ata_delay(ATA_SECONDARY_CTRL);
    272 
    273     /* Clear drive info */
    274     memset(drives, 0, sizeof(drives));
    275     num_drives = 0;
    276 
    277     /* Detect drives on both buses */
    278     for (i = 0; i < 4; i++) {
    279         uint8_t bus = i / 2;
    280         uint8_t drive = i % 2;
    281 
    282         int result = ata_identify_drive(bus, drive, &drives[i]);
    283 
    284         if (result == ATA_OK && drives[i].present) {
    285             num_drives++;
    286 
    287             terminal_writestring("[ATA] Found drive: ");
    288             terminal_writestring(bus == 0 ? "Primary " : "Secondary ");
    289             terminal_writestring(drive == 0 ? "Master\n" : "Slave\n");
    290             terminal_writestring("      Model: ");
    291             terminal_writestring(drives[i].model);
    292             terminal_writestring("\n");
    293 
    294             /* Debug: print drive array address and sectors */
    295             {
    296                 char dbg[100];
    297                 snprintf(dbg, sizeof(dbg), "      [DEBUG] drives[%d] sectors=%d\n",
    298                          i, (int)drives[i].sectors);
    299                 terminal_writestring(dbg);
    300             }
    301 
    302             /* Print capacity */
    303             uint64_t bytes = (uint64_t)drives[i].sectors * ATA_SECTOR_SIZE;
    304             uint32_t mb = (uint32_t)(bytes / (1024 * 1024));
    305 
    306             char size_str[32];
    307             int idx = 0;
    308             if (mb == 0) {
    309                 size_str[idx++] = '0';
    310             } else {
    311                 char temp[16];
    312                 int j = 0;
    313                 while (mb > 0) {
    314                     temp[j++] = '0' + (mb % 10);
    315                     mb /= 10;
    316                 }
    317                 while (j > 0) {
    318                     size_str[idx++] = temp[--j];
    319                 }
    320             }
    321             size_str[idx] = '\0';
    322 
    323             terminal_writestring("      Size: ");
    324             terminal_writestring(size_str);
    325             terminal_writestring(" MB\n");
    326         }
    327     }
    328 
    329     if (num_drives == 0) {
    330         terminal_writestring("[ATA] No drives detected\n");
    331     } else {
    332         terminal_writestring("[ATA] Initialization complete\n");
    333     }
    334 
    335     return num_drives;
    336 }
    337 
    338 /* Get drive information */
    339 ata_drive_info_t* ata_get_drive_info(uint8_t bus, uint8_t drive) {
    340     if (bus > 1 || drive > 1) {
    341         return NULL;
    342     }
    343 
    344     int idx = bus * 2 + drive;
    345     if (!drives[idx].present) {
    346         return NULL;
    347     }
    348 
    349     /* Debug: verify the returned info is valid */
    350     ata_drive_info_t *info = &drives[idx];
    351 
    352     /* Check canaries for memory corruption */
    353     if (drives_canary_before != 0xDEADBEEF || drives_canary_after != 0xCAFEBABE) {
    354         char dbg[100];
    355         snprintf(dbg, sizeof(dbg), "[ATA] CORRUPTION: canaries=0x%x,0x%x\n",
    356                  (int)drives_canary_before, (int)drives_canary_after);
    357         terminal_writestring(dbg);
    358     }
    359 
    360     if (info->sectors == 0 || info->sectors > 0x10000000) {
    361         char dbg[100];
    362         snprintf(dbg, sizeof(dbg), "[ATA] WARNING: suspicious sectors=%d for idx=%d\n",
    363                  (int)info->sectors, idx);
    364         terminal_writestring(dbg);
    365     }
    366 
    367     return info;
    368 }
    369 
    370 /* Read sectors from disk */
    371 int ata_read_sectors(uint8_t bus, uint8_t drive, uint32_t lba, uint8_t count, void* buffer) {
    372     if (bus > 1 || drive > 1) {
    373         return ATA_ERR_INVALID;
    374     }
    375 
    376     int idx = bus * 2 + drive;
    377     if (!drives[idx].present) {
    378         return ATA_ERR_NO_DRIVE;
    379     }
    380 
    381     /* For 28-bit LBA, maximum sector is 0x0FFFFFFF */
    382     if (lba > 0x0FFFFFFF) {
    383         return ATA_ERR_INVALID;
    384     }
    385 
    386     uint16_t base_port = ata_get_base_port(bus);
    387     uint16_t ctrl_port = ata_get_ctrl_port(bus);
    388     uint16_t* buf = (uint16_t*)buffer;
    389     int sectors_to_read = (count == 0) ? 256 : count;
    390     int i, j;
    391 
    392     /* Wait for drive to be ready */
    393     if (ata_wait_bsy(base_port + 7) != ATA_OK) {
    394         return ATA_ERR_TIMEOUT;
    395     }
    396 
    397     /* Select drive and set LBA high bits */
    398     ata_select_drive(base_port, drive, lba);
    399     ata_delay(ctrl_port);
    400 
    401     /* Wait for drive selection to take effect */
    402     if (ata_wait_bsy(base_port + 7) != ATA_OK) {
    403         return ATA_ERR_TIMEOUT;
    404     }
    405 
    406     /* Set up the transfer */
    407     outb(base_port + 2, count);              /* Sector count (0 = 256) */
    408     outb(base_port + 3, lba & 0xFF);         /* LBA low */
    409     outb(base_port + 4, (lba >> 8) & 0xFF);  /* LBA mid */
    410     outb(base_port + 5, (lba >> 16) & 0xFF); /* LBA high */
    411 
    412     /* Send READ command */
    413     outb(base_port + 7, ATA_CMD_READ_PIO);
    414 
    415     /* Read each sector */
    416     for (i = 0; i < sectors_to_read; i++) {
    417         /* Wait for data to be ready */
    418         int result = ata_wait_drq(base_port + 7);
    419         if (result != ATA_OK) {
    420             return result;
    421         }
    422 
    423         /* Read 256 words (512 bytes) */
    424         for (j = 0; j < 256; j++) {
    425             buf[i * 256 + j] = inw(base_port);
    426         }
    427 
    428         /* Small delay between sectors */
    429         ata_delay(ctrl_port);
    430     }
    431 
    432     return ATA_OK;
    433 }
    434 
    435 /* Write sectors to disk */
    436 int ata_write_sectors(uint8_t bus, uint8_t drive, uint32_t lba, uint8_t count, const void* buffer) {
    437     if (bus > 1 || drive > 1) {
    438         return ATA_ERR_INVALID;
    439     }
    440 
    441     int idx = bus * 2 + drive;
    442     if (!drives[idx].present) {
    443         return ATA_ERR_NO_DRIVE;
    444     }
    445 
    446     /* For 28-bit LBA, maximum sector is 0x0FFFFFFF */
    447     if (lba > 0x0FFFFFFF) {
    448         return ATA_ERR_INVALID;
    449     }
    450 
    451     uint16_t base_port = ata_get_base_port(bus);
    452     uint16_t ctrl_port = ata_get_ctrl_port(bus);
    453     const uint16_t* buf = (const uint16_t*)buffer;
    454     int sectors_to_write = (count == 0) ? 256 : count;
    455     int i, j;
    456 
    457     /* Wait for drive to be ready */
    458     if (ata_wait_bsy(base_port + 7) != ATA_OK) {
    459         return ATA_ERR_TIMEOUT;
    460     }
    461 
    462     /* Select drive and set LBA high bits */
    463     ata_select_drive(base_port, drive, lba);
    464     ata_delay(ctrl_port);
    465 
    466     /* Wait for drive selection to take effect */
    467     if (ata_wait_bsy(base_port + 7) != ATA_OK) {
    468         return ATA_ERR_TIMEOUT;
    469     }
    470 
    471     /* Set up the transfer */
    472     outb(base_port + 2, count);              /* Sector count (0 = 256) */
    473     outb(base_port + 3, lba & 0xFF);         /* LBA low */
    474     outb(base_port + 4, (lba >> 8) & 0xFF);  /* LBA mid */
    475     outb(base_port + 5, (lba >> 16) & 0xFF); /* LBA high */
    476 
    477     /* Send WRITE command */
    478     outb(base_port + 7, ATA_CMD_WRITE_PIO);
    479 
    480     /* Write each sector */
    481     for (i = 0; i < sectors_to_write; i++) {
    482         /* Wait for drive to be ready for data */
    483         int result = ata_wait_drq(base_port + 7);
    484         if (result != ATA_OK) {
    485             return result;
    486         }
    487 
    488         /* Write 256 words (512 bytes) */
    489         for (j = 0; j < 256; j++) {
    490             outw(base_port, buf[i * 256 + j]);
    491         }
    492 
    493         /* Small delay between sectors */
    494         ata_delay(ctrl_port);
    495     }
    496 
    497     /* Flush cache after write */
    498     return ata_flush(bus, drive);
    499 }
    500 
    501 /* Flush write cache */
    502 int ata_flush(uint8_t bus, uint8_t drive) {
    503     if (bus > 1 || drive > 1) {
    504         return ATA_ERR_INVALID;
    505     }
    506 
    507     int idx = bus * 2 + drive;
    508     if (!drives[idx].present) {
    509         return ATA_ERR_NO_DRIVE;
    510     }
    511 
    512     uint16_t base_port = ata_get_base_port(bus);
    513     uint16_t ctrl_port = ata_get_ctrl_port(bus);
    514 
    515     /* Select drive */
    516     outb(base_port + 6, 0xE0 | ((drive & 1) << 4));
    517     ata_delay(ctrl_port);
    518 
    519     /* Send CACHE FLUSH command */
    520     outb(base_port + 7, ATA_CMD_CACHE_FLUSH);
    521 
    522     /* Wait for completion */
    523     return ata_wait_bsy(base_port + 7);
    524 }
    525 
    526 /* Get error string */
    527 const char* ata_get_error_string(int error) {
    528     switch (error) {
    529         case ATA_OK:          return "Success";
    530         case ATA_ERR_NO_DRIVE: return "No drive present";
    531         case ATA_ERR_TIMEOUT:  return "Operation timed out";
    532         case ATA_ERR_IO:       return "I/O error";
    533         case ATA_ERR_INVALID:  return "Invalid parameter";
    534         default:              return "Unknown error";
    535     }
    536 }
    537 
    538 /* ========== Lua Bindings ========== */
    539 
    540 #include "lua.h"
    541 #include "lauxlib.h"
    542 
    543 /* ata.init() -> num_drives */
    544 static int lua_ata_init(lua_State* L) {
    545     int num = ata_init();
    546     lua_pushinteger(L, num);
    547     return 1;
    548 }
    549 
    550 /* ata.getDriveInfo(bus, drive) -> table or nil */
    551 static int lua_ata_get_drive_info(lua_State* L) {
    552     int bus = luaL_checkinteger(L, 1);
    553     int drive = luaL_checkinteger(L, 2);
    554 
    555     ata_drive_info_t* info = ata_get_drive_info(bus, drive);
    556     if (!info) {
    557         lua_pushnil(L);
    558         return 1;
    559     }
    560 
    561     lua_newtable(L);
    562 
    563     lua_pushstring(L, "present");
    564     lua_pushboolean(L, info->present);
    565     lua_settable(L, -3);
    566 
    567     lua_pushstring(L, "bus");
    568     lua_pushinteger(L, info->bus);
    569     lua_settable(L, -3);
    570 
    571     lua_pushstring(L, "drive");
    572     lua_pushinteger(L, info->drive);
    573     lua_settable(L, -3);
    574 
    575     lua_pushstring(L, "model");
    576     lua_pushstring(L, info->model);
    577     lua_settable(L, -3);
    578 
    579     lua_pushstring(L, "serial");
    580     lua_pushstring(L, info->serial);
    581     lua_settable(L, -3);
    582 
    583     lua_pushstring(L, "firmware");
    584     lua_pushstring(L, info->firmware);
    585     lua_settable(L, -3);
    586 
    587     lua_pushstring(L, "sectors");
    588     lua_pushinteger(L, info->sectors);
    589     lua_settable(L, -3);
    590 
    591     lua_pushstring(L, "sectorSize");
    592     lua_pushinteger(L, ATA_SECTOR_SIZE);
    593     lua_settable(L, -3);
    594 
    595     lua_pushstring(L, "lba48");
    596     lua_pushboolean(L, info->lba48_supported);
    597     lua_settable(L, -3);
    598 
    599     /* Calculate size in bytes (as string for large disks) */
    600     uint64_t bytes = (uint64_t)info->sectors * ATA_SECTOR_SIZE;
    601     lua_pushstring(L, "bytes");
    602     lua_pushnumber(L, (lua_Number)bytes);
    603     lua_settable(L, -3);
    604 
    605     return 1;
    606 }
    607 
    608 /* ata.readSectors(bus, drive, lba, count) -> data, err */
    609 static int lua_ata_read_sectors(lua_State* L) {
    610     int bus = luaL_checkinteger(L, 1);
    611     int drive = luaL_checkinteger(L, 2);
    612     uint32_t lba = (uint32_t)luaL_checkinteger(L, 3);
    613     int count = luaL_optinteger(L, 4, 1);
    614 
    615     if (count < 1 || count > 256) {
    616         lua_pushnil(L);
    617         lua_pushstring(L, "count must be 1-256");
    618         return 2;
    619     }
    620 
    621     /* Allocate buffer */
    622     size_t buf_size = count * ATA_SECTOR_SIZE;
    623     char* buffer = (char*)malloc(buf_size);
    624     if (!buffer) {
    625         lua_pushnil(L);
    626         lua_pushstring(L, "out of memory");
    627         return 2;
    628     }
    629 
    630     /* Read sectors */
    631     int result = ata_read_sectors(bus, drive, lba, count, buffer);
    632 
    633     if (result != ATA_OK) {
    634         free(buffer);
    635         lua_pushnil(L);
    636         lua_pushstring(L, ata_get_error_string(result));
    637         return 2;
    638     }
    639 
    640     /* Return data as string */
    641     lua_pushlstring(L, buffer, buf_size);
    642     free(buffer);
    643 
    644     lua_pushnil(L);  /* No error */
    645     return 2;
    646 }
    647 
    648 /* ata.writeSectors(bus, drive, lba, data) -> ok, err */
    649 static int lua_ata_write_sectors(lua_State* L) {
    650     int bus = luaL_checkinteger(L, 1);
    651     int drive = luaL_checkinteger(L, 2);
    652     uint32_t lba = (uint32_t)luaL_checkinteger(L, 3);
    653     size_t data_len;
    654     const char* data = luaL_checklstring(L, 4, &data_len);
    655 
    656     /* Data must be a multiple of sector size */
    657     if (data_len == 0 || data_len % ATA_SECTOR_SIZE != 0) {
    658         lua_pushboolean(L, 0);
    659         lua_pushstring(L, "data length must be a multiple of 512 bytes");
    660         return 2;
    661     }
    662 
    663     int count = data_len / ATA_SECTOR_SIZE;
    664     if (count > 256) {
    665         lua_pushboolean(L, 0);
    666         lua_pushstring(L, "cannot write more than 256 sectors at once");
    667         return 2;
    668     }
    669 
    670     /* Write sectors */
    671     int result = ata_write_sectors(bus, drive, lba, count, data);
    672 
    673     if (result != ATA_OK) {
    674         lua_pushboolean(L, 0);
    675         lua_pushstring(L, ata_get_error_string(result));
    676         return 2;
    677     }
    678 
    679     lua_pushboolean(L, 1);
    680     lua_pushnil(L);  /* No error */
    681     return 2;
    682 }
    683 
    684 /* ata.flush(bus, drive) -> ok, err */
    685 static int lua_ata_flush(lua_State* L) {
    686     int bus = luaL_checkinteger(L, 1);
    687     int drive = luaL_checkinteger(L, 2);
    688 
    689     int result = ata_flush(bus, drive);
    690 
    691     if (result != ATA_OK) {
    692         lua_pushboolean(L, 0);
    693         lua_pushstring(L, ata_get_error_string(result));
    694         return 2;
    695     }
    696 
    697     lua_pushboolean(L, 1);
    698     lua_pushnil(L);
    699     return 2;
    700 }
    701 
    702 /* ata.listDrives() -> table of drives */
    703 static int lua_ata_list_drives(lua_State* L) {
    704     lua_newtable(L);
    705 
    706     int idx = 1;
    707     for (int i = 0; i < 4; i++) {
    708         if (drives[i].present) {
    709             lua_newtable(L);
    710 
    711             lua_pushstring(L, "bus");
    712             lua_pushinteger(L, drives[i].bus);
    713             lua_settable(L, -3);
    714 
    715             lua_pushstring(L, "drive");
    716             lua_pushinteger(L, drives[i].drive);
    717             lua_settable(L, -3);
    718 
    719             lua_pushstring(L, "model");
    720             lua_pushstring(L, drives[i].model);
    721             lua_settable(L, -3);
    722 
    723             lua_pushstring(L, "sectors");
    724             lua_pushinteger(L, drives[i].sectors);
    725             lua_settable(L, -3);
    726 
    727             lua_rawseti(L, -2, idx++);
    728         }
    729     }
    730 
    731     return 1;
    732 }
    733 
    734 /* Register ATA module */
    735 static const luaL_Reg ata_funcs[] = {
    736     {"init", lua_ata_init},
    737     {"getDriveInfo", lua_ata_get_drive_info},
    738     {"readSectors", lua_ata_read_sectors},
    739     {"writeSectors", lua_ata_write_sectors},
    740     {"flush", lua_ata_flush},
    741     {"listDrives", lua_ata_list_drives},
    742     {NULL, NULL}
    743 };
    744 
    745 int luaopen_ata(lua_State* L) {
    746     luaL_newlib(L, ata_funcs);
    747 
    748     /* Add constants */
    749     lua_pushinteger(L, ATA_SECTOR_SIZE);
    750     lua_setfield(L, -2, "SECTOR_SIZE");
    751 
    752     lua_pushinteger(L, 0);
    753     lua_setfield(L, -2, "PRIMARY");
    754 
    755     lua_pushinteger(L, 1);
    756     lua_setfield(L, -2, "SECONDARY");
    757 
    758     lua_pushinteger(L, 0);
    759     lua_setfield(L, -2, "MASTER");
    760 
    761     lua_pushinteger(L, 1);
    762     lua_setfield(L, -2, "SLAVE");
    763 
    764     return 1;
    765 }