imxtools: move firmware read/write to library

Split the ugly firmware read/write into a API function and a much simplified code.
Also the code can now report progress.

Change-Id: I3f998eaf0c067c6da42b1d2dd9c5a5bf43c6915d

3 files changed, 259 insertions, 188 deletions

diff --git a/utils/imxtools/scsitools/scsitool.c b/utils/imxtools/scsitools/scsitool.c
index a84b5919d7..8367af6939 100644
--- a/utils/imxtools/scsitools/scsitool.c
+++ b/utils/imxtools/scsitools/scsitool.c
@@ -31,6 +31,8 @@
 #include <sys/stat.h>
 #include <fcntl.h>
 #include <unistd.h>
+#include <inttypes.h>
+#include <errno.h>
 #include "rbscsi.h"
 #include "misc.h"
 #include "stmp_scsi.h"
@@ -147,7 +149,7 @@ static int do_info(void)
             cprintf_field("  Number of drives:", " %u\n", info.nr_drives);
         if(info.has.size)
         {
-            cprintf_field("  Media size:", " %llu ", (unsigned long long)info.size);
+            cprintf_field("  Media size:", " %" PRIu64 " ", info.size);
             cprintf(RED, "(%.3f %s)\n", get_size_natural(info.size), get_size_suffix(info.size));
         }
         if(info.has.alloc_size)
@@ -226,17 +228,17 @@ static int do_info(void)
                 continue;
             if(info.has.sector_size)
             {
-                cprintf_field("  Sector size:", " %llu ", (unsigned long long)info.sector_size);
+                cprintf_field("  Sector size:", " %" PRIu32 " ", info.sector_size);
                 cprintf(RED, "(%.3f %s)\n", get_size_natural(info.sector_size), get_size_suffix(info.sector_size));
             }
             if(info.has.erase_size)
             {
-                cprintf_field("  Erase size:", " %llu ", (unsigned long long)info.erase_size);
+                cprintf_field("  Erase size:", " %" PRIu32 " ", info.erase_size);
                 cprintf(RED, "(%.3f %s)\n", get_size_natural(info.erase_size), get_size_suffix(info.erase_size));
             }
             if(info.has.size)
             {
-                cprintf_field("  Drive size:", " %llu ", (unsigned long long)info.size);
+                cprintf_field("  Drive size:", " %" PRIu64 " ", info.size);
                 cprintf(RED, "(%.3f %s)\n", get_size_natural(info.size), get_size_suffix(info.size));
             }
             if(info.has.sector_count)
@@ -282,92 +284,70 @@ static int do_info(void)
     return 0;
 }
 
