| // |
| // Copyright (C) 2020 The Android Open Source Project |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| // |
| #include <stdio.h> |
| #include <sys/stat.h> |
| #include <sys/types.h> |
| #include <unistd.h> |
| |
| #include <iostream> |
| #include <memory> |
| #include <string> |
| #include <unordered_map> |
| #include <unordered_set> |
| |
| #include <android-base/file.h> |
| #include <android-base/logging.h> |
| #include <android-base/strings.h> |
| #include <android-base/unique_fd.h> |
| #include <gflags/gflags.h> |
| #include <libsnapshot/cow_writer.h> |
| #include <openssl/sha.h> |
| #include <sparse/sparse.h> |
| #include <ziparchive/zip_archive.h> |
| |
| DEFINE_string(source_tf, "", "Source target files (dir or zip file)"); |
| DEFINE_string(ota_tf, "", "Target files of the build for an OTA"); |
| DEFINE_string(compression, "gz", "Compression (options: none, gz, brotli)"); |
| |
| namespace android { |
| namespace snapshot { |
| |
| using android::base::borrowed_fd; |
| using android::base::unique_fd; |
| |
| static constexpr size_t kBlockSize = 4096; |
| |
| void MyLogger(android::base::LogId, android::base::LogSeverity severity, const char*, const char*, |
| unsigned int, const char* message) { |
| if (severity == android::base::ERROR) { |
| fprintf(stderr, "%s\n", message); |
| } else { |
| fprintf(stdout, "%s\n", message); |
| } |
| } |
| |
| class TargetFilesPackage final { |
| public: |
| explicit TargetFilesPackage(const std::string& path); |
| |
| bool Open(); |
| bool HasFile(const std::string& path); |
| std::unordered_set<std::string> GetDynamicPartitionNames(); |
| unique_fd OpenFile(const std::string& path); |
| unique_fd OpenImage(const std::string& path); |
| |
| private: |
| std::string path_; |
| unique_fd fd_; |
| std::unique_ptr<ZipArchive, decltype(&CloseArchive)> zip_; |
| }; |
| |
| TargetFilesPackage::TargetFilesPackage(const std::string& path) |
| : path_(path), zip_(nullptr, &CloseArchive) {} |
| |
| bool TargetFilesPackage::Open() { |
| fd_.reset(open(path_.c_str(), O_RDONLY)); |
| if (fd_ < 0) { |
| PLOG(ERROR) << "open failed: " << path_; |
| return false; |
| } |
| |
| struct stat s; |
| if (fstat(fd_.get(), &s) < 0) { |
| PLOG(ERROR) << "fstat failed: " << path_; |
| return false; |
| } |
| if (S_ISDIR(s.st_mode)) { |
| return true; |
| } |
| |
| // Otherwise, assume it's a zip file. |
| ZipArchiveHandle handle; |
| if (OpenArchiveFd(fd_.get(), path_.c_str(), &handle, false)) { |
| LOG(ERROR) << "Could not open " << path_ << " as a zip archive."; |
| return false; |
| } |
| zip_.reset(handle); |
| return true; |
| } |
| |
| bool TargetFilesPackage::HasFile(const std::string& path) { |
| if (zip_) { |
| ZipEntry64 entry; |
| return !FindEntry(zip_.get(), path, &entry); |
| } |
| |
| auto full_path = path_ + "/" + path; |
| return access(full_path.c_str(), F_OK) == 0; |
| } |
| |
| unique_fd TargetFilesPackage::OpenFile(const std::string& path) { |
| if (!zip_) { |
| auto full_path = path_ + "/" + path; |
| unique_fd fd(open(full_path.c_str(), O_RDONLY)); |
| if (fd < 0) { |
| PLOG(ERROR) << "open failed: " << full_path; |
| return {}; |
| } |
| return fd; |
| } |
| |
| ZipEntry64 entry; |
| if (FindEntry(zip_.get(), path, &entry)) { |
| LOG(ERROR) << path << " not found in archive: " << path_; |
| return {}; |
| } |
| |
| TemporaryFile temp; |
| if (temp.fd < 0) { |
| PLOG(ERROR) << "mkstemp failed"; |
| return {}; |
| } |
| |
| LOG(INFO) << "Extracting " << path << " from " << path_ << " ..."; |
| if (ExtractEntryToFile(zip_.get(), &entry, temp.fd)) { |
| LOG(ERROR) << "could not extract " << path << " from " << path_; |
| return {}; |
| } |
| if (lseek(temp.fd, 0, SEEK_SET) < 0) { |
| PLOG(ERROR) << "lseek failed"; |
| return {}; |
| } |
| return unique_fd{temp.release()}; |
| } |
| |
| unique_fd TargetFilesPackage::OpenImage(const std::string& path) { |
| auto fd = OpenFile(path); |
| if (fd < 0) { |
| return {}; |
| } |
| |
| LOG(INFO) << "Unsparsing " << path << " ..."; |
| std::unique_ptr<struct sparse_file, decltype(&sparse_file_destroy)> s( |
| sparse_file_import(fd.get(), false, false), &sparse_file_destroy); |
| if (!s) { |
| return fd; |
| } |
| |
| TemporaryFile temp; |
| if (temp.fd < 0) { |
| PLOG(ERROR) << "mkstemp failed"; |
| return {}; |
| } |
| if (sparse_file_write(s.get(), temp.fd, false, false, false) < 0) { |
| LOG(ERROR) << "sparse_file_write failed"; |
| return {}; |
| } |
| if (lseek(temp.fd, 0, SEEK_SET) < 0) { |
| PLOG(ERROR) << "lseek failed"; |
| return {}; |
| } |
| |
| fd.reset(temp.release()); |
| return fd; |
| } |
| |
| std::unordered_set<std::string> TargetFilesPackage::GetDynamicPartitionNames() { |
| auto fd = OpenFile("META/misc_info.txt"); |
| if (fd < 0) { |
| return {}; |
| } |
| |
| std::string contents; |
| if (!android::base::ReadFdToString(fd, &contents)) { |
| PLOG(ERROR) << "read failed"; |
| return {}; |
| } |
| |
| std::unordered_set<std::string> set; |
| |
| auto lines = android::base::Split(contents, "\n"); |
| for (const auto& line : lines) { |
| auto parts = android::base::Split(line, "="); |
| if (parts.size() == 2 && parts[0] == "dynamic_partition_list") { |
| auto partitions = android::base::Split(parts[1], " "); |
| for (const auto& name : partitions) { |
| if (!name.empty()) { |
| set.emplace(name); |
| } |
| } |
| break; |
| } |
| } |
| return set; |
| } |
| |
| class NonAbEstimator final { |
| public: |
| NonAbEstimator(const std::string& ota_tf_path, const std::string& source_tf_path) |
| : ota_tf_path_(ota_tf_path), source_tf_path_(source_tf_path) {} |
| |
| bool Run(); |
| |
| private: |
| bool OpenPackages(); |
| bool AnalyzePartition(const std::string& partition_name); |
| std::unordered_map<std::string, uint64_t> GetBlockMap(borrowed_fd fd); |
| |
| std::string ota_tf_path_; |
| std::string source_tf_path_; |
| std::unique_ptr<TargetFilesPackage> ota_tf_; |
| std::unique_ptr<TargetFilesPackage> source_tf_; |
| uint64_t size_ = 0; |
| }; |
| |
| bool NonAbEstimator::Run() { |
| if (!OpenPackages()) { |
| return false; |
| } |
| |
| auto partitions = ota_tf_->GetDynamicPartitionNames(); |
| if (partitions.empty()) { |
| LOG(ERROR) << "No dynamic partitions found in META/misc_info.txt"; |
| return false; |
| } |
| for (const auto& partition : partitions) { |
| if (!AnalyzePartition(partition)) { |
| return false; |
| } |
| } |
| |
| int64_t size_in_mb = int64_t(double(size_) / 1024.0 / 1024.0); |
| |
| std::cout << "Estimated COW size: " << size_ << " (" << size_in_mb << "MiB)\n"; |
| return true; |
| } |
| |
| bool NonAbEstimator::OpenPackages() { |
| ota_tf_ = std::make_unique<TargetFilesPackage>(ota_tf_path_); |
| if (!ota_tf_->Open()) { |
| return false; |
| } |
| if (!source_tf_path_.empty()) { |
| source_tf_ = std::make_unique<TargetFilesPackage>(source_tf_path_); |
| if (!source_tf_->Open()) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| static std::string SHA256(const std::string& input) { |
| std::string hash(32, '\0'); |
| SHA256_CTX c; |
| SHA256_Init(&c); |
| SHA256_Update(&c, input.data(), input.size()); |
| SHA256_Final(reinterpret_cast<unsigned char*>(hash.data()), &c); |
| return hash; |
| } |
| |
| bool NonAbEstimator::AnalyzePartition(const std::string& partition_name) { |
| auto path = "IMAGES/" + partition_name + ".img"; |
| auto fd = ota_tf_->OpenImage(path); |
| if (fd < 0) { |
| return false; |
| } |
| |
| unique_fd source_fd; |
| uint64_t source_size = 0; |
| std::unordered_map<std::string, uint64_t> source_blocks; |
| if (source_tf_) { |
| auto dap = source_tf_->GetDynamicPartitionNames(); |
| |
| source_fd = source_tf_->OpenImage(path); |
| if (source_fd >= 0) { |
| struct stat s; |
| if (fstat(source_fd.get(), &s)) { |
| PLOG(ERROR) << "fstat failed"; |
| return false; |
| } |
| source_size = s.st_size; |
| |
| std::cout << "Hashing blocks for " << partition_name << "...\n"; |
| source_blocks = GetBlockMap(source_fd); |
| if (source_blocks.empty()) { |
| LOG(ERROR) << "Could not build a block map for source partition: " |
| << partition_name; |
| return false; |
| } |
| } else { |
| if (dap.count(partition_name)) { |
| return false; |
| } |
| LOG(ERROR) << "Warning: " << partition_name |
| << " has no incremental diff since it's not in the source image."; |
| } |
| } |
| |
| TemporaryFile cow; |
| if (cow.fd < 0) { |
| PLOG(ERROR) << "mkstemp failed"; |
| return false; |
| } |
| |
| CowOptions options; |
| options.block_size = kBlockSize; |
| options.compression = FLAGS_compression; |
| |
| auto writer = std::make_unique<CowWriter>(options); |
| if (!writer->Initialize(borrowed_fd{cow.fd})) { |
| LOG(ERROR) << "Could not initialize COW writer"; |
| return false; |
| } |
| |
| LOG(INFO) << "Analyzing " << partition_name << " ..."; |
| |
| std::string zeroes(kBlockSize, '\0'); |
| std::string chunk(kBlockSize, '\0'); |
| std::string src_chunk(kBlockSize, '\0'); |
| uint64_t next_block_number = 0; |
| while (true) { |
| if (!android::base::ReadFully(fd, chunk.data(), chunk.size())) { |
| if (errno) { |
| PLOG(ERROR) << "read failed"; |
| return false; |
| } |
| break; |
| } |
| |
| uint64_t block_number = next_block_number++; |
| if (chunk == zeroes) { |
| if (!writer->AddZeroBlocks(block_number, 1)) { |
| LOG(ERROR) << "Could not add zero block"; |
| return false; |
| } |
| continue; |
| } |
| |
| uint64_t source_offset = block_number * kBlockSize; |
| if (source_fd >= 0 && source_offset <= source_size) { |
| off64_t offset = block_number * kBlockSize; |
| if (android::base::ReadFullyAtOffset(source_fd, src_chunk.data(), src_chunk.size(), |
| offset)) { |
| if (chunk == src_chunk) { |
| continue; |
| } |
| } else if (errno) { |
| PLOG(ERROR) << "pread failed"; |
| return false; |
| } |
| } |
| |
| auto hash = SHA256(chunk); |
| if (auto iter = source_blocks.find(hash); iter != source_blocks.end()) { |
| if (!writer->AddCopy(block_number, iter->second)) { |
| return false; |
| } |
| continue; |
| } |
| |
| if (!writer->AddRawBlocks(block_number, chunk.data(), chunk.size())) { |
| return false; |
| } |
| } |
| |
| if (!writer->Finalize()) { |
| return false; |
| } |
| |
| struct stat s; |
| if (fstat(cow.fd, &s) < 0) { |
| PLOG(ERROR) << "fstat failed"; |
| return false; |
| } |
| |
| size_ += s.st_size; |
| return true; |
| } |
| |
| std::unordered_map<std::string, uint64_t> NonAbEstimator::GetBlockMap(borrowed_fd fd) { |
| std::string chunk(kBlockSize, '\0'); |
| |
| std::unordered_map<std::string, uint64_t> block_map; |
| uint64_t block_number = 0; |
| while (true) { |
| if (!android::base::ReadFully(fd, chunk.data(), chunk.size())) { |
| if (errno) { |
| PLOG(ERROR) << "read failed"; |
| return {}; |
| } |
| break; |
| } |
| auto hash = SHA256(chunk); |
| block_map[hash] = block_number; |
| block_number++; |
| } |
| return block_map; |
| } |
| |
| } // namespace snapshot |
| } // namespace android |
| |
| using namespace android::snapshot; |
| |
| int main(int argc, char** argv) { |
| android::base::InitLogging(argv, android::snapshot::MyLogger); |
| gflags::SetUsageMessage("Estimate VAB disk usage from Non A/B builds"); |
| gflags::ParseCommandLineFlags(&argc, &argv, false); |
| |
| if (FLAGS_ota_tf.empty()) { |
| std::cerr << "Must specify -ota_tf on the command-line." << std::endl; |
| return 1; |
| } |
| |
| NonAbEstimator estimator(FLAGS_ota_tf, FLAGS_source_tf); |
| if (!estimator.Run()) { |
| return 1; |
| } |
| return 0; |
| } |