commit 17d469bd579a83f5cad372dd8b1db47d3e874830
author Yifan Hong <elsk@google.com> Thu Feb 18 15:15:46 2021 -0800
committer Yifan Hong <elsk@google.com> Wed Mar 17 15:38:06 2021 -0700
tree 8b081174d97ba4b0a8fb12e7fa187d7c499a7e00
parent bbf374dbd7a9d1e8090d8a9cabcdaf181dee615c

fastboot driver: RunAndReadBuffer don't allocate too much mem

Instead of allocating the buffer for the whole
upload (which can be arbitrary number of bytes as
the device determines), read 1 MiB at a time.

Test: pass
Bug: 173654501
Change-Id: Ib601b0341b10b7dccbb429cd21aad86a2d3bfda8

fastboot/fastboot_driver.cpp

diff --git a/fastboot/fastboot_driver.cpp b/fastboot/fastboot_driver.cpp
index 79305c4..a516a91 100644
--- a/fastboot/fastboot_driver.cpp
+++ b/fastboot/fastboot_driver.cpp
@@ -45,11 +45,13 @@
 #include <android-base/stringprintf.h>
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
+#include <storage_literals/storage_literals.h>
 
 #include "constants.h"
 #include "transport.h"
 
 using android::base::StringPrintf;
+using namespace android::storage_literals;
 
 namespace fastboot {
 
@@ -316,12 +318,20 @@
         return BAD_DEV_RESP;
     }
 
-    std::vector<char> data(dsize);
-    if ((ret = ReadBuffer(data.data(), data.size())) != SUCCESS) {
-        return ret;
-    }
-    if ((ret = write_fn(data.data(), data.size())) != SUCCESS) {
-        return ret;
+    const uint64_t total_size = dsize;
+    const uint64_t buf_size = std::min<uint64_t>(total_size, 1_MiB);
+    std::vector<char> data(buf_size);
+    uint64_t current_offset = 0;
+    while (current_offset < total_size) {
+        uint64_t remaining = total_size - current_offset;
+        uint64_t chunk_size = std::min(buf_size, remaining);
+        if ((ret = ReadBuffer(data.data(), chunk_size)) != SUCCESS) {
+            return ret;
+        }
+        if ((ret = write_fn(data.data(), chunk_size)) != SUCCESS) {
+            return ret;
+        }
+        current_offset += chunk_size;
     }
     return HandleResponse(response, info);
 }