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