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