-void do_extract(const char *file)
+struct rw_fw_context_t
 {
-    FILE *f = NULL;
-    cprintf(BLUE, "Extracting firmware...\n");
-
-    struct stmp_logical_media_table_t *table = NULL;
-    int ret = stmp_get_logical_media_table(g_dev_fd, &table);
-    if(ret)
-    {
-        cprintf(GREY, "Cannot get logical table: %d\n", ret);
-        goto Lend;
-    }
-    int entry = 0;
-    while(entry < table->header.count)
-        if(table->entry[entry].type == SCSI_STMP_DRIVE_TYPE_SYSTEM &&
-                table->entry[entry].tag == SCSI_STMP_DRIVE_TAG_SYSTEM_BOOT)
-            break;
-        else
-            entry++;
-    if(entry == table->header.count)
-    {
-        cprintf(GREY, "Cannot find firmware partition\n");
-        goto Lend;
-    }
-    uint8_t drive_no = table->entry[entry].drive_no;
-    uint64_t drive_sz = table->entry[entry].size;
-    if(g_debug)
-    {
-        cprintf(RED, "* ");
-        cprintf_field("Drive: ", "%#x\n", drive_no);
-        cprintf(RED, "* ");
-        cprintf_field("Size: ", "%#llx\n", (unsigned long long)drive_sz);
-    }
-    struct stmp_logical_drive_info_t info;
-    ret = stmp_get_logical_drive_info(g_dev_fd, drive_no, &info);
-    if(ret || !info.has.sector_size)
-    {
-        cprintf(GREY, "Cannot get sector size\n");
-        goto Lend;
-    }
-    unsigned sector_size = info.sector_size;
-    if(g_debug)
-    {
-        cprintf(RED, "* ");
-        cprintf_field("Sector size: ", "%lu\n", (unsigned long)sector_size);
-    }
-    uint8_t *sector = malloc(sector_size);
-    ret = stmp_read_logical_drive_sectors(g_dev_fd, drive_no, 0, 1, sector, sector_size);
-    if(ret)
-    {
-        cprintf(GREY, "Cannot read first sector: %d\n", ret);
-        goto Lend;
-    }
-    uint32_t fw_size = *(uint32_t *)(sector + 0x1c) * 16;
-    if(g_debug)
+    int tot_size;
+    int cur_size;
+    int last_percent;
+    FILE *f;
+    bool read;
+};
+
+int rw_fw(void *user, void *buf, size_t size)
+{
+    struct rw_fw_context_t *ctx = user;
+    int this_percent = (ctx->cur_size * 100LLU) / ctx->tot_size;
+    if(this_percent != ctx->last_percent && (this_percent % 5) == 0)
     {
-        cprintf(RED, "* ");
-        cprintf_field("Firmware size: ", "%#x\n", fw_size);
+        cprintf(RED, "%d%%", this_percent);
+        cprintf(YELLOW, "...");
+        fflush(stdout);
     }
+    ctx->last_percent = this_percent;
+    int ret = -1;
+    if(ctx->read)
+        ret = fread(buf, size, 1, ctx->f);
+    else
+        ret = fwrite(buf, size, 1, ctx->f);
+    ctx->cur_size += size;
+    if(ret != 1)
+        return -1;
+    else
+        return size;
+}
 
-    f = fopen(file, "wb");
+void rw_finish(struct rw_fw_context_t *ctx)
+{
+    if(ctx->last_percent == 100)
+        return;
+    cprintf(RED, "100%%\n");
+}
+
+void do_extract(const char *file)
+{
+    FILE *f = fopen(file, "wb");
     if(f == NULL)
     {
-        cprintf(GREY, "Cannot open '%s' for writing: %m\n", file);
-        goto Lend;
+        cprintf(GREY, "Cannot open output file: %s\n", strerror(errno));
+        return;
     }
-
-    for(int sec = 0; sec * sector_size < fw_size; sec++)
+    int ret = stmp_read_firmware(g_dev_fd, NULL, NULL);
+    if(ret < 0)
     {
-        ret = stmp_read_logical_drive_sectors(g_dev_fd, drive_no, sec, 1, sector, sector_size);
-        if(ret)
-        {
-            cprintf(GREY, "Cannot read sector %d: %d\n", sec, ret);
-            goto Lend;
-        }
-        if(fwrite(sector, sector_size, 1, f) != 1)
-        {
-            cprintf(GREY, "Write failed: %m\n");
-            goto Lend;
-        }
+        cprintf(GREY, "Cannot get firmware size: %d\n", ret);
+        return;
     }
-    cprintf(BLUE, "Done\n");
-Lend:
-    free(table);
-    if(f)
-        fclose(f);
+    struct rw_fw_context_t ctx;
+    ctx.tot_size = ret;
+    ctx.cur_size = 0;
+    ctx.f = f;
+    ctx.last_percent = -1;
+    ctx.read = false;
+    ret = stmp_read_firmware(g_dev_fd, &ctx, &rw_fw);
+    if(ret < 0)
+        cprintf(GREY, "Cannot read firmware: %d\n", ret);
+    rw_finish(&ctx);
+    fclose(f);
 }
 
 void do_write(const char *file, int want_a_brick)
