249 lines
		
	
	
		
			11 KiB
		
	
	
	
		
			C++
		
	
	
	
			
		
		
	
	
			249 lines
		
	
	
		
			11 KiB
		
	
	
	
		
			C++
		
	
	
	
//
 | 
						|
// Copyright (C) 2021 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 <unistd.h>
 | 
						|
 | 
						|
#include <android-base/file.h>
 | 
						|
#include <android-base/mapped_file.h>
 | 
						|
#include <android-base/properties.h>
 | 
						|
#include <bsdiff/bsdiff.h>
 | 
						|
#include <gtest/gtest.h>
 | 
						|
#include <libsnapshot/cow_writer.h>
 | 
						|
#include <libsnapshot/mock_snapshot_writer.h>
 | 
						|
 | 
						|
#include "update_engine/common/hash_calculator.h"
 | 
						|
#include "update_engine/common/mock_dynamic_partition_control.h"
 | 
						|
#include "update_engine/common/utils.h"
 | 
						|
#include "update_engine/payload_consumer/vabc_partition_writer.h"
 | 
						|
#include "update_engine/payload_generator/delta_diff_generator.h"
 | 
						|
#include "update_engine/update_metadata.pb.h"
 | 
						|
 | 
						|
namespace chromeos_update_engine {
 | 
						|
 | 
						|
using android::snapshot::CowOptions;
 | 
						|
using testing::_;
 | 
						|
using testing::Args;
 | 
						|
using testing::ElementsAreArray;
 | 
						|
using testing::Invoke;
 | 
						|
using testing::Return;
 | 
						|
using testing::Sequence;
 | 
						|
using utils::GetReadonlyZeroBlock;
 | 
						|
 | 
						|
namespace {
 | 
						|
 | 
						|
static constexpr auto& fake_part_name = "fake_part";
 | 
						|
static constexpr size_t FAKE_PART_SIZE = 4096 * 50;
 | 
						|
class VABCPartitionWriterTest : public ::testing::Test {
 | 
						|
 public:
 | 
						|
  void SetUp() override {
 | 
						|
    ftruncate(source_part_.fd, FAKE_PART_SIZE);
 | 
						|
    ON_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
 | 
						|
        .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::NONE)));
 | 
						|
  }
 | 
						|
 | 
						|
 protected:
 | 
						|
  CowMergeOperation* AddMergeOp(PartitionUpdate* partition,
 | 
						|
                                std::array<size_t, 2> src_extent,
 | 
						|
                                std::array<size_t, 2> dst_extent,
 | 
						|
                                CowMergeOperation_Type type) {
 | 
						|
    auto merge_op = partition->add_merge_operations();
 | 
						|
    auto src = merge_op->mutable_src_extent();
 | 
						|
    src->set_start_block(src_extent[0]);
 | 
						|
    src->set_num_blocks(src_extent[1]);
 | 
						|
    auto dst = merge_op->mutable_dst_extent();
 | 
						|
    dst->set_start_block(dst_extent[0]);
 | 
						|
    dst->set_num_blocks(dst_extent[1]);
 | 
						|
    merge_op->set_type(type);
 | 
						|
    return merge_op;
 | 
						|
  }
 | 
						|
 | 
						|
  android::snapshot::CowOptions options_ = {
 | 
						|
      .block_size = static_cast<uint32_t>(kBlockSize)};
 | 
						|
  android::snapshot::MockSnapshotWriter cow_writer_{options_};
 | 
						|
  MockDynamicPartitionControl dynamic_control_;
 | 
						|
  PartitionUpdate partition_update_;
 | 
						|
  InstallPlan install_plan_;
 | 
						|
  TemporaryFile source_part_;
 | 
						|
  InstallPlan::Partition install_part_{.name = fake_part_name,
 | 
						|
                                       .source_path = source_part_.path,
 | 
						|
                                       .source_size = FAKE_PART_SIZE};
 | 
						|
};
 | 
						|
 | 
						|
