fastbootd: Enable erase and flash commands for physical partitions.

Bug: 78793464
Test: adb reboot fastboot && fastboot flashall

Change-Id: Ibe802c36f6efe20111a2315616ef34d3a027950f
diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index e9bb1d7..19f6390 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -94,6 +94,7 @@
     srcs: [
         "device/commands.cpp",
         "device/fastboot_device.cpp",
+        "device/flashing.cpp",
         "device/main.cpp",
         "device/usb_client.cpp",
         "device/utility.cpp",
diff --git a/fastboot/device/commands.cpp b/fastboot/device/commands.cpp
index 0e4a68b..7eaefe6 100644
--- a/fastboot/device/commands.cpp
+++ b/fastboot/device/commands.cpp
@@ -26,9 +26,11 @@
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <cutils/android_reboot.h>
+#include <ext4_utils/wipe.h>
 
 #include "constants.h"
 #include "fastboot_device.h"
+#include "flashing.h"
 #include "utility.h"
 
 using ::android::hardware::hidl_string;
@@ -51,7 +53,8 @@
             {FB_VAR_SLOT_COUNT, GetSlotCount},
             {FB_VAR_HAS_SLOT, GetHasSlot},
             {FB_VAR_SLOT_SUCCESSFUL, GetSlotSuccessful},
-            {FB_VAR_SLOT_UNBOOTABLE, GetSlotUnbootable}};
+            {FB_VAR_SLOT_UNBOOTABLE, GetSlotUnbootable},
+            {FB_VAR_PARTITION_SIZE, GetPartitionSize}};
 
     // args[0] is command name, args[1] is variable.
     auto found_variable = kVariableMap.find(args[1]);
@@ -63,6 +66,20 @@
     return found_variable->second(device, getvar_args);
 }
 