@@ -379,115 +359,25 @@ void do_write(const char *file, int want_a_brick)
         cprintf(GREY, "option on the command line and do not complain if you end up with a brick ;)\n");
         return;
     }
-    FILE *f = NULL;
-    cprintf(BLUE, "Writing firmware...\n");
-
-    struct stmp_logical_media_table_t *table = NULL;
-    int ret = stmp_get_logical_media_table(g_dev_fd, &table);
-    if(ret)
-    {
-        cprintf(GREY, "Cannot get logical table: %d\n", ret);
-        goto Lend;
-    }
-    int entry = 0;
-    while(entry < table->header.count)
-        if(table->entry[entry].type == SCSI_STMP_DRIVE_TYPE_SYSTEM &&
-                table->entry[entry].tag == SCSI_STMP_DRIVE_TAG_SYSTEM_BOOT)
-            break;
-        else
-            entry++;
-    if(entry == table->header.count)
-    {
-        cprintf(GREY, "Cannot find firmware partition\n");
-        goto Lend;
-    }
-    uint8_t drive_no = table->entry[entry].drive_no;
-    uint64_t drive_sz = table->entry[entry].size;
-    if(g_debug)
-    {
-        cprintf(RED, "* ");
-        cprintf_field("Drive: ", "%#x\n", drive_no);
-        cprintf(RED, "* ");
-        cprintf_field("Size: ", "%#llx\n", (unsigned long long)drive_sz);
-    }
-    struct stmp_logical_drive_info_t info;
-    ret = stmp_get_logical_drive_info(g_dev_fd, drive_no, &info);
-    if(ret || !info.has.sector_size)
-    {
-        cprintf(GREY, "Cannot get sector size\n");
-        goto Lend;
-    }
-    unsigned sector_size = info.sector_size;
-    uint8_t *sector = malloc(sector_size);
-
-    /* sanity check by reading first sector */
-    ret = stmp_read_logical_drive_sectors(g_dev_fd, drive_no, 0, 1, sector, sector_size);
-    if(ret)
-    {
-        cprintf(GREY, "Cannot read first sector: %d\n", ret);
-        return;
-    }
-    uint32_t sig = *(uint32_t *)(sector + 0x14);
-    if(sig != 0x504d5453)
-    {
-        cprintf(GREY, "There is something wrong: the first sector doesn't have the STMP signature. Bailing out...\n");
-        return;
-    }
-
-    f = fopen(file, "rb");
+    FILE *f = fopen(file, "rb");
     if(f == NULL)
     {
-        cprintf(GREY, "Cannot open '%s' for writing: %m\n", file);
-        goto Lend;
+        cprintf(GREY, "Cannot open output file: %s\n", strerror(errno));
+        return;
     }
+    struct rw_fw_context_t ctx;
     fseek(f, 0, SEEK_END);
-    int fw_size = ftell(f);
+    ctx.tot_size = ftell(f);
     fseek(f, 0, SEEK_SET);
