1 files changed, 103 insertions, 22 deletions

diff --git a/fastboot/protocol.cpp b/fastboot/protocol.cpp
index bfa83b041..dcdf8f0c1 100644
--- a/fastboot/protocol.cpp
+++ b/fastboot/protocol.cpp
@@ -29,26 +29,34 @@
 #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>
 
 #include "fastboot.h"
 #include "transport.h"
 
 static std::string g_error;
 
+using android::base::unique_fd;
+using android::base::WriteStringToFile;
+
 const std::string fb_get_error() {
     return g_error;
 }
 
-static int check_response(Transport* transport, uint32_t size, char* response) {
+static int64_t check_response(Transport* transport, uint32_t size, char* response) {
     char status[65];
 
     while (true) {
@@ -105,7 +113,7 @@ static int check_response(Transport* transport, uint32_t size, char* response) {
     return -1;
 }
 
-static int _command_start(Transport* transport, const char* cmd, uint32_t size, char* response) {
+static int64_t _command_start(Transport* transport, const char* cmd, uint32_t size, char* response) {
     size_t cmdsize = strlen(cmd);
     if (cmdsize > 64) {
         g_error = android::base::StringPrintf("command too large (%zu)", cmdsize);
@@ -125,37 +133,51 @@ static int _command_start(Transport* transport, const char* cmd, uint32_t size,
     return check_response(transport, size, response);
 }
 
-static int _command_data(Transport* transport, const void* data, uint32_t size) {
-    int r = transport->Write(data, 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 write failure (%s)", strerror(errno));
+        transport->Close();
+        return -1;
+    }
+    if (r != static_cast<int64_t>(size)) {
+        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 transfer failure (%s)", strerror(errno));
+        g_error = android::base::StringPrintf("data read failure (%s)", strerror(errno));
         transport->Close();
         return -1;
     }
-    if (r != ((int) size)) {
-        g_error = "data transfer failure (short transfer)";
+    if (r != (static_cast<int64_t>(size))) {
+        g_error = "data read failure (short transfer)";
         transport->Close();
         return -1;
     }
     return r;
 }
 
-static int _command_end(Transport* transport) {
+static int64_t _command_end(Transport* transport) {
     return check_response(transport, 0, 0) < 0 ? -1 : 0;
 }
 
-static int _command_send(Transport* transport, const char* cmd, const void* data, uint32_t size,
+static int64_t _command_send(Transport* transport, const char* cmd, const void* data, uint32_t size,
                          char* response) {
     if (size == 0) {
         return -1;
     }
 
-    int r = _command_start(transport, cmd, size, response);
+    int64_t r = _command_start(transport, cmd, size, response);
     if (r < 0) {
         return -1;
     }
-
-    r = _command_data(transport, data, size);
+    r = _command_write_data(transport, data, size);
     if (r < 0) {
         return -1;
     }
@@ -168,6 +190,39 @@ static int _command_send(Transport* transport, const char* cmd, const void* data
     return size;
 }
 
+static int64_t _command_send_fd(Transport* transport, const char* cmd, int fd, uint32_t size,
+                                char* response) {
+    static constexpr uint32_t MAX_MAP_SIZE = 512 * 1024 * 1024;
+    off64_t offset = 0;
+    uint32_t remaining = size;
+
+    if (_command_start(transport, cmd, size, response) < 0) {
+        return -1;
+    }
+
+    while (remaining) {
+        android::FileMap filemap;
+        size_t len = std::min(remaining, MAX_MAP_SIZE);
+
+        if (!filemap.create(NULL, fd, offset, len, true)) {
+            return -1;
+        }
+
+        if (_command_write_data(transport, filemap.getDataPtr(), len) < 0) {
+            return -1;
+        }
+
+        remaining -= len;
+        offset += len;
+    }
+
+    if (_command_end(transport) < 0) {
+        return -1;
+    }
+
+    return size;
+}
+
 static int _command_send_no_data(Transport* transport, const char* cmd, char* response) {
     return _command_start(transport, cmd, 0, response);
 }
@@ -180,10 +235,36 @@ int fb_command_response(Transport* transport, const char* cmd, char* response) {
     return _command_send_no_data(transport, cmd, response);
 }
 
-int fb_download_data(Transport* transport, const void* data, uint32_t size) {
-    char cmd[64];
-    snprintf(cmd, sizeof(cmd), "download:%08x", size);
-    return _command_send(transport, cmd, data, size, 0) < 0 ? -1 : 0;
+int64_t fb_download_data(Transport* transport, const void* data, uint32_t size) {
+    std::string cmd(android::base::StringPrintf("download:%08x", size));
+    return _command_send(transport, cmd.c_str(), data, size, 0) < 0 ? -1 : 0;
+}
+
+int64_t fb_download_data_fd(Transport* transport, int fd, uint32_t size) {
+    std::string cmd(android::base::StringPrintf("download:%08x", 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
@@ -207,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;
         }
@@ -220,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;
         }
@@ -242,7 +323,8 @@ 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) {
-        if (_command_data(transport, transport_buf, transport_buf_len) != 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;
         }
         transport_buf_len = 0;
@@ -256,9 +338,8 @@ int fb_download_data_sparse(Transport* transport, struct sparse_file* s) {
         return -1;
     }
 
-    char cmd[64];
-    snprintf(cmd, sizeof(cmd), "download:%08x", size);
-    int r = _command_start(transport, cmd, size, 0);
+    std::string cmd(android::base::StringPrintf("download:%08x", size));
+    int r = _command_start(transport, cmd.c_str(), size, 0);
     if (r < 0) {
         return -1;
     }