1 //
2 // Copyright (C) 2009 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 #ifndef UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_WRITER_H_
18 #define UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_WRITER_H_
19 
20 #include <memory>
21 #include <utility>
22 
23 #include <base/logging.h>
24 #include <brillo/secure_blob.h>
25 
26 #include "update_engine/common/utils.h"
27 #include "update_engine/payload_consumer/file_descriptor.h"
28 #include "update_engine/update_metadata.pb.h"
29 
30 // ExtentWriter is an abstract class which synchronously writes to a given
31 // file descriptor at the extents given.
32 
33 namespace chromeos_update_engine {
34 
35 class ExtentWriter {
36  public:
37   ExtentWriter() = default;
~ExtentWriter()38   virtual ~ExtentWriter() {
39     LOG_IF(ERROR, !end_called_) << "End() not called on ExtentWriter.";
40   }
41 
42   // Returns true on success.
43   virtual bool Init(FileDescriptorPtr fd,
44                     const google::protobuf::RepeatedPtrField<Extent>& extents,
45                     uint32_t block_size) = 0;
46 
47   // Returns true on success.
48   virtual bool Write(const void* bytes, size_t count) = 0;
49 
50   // Should be called when all writing is complete. Returns true on success.
51   // The fd is not closed. Caller is responsible for closing it.
End()52   bool End() {
53     end_called_ = true;
54     return EndImpl();
55   }
56   virtual bool EndImpl() = 0;
57  private:
58   bool end_called_{false};
59 };
60 
61 // DirectExtentWriter is probably the simplest ExtentWriter implementation.
62 // It writes the data directly into the extents.
63 
64 class DirectExtentWriter : public ExtentWriter {
65  public:
66   DirectExtentWriter() = default;
67   ~DirectExtentWriter() override = default;
68 
Init(FileDescriptorPtr fd,const google::protobuf::RepeatedPtrField<Extent> & extents,uint32_t block_size)69   bool Init(FileDescriptorPtr fd,
70             const google::protobuf::RepeatedPtrField<Extent>& extents,
71             uint32_t block_size) override {
72     fd_ = fd;
73     block_size_ = block_size;
74     extents_ = extents;
75     cur_extent_ = extents_.begin();
76     return true;
77   }
78   bool Write(const void* bytes, size_t count) override;
EndImpl()79   bool EndImpl() override { return true; }
80 
81  private:
82   FileDescriptorPtr fd_{nullptr};
83 
84   size_t block_size_{0};
85   // Bytes written into |cur_extent_| thus far.
86   uint64_t extent_bytes_written_{0};
87   google::protobuf::RepeatedPtrField<Extent> extents_;
88   // The next call to write should correspond to |cur_extents_|.
89   google::protobuf::RepeatedPtrField<Extent>::iterator cur_extent_;
90 };
91 
92 // Takes an underlying ExtentWriter to which all operations are delegated.
93 // When End() is called, ZeroPadExtentWriter ensures that the total number
94 // of bytes written is a multiple of block_size_. If not, it writes zeros
95 // to pad as needed.
96 
97 class ZeroPadExtentWriter : public ExtentWriter {
98  public:
ZeroPadExtentWriter(std::unique_ptr<ExtentWriter> underlying_extent_writer)99   explicit ZeroPadExtentWriter(
100       std::unique_ptr<ExtentWriter> underlying_extent_writer)
101       : underlying_extent_writer_(std::move(underlying_extent_writer)) {}
102   ~ZeroPadExtentWriter() override = default;
103 
Init(FileDescriptorPtr fd,const google::protobuf::RepeatedPtrField<Extent> & extents,uint32_t block_size)104   bool Init(FileDescriptorPtr fd,
105             const google::protobuf::RepeatedPtrField<Extent>& extents,
106             uint32_t block_size) override {
107     block_size_ = block_size;
108     return underlying_extent_writer_->Init(fd, extents, block_size);
109   }
Write(const void * bytes,size_t count)110   bool Write(const void* bytes, size_t count) override {
111     if (underlying_extent_writer_->Write(bytes, count)) {
112       bytes_written_mod_block_size_ += count;
113       bytes_written_mod_block_size_ %= block_size_;
114       return true;
115     }
116     return false;
117   }
EndImpl()118   bool EndImpl() override {
119     if (bytes_written_mod_block_size_) {
120       const size_t write_size = block_size_ - bytes_written_mod_block_size_;
121       brillo::Blob zeros(write_size, 0);
122       TEST_AND_RETURN_FALSE(underlying_extent_writer_->Write(zeros.data(),
123                                                              write_size));
124     }
125     return underlying_extent_writer_->End();
126   }
127 
128  private:
129   std::unique_ptr<ExtentWriter> underlying_extent_writer_;
130   size_t block_size_{0};
131   size_t bytes_written_mod_block_size_{0};
132 };
133 
134 }  // namespace chromeos_update_engine
135 
136 #endif  // UPDATE_ENGINE_PAYLOAD_CONSUMER_EXTENT_WRITER_H_
137