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