Parse Android boot.img.

The Android boot.img contains kernel, ramdisk and optionally second
stage bootloader. Both kernel and ramdisk could be compressed by gzip.

This patch add a new boot img filesystem to parse boot.img to split
it into |File|s, and find the deflates to utilize puffin, which could
save us several MBs in delta payload size.

As a side effect, generating delta payload for boot partition is much
faster because it is now splitted into smaller files, and we can make
use of multithreading.

Bug: 110494725
Test: generate a delta payload
Change-Id: If41468b37a407fe1b7a70c2b61f632f7df8176d5
diff --git a/Android.mk b/Android.mk
index 80a8afd..171eb0e 100644
--- a/Android.mk
+++ b/Android.mk
@@ -655,6 +655,7 @@
     payload_generator/annotated_operation.cc \
     payload_generator/blob_file_writer.cc \
     payload_generator/block_mapping.cc \
+    payload_generator/boot_img_filesystem.cc \
     payload_generator/bzip.cc \
     payload_generator/cycle_breaker.cc \
     payload_generator/deflate_utils.cc \
@@ -956,6 +957,7 @@
     payload_generator/ab_generator_unittest.cc \
     payload_generator/blob_file_writer_unittest.cc \
     payload_generator/block_mapping_unittest.cc \
+    payload_generator/boot_img_filesystem_unittest.cc \
     payload_generator/cycle_breaker_unittest.cc \
     payload_generator/deflate_utils_unittest.cc \
     payload_generator/delta_diff_utils_unittest.cc \
diff --git a/common/utils.h b/common/utils.h
index 5c44083..83a63ef 100644
--- a/common/utils.h
+++ b/common/utils.h
@@ -317,6 +317,16 @@
 // reboot. Returns whether it succeeded getting the boot_id.
 bool GetBootId(std::string* boot_id);
 
+// Divide |x| by |y| and round up to the nearest integer.
+constexpr uint64_t DivRoundUp(uint64_t x, uint64_t y) {
+  return (x + y - 1) / y;
+}
+
+// Round |x| up to be a multiple of |y|.
+constexpr uint64_t RoundUp(uint64_t x, uint64_t y) {
+  return DivRoundUp(x, y) * y;
+}
+
 }  // namespace utils
 
 
diff --git a/payload_consumer/bzip_extent_writer_unittest.cc b/payload_consumer/bzip_extent_writer_unittest.cc
index bf050ef..4426876 100644
--- a/payload_consumer/bzip_extent_writer_unittest.cc
+++ b/payload_consumer/bzip_extent_writer_unittest.cc
@@ -100,8 +100,7 @@
   for (size_t i = 0; i < decompressed_data.size(); ++i)
     decompressed_data[i] = static_cast<uint8_t>("ABC\n"[i % 4]);
 
