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> 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). 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. 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. 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
