fastboot: add support for sparse images in flashall and update

Change-Id: I66a73b16a988a65fc91fb22a26d11986025089de
Signed-off-by: Rom Lemarchand <romlem@google.com>
diff --git a/fastboot/fastboot.c b/fastboot/fastboot.c
index e469d97..f186c93 100644
--- a/fastboot/fastboot.c
+++ b/fastboot/fastboot.c
@@ -43,6 +43,7 @@
 
 #include <sys/time.h>
 #include <sys/types.h>
+#include <sys/stat.h>
 
 #include <bootimg.h>
 #include <sparse/sparse.h>
@@ -54,6 +55,8 @@
 #define O_BINARY 0
 #endif
 
+#define ARRAY_SIZE(a) (sizeof(a)/sizeof(*(a)))
+
 char cur_product[FB_RESPONSE_SZ + 1];
 
 void bootimg_set_cmdline(boot_img_hdr *h, const char *cmdline);
@@ -81,6 +84,27 @@
 unsigned second_offset  = 0x00f00000;
 unsigned tags_offset    = 0x00000100;
 
+enum fb_buffer_type {
+    FB_BUFFER,
+    FB_BUFFER_SPARSE,
+};
+
+struct fastboot_buffer {
+    enum fb_buffer_type type;
+    void *data;
+    unsigned int sz;
+};
+
+static struct {
+    char img_name[13];
+    char sig_name[13];
+    char part_name[9];
+    bool is_optional;
+} images[3] = {
+    {"boot.img", "boot.sig", "boot", false},
+    {"recovery.img", "recovery.sig", "recovery", true},
+    {"system.img", "system.sig", "system", false},
+};
 
 void die(const char *fmt, ...)
 {
@@ -135,40 +159,28 @@
     return strdup(path);
 }
 
-#if defined(__APPLE__) && defined(__MACH__)
-#define lseek64 lseek
-#define off64_t off_t
-#endif
-
-static int64_t file_size(const char *fn)
+static int64_t file_size(int fd)
 {
-    off64_t off;
-    int fd;
+    struct stat st;
+    int ret;
 
-    fd = open(fn, O_RDONLY | O_BINARY);
-    if (fd < 0) return -1;
+    ret = fstat(fd, &st);
 
-    off = lseek64(fd, 0, SEEK_END);
-    close(fd);
-
-    return off;
+    return ret ? -1 : st.st_size;
 }
 
-static void *load_file(const char *fn, unsigned *_sz)
+static void *load_fd(int fd, unsigned *_sz)
 {
     char *data;
     int sz;
-    int fd;
     int errno_tmp;
 
     data = 0;
-    fd = open(fn, O_RDONLY | O_BINARY);
-    if(fd < 0) return 0;
 
-    sz = lseek(fd, 0, SEEK_END);
-    if(sz < 0) goto oops;
-
-    if(lseek(fd, 0, SEEK_SET) != 0) goto oops;
+    sz = file_size(fd);
+    if (sz < 0) {
+        goto oops;
+    }
 
     data = (char*) malloc(sz);
     if(data == 0) goto oops;
@@ -187,6 +199,16 @@
     return 0;
 }
 
