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