fastboot: Add 'get_staged' command

(cherry-picked from internal nyc-iot-dev to AOSP)

New user-level command usage:

 * fastboot get_staged <outfile>
   Reads staged data from the last command handled by the device. If the
   last command did not result in staged data, this command will fail.

This enables data staged by OEM commands to be transferred from device
to host. get_staged wraps new device command "upload". Fastboot
clients are not required to support "upload", so get_staged won't
work on all devices.

Bug: 36002804
Test: Implemented "upload" in fastboot on imx6ul. Verified that uploading
      ~100K data from the device works.
Change-Id: I5b1a1ce023f362062505ee62746ea8ab6f36bfbf
(cherry-picked from commit 83a875de994bf48f0faa2a8a23ceb0b8f52b6b04)

5 files changed, 83 insertions, 9 deletions

diff --git a/fastboot/README.md b/fastboot/README.md
index 022d34b91..ec7dcb4bc 100644
--- a/fastboot/README.md
+++ b/fastboot/README.md
@@ -126,6 +126,16 @@ The various currently defined commands are:
                        space in RAM or "FAIL" if not.  The size of
                        the download is remembered.
 
+    upload             Read data from memory which was staged by the last
+                       command, e.g. an oem command.  The client will reply
+                       with "DATA%08x" if it is ready to send %08x bytes of
+                       data.  If no data was staged in the last command,
+                       the client must reply with "FAIL".  After the client
+                       successfully sends %08x bytes, the client shall send
+                       a single packet starting with "OKAY".  Clients
+                       should not support "upload" unless it supports an
+                       oem command that requires "upload" capabilities.
+
     verify:%08x        Send a digital signature to verify the downloaded
                        data.  Required if the bootloader is "secure"
                        otherwise "flash" and "boot" will be ignored.
diff --git a/fastboot/engine.cpp b/fastboot/engine.cpp
index 4cf89ca85..56ee81637 100644
--- a/fastboot/engine.cpp
+++ b/fastboot/engine.cpp
@@ -45,6 +45,7 @@
 #define OP_DOWNLOAD_SPARSE 5
 #define OP_WAIT_FOR_DISCONNECT 6
 #define OP_DOWNLOAD_FD 7
+#define OP_UPLOAD 8
 
 typedef struct Action Action;
 
@@ -341,6 +342,13 @@ void fb_queue_download_fd(const char *name, int fd, uint32_t sz)
     a->msg = mkmsg("sending '%s' (%d KB)", name, sz / 1024);
 }
 
+void fb_queue_upload(char *outfile)
+{
+    Action *a = queue_action(OP_UPLOAD, "");
+    a->data = outfile;
+    a->msg = mkmsg("uploading '%s'", outfile);
+}
+
 void fb_queue_notice(const char *notice)
 {
     Action *a = queue_action(OP_NOTICE, "");
@@ -395,6 +403,9 @@ int64_t fb_execute_queue(Transport* transport)
             if (status) break;
         } else if (a->op == OP_WAIT_FOR_DISCONNECT) {
             transport->WaitForDisconnect();
+        } else if (a->op == OP_UPLOAD) {
+            status = fb_upload_data(transport, reinterpret_cast<char*>(a->data));
+            status = a->func(a, status, status ? fb_get_error().c_str() : "");
         } else {
             die("bogus action");
         }
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index e5ba93a59..0d59901b6 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -375,6 +375,9 @@ static void usage() {
             "  stage <infile>                           Sends contents of <infile> to stage for\n"
             "                                           the next command. Supported only on\n"
             "                                           Android Things devices.\n"
+            "  get_staged <outfile>                     Receives data to <outfile> staged by the\n"
+            "                                           last command. Supported only on Android\n"
+            "                                           Things devices.\n"
             "  help                                     Show this help message.\n"
             "\n"
             "options:\n"
@@ -1819,6 +1822,11 @@ int main(int argc, char **argv)
                 die("cannot load '%s'", infile.c_str());
             }
             fb_queue_download_fd(infile.c_str(), buf.fd, buf.sz);
