payload_consumer/partition_writer_unittest.cc - platform/system/update_engine - Git at Google

 //
 // Copyright (C) 2020 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 <memory>
 #include <vector>

 #include <brillo/secure_blob.h>
 #include <gtest/gtest.h>

 #include "update_engine/common/dynamic_partition_control_stub.h"
 #include "update_engine/common/error_code.h"
 #include "update_engine/common/fake_prefs.h"
 #include "update_engine/common/hash_calculator.h"
 #include "update_engine/common/test_utils.h"
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/extent_writer.h"
 #include "update_engine/payload_consumer/fake_file_descriptor.h"
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
 #include "update_engine/payload_consumer/payload_constants.h"
 #include "update_engine/payload_generator/annotated_operation.h"
 #include "update_engine/payload_generator/delta_diff_generator.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/payload_generator/payload_generation_config.h"
 #include "update_engine/update_metadata.pb.h"

 namespace chromeos_update_engine {

 class PartitionWriterTest : public testing::Test {
  public:
   // Helper function to pretend that the ECC file descriptor was already opened.
   // Returns a pointer to the created file descriptor.
   FakeFileDescriptor* SetFakeECCFile(size_t size) {
     EXPECT_FALSE(writer_.verified_source_fd_.source_ecc_fd_)
         << "source_ecc_fdb already open.";
     FakeFileDescriptor* ret = new FakeFileDescriptor();
     fake_ecc_fd_.reset(ret);
     // Call open to simulate it was already opened.
     ret->Open("", 0);
     ret->SetFileSize(size);
     writer_.verified_source_fd_.source_ecc_fd_ = fake_ecc_fd_;
     return ret;
   }

   uint64_t GetSourceEccRecoveredFailures() const {
     return writer_.verified_source_fd_.source_ecc_recovered_failures_;
   }

   AnnotatedOperation GenerateSourceCopyOp(const brillo::Blob& copied_data,
                                           bool add_hash,
                                           PartitionConfig* old_part = nullptr) {
     PayloadGenerationConfig config;
     const uint64_t kDefaultBlockSize = config.block_size;
     EXPECT_EQ(0U, copied_data.size() % kDefaultBlockSize);
     uint64_t num_blocks = copied_data.size() / kDefaultBlockSize;
     AnnotatedOperation aop;
     *(aop.op.add_src_extents()) = ExtentForRange(0, num_blocks);
     *(aop.op.add_dst_extents()) = ExtentForRange(0, num_blocks);
     aop.op.set_type(InstallOperation::SOURCE_COPY);
     brillo::Blob src_hash;
     EXPECT_TRUE(HashCalculator::RawHashOfData(copied_data, &src_hash));
     if (add_hash)
       aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());

     return aop;
   }

   brillo::Blob PerformSourceCopyOp(const InstallOperation& op,
                                    const brillo::Blob blob_data) {
     LOG(INFO) << "Using source part " << source_partition.path();
     FileDescriptorPtr fd(new EintrSafeFileDescriptor());
     DirectExtentWriter extent_writer{fd};
     EXPECT_TRUE(fd->Open(source_partition.path().c_str(), O_RDWR));
     if (HasFailure()) {
       return {};
     }
     EXPECT_TRUE(extent_writer.Init(op.src_extents(), kBlockSize));
     if (HasFailure()) {
       return {};
     }
     EXPECT_TRUE(extent_writer.Write(blob_data.data(), blob_data.size()));
     if (HasFailure()) {
       return {};
     }
     fd->Flush();

     install_part_.source_size = blob_data.size();
     install_part_.target_size = blob_data.size();

     ErrorCode error;
     EXPECT_TRUE(writer_.Init(&install_plan_, true, 0));
     if (HasFailure()) {
       return {};
     }
     EXPECT_TRUE(writer_.PerformSourceCopyOperation(op, &error));
     writer_.CheckpointUpdateProgress(1);

     brillo::Blob output_data;
     EXPECT_TRUE(utils::ReadFile(target_partition.path(), &output_data));
     return output_data;
   }

   FakePrefs prefs_{};
   InstallPlan install_plan_{};
   InstallPlan::Payload payload_{};
   DynamicPartitionControlStub dynamic_control_{};
   FileDescriptorPtr fake_ecc_fd_{};
   DeltaArchiveManifest manifest_{};
   ScopedTempFile source_partition{"source-part-XXXXXX"};
   ScopedTempFile target_partition{"target-part-XXXXXX"};
   InstallPlan::Partition install_part_{.source_path = source_partition.path(),
                                        .target_path = target_partition.path()};
   PartitionUpdate partition_update_{};
   PartitionWriter writer_{
       partition_update_, install_part_, &dynamic_control_, kBlockSize, false};
 };
 // Test that the error-corrected file descriptor is used to read a partition
 // when no hash is available for SOURCE_COPY but it falls back to the normal
 // file descriptor when the size of the error corrected one is too small.
 TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyWhenNoHashFallbackTest) {
   constexpr size_t kCopyOperationSize = 4 * 4096;
   ScopedTempFile source("Source-XXXXXX");
   // Setup the source path with the right expected data.
   brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
   ASSERT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));

   // Setup the fec file descriptor as the fake stream, with smaller data than
   // the expected.
   FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize / 2);

   PartitionConfig old_part(kPartitionNameRoot);
   old_part.path = source.path();
   old_part.size = expected_data.size();

   // The payload operation doesn't include an operation hash.
   auto source_copy_op = GenerateSourceCopyOp(expected_data, false, &old_part);
   ASSERT_NO_FATAL_FAILURE();
   auto output_data = PerformSourceCopyOp(source_copy_op.op, expected_data);
   ASSERT_NO_FATAL_FAILURE();
   ASSERT_EQ(output_data, expected_data);

   // Verify that the fake_fec was attempted to be used. Since the file
   // descriptor is shorter it can actually do more than one read to realize it
   // reached the EOF.
   ASSERT_LE(1U, fake_fec->GetReadOps().size());
   // This fallback doesn't count as an error-corrected operation since the
   // operation hash was not available.
   ASSERT_EQ(0U, GetSourceEccRecoveredFailures());
 }

 // Test that the error-corrected file descriptor is used to read the partition
 // since the source partition doesn't match the operation hash.
 TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyFallbackTest) {
   constexpr size_t kCopyOperationSize = 4 * 4096;
   // Write invalid data to the source image, which doesn't match the expected
   // hash.
   brillo::Blob invalid_data(kCopyOperationSize, 0x55);

   // Setup the fec file descriptor as the fake stream, which matches
   // |expected_data|.
   FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize);
   brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);

   auto source_copy_op = GenerateSourceCopyOp(expected_data, true);
   ASSERT_NO_FATAL_FAILURE();
   auto output_data = PerformSourceCopyOp(source_copy_op.op, invalid_data);
   ASSERT_NO_FATAL_FAILURE();
   ASSERT_EQ(output_data, expected_data);

   // Verify that the fake_fec was actually used.
   EXPECT_GE(fake_fec->GetReadOps().size(), 1U);
   EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
 }

 TEST_F(PartitionWriterTest, ChooseSourceFDTest) {
   constexpr size_t kSourceSize = 4 * 4096;
   ScopedTempFile source("Source-XXXXXX");
   // Write invalid data to the source image, which doesn't match the expected
   // hash.
   brillo::Blob invalid_data(kSourceSize, 0x55);
   ASSERT_TRUE(test_utils::WriteFileVector(source.path(), invalid_data));

   writer_.verified_source_fd_.source_fd_ =
       std::make_shared<EintrSafeFileDescriptor>();
   writer_.verified_source_fd_.source_fd_->Open(source.path().c_str(), O_RDONLY);

   // Setup the fec file descriptor as the fake stream, which matches
   // |expected_data|.
   FakeFileDescriptor* fake_fec = SetFakeECCFile(kSourceSize);
   brillo::Blob expected_data = FakeFileDescriptorData(kSourceSize);

   InstallOperation op;
   *(op.add_src_extents()) = ExtentForRange(0, kSourceSize / 4096);
   brillo::Blob src_hash;
   ASSERT_TRUE(HashCalculator::RawHashOfData(expected_data, &src_hash));
   op.set_src_sha256_hash(src_hash.data(), src_hash.size());

   ErrorCode error = ErrorCode::kSuccess;
   ASSERT_NE(writer_.ChooseSourceFD(op, &error), nullptr);
   ASSERT_EQ(ErrorCode::kSuccess, error);
   // Verify that the fake_fec was actually used.
   ASSERT_EQ(1U, fake_fec->GetReadOps().size());
   ASSERT_EQ(1U, GetSourceEccRecoveredFailures());
 }

 }  // namespace chromeos_update_engine
	//
	// Copyright (C) 2020 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 <memory>
	#include <vector>

	#include <brillo/secure_blob.h>
	#include <gtest/gtest.h>

	#include "update_engine/common/dynamic_partition_control_stub.h"
	#include "update_engine/common/error_code.h"
	#include "update_engine/common/fake_prefs.h"
	#include "update_engine/common/hash_calculator.h"
	#include "update_engine/common/test_utils.h"
	#include "update_engine/common/utils.h"
	#include "update_engine/payload_consumer/extent_writer.h"
	#include "update_engine/payload_consumer/fake_file_descriptor.h"
	#include "update_engine/payload_consumer/file_descriptor.h"
	#include "update_engine/payload_consumer/install_plan.h"
	#include "update_engine/payload_consumer/partition_writer.h"
	#include "update_engine/payload_consumer/payload_constants.h"
	#include "update_engine/payload_generator/annotated_operation.h"
	#include "update_engine/payload_generator/delta_diff_generator.h"
	#include "update_engine/payload_generator/extent_ranges.h"
	#include "update_engine/payload_generator/payload_generation_config.h"
	#include "update_engine/update_metadata.pb.h"

	namespace chromeos_update_engine {

	class PartitionWriterTest : public testing::Test {
	public:
	// Helper function to pretend that the ECC file descriptor was already opened.
	// Returns a pointer to the created file descriptor.
	FakeFileDescriptor* SetFakeECCFile(size_t size) {
	EXPECT_FALSE(writer_.verified_source_fd_.source_ecc_fd_)
	<< "source_ecc_fdb already open.";
	FakeFileDescriptor* ret = new FakeFileDescriptor();
	fake_ecc_fd_.reset(ret);
	// Call open to simulate it was already opened.
	ret->Open("", 0);
	ret->SetFileSize(size);
	writer_.verified_source_fd_.source_ecc_fd_ = fake_ecc_fd_;
	return ret;
	}

	uint64_t GetSourceEccRecoveredFailures() const {
	return writer_.verified_source_fd_.source_ecc_recovered_failures_;
	}

	AnnotatedOperation GenerateSourceCopyOp(const brillo::Blob& copied_data,
	bool add_hash,
	PartitionConfig* old_part = nullptr) {
	PayloadGenerationConfig config;
	const uint64_t kDefaultBlockSize = config.block_size;
	EXPECT_EQ(0U, copied_data.size() % kDefaultBlockSize);
	uint64_t num_blocks = copied_data.size() / kDefaultBlockSize;
	AnnotatedOperation aop;
	*(aop.op.add_src_extents()) = ExtentForRange(0, num_blocks);
	*(aop.op.add_dst_extents()) = ExtentForRange(0, num_blocks);
	aop.op.set_type(InstallOperation::SOURCE_COPY);
	brillo::Blob src_hash;
	EXPECT_TRUE(HashCalculator::RawHashOfData(copied_data, &src_hash));
	if (add_hash)
	aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());

	return aop;
	}

	brillo::Blob PerformSourceCopyOp(const InstallOperation& op,
	const brillo::Blob blob_data) {
	LOG(INFO) << "Using source part " << source_partition.path();
	FileDescriptorPtr fd(new EintrSafeFileDescriptor());
	DirectExtentWriter extent_writer{fd};
	EXPECT_TRUE(fd->Open(source_partition.path().c_str(), O_RDWR));
	if (HasFailure()) {
	return {};
	}
	EXPECT_TRUE(extent_writer.Init(op.src_extents(), kBlockSize));
	if (HasFailure()) {
	return {};
	}
	EXPECT_TRUE(extent_writer.Write(blob_data.data(), blob_data.size()));
	if (HasFailure()) {
	return {};
	}
	fd->Flush();

	install_part_.source_size = blob_data.size();
	install_part_.target_size = blob_data.size();

	ErrorCode error;
	EXPECT_TRUE(writer_.Init(&install_plan_, true, 0));
	if (HasFailure()) {
	return {};
	}
	EXPECT_TRUE(writer_.PerformSourceCopyOperation(op, &error));
	writer_.CheckpointUpdateProgress(1);

	brillo::Blob output_data;
	EXPECT_TRUE(utils::ReadFile(target_partition.path(), &output_data));
	return output_data;
	}

	FakePrefs prefs_{};
	InstallPlan install_plan_{};
	InstallPlan::Payload payload_{};
	DynamicPartitionControlStub dynamic_control_{};
	FileDescriptorPtr fake_ecc_fd_{};
	DeltaArchiveManifest manifest_{};
	ScopedTempFile source_partition{"source-part-XXXXXX"};
	ScopedTempFile target_partition{"target-part-XXXXXX"};
	InstallPlan::Partition install_part_{.source_path = source_partition.path(),
	.target_path = target_partition.path()};
	PartitionUpdate partition_update_{};
	PartitionWriter writer_{
	partition_update_, install_part_, &dynamic_control_, kBlockSize, false};
	};
	// Test that the error-corrected file descriptor is used to read a partition
	// when no hash is available for SOURCE_COPY but it falls back to the normal
	// file descriptor when the size of the error corrected one is too small.
	TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyWhenNoHashFallbackTest) {
	constexpr size_t kCopyOperationSize = 4 * 4096;
	ScopedTempFile source("Source-XXXXXX");
	// Setup the source path with the right expected data.
	brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
	ASSERT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));

	// Setup the fec file descriptor as the fake stream, with smaller data than
	// the expected.
	FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize / 2);

	PartitionConfig old_part(kPartitionNameRoot);
	old_part.path = source.path();
	old_part.size = expected_data.size();

	// The payload operation doesn't include an operation hash.
	auto source_copy_op = GenerateSourceCopyOp(expected_data, false, &old_part);
	ASSERT_NO_FATAL_FAILURE();
	auto output_data = PerformSourceCopyOp(source_copy_op.op, expected_data);
	ASSERT_NO_FATAL_FAILURE();
	ASSERT_EQ(output_data, expected_data);

	// Verify that the fake_fec was attempted to be used. Since the file
	// descriptor is shorter it can actually do more than one read to realize it
	// reached the EOF.
	ASSERT_LE(1U, fake_fec->GetReadOps().size());
	// This fallback doesn't count as an error-corrected operation since the
	// operation hash was not available.
	ASSERT_EQ(0U, GetSourceEccRecoveredFailures());
	}

	// Test that the error-corrected file descriptor is used to read the partition
	// since the source partition doesn't match the operation hash.
	TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyFallbackTest) {
	constexpr size_t kCopyOperationSize = 4 * 4096;
	// Write invalid data to the source image, which doesn't match the expected
	// hash.
	brillo::Blob invalid_data(kCopyOperationSize, 0x55);

	// Setup the fec file descriptor as the fake stream, which matches
	// \|expected_data\|.
	FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize);
	brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);

	auto source_copy_op = GenerateSourceCopyOp(expected_data, true);
	ASSERT_NO_FATAL_FAILURE();
	auto output_data = PerformSourceCopyOp(source_copy_op.op, invalid_data);
	ASSERT_NO_FATAL_FAILURE();
	ASSERT_EQ(output_data, expected_data);

	// Verify that the fake_fec was actually used.
	EXPECT_GE(fake_fec->GetReadOps().size(), 1U);
	EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
	}

	TEST_F(PartitionWriterTest, ChooseSourceFDTest) {
	constexpr size_t kSourceSize = 4 * 4096;
	ScopedTempFile source("Source-XXXXXX");
	// Write invalid data to the source image, which doesn't match the expected
	// hash.
	brillo::Blob invalid_data(kSourceSize, 0x55);
	ASSERT_TRUE(test_utils::WriteFileVector(source.path(), invalid_data));

	writer_.verified_source_fd_.source_fd_ =
	std::make_shared<EintrSafeFileDescriptor>();
	writer_.verified_source_fd_.source_fd_->Open(source.path().c_str(), O_RDONLY);

	// Setup the fec file descriptor as the fake stream, which matches
	// \|expected_data\|.
	FakeFileDescriptor* fake_fec = SetFakeECCFile(kSourceSize);
	brillo::Blob expected_data = FakeFileDescriptorData(kSourceSize);

	InstallOperation op;
	*(op.add_src_extents()) = ExtentForRange(0, kSourceSize / 4096);
	brillo::Blob src_hash;
	ASSERT_TRUE(HashCalculator::RawHashOfData(expected_data, &src_hash));
	op.set_src_sha256_hash(src_hash.data(), src_hash.size());

	ErrorCode error = ErrorCode::kSuccess;
	ASSERT_NE(writer_.ChooseSourceFD(op, &error), nullptr);
	ASSERT_EQ(ErrorCode::kSuccess, error);
	// Verify that the fake_fec was actually used.
	ASSERT_EQ(1U, fake_fec->GetReadOps().size());
	ASSERT_EQ(1U, GetSourceEccRecoveredFailures());
	}

	} // namespace chromeos_update_engine