1 // Copyright (C) 2019 The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #pragma once
16 
17 #include <memory>
18 #include <optional>
19 #include <string>
20 #include <unordered_map>
21 #include <unordered_set>
22 
23 #include <android-base/file.h>
24 #include <gmock/gmock.h>
25 #include <gtest/gtest.h>
26 #include <libfiemap/image_manager.h>
27 #include <liblp/mock_property_fetcher.h>
28 #include <liblp/partition_opener.h>
29 #include <libsnapshot/snapshot.h>
30 #include <storage_literals/storage_literals.h>
31 #include <update_engine/update_metadata.pb.h>
32 
33 namespace android {
34 namespace snapshot {
35 
36 using aidl::android::hardware::boot::MergeStatus;
37 using android::fs_mgr::IPropertyFetcher;
38 using android::fs_mgr::MetadataBuilder;
39 using android::fs_mgr::testing::MockPropertyFetcher;
40 using chromeos_update_engine::DeltaArchiveManifest;
41 using chromeos_update_engine::PartitionUpdate;
42 using testing::_;
43 using testing::AssertionResult;
44 using testing::NiceMock;
45 
46 using namespace android::storage_literals;
47 using namespace std::string_literals;
48 
49 // These are not reset between each test because it's expensive to create
50 // these resources (starting+connecting to gsid, zero-filling images).
51 extern std::unique_ptr<SnapshotManager> sm;
52 extern class TestDeviceInfo* test_device;
53 extern std::string fake_super;
54 static constexpr uint64_t kSuperSize = 32_MiB + 4_KiB;
55 static constexpr uint64_t kGroupSize = 32_MiB;
56 
57 // Redirect requests for "super" to our fake super partition.
58 class TestPartitionOpener final : public android::fs_mgr::PartitionOpener {
59   public:
TestPartitionOpener(const std::string & fake_super_path)60     explicit TestPartitionOpener(const std::string& fake_super_path)
61         : fake_super_path_(fake_super_path) {}
62 
63     android::base::unique_fd Open(const std::string& partition_name, int flags) const override;
64     bool GetInfo(const std::string& partition_name,
65                  android::fs_mgr::BlockDeviceInfo* info) const override;
66     std::string GetDeviceString(const std::string& partition_name) const override;
67 
68   private:
69     std::string fake_super_path_;
70 };
71 
72 class TestDeviceInfo : public SnapshotManager::IDeviceInfo {
73   public:
TestDeviceInfo()74     TestDeviceInfo() {}
TestDeviceInfo(const std::string & fake_super)75     explicit TestDeviceInfo(const std::string& fake_super) { set_fake_super(fake_super); }
TestDeviceInfo(const std::string & fake_super,const std::string & slot_suffix)76     TestDeviceInfo(const std::string& fake_super, const std::string& slot_suffix)
77         : TestDeviceInfo(fake_super) {
78         set_slot_suffix(slot_suffix);
79     }
GetMetadataDir()80     std::string GetMetadataDir() const override { return "/metadata/ota/test"s; }
GetSlotSuffix()81     std::string GetSlotSuffix() const override { return slot_suffix_; }
GetOtherSlotSuffix()82     std::string GetOtherSlotSuffix() const override { return slot_suffix_ == "_a" ? "_b" : "_a"; }
GetSuperDevice(uint32_t slot)83     std::string GetSuperDevice([[maybe_unused]] uint32_t slot) const override { return "super"; }
GetPartitionOpener()84     const android::fs_mgr::IPartitionOpener& GetPartitionOpener() const override {
85         return *opener_.get();
86     }
SetBootControlMergeStatus(MergeStatus status)87     bool SetBootControlMergeStatus(MergeStatus status) override {
88         merge_status_ = status;
89         return true;
90     }
IsOverlayfsSetup()91     bool IsOverlayfsSetup() const override { return false; }
IsRecovery()92     bool IsRecovery() const override { return recovery_; }
SetSlotAsUnbootable(unsigned int slot)93     bool SetSlotAsUnbootable(unsigned int slot) override {
94         unbootable_slots_.insert(slot);
95         return true;
96     }
IsTestDevice()97     bool IsTestDevice() const override { return true; }
IsFirstStageInit()98     bool IsFirstStageInit() const override { return first_stage_init_; }
OpenImageManager()99     std::unique_ptr<IImageManager> OpenImageManager() const override {
100         return IDeviceInfo::OpenImageManager("ota/test");
101     }
GetDeviceMapper()102     android::dm::IDeviceMapper& GetDeviceMapper() override {
103         if (dm_) {
104             return *dm_;
105         }
106         return android::dm::DeviceMapper::Instance();
107     }
108 
IsSlotUnbootable(uint32_t slot)109     bool IsSlotUnbootable(uint32_t slot) { return unbootable_slots_.count(slot) != 0; }
110 
set_slot_suffix(const std::string & suffix)111     void set_slot_suffix(const std::string& suffix) { slot_suffix_ = suffix; }
set_fake_super(const std::string & path)112     void set_fake_super(const std::string& path) {
113         opener_ = std::make_unique<TestPartitionOpener>(path);
114     }
set_recovery(bool value)115     void set_recovery(bool value) { recovery_ = value; }
set_first_stage_init(bool value)116     void set_first_stage_init(bool value) { first_stage_init_ = value; }
set_dm(android::dm::IDeviceMapper * dm)117     void set_dm(android::dm::IDeviceMapper* dm) { dm_ = dm; }
118 
merge_status()119     MergeStatus merge_status() const { return merge_status_; }
120 
121   private:
122     std::string slot_suffix_ = "_a";
123     std::unique_ptr<TestPartitionOpener> opener_;
124     MergeStatus merge_status_;
125     bool recovery_ = false;
126     bool first_stage_init_ = false;
127     std::unordered_set<uint32_t> unbootable_slots_;
128     android::dm::IDeviceMapper* dm_ = nullptr;
129 };
130 
131 class DeviceMapperWrapper : public android::dm::IDeviceMapper {
132     using DmDeviceState = android::dm::DmDeviceState;
133     using DmTable = android::dm::DmTable;
134 
135   public:
DeviceMapperWrapper()136     DeviceMapperWrapper() : impl_(android::dm::DeviceMapper::Instance()) {}
DeviceMapperWrapper(android::dm::IDeviceMapper & impl)137     explicit DeviceMapperWrapper(android::dm::IDeviceMapper& impl) : impl_(impl) {}
138 
CreateDevice(const std::string & name,const DmTable & table,std::string * path,const std::chrono::milliseconds & timeout_ms)139     virtual bool CreateDevice(const std::string& name, const DmTable& table, std::string* path,
140                               const std::chrono::milliseconds& timeout_ms) override {
141         return impl_.CreateDevice(name, table, path, timeout_ms);
142     }
GetState(const std::string & name)143     virtual DmDeviceState GetState(const std::string& name) const override {
144         return impl_.GetState(name);
145     }
LoadTable(const std::string & name,const DmTable & table)146     virtual bool LoadTable(const std::string& name, const DmTable& table) {
147         return impl_.LoadTable(name, table);
148     }
LoadTableAndActivate(const std::string & name,const DmTable & table)149     virtual bool LoadTableAndActivate(const std::string& name, const DmTable& table) {
150         return impl_.LoadTableAndActivate(name, table);
151     }
GetTableInfo(const std::string & name,std::vector<TargetInfo> * table)152     virtual bool GetTableInfo(const std::string& name, std::vector<TargetInfo>* table) {
153         return impl_.GetTableInfo(name, table);
154     }
GetTableStatus(const std::string & name,std::vector<TargetInfo> * table)155     virtual bool GetTableStatus(const std::string& name, std::vector<TargetInfo>* table) {
156         return impl_.GetTableStatus(name, table);
157     }
GetTableStatusIma(const std::string & name,std::vector<TargetInfo> * table)158     virtual bool GetTableStatusIma(const std::string& name, std::vector<TargetInfo>* table) {
159         return impl_.GetTableStatusIma(name, table);
160     }
GetDmDevicePathByName(const std::string & name,std::string * path)161     virtual bool GetDmDevicePathByName(const std::string& name, std::string* path) {
162         return impl_.GetDmDevicePathByName(name, path);
163     }
GetDeviceString(const std::string & name,std::string * dev)164     virtual bool GetDeviceString(const std::string& name, std::string* dev) {
165         return impl_.GetDeviceString(name, dev);
166     }
DeleteDeviceIfExists(const std::string & name)167     virtual bool DeleteDeviceIfExists(const std::string& name) {
168         return impl_.DeleteDeviceIfExists(name);
169     }
170 
171   private:
172     android::dm::IDeviceMapper& impl_;
173 };
174 
175 class SnapshotTestPropertyFetcher : public android::fs_mgr::IPropertyFetcher {
176   public:
177     explicit SnapshotTestPropertyFetcher(const std::string& slot_suffix,
178                                          std::unordered_map<std::string, std::string>&& props = {});
179 
180     std::string GetProperty(const std::string& key, const std::string& defaultValue) override;
181     bool GetBoolProperty(const std::string& key, bool defaultValue) override;
182 
183     static void SetUp(const std::string& slot_suffix = "_a") { Reset(slot_suffix); }
TearDown()184     static void TearDown() { Reset("_a"); }
185 
186   private:
Reset(const std::string & slot_suffix)187     static void Reset(const std::string& slot_suffix) {
188         IPropertyFetcher::OverrideForTesting(
189                 std::make_unique<SnapshotTestPropertyFetcher>(slot_suffix));
190     }
191 
192   private:
193     std::unordered_map<std::string, std::string> properties_;
194 };
195 
196 // Helper for error-spam-free cleanup.
197 void DeleteBackingImage(android::fiemap::IImageManager* manager, const std::string& name);
198 
199 // Write some random data to the given device.
200 // If expect_size is not specified, will write until reaching end of the device.
201 // Expect space of |path| is multiple of 4K.
202 bool WriteRandomData(const std::string& path, std::optional<size_t> expect_size = std::nullopt,
203                      std::string* hash = nullptr);
204 std::string HashSnapshot(ICowWriter::FileDescriptor* writer);
205 
206 std::string ToHexString(const uint8_t* buf, size_t len);
207 
208 std::optional<std::string> GetHash(const std::string& path);
209 
210 // Add partitions and groups described by |manifest|.
211 AssertionResult FillFakeMetadata(MetadataBuilder* builder, const DeltaArchiveManifest& manifest,
212                                  const std::string& suffix);
213 
214 // In the update package metadata, set a partition with the given size.
215 void SetSize(PartitionUpdate* partition_update, uint64_t size);
216 
217 // Get partition size from update package metadata.
218 uint64_t GetSize(PartitionUpdate* partition_update);
219 
220 bool IsVirtualAbEnabled();
221 
222 #define SKIP_IF_NON_VIRTUAL_AB()                                                        \
223     do {                                                                                \
224         if (!IsVirtualAbEnabled()) GTEST_SKIP() << "Test for Virtual A/B devices only"; \
225     } while (0)
226 
227 #define RETURN_IF_NON_VIRTUAL_AB_MSG(msg) \
228     do {                                  \
229         if (!IsVirtualAbEnabled()) {      \
230             std::cerr << (msg);           \
231             return;                       \
232         }                                 \
233     } while (0)
234 
235 #define RETURN_IF_NON_VIRTUAL_AB() RETURN_IF_NON_VIRTUAL_AB_MSG("")
236 
237 }  // namespace snapshot
238 }  // namespace android
239