1 // Copyright 2017 The Chromium OS Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "bsdiff/patch_writer.h"
6
7 #include <string.h>
8 #include <limits>
9
10 #include "bsdiff/brotli_compressor.h"
11 #include "bsdiff/bz2_compressor.h"
12 #include "bsdiff/constants.h"
13 #include "bsdiff/control_entry.h"
14 #include "bsdiff/logging.h"
15
16 namespace {
17
18
19 } // namespace
20
21 namespace bsdiff {
22
BsdiffPatchWriter(const std::string & patch_filename)23 BsdiffPatchWriter::BsdiffPatchWriter(const std::string& patch_filename)
24 : patch_filename_(patch_filename),
25 format_(BsdiffFormat::kLegacy),
26 brotli_quality_(-1) {
27 types_.emplace_back(CompressorType::kBZ2);
28 }
29
30
BsdiffPatchWriter(const std::string & patch_filename,const std::vector<CompressorType> & types,int brotli_quality)31 BsdiffPatchWriter::BsdiffPatchWriter(const std::string& patch_filename,
32 const std::vector<CompressorType>& types,
33 int brotli_quality)
34 : patch_filename_(patch_filename),
35 format_(BsdiffFormat::kBsdf2),
36 types_(types),
37 brotli_quality_(brotli_quality) {}
38
InitializeCompressorList(std::vector<std::unique_ptr<bsdiff::CompressorInterface>> * compressor_list)39 bool BsdiffPatchWriter::InitializeCompressorList(
40 std::vector<std::unique_ptr<bsdiff::CompressorInterface>>*
41 compressor_list) {
42 if (types_.empty()) {
43 LOG(ERROR) << "Patch writer expects at least one compressor.";
44 return false;
45 }
46
47 for (const auto& type : types_) {
48 switch (type) {
49 case CompressorType::kBZ2:
50 compressor_list->emplace_back(new BZ2Compressor());
51 break;
52 case CompressorType::kBrotli:
53 compressor_list->emplace_back(new BrotliCompressor(brotli_quality_));
54 break;
55 case CompressorType::kNoCompression:
56 LOG(ERROR) << "Unsupported compression type " << static_cast<int>(type);
57 return false;
58 }
59 }
60
61 for (const auto& compressor : *compressor_list) {
62 if (!compressor) {
63 return false;
64 }
65 }
66
67 return true;
68 }
69
Init(size_t)70 bool BsdiffPatchWriter::Init(size_t /* new_size */) {
71 if (!InitializeCompressorList(&ctrl_stream_list_)) {
72 LOG(ERROR) << "Failed to initialize control stream compressors.";
73 return false;
74 }
75
76 if (!InitializeCompressorList(&diff_stream_list_)) {
77 LOG(ERROR) << "Failed to initialize diff stream compressors.";
78 return false;
79 }
80
81 if (!InitializeCompressorList(&extra_stream_list_)) {
82 LOG(ERROR) << "Failed to initialize extra stream compressors.";
83 return false;
84 }
85
86 fp_ = fopen(patch_filename_.c_str(), "w");
87 if (!fp_) {
88 LOG(ERROR) << "Opening " << patch_filename_;
89 return false;
90 }
91 return true;
92 }
93
WriteDiffStream(const uint8_t * data,size_t size)94 bool BsdiffPatchWriter::WriteDiffStream(const uint8_t* data, size_t size) {
95 for (const auto& compressor : diff_stream_list_) {
96 if (!compressor->Write(data, size)) {
97 return false;
98 }
99 }
100
101 return true;
102 }
103
WriteExtraStream(const uint8_t * data,size_t size)104 bool BsdiffPatchWriter::WriteExtraStream(const uint8_t* data, size_t size) {
105 for (const auto& compressor : extra_stream_list_) {
106 if (!compressor->Write(data, size)) {
107 return false;
108 }
109 }
110
111 return true;
112 }
113
AddControlEntry(const ControlEntry & entry)114 bool BsdiffPatchWriter::AddControlEntry(const ControlEntry& entry) {
115 // Generate the 24 byte control entry.
116 uint8_t buf[24];
117 EncodeInt64(entry.diff_size, buf);
118 EncodeInt64(entry.extra_size, buf + 8);
119 EncodeInt64(entry.offset_increment, buf + 16);
120
121 for (const auto& compressor : ctrl_stream_list_) {
122 if (!compressor->Write(buf, sizeof(buf))) {
123 return false;
124 }
125 }
126
127 written_output_ += entry.diff_size + entry.extra_size;
128 return true;
129 }
130
SelectSmallestResult(const std::vector<std::unique_ptr<CompressorInterface>> & compressor_list,CompressorInterface ** smallest_compressor)131 bool BsdiffPatchWriter::SelectSmallestResult(
132 const std::vector<std::unique_ptr<CompressorInterface>>& compressor_list,
133 CompressorInterface** smallest_compressor) {
134 size_t min_size = std::numeric_limits<size_t>::max();
135 for (const auto& compressor : compressor_list) {
136 if (!compressor->Finish()) {
137 LOG(ERROR) << "Failed to finalize compressed streams.";
138 return false;
139 }
140
141 // Update |smallest_compressor| if the current compressor produces a
142 // smaller result.
143 if (compressor->GetCompressedData().size() < min_size) {
144 min_size = compressor->GetCompressedData().size();
145 *smallest_compressor = compressor.get();
146 }
147 }
148
149 return true;
150 }
151
Close()152 bool BsdiffPatchWriter::Close() {
153 if (!fp_) {
154 LOG(ERROR) << "File not open.";
155 return false;
156 }
157
158 CompressorInterface* ctrl_stream = nullptr;
159 if (!SelectSmallestResult(ctrl_stream_list_, &ctrl_stream) || !ctrl_stream) {
160 return false;
161 }
162
163 CompressorInterface* diff_stream = nullptr;
164 if (!SelectSmallestResult(diff_stream_list_, &diff_stream) || !diff_stream) {
165 return false;
166 }
167
168 CompressorInterface* extra_stream = nullptr;
169 if (!SelectSmallestResult(extra_stream_list_, &extra_stream) ||
170 !extra_stream) {
171 return false;
172 }
173
174 auto ctrl_data = ctrl_stream->GetCompressedData();
175 auto diff_data = diff_stream->GetCompressedData();
176 auto extra_data = extra_stream->GetCompressedData();
177
178 uint8_t types[3] = {static_cast<uint8_t>(ctrl_stream->Type()),
179 static_cast<uint8_t>(diff_stream->Type()),
180 static_cast<uint8_t>(extra_stream->Type())};
181
182 if (!WriteHeader(types, ctrl_data.size(), diff_data.size()))
183 return false;
184
185 if (fwrite(ctrl_data.data(), 1, ctrl_data.size(), fp_) != ctrl_data.size()) {
186 LOG(ERROR) << "Writing ctrl_data.";
187 return false;
188 }
189 if (fwrite(diff_data.data(), 1, diff_data.size(), fp_) != diff_data.size()) {
190 LOG(ERROR) << "Writing diff_data.";
191 return false;
192 }
193 if (fwrite(extra_data.data(), 1, extra_data.size(), fp_) !=
194 extra_data.size()) {
195 LOG(ERROR) << "Writing extra_data.";
196 return false;
197 }
198 if (fclose(fp_) != 0) {
199 LOG(ERROR) << "Closing the patch file.";
200 return false;
201 }
202 fp_ = nullptr;
203 return true;
204 }
205
WriteHeader(uint8_t types[3],uint64_t ctrl_size,uint64_t diff_size)206 bool BsdiffPatchWriter::WriteHeader(uint8_t types[3],
207 uint64_t ctrl_size,
208 uint64_t diff_size) {
209 /* Header format is
210 * 0 8 magic header
211 * 8 8 length of compressed ctrl block
212 * 16 8 length of compressed diff block
213 * 24 8 length of new file
214 *
215 * File format is
216 * 0 32 Header
217 * 32 ?? compressed ctrl block
218 * ?? ?? compressed diff block
219 * ?? ?? compressed extra block
220 */
221 uint8_t header[32];
222 if (format_ == BsdiffFormat::kLegacy) {
223 // The magic header is "BSDIFF40" for legacy format.
224 memcpy(header, kLegacyMagicHeader, 8);
225 } else if (format_ == BsdiffFormat::kBsdf2) {
226 // The magic header for BSDF2 format:
227 // 0 5 BSDF2
228 // 5 1 compressed type for control stream
229 // 6 1 compressed type for diff stream
230 // 7 1 compressed type for extra stream
231 memcpy(header, kBSDF2MagicHeader, 5);
232 memcpy(header + 5, types, 3);
233 } else {
234 LOG(ERROR) << "Unsupported bsdiff format.";
235 return false;
236 }
237
238 EncodeInt64(ctrl_size, header + 8);
239 EncodeInt64(diff_size, header + 16);
240 EncodeInt64(written_output_, header + 24);
241 if (fwrite(header, sizeof(header), 1, fp_) != 1) {
242 LOG(ERROR) << "writing to the patch file";
243 return false;
244 }
245 return true;
246 }
247
248 } // namespace bsdiff
249