Merge "Add support to update the DTB when flashing a vendor_boot ramdisk" into main
diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index bfe0768..b61fbd4 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -430,6 +430,7 @@
     ],
     data: [
         ":fastboot_test_dtb",
+        ":fastboot_test_dtb_replace",
         ":fastboot_test_bootconfig",
         ":fastboot_test_vendor_ramdisk_none",
         ":fastboot_test_vendor_ramdisk_platform",
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 6b9e493..156dc3b 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -552,6 +552,12 @@
             "                            Secondary images may be flashed to inactive slot.\n"
             " flash PARTITION [FILENAME] Flash given partition, using the image from\n"
             "                            $ANDROID_PRODUCT_OUT if no filename is given.\n"
+            " flash vendor_boot:RAMDISK [FILENAME]\n"
+            "                            Flash vendor_boot ramdisk, fetching the existing\n"
+            "                            vendor_boot image and repackaging it with the new\n"
+            "                            ramdisk.\n"
+            " --dtb DTB                  If set with flash vendor_boot:RAMDISK, then\n"
+            "                            update the vendor_boot image with provided DTB.\n"
             "\n"
             "basics:\n"
             " devices [-l]               List devices in bootloader (-l: with device paths).\n"
@@ -1020,6 +1026,8 @@
 }
 
 int64_t get_sparse_limit(int64_t size, const FlashingPlan* fp) {
+    if (!fp) return 0;
+
     int64_t limit = int64_t(fp->sparse_limit);
     if (limit == 0) {
         // Unlimited, so see what the target device's limit is.
@@ -1465,6 +1473,7 @@
 static std::string repack_ramdisk(const char* pname, struct fastboot_buffer* buf,
                                   fastboot::IFastBootDriver* fb) {
     std::string_view pname_sv{pname};
+    struct fastboot_buffer dtb_buf = {.sz = 0, .fd = unique_fd(-1)};
 
     if (!android::base::StartsWith(pname_sv, "vendor_boot:") &&
         !android::base::StartsWith(pname_sv, "vendor_boot_a:") &&
@@ -1480,10 +1489,25 @@
     std::string partition(pname_sv.substr(0, pname_sv.find(':')));
     std::string ramdisk(pname_sv.substr(pname_sv.find(':') + 1));
 
+    if (!g_dtb_path.empty()) {
+        if (!load_buf(g_dtb_path.c_str(), &dtb_buf, nullptr)) {
+            die("cannot load '%s': %s", g_dtb_path.c_str(), strerror(errno));
+        }
+
+        if (dtb_buf.type != FB_BUFFER_FD) {
+            die("Flashing sparse vendor ramdisk image with dtb is not supported.");
+        }
+        if (dtb_buf.sz <= 0) {
+            die("repack_ramdisk() sees invalid dtb size: %" PRId64, buf->sz);
+        }
+        verbose("Updating DTB with %s", pname_sv.data());
+    }
+
     unique_fd vendor_boot(make_temporary_fd("vendor boot repack"));
     uint64_t vendor_boot_size = fetch_partition(partition, vendor_boot, fb);
     auto repack_res = replace_vendor_ramdisk(vendor_boot, vendor_boot_size, ramdisk, buf->fd,
-                                             static_cast<uint64_t>(buf->sz));
+                                             static_cast<uint64_t>(buf->sz), dtb_buf.fd,
+                                             static_cast<uint64_t>(dtb_buf.sz));
     if (!repack_res.ok()) {
         die("%s", repack_res.error().message().c_str());
     }
diff --git a/fastboot/fuzzer/fastboot_fuzzer.cpp b/fastboot/fuzzer/fastboot_fuzzer.cpp
index 60940fe..4594a8a 100644
--- a/fastboot/fuzzer/fastboot_fuzzer.cpp
+++ b/fastboot/fuzzer/fastboot_fuzzer.cpp
@@ -15,6 +15,7 @@
  *
  */
 #include <android-base/file.h>
+#include <android-base/unique_fd.h>
 #include "fastboot.h"
 #include "socket.h"
 #include "socket_mock_fuzz.h"
@@ -25,6 +26,7 @@
 #include <fuzzer/FuzzedDataProvider.h>
 
 using namespace std;
+using android::base::unique_fd;
 
 const size_t kYearMin = 2000;
 const size_t kYearMax = 2127;
@@ -255,7 +257,7 @@
     uint64_t ramdisk_size =
             fdp_->ConsumeBool() ? content_ramdisk_fd.size() : fdp_->ConsumeIntegral<uint64_t>();
     (void)replace_vendor_ramdisk(vendor_boot_fd, vendor_boot_size, ramdisk_name, ramdisk_fd,
-                                 ramdisk_size);
+                                 ramdisk_size, unique_fd(-1), 0);
     close(vendor_boot_fd);
     close(ramdisk_fd);
 }
diff --git a/fastboot/testdata/Android.bp b/fastboot/testdata/Android.bp
index a490fe2..47bf095 100644
--- a/fastboot/testdata/Android.bp
+++ b/fastboot/testdata/Android.bp
@@ -40,6 +40,14 @@
     cmd: "$(location fastboot_gen_rand) --seed dtb --length 1024 > $(out)",
 }
 
