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 "update_engine/payload_consumer/partition_update_generator_android.h"
18 
19 #include <map>
20 #include <memory>
21 #include <set>
22 #include <utility>
23 #include <vector>
24 
25 #include <android-base/strings.h>
26 #include <brillo/secure_blob.h>
27 #include <gtest/gtest.h>
28 
29 #include "update_engine/common/boot_control_interface.h"
30 #include "update_engine/common/fake_boot_control.h"
31 #include "update_engine/common/hash_calculator.h"
32 #include "update_engine/common/test_utils.h"
33 #include "update_engine/common/utils.h"
34 
35 namespace chromeos_update_engine {
36 
37 class FakePartitionUpdateGenerator : public PartitionUpdateGeneratorAndroid {
38  public:
GetAbPartitionsOnDevice() const39   std::vector<std::string> GetAbPartitionsOnDevice() const {
40     return ab_partitions_;
41   }
42   using PartitionUpdateGeneratorAndroid::PartitionUpdateGeneratorAndroid;
43   std::vector<std::string> ab_partitions_;
44 };
45 
46 class PartitionUpdateGeneratorAndroidTest : public ::testing::Test {
47  protected:
SetUp()48   void SetUp() override {
49     ASSERT_TRUE(device_dir_.CreateUniqueTempDir());
50     boot_control_ = std::make_unique<FakeBootControl>();
51     ASSERT_TRUE(boot_control_);
52     boot_control_->SetNumSlots(2);
53     generator_ = std::make_unique<FakePartitionUpdateGenerator>(
54         boot_control_.get(), 4096);
55     ASSERT_TRUE(generator_);
56   }
57 
58   std::unique_ptr<FakePartitionUpdateGenerator> generator_;
59   std::unique_ptr<FakeBootControl> boot_control_;
60 
61   base::ScopedTempDir device_dir_;
62   std::map<std::string, std::string> device_map_;
63 
SetUpBlockDevice(const std::map<std::string,std::string> & contents)64   void SetUpBlockDevice(const std::map<std::string, std::string>& contents) {
65     std::set<std::string> partition_base_names;
66     for (const auto& [name, content] : contents) {
67       auto path = device_dir_.GetPath().value() + "/" + name;
68       ASSERT_TRUE(
69           utils::WriteFile(path.c_str(), content.data(), content.size()));
70 
71       if (android::base::EndsWith(name, "_a")) {
72         auto prefix = name.substr(0, name.size() - 2);
73         boot_control_->SetPartitionDevice(prefix, 0, path);
74         partition_base_names.emplace(prefix);
75       } else if (android::base::EndsWith(name, "_b")) {
76         auto prefix = name.substr(0, name.size() - 2);
77         boot_control_->SetPartitionDevice(prefix, 1, path);
78         partition_base_names.emplace(prefix);
79       }
80       device_map_[name] = std::move(path);
81     }
82     generator_->ab_partitions_ = {partition_base_names.begin(),
83                                   partition_base_names.end()};
84   }
85 
CheckPartitionUpdate(const std::string & name,const std::string & content,const PartitionUpdate & partition_update)86   void CheckPartitionUpdate(const std::string& name,
87                             const std::string& content,
88                             const PartitionUpdate& partition_update) {
89     ASSERT_EQ(name, partition_update.partition_name());
90 
91     brillo::Blob out_hash;
92     ASSERT_TRUE(HashCalculator::RawHashOfBytes(
93         content.data(), content.size(), &out_hash));
94     ASSERT_EQ(std::string(out_hash.begin(), out_hash.end()),
95               partition_update.old_partition_info().hash());
96     ASSERT_EQ(std::string(out_hash.begin(), out_hash.end()),
97               partition_update.new_partition_info().hash());
98 
99     ASSERT_EQ(1, partition_update.operations_size());
100     const auto& operation = partition_update.operations(0);
101     ASSERT_EQ(InstallOperation::SOURCE_COPY, operation.type());
102 
103     ASSERT_EQ(1, operation.src_extents_size());
104     ASSERT_EQ(0u, operation.src_extents(0).start_block());
105     ASSERT_EQ(content.size() / 4096, operation.src_extents(0).num_blocks());
106 
107     ASSERT_EQ(1, operation.dst_extents_size());
108     ASSERT_EQ(0u, operation.dst_extents(0).start_block());
109     ASSERT_EQ(content.size() / 4096, operation.dst_extents(0).num_blocks());
110   }
111 };
112 
TEST_F(PartitionUpdateGeneratorAndroidTest,CreatePartitionUpdate)113 TEST_F(PartitionUpdateGeneratorAndroidTest, CreatePartitionUpdate) {
114   auto system_contents = std::string(4096 * 2, '1');
115   auto boot_contents = std::string(4096 * 5, 'b');
116   std::map<std::string, std::string> contents = {
117       {"system_a", system_contents},
118       {"system_b", std::string(4096 * 2, 0)},
119       {"boot_a", boot_contents},
120       {"boot_b", std::string(4096 * 5, 0)},
121   };
122   SetUpBlockDevice(contents);
123 
124   auto system_partition_update = generator_->CreatePartitionUpdate(
125       "system", device_map_["system_a"], device_map_["system_b"], 4096 * 2);
126   ASSERT_TRUE(system_partition_update.has_value());
127   CheckPartitionUpdate(
128       "system", system_contents, system_partition_update.value());
129 
130   auto boot_partition_update = generator_->CreatePartitionUpdate(
131       "boot", device_map_["boot_a"], device_map_["boot_b"], 4096 * 5);
132   ASSERT_TRUE(boot_partition_update.has_value());
133   CheckPartitionUpdate("boot", boot_contents, boot_partition_update.value());
134 }
135 
TEST_F(PartitionUpdateGeneratorAndroidTest,GenerateOperations)136 TEST_F(PartitionUpdateGeneratorAndroidTest, GenerateOperations) {
137   auto system_contents = std::string(4096 * 10, '2');
138   auto boot_contents = std::string(4096 * 5, 'b');
139   std::map<std::string, std::string> contents = {
140       {"system_a", system_contents},
141       {"system_b", std::string(4096 * 10, 0)},
142       {"boot_a", boot_contents},
143       {"boot_b", std::string(4096 * 5, 0)},
144       {"vendor_a", ""},
145       {"vendor_b", ""},
146       {"persist", ""},
147   };
148   SetUpBlockDevice(contents);
149 
150   std::vector<PartitionUpdate> update_list;
151   ASSERT_TRUE(generator_->GenerateOperationsForPartitionsNotInPayload(
152       0, 1, std::set<std::string>{"vendor"}, &update_list));
153 
154   ASSERT_EQ(2u, update_list.size());
155   CheckPartitionUpdate("boot", boot_contents, update_list[0]);
156   CheckPartitionUpdate("system", system_contents, update_list[1]);
157 }
158 
159 }  // namespace chromeos_update_engine
160