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 <algorithm>
6 #include <fstream>
7 #include <iostream>
8 #include <sstream>
9
10 #ifdef USE_BRILLO
11 #include "brillo/flag_helper.h"
12 #else
13 #include "gflags/gflags.h"
14 #endif
15
16 #include "puffin/src/extent_stream.h"
17 #include "puffin/src/file_stream.h"
18 #include "puffin/src/include/puffin/common.h"
19 #include "puffin/src/include/puffin/huffer.h"
20 #include "puffin/src/include/puffin/puffdiff.h"
21 #include "puffin/src/include/puffin/puffer.h"
22 #include "puffin/src/include/puffin/puffpatch.h"
23 #include "puffin/src/include/puffin/utils.h"
24 #include "puffin/src/memory_stream.h"
25 #include "puffin/src/puffin_stream.h"
26 #include "puffin/src/set_errors.h"
27
28 using puffin::BitExtent;
29 using puffin::Buffer;
30 using puffin::ByteExtent;
31 using puffin::Error;
32 using puffin::ExtentStream;
33 using puffin::FileStream;
34 using puffin::Huffer;
35 using puffin::MemoryStream;
36 using puffin::Puffer;
37 using puffin::PuffinStream;
38 using puffin::UniqueStreamPtr;
39 using std::string;
40 using std::vector;
41
42 namespace {
43
44 constexpr char kExtentDelimeter = ',';
45 constexpr char kOffsetLengthDelimeter = ':';
46
47 template <typename T>
StringToExtents(const string & str)48 vector<T> StringToExtents(const string& str) {
49 vector<T> extents;
50 if (!str.empty()) {
51 std::stringstream ss(str);
52 string extent_str;
53 while (getline(ss, extent_str, kExtentDelimeter)) {
54 std::stringstream extent_ss(extent_str);
55 string offset_str, length_str;
56 getline(extent_ss, offset_str, kOffsetLengthDelimeter);
57 getline(extent_ss, length_str, kOffsetLengthDelimeter);
58 extents.emplace_back(stoull(offset_str), stoull(length_str));
59 }
60 }
61 return extents;
62 }
63
64 const uint64_t kDefaultPuffCacheSize = 50 * 1024 * 1024; // 50 MB
65
66 // An enum representing the type of compressed files.
67 enum class FileType { kDeflate, kZlib, kGzip, kZip, kRaw, kUnknown };
68
69 // Returns a file type based on the input string |file_type| (normally the final
70 // extension of the file).
StringToFileType(const string & file_type)71 FileType StringToFileType(const string& file_type) {
72 if (file_type == "raw") {
73 return FileType::kRaw;
74 }
75 if (file_type == "deflate") {
76 return FileType::kDeflate;
77 } else if (file_type == "zlib") {
78 return FileType::kZlib;
79 } else if (file_type == "gzip" || file_type == "gz" || file_type == "tgz") {
80 return FileType::kGzip;
81 } else if (file_type == "zip" || file_type == "apk" || file_type == "jar") {
82 return FileType::kZip;
83 }
84 return FileType::kUnknown;
85 }
86
87 // Finds the location of deflates in |stream|. If |file_type_to_override| is
88 // non-empty, it infers the file type based on that, otherwise, it infers the
89 // file type based on the final extension of |file_name|. It returns false if
90 // file type cannot be inferred from any of the input arguments. |deflates|
91 // is filled with byte-aligned location of deflates.
LocateDeflatesBasedOnFileType(const UniqueStreamPtr & stream,const string & file_name,const string & file_type_to_override,vector<ByteExtent> * deflates)92 bool LocateDeflatesBasedOnFileType(const UniqueStreamPtr& stream,
93 const string& file_name,
94 const string& file_type_to_override,
95 vector<ByteExtent>* deflates) {
96 auto file_type = FileType::kUnknown;
97
98 auto last_dot = file_name.find_last_of(".");
99 if (last_dot == string::npos) {
100 // Could not find a dot so we assume there is no extension.
101 return false;
102 }
103 auto extension = file_name.substr(last_dot + 1);
104 file_type = StringToFileType(extension);
105
106 if (!file_type_to_override.empty()) {
107 auto override_file_type = StringToFileType(file_type_to_override);
108 if (override_file_type == FileType::kUnknown) {
109 LOG(ERROR) << "Overriden file type " << file_type_to_override
110 << " does not exist.";
111 return false;
112 }
113 if (file_type != FileType::kUnknown && file_type != override_file_type) {
114 LOG(WARNING) << "Based on the file name, the file type is " << extension
115 << ", But the overriden file type is "
116 << file_type_to_override << ". Is this intentional?";
117 }
118 file_type = override_file_type;
119 }
120
121 if (file_type == FileType::kRaw) {
122 // Do not need to populate |deflates|.
123 return true;
124 }
125
126 uint64_t stream_size;
127 TEST_AND_RETURN_FALSE(stream->GetSize(&stream_size));
128 if (file_type == FileType::kDeflate) {
129 // Assume the whole stream is a deflate block.
130 *deflates = {ByteExtent(0, stream_size)};
131 return true;
132 }
133
134 Buffer data(stream_size);
135 TEST_AND_RETURN_FALSE(stream->Read(data.data(), data.size()));
136 switch (file_type) {
137 case FileType::kZlib:
138 TEST_AND_RETURN_FALSE(puffin::LocateDeflatesInZlib(data, deflates));
139 break;
140 case FileType::kGzip:
141 TEST_AND_RETURN_FALSE(puffin::LocateDeflatesInGzip(data, deflates));
142 break;
143 case FileType::kZip:
144 TEST_AND_RETURN_FALSE(puffin::LocateDeflatesInZipArchive(data, deflates));
145 break;
146 default:
147 LOG(ERROR) << "Unknown file type: (" << file_type_to_override << ") nor ("
148 << extension << ").";
149 return false;
150 }
151 // Return the stream to its zero offset in case we used it.
152 TEST_AND_RETURN_FALSE(stream->Seek(0));
153
154 return true;
155 }
156
157 } // namespace
158
159 #define SETUP_FLAGS \
160 DEFINE_string(src_file, "", "Source file"); \
161 DEFINE_string(dst_file, "", "Target file"); \
162 DEFINE_string(patch_file, "", "patch file"); \
163 DEFINE_string( \
164 src_deflates_byte, "", \
165 "Source deflate byte locations in the format offset:length,..."); \
166 DEFINE_string( \
167 dst_deflates_byte, "", \
168 "Target deflate byte locations in the format offset:length,..."); \
169 DEFINE_string( \
170 src_deflates_bit, "", \
171 "Source deflate bit locations in the format offset:length,..."); \
172 DEFINE_string( \
173 dst_deflates_bit, "", \
174 "Target deflatebit locations in the format offset:length,..."); \
175 DEFINE_string(src_puffs, "", \
176 "Source puff locations in the format offset:length,..."); \
177 DEFINE_string(dst_puffs, "", \
178 "Target puff locations in the format offset:length,..."); \
179 DEFINE_string(src_extents, "", \
180 "Source extents in the format of offset:length,..."); \
181 DEFINE_string(dst_extents, "", \
182 "Target extents in the format of offset:length,..."); \
183 DEFINE_string(operation, "", \
184 "Type of the operation: puff, huff, puffdiff, puffpatch, " \
185 "puffhuff"); \
186 DEFINE_string(src_file_type, "", \
187 "Type of the input source file: deflate, gzip, " \
188 "zlib or zip"); \
189 DEFINE_string(dst_file_type, "", \
190 "Same as src_file_type but for the target file"); \
191 DEFINE_bool(verbose, false, \
192 "Logs all the given parameters including internally " \
193 "generated ones"); \
194 DEFINE_uint64(cache_size, kDefaultPuffCacheSize, \
195 "Maximum size to cache the puff stream. Used in puffpatch");
196
197 #ifndef USE_BRILLO
198 SETUP_FLAGS;
199 #endif
200
201 // Main entry point to the application.
main(int argc,char ** argv)202 int main(int argc, char** argv) {
203 #ifdef USE_BRILLO
204 SETUP_FLAGS;
205 brillo::FlagHelper::Init(argc, argv, "Puffin tool");
206 #else
207 // google::InitGoogleLogging(argv[0]);
208 google::ParseCommandLineFlags(&argc, &argv, true);
209 #endif
210
211 TEST_AND_RETURN_VALUE(!FLAGS_operation.empty(), -1);
212 TEST_AND_RETURN_VALUE(!FLAGS_src_file.empty(), -1);
213 TEST_AND_RETURN_VALUE(!FLAGS_dst_file.empty(), -1);
214
215 auto src_deflates_byte = StringToExtents<ByteExtent>(FLAGS_src_deflates_byte);
216 auto dst_deflates_byte = StringToExtents<ByteExtent>(FLAGS_dst_deflates_byte);
217 auto src_deflates_bit = StringToExtents<BitExtent>(FLAGS_src_deflates_bit);
218 auto dst_deflates_bit = StringToExtents<BitExtent>(FLAGS_dst_deflates_bit);
219 auto src_puffs = StringToExtents<ByteExtent>(FLAGS_src_puffs);
220 auto dst_puffs = StringToExtents<ByteExtent>(FLAGS_dst_puffs);
221 auto src_extents = StringToExtents<ByteExtent>(FLAGS_src_extents);
222 auto dst_extents = StringToExtents<ByteExtent>(FLAGS_dst_extents);
223
224 auto src_stream = FileStream::Open(FLAGS_src_file, true, false);
225 TEST_AND_RETURN_VALUE(src_stream, -1);
226 if (!src_extents.empty()) {
227 src_stream =
228 ExtentStream::CreateForRead(std::move(src_stream), src_extents);
229 TEST_AND_RETURN_VALUE(src_stream, -1);
230 }
231
232 if (FLAGS_operation == "puff" || FLAGS_operation == "puffhuff") {
233 TEST_AND_RETURN_VALUE(
234 LocateDeflatesBasedOnFileType(src_stream, FLAGS_src_file,
235 FLAGS_src_file_type, &src_deflates_byte),
236 -1);
237
238 if (src_deflates_bit.empty() && src_deflates_byte.empty()) {
239 LOG(WARNING) << "You should pass source deflates, is this intentional?";
240 }
241 if (src_deflates_bit.empty()) {
242 TEST_AND_RETURN_VALUE(FindDeflateSubBlocks(src_stream, src_deflates_byte,
243 &src_deflates_bit),
244 -1);
245 }
246 TEST_AND_RETURN_VALUE(dst_puffs.empty(), -1);
247 uint64_t dst_puff_size;
248 TEST_AND_RETURN_VALUE(FindPuffLocations(src_stream, src_deflates_bit,
249 &dst_puffs, &dst_puff_size),
250 -1);
251
252 auto dst_stream = FileStream::Open(FLAGS_dst_file, false, true);
253 TEST_AND_RETURN_VALUE(dst_stream, -1);
254 auto puffer = std::make_shared<Puffer>();
255 auto reader =
256 PuffinStream::CreateForPuff(std::move(src_stream), puffer,
257 dst_puff_size, src_deflates_bit, dst_puffs);
258
259 Buffer puff_buffer;
260 auto writer = FLAGS_operation == "puffhuff"
261 ? MemoryStream::CreateForWrite(&puff_buffer)
262 : std::move(dst_stream);
263
264 Buffer buffer(1024 * 1024);
265 uint64_t bytes_wrote = 0;
266 while (bytes_wrote < dst_puff_size) {
267 auto write_size = std::min(static_cast<uint64_t>(buffer.size()),
268 dst_puff_size - bytes_wrote);
269 TEST_AND_RETURN_VALUE(reader->Read(buffer.data(), write_size), -1);
270 TEST_AND_RETURN_VALUE(writer->Write(buffer.data(), write_size), -1);
271 bytes_wrote += write_size;
272 }
273
274 // puffhuff operation puffs a stream and huffs it back to the target stream
275 // to make sure we can get to the original stream.
276 if (FLAGS_operation == "puffhuff") {
277 src_puffs = dst_puffs;
278 dst_deflates_byte = src_deflates_byte;
279 dst_deflates_bit = src_deflates_bit;
280
281 auto read_puff_stream = MemoryStream::CreateForRead(puff_buffer);
282 auto huffer = std::make_shared<Huffer>();
283 auto huff_writer = PuffinStream::CreateForHuff(
284 std::move(dst_stream), huffer, dst_puff_size, dst_deflates_bit,
285 src_puffs);
286
287 uint64_t bytes_read = 0;
288 while (bytes_read < dst_puff_size) {
289 auto read_size = std::min(static_cast<uint64_t>(buffer.size()),
290 dst_puff_size - bytes_read);
291 TEST_AND_RETURN_VALUE(read_puff_stream->Read(buffer.data(), read_size),
292 -1);
293 TEST_AND_RETURN_VALUE(huff_writer->Write(buffer.data(), read_size), -1);
294 bytes_read += read_size;
295 }
296 }
297 } else if (FLAGS_operation == "huff") {
298 if (dst_deflates_bit.empty() && src_puffs.empty()) {
299 LOG(WARNING) << "You should pass source puffs and destination deflates"
300 << ", is this intentional?";
301 }
302 TEST_AND_RETURN_VALUE(src_puffs.size() == dst_deflates_bit.size(), -1);
303 uint64_t src_stream_size;
304 TEST_AND_RETURN_VALUE(src_stream->GetSize(&src_stream_size), -1);
305 auto dst_file = FileStream::Open(FLAGS_dst_file, false, true);
306 TEST_AND_RETURN_VALUE(dst_file, -1);
307
308 auto huffer = std::make_shared<Huffer>();
309 auto dst_stream = PuffinStream::CreateForHuff(std::move(dst_file), huffer,
310 src_stream_size,
311 dst_deflates_bit, src_puffs);
312
313 Buffer buffer(1024 * 1024);
314 uint64_t bytes_read = 0;
315 while (bytes_read < src_stream_size) {
316 auto read_size = std::min(static_cast<uint64_t>(buffer.size()),
317 src_stream_size - bytes_read);
318 TEST_AND_RETURN_VALUE(src_stream->Read(buffer.data(), read_size), -1);
319 TEST_AND_RETURN_VALUE(dst_stream->Write(buffer.data(), read_size), -1);
320 bytes_read += read_size;
321 }
322 } else if (FLAGS_operation == "puffdiff") {
323 auto dst_stream = FileStream::Open(FLAGS_dst_file, true, false);
324 TEST_AND_RETURN_VALUE(dst_stream, -1);
325
326 TEST_AND_RETURN_VALUE(
327 LocateDeflatesBasedOnFileType(src_stream, FLAGS_src_file,
328 FLAGS_src_file_type, &src_deflates_byte),
329 -1);
330 TEST_AND_RETURN_VALUE(
331 LocateDeflatesBasedOnFileType(dst_stream, FLAGS_dst_file,
332 FLAGS_dst_file_type, &dst_deflates_byte),
333 -1);
334
335 if (src_deflates_bit.empty() && src_deflates_byte.empty()) {
336 LOG(WARNING) << "You should pass source deflates, is this intentional?";
337 }
338 if (dst_deflates_bit.empty() && dst_deflates_byte.empty()) {
339 LOG(WARNING) << "You should pass target deflates, is this intentional?";
340 }
341 if (!dst_extents.empty()) {
342 dst_stream =
343 ExtentStream::CreateForWrite(std::move(dst_stream), dst_extents);
344 TEST_AND_RETURN_VALUE(dst_stream, -1);
345 }
346
347 if (src_deflates_bit.empty()) {
348 TEST_AND_RETURN_VALUE(FindDeflateSubBlocks(src_stream, src_deflates_byte,
349 &src_deflates_bit),
350 -1);
351 }
352
353 if (dst_deflates_bit.empty()) {
354 TEST_AND_RETURN_VALUE(FindDeflateSubBlocks(dst_stream, dst_deflates_byte,
355 &dst_deflates_bit),
356 -1);
357 }
358
359 Buffer puffdiff_delta;
360 TEST_AND_RETURN_VALUE(
361 puffin::PuffDiff(std::move(src_stream), std::move(dst_stream),
362 src_deflates_bit, dst_deflates_bit, "/tmp/patch.tmp",
363 &puffdiff_delta),
364 -1);
365 if (FLAGS_verbose) {
366 LOG(INFO) << "patch_size: " << puffdiff_delta.size();
367 }
368 auto patch_stream = FileStream::Open(FLAGS_patch_file, false, true);
369 TEST_AND_RETURN_VALUE(patch_stream, -1);
370 TEST_AND_RETURN_VALUE(
371 patch_stream->Write(puffdiff_delta.data(), puffdiff_delta.size()), -1);
372 } else if (FLAGS_operation == "puffpatch") {
373 auto patch_stream = FileStream::Open(FLAGS_patch_file, true, false);
374 TEST_AND_RETURN_VALUE(patch_stream, -1);
375 uint64_t patch_size;
376 TEST_AND_RETURN_VALUE(patch_stream->GetSize(&patch_size), -1);
377
378 Buffer puffdiff_delta(patch_size);
379 TEST_AND_RETURN_VALUE(
380 patch_stream->Read(puffdiff_delta.data(), puffdiff_delta.size()), -1);
381 auto dst_stream = FileStream::Open(FLAGS_dst_file, false, true);
382 TEST_AND_RETURN_VALUE(dst_stream, -1);
383 if (!dst_extents.empty()) {
384 dst_stream =
385 ExtentStream::CreateForWrite(std::move(dst_stream), dst_extents);
386 TEST_AND_RETURN_VALUE(dst_stream, -1);
387 }
388 // Apply the patch. Use 50MB cache, it should be enough for most of the
389 // operations.
390 TEST_AND_RETURN_VALUE(
391 puffin::PuffPatch(std::move(src_stream), std::move(dst_stream),
392 puffdiff_delta.data(), puffdiff_delta.size(),
393 FLAGS_cache_size), // max_cache_size
394 -1);
395 }
396
397 if (FLAGS_verbose) {
398 LOG(INFO) << "src_deflates_byte: "
399 << puffin::ExtentsToString(src_deflates_byte);
400 LOG(INFO) << "dst_deflates_byte: "
401 << puffin::ExtentsToString(dst_deflates_byte);
402 LOG(INFO) << "src_deflates_bit: "
403 << puffin::ExtentsToString(src_deflates_bit);
404 LOG(INFO) << "dst_deflates_bit: "
405 << puffin::ExtentsToString(dst_deflates_bit);
406 LOG(INFO) << "src_puffs: " << puffin::ExtentsToString(src_puffs);
407 LOG(INFO) << "dst_puffs: " << puffin::ExtentsToString(dst_puffs);
408 LOG(INFO) << "src_extents: " << puffin::ExtentsToString(src_extents);
409 LOG(INFO) << "dst_extents: " << puffin::ExtentsToString(dst_extents);
410 }
411 return 0;
412 }
413