+// Fake dtb image for replacement.
+genrule {
+    name: "fastboot_test_dtb_replace",
+    defaults: ["fastboot_test_data_gen_defaults"],
+    out: ["dtb_replace.img"],
+    cmd: "$(location fastboot_gen_rand) --seed dtb --length 2048 > $(out)",
+}
+
 // Fake bootconfig image.
 genrule {
     name: "fastboot_test_bootconfig",
diff --git a/fastboot/vendor_boot_img_utils.cpp b/fastboot/vendor_boot_img_utils.cpp
index 9f05253..da547f1 100644
--- a/fastboot/vendor_boot_img_utils.cpp
+++ b/fastboot/vendor_boot_img_utils.cpp
@@ -209,7 +209,8 @@
 
 // Replace the vendor ramdisk as a whole.
 [[nodiscard]] Result<std::string> replace_default_vendor_ramdisk(const std::string& vendor_boot,
-                                                                 const std::string& new_ramdisk) {
+                                                                 const std::string& new_ramdisk,
+                                                                 const std::string& new_dtb) {
     if (auto res = check_vendor_boot_hdr(vendor_boot, 3); !res.ok()) return res.error();
     auto hdr = reinterpret_cast<const vendor_boot_img_hdr_v3*>(vendor_boot.data());
     auto hdr_size = get_vendor_boot_header_size(hdr);
@@ -244,8 +245,19 @@
         return res.error();
     if (auto res = updater.CheckOffset(o + p, o + new_p); !res.ok()) return res.error();
 
-    // Copy DTB (Q bytes).
-    if (auto res = updater.Copy(q); !res.ok()) return res.error();
+    // Copy DTB (Q bytes). Replace if a new one was provided.
+    new_hdr->dtb_size = !new_dtb.empty() ? new_dtb.size() : hdr->dtb_size;
+    const uint32_t new_q = round_up(new_hdr->dtb_size, new_hdr->page_size);
+    if (new_dtb.empty()) {
+        if (auto res = updater.Copy(q); !res.ok()) return res.error();
+    } else {
+        if (auto res = updater.Replace(hdr->dtb_size, new_dtb); !res.ok()) return res.error();
+        if (auto res = updater.Skip(q - hdr->dtb_size, new_q - new_hdr->dtb_size); !res.ok())
+            return res.error();
+    }
+    if (auto res = updater.CheckOffset(o + p + q, o + new_p + new_q); !res.ok()) {
+        return res.error();
+    }
 
     if (new_hdr->header_version >= 4) {
         auto hdr_v4 = static_cast<const vendor_boot_img_hdr_v4*>(hdr);
@@ -256,7 +268,7 @@
         auto new_hdr_v4 = static_cast<const vendor_boot_img_hdr_v4*>(new_hdr);
         auto new_r = round_up(new_hdr_v4->vendor_ramdisk_table_size, new_hdr->page_size);
         if (auto res = updater.Skip(r, new_r); !res.ok()) return res.error();
-        if (auto res = updater.CheckOffset(o + p + q + r, o + new_p + q + new_r); !res.ok())
+        if (auto res = updater.CheckOffset(o + p + q + r, o + new_p + new_q + new_r); !res.ok())
             return res.error();
 
         // Replace table with single entry representing the full ramdisk.
@@ -303,7 +315,8 @@
 // replace it with the content of |new_ramdisk|.
 [[nodiscard]] Result<std::string> replace_vendor_ramdisk_fragment(const std::string& ramdisk_name,
                                                                   const std::string& vendor_boot,
-                                                                  const std::string& new_ramdisk) {
+                                                                  const std::string& new_ramdisk,
+                                                                  const std::string& new_dtb) {
     if (auto res = check_vendor_boot_hdr(vendor_boot, 4); !res.ok()) return res.error();
     auto hdr = reinterpret_cast<const vendor_boot_img_hdr_v4*>(vendor_boot.data());
     auto hdr_size = get_vendor_boot_header_size(hdr);
@@ -368,8 +381,19 @@
         return res.error();
     if (auto res = updater.CheckOffset(o + p, o + new_p); !res.ok()) return res.error();
 
-    // Copy DTB (Q bytes).
-    if (auto res = updater.Copy(q); !res.ok()) return res.error();
+    // Copy DTB (Q bytes). Replace if a new one was provided.
+    new_hdr->dtb_size = !new_dtb.empty() ? new_dtb.size() : hdr->dtb_size;
+    const uint32_t new_q = round_up(new_hdr->dtb_size, new_hdr->page_size);
+    if (new_dtb.empty()) {
+        if (auto res = updater.Copy(q); !res.ok()) return res.error();
+    } else {
+        if (auto res = updater.Replace(hdr->dtb_size, new_dtb); !res.ok()) return res.error();
+        if (auto res = updater.Skip(q - hdr->dtb_size, new_q - new_hdr->dtb_size); !res.ok())
+            return res.error();
+    }
+    if (auto res = updater.CheckOffset(o + p + q, o + new_p + new_q); !res.ok()) {
+        return res.error();
+    }
 
     // Copy table, but with corresponding entries modified, including:
     // - ramdisk_size of the entry replaced
@@ -392,7 +416,7 @@
                                             hdr->vendor_ramdisk_table_entry_size);
         !res.ok())
         return res.error();
-    if (auto res = updater.CheckOffset(o + p + q + r, o + new_p + q + r); !res.ok())
+    if (auto res = updater.CheckOffset(o + p + q + r, o + new_p + new_q + r); !res.ok())
         return res.error();
 
     // Copy bootconfig (S bytes).
@@ -404,11 +428,11 @@
 
 }  // namespace
 
-[[nodiscard]] Result<void> replace_vendor_ramdisk(android::base::borrowed_fd vendor_boot_fd,
-                                                  uint64_t vendor_boot_size,
-                                                  const std::string& ramdisk_name,
-                                                  android::base::borrowed_fd new_ramdisk_fd,
-                                                  uint64_t new_ramdisk_size) {
+[[nodiscard]] Result<void> replace_vendor_ramdisk(
+        android::base::borrowed_fd vendor_boot_fd, uint64_t vendor_boot_size,
+        const std::string& ramdisk_name, android::base::borrowed_fd new_ramdisk_fd,
+        uint64_t new_ramdisk_size, android::base::borrowed_fd new_dtb_fd, uint64_t new_dtb_size) {
+    Result<std::string> new_dtb = {""};
     if (new_ramdisk_size > std::numeric_limits<uint32_t>::max()) {
         return Errorf("New vendor ramdisk is too big");
     }
@@ -417,12 +441,17 @@
     if (!vendor_boot.ok()) return vendor_boot.error();
     auto new_ramdisk = load_file(new_ramdisk_fd, new_ramdisk_size, "new vendor ramdisk");
     if (!new_ramdisk.ok()) return new_ramdisk.error();
+    if (new_dtb_size > 0 && new_dtb_fd >= 0) {
+        new_dtb = load_file(new_dtb_fd, new_dtb_size, "new dtb");
+        if (!new_dtb.ok()) return new_dtb.error();
+    }
 
     Result<std::string> new_vendor_boot;
     if (ramdisk_name == "default") {
-        new_vendor_boot = replace_default_vendor_ramdisk(*vendor_boot, *new_ramdisk);
+        new_vendor_boot = replace_default_vendor_ramdisk(*vendor_boot, *new_ramdisk, *new_dtb);
     } else {
-        new_vendor_boot = replace_vendor_ramdisk_fragment(ramdisk_name, *vendor_boot, *new_ramdisk);
+        new_vendor_boot =
+                replace_vendor_ramdisk_fragment(ramdisk_name, *vendor_boot, *new_ramdisk, *new_dtb);
     }
     if (!new_vendor_boot.ok()) return new_vendor_boot.error();
     if (auto res = store_file(vendor_boot_fd, *new_vendor_boot, "new vendor boot image"); !res.ok())
diff --git a/fastboot/vendor_boot_img_utils.h b/fastboot/vendor_boot_img_utils.h
index 0b702bc..0ca78da 100644
--- a/fastboot/vendor_boot_img_utils.h
+++ b/fastboot/vendor_boot_img_utils.h
@@ -31,4 +31,4 @@
 [[nodiscard]] android::base::Result<void> replace_vendor_ramdisk(
         android::base::borrowed_fd vendor_boot_fd, uint64_t vendor_boot_size,
         const std::string& ramdisk_name, android::base::borrowed_fd new_ramdisk_fd,
-        uint64_t new_ramdisk_size);
+        uint64_t new_ramdisk_size, android::base::borrowed_fd new_dtb_fd, uint64_t new_dtb_size);
diff --git a/fastboot/vendor_boot_img_utils_test.cpp b/fastboot/vendor_boot_img_utils_test.cpp
index 8107270..841e532 100644
--- a/fastboot/vendor_boot_img_utils_test.cpp
+++ b/fastboot/vendor_boot_img_utils_test.cpp
@@ -241,6 +241,7 @@
 
 struct RepackVendorBootImgTestParam {
     std::string vendor_boot_file_name;
+    std::string dtb_file_name;
     uint32_t expected_header_version;
     friend std::ostream& operator<<(std::ostream& os, const RepackVendorBootImgTestParam& param) {
         return os << param.vendor_boot_file_name;
@@ -252,22 +253,50 @@
     virtual void SetUp() {
         vboot = std::make_unique<ReadWriteTestFileHandle>(GetParam().vendor_boot_file_name);
         ASSERT_RESULT_OK(vboot->Open());
+
+        if (!GetParam().dtb_file_name.empty()) {
+            dtb_replacement = std::make_unique<ReadOnlyTestFileHandle>(GetParam().dtb_file_name);
+            ASSERT_RESULT_OK(dtb_replacement->Open());
+        }
     }
     std::unique_ptr<TestFileHandle> vboot;
+    std::unique_ptr<TestFileHandle> dtb_replacement;
 };
 
 TEST_P(RepackVendorBootImgTest, InvalidSize) {
-    EXPECT_ERROR(replace_vendor_ramdisk(vboot->fd(), vboot->size() + 1, "default",
-                                        env->replace->fd(), env->replace->size()),
-                 HasSubstr("Size of vendor boot does not match"));
-    EXPECT_ERROR(replace_vendor_ramdisk(vboot->fd(), vboot->size(), "default", env->replace->fd(),
-                                        env->replace->size() + 1),
-                 HasSubstr("Size of new vendor ramdisk does not match"));
+    EXPECT_ERROR(
+            replace_vendor_ramdisk(vboot->fd(), vboot->size() + 1, "default", env->replace->fd(),
+                                   env->replace->size(),
+                                   !GetParam().dtb_file_name.empty() ? dtb_replacement->fd()
+                                                                     : android::base::unique_fd(-1),
+                                   !GetParam().dtb_file_name.empty() ? dtb_replacement->size() : 0),
+            HasSubstr("Size of vendor boot does not match"));
+    EXPECT_ERROR(
+            replace_vendor_ramdisk(vboot->fd(), vboot->size(), "default", env->replace->fd(),
+                                   env->replace->size() + 1,
+                                   !GetParam().dtb_file_name.empty() ? dtb_replacement->fd()
+                                                                     : android::base::unique_fd(-1),
+                                   !GetParam().dtb_file_name.empty() ? dtb_replacement->size() : 0),
+            HasSubstr("Size of new vendor ramdisk does not match"));
+    if (!GetParam().dtb_file_name.empty()) {
+        EXPECT_ERROR(replace_vendor_ramdisk(vboot->fd(), vboot->size(), "default",
+                                            env->replace->fd(), env->replace->size(),
+                                            dtb_replacement->fd(), dtb_replacement->size() + 1),
+                     HasSubstr("Size of new dtb does not match"));
+    }
+    EXPECT_ERROR(
+            replace_vendor_ramdisk(
+                    vboot->fd(), vboot->size(), "default", env->replace->fd(), env->replace->size(),
+                    android::base::unique_fd(std::numeric_limits<int32_t>::max()), 1),
+            HasSubstr("Can't seek to the beginning of new dtb image"));
 }
 
 TEST_P(RepackVendorBootImgTest, ReplaceUnknown) {
-    auto res = replace_vendor_ramdisk(vboot->fd(), vboot->size(), "unknown", env->replace->fd(),
-                                      env->replace->size());
+    auto res = replace_vendor_ramdisk(
+            vboot->fd(), vboot->size(), "unknown", env->replace->fd(), env->replace->size(),
+            !GetParam().dtb_file_name.empty() ? dtb_replacement->fd()
+                                              : android::base::unique_fd(-1),
+            !GetParam().dtb_file_name.empty() ? dtb_replacement->size() : 0);
     if (GetParam().expected_header_version == 3) {
         EXPECT_ERROR(res, Eq("Require vendor boot header V4 but is V3"));
     } else if (GetParam().expected_header_version == 4) {
@@ -279,8 +308,11 @@
     auto old_content = vboot->Read();
     ASSERT_RESULT_OK(old_content);
 
-    ASSERT_RESULT_OK(replace_vendor_ramdisk(vboot->fd(), vboot->size(), "default",
-                                            env->replace->fd(), env->replace->size()));
+    ASSERT_RESULT_OK(replace_vendor_ramdisk(
+            vboot->fd(), vboot->size(), "default", env->replace->fd(), env->replace->size(),
+            !GetParam().dtb_file_name.empty() ? dtb_replacement->fd()
+                                              : android::base::unique_fd(-1),
+            !GetParam().dtb_file_name.empty() ? dtb_replacement->size() : 0));
     EXPECT_RESULT(vboot->fsize(), vboot->size()) << "File size should not change after repack";
 
     auto new_content_res = vboot->Read();
@@ -291,14 +323,23 @@
     ASSERT_EQ(0, memcmp(VENDOR_BOOT_MAGIC, hdr->magic, VENDOR_BOOT_MAGIC_SIZE));
     ASSERT_EQ(GetParam().expected_header_version, hdr->header_version);
     EXPECT_EQ(hdr->vendor_ramdisk_size, env->replace->size());
-    EXPECT_EQ(hdr->dtb_size, env->dtb->size());
+    if (GetParam().dtb_file_name.empty()) {
+        EXPECT_EQ(hdr->dtb_size, env->dtb->size());
+    } else {
+        EXPECT_EQ(hdr->dtb_size, dtb_replacement->size());
+    }
 
     auto o = round_up(sizeof(vendor_boot_img_hdr_v3), hdr->page_size);
     auto p = round_up(hdr->vendor_ramdisk_size, hdr->page_size);
     auto q = round_up(hdr->dtb_size, hdr->page_size);
 
     EXPECT_THAT(new_content.substr(o, p), IsPadded(env->replace_content));
-    EXPECT_THAT(new_content.substr(o + p, q), IsPadded(env->dtb_content));
+    if (GetParam().dtb_file_name.empty()) {
+        EXPECT_THAT(new_content.substr(o + p, q), IsPadded(env->dtb_content));
+    } else {
+        auto dtb_content_res = dtb_replacement->Read();
+        EXPECT_THAT(new_content.substr(o + p, q), IsPadded(*dtb_content_res));
+    }
 
     if (hdr->header_version < 4) return;
 
@@ -321,11 +362,17 @@
 
 INSTANTIATE_TEST_SUITE_P(
         RepackVendorBootImgTest, RepackVendorBootImgTest,
-        ::testing::Values(RepackVendorBootImgTestParam{"vendor_boot_v3.img", 3},
-                          RepackVendorBootImgTestParam{"vendor_boot_v4_with_frag.img", 4},
-                          RepackVendorBootImgTestParam{"vendor_boot_v4_without_frag.img", 4}),
+        ::testing::Values(RepackVendorBootImgTestParam{"vendor_boot_v3.img", "", 3},
+                          RepackVendorBootImgTestParam{"vendor_boot_v4_with_frag.img", "", 4},
+                          RepackVendorBootImgTestParam{"vendor_boot_v4_without_frag.img", "", 4},
+                          RepackVendorBootImgTestParam{"vendor_boot_v4_with_frag.img",
+                                                       "dtb_replace.img", 4},
+                          RepackVendorBootImgTestParam{"vendor_boot_v4_without_frag.img",
+                                                       "dtb_replace.img", 4}),
         [](const auto& info) {
-            return android::base::StringReplace(info.param.vendor_boot_file_name, ".", "_", false);
+            std::string test_name =
+                    android::base::StringReplace(info.param.vendor_boot_file_name, ".", "_", false);
+            return test_name + (!info.param.dtb_file_name.empty() ? "_replace_dtb" : "");
         });
 
 std::string_view GetRamdiskName(const vendor_ramdisk_table_entry_v4* entry) {
@@ -368,7 +415,8 @@
     ASSERT_RESULT_OK(old_content);
 
     ASSERT_RESULT_OK(replace_vendor_ramdisk(vboot->fd(), vboot->size(), replace_ramdisk_name,
-                                            env->replace->fd(), env->replace->size()));
+                                            env->replace->fd(), env->replace->size(),
+                                            android::base::unique_fd(-1), 0));
     EXPECT_RESULT(vboot->fsize(), vboot->size()) << "File size should not change after repack";
 
     auto new_content_res = vboot->Read();