1 // copyright (c) 2019 the android open source project
2 //
3 // licensed under the apache license, version 2.0 (the "license");
4 // you may not use this file except in compliance with the license.
5 // you may obtain a copy of the license at
6 //
7 //      http://www.apache.org/licenses/license-2.0
8 //
9 // unless required by applicable law or agreed to in writing, software
10 // distributed under the license is distributed on an "as is" basis,
11 // without warranties or conditions of any kind, either express or implied.
12 // see the license for the specific language governing permissions and
13 // limitations under the license.
14 
15 #pragma once
16 
17 #include <libsnapshot/cow_compress.h>
18 
19 #include <stdint.h>
20 
21 #include <condition_variable>
22 #include <cstdint>
23 #include <memory>
24 #include <mutex>
25 #include <optional>
26 #include <queue>
27 #include <string>
28 #include <vector>
29 
30 #include <android-base/unique_fd.h>
31 #include <libsnapshot/cow_format.h>
32 #include <libsnapshot/cow_reader.h>
33 
34 namespace android {
35 namespace snapshot {
36 struct CowSizeInfo {
37     uint64_t cow_size;
38     uint64_t op_count_max;
39 };
40 struct CowOptions {
41     uint32_t block_size = 4096;
42     std::string compression;
43 
44     // Maximum number of blocks that can be written.
45     std::optional<uint64_t> max_blocks;
46 
47     // Number of CowOperations in a cluster. 0 for no clustering. Cannot be 1.
48     uint32_t cluster_ops = 1024;
49 
50     bool scratch_space = true;
51 
52     // Preset the number of merged ops. Only useful for testing.
53     uint64_t num_merge_ops = 0;
54 
55     // Number of threads for compression
56     uint16_t num_compress_threads = 0;
57 
58     // Batch write cluster ops
59     bool batch_write = false;
60 
61     // Size of the cow operation buffer; used in v3 only.
62     uint64_t op_count_max = 0;
63 
64     // Compression factor
65     uint64_t compression_factor = 4096;
66 };
67 
68 // Interface for writing to a snapuserd COW. All operations are ordered; merges
69 // will occur in the sequence they were added to the COW.
70 class ICowWriter {
71   public:
72     using FileDescriptor = chromeos_update_engine::FileDescriptor;
73 
~ICowWriter()74     virtual ~ICowWriter() {}
75 
76     // Encode an operation that copies the contents of |old_block| to the
77     // location of |new_block|. 'num_blocks' is the number of contiguous
78     // COPY operations from |old_block| to |new_block|.
79     virtual bool AddCopy(uint64_t new_block, uint64_t old_block, uint64_t num_blocks = 1) = 0;
80 
81     // Encode a sequence of raw blocks. |size| must be a multiple of the block size.
82     virtual bool AddRawBlocks(uint64_t new_block_start, const void* data, size_t size) = 0;
83 
84     // Add a sequence of xor'd blocks. |size| must be a multiple of the block size.
85     virtual bool AddXorBlocks(uint32_t new_block_start, const void* data, size_t size,
86                               uint32_t old_block, uint16_t offset) = 0;
87 
88     // Encode a sequence of zeroed blocks. |size| must be a multiple of the block size.
89     virtual bool AddZeroBlocks(uint64_t new_block_start, uint64_t num_blocks) = 0;
90 
91     // Add a label to the op sequence.
92     virtual bool AddLabel(uint64_t label) = 0;
93 
94     // Add sequence data for op merging. Data is a list of the destination block numbers.
95     virtual bool AddSequenceData(size_t num_ops, const uint32_t* data) = 0;
96 
97     // Flush all pending writes. This must be called before closing the writer
98     // to ensure that the correct headers and footers are written.
99     virtual bool Finalize() = 0;
100 
101     // Return number of bytes the cow image occupies on disk + the size of sequence && ops buffer
102     // The latter two fields are used in v3 cow format and left as 0 for v2 cow format
103     virtual CowSizeInfo GetCowSizeInfo() const = 0;
104 
105     virtual uint32_t GetBlockSize() const = 0;
106     virtual std::optional<uint32_t> GetMaxBlocks() const = 0;
107 
108     // Open an ICowReader for this writer. The reader will be a snapshot of the
109     // current operations in the writer; new writes after OpenReader() will not
110     // be reflected.
111     virtual std::unique_ptr<ICowReader> OpenReader() = 0;
112 
113     // Open a file descriptor. This allows reading and seeing through the cow
114     // as if it were a normal file. The optional source_device must be a valid
115     // path if the CowReader contains any copy or xor operations.
116     virtual std::unique_ptr<FileDescriptor> OpenFileDescriptor(
117             const std::optional<std::string>& source_device) = 0;
118 };
119 
120 class CompressWorker {
121   public:
122     CompressWorker(std::unique_ptr<ICompressor>&& compressor);
123     bool RunThread();
124     void EnqueueCompressBlocks(const void* buffer, size_t block_size, size_t num_blocks);
125     bool GetCompressedBuffers(std::vector<std::vector<uint8_t>>* compressed_buf);
126     void Finalize();
127     static uint32_t GetDefaultCompressionLevel(CowCompressionAlgorithm compression);
128 
129     static bool CompressBlocks(ICompressor* compressor, size_t block_size, const void* buffer,
130                                size_t num_blocks,
131                                std::vector<std::vector<uint8_t>>* compressed_data);
132 
133   private:
134     struct CompressWork {
135         const void* buffer;
136         size_t num_blocks;
137         size_t block_size;
138         bool compression_status = false;
139         std::vector<std::vector<uint8_t>> compressed_data;
140     };
141 
142     std::unique_ptr<ICompressor> compressor_;
143 
144     std::queue<CompressWork> work_queue_;
145     std::queue<CompressWork> compressed_queue_;
146     std::mutex lock_;
147     std::condition_variable cv_;
148     bool stopped_ = false;
149     size_t total_submitted_ = 0;
150     size_t total_processed_ = 0;
151 
152     bool CompressBlocks(const void* buffer, size_t num_blocks, size_t block_size,
153                         std::vector<std::vector<uint8_t>>* compressed_data);
154 };
155 
156 // Create an ICowWriter not backed by any file. This is useful for estimating
157 // the final size of a cow file.
158 std::unique_ptr<ICowWriter> CreateCowEstimator(uint32_t version, const CowOptions& options);
159 
160 // Create an ICowWriter of the given version and options. If a label is given,
161 // the writer is opened in append mode.
162 std::unique_ptr<ICowWriter> CreateCowWriter(uint32_t version, const CowOptions& options,
163                                             android::base::unique_fd&& fd,
164                                             std::optional<uint64_t> label = {});
165 
166 }  // namespace snapshot
167 }  // namespace android
168