| // Copyright 2021 The Fuchsia Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #include <fcntl.h> |
| #include <sys/mman.h> |
| #include <sys/stat.h> |
| |
| #include <cstdio> |
| #include <set> |
| #include <string> |
| #include <string_view> |
| #include <type_traits> |
| |
| #include <fbl/unique_fd.h> |
| #include <safemath/safe_math.h> |
| #include <src/lib/digest/merkle-tree.h> |
| |
| #include "src/lib/fxl/command_line.h" |
| #include "src/lib/fxl/log_settings_command_line.h" |
| #include "src/storage/blobfs/compression/configs/chunked_compression_params.h" |
| #include "src/storage/tools/blobfs-compression/blobfs-compression.h" |
| |
| namespace { |
| |
| using blobfs::DeliveryBlobType; |
| using ::chunked_compression::CompressionParams; |
| |
| const std::set<std::string_view> kCliOptions = { |
| "source_file", "compressed_file", "type", "disable_size_alignment", "calculate_digest", |
| "help", "verbose", |
| }; |
| |
| zx::result<DeliveryBlobType> DeliveryTypeFromString(const std::string& delivery_type_str) { |
| using DeliveryBlobTypeRaw = std::underlying_type_t<DeliveryBlobType>; |
| const std::set<DeliveryBlobTypeRaw> kSupportedBlobTypes = { |
| static_cast<DeliveryBlobTypeRaw>(DeliveryBlobType::kType1), |
| }; |
| const DeliveryBlobTypeRaw type_raw = |
| safemath::checked_cast<DeliveryBlobTypeRaw>(std::stoul(delivery_type_str)); |
| if (kSupportedBlobTypes.find(type_raw) == kSupportedBlobTypes.cend()) { |
| return zx::error(ZX_ERR_INVALID_ARGS); |
| } |
| return zx::ok(DeliveryBlobType{type_raw}); |
| } |
| |
| void usage(const char* fname) { |
| fprintf(stderr, "Usage: %s [--option1=value --option2 ...]\n\n", fname); |
| fprintf( |
| stderr, |
| "The tool will output the maximum possible compressed file size using the exact same \n" |
| "compression implementation in blobfs. The merkle tree used here is a non-compact merkle \n" |
| "tree as it contributes to a bigger size than a compact merkle tree.\n\n"); |
| fprintf(stderr, "Options:\n"); |
| fprintf(stderr, "--%s=/path/to/file\n %s\n", "source_file", "The file to be compressed."); |
| fprintf(stderr, "--%s=/path/to/file\n %s\n", "compressed_file", |
| "The compressed file output path (override if existing). Unless --type is " |
| "specified, will contain compressed data with zero-padding at the end to ensure the " |
| "compressed file size matches the size in stdout."); |
| fprintf(stderr, "--%s=TYPE\n %s\n", "type", |
| "If specified, use specified type when generating the output. Supported types:" |
| "\n\t1 - Type A: zstd-chunked, default compression level"); |
| fprintf(stderr, "--%s\n %s\n", "disable_size_alignment", |
| "Do not align compressed output with block size. Incompatible with --type."); |
| fprintf(stderr, "--%s=/path/to/file\n %s\n", "calculate_digest", |
| "Calculate the Merkle root/digest of a delivery blob (i.e. one created with --type). " |
| "Other options are ignored if set."); |
| fprintf(stderr, "--%s\n %s\n", "help", "print this usage message."); |
| fprintf(stderr, "--%s\n %s\n", "verbose", "show debugging information."); |
| } |
| |
| // Truncates |fd| to |write_size|, and mmaps the file for writing. |
| // Returns the mapped buffer in |out_write_buf| of length |write_size|. |
| // This method can fail only with user-input-irrelevant errors. |
| zx_status_t MapFileForWriting(const fbl::unique_fd& fd, const char* file, size_t write_size, |
| uint8_t** out_write_buf) { |
| off_t trunc_size; |
| if (!safemath::MakeCheckedNum<size_t>(write_size).Cast<off_t>().AssignIfValid(&trunc_size)) { |
| return ZX_ERR_OUT_OF_RANGE; |
| } |
| |
| if (ftruncate(fd.get(), trunc_size)) { |
| fprintf(stderr, "Failed to truncate '%s': %s\n", file, strerror(errno)); |
| return ZX_ERR_NO_SPACE; |
| } |
| |
| void* data = nullptr; |
| if (write_size > 0) { |
| data = mmap(nullptr, write_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd.get(), 0); |
| if (data == MAP_FAILED) { |
| fprintf(stderr, "mmap failed: %s\n", strerror(errno)); |
| return ZX_ERR_NO_MEMORY; |
| } |
| } |
| |
| *out_write_buf = static_cast<uint8_t*>(data); |
| return ZX_OK; |
| } |
| |
| // Mmaps the |fd| for reading, returning the mapping as a span. |
| zx::result<cpp20::span<const uint8_t>> MapFileForReading(const fbl::unique_fd& fd) { |
| struct stat info; |
| if (fstat(fd.get(), &info) < 0) { |
| fprintf(stderr, "fstat failed: %s\n", strerror(errno)); |
| return zx::error(ZX_ERR_BAD_STATE); |
| } |
| if (!S_ISREG(info.st_mode)) { |
| fprintf(stderr, "Cannot map input: can only map regular files\n"); |
| return zx::error(ZX_ERR_NOT_FILE); |
| } |
| |
| ZX_ASSERT(info.st_size >= 0); |
| const size_t size = info.st_size; |
| const void* data = nullptr; |
| |
| if (size > 0) { |
| data = mmap(nullptr, size, PROT_READ, MAP_SHARED, fd.get(), 0); |
| if (data == MAP_FAILED) { |
| fprintf(stderr, "mmap failed: %s\n", strerror(errno)); |
| return zx::error(ZX_ERR_NO_MEMORY); |
| } |
| } |
| |
| return zx::ok(cpp20::span{static_cast<const uint8_t*>(data), size}); |
| } |
| |
| zx::result<> WriteDataToFile(const fbl::unique_fd& fd, cpp20::span<const uint8_t> data) { |
| size_t written_bytes = 0; |
| ssize_t write_result = 0; |
| while (written_bytes < data.size_bytes()) { |
| write_result = write(fd.get(), data.data() + written_bytes, data.size_bytes() - written_bytes); |
| if (write_result < 0) { |
| fprintf(stderr, "Failed to write blob: %s\n", strerror(errno)); |
| return zx::error(ZX_ERR_IO); |
| } |
| written_bytes += write_result; |
| } |
| return zx::ok(); |
| } |
| |
| zx::result<digest::Digest> HandleCalculateDigest(std::string_view path) { |
| fbl::unique_fd fd(open(path.data(), O_RDONLY)); |
| if (!fd) { |
| fprintf(stderr, "Failed to open %s: %s", path.data(), strerror(errno)); |
| return zx::error(ZX_ERR_BAD_STATE); |
| } |
| zx::result data = MapFileForReading(fd); |
| if (data.is_error()) { |
| fprintf(stderr, "Failed to map delivery blob for reading!\n"); |
| return data.take_error(); |
| } |
| zx::result digest = blobfs::CalculateDeliveryBlobDigest(*data); |
| if (digest.is_error()) { |
| fprintf(stderr, |
| "Failed to calculate delivery blob digest. Ensure the file is a valid delivery blob, " |
| "and that the file is not corrupted.\n"); |
| return digest.take_error(); |
| } |
| return digest; |
| } |
| |
| } // namespace |
| |
| int main(int argc, char** argv) { |
| const auto cl = fxl::CommandLineFromArgcArgv(argc, argv); |
| if (!fxl::SetLogSettingsFromCommandLine(cl)) { |
| return 1; |
| } |
| |
| const bool verbose = cl.HasOption("verbose", nullptr); |
| if (verbose) { |
| printf("Received flags:\n"); |
| for (const auto& option : cl.options()) { |
| printf(" %s = \"%s\"\n", option.name.c_str(), option.value.c_str()); |
| } |
| printf("\n"); |
| } |
| |
| // Check unknown input options. |
| bool printHelp = cl.HasOption("help"); |
| for (const auto& option : cl.options()) { |
| if (kCliOptions.find(option.name) == kCliOptions.end()) { |
| fprintf(stderr, "Error: unknown option \"%s\".\n", option.name.c_str()); |
| printHelp = true; |
| } |
| } |
| if (printHelp) { |
| usage(argv[0]); |
| return 0; |
| } |
| |
| // Handle case where --calculate_digest is specified (other options will be ignored). |
| if (cl.HasOption("calculate_digest")) { |
| std::string delivery_blob_path; |
| ZX_ASSERT(cl.GetOptionValue("calculate_digest", &delivery_blob_path)); |
| zx::result digest = HandleCalculateDigest(delivery_blob_path); |
| if (digest.is_error()) { |
| return digest.error_value(); |
| } |
| fbl::String digest_str = digest->ToString(); |
| fprintf(stdout, "%s\n", digest_str.c_str()); |
| return ZX_OK; |
| } |
| |
| blobfs_compress::CompressionCliOptionStruct options; |
| |
| // Parse required args. |
| if (!cl.HasOption("source_file")) { |
| fprintf(stderr, "Error: missing required option: --source_file\n"); |
| usage(argv[0]); |
| return ZX_ERR_INVALID_ARGS; |
| } |
| ZX_ASSERT(cl.GetOptionValue("source_file", &options.source_file)); |
| options.source_file_fd.reset(open(options.source_file.c_str(), O_RDONLY)); |
| |
| // Parse optional args. |
| if (cl.HasOption("disable_size_alignment") && cl.HasOption("type")) { |
| usage(argv[0]); |
| return ZX_ERR_INVALID_ARGS; |
| } |
| options.disable_size_alignment = cl.HasOption("disable_size_alignment"); |
| |
| if (cl.HasOption("type")) { |
| if (!cl.HasOption("compressed_file")) { |
| fprintf(stderr, "Error: --compressed-file must be specified with --type.\n"); |
| usage(argv[0]); |
| return ZX_ERR_INVALID_ARGS; |
| } |
| std::string delivery_blob_type_option; |
| ZX_ASSERT(cl.GetOptionValue("type", &delivery_blob_type_option)); |
| zx::result delivery_type = DeliveryTypeFromString(delivery_blob_type_option); |
| if (delivery_type.is_error()) { |
| fprintf(stderr, |
| "Error: unrecognized or invalid value for --type. See usage " |
| "for list of supported delivery types.\n"); |
| usage(argv[0]); |
| return delivery_type.error_value(); |
| } |
| options.type = delivery_type.value(); |
| } |
| |
| if (cl.HasOption("compressed_file")) { |
| ZX_ASSERT(cl.GetOptionValue("compressed_file", &options.compressed_file)); |
| options.compressed_file_fd.reset( |
| open(options.compressed_file.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0644)); |
| } |
| |
| zx_status_t error_code = blobfs_compress::ValidateCliOptions(options); |
| if (error_code) { |
| usage(argv[0]); |
| return error_code; |
| } |
| |
| zx::result source = MapFileForReading(options.source_file_fd); |
| if (source.is_error()) { |
| return source.error_value(); |
| } |
| |
| // We need to generate a delivery blob. |
| if (options.type.has_value()) { |
| ZX_ASSERT(!options.compressed_file.empty() && options.compressed_file_fd.is_valid()); |
| const zx::result delivery_blob = blobfs_compress::GenerateDeliveryBlob(*source, *options.type); |
| if (delivery_blob.is_error()) { |
| fprintf(stderr, "Error generating delivery blob.\n"); |
| return delivery_blob.error_value(); |
| } |
| const zx::result write_result = |
| WriteDataToFile(options.compressed_file_fd, {delivery_blob->data(), delivery_blob->size()}); |
| return write_result.status_value(); |
| } |
| |
| uint8_t* dest_data = nullptr; |
| CompressionParams params = blobfs::GetDefaultChunkedCompressionParams(source->size()); |
| if (!options.compressed_file.empty()) { |
| const size_t dest_buffer_size = |
| params.ComputeOutputSizeLimit(source->size()) + |
| digest::CalculateMerkleTreeSize(source->size(), digest::kDefaultNodeSize, false); |
| error_code = MapFileForWriting(options.compressed_file_fd, options.compressed_file.c_str(), |
| dest_buffer_size, &dest_data); |
| if (error_code) { |
| return error_code; |
| } |
| } |
| |
| // Compress the blob and output compressed size, optionally writing data into the mapped buffer. |
| |
| size_t dest_size; |
| if (blobfs_compress::BlobfsCompress(source->data(), source->size(), dest_data, &dest_size, params, |
| options)) { |
| return ZX_ERR_INTERNAL; |
| } |
| |
| if (!options.compressed_file.empty()) { |
| off_t trunc_size; |
| if (!safemath::MakeCheckedNum<size_t>(dest_size).Cast<off_t>().AssignIfValid(&trunc_size)) { |
| return ZX_ERR_OUT_OF_RANGE; |
| } |
| |
| ftruncate(options.compressed_file_fd.get(), trunc_size); |
| } |
| return ZX_OK; |
| } |