-    if(g_debug)
-    {
-        cprintf(RED, "* ");
-        cprintf_field("Firmware size: ", "%#x\n", fw_size);
-    }
-    /* sanity check size */
-    if((uint64_t)fw_size > drive_sz)
-    {
-        cprintf(GREY, "You cannot write a firmware greater than the partition size.\n");
-        goto Lend;
-    }
-
-    int percent = -1;
-    for(int off = 0; off < fw_size; off += sector_size)
-    {
-        int sec = off / sector_size;
-        int this_percent = (sec * 100) / (fw_size / sector_size);
-        if(this_percent != percent && (this_percent % 5) == 0)
-        {
-            cprintf(RED, "%d%%", this_percent);
-            cprintf(YELLOW, "...");
-            fflush(stdout);
-        }
-        percent = this_percent;
-        int xfer_len = MIN(fw_size - off, (int)sector_size);
-        if(fread(sector, xfer_len, 1, f) != 1)
-        {
-            cprintf(GREY, "Read failed: %m\n");
-            goto Lend;
-        }
-        /* NOTE transfer a whole sector even if incomplete, the device won't access
-         * partial sectors */
-        if(xfer_len < (int)sector_size)
-            memset(sector + xfer_len, 0, sector_size - xfer_len);
-        ret = stmp_write_logical_drive_sectors(g_dev_fd, drive_no, sec, 1, sector, sector_size);
-        if(ret)
-        {
-            cprintf(GREY, "Cannot write sector %d: %d\n", sec, ret);
-            goto Lend;
-        }
-    }
-    cprintf(BLUE, "Done\n");
-Lend:
-    if(f)
-        fclose(f);
+    ctx.cur_size = 0;
+    ctx.f = f;
+    ctx.last_percent = -1;
+    ctx.read = true;
+    int ret = stmp_write_firmware(g_dev_fd, &ctx, &rw_fw);
+    if(ret < 0)
+        cprintf(GREY, "Cannot write firmware: %d\n", ret);
+    rw_finish(&ctx);
+    fclose(f);
 }
 
 static void usage(void)
@@ -586,14 +476,14 @@ int main(int argc, char **argv)
     rb_scsi_device_t scsi_dev = rb_scsi_open(argv[optind], g_debug ? RB_SCSI_DEBUG : 0, NULL, scsi_printf);
     if(scsi_dev == 0)
     {
-        cprintf(GREY, "Cannot open device: %m\n");
+        cprintf(GREY, "Cannot open device\n");
         ret = 1;
         goto Lend;
     }
     g_dev_fd = stmp_open(scsi_dev, g_debug ? STMP_DEBUG : 0, NULL, scsi_printf);
     if(g_dev_fd == 0)
     {
-        cprintf(GREY, "Cannot open stmp device: %m\n");
+        cprintf(GREY, "Cannot open stmp device\n");
         ret = 2;
         goto Lend;
     }
diff --git a/utils/imxtools/scsitools/stmp_scsi.c b/utils/imxtools/scsitools/stmp_scsi.c
index dcda91e295..e1d0dc4c77 100644
--- a/utils/imxtools/scsitools/stmp_scsi.c
+++ b/utils/imxtools/scsitools/stmp_scsi.c
@@ -23,7 +23,6 @@
 #include <stdarg.h>
 #include <stdio.h>
 #define _BSD_SOURCE
-#include <endian.h>
 #include "stmp_scsi.h"
 
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -131,6 +130,7 @@ int stmp_sense_analysis(stmp_device_t dev, int status, uint8_t *sense, int sense
         for(int i = 0; i < sense_size; i++)
             stmp_printf(dev, " %02x", sense[i]);
         stmp_printf(dev, "\n");
+        rb_scsi_decode_sense(dev->dev, sense, sense_size);
     }
     return status;
 }
@@ -614,3 +614,176 @@ int stmp_get_logical_drive_info(stmp_device_t dev, uint8_t drive, struct stmp_lo
     }
     return 0;
 }
