Download progress indicator for FTP and HTTP in stage 1.

Patch from Ricky Zhou to add a pinwheel and progress bar for
FTP and HTTP downloads from loader.

11 files changed, 163 insertions, 24 deletions

diff --git a/loader/ftp.c b/loader/ftp.c
index ce6deb5d4..e84237ce7 100644
--- a/loader/ftp.c
+++ b/loader/ftp.c
@@ -346,13 +346,15 @@ int ftpOpen(char *host, int family, char *name, char *password,
  * RFC 2428 FTP Extensions for IPv6 and NATs
  */
 int ftpGetFileDesc(int sock, struct in6_addr host, int family,
-                   char * remotename) {
+                   char * remotename, long long *size) {
     int dataSocket;
     struct sockaddr_in dataAddress;
     struct sockaddr_in6 dataAddress6;
     int i, j;
     char * passReply;
     char * chptr;
+    char * sizeCommand;
+    char * sizeReply;
     char * retrCommand;
     int rc;
 
@@ -458,6 +460,33 @@ int ftpGetFileDesc(int sock, struct in6_addr host, int family,
         return FTPERR_FAILED_CONNECT;
     }
 
+    sizeCommand = alloca(strlen(remotename) + 20);
+    sprintf(sizeCommand, "SIZE %s\r\n", remotename);
+    i = strlen(sizeCommand);
+    if (write(sock, sizeCommand, i) != i) {
+        return FTPERR_SERVER_IO_ERROR;
+    }
+
+    if (ftpCheckResponse(sock, &sizeReply)) {
+        /* No worries, the SIZE command isn't in RFC 959 anyway. */
+        *size = 0;
+    } else {
+        /* We have a SIZE response of the form:
+         * 213 95600640
+         * where 95600640 is the size in bytes.
+         */
+
+        /* Skip to first non-space character */
+        while (isspace(*sizeReply) && *sizeReply) sizeReply++;
+        /* Skip reply code */
+        while (!isspace(*sizeReply) && *sizeReply) sizeReply++;
+        /* Skip any remaining whitespace */
+        while (isspace(*sizeReply) && *sizeReply) sizeReply++;
+        /* sizeReply now points to the beginning of the size */
+        if (sscanf(sizeReply, "%lld", size) != 1) *size = 0;
+        if (*size < 0) *size = 0;
+    }
+
     retrCommand = alloca(strlen(remotename) + 20);
     sprintf(retrCommand, "RETR %s\r\n", remotename);
     i = strlen(retrCommand);
@@ -668,10 +697,11 @@ static char *find_status_code (char *headers)
  * by '\r\n', ending with '\r\n'.
  */
 int httpGetFileDesc(char * hostname, int port, char * remotename,
-                    char *extraHeaders) {
+                    char *extraHeaders, long long *size) {
     char * buf, *headers = NULL;
     char *status;
     char *hstr;
+    char *contlen;
     int family;
     struct in_addr addr;
     struct in6_addr addr6;
@@ -740,6 +770,22 @@ int httpGetFileDesc(char * hostname, int port, char * remotename,
         if (headers) free(headers);
         return FTPERR_SERVER_IO_ERROR;
     } else if (!strncmp(status, "200", 3)) {
+        contlen = find_header(headers, "Content-Length");
+
+        if (contlen == NULL) {
+            *size = 0;
+        } else {
+            errno = 0;
+            *size = strtoll(contlen, NULL, 10);
+
+            if ((errno == ERANGE && (*size == LLONG_MIN || *size == LLONG_MAX)) ||
+                (errno != 0 && *size == 0)) {
+                logMessage(ERROR, "%s: %d: %m", __func__, __LINE__);
+                abort();
+            }
+        }
+
+        if (*size < 0) *size = 0;
         if (status) free(status);
         if (headers) free(headers);
         return sock;
@@ -760,7 +806,7 @@ int httpGetFileDesc(char * hostname, int port, char * remotename,
 
         logMessage(INFO, "redirecting to %s", redir_loc);
         convertURLToUI(redir_loc, &ui);
-        retval = httpGetFileDesc (ui.address, -1, ui.prefix, extraHeaders);
+        retval = httpGetFileDesc (ui.address, -1, ui.prefix, extraHeaders, size);
         free(redir_loc);
         return retval;
     } else if (!strncmp(status, "403", 3)) {
diff --git a/loader/ftp.h b/loader/ftp.h
index c1e7fcb87..9186e880c 100644
--- a/loader/ftp.h
+++ b/loader/ftp.h
@@ -45,9 +45,10 @@ int ftpOpen(char * host, int family, char * name, char * password,
             int port);
 int ftpGetFile(int sock, char * remotename, int dest);
 int ftpGetFileDesc(int sock, struct in6_addr host, int family,
-                   char * remotename);
+                   char * remotename, long long *size);
 int ftpGetFileDone(int sock);
 
-int httpGetFileDesc(char * hostname, int port, char * remotename, char *extraHeaders);
+int httpGetFileDesc(char * hostname, int port, char * remotename,
+                    char *extraHeaders, long long *size);
 
 #endif
diff --git a/loader/loadermisc.c b/loader/loadermisc.c
index 1fd77093f..58812c376 100644
--- a/loader/loadermisc.c
+++ b/loader/loadermisc.c
@@ -33,12 +33,15 @@
 #include <stdlib.h>
 
 #include "log.h"
+#include "windows.h"
 
-int copyFileFd(int infd, char * dest) {
+int copyFileFd(int infd, char * dest, progressCB pbcb,
+               struct progressCBdata *data, long long total) {
     int outfd;
     char buf[4096];
     int i;
     int rc = 0;
+    long long count = 0;
 
     outfd = open(dest, O_CREAT | O_RDWR, 0666);
 
@@ -52,6 +55,12 @@ int copyFileFd(int infd, char * dest) {
             rc = 1;
             break;
         }
+
+        count += 1;
+
+        if (pbcb && data && total) {
+            pbcb(data, count, total);
+        }
     }
 
     close(outfd);
@@ -70,7 +79,7 @@ int copyFile(char * source, char * dest) {
         return 1;
     }
 
-    rc = copyFileFd(infd, dest);
+    rc = copyFileFd(infd, dest, NULL, NULL, 0);
 
     close(infd);
 
diff --git a/loader/loadermisc.h b/loader/loadermisc.h
index 8094fa70f..c7e2b527a 100644
--- a/loader/loadermisc.h
+++ b/loader/loadermisc.h
@@ -22,8 +22,11 @@
 #include <stdio.h>
 #include <stdarg.h>
 
+#include "windows.h"
+
 int copyFile(char * source, char * dest);
-int copyFileFd(int infd, char * dest);
+int copyFileFd(int infd, char * dest, progressCB pbcb,
+               struct progressCBdata *data, long long total);
 char * readLine(FILE * f);
 int simpleStringCmp(const void * a, const void * b);
 int totalMemory(void);
diff --git a/loader/method.c b/loader/method.c
index 26dc2c3fa..2313cee43 100644
--- a/loader/method.c
+++ b/loader/method.c
@@ -441,12 +441,13 @@ int mountStage2(char *stage2path) {
 
 
 /* copies a second stage from fd to dest and mounts on mntpoint */
-int copyFileAndLoopbackMount(int fd, char * dest,
-                             char * device, char * mntpoint) {
+int copyFileAndLoopbackMount(int fd, char * dest, char * device, char * mntpoint,
+                             progressCB pbcb, struct progressCBdata *data,
+                             long long total) {
     int rc;
     struct stat sb;
 
-    rc = copyFileFd(fd, dest);
+    rc = copyFileFd(fd, dest, pbcb, data, total);
     stat(dest, &sb);
     logMessage(DEBUGLVL, "copied %" PRId64 " bytes to %s (%s)", sb.st_size, dest, 
                ((rc) ? " incomplete" : "complete"));
diff --git a/loader/method.h b/loader/method.h
index 5f6c7a83d..1b5e2d3f7 100644
--- a/loader/method.h
+++ b/loader/method.h
@@ -21,6 +21,7 @@
 #define H_METHOD
 
 #include "loader.h"
+#include "windows.h"
 
 /* method identifiers, needs to match struct installMethod order in loader.c */
 enum {
@@ -46,7 +47,8 @@ void queryIsoMediaCheck(char * isoDir);
 
 void umountStage2(void);
 int mountStage2(char *stage2path);
-int copyFileAndLoopbackMount(int fd, char *dest, char *device, char *mntpoint);
+int copyFileAndLoopbackMount(int fd, char *dest, char *device, char *mntpoint,
+                             progressCB pbcb, struct progressCBdata *data, long long total);
 int getFileFromBlockDevice(char *device, char *path, char * dest);
 
 int unpackCpioBall(char * ballPath, char * rootDir);
diff --git a/loader/urlinstall.c b/loader/urlinstall.c
index 177716612..cd859d0f7 100644
--- a/loader/urlinstall.c
+++ b/loader/urlinstall.c
@@ -56,6 +56,8 @@ static int loadSingleUrlImage(struct iurlinfo * ui, char *path,
     int fd;
     int rc = 0;
     char *ehdrs = NULL;
+    long long size;
+    struct progressCBdata *data = NULL;
 
     if (ui->protocol == URL_METHOD_HTTP) {
         char *arch = getProductArch();
@@ -70,7 +72,7 @@ static int loadSingleUrlImage(struct iurlinfo * ui, char *path,
         }
     }
 
-    fd = urlinstStartTransfer(ui, path, ehdrs);
+    fd = urlinstStartTransfer(ui, path, ehdrs, &data, &size);
 
     if (fd == -2) {
         if (ehdrs) free (ehdrs);
@@ -89,10 +91,11 @@ static int loadSingleUrlImage(struct iurlinfo * ui, char *path,
     }
 
     if (dest != NULL) {
-        rc = copyFileAndLoopbackMount(fd, dest, device, mntpoint);
+        rc = copyFileAndLoopbackMount(fd, dest, device, mntpoint,
+                                      progressCallback, data, size);
     }
 
-    urlinstFinishTransfer(ui, fd);
+    urlinstFinishTransfer(ui, fd, &data);
     return rc;
 }
 
@@ -283,6 +286,8 @@ int getFileFromUrl(char * url, char * dest,
     int fd, rc;
     iface_t iface;
     char *ehdrs = NULL, *ip = NULL;
+    long long size;
+    struct progressCBdata *data = NULL;
 
     iface_init_iface_t(&iface);
 
@@ -366,14 +371,14 @@ int getFileFromUrl(char * url, char * dest,
         }
     }
 
-    fd = urlinstStartTransfer(&ui, file, ehdrs);
+    fd = urlinstStartTransfer(&ui, file, ehdrs, &data, &size);
     if (fd < 0) {
         logMessage(ERROR, "failed to retrieve http://%s/%s%s", ui.address, ui.prefix, file);
         retval = 1;
         goto err;
     }
 
-    rc = copyFileFd(fd, dest);
+    rc = copyFileFd(fd, dest, progressCallback, data, size);
     if (rc) {
         unlink (dest);
         logMessage(ERROR, "failed to copy file to %s", dest);
@@ -381,7 +386,7 @@ int getFileFromUrl(char * url, char * dest,
         goto err;
     }
 
-    urlinstFinishTransfer(&ui, fd);
+    urlinstFinishTransfer(&ui, fd, &data);
 
 err:
     if (file) free(file);
diff --git a/loader/urls.c b/loader/urls.c
index 3cd9a712b..a46e0342a 100644
--- a/loader/urls.c
+++ b/loader/urls.c
@@ -164,7 +164,8 @@ char *convertUIToURL(struct iurlinfo *ui) {
 /* see ftp.c:httpGetFileDesc() for details */
 /* set to NULL if not needed */
 int urlinstStartTransfer(struct iurlinfo * ui, char *path,
-                         char *extraHeaders) {
+                         char *extraHeaders, struct progressCBdata **data,
+                         long long *size) {
     int fd, port;
     int family = -1;
     struct in_addr addr;
@@ -214,14 +215,14 @@ int urlinstStartTransfer(struct iurlinfo * ui, char *path,
             return -2;
         }
 
-        fd = ftpGetFileDesc(ui->ftpPort, addr6, family, path);
+        fd = ftpGetFileDesc(ui->ftpPort, addr6, family, path, size);
         if (fd < 0) {
             close(ui->ftpPort);
             if (hostname) free(hostname);
             return -1;
         }
     } else {
-        fd = httpGetFileDesc(hostname, port, path, extraHeaders);
+        fd = httpGetFileDesc(hostname, port, path, extraHeaders, size);
         if (fd < 0) {
             if (portstr) free(portstr);
             return -1;
@@ -232,6 +233,9 @@ int urlinstStartTransfer(struct iurlinfo * ui, char *path,
 
     if (FL_CMDLINE(flags)) {
         printf("%s %s...\n", _("Retrieving"), fileName+1);
+    } else if (*size) {
+        *data = winProgressBar(70, 5, _("Retrieving"), "%s %s...", _("Retrieving"),
+                               fileName+1);
     } else {
         winStatus(70, 3, _("Retrieving"), "%s %s...", _("Retrieving"), fileName+1);
     }
@@ -240,7 +244,8 @@ int urlinstStartTransfer(struct iurlinfo * ui, char *path,
     return fd;
 }
 
-int urlinstFinishTransfer(struct iurlinfo * ui, int fd) {
+int urlinstFinishTransfer(struct iurlinfo * ui, int fd, struct progressCBdata **data) {
+    if (*data) free(*data);
     if (ui->protocol == URL_METHOD_FTP)
         close(ui->ftpPort);
     close(fd);
diff --git a/loader/urls.h b/loader/urls.h
index 75761ac76..c053edf4a 100644
--- a/loader/urls.h
+++ b/loader/urls.h
@@ -20,6 +20,8 @@
 #ifndef H_LOADER_URLS
 #define H_LOADER_URLS
 
+#include "windows.h"
+
 enum urlprotocol_t { URL_METHOD_FTP, URL_METHOD_HTTP };
 typedef enum urlprotocol_t urlprotocol;
 
@@ -40,7 +42,9 @@ char *convertUIToURL(struct iurlinfo *ui);
 int setupRemote(struct iurlinfo * ui);
 int urlMainSetupPanel(struct iurlinfo * ui);
 int urlSecondarySetupPanel(struct iurlinfo * ui);
-int urlinstStartTransfer(struct iurlinfo * ui, char *path, char *extraHeaders);
-int urlinstFinishTransfer(struct iurlinfo * ui, int fd);
+int urlinstStartTransfer(struct iurlinfo * ui, char *path, char *extraHeaders,
+                         struct progressCBdata **data, long long *size);
+int urlinstFinishTransfer(struct iurlinfo * ui, int fd,
+                          struct progressCBdata **data);
 
 #endif
diff --git a/loader/windows.c b/loader/windows.c
index 424f8a71e..6adfd3dcf 100644
--- a/loader/windows.c
+++ b/loader/windows.c
@@ -29,7 +29,9 @@
 #include <string.h>
 #include <stdio.h>
 #include <stdarg.h>
+#include <math.h>
 
+#include "log.h"
 #include "windows.h"
 
 void winStatus(int width, int height, char * title, char * text, ...) {
@@ -64,4 +66,53 @@ void scsiWindow(const char * driver) {
         _("Loading %s driver..."), driver);
 }
 
+void progressCallback(void *pbdata, long long pos, long long total) {
+    struct progressCBdata *data = pbdata;
+    char tickmark[2] = "-";
+    char *ticks = "-\\|/";
+
+    newtScaleSet(data->scale, ceil(pos * 100.0 / total));
+    *tickmark = ticks[(total / (pos + 1)) % 5];
+
+    newtLabelSetText(data->label, tickmark);
+    newtRefresh();
+}
+
+struct progressCBdata *winProgressBar(int width, int height, char *title, char *text, ...) {
+    struct progressCBdata *data;
+    char *buf = NULL;
+    va_list args;
+    int llen;
+    newtComponent t, f, scale, label;
+
+    va_start(args, text);
+
+    if (vasprintf(&buf, text, args) != -1) {
+        va_end(args);
+        newtCenteredWindow(width, height, title);
+        t = newtTextbox(1, 1, width - 2, height - 2, NEWT_TEXTBOX_WRAP);
+        newtTextboxSetText(t, buf);
+        llen = strlen(buf);
+        free(buf);
+        label = newtLabel(llen + 1, 1, "-");
+        f = newtForm(NULL, NULL, 0);
+        newtFormAddComponent(f, t);
+        scale = newtScale(3, 3, width - 6, 100);
+        newtFormAddComponent(f, scale);
+        newtDrawForm(f);
+        newtRefresh();
+
+        if ((data = malloc(sizeof(struct progressCBdata))) == NULL) {
+            logMessage(ERROR, "%s: %d: %m", __func__, __LINE__);
+            abort();
+        }
+
+        data->scale = scale;
+        data->label = label;
+        return data;
+    }
+
+    return NULL;
+}
+
 /* vim:set shiftwidth=4 softtabstop=4: */
diff --git a/loader/windows.h b/loader/windows.h
index 2400c9173..f635102f7 100644
--- a/loader/windows.h
+++ b/loader/windows.h
@@ -20,6 +20,8 @@
 #ifndef _WINDOWS_H_
 #define _WINDOWS_H_
 
+#include <newt.h>
+
 #include "lang.h"
 
 void winStatus(int width, int height, char * title, char * text, ...);
@@ -28,4 +30,14 @@ void scsiWindow(const char * driver);
 #define errorWindow(String) \
 	newtWinMessage(_("Error"), _("OK"), String, strerror (errno));
 
+typedef void (*progressCB) (void *pbdata, long long offset, long long total);
+
+struct progressCBdata {
+    newtComponent scale;
+    newtComponent label;
+};
+
+void progressCallback(void *pbdata, long long pos, long long total);
+struct progressCBdata *winProgressBar(int width, int height, char *title, char *text, ...);
+
 #endif /* _WINDOWS_H_ */