1 //
2 // Copyright (C) 2020 The Android Open Source Project
3 //
4 // Licensed under the Apache License, Version 2.0 (the "License");
5 // you may not use this file except in compliance with the License.
6 // You may obtain a copy of the License at
7 //
8 //      http://www.apache.org/licenses/LICENSE-2.0
9 //
10 // Unless required by applicable law or agreed to in writing, software
11 // distributed under the License is distributed on an "AS IS" BASIS,
12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 // See the License for the specific language governing permissions and
14 // limitations under the License.
15 //
16 
17 #include <memory>
18 #include <vector>
19 
20 #include <brillo/secure_blob.h>
21 #include <gtest/gtest.h>
22 
23 #include "update_engine/common/dynamic_partition_control_stub.h"
24 #include "update_engine/common/error_code.h"
25 #include "update_engine/common/fake_prefs.h"
26 #include "update_engine/common/hash_calculator.h"
27 #include "update_engine/common/test_utils.h"
28 #include "update_engine/common/utils.h"
29 #include "update_engine/payload_consumer/delta_performer.h"
30 #include "update_engine/payload_consumer/extent_reader.h"
31 #include "update_engine/payload_consumer/extent_writer.h"
32 #include "update_engine/payload_consumer/fake_file_descriptor.h"
33 #include "update_engine/payload_consumer/file_descriptor.h"
34 #include "update_engine/payload_consumer/install_plan.h"
35 #include "update_engine/payload_generator/annotated_operation.h"
36 #include "update_engine/payload_generator/delta_diff_generator.h"
37 #include "update_engine/payload_generator/extent_ranges.h"
38 #include "update_engine/payload_generator/payload_file.h"
39 #include "update_engine/payload_generator/payload_generation_config.h"
40 #include "update_engine/update_metadata.pb.h"
41 
42 namespace chromeos_update_engine {
43 
44 class PartitionWriterTest : public testing::Test {
45  public:
46   // Helper function to pretend that the ECC file descriptor was already opened.
47   // Returns a pointer to the created file descriptor.
SetFakeECCFile(size_t size)48   FakeFileDescriptor* SetFakeECCFile(size_t size) {
49     EXPECT_FALSE(writer_.source_ecc_fd_) << "source_ecc_fd_ already open.";
50     FakeFileDescriptor* ret = new FakeFileDescriptor();
51     fake_ecc_fd_.reset(ret);
52     // Call open to simulate it was already opened.
53     ret->Open("", 0);
54     ret->SetFileSize(size);
55     writer_.source_ecc_fd_ = fake_ecc_fd_;
56     return ret;
57   }
58 
GetSourceEccRecoveredFailures() const59   uint64_t GetSourceEccRecoveredFailures() const {
60     return writer_.source_ecc_recovered_failures_;
61   }
62 
GenerateSourceCopyOp(const brillo::Blob & copied_data,bool add_hash,PartitionConfig * old_part=nullptr)63   AnnotatedOperation GenerateSourceCopyOp(const brillo::Blob& copied_data,
64                                           bool add_hash,
65                                           PartitionConfig* old_part = nullptr) {
66     PayloadGenerationConfig config;
67     const uint64_t kDefaultBlockSize = config.block_size;
68     EXPECT_EQ(0U, copied_data.size() % kDefaultBlockSize);
69     uint64_t num_blocks = copied_data.size() / kDefaultBlockSize;
70     AnnotatedOperation aop;
71     *(aop.op.add_src_extents()) = ExtentForRange(0, num_blocks);
72     *(aop.op.add_dst_extents()) = ExtentForRange(0, num_blocks);
73     aop.op.set_type(InstallOperation::SOURCE_COPY);
74     brillo::Blob src_hash;
75     EXPECT_TRUE(HashCalculator::RawHashOfData(copied_data, &src_hash));
76     if (add_hash)
77       aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());
78 
79     return aop;
80   }
81 
PerformSourceCopyOp(const InstallOperation & op,const brillo::Blob blob_data)82   brillo::Blob PerformSourceCopyOp(const InstallOperation& op,
83                                    const brillo::Blob blob_data) {
84     ScopedTempFile source_partition("Blob-XXXXXX");
85     FileDescriptorPtr fd(new EintrSafeFileDescriptor());
86     DirectExtentWriter extent_writer{fd};
87     EXPECT_TRUE(fd->Open(source_partition.path().c_str(), O_RDWR));
88     EXPECT_TRUE(extent_writer.Init(op.src_extents(), kBlockSize));
89     EXPECT_TRUE(extent_writer.Write(blob_data.data(), blob_data.size()));
90 
91     ScopedTempFile target_partition("Blob-XXXXXX");
92 
93     install_part_.source_path = source_partition.path();
94     install_part_.target_path = target_partition.path();
95     install_part_.source_size = blob_data.size();
96     install_part_.target_size = blob_data.size();
97 
98     ErrorCode error;
99     EXPECT_TRUE(writer_.Init(&install_plan_, true, 0));
100     EXPECT_TRUE(writer_.PerformSourceCopyOperation(op, &error));
101     writer_.CheckpointUpdateProgress(1);
102 
103     brillo::Blob output_data;
104     EXPECT_TRUE(utils::ReadFile(target_partition.path(), &output_data));
105     return output_data;
106   }
107 
108   FakePrefs prefs_{};
109   InstallPlan install_plan_{};
110   InstallPlan::Payload payload_{};
111   DynamicPartitionControlStub dynamic_control_{};
112   FileDescriptorPtr fake_ecc_fd_{};
113   DeltaArchiveManifest manifest_{};
114   PartitionUpdate partition_update_{};
115   InstallPlan::Partition install_part_{};
116   PartitionWriter writer_{
117       partition_update_, install_part_, &dynamic_control_, kBlockSize, false};
118 };
119 // Test that the error-corrected file descriptor is used to read a partition
120 // when no hash is available for SOURCE_COPY but it falls back to the normal
121 // file descriptor when the size of the error corrected one is too small.
TEST_F(PartitionWriterTest,ErrorCorrectionSourceCopyWhenNoHashFallbackTest)122 TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyWhenNoHashFallbackTest) {
123   constexpr size_t kCopyOperationSize = 4 * 4096;
124   ScopedTempFile source("Source-XXXXXX");
125   // Setup the source path with the right expected data.
126   brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
127   EXPECT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));
128 
129   // Setup the fec file descriptor as the fake stream, with smaller data than
130   // the expected.
131   FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize / 2);
132 
133   PartitionConfig old_part(kPartitionNameRoot);
134   old_part.path = source.path();
135   old_part.size = expected_data.size();
136 
137   // The payload operation doesn't include an operation hash.
138   auto source_copy_op = GenerateSourceCopyOp(expected_data, false, &old_part);
139 
140   auto output_data = PerformSourceCopyOp(source_copy_op.op, expected_data);
141   ASSERT_EQ(output_data, expected_data);
142 
143   // Verify that the fake_fec was attempted to be used. Since the file
144   // descriptor is shorter it can actually do more than one read to realize it
145   // reached the EOF.
146   EXPECT_LE(1U, fake_fec->GetReadOps().size());
147   // This fallback doesn't count as an error-corrected operation since the
148   // operation hash was not available.
149   EXPECT_EQ(0U, GetSourceEccRecoveredFailures());
150 }
151 
152 // Test that the error-corrected file descriptor is used to read the partition
153 // since the source partition doesn't match the operation hash.
TEST_F(PartitionWriterTest,ErrorCorrectionSourceCopyFallbackTest)154 TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyFallbackTest) {
155   constexpr size_t kCopyOperationSize = 4 * 4096;
156   // Write invalid data to the source image, which doesn't match the expected
157   // hash.
158   brillo::Blob invalid_data(kCopyOperationSize, 0x55);
159 
160   // Setup the fec file descriptor as the fake stream, which matches
161   // |expected_data|.
162   FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize);
163   brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
164 
165   auto source_copy_op = GenerateSourceCopyOp(expected_data, true);
166   auto output_data = PerformSourceCopyOp(source_copy_op.op, invalid_data);
167   ASSERT_EQ(output_data, expected_data);
168 
169   // Verify that the fake_fec was actually used.
170   EXPECT_EQ(1U, fake_fec->GetReadOps().size());
171   EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
172 }
173 
TEST_F(PartitionWriterTest,ChooseSourceFDTest)174 TEST_F(PartitionWriterTest, ChooseSourceFDTest) {
175   constexpr size_t kSourceSize = 4 * 4096;
176   ScopedTempFile source("Source-XXXXXX");
177   // Write invalid data to the source image, which doesn't match the expected
178   // hash.
179   brillo::Blob invalid_data(kSourceSize, 0x55);
180   EXPECT_TRUE(test_utils::WriteFileVector(source.path(), invalid_data));
181 
182   writer_.source_fd_ = std::make_shared<EintrSafeFileDescriptor>();
183   writer_.source_fd_->Open(source.path().c_str(), O_RDONLY);
184 
185   // Setup the fec file descriptor as the fake stream, which matches
186   // |expected_data|.
187   FakeFileDescriptor* fake_fec = SetFakeECCFile(kSourceSize);
188   brillo::Blob expected_data = FakeFileDescriptorData(kSourceSize);
189 
190   InstallOperation op;
191   *(op.add_src_extents()) = ExtentForRange(0, kSourceSize / 4096);
192   brillo::Blob src_hash;
193   EXPECT_TRUE(HashCalculator::RawHashOfData(expected_data, &src_hash));
194   op.set_src_sha256_hash(src_hash.data(), src_hash.size());
195 
196   ErrorCode error = ErrorCode::kSuccess;
197   EXPECT_EQ(writer_.source_ecc_fd_, writer_.ChooseSourceFD(op, &error));
198   EXPECT_EQ(ErrorCode::kSuccess, error);
199   // Verify that the fake_fec was actually used.
200   EXPECT_EQ(1U, fake_fec->GetReadOps().size());
201   EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
202 }
203 
204 }  // namespace chromeos_update_engine
205