+static void *load_file(const char *fn, unsigned *_sz)
+{
+    int fd;
+
+    fd = open(fn, O_RDONLY | O_BINARY);
+    if(fd < 0) return 0;
+
+    return load_fd(fd, _sz);
+}
+
 int match_fastboot_with_serial(usb_ifc_info *info, const char *local_serial)
 {
     if(!(vendor_id && (info->dev_vendor == vendor_id)) &&
@@ -392,6 +414,31 @@
     return data;
 }
 
+static int unzip_to_file(zipfile_t zip, char *name)
+{
+    int fd;
+    char *data;
+    unsigned sz;
+
+    fd = fileno(tmpfile());
+    if (fd < 0) {
+        return -1;
+    }
+
+    data = unzip_file(zip, name, &sz);
+    if (data == 0) {
+        return -1;
+    }
+
+    if (write(fd, data, sz) != sz) {
+        fd = -1;
+    }
+
+    free(data);
+    lseek(fd, 0, SEEK_SET);
+    return fd;
+}
+
 static char *strip(char *s)
 {
     int n;
@@ -490,27 +537,20 @@
     fb_queue_notice("--------------------------------------------");
 }
 
-
-struct sparse_file **load_sparse_files(const char *fname, int max_size)
+static struct sparse_file **load_sparse_files(int fd, int max_size)
 {
-    int fd;
     struct sparse_file *s;
     int files;
     struct sparse_file **out_s;
 
-    fd = open(fname, O_RDONLY | O_BINARY);
-    if (fd < 0) {
-        die("cannot open '%s'\n", fname);
-    }
-
     s = sparse_file_import_auto(fd, false);
     if (!s) {
-        die("cannot sparse read file '%s'\n", fname);
+        die("cannot sparse read file\n");
     }
 
     files = sparse_file_resparse(s, max_size, NULL, 0);
     if (files < 0) {
-        die("Failed to resparse '%s'\n", fname);
+        die("Failed to resparse\n");
     }
 
     out_s = calloc(sizeof(struct sparse_file *), files + 1);
@@ -520,7 +560,7 @@
 
     files = sparse_file_resparse(s, max_size, out_s, files);
     if (files < 0) {
-        die("Failed to resparse '%s'\n", fname);
+        die("Failed to resparse\n");
     }
 
     return out_s;
@@ -581,29 +621,78 @@
      return fb_format_supported(usb, part);
 }
 
-void do_flash(usb_handle *usb, const char *pname, const char *fname)
+static int load_buf_fd(usb_handle *usb, int fd,
+        struct fastboot_buffer *buf)
 {
     int64_t sz64;
     void *data;
     int64_t limit;
 
-    sz64 = file_size(fname);
+    sz64 = file_size(fd);
+    if (sz64 < 0) {
+        return -1;
+    }
     limit = get_sparse_limit(usb, sz64);
     if (limit) {
-        struct sparse_file **s = load_sparse_files(fname, limit);
+        struct sparse_file **s = load_sparse_files(fd, limit);
         if (s == NULL) {
-            die("cannot sparse load '%s'\n", fname);
+            return -1;
         }
-        while (*s) {
-            sz64 = sparse_file_len(*s, true, false);
-            fb_queue_flash_sparse(pname, *s++, sz64);
-        }
+        buf->type = FB_BUFFER_SPARSE;
+        buf->data = s;
     } else {
         unsigned int sz;
-        data = load_file(fname, &sz);
-        if (data == 0) die("cannot load '%s': %s\n", fname, strerror(errno));
-        fb_queue_flash(pname, data, sz);
+        data = load_fd(fd, &sz);
+        if (data == 0) return -1;
+        buf->type = FB_BUFFER;
+        buf->data = data;
+        buf->sz = sz;
     }
+
+    return 0;
+}
+
+static int load_buf(usb_handle *usb, const char *fname,
+        struct fastboot_buffer *buf)
+{
+    int fd;
+
+    fd = open(fname, O_RDONLY | O_BINARY);
+    if (fd < 0) {
+        die("cannot open '%s'\n", fname);
+    }
+
+    return load_buf_fd(usb, fd, buf);
+}
+
+static void flash_buf(const char *pname, struct fastboot_buffer *buf)
+{
+    struct sparse_file **s;
+
+    switch (buf->type) {
+        case FB_BUFFER_SPARSE:
+            s = buf->data;
+            while (*s) {
+                int64_t sz64 = sparse_file_len(*s, true, false);
+                fb_queue_flash_sparse(pname, *s++, sz64);
+            }
+            break;
+        case FB_BUFFER:
+            fb_queue_flash(pname, buf->data, buf->sz);
+            break;
+        default:
+            die("unknown buffer type: %d", buf->type);
+    }
+}
+
+void do_flash(usb_handle *usb, const char *pname, const char *fname)
+{
+    struct fastboot_buffer buf;
+
+    if (load_buf(usb, fname, &buf)) {
+        die("cannot load '%s'", fname);
+    }
+    flash_buf(pname, &buf);
 }
 
 void do_update_signature(zipfile_t zip, char *fn)
@@ -616,13 +705,17 @@
     fb_queue_command("signature", "installing signature");
 }
 