-  vector<Extent> extents = {
-      ExtentForRange(0, (kDecompressedLength + kBlockSize - 1) / kBlockSize)};
+  vector<Extent> extents = {ExtentForBytes(kBlockSize, 0, kDecompressedLength)};
 
   BzipExtentWriter bzip_writer(std::make_unique<DirectExtentWriter>());
   EXPECT_TRUE(
diff --git a/payload_generator/boot_img_filesystem.cc b/payload_generator/boot_img_filesystem.cc
new file mode 100644
index 0000000..90be966
--- /dev/null
+++ b/payload_generator/boot_img_filesystem.cc
@@ -0,0 +1,110 @@
+//
+// 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 "update_engine/payload_generator/boot_img_filesystem.h"
+
+#include <base/logging.h>
+#include <brillo/secure_blob.h>
+#include <puffin/utils.h>
+
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+
+using std::string;
+using std::unique_ptr;
+using std::vector;
+
+namespace chromeos_update_engine {
+
+unique_ptr<BootImgFilesystem> BootImgFilesystem::CreateFromFile(
+    const string& filename) {
+  if (filename.empty())
+    return nullptr;
+
+  brillo::Blob header;
+  if (!utils::ReadFileChunk(filename, 0, sizeof(boot_img_hdr), &header) ||
+      header.size() != sizeof(boot_img_hdr) ||
+      memcmp(header.data(), BOOT_MAGIC, BOOT_MAGIC_SIZE) != 0) {
+    return nullptr;
+  }
+
+  unique_ptr<BootImgFilesystem> result(new BootImgFilesystem());
+  result->filename_ = filename;
+  memcpy(&result->hdr_, header.data(), header.size());
+  return result;
+}
+
+size_t BootImgFilesystem::GetBlockSize() const {
+  // Page size may not be 4K, but we currently only support 4K block size.
+  return kBlockSize;
+}
+
+size_t BootImgFilesystem::GetBlockCount() const {
+  return utils::DivRoundUp(utils::FileSize(filename_), kBlockSize);
+}
+
+FilesystemInterface::File BootImgFilesystem::GetFile(const string& name,
+                                                     uint64_t offset,
+                                                     uint64_t size) const {
+  File file;
+  file.name = name;
+  file.extents = {ExtentForBytes(kBlockSize, offset, size)};
+
+  brillo::Blob data;
+  if (utils::ReadFileChunk(filename_, offset, size, &data)) {
+    constexpr size_t kGZipHeaderSize = 10;
+    // Check GZip header magic.
+    if (data.size() > kGZipHeaderSize && data[0] == 0x1F && data[1] == 0x8B) {
+      if (!puffin::LocateDeflatesInGzip(data, &file.deflates)) {
+        // We still use the deflates found even if LocateDeflatesInGzip() fails,
+        // if any deflates are returned, they should be correct, it's possible
+        // something went wrong later but it shouldn't stop us from using the
+        // previous deflates. Another common case is if there's more data after
+        // the gzip, the function will try to parse that as another gzip and
+        // will fail, but we still want the deflates from the first gzip.
+        LOG(WARNING) << "Error occurred parsing gzip " << name << " at offset "
+                     << offset << " of " << filename_ << ", found "
+                     << file.deflates.size() << " deflates.";
+      }
+      for (auto& deflate : file.deflates) {
+        deflate.offset += offset * 8;
+      }
+    }
+  }
+  return file;
+}
+
+bool BootImgFilesystem::GetFiles(vector<File>* files) const {
+  files->clear();
+  const uint64_t file_size = utils::FileSize(filename_);
+  // The first page is header.
+  uint64_t offset = hdr_.page_size;
+  if (hdr_.kernel_size > 0 && offset + hdr_.kernel_size <= file_size) {
+    files->emplace_back(GetFile("<kernel>", offset, hdr_.kernel_size));
+  }
+  offset += utils::RoundUp(hdr_.kernel_size, hdr_.page_size);
+  if (hdr_.ramdisk_size > 0 && offset + hdr_.ramdisk_size <= file_size) {
+    files->emplace_back(GetFile("<ramdisk>", offset, hdr_.ramdisk_size));
+  }
+  return true;
+}
+
+bool BootImgFilesystem::LoadSettings(brillo::KeyValueStore* store) const {
+  return false;
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/boot_img_filesystem.h b/payload_generator/boot_img_filesystem.h
new file mode 100644
index 0000000..87725d4
--- /dev/null
+++ b/payload_generator/boot_img_filesystem.h
@@ -0,0 +1,78 @@
+//
+// 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.
+//
+
+#ifndef UPDATE_ENGINE_PAYLOAD_GENERATOR_BOOT_IMG_FILESYSTEM_H_
+#define UPDATE_ENGINE_PAYLOAD_GENERATOR_BOOT_IMG_FILESYSTEM_H_
+
+#include "update_engine/payload_generator/filesystem_interface.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace chromeos_update_engine {
+
+class BootImgFilesystem : public FilesystemInterface {
+ public:
+  // Creates an BootImgFilesystem from an Android boot.img file.
+  static std::unique_ptr<BootImgFilesystem> CreateFromFile(
+      const std::string& filename);
+  ~BootImgFilesystem() override = default;
+
+  // FilesystemInterface overrides.
+  size_t GetBlockSize() const override;
+  size_t GetBlockCount() const override;
+
+  // GetFiles will return one FilesystemInterface::File for kernel and one for
+  // ramdisk.
+  bool GetFiles(std::vector<File>* files) const override;
+
+  bool LoadSettings(brillo::KeyValueStore* store) const override;
+
+ private:
+  friend class BootImgFilesystemTest;
+
+  BootImgFilesystem() = default;
+
+  File GetFile(const std::string& name, uint64_t offset, uint64_t size) const;
+
+  // The boot.img file path.
+  std::string filename_;
+
+// https://android.googlesource.com/platform/system/core/+/master/mkbootimg/include/bootimg/bootimg.h
+#define BOOT_MAGIC "ANDROID!"
+#define BOOT_MAGIC_SIZE 8
+  struct boot_img_hdr {
+    // Must be BOOT_MAGIC.
+    uint8_t magic[BOOT_MAGIC_SIZE];
+    uint32_t kernel_size;  /* size in bytes */
+    uint32_t kernel_addr;  /* physical load addr */
+    uint32_t ramdisk_size; /* size in bytes */
+    uint32_t ramdisk_addr; /* physical load addr */
+    uint32_t second_size;  /* size in bytes */
+    uint32_t second_addr;  /* physical load addr */
+    uint32_t tags_addr;    /* physical addr for kernel tags */
+    uint32_t page_size;    /* flash page size we assume */
+  } __attribute__((packed));
+  // The boot image header.
+  boot_img_hdr hdr_;
+
+  DISALLOW_COPY_AND_ASSIGN(BootImgFilesystem);
+};
+
+}  // namespace chromeos_update_engine
+
+#endif  // UPDATE_ENGINE_PAYLOAD_GENERATOR_BOOT_IMG_FILESYSTEM_H_
diff --git a/payload_generator/boot_img_filesystem_unittest.cc b/payload_generator/boot_img_filesystem_unittest.cc
new file mode 100644
index 0000000..b1e0d99
--- /dev/null
+++ b/payload_generator/boot_img_filesystem_unittest.cc
@@ -0,0 +1,117 @@
+//
+// 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 "update_engine/payload_generator/boot_img_filesystem.h"
+
+#include <vector>
+
+#include <brillo/secure_blob.h>
+#include <gtest/gtest.h>
+
+#include "update_engine/common/test_utils.h"
+#include "update_engine/common/utils.h"
+
+namespace chromeos_update_engine {
+
+using std::unique_ptr;
+using std::vector;
+
+class BootImgFilesystemTest : public ::testing::Test {
+ protected:
+  brillo::Blob GetBootImg(const brillo::Blob& kernel,
+                          const brillo::Blob& ramdisk) {
+    brillo::Blob boot_img(16 * 1024);
+    BootImgFilesystem::boot_img_hdr hdr;
+    memcpy(hdr.magic, BOOT_MAGIC, BOOT_MAGIC_SIZE);
+    hdr.kernel_size = kernel.size();
+    hdr.ramdisk_size = ramdisk.size();
+    hdr.page_size = 4096;
+    size_t offset = 0;
+    memcpy(boot_img.data() + offset, &hdr, sizeof(hdr));
+    offset += utils::RoundUp(sizeof(hdr), hdr.page_size);
+    memcpy(boot_img.data() + offset, kernel.data(), kernel.size());
+    offset += utils::RoundUp(kernel.size(), hdr.page_size);
+    memcpy(boot_img.data() + offset, ramdisk.data(), ramdisk.size());
+    return boot_img;
+  }
+
+  test_utils::ScopedTempFile boot_file_;
+};
+
+TEST_F(BootImgFilesystemTest, SimpleTest) {
+  test_utils::WriteFileVector(
+      boot_file_.path(),
+      GetBootImg(brillo::Blob(1234, 'k'), brillo::Blob(5678, 'r')));
+  unique_ptr<BootImgFilesystem> fs =
+      BootImgFilesystem::CreateFromFile(boot_file_.path());
+  EXPECT_NE(nullptr, fs);
+
+  vector<FilesystemInterface::File> files;
+  EXPECT_TRUE(fs->GetFiles(&files));
+  ASSERT_EQ(2u, files.size());
+
+  EXPECT_EQ("<kernel>", files[0].name);
+  EXPECT_EQ(1u, files[0].extents.size());
+  EXPECT_EQ(1u, files[0].extents[0].start_block());
+  EXPECT_EQ(1u, files[0].extents[0].num_blocks());
+  EXPECT_TRUE(files[0].deflates.empty());
+
+  EXPECT_EQ("<ramdisk>", files[1].name);
+  EXPECT_EQ(1u, files[1].extents.size());
+  EXPECT_EQ(2u, files[1].extents[0].start_block());
+  EXPECT_EQ(2u, files[1].extents[0].num_blocks());
+  EXPECT_TRUE(files[1].deflates.empty());
+}
+
+TEST_F(BootImgFilesystemTest, BadImageTest) {
+  brillo::Blob boot_img = GetBootImg({}, {});
+  boot_img[7] = '?';
+  test_utils::WriteFileVector(boot_file_.path(), boot_img);
+  unique_ptr<BootImgFilesystem> fs =
+      BootImgFilesystem::CreateFromFile(boot_file_.path());
+  EXPECT_EQ(nullptr, fs);
+}
+
+TEST_F(BootImgFilesystemTest, GZipRamdiskTest) {
+  // echo ramdisk | gzip | hexdump -v -e '/1 "0x%02x, "'
+  const brillo::Blob ramdisk = {0x1f, 0x8b, 0x08, 0x00, 0x3a, 0x83, 0x35,
+                                0x5b, 0x00, 0x03, 0x2b, 0x4a, 0xcc, 0x4d,
+                                0xc9, 0x2c, 0xce, 0xe6, 0x02, 0x00, 0x2e,
+                                0xf6, 0x0b, 0x08, 0x08, 0x00, 0x00, 0x00};
+  test_utils::WriteFileVector(boot_file_.path(),
+                              GetBootImg(brillo::Blob(5678, 'k'), ramdisk));
+  unique_ptr<BootImgFilesystem> fs =
+      BootImgFilesystem::CreateFromFile(boot_file_.path());
+  EXPECT_NE(nullptr, fs);
+
+  vector<FilesystemInterface::File> files;
+  EXPECT_TRUE(fs->GetFiles(&files));
+  ASSERT_EQ(2u, files.size());
+
+  EXPECT_EQ("<kernel>", files[0].name);
+  EXPECT_EQ(1u, files[0].extents.size());
+  EXPECT_EQ(1u, files[0].extents[0].start_block());
+  EXPECT_EQ(2u, files[0].extents[0].num_blocks());
+  EXPECT_TRUE(files[0].deflates.empty());
+
+  EXPECT_EQ("<ramdisk>", files[1].name);
+  EXPECT_EQ(1u, files[1].extents.size());
+  EXPECT_EQ(3u, files[1].extents[0].start_block());
+  EXPECT_EQ(1u, files[1].extents[0].num_blocks());
+  EXPECT_EQ(1u, files[1].deflates.size());
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_ranges.cc b/payload_generator/extent_ranges.cc
index c1d3d63..41e8f76 100644
--- a/payload_generator/extent_ranges.cc
+++ b/payload_generator/extent_ranges.cc
@@ -227,6 +227,14 @@
   return ret;
 }
 
+Extent ExtentForBytes(uint64_t block_size,
+                      uint64_t start_bytes,
+                      uint64_t size_bytes) {
+  uint64_t start_block = start_bytes / block_size;
+  uint64_t end_block = utils::DivRoundUp(start_bytes + size_bytes, block_size);
+  return ExtentForRange(start_block, end_block - start_block);
+}
+
 vector<Extent> ExtentRanges::GetExtentsForBlockCount(
     uint64_t count) const {
   vector<Extent> out;
diff --git a/payload_generator/extent_ranges.h b/payload_generator/extent_ranges.h
index 198c834..02cf8fc 100644
--- a/payload_generator/extent_ranges.h
+++ b/payload_generator/extent_ranges.h
@@ -41,6 +41,9 @@
 };
 
 Extent ExtentForRange(uint64_t start_block, uint64_t num_blocks);
+Extent ExtentForBytes(uint64_t block_size,
+                      uint64_t start_bytes,
+                      uint64_t size_bytes);
 
 class ExtentRanges {
  public:
diff --git a/payload_generator/full_update_generator.cc b/payload_generator/full_update_generator.cc
index 482a789..98bb0f3 100644
--- a/payload_generator/full_update_generator.cc
+++ b/payload_generator/full_update_generator.cc
@@ -152,7 +152,7 @@
   // We potentially have all the ChunkProcessors in memory but only
   // |max_threads| will actually hold a block in memory while we process.
   size_t partition_blocks = new_part.size / config.block_size;
-  size_t num_chunks = (partition_blocks + chunk_blocks - 1) / chunk_blocks;
+  size_t num_chunks = utils::DivRoundUp(partition_blocks, chunk_blocks);
   aops->resize(num_chunks);
   vector<ChunkProcessor> chunk_processors;
   chunk_processors.reserve(num_chunks);
diff --git a/payload_generator/payload_generation_config.cc b/payload_generator/payload_generation_config.cc
index 15d4ab5..212d45f 100644
--- a/payload_generator/payload_generation_config.cc
+++ b/payload_generator/payload_generation_config.cc
@@ -20,6 +20,7 @@
 
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/delta_performer.h"
+#include "update_engine/payload_generator/boot_img_filesystem.h"
 #include "update_engine/payload_generator/delta_diff_generator.h"
 #include "update_engine/payload_generator/delta_diff_utils.h"
 #include "update_engine/payload_generator/ext2_filesystem.h"
@@ -64,6 +65,12 @@
     }
   }
 
+  fs_interface = BootImgFilesystem::CreateFromFile(path);
+  if (fs_interface) {
+    TEST_AND_RETURN_FALSE(fs_interface->GetBlockSize() == kBlockSize);
+    return true;
+  }
+
   // Fall back to a RAW filesystem.
   TEST_AND_RETURN_FALSE(size % kBlockSize == 0);
   fs_interface = RawFilesystem::Create(
diff --git a/payload_generator/payload_signer.cc b/payload_generator/payload_signer.cc
index 0b47dd4..8eba2dc 100644
--- a/payload_generator/payload_signer.cc
+++ b/payload_generator/payload_signer.cc
@@ -231,8 +231,8 @@
     Extent* dummy_extent = dummy_op->add_dst_extents();
     // Tell the dummy op to write this data to a big sparse hole
     dummy_extent->set_start_block(kSparseHole);
-    dummy_extent->set_num_blocks((signature_blob_length + kBlockSize - 1) /
-                                 kBlockSize);
+    dummy_extent->set_num_blocks(
+        utils::DivRoundUp(signature_blob_length, kBlockSize));
   }
 }
 
diff --git a/payload_generator/squashfs_filesystem.cc b/payload_generator/squashfs_filesystem.cc
index c98ad12..6c892f5 100644
--- a/payload_generator/squashfs_filesystem.cc
+++ b/payload_generator/squashfs_filesystem.cc
@@ -44,14 +44,6 @@
 
 namespace {
 
-Extent ExtentForBytes(uint64_t block_size,
-                      uint64_t start_bytes,
-                      uint64_t size_bytes) {
-  uint64_t start_block = start_bytes / block_size;
-  uint64_t end_block = (start_bytes + size_bytes + block_size - 1) / block_size;
-  return ExtentForRange(start_block, end_block - start_block);
-}
-
 // The size of the squashfs super block.
 constexpr size_t kSquashfsSuperBlockSize = 96;
 constexpr uint64_t kSquashfsCompressedBit = 1 << 24;
@@ -192,8 +184,7 @@
   for (const auto& file : files_) {
     file_extents.AddExtents(file.extents);
   }
-  vector<Extent> full = {
-      ExtentForRange(0, (size_ + kBlockSize - 1) / kBlockSize)};
+  vector<Extent> full = {ExtentForBytes(kBlockSize, 0, size_)};
   auto metadata_extents = FilterExtentRanges(full, file_extents);
   // For now there should be at most two extents. One for superblock and one for
   // metadata at the end. Just create appropriate files with <metadata-i> name.
diff --git a/update_engine.gyp b/update_engine.gyp
index 7cb2a7d..85632a2 100644
--- a/update_engine.gyp
+++ b/update_engine.gyp
@@ -391,6 +391,7 @@
         'payload_generator/annotated_operation.cc',
         'payload_generator/blob_file_writer.cc',
         'payload_generator/block_mapping.cc',
+        'payload_generator/boot_img_filesystem.cc',
         'payload_generator/bzip.cc',
         'payload_generator/cycle_breaker.cc',
         'payload_generator/deflate_utils.cc',
@@ -544,6 +545,7 @@
             'payload_generator/ab_generator_unittest.cc',
             'payload_generator/blob_file_writer_unittest.cc',
             'payload_generator/block_mapping_unittest.cc',
+            'payload_generator/boot_img_filesystem_unittest.cc',
             'payload_generator/cycle_breaker_unittest.cc',
             'payload_generator/deflate_utils_unittest.cc',
             'payload_generator/delta_diff_utils_unittest.cc',