+
+int stmp_read_firmware(stmp_device_t dev, void *user, stmp_fw_rw_fn_t fn)
+{
+    /* read logicial table */
+    uint8_t *sector = NULL;
+    struct stmp_logical_media_table_t *table = NULL;
+    int ret = stmp_get_logical_media_table(dev, &table);
+    if(ret)
+    {
+        stmp_printf(dev, "Cannot get logical table: %d\n", ret);
+        return -1;
+    }
+    /* locate firmware partition */
+    int entry = 0;
+    while(entry < table->header.count)
+        if(table->entry[entry].type == SCSI_STMP_DRIVE_TYPE_SYSTEM &&
+                table->entry[entry].tag == SCSI_STMP_DRIVE_TAG_SYSTEM_BOOT)
+            break;
+        else
+            entry++;
+    if(entry == table->header.count)
+    {
+        stmp_printf(dev, "Cannot find firmware partition\n");
+        goto Lerr;
+    }
+    uint8_t drive_no = table->entry[entry].drive_no;
+    uint64_t drive_sz = table->entry[entry].size;
+    stmp_debugf(dev, "Firmware drive: %#x\n", drive_no);
+    stmp_debugf(dev, "Firmware max size: %#llx\n", (unsigned long long)drive_sz);
+    /* get drive info */
+    struct stmp_logical_drive_info_t info;
+    ret = stmp_get_logical_drive_info(dev, drive_no, &info);
+    if(ret || !info.has.sector_size)
+    {
+        stmp_printf(dev, "Cannot get sector size\n");
+        goto Lerr;
+    }
+    unsigned sector_size = info.sector_size;
+    stmp_debugf(dev, "Firmware sector size: %lu\n", (unsigned long)sector_size);
+    /* allocate a buffer for one sector */
+    sector = malloc(sector_size);
+    /* read the first sector to check it is correct and get the total size */
+    ret = stmp_read_logical_drive_sectors(dev, drive_no, 0, 1, sector, sector_size);
+    if(ret)
+    {
+        stmp_printf(dev, "Cannot read first sector: %d\n", ret);
+        goto Lerr;
+    }
+    uint32_t sig = *(uint32_t *)(sector + 0x14);
+    if(sig != 0x504d5453)
+    {
+        stmp_printf(dev, "There is something wrong: the first sector doesn't have the STMP signature.\n");
+        goto Lerr;
+    }
+    uint32_t fw_size = *(uint32_t *)(sector + 0x1c) * 16; /* see SB file format */
+    stmp_debugf(dev, "Firmware size: %#x\n", fw_size);
+    /* if fn is NULL, just return the size immediately */
+    if(fn != NULL)
+    {
+        /* read all sectors one by one */
+        for(int sec = 0; sec * sector_size < fw_size; sec++)
+        {
+            ret = stmp_read_logical_drive_sectors(dev, drive_no, sec, 1, sector, sector_size);
+            if(ret)
+            {
+                stmp_printf(dev, "Cannot read sector %d: %d\n", sec, ret);
+                goto Lerr;
+            }
+            int xfer_len = MIN(sector_size, fw_size - sec * sector_size);
+            ret = fn(user, sector, xfer_len);
+            if(ret != xfer_len)
+            {
+                stmp_printf(dev, "User write failed: %d\n", ret);
+                goto Lerr;
+            }
+        }
+    }
+    ret = fw_size;
+Lend:
+    free(table);
+    if(sector)
+        free(sector);
+    return ret;
+Lerr:
+    ret = -1;
+    goto Lend;
+}
+
+int stmp_write_firmware(stmp_device_t dev, void *user, stmp_fw_rw_fn_t fn)
+{
+    /* read logicial table */
+    struct stmp_logical_media_table_t *table = NULL;
+    int ret = stmp_get_logical_media_table(dev, &table);
+    if(ret)
+    {
+        stmp_printf(dev, "Cannot get logical table: %d\n", ret);
+        return -1;
+    }
+    /* locate firmware partition */
+    int entry = 0;
+    while(entry < table->header.count)
+        if(table->entry[entry].type == SCSI_STMP_DRIVE_TYPE_SYSTEM &&
+                table->entry[entry].tag == SCSI_STMP_DRIVE_TAG_SYSTEM_BOOT)
+            break;
+        else
+            entry++;
+    if(entry == table->header.count)
+    {
+        stmp_printf(dev, "Cannot find firmware partition\n");
+        goto Lerr;
+    }
+    uint8_t drive_no = table->entry[entry].drive_no;
+    uint64_t drive_sz = table->entry[entry].size;
+    stmp_debugf(dev, "Firmware drive: %#x\n", drive_no);
+    stmp_debugf(dev, "Firmware max size: %#llx\n", (unsigned long long)drive_sz);
+    /* get drive info */
+    struct stmp_logical_drive_info_t info;
+    ret = stmp_get_logical_drive_info(dev, drive_no, &info);
+    if(ret || !info.has.sector_size)
+    {
+        stmp_printf(dev, "Cannot get sector size\n");
+        goto Lerr;
+    }
+    unsigned sector_size = info.sector_size;
+    stmp_debugf(dev, "Firmware sector size: %lu\n", (unsigned long)sector_size);
+    /* allocate a buffer for one sector */
+    uint8_t *sector = malloc(sector_size);
+    /* read the first sector to check it is correct and get the total size */
+    ret = fn(user, sector, sector_size);
+    /* the whole file could be smaller than one sector, but it must be greater
+     * then the header size */
+    if(ret < 0x20)
+    {
+        stmp_printf(dev, "User read failed: %d\n", ret);
+        goto Lerr;
+    }
+    uint32_t sig = *(uint32_t *)(sector + 0x14);
+    if(sig != 0x504d5453)
+    {
+        stmp_printf(dev, "There is something wrong: the first sector doesn't have the STMP signature.\n");
+        goto Lerr;
+    }
+    uint32_t fw_size = *(uint32_t *)(sector + 0x1c) * 16; /* see SB file format */
+    stmp_debugf(dev, "Firmware size: %#x\n", fw_size);
+    /* write all sectors one by one */
+    for(int sec = 0; sec * sector_size < fw_size; sec++)
+    {
+        int xfer_len = MIN(sector_size, fw_size - sec * sector_size);
+        /* avoid rereading the first sector */
+        if(sec != 0)
+            ret = fn(user, sector, xfer_len);
+        if(ret != xfer_len)
+        {
+            stmp_printf(dev, "User read failed: %d\n", ret);
+            goto Lerr;
+        }
+        if(ret < (int)sector_size)
+            memset(sector + ret, 0, sector_size - ret);
+        ret = stmp_write_logical_drive_sectors(dev, drive_no, sec, 1, sector, sector_size);
+        if(ret)
+        {
+            stmp_printf(dev, "Cannot write sector %d: %d\n", sec, ret);
+            goto Lerr;
+        }
+    }
+    ret = fw_size;
+Lend:
+    free(table);
+    return ret;
+Lerr:
+    ret = -1;
+    goto Lend;
+}
diff --git a/utils/imxtools/scsitools/stmp_scsi.h b/utils/imxtools/scsitools/stmp_scsi.h
index 68d77daeec..56068b5d4c 100644
--- a/utils/imxtools/scsitools/stmp_scsi.h
+++ b/utils/imxtools/scsitools/stmp_scsi.h
@@ -355,6 +355,14 @@ int stmp_read_logical_drive_sectors(stmp_device_t dev, uint8_t drive, uint64_t a
     uint32_t count, void *buffer, int buffer_size);
 int stmp_write_logical_drive_sectors(stmp_device_t dev, uint8_t drive, uint64_t address,
     uint32_t count, void *buffer, int buffer_size);
+/* return <0 on error, or firmware size in bytes otherwise,
+ * if not NULL, the read/write function will be called as many times as needed to provide
+ * the entire firmware, it should return number of bytes read/written on success or -1 on error
+ * in all cases, the total size of the firmware is based on the header
+ * if NULL for read, return firmware size */
+typedef int (*stmp_fw_rw_fn_t)(void *user, void *buf, size_t size);
+int stmp_read_firmware(stmp_device_t dev, void *user, stmp_fw_rw_fn_t fn);
+int stmp_write_firmware(stmp_device_t dev, void *user, stmp_fw_rw_fn_t fn);
 /* string helpers */
 const char *stmp_get_logical_media_type_string(uint32_t type);
 const char *stmp_get_logical_media_vendor_string(uint32_t type);