• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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