+bool EraseHandler(FastbootDevice* device, const std::vector<std::string>& args) {
+    if (args.size() < 2) {
+        return device->WriteStatus(FastbootResult::FAIL, "Invalid arguments");
+    }
+    PartitionHandle handle;
+    if (!OpenPartition(device, args[1], &handle)) {
+        return device->WriteStatus(FastbootResult::FAIL, "Partition doesn't exist");
+    }
+    if (wipe_block_device(handle.fd(), get_block_device_size(handle.fd())) == 0) {
+        return device->WriteStatus(FastbootResult::OKAY, "Erasing succeeded");
+    }
+    return device->WriteStatus(FastbootResult::FAIL, "Erasing failed");
+}
+
 bool DownloadHandler(FastbootDevice* device, const std::vector<std::string>& args) {
     if (args.size() < 2) {
         return device->WriteStatus(FastbootResult::FAIL, "size argument unspecified");
@@ -72,12 +89,12 @@
     if (!android::base::ParseUint("0x" + args[1], &size, UINT_MAX)) {
         return device->WriteStatus(FastbootResult::FAIL, "Invalid size");
     }
-    device->get_download_data().resize(size);
+    device->download_data().resize(size);
     if (!device->WriteStatus(FastbootResult::DATA, android::base::StringPrintf("%08x", size))) {
         return false;
     }
 
-    if (device->HandleData(true, &device->get_download_data())) {
+    if (device->HandleData(true, &device->download_data())) {
         return device->WriteStatus(FastbootResult::OKAY, "");
     }
 
@@ -85,6 +102,17 @@
     return device->WriteStatus(FastbootResult::FAIL, "Couldn't download data");
 }
 
+bool FlashHandler(FastbootDevice* device, const std::vector<std::string>& args) {
+    if (args.size() < 2) {
+        return device->WriteStatus(FastbootResult::FAIL, "Invalid arguments");
+    }
+    int ret = Flash(device, args[1]);
+    if (ret < 0) {
+        return device->WriteStatus(FastbootResult::FAIL, strerror(-ret));
+    }
+    return device->WriteStatus(FastbootResult::OKAY, "Flashing succeeded");
+}
+
 bool SetActiveHandler(FastbootDevice* device, const std::vector<std::string>& args) {
     if (args.size() < 2) {
         return device->WriteStatus(FastbootResult::FAIL, "Missing slot argument");
diff --git a/fastboot/device/commands.h b/fastboot/device/commands.h
index 8785b91..830eb55 100644
--- a/fastboot/device/commands.h
+++ b/fastboot/device/commands.h
@@ -39,3 +39,5 @@
 bool RebootFastbootHandler(FastbootDevice* device, const std::vector<std::string>& args);
 bool RebootRecoveryHandler(FastbootDevice* device, const std::vector<std::string>& args);
 bool GetVarHandler(FastbootDevice* device, const std::vector<std::string>& args);
+bool EraseHandler(FastbootDevice* device, const std::vector<std::string>& args);
+bool FlashHandler(FastbootDevice* device, const std::vector<std::string>& args);
diff --git a/fastboot/device/fastboot_device.cpp b/fastboot/device/fastboot_device.cpp
index a225bf8..b94fbb0 100644
--- a/fastboot/device/fastboot_device.cpp
+++ b/fastboot/device/fastboot_device.cpp
@@ -23,6 +23,7 @@
 #include <algorithm>
 
 #include "constants.h"
+#include "flashing.h"
 #include "usb_client.h"
 
 using ::android::hardware::hidl_string;
@@ -40,6 +41,8 @@
               {FB_CMD_REBOOT_BOOTLOADER, RebootBootloaderHandler},
               {FB_CMD_REBOOT_FASTBOOT, RebootFastbootHandler},
               {FB_CMD_REBOOT_RECOVERY, RebootRecoveryHandler},
+              {FB_CMD_ERASE, EraseHandler},
+              {FB_CMD_FLASH, FlashHandler},
       }),
       transport_(std::make_unique<ClientUsbTransport>()),
       boot_control_hal_(IBootControl::getService()) {}
diff --git a/fastboot/device/fastboot_device.h b/fastboot/device/fastboot_device.h
index 7517120..addc2ef 100644
--- a/fastboot/device/fastboot_device.h
+++ b/fastboot/device/fastboot_device.h
@@ -43,9 +43,7 @@
     bool WriteOkay(const std::string& message);
     bool WriteFail(const std::string& message);
 
-    std::vector<char>& get_download_data() { return download_data_; }
-    void set_upload_data(const std::vector<char>& data) { upload_data_ = data; }
-    void set_upload_data(std::vector<char>&& data) { upload_data_ = std::move(data); }
+    std::vector<char>& download_data() { return download_data_; }
     Transport* get_transport() { return transport_.get(); }
     android::sp<android::hardware::boot::V1_0::IBootControl> boot_control_hal() {
         return boot_control_hal_;
@@ -57,5 +55,4 @@
     std::unique_ptr<Transport> transport_;
     android::sp<android::hardware::boot::V1_0::IBootControl> boot_control_hal_;
     std::vector<char> download_data_;
-    std::vector<char> upload_data_;
 };
diff --git a/fastboot/device/flashing.cpp b/fastboot/device/flashing.cpp
new file mode 100644
index 0000000..d3dd82c
--- /dev/null
+++ b/fastboot/device/flashing.cpp
@@ -0,0 +1,102 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "flashing.h"
+
+#include <fcntl.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <memory>
+
+#include <android-base/logging.h>
+#include <android-base/strings.h>
+#include <ext4_utils/ext4_utils.h>
+#include <fs_mgr.h>
+#include <sparse/sparse.h>
+
+#include "fastboot_device.h"
+#include "utility.h"
+
+namespace {
+
+constexpr uint32_t SPARSE_HEADER_MAGIC = 0xed26ff3a;
+
+}  // namespace
+
+int FlashRawDataChunk(int fd, const char* data, size_t len) {
+    size_t ret = 0;
+    while (ret < len) {
+        int this_len = std::min(static_cast<size_t>(1048576UL * 8), len - ret);
+        int this_ret = write(fd, data, this_len);
+        if (this_ret < 0) {
+            PLOG(ERROR) << "Failed to flash data of len " << len;
+            return -1;
+        }
+        data += this_ret;
+        ret += this_ret;
+    }
+    return 0;
+}
+
+int FlashRawData(int fd, const std::vector<char>& downloaded_data) {
+    int ret = FlashRawDataChunk(fd, downloaded_data.data(), downloaded_data.size());
+    if (ret < 0) {
+        return -errno;
+    }
+    return ret;
+}
+
+int WriteCallback(void* priv, const void* data, size_t len) {
+    int fd = reinterpret_cast<long long>(priv);
+    if (!data) {
+        return lseek64(fd, len, SEEK_CUR) >= 0 ? 0 : -errno;
+    }
+    return FlashRawDataChunk(fd, reinterpret_cast<const char*>(data), len);
+}
+
+int FlashSparseData(int fd, std::vector<char>& downloaded_data) {
+    struct sparse_file* file = sparse_file_import_buf(downloaded_data.data(), true, false);
+    if (!file) {
+        return -ENOENT;
+    }
+    return sparse_file_callback(file, false, false, WriteCallback, reinterpret_cast<void*>(fd));
+}
+
+int FlashBlockDevice(int fd, std::vector<char>& downloaded_data) {
+    lseek64(fd, 0, SEEK_SET);
+    if (downloaded_data.size() >= sizeof(SPARSE_HEADER_MAGIC) &&
+        *reinterpret_cast<uint32_t*>(downloaded_data.data()) == SPARSE_HEADER_MAGIC) {
+        return FlashSparseData(fd, downloaded_data);
+    } else {
+        return FlashRawData(fd, downloaded_data);
+    }
+}
+
+int Flash(FastbootDevice* device, const std::string& partition_name) {
+    PartitionHandle handle;
+    if (!OpenPartition(device, partition_name, &handle)) {
+        return -ENOENT;
+    }
+
+    std::vector<char> data = std::move(device->download_data());
+    if (data.size() == 0) {
+        return -EINVAL;
+    } else if (data.size() > get_block_device_size(handle.fd())) {
+        return -EOVERFLOW;
+    }
+    return FlashBlockDevice(handle.fd(), data);
+}
diff --git a/fastboot/device/flashing.h b/fastboot/device/flashing.h
new file mode 100644
index 0000000..206a407
--- /dev/null
+++ b/fastboot/device/flashing.h
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+class FastbootDevice;
+
+int Flash(FastbootDevice* device, const std::string& partition_name);
diff --git a/fastboot/device/utility.cpp b/fastboot/device/utility.cpp
index c8d2b3e..73cf1bf 100644
--- a/fastboot/device/utility.cpp
+++ b/fastboot/device/utility.cpp
@@ -16,8 +16,45 @@
 
 #include "utility.h"
 
+#include <android-base/logging.h>
+
+#include "fastboot_device.h"
+
+using android::base::unique_fd;
 using android::hardware::boot::V1_0::Slot;
 
+static bool OpenPhysicalPartition(const std::string& name, PartitionHandle* handle) {
+    std::optional<std::string> path = FindPhysicalPartition(name);
+    if (!path) {
+        return false;
+    }
+    *handle = PartitionHandle(*path);
+    return true;
+}
+
+bool OpenPartition(FastbootDevice* /* device */, const std::string& name, PartitionHandle* handle) {
+    if (!OpenPhysicalPartition(name, handle)) {
+        LOG(ERROR) << "No such partition: " << name;
+        return false;
+    }
+
+    unique_fd fd(TEMP_FAILURE_RETRY(open(handle->path().c_str(), O_WRONLY | O_EXCL)));
+    if (fd < 0) {
+        PLOG(ERROR) << "Failed to open block device: " << handle->path();
+        return false;
+    }
+    handle->set_fd(std::move(fd));
+    return true;
+}
+
+std::optional<std::string> FindPhysicalPartition(const std::string& name) {
+    std::string path = "/dev/block/by-name/" + name;
+    if (access(path.c_str(), R_OK | W_OK) < 0) {
+        return {};
+    }
+    return path;
+}
+
 bool GetSlotNumber(const std::string& slot, Slot* number) {
     if (slot.size() != 1) {
         return false;
diff --git a/fastboot/device/utility.h b/fastboot/device/utility.h
index 867d693..26f486b 100644
--- a/fastboot/device/utility.h
+++ b/fastboot/device/utility.h
@@ -15,8 +15,32 @@
  */
 #pragma once
 
+#include <optional>
 #include <string>
 
+#include <android-base/unique_fd.h>
 #include <android/hardware/boot/1.0/IBootControl.h>
 
+// Logical partitions are only mapped to a block device as needed, and
+// immediately unmapped when no longer needed. In order to enforce this we
+// require accessing partitions through a Handle abstraction, which may perform
+// additional operations after closing its file descriptor.
+class PartitionHandle {
+  public:
+    PartitionHandle() {}
+    explicit PartitionHandle(const std::string& path) : path_(path) {}
+    const std::string& path() const { return path_; }
+    int fd() const { return fd_.get(); }
+    void set_fd(android::base::unique_fd&& fd) { fd_ = std::move(fd); }
+
+  private:
+    std::string path_;
+    android::base::unique_fd fd_;
+};
+
+class FastbootDevice;
+
+std::optional<std::string> FindPhysicalPartition(const std::string& name);
+bool OpenPartition(FastbootDevice* device, const std::string& name, PartitionHandle* handle);
+
 bool GetSlotNumber(const std::string& slot, android::hardware::boot::V1_0::Slot* number);
diff --git a/fastboot/device/variables.cpp b/fastboot/device/variables.cpp
index 33f7f74..8f66fea 100644
--- a/fastboot/device/variables.cpp
+++ b/fastboot/device/variables.cpp
@@ -16,6 +16,8 @@
 
 #include "variables.h"
 
+#include <inttypes.h>
+
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/properties.h>
@@ -24,6 +26,7 @@
 #include <ext4_utils/ext4_utils.h>
 
 #include "fastboot_device.h"
+#include "flashing.h"
 #include "utility.h"
 
 using ::android::hardware::boot::V1_0::BoolResult;
@@ -125,3 +128,15 @@
     std::string result = (args[0] == "userdata" ? "no" : "yes");
     return device->WriteOkay(result);
 }
+
+bool GetPartitionSize(FastbootDevice* device, const std::vector<std::string>& args) {
+    if (args.size() < 1) {
+        return device->WriteFail("Missing argument");
+    }
+    PartitionHandle handle;
+    if (!OpenPartition(device, args[0], &handle)) {
+        return device->WriteFail("Could not open partition");
+    }
+    uint64_t size = get_block_device_size(handle.fd());
+    return device->WriteOkay(android::base::StringPrintf("%" PRIX64, size));
+}
diff --git a/fastboot/device/variables.h b/fastboot/device/variables.h
index 45c6dc9..88947e0 100644
--- a/fastboot/device/variables.h
+++ b/fastboot/device/variables.h
@@ -34,3 +34,4 @@
 bool GetMaxDownloadSize(FastbootDevice* device, const std::vector<std::string>& args);
 bool GetUnlocked(FastbootDevice* device, const std::vector<std::string>& args);
 bool GetHasSlot(FastbootDevice* device, const std::vector<std::string>& args);
+bool GetPartitionSize(FastbootDevice* device, const std::vector<std::string>& args);