-void do_update(char *fn, int erase_first)
+void do_update(usb_handle *usb, char *fn, int erase_first)
 {
     void *zdata;
     unsigned zsize;
     void *data;
     unsigned sz;
     zipfile_t zip;
+    int fd;
+    int rc;
+    struct fastboot_buffer buf;
+    int i;
 
     queue_info_dump();
 
@@ -651,30 +744,25 @@
 
     setup_requirements(data, sz);
 
-    data = unzip_file(zip, "boot.img", &sz);
-    if (data == 0) die("update package missing boot.img");
-    do_update_signature(zip, "boot.sig");
-    if (erase_first && needs_erase("boot")) {
-        fb_queue_erase("boot");
-    }
-    fb_queue_flash("boot", data, sz);
-
-    data = unzip_file(zip, "recovery.img", &sz);
-    if (data != 0) {
-        do_update_signature(zip, "recovery.sig");
-        if (erase_first && needs_erase("recovery")) {
-            fb_queue_erase("recovery");
+    for (i = 0; i < ARRAY_SIZE(images); i++) {
+        fd = unzip_to_file(zip, images[i].img_name);
+        if (fd < 0) {
+            if (images[i].is_optional)
+                continue;
+            die("update package missing %s", images[i].img_name);
         }
-        fb_queue_flash("recovery", data, sz);
+        rc = load_buf_fd(usb, fd, &buf);
+        if (rc) die("cannot load %s from flash", images[i].img_name);
+        do_update_signature(zip, images[i].sig_name);
+        if (erase_first && needs_erase(images[i].part_name)) {
+            fb_queue_erase(images[i].part_name);
+        }
+        flash_buf(images[i].part_name, &buf);
+        /* not closing the fd here since the sparse code keeps the fd around
+         * but hasn't mmaped data yet. The tmpfile will get cleaned up when the
+         * program exits.
+         */
     }
-
-    data = unzip_file(zip, "system.img", &sz);
-    if (data == 0) die("update package missing system.img");
-    do_update_signature(zip, "system.sig");
-    if (erase_first && needs_erase("system")) {
-        fb_queue_erase("system");
-    }
-    fb_queue_flash("system", data, sz);
 }
 
 void do_send_signature(char *fn)
@@ -695,11 +783,13 @@
     fb_queue_command("signature", "installing signature");
 }
 
-void do_flashall(int erase_first)
+void do_flashall(usb_handle *usb, int erase_first)
 {
     char *fname;
     void *data;
     unsigned sz;
+    struct fastboot_buffer buf;
+    int i;
 
     queue_info_dump();
 
@@ -711,33 +801,19 @@
     if (data == 0) die("could not load android-info.txt: %s", strerror(errno));
     setup_requirements(data, sz);
 
-    fname = find_item("boot", product);
-    data = load_file(fname, &sz);
-    if (data == 0) die("could not load boot.img: %s", strerror(errno));
-    do_send_signature(fname);
-    if (erase_first && needs_erase("boot")) {
-        fb_queue_erase("boot");
-    }
-    fb_queue_flash("boot", data, sz);
-
-    fname = find_item("recovery", product);
-    data = load_file(fname, &sz);
-    if (data != 0) {
-        do_send_signature(fname);
-        if (erase_first && needs_erase("recovery")) {
-            fb_queue_erase("recovery");
+    for (i = 0; i < ARRAY_SIZE(images); i++) {
+        fname = find_item(images[i].part_name, product);
+        if (load_buf(usb, fname, &buf)) {
+            if (images[i].is_optional)
+                continue;
+            die("could not load %s\n", images[i].img_name);
         }
-        fb_queue_flash("recovery", data, sz);
+        do_send_signature(fname);
+        if (erase_first && needs_erase(images[i].part_name)) {
+            fb_queue_erase(images[i].part_name);
+        }
+        flash_buf(images[i].part_name, &buf);
     }
-
-    fname = find_item("system", product);
-    data = load_file(fname, &sz);
-    if (data == 0) die("could not load system.img: %s", strerror(errno));
-    do_send_signature(fname);
-    if (erase_first && needs_erase("system")) {
-        fb_queue_erase("system");
-    }
-    fb_queue_flash("system", data, sz);
 }
 
 #define skip(n) do { argc -= (n); argv += (n); } while (0)
@@ -997,14 +1073,14 @@
             fb_queue_flash(pname, data, sz);
         } else if(!strcmp(*argv, "flashall")) {
             skip(1);
-            do_flashall(erase_first);
+            do_flashall(usb, erase_first);
             wants_reboot = 1;
         } else if(!strcmp(*argv, "update")) {
             if (argc > 1) {
-                do_update(argv[1], erase_first);
+                do_update(usb, argv[1], erase_first);
                 skip(2);
             } else {
-                do_update("update.zip", erase_first);
+                do_update(usb, "update.zip", erase_first);
                 skip(1);
             }
             wants_reboot = 1;