+        } else if(!strcmp(*argv, "get_staged")) {
+            require(2);
+            char *outfile = argv[1];
+            skip(2);
+            fb_queue_upload(outfile);
         } else if(!strcmp(*argv, "oem")) {
             argc = do_oem_command(argc, argv);
         } else if(!strcmp(*argv, "flashing")) {
diff --git a/fastboot/fastboot.h b/fastboot/fastboot.h
index 5e9dba948..e30c6ded3 100644
--- a/fastboot/fastboot.h
+++ b/fastboot/fastboot.h
@@ -44,6 +44,7 @@ int fb_command_response(Transport* transport, const char* cmd, char* response);
 int64_t fb_download_data(Transport* transport, const void* data, uint32_t size);
 int64_t fb_download_data_fd(Transport* transport, int fd, uint32_t size);
 int fb_download_data_sparse(Transport* transport, struct sparse_file* s);
+int64_t fb_upload_data(Transport* transport, const char* outfile);
 const std::string fb_get_error();
 
 #define FB_COMMAND_SZ 64
@@ -65,6 +66,7 @@ void fb_queue_reboot(void);
 void fb_queue_command(const char *cmd, const char *msg);
 void fb_queue_download(const char *name, void *data, uint32_t size);
 void fb_queue_download_fd(const char *name, int fd, uint32_t sz);
+void fb_queue_upload(char *outfile);
 void fb_queue_notice(const char *notice);
 void fb_queue_wait_for_disconnect(void);
 int64_t fb_execute_queue(Transport* transport);
diff --git a/fastboot/protocol.cpp b/fastboot/protocol.cpp
index 334f81ffb..dcdf8f0c1 100644
--- a/fastboot/protocol.cpp
+++ b/fastboot/protocol.cpp
@@ -29,14 +29,18 @@
 #define round_down(a, b) \
     ({ typeof(a) _a = (a); typeof(b) _b = (b); _a - (_a % _b); })
 
+#include <fcntl.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 #include <errno.h>
 
 #include <algorithm>
+#include <vector>
 
+#include <android-base/file.h>
 #include <android-base/stringprintf.h>
+#include <android-base/unique_fd.h>
 #include <sparse/sparse.h>
 #include <utils/FileMap.h>
 
@@ -45,6 +49,9 @@
 
 static std::string g_error;
 
+using android::base::unique_fd;
+using android::base::WriteStringToFile;
+
 const std::string fb_get_error() {
     return g_error;
 }
@@ -126,15 +133,30 @@ static int64_t _command_start(Transport* transport, const char* cmd, uint32_t si
     return check_response(transport, size, response);
 }
 
-static int64_t _command_data(Transport* transport, const void* data, uint32_t size) {
+static int64_t _command_write_data(Transport* transport, const void* data, uint32_t size) {
     int64_t r = transport->Write(data, size);
     if (r < 0) {
-        g_error = android::base::StringPrintf("data transfer failure (%s)", strerror(errno));
+        g_error = android::base::StringPrintf("data write failure (%s)", strerror(errno));
         transport->Close();
         return -1;
     }
     if (r != static_cast<int64_t>(size)) {
-        g_error = "data transfer failure (short transfer)";
+        g_error = "data write failure (short transfer)";
+        transport->Close();
+        return -1;
+    }
+    return r;
+}
+
+static int64_t _command_read_data(Transport* transport, void* data, uint32_t size) {
+    int64_t r = transport->Read(data, size);
+    if (r < 0) {
+        g_error = android::base::StringPrintf("data read failure (%s)", strerror(errno));
+        transport->Close();
+        return -1;
+    }
+    if (r != (static_cast<int64_t>(size))) {
+        g_error = "data read failure (short transfer)";
         transport->Close();
         return -1;
     }
@@ -155,8 +177,7 @@ static int64_t _command_send(Transport* transport, const char* cmd, const void*
     if (r < 0) {
         return -1;
     }
-
-    r = _command_data(transport, data, size);
+    r = _command_write_data(transport, data, size);
     if (r < 0) {
         return -1;
     }
@@ -187,7 +208,7 @@ static int64_t _command_send_fd(Transport* transport, const char* cmd, int fd, u
             return -1;
         }
 
-        if (_command_data(transport, filemap.getDataPtr(), len) < 0) {
+        if (_command_write_data(transport, filemap.getDataPtr(), len) < 0) {
             return -1;
         }
 
@@ -224,6 +245,28 @@ int64_t fb_download_data_fd(Transport* transport, int fd, uint32_t size) {
     return _command_send_fd(transport, cmd.c_str(), fd, size, 0) < 0 ? -1 : 0;
 }
 
+int64_t fb_upload_data(Transport* transport, const char* outfile) {
+    // positive return value is the upload size sent by the device
+    int64_t r = _command_start(transport, "upload", std::numeric_limits<int32_t>::max(), nullptr);
+    if (r <= 0) {
+        g_error = android::base::StringPrintf("command start failed (%s)", strerror(errno));
+        return r;
+    }
+
+    std::string data;
+    data.resize(r);
+    if ((r = _command_read_data(transport, &data[0], data.size())) == -1) {
+        return r;
+    }
+
+    if (!WriteStringToFile(data, outfile, true)) {
+        g_error = android::base::StringPrintf("write to '%s' failed", outfile);
+        return -1;
+    }
+
+    return _command_end(transport);
+}
+
 #define TRANSPORT_BUF_SIZE 1024
 static char transport_buf[TRANSPORT_BUF_SIZE];
 static int transport_buf_len;
@@ -245,7 +288,7 @@ static int fb_download_data_sparse_write(void *priv, const void *data, int len)
     }
 
     if (transport_buf_len == TRANSPORT_BUF_SIZE) {
-        r = _command_data(transport, transport_buf, TRANSPORT_BUF_SIZE);
+        r = _command_write_data(transport, transport_buf, TRANSPORT_BUF_SIZE);
         if (r != TRANSPORT_BUF_SIZE) {
             return -1;
         }
@@ -258,7 +301,7 @@ static int fb_download_data_sparse_write(void *priv, const void *data, int len)
             return -1;
         }
         to_write = round_down(len, TRANSPORT_BUF_SIZE);
-        r = _command_data(transport, ptr, to_write);
+        r = _command_write_data(transport, ptr, to_write);
         if (r != to_write) {
             return -1;
         }
@@ -280,7 +323,7 @@ static int fb_download_data_sparse_write(void *priv, const void *data, int len)
 
 static int fb_download_data_sparse_flush(Transport* transport) {
     if (transport_buf_len > 0) {
-        int64_t r = _command_data(transport, transport_buf, transport_buf_len);
+        int64_t r = _command_write_data(transport, transport_buf, transport_buf_len);
         if (r != static_cast<int64_t>(transport_buf_len)) {
             return -1;
         }