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_.verified_source_fd_.source_ecc_fd_)
50 << "source_ecc_fdb already open.";
51 FakeFileDescriptor* ret = new FakeFileDescriptor();
52 fake_ecc_fd_.reset(ret);
53 // Call open to simulate it was already opened.
54 ret->Open("", 0);
55 ret->SetFileSize(size);
56 writer_.verified_source_fd_.source_ecc_fd_ = fake_ecc_fd_;
57 return ret;
58 }
59
GetSourceEccRecoveredFailures() const60 uint64_t GetSourceEccRecoveredFailures() const {
61 return writer_.verified_source_fd_.source_ecc_recovered_failures_;
62 }
63
GenerateSourceCopyOp(const brillo::Blob & copied_data,bool add_hash,PartitionConfig * old_part=nullptr)64 AnnotatedOperation GenerateSourceCopyOp(const brillo::Blob& copied_data,
65 bool add_hash,
66 PartitionConfig* old_part = nullptr) {
67 PayloadGenerationConfig config;
68 const uint64_t kDefaultBlockSize = config.block_size;
69 EXPECT_EQ(0U, copied_data.size() % kDefaultBlockSize);
70 uint64_t num_blocks = copied_data.size() / kDefaultBlockSize;
71 AnnotatedOperation aop;
72 *(aop.op.add_src_extents()) = ExtentForRange(0, num_blocks);
73 *(aop.op.add_dst_extents()) = ExtentForRange(0, num_blocks);
74 aop.op.set_type(InstallOperation::SOURCE_COPY);
75 brillo::Blob src_hash;
76 EXPECT_TRUE(HashCalculator::RawHashOfData(copied_data, &src_hash));
77 if (add_hash)
78 aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());
79
80 return aop;
81 }
82
PerformSourceCopyOp(const InstallOperation & op,const brillo::Blob blob_data)83 brillo::Blob PerformSourceCopyOp(const InstallOperation& op,
84 const brillo::Blob blob_data) {
85 LOG(INFO) << "Using source part " << source_partition.path();
86 FileDescriptorPtr fd(new EintrSafeFileDescriptor());
87 DirectExtentWriter extent_writer{fd};
88 EXPECT_TRUE(fd->Open(source_partition.path().c_str(), O_RDWR));
89 if (HasFailure()) {
90 return {};
91 }
92 EXPECT_TRUE(extent_writer.Init(op.src_extents(), kBlockSize));
93 if (HasFailure()) {
94 return {};
95 }
96 EXPECT_TRUE(extent_writer.Write(blob_data.data(), blob_data.size()));
97 if (HasFailure()) {
98 return {};
99 }
100 fd->Flush();
101
102 install_part_.source_size = blob_data.size();
103 install_part_.target_size = blob_data.size();
104
105 ErrorCode error;
106 EXPECT_TRUE(writer_.Init(&install_plan_, true, 0));
107 if (HasFailure()) {
108 return {};
109 }
110 EXPECT_TRUE(writer_.PerformSourceCopyOperation(op, &error));
111 writer_.CheckpointUpdateProgress(1);
112
113 brillo::Blob output_data;
114 EXPECT_TRUE(utils::ReadFile(target_partition.path(), &output_data));
115 return output_data;
116 }
117
118 FakePrefs prefs_{};
119 InstallPlan install_plan_{};
120 InstallPlan::Payload payload_{};
121 DynamicPartitionControlStub dynamic_control_{};
122 FileDescriptorPtr fake_ecc_fd_{};
123 DeltaArchiveManifest manifest_{};
124 ScopedTempFile source_partition{"source-part-XXXXXX"};
125 ScopedTempFile target_partition{"target-part-XXXXXX"};
126 InstallPlan::Partition install_part_{.source_path = source_partition.path(),
127 .target_path = target_partition.path()};
128 PartitionUpdate partition_update_{};
129 PartitionWriter writer_{
130 partition_update_, install_part_, &dynamic_control_, kBlockSize, false};
131 };
132 // Test that the error-corrected file descriptor is used to read a partition
133 // when no hash is available for SOURCE_COPY but it falls back to the normal
134 // file descriptor when the size of the error corrected one is too small.
TEST_F(PartitionWriterTest,ErrorCorrectionSourceCopyWhenNoHashFallbackTest)135 TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyWhenNoHashFallbackTest) {
136 constexpr size_t kCopyOperationSize = 4 * 4096;
137 ScopedTempFile source("Source-XXXXXX");
138 // Setup the source path with the right expected data.
139 brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
140 ASSERT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));
141
142 // Setup the fec file descriptor as the fake stream, with smaller data than
143 // the expected.
144 FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize / 2);
145
146 PartitionConfig old_part(kPartitionNameRoot);
147 old_part.path = source.path();
148 old_part.size = expected_data.size();
149
150 // The payload operation doesn't include an operation hash.
151 auto source_copy_op = GenerateSourceCopyOp(expected_data, false, &old_part);
152 ASSERT_NO_FATAL_FAILURE();
153 auto output_data = PerformSourceCopyOp(source_copy_op.op, expected_data);
154 ASSERT_NO_FATAL_FAILURE();
155 ASSERT_EQ(output_data, expected_data);
156
157 // Verify that the fake_fec was attempted to be used. Since the file
158 // descriptor is shorter it can actually do more than one read to realize it
159 // reached the EOF.
160 ASSERT_LE(1U, fake_fec->GetReadOps().size());
161 // This fallback doesn't count as an error-corrected operation since the
162 // operation hash was not available.
163 ASSERT_EQ(0U, GetSourceEccRecoveredFailures());
164 }
165
166 // Test that the error-corrected file descriptor is used to read the partition
167 // since the source partition doesn't match the operation hash.
TEST_F(PartitionWriterTest,ErrorCorrectionSourceCopyFallbackTest)168 TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyFallbackTest) {
169 constexpr size_t kCopyOperationSize = 4 * 4096;
170 // Write invalid data to the source image, which doesn't match the expected
171 // hash.
172 brillo::Blob invalid_data(kCopyOperationSize, 0x55);
173
174 // Setup the fec file descriptor as the fake stream, which matches
175 // |expected_data|.
176 FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize);
177 brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
178
179 auto source_copy_op = GenerateSourceCopyOp(expected_data, true);
180 ASSERT_NO_FATAL_FAILURE();
181 auto output_data = PerformSourceCopyOp(source_copy_op.op, invalid_data);
182 ASSERT_NO_FATAL_FAILURE();
183 ASSERT_EQ(output_data, expected_data);
184
185 // Verify that the fake_fec was actually used.
186 EXPECT_GE(fake_fec->GetReadOps().size(), 1U);
187 EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
188 }
189
TEST_F(PartitionWriterTest,ChooseSourceFDTest)190 TEST_F(PartitionWriterTest, ChooseSourceFDTest) {
191 constexpr size_t kSourceSize = 4 * 4096;
192 ScopedTempFile source("Source-XXXXXX");
193 // Write invalid data to the source image, which doesn't match the expected
194 // hash.
195 brillo::Blob invalid_data(kSourceSize, 0x55);
196 ASSERT_TRUE(test_utils::WriteFileVector(source.path(), invalid_data));
197
198 writer_.verified_source_fd_.source_fd_ =
199 std::make_shared<EintrSafeFileDescriptor>();
200 writer_.verified_source_fd_.source_fd_->Open(source.path().c_str(), O_RDONLY);
201
202 // Setup the fec file descriptor as the fake stream, which matches
203 // |expected_data|.
204 FakeFileDescriptor* fake_fec = SetFakeECCFile(kSourceSize);
205 brillo::Blob expected_data = FakeFileDescriptorData(kSourceSize);
206
207 InstallOperation op;
208 *(op.add_src_extents()) = ExtentForRange(0, kSourceSize / 4096);
209 brillo::Blob src_hash;
210 ASSERT_TRUE(HashCalculator::RawHashOfData(expected_data, &src_hash));
211 op.set_src_sha256_hash(src_hash.data(), src_hash.size());
212
213 ErrorCode error = ErrorCode::kSuccess;
214 ASSERT_EQ(writer_.verified_source_fd_.source_ecc_fd_,
215 writer_.ChooseSourceFD(op, &error));
216 ASSERT_EQ(ErrorCode::kSuccess, error);
217 // Verify that the fake_fec was actually used.
218 ASSERT_EQ(1U, fake_fec->GetReadOps().size());
219 ASSERT_EQ(1U, GetSourceEccRecoveredFailures());
220 }
221
222 } // namespace chromeos_update_engine
223