TEST_F(VABCPartitionWriterTest, MergeSequenceWriteTest) {
 | 
						|
  AddMergeOp(&partition_update_, {5, 1}, {10, 1}, CowMergeOperation::COW_COPY);
 | 
						|
  AddMergeOp(&partition_update_, {12, 2}, {13, 2}, CowMergeOperation::COW_XOR);
 | 
						|
  AddMergeOp(&partition_update_, {15, 1}, {20, 1}, CowMergeOperation::COW_COPY);
 | 
						|
  AddMergeOp(&partition_update_, {20, 1}, {25, 1}, CowMergeOperation::COW_COPY);
 | 
						|
  AddMergeOp(&partition_update_, {42, 5}, {40, 5}, CowMergeOperation::COW_XOR);
 | 
						|
  VABCPartitionWriter writer_{
 | 
						|
      partition_update_, install_part_, &dynamic_control_, kBlockSize};
 | 
						|
  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
 | 
						|
      .WillOnce(Invoke([](const std::string&,
 | 
						|
                          const std::optional<std::string>&,
 | 
						|
                          bool) {
 | 
						|
        auto cow_writer =
 | 
						|
            std::make_unique<android::snapshot::MockSnapshotWriter>(
 | 
						|
                android::snapshot::CowOptions{});
 | 
						|
        auto expected_merge_sequence = {10, 14, 13, 20, 25, 40, 41, 42, 43, 44};
 | 
						|
        EXPECT_CALL(*cow_writer, Initialize()).WillOnce(Return(true));
 | 
						|
        EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
 | 
						|
            .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
 | 
						|
            .WillOnce(Return(true));
 | 
						|
        ON_CALL(*cow_writer, EmitCopy(_, _)).WillByDefault(Return(true));
 | 
						|
        ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
 | 
						|
        return cow_writer;
 | 
						|
      }));
 | 
						|
  EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
 | 
						|
      .WillRepeatedly(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
 | 
						|
  ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
 | 
						|
}
 | 
						|
 | 
						|
TEST_F(VABCPartitionWriterTest, MergeSequenceXorSameBlock) {
 | 
						|
  AddMergeOp(&partition_update_, {19, 4}, {19, 3}, CowMergeOperation::COW_XOR)
 | 
						|
      ->set_src_offset(1);
 | 
						|
  VABCPartitionWriter writer_{
 | 
						|
      partition_update_, install_part_, &dynamic_control_, kBlockSize};
 | 
						|
  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
 | 
						|
      .WillOnce(Invoke(
 | 
						|
          [](const std::string&, const std::optional<std::string>&, bool) {
 | 
						|
            auto cow_writer =
 | 
						|
                std::make_unique<android::snapshot::MockSnapshotWriter>(
 | 
						|
                    android::snapshot::CowOptions{});
 | 
						|
            auto expected_merge_sequence = {19, 20, 21};
 | 
						|
            EXPECT_CALL(*cow_writer, Initialize()).WillOnce(Return(true));
 | 
						|
            EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
 | 
						|
                .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
 | 
						|
                .WillOnce(Return(true));
 | 
						|
            ON_CALL(*cow_writer, EmitCopy(_, _)).WillByDefault(Return(true));
 | 
						|
            ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
 | 
						|
            return cow_writer;
 | 
						|
          }));
 | 
						|
  EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
 | 
						|
      .WillRepeatedly(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
 | 
						|
  ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
 | 
						|
}
 | 
						|
 | 
						|
TEST_F(VABCPartitionWriterTest, EmitBlockTest) {
 | 
						|
  AddMergeOp(&partition_update_, {5, 1}, {10, 1}, CowMergeOperation::COW_COPY);
 | 
						|
  AddMergeOp(&partition_update_, {10, 1}, {15, 1}, CowMergeOperation::COW_COPY);
 | 
						|
  AddMergeOp(&partition_update_, {15, 2}, {20, 2}, CowMergeOperation::COW_COPY);
 | 
						|
  AddMergeOp(&partition_update_, {20, 1}, {25, 1}, CowMergeOperation::COW_COPY);
 | 
						|
  VABCPartitionWriter writer_{
 | 
						|
      partition_update_, install_part_, &dynamic_control_, kBlockSize};
 | 
						|
  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
 | 
						|
      .WillOnce(Invoke(
 | 
						|
          [](const std::string&, const std::optional<std::string>&, bool) {
 | 
						|
            auto cow_writer =
 | 
						|
                std::make_unique<android::snapshot::MockSnapshotWriter>(
 | 
						|
                    android::snapshot::CowOptions{});
 | 
						|
            Sequence s;
 | 
						|
            ON_CALL(*cow_writer, EmitCopy(_, _)).WillByDefault(Return(true));
 | 
						|
            ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
 | 
						|
            ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
 | 
						|
            EXPECT_CALL(*cow_writer, Initialize()).InSequence(s);
 | 
						|
            EXPECT_CALL(*cow_writer, EmitCopy(10, 5)).InSequence(s);
 | 
						|
            EXPECT_CALL(*cow_writer, EmitCopy(15, 10)).InSequence(s);
 | 
						|
            // libsnapshot want blocks in reverser order, so 21 goes before 20
 | 
						|
            EXPECT_CALL(*cow_writer, EmitCopy(21, 16)).InSequence(s);
 | 
						|
            EXPECT_CALL(*cow_writer, EmitCopy(20, 15)).InSequence(s);
 | 
						|
 | 
						|
            EXPECT_CALL(*cow_writer, EmitCopy(25, 20)).InSequence(s);
 | 
						|
            return cow_writer;
 | 
						|
          }));
 | 
						|
  ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
 | 
						|
}
 | 
						|
 | 
						|
std::string GetNoopBSDIFF(size_t data_size) {
 | 
						|
  auto zeros = GetReadonlyZeroBlock(data_size);
 | 
						|
  TemporaryFile patch_file;
 | 
						|
  int error = bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(zeros->data()),
 | 
						|
                             zeros->size(),
 | 
						|
                             reinterpret_cast<const uint8_t*>(zeros->data()),
 | 
						|
                             zeros->size(),
 | 
						|
                             patch_file.path,
 | 
						|
                             nullptr);
 | 
						|
  if (error) {
 | 
						|
    LOG(ERROR) << "Failed to generate BSDIFF patch " << error;
 | 
						|
    return {};
 | 
						|
  }
 | 
						|
  std::string patch_data;
 | 
						|
  if (!utils::ReadFile(patch_file.path, &patch_data)) {
 | 
						|
    return {};
 | 
						|
  }
 | 
						|
  return patch_data;
 | 
						|
}
 | 
						|
 | 
						|
TEST_F(VABCPartitionWriterTest, StreamXORBlockTest) {
 | 
						|
  AddMergeOp(&partition_update_, {5, 2}, {10, 2}, CowMergeOperation::COW_XOR);
 | 
						|
  AddMergeOp(&partition_update_, {8, 2}, {13, 2}, CowMergeOperation::COW_XOR);
 | 
						|
  auto install_op = partition_update_.add_operations();
 | 
						|
  *install_op->add_src_extents() = ExtentForRange(5, 5);
 | 
						|
  *install_op->add_dst_extents() = ExtentForRange(10, 5);
 | 
						|
  install_op->set_type(InstallOperation::SOURCE_BSDIFF);
 | 
						|
  auto data_hash = install_op->mutable_src_sha256_hash();
 | 
						|
  auto zeros = GetReadonlyZeroBlock(kBlockSize * 5);
 | 
						|
  brillo::Blob expected_hash;
 | 
						|
  truncate64(source_part_.path, kBlockSize * 20);
 | 
						|
  HashCalculator::RawHashOfBytes(zeros->data(), zeros->size(), &expected_hash);
 | 
						|
  data_hash->assign(reinterpret_cast<const char*>(expected_hash.data()),
 | 
						|
                    expected_hash.size());
 | 
						|
 | 
						|
  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
 | 
						|
      .WillOnce(Invoke([](const std::string&,
 | 
						|
                          const std::optional<std::string>&,
 | 
						|
                          bool) {
 | 
						|
        auto cow_writer =
 | 
						|
            std::make_unique<android::snapshot::MockSnapshotWriter>(
 | 
						|
                android::snapshot::CowOptions{});
 | 
						|
        ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
 | 
						|
        auto expected_merge_sequence = {10, 11, 13, 14};
 | 
						|
        auto expected_merge_sequence_rev = {11, 10, 14, 13};
 | 
						|
        const bool is_ascending = android::base::GetBoolProperty(
 | 
						|
            "ro.virtual_ab.userspace.snapshots.enabled", false);
 | 
						|
        ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
 | 
						|
        if (!is_ascending) {
 | 
						|
          EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
 | 
						|
              .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence_rev)))
 | 
						|
              .WillOnce(Return(true));
 | 
						|
        } else {
 | 
						|
          EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
 | 
						|
              .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
 | 
						|
              .WillOnce(Return(true));
 | 
						|
        }
 | 
						|
        EXPECT_CALL(*cow_writer, Initialize()).Times(1);
 | 
						|
        EXPECT_CALL(*cow_writer, EmitCopy(_, _)).Times(0);
 | 
						|
        EXPECT_CALL(*cow_writer, EmitRawBlocks(_, _, _)).WillOnce(Return(true));
 | 
						|
        EXPECT_CALL(*cow_writer, EmitXorBlocks(10, _, kBlockSize * 2, 5, 0))
 | 
						|
            .WillOnce(Return(true));
 | 
						|
        EXPECT_CALL(*cow_writer, EmitXorBlocks(13, _, kBlockSize * 2, 8, 0))
 | 
						|
            .WillOnce(Return(true));
 | 
						|
        return cow_writer;
 | 
						|
      }));
 | 
						|
  EXPECT_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
 | 
						|
      .WillRepeatedly(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
 | 
						|
  VABCPartitionWriter writer_{
 | 
						|
      partition_update_, install_part_, &dynamic_control_, kBlockSize};
 | 
						|
  ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
 | 
						|
  const auto patch_data = GetNoopBSDIFF(kBlockSize * 5);
 | 
						|
  ASSERT_GT(patch_data.size(), 0UL);
 | 
						|
  ASSERT_TRUE(writer_.PerformDiffOperation(
 | 
						|
      *install_op, nullptr, patch_data.data(), patch_data.size()));
 | 
						|
}
 | 
						|
 | 
						|
}  // namespace
 | 
						|
 | 
						|
}  // namespace chromeos_update_engine
 |