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/snapshot_extent_writer.h"
18
19 #include <algorithm>
20 #include <cstdint>
21
22 #include <libsnapshot/cow_writer.h>
23
24 #include "update_engine/update_metadata.pb.h"
25
26 namespace chromeos_update_engine {
27
SnapshotExtentWriter(android::snapshot::ICowWriter * cow_writer)28 SnapshotExtentWriter::SnapshotExtentWriter(
29 android::snapshot::ICowWriter* cow_writer)
30 : cow_writer_(cow_writer) {
31 CHECK_NE(cow_writer, nullptr);
32 }
33
~SnapshotExtentWriter()34 SnapshotExtentWriter::~SnapshotExtentWriter() {
35 CHECK(buffer_.empty()) << buffer_.size();
36 }
37
Init(const google::protobuf::RepeatedPtrField<Extent> & extents,uint32_t block_size)38 bool SnapshotExtentWriter::Init(
39 const google::protobuf::RepeatedPtrField<Extent>& extents,
40 uint32_t block_size) {
41 extents_ = extents;
42 cur_extent_idx_ = 0;
43 buffer_.clear();
44 buffer_.reserve(block_size);
45 block_size_ = block_size;
46 return true;
47 }
48
ConsumeWithBuffer(const uint8_t * data,size_t count)49 size_t SnapshotExtentWriter::ConsumeWithBuffer(const uint8_t* data,
50 size_t count) {
51 CHECK_LT(cur_extent_idx_, static_cast<size_t>(extents_.size()));
52 const auto& cur_extent = extents_[cur_extent_idx_];
53 const auto cur_extent_size = cur_extent.num_blocks() * block_size_;
54
55 if (buffer_.empty() && count >= cur_extent_size) {
56 if (!cow_writer_->AddRawBlocks(
57 cur_extent.start_block(), data, cur_extent_size)) {
58 LOG(ERROR) << "AddRawBlocks(" << cur_extent.start_block() << ", " << data
59 << ", " << cur_extent_size << ") failed.";
60 // return value is expected to be greater than 0. Return 0 to signal error
61 // condition
62 return 0;
63 }
64 if (!next_extent()) {
65 CHECK_EQ(count, cur_extent_size)
66 << "Exhausted all blocks, but still have " << count - cur_extent_size
67 << " bytes left";
68 }
69 return cur_extent_size;
70 }
71 CHECK_LT(buffer_.size(), cur_extent_size)
72 << "Data left in buffer should never be >= cur_extent_size, otherwise "
73 "we should have send that data to CowWriter. Buffer size: "
74 << buffer_.size() << " current extent size: " << cur_extent_size;
75 size_t bytes_to_copy =
76 std::min<size_t>(count, cur_extent_size - buffer_.size());
77 CHECK_GT(bytes_to_copy, 0U);
78
79 buffer_.insert(buffer_.end(), data, data + bytes_to_copy);
80 CHECK_LE(buffer_.size(), cur_extent_size);
81
82 if (buffer_.size() == cur_extent_size) {
83 if (!cow_writer_->AddRawBlocks(
84 cur_extent.start_block(), buffer_.data(), buffer_.size())) {
85 LOG(ERROR) << "AddRawBlocks(" << cur_extent.start_block() << ", "
86 << buffer_.data() << ", " << buffer_.size() << ") failed.";
87 return 0;
88 }
89 buffer_.clear();
90 if (!next_extent()) {
91 CHECK_EQ(count, bytes_to_copy) << "Exhausted all blocks, but still have "
92 << count - bytes_to_copy << " bytes left";
93 }
94 }
95 return bytes_to_copy;
96 }
97
98 // Returns true on success.
99 // This will construct a COW_REPLACE operation and forward it to CowWriter. It
100 // is important that caller does not perform SOURCE_COPY operation on this
101 // class, otherwise raw data will be stored. Caller should find ways to use
102 // COW_COPY whenever possible.
Write(const void * bytes,size_t count)103 bool SnapshotExtentWriter::Write(const void* bytes, size_t count) {
104 if (count == 0) {
105 return true;
106 }
107 CHECK_NE(extents_.size(), 0);
108
109 auto data = static_cast<const uint8_t*>(bytes);
110 while (count > 0) {
111 auto bytes_written = ConsumeWithBuffer(data, count);
112 TEST_AND_RETURN_FALSE(bytes_written > 0);
113 data += bytes_written;
114 count -= bytes_written;
115 }
116 return true;
117 }
118
next_extent()119 bool SnapshotExtentWriter::next_extent() {
120 cur_extent_idx_++;
121 return cur_extent_idx_ < static_cast<size_t>(extents_.size());
122 }
123 } // namespace chromeos_update_engine
124