| Kelvin Zhang | 35cff4f | 2021-12-08 16:06:00 -0800 | [diff] [blame] | 1 | // | 
|  | 2 | // Copyright (C) 2021 The Android Open Source Project | 
|  | 3 | // | 
|  | 4 | // Licensed under the Apache License, Version 2.0 (the "License"); | 
|  | 5 | // you may not use this file except in compliance with the License. | 
|  | 6 | // You may obtain a copy of the License at | 
|  | 7 | // | 
|  | 8 | //      http://www.apache.org/licenses/LICENSE-2.0 | 
|  | 9 | // | 
|  | 10 | // Unless required by applicable law or agreed to in writing, software | 
|  | 11 | // distributed under the License is distributed on an "AS IS" BASIS, | 
|  | 12 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | 
|  | 13 | // See the License for the specific language governing permissions and | 
|  | 14 | // limitations under the License. | 
|  | 15 | // | 
|  | 16 |  | 
|  | 17 | #include "lz4diff.h" | 
|  | 18 | #include "lz4diff/lz4patch.h" | 
|  | 19 | #include "lz4diff_compress.h" | 
|  | 20 | #include "update_engine/payload_generator/filesystem_interface.h" | 
|  | 21 | #include "update_engine/payload_generator/erofs_filesystem.h" | 
|  | 22 | #include "update_engine/common/utils.h" | 
|  | 23 |  | 
|  | 24 | using namespace chromeos_update_engine; | 
|  | 25 |  | 
|  | 26 | template <typename T> | 
|  | 27 | std::ostream& operator<<(std::ostream& out, const std::vector<T>& vec) { | 
|  | 28 | if (vec.begin() == vec.end()) { | 
|  | 29 | out << "{}"; | 
|  | 30 | return out; | 
|  | 31 | } | 
|  | 32 | out << "{"; | 
|  | 33 | auto begin = vec.begin(); | 
|  | 34 | out << *begin; | 
|  | 35 | for (const auto& ext : Range{++begin, vec.end()}) { | 
|  | 36 | out << ", " << ext; | 
|  | 37 | } | 
|  | 38 | out << "}"; | 
|  | 39 | return out; | 
|  | 40 | } | 
|  | 41 |  | 
|  | 42 | enum Lz4DiffOp { DIFF, PATCH, TEST }; | 
|  | 43 |  | 
|  | 44 | int ExecuteLz4diff(const char* src_image_path, | 
|  | 45 | const char* dst_image_path, | 
|  | 46 | const FilesystemInterface::File& src_file, | 
|  | 47 | const FilesystemInterface::File& dst_file, | 
|  | 48 | const char* patch_file, | 
|  | 49 | Lz4DiffOp op) { | 
|  | 50 | brillo::Blob src_blob; | 
|  | 51 | CHECK(utils::ReadExtents( | 
|  | 52 | src_image_path, src_file.extents, &src_blob, kBlockSize)); | 
|  | 53 | brillo::Blob dst_blob; | 
|  | 54 | CHECK(utils::ReadExtents( | 
|  | 55 | dst_image_path, dst_file.extents, &dst_blob, kBlockSize)); | 
|  | 56 |  | 
|  | 57 | brillo::Blob lz4diff_patch; | 
|  | 58 | if (op == DIFF || op == TEST) { | 
|  | 59 | Lz4Diff(src_blob, | 
|  | 60 | dst_blob, | 
|  | 61 | src_file.compressed_file_info, | 
|  | 62 | dst_file.compressed_file_info, | 
|  | 63 | &lz4diff_patch); | 
|  | 64 | if (patch_file) { | 
|  | 65 | CHECK(utils::WriteFile( | 
|  | 66 | patch_file, lz4diff_patch.data(), lz4diff_patch.size())); | 
|  | 67 | } | 
|  | 68 | } | 
|  | 69 | if (op == PATCH || op == TEST) { | 
|  | 70 | utils::ReadFile(patch_file, &lz4diff_patch); | 
|  | 71 | Blob actual_target; | 
|  | 72 | CHECK(Lz4Patch( | 
|  | 73 | ToStringView(src_blob), ToStringView(lz4diff_patch), &actual_target)); | 
|  | 74 | if (actual_target != dst_blob) { | 
|  | 75 | LOG(ERROR) << "Final postfixed blob mismatch. " << src_file.name; | 
|  | 76 | } else { | 
|  | 77 | LOG(INFO) << "LZ4patch success. Final blob matches. " << src_file.name; | 
|  | 78 | } | 
|  | 79 | } | 
|  | 80 | return 0; | 
|  | 81 | } | 
|  | 82 |  | 
|  | 83 | int ExecuteLz4diffOp(const char* src_image_path, | 
|  | 84 | const char* dst_image_path, | 
|  | 85 | const char* inode_path, | 
|  | 86 | const char* patch_file, | 
|  | 87 | Lz4DiffOp op) { | 
|  | 88 | auto src_fs = ErofsFilesystem::CreateFromFile(src_image_path); | 
|  | 89 | CHECK_NE(src_fs, nullptr); | 
|  | 90 | auto dst_fs = ErofsFilesystem::CreateFromFile(dst_image_path); | 
|  | 91 | CHECK_NE(dst_fs, nullptr); | 
|  | 92 | std::vector<FilesystemInterface::File> src_files; | 
|  | 93 | CHECK(src_fs->GetFiles(&src_files)); | 
|  | 94 | std::vector<FilesystemInterface::File> dst_files; | 
|  | 95 | CHECK(dst_fs->GetFiles(&dst_files)); | 
|  | 96 | ScopedTempFile temp_patch; | 
|  | 97 | if (patch_file == nullptr && op == TEST) { | 
|  | 98 | patch_file = temp_patch.path().c_str(); | 
|  | 99 | } | 
|  | 100 | if (inode_path == nullptr) { | 
|  | 101 | for (const auto& src_file : src_files) { | 
|  | 102 | auto dst_file = std::find_if( | 
|  | 103 | dst_files.begin(), | 
|  | 104 | dst_files.end(), | 
|  | 105 | [path(src_file.name)](auto&& file) { return file.name == path; }); | 
|  | 106 | int err = ExecuteLz4diff( | 
|  | 107 | src_image_path, dst_image_path, src_file, *dst_file, patch_file, op); | 
|  | 108 | if (err) { | 
|  | 109 | return err; | 
|  | 110 | } | 
|  | 111 | } | 
|  | 112 | return 0; | 
|  | 113 | } | 
|  | 114 | auto src_file = std::find_if( | 
|  | 115 | src_files.begin(), src_files.end(), [inode_path](auto&& file) { | 
|  | 116 | return file.name == inode_path; | 
|  | 117 | }); | 
|  | 118 | if (src_file == src_files.end()) { | 
|  | 119 | LOG(ERROR) << "Failed to find " << inode_path << " in EROFS image" | 
|  | 120 | << src_image_path; | 
|  | 121 | return 2; | 
|  | 122 | } | 
|  | 123 | auto dst_file = std::find_if( | 
|  | 124 | dst_files.begin(), dst_files.end(), [inode_path](auto&& file) { | 
|  | 125 | return file.name == inode_path; | 
|  | 126 | }); | 
|  | 127 | if (dst_file == dst_files.end()) { | 
|  | 128 | LOG(ERROR) << "Failed to find " << inode_path << " in EROFS image" | 
|  | 129 | << dst_image_path; | 
|  | 130 | return 3; | 
|  | 131 | } | 
|  | 132 | return ExecuteLz4diff( | 
|  | 133 | src_image_path, dst_image_path, *src_file, *dst_file, patch_file, op); | 
|  | 134 | } | 
|  | 135 |  | 
|  | 136 | int main(int argc, const char** argv) { | 
|  | 137 | if (argc < 4) { | 
|  | 138 | printf( | 
|  | 139 | "Usage: %s <diff/patch/test> <src EROFS image> <dst EROFS image> " | 
|  | 140 | "...args\n", | 
|  | 141 | argv[0]); | 
|  | 142 | return 2; | 
|  | 143 | } | 
|  | 144 | const char* src_image_path = argv[2]; | 
|  | 145 | const char* dst_image_path = argv[3]; | 
|  | 146 | auto src_fs = ErofsFilesystem::CreateFromFile(src_image_path); | 
|  | 147 | CHECK_NE(src_fs, nullptr); | 
|  | 148 | auto dst_fs = ErofsFilesystem::CreateFromFile(dst_image_path); | 
|  | 149 | CHECK_NE(dst_fs, nullptr); | 
|  | 150 | std::vector<FilesystemInterface::File> src_files; | 
|  | 151 | CHECK(src_fs->GetFiles(&src_files)); | 
|  | 152 | std::vector<FilesystemInterface::File> dst_files; | 
|  | 153 | CHECK(dst_fs->GetFiles(&dst_files)); | 
|  | 154 | std::string_view op = argv[1]; | 
|  | 155 | if (op == "diff") { | 
|  | 156 | if (argc != 6 && argc != 5) { | 
|  | 157 | printf( | 
|  | 158 | "Usage: %s diff <path to src erofs image> <path to dst erofs imaeg> " | 
|  | 159 | "<path " | 
|  | 160 | "of file inside erofs image> [output path]\n", | 
|  | 161 | argv[0]); | 
|  | 162 | return 1; | 
|  | 163 | } | 
|  | 164 | const char* path = argv[4]; | 
|  | 165 | const char* patch_file = argc == 6 ? argv[5] : nullptr; | 
|  | 166 | return ExecuteLz4diffOp( | 
|  | 167 | src_image_path, dst_image_path, path, patch_file, DIFF); | 
|  | 168 | } else if (op == "patch") { | 
|  | 169 | if (argc != 6 && argc != 7) { | 
|  | 170 | printf( | 
|  | 171 | "Usage: %s patch <path to src erofs image> <path to dst erofs imaeg> " | 
|  | 172 | "<path " | 
|  | 173 | "of file inside erofs image> <patch file>\n", | 
|  | 174 | argv[0]); | 
|  | 175 | return 3; | 
|  | 176 | } | 
|  | 177 | const char* inode_path = argv[4]; | 
|  | 178 | const char* patch_file = argv[5]; | 
|  | 179 | return ExecuteLz4diffOp( | 
|  | 180 | src_image_path, dst_image_path, inode_path, patch_file, PATCH); | 
|  | 181 | } else if (op == "test") { | 
|  | 182 | if (argc != 4) { | 
|  | 183 | printf( | 
|  | 184 | "Usage: %s test <path to src erofs image> <path to dst erofs imaeg> " | 
|  | 185 | "\n", | 
|  | 186 | argv[0]); | 
|  | 187 | return 4; | 
|  | 188 | } | 
|  | 189 | return ExecuteLz4diffOp( | 
|  | 190 | src_image_path, dst_image_path, nullptr, nullptr, TEST); | 
|  | 191 | } else { | 
|  | 192 | LOG(ERROR) << "Unrecognized op " << op; | 
|  | 193 | return 4; | 
|  | 194 | } | 
|  | 195 |  | 
|  | 196 | return 0; | 
|  | 197 | } |