WIABlob.cpp 77 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088
  1. // Copyright 2018 Dolphin Emulator Project
  2. // SPDX-License-Identifier: GPL-2.0-or-later
  3. #include "DiscIO/WIABlob.h"
  4. #include <algorithm>
  5. #include <array>
  6. #include <cstring>
  7. #include <limits>
  8. #include <map>
  9. #include <memory>
  10. #include <mutex>
  11. #include <optional>
  12. #include <type_traits>
  13. #include <utility>
  14. #include <fmt/format.h>
  15. #include <zstd.h>
  16. #include "Common/Align.h"
  17. #include "Common/Assert.h"
  18. #include "Common/CommonTypes.h"
  19. #include "Common/Crypto/SHA1.h"
  20. #include "Common/FileUtil.h"
  21. #include "Common/IOFile.h"
  22. #include "Common/Logging/Log.h"
  23. #include "Common/MsgHandler.h"
  24. #include "Common/ScopeGuard.h"
  25. #include "Common/Swap.h"
  26. #include "DiscIO/Blob.h"
  27. #include "DiscIO/DiscUtils.h"
  28. #include "DiscIO/Filesystem.h"
  29. #include "DiscIO/LaggedFibonacciGenerator.h"
  30. #include "DiscIO/MultithreadedCompressor.h"
  31. #include "DiscIO/Volume.h"
  32. #include "DiscIO/VolumeWii.h"
  33. #include "DiscIO/WIACompression.h"
  34. #include "DiscIO/WiiEncryptionCache.h"
  35. namespace DiscIO
  36. {
  37. static void PushBack(std::vector<u8>* vector, const u8* begin, const u8* end)
  38. {
  39. const size_t offset_in_vector = vector->size();
  40. vector->resize(offset_in_vector + (end - begin));
  41. std::copy(begin, end, vector->data() + offset_in_vector);
  42. }
  43. template <typename T>
  44. static void PushBack(std::vector<u8>* vector, const T& x)
  45. {
  46. static_assert(std::is_trivially_copyable_v<T>);
  47. const u8* x_ptr = reinterpret_cast<const u8*>(&x);
  48. PushBack(vector, x_ptr, x_ptr + sizeof(T));
  49. }
  50. std::pair<int, int> GetAllowedCompressionLevels(WIARVZCompressionType compression_type, bool gui)
  51. {
  52. switch (compression_type)
  53. {
  54. case WIARVZCompressionType::Bzip2:
  55. case WIARVZCompressionType::LZMA:
  56. case WIARVZCompressionType::LZMA2:
  57. return {1, 9};
  58. case WIARVZCompressionType::Zstd:
  59. // The actual minimum level can be gotten by calling ZSTD_minCLevel(). However, returning that
  60. // would make the UI rather weird, because it is a negative number with very large magnitude.
  61. // Note: Level 0 is a special number which means "default level" (level 3 as of this writing).
  62. if (gui)
  63. return {1, ZSTD_maxCLevel()};
  64. else
  65. return {ZSTD_minCLevel(), ZSTD_maxCLevel()};
  66. default:
  67. return {0, -1};
  68. }
  69. }
  70. template <bool RVZ>
  71. WIARVZFileReader<RVZ>::WIARVZFileReader(File::IOFile file, const std::string& path)
  72. : m_file(std::move(file)), m_path(path), m_encryption_cache(this)
  73. {
  74. m_valid = Initialize(path);
  75. }
  76. template <bool RVZ>
  77. WIARVZFileReader<RVZ>::~WIARVZFileReader() = default;
  78. template <bool RVZ>
  79. bool WIARVZFileReader<RVZ>::Initialize(const std::string& path)
  80. {
  81. if (!m_file.Seek(0, File::SeekOrigin::Begin) || !m_file.ReadArray(&m_header_1, 1))
  82. return false;
  83. if ((!RVZ && m_header_1.magic != WIA_MAGIC) || (RVZ && m_header_1.magic != RVZ_MAGIC))
  84. return false;
  85. const u32 version = RVZ ? RVZ_VERSION : WIA_VERSION;
  86. const u32 version_read_compatible =
  87. RVZ ? RVZ_VERSION_READ_COMPATIBLE : WIA_VERSION_READ_COMPATIBLE;
  88. const u32 file_version = Common::swap32(m_header_1.version);
  89. const u32 file_version_compatible = Common::swap32(m_header_1.version_compatible);
  90. if (version < file_version_compatible || version_read_compatible > file_version)
  91. {
  92. ERROR_LOG_FMT(DISCIO, "Unsupported version {} in {}", VersionToString(file_version), path);
  93. return false;
  94. }
  95. const auto header_1_actual_hash = Common::SHA1::CalculateDigest(
  96. reinterpret_cast<const u8*>(&m_header_1), sizeof(m_header_1) - Common::SHA1::DIGEST_LEN);
  97. if (m_header_1.header_1_hash != header_1_actual_hash)
  98. return false;
  99. if (Common::swap64(m_header_1.wia_file_size) != m_file.GetSize())
  100. {
  101. ERROR_LOG_FMT(DISCIO, "File size is incorrect for {}", path);
  102. return false;
  103. }
  104. const u32 header_2_size = Common::swap32(m_header_1.header_2_size);
  105. const u32 header_2_min_size = sizeof(WIAHeader2) - sizeof(WIAHeader2::compressor_data);
  106. if (header_2_size < header_2_min_size)
  107. return false;
  108. std::vector<u8> header_2(header_2_size);
  109. if (!m_file.ReadBytes(header_2.data(), header_2.size()))
  110. return false;
  111. const auto header_2_actual_hash = Common::SHA1::CalculateDigest(header_2);
  112. if (m_header_1.header_2_hash != header_2_actual_hash)
  113. return false;
  114. std::memcpy(&m_header_2, header_2.data(), std::min(header_2.size(), sizeof(WIAHeader2)));
  115. if (m_header_2.compressor_data_size > sizeof(WIAHeader2::compressor_data) ||
  116. header_2_size < header_2_min_size + m_header_2.compressor_data_size)
  117. {
  118. return false;
  119. }
  120. const u32 chunk_size = Common::swap32(m_header_2.chunk_size);
  121. const auto is_power_of_two = [](u32 x) { return (x & (x - 1)) == 0; };
  122. if ((!RVZ || chunk_size < VolumeWii::BLOCK_TOTAL_SIZE || !is_power_of_two(chunk_size)) &&
  123. chunk_size % VolumeWii::GROUP_TOTAL_SIZE != 0)
  124. {
  125. return false;
  126. }
  127. const u32 compression_type = Common::swap32(m_header_2.compression_type);
  128. m_compression_type = static_cast<WIARVZCompressionType>(compression_type);
  129. if (m_compression_type > (RVZ ? WIARVZCompressionType::Zstd : WIARVZCompressionType::LZMA2) ||
  130. (RVZ && m_compression_type == WIARVZCompressionType::Purge))
  131. {
  132. ERROR_LOG_FMT(DISCIO, "Unsupported compression type {} in {}", compression_type, path);
  133. return false;
  134. }
  135. const size_t number_of_partition_entries = Common::swap32(m_header_2.number_of_partition_entries);
  136. const size_t partition_entry_size = Common::swap32(m_header_2.partition_entry_size);
  137. std::vector<u8> partition_entries(partition_entry_size * number_of_partition_entries);
  138. if (!m_file.Seek(Common::swap64(m_header_2.partition_entries_offset), File::SeekOrigin::Begin))
  139. return false;
  140. if (!m_file.ReadBytes(partition_entries.data(), partition_entries.size()))
  141. return false;
  142. const auto partition_entries_actual_hash = Common::SHA1::CalculateDigest(partition_entries);
  143. if (m_header_2.partition_entries_hash != partition_entries_actual_hash)
  144. return false;
  145. const size_t copy_length = std::min(partition_entry_size, sizeof(PartitionEntry));
  146. const size_t memset_length = sizeof(PartitionEntry) - copy_length;
  147. u8* ptr = partition_entries.data();
  148. m_partition_entries.resize(number_of_partition_entries);
  149. for (size_t i = 0; i < number_of_partition_entries; ++i, ptr += partition_entry_size)
  150. {
  151. std::memcpy(&m_partition_entries[i], ptr, copy_length);
  152. std::memset(reinterpret_cast<u8*>(&m_partition_entries[i]) + copy_length, 0, memset_length);
  153. }
  154. for (size_t i = 0; i < m_partition_entries.size(); ++i)
  155. {
  156. const std::array<PartitionDataEntry, 2>& entries = m_partition_entries[i].data_entries;
  157. size_t non_empty_entries = 0;
  158. for (size_t j = 0; j < entries.size(); ++j)
  159. {
  160. const u32 number_of_sectors = Common::swap32(entries[j].number_of_sectors);
  161. if (number_of_sectors != 0)
  162. {
  163. ++non_empty_entries;
  164. const u32 last_sector = Common::swap32(entries[j].first_sector) + number_of_sectors;
  165. m_data_entries.emplace(last_sector * VolumeWii::BLOCK_TOTAL_SIZE, DataEntry(i, j));
  166. }
  167. }
  168. if (non_empty_entries > 1)
  169. {
  170. if (Common::swap32(entries[0].first_sector) > Common::swap32(entries[1].first_sector))
  171. return false;
  172. }
  173. }
  174. const u32 number_of_raw_data_entries = Common::swap32(m_header_2.number_of_raw_data_entries);
  175. m_raw_data_entries.resize(number_of_raw_data_entries);
  176. Chunk& raw_data_entries =
  177. ReadCompressedData(Common::swap64(m_header_2.raw_data_entries_offset),
  178. Common::swap32(m_header_2.raw_data_entries_size),
  179. number_of_raw_data_entries * sizeof(RawDataEntry), m_compression_type);
  180. if (!raw_data_entries.ReadAll(&m_raw_data_entries))
  181. return false;
  182. for (size_t i = 0; i < m_raw_data_entries.size(); ++i)
  183. {
  184. const RawDataEntry& entry = m_raw_data_entries[i];
  185. const u64 data_size = Common::swap64(entry.data_size);
  186. if (data_size != 0)
  187. m_data_entries.emplace(Common::swap64(entry.data_offset) + data_size, DataEntry(i));
  188. }
  189. const u32 number_of_group_entries = Common::swap32(m_header_2.number_of_group_entries);
  190. m_group_entries.resize(number_of_group_entries);
  191. Chunk& group_entries =
  192. ReadCompressedData(Common::swap64(m_header_2.group_entries_offset),
  193. Common::swap32(m_header_2.group_entries_size),
  194. number_of_group_entries * sizeof(GroupEntry), m_compression_type);
  195. if (!group_entries.ReadAll(&m_group_entries))
  196. return false;
  197. if (HasDataOverlap())
  198. return false;
  199. return true;
  200. }
  201. template <bool RVZ>
  202. bool WIARVZFileReader<RVZ>::HasDataOverlap() const
  203. {
  204. for (size_t i = 0; i < m_partition_entries.size(); ++i)
  205. {
  206. const std::array<PartitionDataEntry, 2>& entries = m_partition_entries[i].data_entries;
  207. for (size_t j = 0; j < entries.size(); ++j)
  208. {
  209. if (Common::swap32(entries[j].number_of_sectors) == 0)
  210. continue;
  211. const u64 data_offset = Common::swap32(entries[j].first_sector) * VolumeWii::BLOCK_TOTAL_SIZE;
  212. const auto it = m_data_entries.upper_bound(data_offset);
  213. if (it == m_data_entries.end())
  214. return true; // Not an overlap, but an error nonetheless
  215. if (!it->second.is_partition || it->second.index != i || it->second.partition_data_index != j)
  216. return true; // Overlap
  217. }
  218. }
  219. for (size_t i = 0; i < m_raw_data_entries.size(); ++i)
  220. {
  221. if (Common::swap64(m_raw_data_entries[i].data_size) == 0)
  222. continue;
  223. const u64 data_offset = Common::swap64(m_raw_data_entries[i].data_offset);
  224. const auto it = m_data_entries.upper_bound(data_offset);
  225. if (it == m_data_entries.end())
  226. return true; // Not an overlap, but an error nonetheless
  227. if (it->second.is_partition || it->second.index != i)
  228. return true; // Overlap
  229. }
  230. return false;
  231. }
  232. template <bool RVZ>
  233. std::unique_ptr<WIARVZFileReader<RVZ>> WIARVZFileReader<RVZ>::Create(File::IOFile file,
  234. const std::string& path)
  235. {
  236. std::unique_ptr<WIARVZFileReader> blob(new WIARVZFileReader(std::move(file), path));
  237. return blob->m_valid ? std::move(blob) : nullptr;
  238. }
  239. template <bool RVZ>
  240. BlobType WIARVZFileReader<RVZ>::GetBlobType() const
  241. {
  242. return RVZ ? BlobType::RVZ : BlobType::WIA;
  243. }
  244. template <bool RVZ>
  245. std::unique_ptr<BlobReader> WIARVZFileReader<RVZ>::CopyReader() const
  246. {
  247. return Create(m_file.Duplicate("rb"), m_path);
  248. }
  249. template <bool RVZ>
  250. std::string WIARVZFileReader<RVZ>::GetCompressionMethod() const
  251. {
  252. switch (m_compression_type)
  253. {
  254. case WIARVZCompressionType::Purge:
  255. return "Purge";
  256. case WIARVZCompressionType::Bzip2:
  257. return "bzip2";
  258. case WIARVZCompressionType::LZMA:
  259. return "LZMA";
  260. case WIARVZCompressionType::LZMA2:
  261. return "LZMA2";
  262. case WIARVZCompressionType::Zstd:
  263. return "Zstandard";
  264. default:
  265. return {};
  266. }
  267. }
  268. template <bool RVZ>
  269. bool WIARVZFileReader<RVZ>::Read(u64 offset, u64 size, u8* out_ptr)
  270. {
  271. if (offset + size > Common::swap64(m_header_1.iso_file_size))
  272. return false;
  273. if (offset < sizeof(WIAHeader2::disc_header))
  274. {
  275. const u64 bytes_to_read = std::min(sizeof(WIAHeader2::disc_header) - offset, size);
  276. std::memcpy(out_ptr, m_header_2.disc_header.data() + offset, bytes_to_read);
  277. offset += bytes_to_read;
  278. size -= bytes_to_read;
  279. out_ptr += bytes_to_read;
  280. }
  281. const u32 chunk_size = Common::swap32(m_header_2.chunk_size);
  282. while (size > 0)
  283. {
  284. const auto it = m_data_entries.upper_bound(offset);
  285. if (it == m_data_entries.end())
  286. return false;
  287. const DataEntry& data = it->second;
  288. if (data.is_partition)
  289. {
  290. const PartitionEntry& partition = m_partition_entries[it->second.index];
  291. const u32 partition_first_sector = Common::swap32(partition.data_entries[0].first_sector);
  292. const u64 partition_data_offset = partition_first_sector * VolumeWii::BLOCK_TOTAL_SIZE;
  293. const u32 second_number_of_sectors =
  294. Common::swap32(partition.data_entries[1].number_of_sectors);
  295. const u32 partition_total_sectors =
  296. second_number_of_sectors ? Common::swap32(partition.data_entries[1].first_sector) -
  297. partition_first_sector + second_number_of_sectors :
  298. Common::swap32(partition.data_entries[0].number_of_sectors);
  299. for (const PartitionDataEntry& partition_data : partition.data_entries)
  300. {
  301. if (size == 0)
  302. return true;
  303. const u32 first_sector = Common::swap32(partition_data.first_sector);
  304. const u32 number_of_sectors = Common::swap32(partition_data.number_of_sectors);
  305. const u64 data_offset = first_sector * VolumeWii::BLOCK_TOTAL_SIZE;
  306. const u64 data_size = number_of_sectors * VolumeWii::BLOCK_TOTAL_SIZE;
  307. if (data_size == 0)
  308. continue;
  309. if (data_offset + data_size <= offset)
  310. continue;
  311. if (offset < data_offset)
  312. return false;
  313. const u64 bytes_to_read = std::min(data_size - (offset - data_offset), size);
  314. m_exception_list.clear();
  315. m_write_to_exception_list = true;
  316. m_exception_list_last_group_index = std::numeric_limits<u64>::max();
  317. Common::ScopeGuard guard([this] { m_write_to_exception_list = false; });
  318. bool hash_exception_error = false;
  319. if (!m_encryption_cache.EncryptGroups(
  320. offset - partition_data_offset, bytes_to_read, out_ptr, partition_data_offset,
  321. partition_total_sectors * VolumeWii::BLOCK_DATA_SIZE, partition.partition_key,
  322. [this, &hash_exception_error](
  323. VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP], u64 offset_) {
  324. // EncryptGroups calls ReadWiiDecrypted, which calls ReadFromGroups,
  325. // which populates m_exception_list when m_write_to_exception_list == true
  326. if (!ApplyHashExceptions(m_exception_list, hash_blocks))
  327. hash_exception_error = true;
  328. }))
  329. {
  330. return false;
  331. }
  332. if (hash_exception_error)
  333. return false;
  334. offset += bytes_to_read;
  335. size -= bytes_to_read;
  336. out_ptr += bytes_to_read;
  337. }
  338. }
  339. else
  340. {
  341. const RawDataEntry& raw_data = m_raw_data_entries[data.index];
  342. if (!ReadFromGroups(&offset, &size, &out_ptr, chunk_size, VolumeWii::BLOCK_TOTAL_SIZE,
  343. Common::swap64(raw_data.data_offset), Common::swap64(raw_data.data_size),
  344. Common::swap32(raw_data.group_index),
  345. Common::swap32(raw_data.number_of_groups), 0))
  346. {
  347. return false;
  348. }
  349. }
  350. }
  351. return true;
  352. }
  353. template <bool RVZ>
  354. const typename WIARVZFileReader<RVZ>::PartitionEntry*
  355. WIARVZFileReader<RVZ>::GetPartition(u64 partition_data_offset, u32* partition_first_sector) const
  356. {
  357. const auto it = m_data_entries.upper_bound(partition_data_offset);
  358. if (it == m_data_entries.end() || !it->second.is_partition)
  359. return nullptr;
  360. const PartitionEntry* partition = &m_partition_entries[it->second.index];
  361. *partition_first_sector = Common::swap32(partition->data_entries[0].first_sector);
  362. if (partition_data_offset != *partition_first_sector * VolumeWii::BLOCK_TOTAL_SIZE)
  363. return nullptr;
  364. return partition;
  365. }
  366. template <bool RVZ>
  367. bool WIARVZFileReader<RVZ>::SupportsReadWiiDecrypted(u64 offset, u64 size,
  368. u64 partition_data_offset) const
  369. {
  370. u32 partition_first_sector;
  371. const PartitionEntry* partition = GetPartition(partition_data_offset, &partition_first_sector);
  372. if (!partition)
  373. return false;
  374. for (const PartitionDataEntry& data : partition->data_entries)
  375. {
  376. const u32 start_sector = Common::swap32(data.first_sector) - partition_first_sector;
  377. const u32 end_sector = start_sector + Common::swap32(data.number_of_sectors);
  378. if (offset + size <= end_sector * VolumeWii::BLOCK_DATA_SIZE)
  379. return true;
  380. }
  381. return false;
  382. }
  383. template <bool RVZ>
  384. bool WIARVZFileReader<RVZ>::ReadWiiDecrypted(u64 offset, u64 size, u8* out_ptr,
  385. u64 partition_data_offset)
  386. {
  387. u32 partition_first_sector;
  388. const PartitionEntry* partition = GetPartition(partition_data_offset, &partition_first_sector);
  389. if (!partition)
  390. return false;
  391. const u64 chunk_size = Common::swap32(m_header_2.chunk_size) * VolumeWii::BLOCK_DATA_SIZE /
  392. VolumeWii::BLOCK_TOTAL_SIZE;
  393. for (const PartitionDataEntry& data : partition->data_entries)
  394. {
  395. if (size == 0)
  396. return true;
  397. const u64 data_offset =
  398. (Common::swap32(data.first_sector) - partition_first_sector) * VolumeWii::BLOCK_DATA_SIZE;
  399. const u64 data_size = Common::swap32(data.number_of_sectors) * VolumeWii::BLOCK_DATA_SIZE;
  400. if (!ReadFromGroups(
  401. &offset, &size, &out_ptr, chunk_size, VolumeWii::BLOCK_DATA_SIZE, data_offset,
  402. data_size, Common::swap32(data.group_index), Common::swap32(data.number_of_groups),
  403. std::max<u32>(1, static_cast<u32>(chunk_size / VolumeWii::GROUP_DATA_SIZE))))
  404. {
  405. return false;
  406. }
  407. }
  408. return size == 0;
  409. }
  410. template <bool RVZ>
  411. bool WIARVZFileReader<RVZ>::ReadFromGroups(u64* offset, u64* size, u8** out_ptr, u64 chunk_size,
  412. u32 sector_size, u64 data_offset, u64 data_size,
  413. u32 group_index, u32 number_of_groups,
  414. u32 exception_lists)
  415. {
  416. if (data_offset + data_size <= *offset)
  417. return true;
  418. if (*offset < data_offset)
  419. return false;
  420. const u64 skipped_data = data_offset % sector_size;
  421. data_offset -= skipped_data;
  422. data_size += skipped_data;
  423. const u64 start_group_index = (*offset - data_offset) / chunk_size;
  424. for (u64 i = start_group_index; i < number_of_groups && (*size) > 0; ++i)
  425. {
  426. const u64 total_group_index = group_index + i;
  427. if (total_group_index >= m_group_entries.size())
  428. return false;
  429. const GroupEntry group = m_group_entries[total_group_index];
  430. const u64 group_offset_in_data = i * chunk_size;
  431. const u64 offset_in_group = *offset - group_offset_in_data - data_offset;
  432. chunk_size = std::min(chunk_size, data_size - group_offset_in_data);
  433. const u64 bytes_to_read = std::min(chunk_size - offset_in_group, *size);
  434. u32 group_data_size = Common::swap32(group.data_size);
  435. WIARVZCompressionType compression_type = m_compression_type;
  436. u32 rvz_packed_size = 0;
  437. if constexpr (RVZ)
  438. {
  439. if ((group_data_size & 0x80000000) == 0)
  440. compression_type = WIARVZCompressionType::None;
  441. group_data_size &= 0x7FFFFFFF;
  442. rvz_packed_size = Common::swap32(group.rvz_packed_size);
  443. }
  444. if (group_data_size == 0)
  445. {
  446. std::memset(*out_ptr, 0, bytes_to_read);
  447. }
  448. else
  449. {
  450. const u64 group_offset_in_file = static_cast<u64>(Common::swap32(group.data_offset)) << 2;
  451. Chunk& chunk =
  452. ReadCompressedData(group_offset_in_file, group_data_size, chunk_size, compression_type,
  453. exception_lists, rvz_packed_size, group_offset_in_data);
  454. if (!chunk.Read(offset_in_group, bytes_to_read, *out_ptr))
  455. {
  456. m_cached_chunk_offset = std::numeric_limits<u64>::max(); // Invalidate the cache
  457. return false;
  458. }
  459. if (m_write_to_exception_list && m_exception_list_last_group_index != total_group_index)
  460. {
  461. const u64 exception_list_index = offset_in_group / VolumeWii::GROUP_DATA_SIZE;
  462. const u16 additional_offset =
  463. static_cast<u16>(group_offset_in_data % VolumeWii::GROUP_DATA_SIZE /
  464. VolumeWii::BLOCK_DATA_SIZE * VolumeWii::BLOCK_HEADER_SIZE);
  465. chunk.GetHashExceptions(&m_exception_list, exception_list_index, additional_offset);
  466. m_exception_list_last_group_index = total_group_index;
  467. }
  468. }
  469. *offset += bytes_to_read;
  470. *size -= bytes_to_read;
  471. *out_ptr += bytes_to_read;
  472. }
  473. return true;
  474. }
  475. template <bool RVZ>
  476. typename WIARVZFileReader<RVZ>::Chunk&
  477. WIARVZFileReader<RVZ>::ReadCompressedData(u64 offset_in_file, u64 compressed_size,
  478. u64 decompressed_size,
  479. WIARVZCompressionType compression_type,
  480. u32 exception_lists, u32 rvz_packed_size, u64 data_offset)
  481. {
  482. if (offset_in_file == m_cached_chunk_offset)
  483. return m_cached_chunk;
  484. std::unique_ptr<Decompressor> decompressor;
  485. switch (compression_type)
  486. {
  487. case WIARVZCompressionType::None:
  488. decompressor = std::make_unique<NoneDecompressor>();
  489. break;
  490. case WIARVZCompressionType::Purge:
  491. decompressor = std::make_unique<PurgeDecompressor>(rvz_packed_size == 0 ? decompressed_size :
  492. rvz_packed_size);
  493. break;
  494. case WIARVZCompressionType::Bzip2:
  495. decompressor = std::make_unique<Bzip2Decompressor>();
  496. break;
  497. case WIARVZCompressionType::LZMA:
  498. decompressor = std::make_unique<LZMADecompressor>(false, m_header_2.compressor_data,
  499. m_header_2.compressor_data_size);
  500. break;
  501. case WIARVZCompressionType::LZMA2:
  502. decompressor = std::make_unique<LZMADecompressor>(true, m_header_2.compressor_data,
  503. m_header_2.compressor_data_size);
  504. break;
  505. case WIARVZCompressionType::Zstd:
  506. decompressor = std::make_unique<ZstdDecompressor>();
  507. break;
  508. }
  509. const bool compressed_exception_lists = compression_type > WIARVZCompressionType::Purge;
  510. m_cached_chunk =
  511. Chunk(&m_file, offset_in_file, compressed_size, decompressed_size, exception_lists,
  512. compressed_exception_lists, rvz_packed_size, data_offset, std::move(decompressor));
  513. m_cached_chunk_offset = offset_in_file;
  514. return m_cached_chunk;
  515. }
  516. template <bool RVZ>
  517. std::string WIARVZFileReader<RVZ>::VersionToString(u32 version)
  518. {
  519. const u8 a = version >> 24;
  520. const u8 b = (version >> 16) & 0xff;
  521. const u8 c = (version >> 8) & 0xff;
  522. const u8 d = version & 0xff;
  523. if (d == 0 || d == 0xff)
  524. return fmt::format("{}.{:02x}.{:02x}", a, b, c);
  525. else
  526. return fmt::format("{}.{:02x}.{:02x}.beta{}", a, b, c, d);
  527. }
  528. template <bool RVZ>
  529. WIARVZFileReader<RVZ>::Chunk::Chunk() = default;
  530. template <bool RVZ>
  531. WIARVZFileReader<RVZ>::Chunk::Chunk(File::IOFile* file, u64 offset_in_file, u64 compressed_size,
  532. u64 decompressed_size, u32 exception_lists,
  533. bool compressed_exception_lists, u32 rvz_packed_size,
  534. u64 data_offset, std::unique_ptr<Decompressor> decompressor)
  535. : m_decompressor(std::move(decompressor)), m_file(file), m_offset_in_file(offset_in_file),
  536. m_exception_lists(exception_lists), m_compressed_exception_lists(compressed_exception_lists),
  537. m_rvz_packed_size(rvz_packed_size), m_data_offset(data_offset)
  538. {
  539. constexpr size_t MAX_SIZE_PER_EXCEPTION_LIST =
  540. Common::AlignUp(VolumeWii::BLOCK_HEADER_SIZE, Common::SHA1::DIGEST_LEN) /
  541. Common::SHA1::DIGEST_LEN * VolumeWii::BLOCKS_PER_GROUP * sizeof(HashExceptionEntry) +
  542. sizeof(u16);
  543. m_out_bytes_allocated_for_exceptions =
  544. m_compressed_exception_lists ? MAX_SIZE_PER_EXCEPTION_LIST * m_exception_lists : 0;
  545. m_in.data.resize(compressed_size);
  546. m_out.data.resize(decompressed_size + m_out_bytes_allocated_for_exceptions);
  547. }
  548. template <bool RVZ>
  549. bool WIARVZFileReader<RVZ>::Chunk::Read(u64 offset, u64 size, u8* out_ptr)
  550. {
  551. if (!m_decompressor || !m_file ||
  552. offset + size > m_out.data.size() - m_out_bytes_allocated_for_exceptions)
  553. {
  554. return false;
  555. }
  556. while (offset + size > GetOutBytesWrittenExcludingExceptions())
  557. {
  558. u64 bytes_to_read;
  559. if (offset + size == m_out.data.size())
  560. {
  561. // Read all the remaining data.
  562. bytes_to_read = m_in.data.size() - m_in.bytes_written;
  563. }
  564. else
  565. {
  566. // Pick a suitable amount of compressed data to read. We have to ensure that bytes_to_read
  567. // is larger than 0 and smaller than or equal to the number of bytes available to read,
  568. // but the rest is a bit arbitrary and could be changed.
  569. // The compressed data is probably not much bigger than the decompressed data.
  570. // Add a few bytes for possible compression overhead and for any hash exceptions.
  571. bytes_to_read = offset + size - GetOutBytesWrittenExcludingExceptions() + 0x100;
  572. // Align the access in an attempt to gain speed. But we don't actually know the
  573. // block size of the underlying storage device, so we just use the Wii block size.
  574. bytes_to_read =
  575. Common::AlignUp(bytes_to_read + m_offset_in_file, VolumeWii::BLOCK_TOTAL_SIZE) -
  576. m_offset_in_file;
  577. // Ensure we don't read too much.
  578. bytes_to_read = std::min<u64>(m_in.data.size() - m_in.bytes_written, bytes_to_read);
  579. }
  580. if (bytes_to_read == 0)
  581. {
  582. // Compressed size is larger than expected or decompressed size is smaller than expected
  583. return false;
  584. }
  585. if (!m_file->Seek(m_offset_in_file, File::SeekOrigin::Begin))
  586. return false;
  587. if (!m_file->ReadBytes(m_in.data.data() + m_in.bytes_written, bytes_to_read))
  588. return false;
  589. m_offset_in_file += bytes_to_read;
  590. m_in.bytes_written += bytes_to_read;
  591. if (m_exception_lists > 0 && !m_compressed_exception_lists)
  592. {
  593. if (!HandleExceptions(m_in.data.data(), m_in.data.size(), m_in.bytes_written,
  594. &m_in_bytes_used_for_exceptions, true))
  595. {
  596. return false;
  597. }
  598. m_in_bytes_read = m_in_bytes_used_for_exceptions;
  599. }
  600. if (m_exception_lists == 0 || m_compressed_exception_lists)
  601. {
  602. if (!Decompress())
  603. return false;
  604. }
  605. if (m_exception_lists > 0 && m_compressed_exception_lists)
  606. {
  607. if (!HandleExceptions(m_out.data.data(), m_out_bytes_allocated_for_exceptions,
  608. m_out.bytes_written, &m_out_bytes_used_for_exceptions, false))
  609. {
  610. return false;
  611. }
  612. if (m_rvz_packed_size != 0 && m_exception_lists == 0)
  613. {
  614. if (!Decompress())
  615. return false;
  616. }
  617. }
  618. if (m_exception_lists == 0)
  619. {
  620. const size_t expected_out_bytes = m_out.data.size() - m_out_bytes_allocated_for_exceptions +
  621. m_out_bytes_used_for_exceptions;
  622. if (m_out.bytes_written > expected_out_bytes)
  623. return false; // Decompressed size is larger than expected
  624. // The reason why we need the m_in.bytes_written == m_in.data.size() check as part of
  625. // this conditional is because (for example) zstd can finish writing all data to m_out
  626. // before becoming done if we've given it all input data except the checksum at the end.
  627. if (m_out.bytes_written == expected_out_bytes && !m_decompressor->Done() &&
  628. m_in.bytes_written == m_in.data.size())
  629. {
  630. return false; // Decompressed size is larger than expected
  631. }
  632. if (m_decompressor->Done() && m_in_bytes_read != m_in.data.size())
  633. return false; // Compressed size is smaller than expected
  634. }
  635. }
  636. std::memcpy(out_ptr, m_out.data.data() + offset + m_out_bytes_used_for_exceptions, size);
  637. return true;
  638. }
  639. template <bool RVZ>
  640. bool WIARVZFileReader<RVZ>::Chunk::Decompress()
  641. {
  642. if (m_rvz_packed_size != 0 && m_exception_lists == 0)
  643. {
  644. const size_t bytes_to_move = m_out.bytes_written - m_out_bytes_used_for_exceptions;
  645. DecompressionBuffer in{std::vector<u8>(bytes_to_move), bytes_to_move};
  646. // Copying to a null pointer is undefined behaviour, so only copy when we
  647. // actually have data to copy.
  648. if (bytes_to_move > 0)
  649. {
  650. std::memcpy(in.data.data(), m_out.data.data() + m_out_bytes_used_for_exceptions,
  651. bytes_to_move);
  652. }
  653. m_out.bytes_written = m_out_bytes_used_for_exceptions;
  654. m_decompressor = std::make_unique<RVZPackDecompressor>(std::move(m_decompressor), std::move(in),
  655. m_data_offset, m_rvz_packed_size);
  656. m_rvz_packed_size = 0;
  657. }
  658. return m_decompressor->Decompress(m_in, &m_out, &m_in_bytes_read);
  659. }
  660. template <bool RVZ>
  661. bool WIARVZFileReader<RVZ>::Chunk::HandleExceptions(const u8* data, size_t bytes_allocated,
  662. size_t bytes_written, size_t* bytes_used,
  663. bool align)
  664. {
  665. while (m_exception_lists > 0)
  666. {
  667. if (sizeof(u16) + *bytes_used > bytes_allocated)
  668. {
  669. ERROR_LOG_FMT(DISCIO, "More hash exceptions than expected");
  670. return false;
  671. }
  672. if (sizeof(u16) + *bytes_used > bytes_written)
  673. return true;
  674. const u16 exceptions = Common::swap16(data + *bytes_used);
  675. size_t exception_list_size = exceptions * sizeof(HashExceptionEntry) + sizeof(u16);
  676. if (align && m_exception_lists == 1)
  677. exception_list_size = Common::AlignUp(*bytes_used + exception_list_size, 4) - *bytes_used;
  678. if (exception_list_size + *bytes_used > bytes_allocated)
  679. {
  680. ERROR_LOG_FMT(DISCIO, "More hash exceptions than expected");
  681. return false;
  682. }
  683. if (exception_list_size + *bytes_used > bytes_written)
  684. return true;
  685. *bytes_used += exception_list_size;
  686. --m_exception_lists;
  687. }
  688. return true;
  689. }
  690. template <bool RVZ>
  691. void WIARVZFileReader<RVZ>::Chunk::GetHashExceptions(
  692. std::vector<HashExceptionEntry>* exception_list, u64 exception_list_index,
  693. u16 additional_offset) const
  694. {
  695. ASSERT(m_exception_lists == 0);
  696. const u8* data_start = m_compressed_exception_lists ? m_out.data.data() : m_in.data.data();
  697. const u8* data = data_start;
  698. for (u64 i = exception_list_index; i > 0; --i)
  699. data += Common::swap16(data) * sizeof(HashExceptionEntry) + sizeof(u16);
  700. const u16 exceptions = Common::swap16(data);
  701. data += sizeof(u16);
  702. for (size_t i = 0; i < exceptions; ++i)
  703. {
  704. std::memcpy(&exception_list->emplace_back(), data, sizeof(HashExceptionEntry));
  705. data += sizeof(HashExceptionEntry);
  706. u16& offset = exception_list->back().offset;
  707. offset = Common::swap16(Common::swap16(offset) + additional_offset);
  708. }
  709. ASSERT(data <= data_start + (m_compressed_exception_lists ? m_out_bytes_used_for_exceptions :
  710. m_in_bytes_used_for_exceptions));
  711. }
  712. template <bool RVZ>
  713. size_t WIARVZFileReader<RVZ>::Chunk::GetOutBytesWrittenExcludingExceptions() const
  714. {
  715. return m_exception_lists == 0 ? m_out.bytes_written - m_out_bytes_used_for_exceptions : 0;
  716. }
  717. template <bool RVZ>
  718. bool WIARVZFileReader<RVZ>::ApplyHashExceptions(
  719. const std::vector<HashExceptionEntry>& exception_list,
  720. VolumeWii::HashBlock hash_blocks[VolumeWii::BLOCKS_PER_GROUP])
  721. {
  722. for (const HashExceptionEntry& exception : exception_list)
  723. {
  724. const u16 offset = Common::swap16(exception.offset);
  725. const size_t block_index = offset / VolumeWii::BLOCK_HEADER_SIZE;
  726. if (block_index > VolumeWii::BLOCKS_PER_GROUP)
  727. return false;
  728. const size_t offset_in_block = offset % VolumeWii::BLOCK_HEADER_SIZE;
  729. if (offset_in_block + Common::SHA1::DIGEST_LEN > VolumeWii::BLOCK_HEADER_SIZE)
  730. return false;
  731. std::memcpy(reinterpret_cast<u8*>(&hash_blocks[block_index]) + offset_in_block, &exception.hash,
  732. Common::SHA1::DIGEST_LEN);
  733. }
  734. return true;
  735. }
  736. template <bool RVZ>
  737. bool WIARVZFileReader<RVZ>::PadTo4(File::IOFile* file, u64* bytes_written)
  738. {
  739. constexpr u32 ZEROES = 0;
  740. const u64 bytes_to_write = Common::AlignUp(*bytes_written, 4) - *bytes_written;
  741. if (bytes_to_write == 0)
  742. return true;
  743. *bytes_written += bytes_to_write;
  744. return file->WriteBytes(&ZEROES, bytes_to_write);
  745. }
  746. template <bool RVZ>
  747. void WIARVZFileReader<RVZ>::AddRawDataEntry(u64 offset, u64 size, int chunk_size, u32* total_groups,
  748. std::vector<RawDataEntry>* raw_data_entries,
  749. std::vector<DataEntry>* data_entries)
  750. {
  751. constexpr size_t SKIP_SIZE = sizeof(WIAHeader2::disc_header);
  752. const u64 skip = offset < SKIP_SIZE ? std::min(SKIP_SIZE - offset, size) : 0;
  753. offset += skip;
  754. size -= skip;
  755. if (size == 0)
  756. return;
  757. const u32 group_index = *total_groups;
  758. const u32 groups = static_cast<u32>(Common::AlignUp(size, chunk_size) / chunk_size);
  759. *total_groups += groups;
  760. data_entries->emplace_back(raw_data_entries->size());
  761. raw_data_entries->emplace_back(RawDataEntry{Common::swap64(offset), Common::swap64(size),
  762. Common::swap32(group_index), Common::swap32(groups)});
  763. }
  764. template <bool RVZ>
  765. typename WIARVZFileReader<RVZ>::PartitionDataEntry WIARVZFileReader<RVZ>::CreatePartitionDataEntry(
  766. u64 offset, u64 size, u32 index, int chunk_size, u32* total_groups,
  767. const std::vector<PartitionEntry>& partition_entries, std::vector<DataEntry>* data_entries)
  768. {
  769. const u32 group_index = *total_groups;
  770. const u64 rounded_size = Common::AlignDown(size, VolumeWii::BLOCK_TOTAL_SIZE);
  771. const u32 groups = static_cast<u32>(Common::AlignUp(rounded_size, chunk_size) / chunk_size);
  772. *total_groups += groups;
  773. data_entries->emplace_back(partition_entries.size(), index);
  774. return PartitionDataEntry{Common::swap32(offset / VolumeWii::BLOCK_TOTAL_SIZE),
  775. Common::swap32(size / VolumeWii::BLOCK_TOTAL_SIZE),
  776. Common::swap32(group_index), Common::swap32(groups)};
  777. }
  778. template <bool RVZ>
  779. ConversionResultCode WIARVZFileReader<RVZ>::SetUpDataEntriesForWriting(
  780. const VolumeDisc* volume, int chunk_size, u64 iso_size, u32* total_groups,
  781. std::vector<PartitionEntry>* partition_entries, std::vector<RawDataEntry>* raw_data_entries,
  782. std::vector<DataEntry>* data_entries, std::vector<const FileSystem*>* partition_file_systems)
  783. {
  784. std::vector<Partition> partitions;
  785. if (volume && volume->HasWiiHashes() && volume->HasWiiEncryption())
  786. partitions = volume->GetPartitions();
  787. std::sort(partitions.begin(), partitions.end(),
  788. [](const Partition& a, const Partition& b) { return a.offset < b.offset; });
  789. *total_groups = 0;
  790. u64 last_partition_end_offset = 0;
  791. const auto add_raw_data_entry = [&](u64 offset, u64 size) {
  792. return AddRawDataEntry(offset, size, chunk_size, total_groups, raw_data_entries, data_entries);
  793. };
  794. const auto create_partition_data_entry = [&](u64 offset, u64 size, u32 index) {
  795. return CreatePartitionDataEntry(offset, size, index, chunk_size, total_groups,
  796. *partition_entries, data_entries);
  797. };
  798. for (const Partition& partition : partitions)
  799. {
  800. // If a partition is odd in some way that prevents us from encoding it as a partition,
  801. // we encode it as raw data instead by skipping the current loop iteration.
  802. // Partitions can always be encoded as raw data, but it is less space efficient.
  803. if (partition.offset < last_partition_end_offset)
  804. {
  805. WARN_LOG_FMT(DISCIO, "Overlapping partitions at {:x}", partition.offset);
  806. continue;
  807. }
  808. if (volume->ReadSwapped<u32>(partition.offset, PARTITION_NONE) != 0x10001U)
  809. {
  810. // This looks more like garbage data than an actual partition.
  811. // The values of data_offset and data_size will very likely also be garbage.
  812. // Some WBFS writing programs scrub the SSBB Masterpiece partitions without
  813. // removing them from the partition table, causing this problem.
  814. WARN_LOG_FMT(DISCIO, "Invalid partition at {:x}", partition.offset);
  815. continue;
  816. }
  817. std::optional<u64> data_offset =
  818. volume->ReadSwappedAndShifted(partition.offset + 0x2b8, PARTITION_NONE);
  819. std::optional<u64> data_size =
  820. volume->ReadSwappedAndShifted(partition.offset + 0x2bc, PARTITION_NONE);
  821. if (!data_offset || !data_size)
  822. return ConversionResultCode::ReadFailed;
  823. const u64 data_start = partition.offset + *data_offset;
  824. const u64 data_end = data_start + *data_size;
  825. if (data_start % VolumeWii::BLOCK_TOTAL_SIZE != 0)
  826. {
  827. WARN_LOG_FMT(DISCIO, "Misaligned partition at {:x}", partition.offset);
  828. continue;
  829. }
  830. if (*data_size < VolumeWii::BLOCK_TOTAL_SIZE)
  831. {
  832. WARN_LOG_FMT(DISCIO, "Very small partition at {:x}", partition.offset);
  833. continue;
  834. }
  835. if (data_end > iso_size)
  836. {
  837. WARN_LOG_FMT(DISCIO, "Too large partition at {:x}", partition.offset);
  838. *data_size = iso_size - *data_offset - partition.offset;
  839. }
  840. const std::optional<u64> fst_offset = GetFSTOffset(*volume, partition);
  841. const std::optional<u64> fst_size = GetFSTSize(*volume, partition);
  842. if (!fst_offset || !fst_size)
  843. return ConversionResultCode::ReadFailed;
  844. const IOS::ES::TicketReader& ticket = volume->GetTicket(partition);
  845. if (!ticket.IsValid())
  846. return ConversionResultCode::ReadFailed;
  847. add_raw_data_entry(last_partition_end_offset, partition.offset - last_partition_end_offset);
  848. add_raw_data_entry(partition.offset, *data_offset);
  849. const u64 fst_end = volume->PartitionOffsetToRawOffset(*fst_offset + *fst_size, partition);
  850. const u64 split_point = std::min(
  851. data_end, Common::AlignUp(fst_end - data_start, VolumeWii::GROUP_TOTAL_SIZE) + data_start);
  852. PartitionEntry partition_entry;
  853. partition_entry.partition_key = ticket.GetTitleKey();
  854. partition_entry.data_entries[0] =
  855. create_partition_data_entry(data_start, split_point - data_start, 0);
  856. partition_entry.data_entries[1] =
  857. create_partition_data_entry(split_point, data_end - split_point, 1);
  858. // Note: We can't simply set last_partition_end_offset to data_end,
  859. // because construct_partition_data_entry may have rounded it
  860. last_partition_end_offset =
  861. (Common::swap32(partition_entry.data_entries[1].first_sector) +
  862. Common::swap32(partition_entry.data_entries[1].number_of_sectors)) *
  863. VolumeWii::BLOCK_TOTAL_SIZE;
  864. partition_entries->emplace_back(std::move(partition_entry));
  865. partition_file_systems->emplace_back(volume->GetFileSystem(partition));
  866. }
  867. add_raw_data_entry(last_partition_end_offset, iso_size - last_partition_end_offset);
  868. return ConversionResultCode::Success;
  869. }
  870. template <bool RVZ>
  871. std::optional<std::vector<u8>> WIARVZFileReader<RVZ>::Compress(Compressor* compressor,
  872. const u8* data, size_t size)
  873. {
  874. if (compressor)
  875. {
  876. if (!compressor->Start(size) || !compressor->Compress(data, size) || !compressor->End())
  877. return std::nullopt;
  878. data = compressor->GetData();
  879. size = compressor->GetSize();
  880. }
  881. return std::vector<u8>(data, data + size);
  882. }
  883. template <bool RVZ>
  884. void WIARVZFileReader<RVZ>::SetUpCompressor(std::unique_ptr<Compressor>* compressor,
  885. WIARVZCompressionType compression_type,
  886. int compression_level, WIAHeader2* header_2)
  887. {
  888. switch (compression_type)
  889. {
  890. case WIARVZCompressionType::None:
  891. *compressor = nullptr;
  892. break;
  893. case WIARVZCompressionType::Purge:
  894. *compressor = std::make_unique<PurgeCompressor>();
  895. break;
  896. case WIARVZCompressionType::Bzip2:
  897. *compressor = std::make_unique<Bzip2Compressor>(compression_level);
  898. break;
  899. case WIARVZCompressionType::LZMA:
  900. case WIARVZCompressionType::LZMA2:
  901. {
  902. u8* compressor_data = nullptr;
  903. u8* compressor_data_size = nullptr;
  904. if (header_2)
  905. {
  906. compressor_data = header_2->compressor_data;
  907. compressor_data_size = &header_2->compressor_data_size;
  908. }
  909. const bool lzma2 = compression_type == WIARVZCompressionType::LZMA2;
  910. *compressor = std::make_unique<LZMACompressor>(lzma2, compression_level, compressor_data,
  911. compressor_data_size);
  912. break;
  913. }
  914. case WIARVZCompressionType::Zstd:
  915. *compressor = std::make_unique<ZstdCompressor>(compression_level);
  916. break;
  917. }
  918. }
  919. template <bool RVZ>
  920. bool WIARVZFileReader<RVZ>::TryReuse(std::map<ReuseID, GroupEntry>* reusable_groups,
  921. std::mutex* reusable_groups_mutex,
  922. OutputParametersEntry* entry)
  923. {
  924. if (entry->reused_group)
  925. return true;
  926. if (!entry->reuse_id)
  927. return false;
  928. std::lock_guard guard(*reusable_groups_mutex);
  929. const auto it = reusable_groups->find(*entry->reuse_id);
  930. if (it == reusable_groups->end())
  931. return false;
  932. entry->reused_group = it->second;
  933. return true;
  934. }
  935. static bool AllAre(const std::vector<u8>& data, u8 x)
  936. {
  937. return std::all_of(data.begin(), data.end(), [x](u8 y) { return x == y; });
  938. }
  939. static bool AllAre(const u8* begin, const u8* end, u8 x)
  940. {
  941. return std::all_of(begin, end, [x](u8 y) { return x == y; });
  942. }
  943. static bool AllZero(const std::vector<u8>& data)
  944. {
  945. return AllAre(data, 0);
  946. }
  947. static bool AllSame(const std::vector<u8>& data)
  948. {
  949. return AllAre(data, data.front());
  950. }
  951. static bool AllSame(const u8* begin, const u8* end)
  952. {
  953. return AllAre(begin, end, *begin);
  954. }
  955. template <typename OutputParametersEntry>
  956. static void RVZPack(const u8* in, OutputParametersEntry* out, u64 bytes_per_chunk, size_t chunks,
  957. u64 total_size, u64 data_offset, bool multipart, bool allow_junk_reuse,
  958. bool compression, const FileSystem* file_system)
  959. {
  960. using Seed = std::array<u32, LaggedFibonacciGenerator::SEED_SIZE>;
  961. struct JunkInfo
  962. {
  963. size_t start_offset;
  964. Seed seed;
  965. };
  966. constexpr size_t SEED_SIZE = LaggedFibonacciGenerator::SEED_SIZE * sizeof(u32);
  967. // Maps end_offset -> (start_offset, seed)
  968. std::map<size_t, JunkInfo> junk_info;
  969. size_t position = 0;
  970. while (position < total_size)
  971. {
  972. // Skip the 0 to 32 zero bytes that typically come after a file
  973. size_t zeroes = 0;
  974. while (position + zeroes < total_size && in[position + zeroes] == 0)
  975. ++zeroes;
  976. // If there are very many zero bytes (perhaps the PRNG junk data has been scrubbed?)
  977. // and we aren't using compression, it makes sense to encode the zero bytes as junk.
  978. // If we are using compression, the compressor will likely encode zeroes better than we can
  979. if (!compression && zeroes > SEED_SIZE)
  980. junk_info.emplace(position + zeroes, JunkInfo{position, {}});
  981. position += zeroes;
  982. data_offset += zeroes;
  983. const size_t bytes_to_read =
  984. std::min(Common::AlignUp(data_offset + 1, VolumeWii::BLOCK_TOTAL_SIZE) - data_offset,
  985. total_size - position);
  986. const size_t data_offset_mod = static_cast<size_t>(data_offset % VolumeWii::BLOCK_TOTAL_SIZE);
  987. Seed seed;
  988. const size_t bytes_reconstructed = LaggedFibonacciGenerator::GetSeed(
  989. in + position, bytes_to_read, data_offset_mod, seed.data());
  990. if (bytes_reconstructed > 0)
  991. junk_info.emplace(position + bytes_reconstructed, JunkInfo{position, seed});
  992. if (file_system)
  993. {
  994. const std::unique_ptr<DiscIO::FileInfo> file_info =
  995. file_system->FindFileInfo(data_offset + bytes_reconstructed);
  996. // If we're at a file and there's more space in this block after the file,
  997. // continue after the file instead of skipping to the next block
  998. if (file_info)
  999. {
  1000. const u64 file_end_offset = file_info->GetOffset() + file_info->GetSize();
  1001. if (file_end_offset < data_offset + bytes_to_read)
  1002. {
  1003. position += file_end_offset - data_offset;
  1004. data_offset = file_end_offset;
  1005. continue;
  1006. }
  1007. }
  1008. }
  1009. position += bytes_to_read;
  1010. data_offset += bytes_to_read;
  1011. }
  1012. for (size_t i = 0; i < chunks; ++i)
  1013. {
  1014. OutputParametersEntry& entry = out[i];
  1015. if (entry.reused_group)
  1016. continue;
  1017. u64 current_offset = i * bytes_per_chunk;
  1018. const u64 end_offset = std::min(current_offset + bytes_per_chunk, total_size);
  1019. const bool store_junk_efficiently = allow_junk_reuse || !entry.reuse_id;
  1020. // TODO: It would be possible to support skipping RVZ packing even when the chunk size is larger
  1021. // than 2 MiB (multipart == true), but it would be more effort than it's worth since Dolphin's
  1022. // converter doesn't expose chunk sizes larger than 2 MiB to the user anyway
  1023. bool first_loop_iteration = !multipart;
  1024. while (current_offset < end_offset)
  1025. {
  1026. u64 next_junk_start = end_offset;
  1027. u64 next_junk_end = end_offset;
  1028. Seed* seed = nullptr;
  1029. if (store_junk_efficiently && end_offset - current_offset > SEED_SIZE)
  1030. {
  1031. const auto next_junk_it = junk_info.upper_bound(current_offset + SEED_SIZE);
  1032. if (next_junk_it != junk_info.end() &&
  1033. next_junk_it->second.start_offset + SEED_SIZE < end_offset)
  1034. {
  1035. next_junk_start = std::max<u64>(current_offset, next_junk_it->second.start_offset);
  1036. next_junk_end = std::min<u64>(end_offset, next_junk_it->first);
  1037. seed = &next_junk_it->second.seed;
  1038. }
  1039. }
  1040. if (first_loop_iteration)
  1041. {
  1042. if (next_junk_start == end_offset)
  1043. {
  1044. // Storing this chunk with RVZ packing would be inefficient, so store it without
  1045. PushBack(&entry.main_data, in + current_offset, in + end_offset);
  1046. break;
  1047. }
  1048. first_loop_iteration = false;
  1049. }
  1050. const u64 non_junk_bytes = next_junk_start - current_offset;
  1051. if (non_junk_bytes > 0)
  1052. {
  1053. const u8* ptr = in + current_offset;
  1054. PushBack(&entry.main_data, Common::swap32(static_cast<u32>(non_junk_bytes)));
  1055. PushBack(&entry.main_data, ptr, ptr + non_junk_bytes);
  1056. current_offset += non_junk_bytes;
  1057. entry.rvz_packed_size += sizeof(u32) + non_junk_bytes;
  1058. }
  1059. const u64 junk_bytes = next_junk_end - current_offset;
  1060. if (junk_bytes > 0)
  1061. {
  1062. PushBack(&entry.main_data, Common::swap32(static_cast<u32>(junk_bytes) | 0x80000000));
  1063. PushBack(&entry.main_data, *seed);
  1064. current_offset += junk_bytes;
  1065. entry.rvz_packed_size += sizeof(u32) + SEED_SIZE;
  1066. }
  1067. }
  1068. }
  1069. }
  1070. template <typename OutputParametersEntry>
  1071. static void RVZPack(const u8* in, OutputParametersEntry* out, u64 size, u64 data_offset,
  1072. bool allow_junk_reuse, bool compression, const FileSystem* file_system)
  1073. {
  1074. RVZPack(in, out, size, 1, size, data_offset, false, allow_junk_reuse, compression, file_system);
  1075. }
  1076. template <bool RVZ>
  1077. ConversionResult<typename WIARVZFileReader<RVZ>::OutputParameters>
  1078. WIARVZFileReader<RVZ>::ProcessAndCompress(CompressThreadState* state, CompressParameters parameters,
  1079. const std::vector<PartitionEntry>& partition_entries,
  1080. const std::vector<DataEntry>& data_entries,
  1081. const FileSystem* file_system,
  1082. std::map<ReuseID, GroupEntry>* reusable_groups,
  1083. std::mutex* reusable_groups_mutex,
  1084. u64 chunks_per_wii_group, u64 exception_lists_per_chunk,
  1085. bool compressed_exception_lists, bool compression)
  1086. {
  1087. std::vector<OutputParametersEntry> output_entries;
  1088. if (!parameters.data_entry->is_partition)
  1089. {
  1090. OutputParametersEntry& entry = output_entries.emplace_back();
  1091. std::vector<u8>& data = parameters.data;
  1092. if (AllSame(data))
  1093. entry.reuse_id = ReuseID{WiiKey{}, data.size(), false, data.front()};
  1094. if constexpr (RVZ)
  1095. {
  1096. RVZPack(data.data(), output_entries.data(), data.size(), parameters.data_offset, true,
  1097. compression, file_system);
  1098. }
  1099. else
  1100. {
  1101. entry.main_data = std::move(data);
  1102. }
  1103. }
  1104. else
  1105. {
  1106. const PartitionEntry& partition_entry = partition_entries[parameters.data_entry->index];
  1107. auto aes_context = Common::AES::CreateContextDecrypt(partition_entry.partition_key.data());
  1108. const u64 groups = Common::AlignUp(parameters.data.size(), VolumeWii::GROUP_TOTAL_SIZE) /
  1109. VolumeWii::GROUP_TOTAL_SIZE;
  1110. ASSERT(parameters.data.size() % VolumeWii::BLOCK_TOTAL_SIZE == 0);
  1111. const u64 blocks = parameters.data.size() / VolumeWii::BLOCK_TOTAL_SIZE;
  1112. const u64 blocks_per_chunk = chunks_per_wii_group == 1 ?
  1113. exception_lists_per_chunk * VolumeWii::BLOCKS_PER_GROUP :
  1114. VolumeWii::BLOCKS_PER_GROUP / chunks_per_wii_group;
  1115. const u64 chunks = Common::AlignUp(blocks, blocks_per_chunk) / blocks_per_chunk;
  1116. const u64 in_data_per_chunk = blocks_per_chunk * VolumeWii::BLOCK_TOTAL_SIZE;
  1117. const u64 out_data_per_chunk = blocks_per_chunk * VolumeWii::BLOCK_DATA_SIZE;
  1118. const size_t first_chunk = output_entries.size();
  1119. const auto create_reuse_id = [&partition_entry, blocks,
  1120. blocks_per_chunk](u8 value, bool encrypted, u64 block) {
  1121. const u64 size = std::min(blocks - block, blocks_per_chunk) * VolumeWii::BLOCK_DATA_SIZE;
  1122. return ReuseID{partition_entry.partition_key, size, encrypted, value};
  1123. };
  1124. const u8* parameters_data_end = parameters.data.data() + parameters.data.size();
  1125. for (u64 i = 0; i < chunks; ++i)
  1126. {
  1127. const u64 block_index = i * blocks_per_chunk;
  1128. OutputParametersEntry& entry = output_entries.emplace_back();
  1129. std::optional<ReuseID>& reuse_id = entry.reuse_id;
  1130. // Set this chunk as reusable if the encrypted data is AllSame
  1131. const u8* data = parameters.data.data() + block_index * VolumeWii::BLOCK_TOTAL_SIZE;
  1132. if (AllSame(data, std::min(parameters_data_end, data + in_data_per_chunk)))
  1133. reuse_id = create_reuse_id(parameters.data.front(), true, i * blocks_per_chunk);
  1134. TryReuse(reusable_groups, reusable_groups_mutex, &entry);
  1135. if (!entry.reused_group && reuse_id)
  1136. {
  1137. const auto it = std::find_if(output_entries.begin(), output_entries.begin() + i,
  1138. [reuse_id](const auto& e) { return e.reuse_id == reuse_id; });
  1139. if (it != output_entries.begin() + i)
  1140. entry.reused_group = it->reused_group;
  1141. }
  1142. }
  1143. if (!std::all_of(output_entries.begin(), output_entries.end(),
  1144. [](const OutputParametersEntry& entry) { return entry.reused_group; }))
  1145. {
  1146. const u64 number_of_exception_lists =
  1147. chunks_per_wii_group == 1 ? exception_lists_per_chunk : chunks;
  1148. std::vector<std::vector<HashExceptionEntry>> exception_lists(number_of_exception_lists);
  1149. for (u64 i = 0; i < groups; ++i)
  1150. {
  1151. const u64 offset_of_group = i * VolumeWii::GROUP_TOTAL_SIZE;
  1152. const u64 write_offset_of_group = i * VolumeWii::GROUP_DATA_SIZE;
  1153. const u64 blocks_in_this_group =
  1154. std::min<u64>(VolumeWii::BLOCKS_PER_GROUP, blocks - i * VolumeWii::BLOCKS_PER_GROUP);
  1155. for (u32 j = 0; j < VolumeWii::BLOCKS_PER_GROUP; ++j)
  1156. {
  1157. if (j < blocks_in_this_group)
  1158. {
  1159. const u64 offset_of_block = offset_of_group + j * VolumeWii::BLOCK_TOTAL_SIZE;
  1160. VolumeWii::DecryptBlockData(parameters.data.data() + offset_of_block,
  1161. state->decryption_buffer[j].data(), aes_context.get());
  1162. }
  1163. else
  1164. {
  1165. state->decryption_buffer[j].fill(0);
  1166. }
  1167. }
  1168. VolumeWii::HashGroup(state->decryption_buffer.data(), state->hash_buffer.data());
  1169. for (u64 j = 0; j < blocks_in_this_group; ++j)
  1170. {
  1171. const u64 chunk_index = j / blocks_per_chunk;
  1172. const u64 block_index_in_chunk = j % blocks_per_chunk;
  1173. if (output_entries[chunk_index].reused_group)
  1174. continue;
  1175. const u64 exception_list_index = chunks_per_wii_group == 1 ? i : chunk_index;
  1176. const u64 offset_of_block = offset_of_group + j * VolumeWii::BLOCK_TOTAL_SIZE;
  1177. const u64 hash_offset_of_block = block_index_in_chunk * VolumeWii::BLOCK_HEADER_SIZE;
  1178. VolumeWii::HashBlock hashes;
  1179. VolumeWii::DecryptBlockHashes(parameters.data.data() + offset_of_block, &hashes,
  1180. aes_context.get());
  1181. const auto compare_hash = [&](size_t offset_in_block) {
  1182. ASSERT(offset_in_block + Common::SHA1::DIGEST_LEN <= VolumeWii::BLOCK_HEADER_SIZE);
  1183. const u8* desired_hash = reinterpret_cast<u8*>(&hashes) + offset_in_block;
  1184. const u8* computed_hash =
  1185. reinterpret_cast<u8*>(&state->hash_buffer[j]) + offset_in_block;
  1186. // We want to store a hash exception either if there is a hash mismatch, or if this
  1187. // chunk might get reused in a context where it is paired up (within a 2 MiB Wii group)
  1188. // with chunks that are different from the chunks it currently is paired up with, since
  1189. // that affects the recalculated hashes. Chunks which have been marked as reusable at
  1190. // this point normally have zero matching hashes anyway, so this shouldn't waste space.
  1191. if ((chunks_per_wii_group != 1 && output_entries[chunk_index].reuse_id) ||
  1192. !std::equal(desired_hash, desired_hash + Common::SHA1::DIGEST_LEN, computed_hash))
  1193. {
  1194. const u64 hash_offset = hash_offset_of_block + offset_in_block;
  1195. ASSERT(hash_offset <= std::numeric_limits<u16>::max());
  1196. HashExceptionEntry& exception = exception_lists[exception_list_index].emplace_back();
  1197. exception.offset = static_cast<u16>(Common::swap16(hash_offset));
  1198. std::memcpy(exception.hash.data(), desired_hash, Common::SHA1::DIGEST_LEN);
  1199. }
  1200. };
  1201. const auto compare_hashes = [&compare_hash](size_t offset, size_t size) {
  1202. for (size_t l = 0; l < size; l += Common::SHA1::DIGEST_LEN)
  1203. // The std::min is to ensure that we don't go beyond the end of HashBlock with
  1204. // padding_2, which is 32 bytes long (not divisible by SHA1::DIGEST_LEN, which is 20).
  1205. compare_hash(offset + std::min(l, size - Common::SHA1::DIGEST_LEN));
  1206. };
  1207. using HashBlock = VolumeWii::HashBlock;
  1208. compare_hashes(offsetof(HashBlock, h0), sizeof(HashBlock::h0));
  1209. compare_hashes(offsetof(HashBlock, padding_0), sizeof(HashBlock::padding_0));
  1210. compare_hashes(offsetof(HashBlock, h1), sizeof(HashBlock::h1));
  1211. compare_hashes(offsetof(HashBlock, padding_1), sizeof(HashBlock::padding_1));
  1212. compare_hashes(offsetof(HashBlock, h2), sizeof(HashBlock::h2));
  1213. compare_hashes(offsetof(HashBlock, padding_2), sizeof(HashBlock::padding_2));
  1214. }
  1215. static_assert(std::is_trivially_copyable_v<
  1216. typename decltype(CompressThreadState::decryption_buffer)::value_type>);
  1217. if constexpr (RVZ)
  1218. {
  1219. // We must not store junk efficiently for chunks that may get reused at a position
  1220. // which has a different value of data_offset % VolumeWii::BLOCK_TOTAL_SIZE
  1221. const bool allow_junk_reuse = chunks_per_wii_group == 1;
  1222. const u64 bytes_per_chunk = std::min(out_data_per_chunk, VolumeWii::GROUP_DATA_SIZE);
  1223. const u64 total_size = blocks_in_this_group * VolumeWii::BLOCK_DATA_SIZE;
  1224. const u64 data_offset = parameters.data_offset + write_offset_of_group;
  1225. RVZPack(state->decryption_buffer[0].data(), output_entries.data() + first_chunk,
  1226. bytes_per_chunk, chunks, total_size, data_offset, groups > 1, allow_junk_reuse,
  1227. compression, file_system);
  1228. }
  1229. else
  1230. {
  1231. const u8* in_ptr = state->decryption_buffer[0].data();
  1232. for (u64 j = 0; j < chunks; ++j)
  1233. {
  1234. OutputParametersEntry& entry = output_entries[first_chunk + j];
  1235. if (!entry.reused_group)
  1236. {
  1237. const u64 bytes_left = (blocks - j * blocks_per_chunk) * VolumeWii::BLOCK_DATA_SIZE;
  1238. const u64 bytes_to_write_total = std::min(out_data_per_chunk, bytes_left);
  1239. if (i == 0)
  1240. entry.main_data.resize(bytes_to_write_total);
  1241. const u64 bytes_to_write = std::min(bytes_to_write_total, VolumeWii::GROUP_DATA_SIZE);
  1242. std::memcpy(entry.main_data.data() + write_offset_of_group, in_ptr, bytes_to_write);
  1243. // Set this chunk as reusable if the decrypted data is AllSame.
  1244. // There is also a requirement that it lacks exceptions, but this is checked later
  1245. if (i == 0 && !entry.reuse_id)
  1246. {
  1247. if (AllSame(in_ptr, in_ptr + bytes_to_write))
  1248. entry.reuse_id = create_reuse_id(*in_ptr, false, j * blocks_per_chunk);
  1249. }
  1250. else
  1251. {
  1252. if (entry.reuse_id && !entry.reuse_id->encrypted &&
  1253. (!AllSame(in_ptr, in_ptr + bytes_to_write) || entry.reuse_id->value != *in_ptr))
  1254. {
  1255. entry.reuse_id.reset();
  1256. }
  1257. }
  1258. }
  1259. in_ptr += out_data_per_chunk;
  1260. }
  1261. }
  1262. }
  1263. for (size_t i = 0; i < exception_lists.size(); ++i)
  1264. {
  1265. OutputParametersEntry& entry = output_entries[chunks_per_wii_group == 1 ? 0 : i];
  1266. if (entry.reused_group)
  1267. continue;
  1268. const std::vector<HashExceptionEntry>& in = exception_lists[i];
  1269. std::vector<u8>& out = entry.exception_lists;
  1270. const u16 exceptions = Common::swap16(static_cast<u16>(in.size()));
  1271. PushBack(&out, exceptions);
  1272. for (const HashExceptionEntry& exception : in)
  1273. PushBack(&out, exception);
  1274. }
  1275. for (u64 i = 0; i < output_entries.size(); ++i)
  1276. {
  1277. OutputParametersEntry& entry = output_entries[i];
  1278. // If this chunk was set as reusable because the decrypted data is AllSame,
  1279. // but it has exceptions, unmark it as reusable
  1280. if (entry.reuse_id && !entry.reuse_id->encrypted && !AllZero(entry.exception_lists))
  1281. entry.reuse_id.reset();
  1282. }
  1283. }
  1284. }
  1285. for (OutputParametersEntry& entry : output_entries)
  1286. {
  1287. TryReuse(reusable_groups, reusable_groups_mutex, &entry);
  1288. if (entry.reused_group)
  1289. continue;
  1290. // Special case - a compressed size of zero is treated by WIA as meaning the data is all zeroes
  1291. if (entry.reuse_id && !entry.reuse_id->encrypted && entry.reuse_id->value == 0)
  1292. {
  1293. entry.exception_lists.clear();
  1294. entry.main_data.clear();
  1295. if constexpr (RVZ)
  1296. {
  1297. entry.rvz_packed_size = 0;
  1298. entry.compressed = false;
  1299. }
  1300. continue;
  1301. }
  1302. const auto pad_exception_lists = [&entry]() {
  1303. while (entry.exception_lists.size() % 4 != 0)
  1304. entry.exception_lists.push_back(0);
  1305. };
  1306. if (state->compressor)
  1307. {
  1308. if (!state->compressor->Start(entry.exception_lists.size() + entry.main_data.size()))
  1309. return ConversionResultCode::InternalError;
  1310. }
  1311. if (!entry.exception_lists.empty())
  1312. {
  1313. if (compressed_exception_lists && state->compressor)
  1314. {
  1315. if (!state->compressor->Compress(entry.exception_lists.data(),
  1316. entry.exception_lists.size()))
  1317. {
  1318. return ConversionResultCode::InternalError;
  1319. }
  1320. }
  1321. else
  1322. {
  1323. if (!compressed_exception_lists)
  1324. pad_exception_lists();
  1325. if (state->compressor)
  1326. {
  1327. if (!state->compressor->AddPrecedingDataOnlyForPurgeHashing(entry.exception_lists.data(),
  1328. entry.exception_lists.size()))
  1329. {
  1330. return ConversionResultCode::InternalError;
  1331. }
  1332. }
  1333. }
  1334. }
  1335. if (state->compressor)
  1336. {
  1337. if (!state->compressor->Compress(entry.main_data.data(), entry.main_data.size()))
  1338. return ConversionResultCode::InternalError;
  1339. if (!state->compressor->End())
  1340. return ConversionResultCode::InternalError;
  1341. }
  1342. bool compressed = !!state->compressor;
  1343. if constexpr (RVZ)
  1344. {
  1345. size_t uncompressed_size = entry.main_data.size();
  1346. if (compressed_exception_lists)
  1347. uncompressed_size += Common::AlignUp(entry.exception_lists.size(), 4);
  1348. compressed = state->compressor && state->compressor->GetSize() < uncompressed_size;
  1349. entry.compressed = compressed;
  1350. if (!compressed)
  1351. pad_exception_lists();
  1352. }
  1353. if (compressed)
  1354. {
  1355. const u8* data = state->compressor->GetData();
  1356. const size_t size = state->compressor->GetSize();
  1357. entry.main_data.resize(size);
  1358. std::copy_n(data, size, entry.main_data.data());
  1359. if (compressed_exception_lists)
  1360. entry.exception_lists.clear();
  1361. }
  1362. }
  1363. return OutputParameters{std::move(output_entries), parameters.bytes_read, parameters.group_index};
  1364. }
  1365. template <bool RVZ>
  1366. ConversionResultCode WIARVZFileReader<RVZ>::Output(std::vector<OutputParametersEntry>* entries,
  1367. File::IOFile* outfile,
  1368. std::map<ReuseID, GroupEntry>* reusable_groups,
  1369. std::mutex* reusable_groups_mutex,
  1370. GroupEntry* group_entry, u64* bytes_written)
  1371. {
  1372. for (OutputParametersEntry& entry : *entries)
  1373. {
  1374. TryReuse(reusable_groups, reusable_groups_mutex, &entry);
  1375. if (entry.reused_group)
  1376. {
  1377. *group_entry = *entry.reused_group;
  1378. ++group_entry;
  1379. continue;
  1380. }
  1381. if (*bytes_written >> 2 > std::numeric_limits<u32>::max())
  1382. return ConversionResultCode::InternalError;
  1383. ASSERT((*bytes_written & 3) == 0);
  1384. group_entry->data_offset = Common::swap32(static_cast<u32>(*bytes_written >> 2));
  1385. u32 data_size = static_cast<u32>(entry.exception_lists.size() + entry.main_data.size());
  1386. if constexpr (RVZ)
  1387. {
  1388. data_size = (data_size & 0x7FFFFFFF) | (static_cast<u32>(entry.compressed) << 31);
  1389. group_entry->rvz_packed_size = Common::swap32(static_cast<u32>(entry.rvz_packed_size));
  1390. }
  1391. group_entry->data_size = Common::swap32(data_size);
  1392. if (!outfile->WriteArray(entry.exception_lists.data(), entry.exception_lists.size()))
  1393. return ConversionResultCode::WriteFailed;
  1394. if (!outfile->WriteArray(entry.main_data.data(), entry.main_data.size()))
  1395. return ConversionResultCode::WriteFailed;
  1396. *bytes_written += entry.exception_lists.size() + entry.main_data.size();
  1397. if (entry.reuse_id)
  1398. {
  1399. std::lock_guard guard(*reusable_groups_mutex);
  1400. reusable_groups->emplace(*entry.reuse_id, *group_entry);
  1401. }
  1402. if (!PadTo4(outfile, bytes_written))
  1403. return ConversionResultCode::WriteFailed;
  1404. ++group_entry;
  1405. }
  1406. return ConversionResultCode::Success;
  1407. }
  1408. template <bool RVZ>
  1409. ConversionResultCode WIARVZFileReader<RVZ>::RunCallback(size_t groups_written, u64 bytes_read,
  1410. u64 bytes_written, u32 total_groups,
  1411. u64 iso_size, CompressCB callback)
  1412. {
  1413. int ratio = 0;
  1414. if (bytes_read != 0)
  1415. ratio = static_cast<int>(100 * bytes_written / bytes_read);
  1416. const std::string text = Common::FmtFormatT("{0} of {1} blocks. Compression ratio {2}%",
  1417. groups_written, total_groups, ratio);
  1418. const float completion = static_cast<float>(bytes_read) / iso_size;
  1419. return callback(text, completion) ? ConversionResultCode::Success :
  1420. ConversionResultCode::Canceled;
  1421. }
  1422. template <bool RVZ>
  1423. bool WIARVZFileReader<RVZ>::WriteHeader(File::IOFile* file, const u8* data, size_t size,
  1424. u64 upper_bound, u64* bytes_written, u64* offset_out)
  1425. {
  1426. // The first part of the check is to prevent this from running more than once. If *bytes_written
  1427. // is past the upper bound, we are already at the end of the file, so we don't need to do anything
  1428. if (*bytes_written <= upper_bound && *bytes_written + size > upper_bound)
  1429. {
  1430. WARN_LOG_FMT(DISCIO,
  1431. "Headers did not fit in the allocated space. Writing to end of file instead");
  1432. if (!file->Seek(0, File::SeekOrigin::End))
  1433. return false;
  1434. *bytes_written = file->Tell();
  1435. }
  1436. *offset_out = *bytes_written;
  1437. if (!file->WriteArray(data, size))
  1438. return false;
  1439. *bytes_written += size;
  1440. return PadTo4(file, bytes_written);
  1441. }
  1442. template <bool RVZ>
  1443. ConversionResultCode
  1444. WIARVZFileReader<RVZ>::Convert(BlobReader* infile, const VolumeDisc* infile_volume,
  1445. File::IOFile* outfile, WIARVZCompressionType compression_type,
  1446. int compression_level, int chunk_size, CompressCB callback)
  1447. {
  1448. ASSERT(infile->GetDataSizeType() == DataSizeType::Accurate);
  1449. ASSERT(chunk_size > 0);
  1450. const u64 iso_size = infile->GetDataSize();
  1451. const u64 chunks_per_wii_group = std::max<u64>(1, VolumeWii::GROUP_TOTAL_SIZE / chunk_size);
  1452. const u64 exception_lists_per_chunk = std::max<u64>(1, chunk_size / VolumeWii::GROUP_TOTAL_SIZE);
  1453. const bool compressed_exception_lists = compression_type > WIARVZCompressionType::Purge;
  1454. u64 bytes_read = 0;
  1455. u64 bytes_written = 0;
  1456. size_t groups_processed = 0;
  1457. WIAHeader1 header_1{};
  1458. WIAHeader2 header_2{};
  1459. std::vector<PartitionEntry> partition_entries;
  1460. std::vector<RawDataEntry> raw_data_entries;
  1461. std::vector<GroupEntry> group_entries;
  1462. u32 total_groups;
  1463. std::vector<DataEntry> data_entries;
  1464. const FileSystem* non_partition_file_system =
  1465. infile_volume ? infile_volume->GetFileSystem(PARTITION_NONE) : nullptr;
  1466. std::vector<const FileSystem*> partition_file_systems;
  1467. const ConversionResultCode set_up_data_entries_result = SetUpDataEntriesForWriting(
  1468. infile_volume, chunk_size, iso_size, &total_groups, &partition_entries, &raw_data_entries,
  1469. &data_entries, &partition_file_systems);
  1470. if (set_up_data_entries_result != ConversionResultCode::Success)
  1471. return set_up_data_entries_result;
  1472. group_entries.resize(total_groups);
  1473. const size_t partition_entries_size = partition_entries.size() * sizeof(PartitionEntry);
  1474. const size_t raw_data_entries_size = raw_data_entries.size() * sizeof(RawDataEntry);
  1475. const size_t group_entries_size = group_entries.size() * sizeof(GroupEntry);
  1476. // An estimate for how much space will be taken up by headers.
  1477. // We will reserve this much space at the beginning of the file, and if the headers don't
  1478. // fit in that space, we will need to write them at the end of the file instead.
  1479. const u64 headers_size_upper_bound = [&] {
  1480. // 0x100 is added to account for compression overhead (in particular for Purge).
  1481. u64 upper_bound = sizeof(WIAHeader1) + sizeof(WIAHeader2) + partition_entries_size +
  1482. raw_data_entries_size + 0x100;
  1483. // Compared to WIA, RVZ adds an extra member to the GroupEntry struct. This added data usually
  1484. // compresses well, so we'll assume the compression ratio for RVZ GroupEntries is 9 / 16 or
  1485. // better. This constant is somehwat arbitrarily chosen, but no games were found that get a
  1486. // worse compression ratio than that. There are some games that get a worse ratio than 1 / 2,
  1487. // such as Metroid: Other M (PAL) with the default settings.
  1488. if (RVZ && compression_type > WIARVZCompressionType::Purge)
  1489. upper_bound += static_cast<u64>(group_entries_size) * 9 / 16;
  1490. else
  1491. upper_bound += group_entries_size;
  1492. // This alignment is also somewhat arbitrary.
  1493. return Common::AlignUp(upper_bound, VolumeWii::BLOCK_TOTAL_SIZE);
  1494. }();
  1495. std::vector<u8> buffer;
  1496. buffer.resize(headers_size_upper_bound);
  1497. outfile->WriteBytes(buffer.data(), buffer.size());
  1498. bytes_written = headers_size_upper_bound;
  1499. if (!infile->Read(0, header_2.disc_header.size(), header_2.disc_header.data()))
  1500. return ConversionResultCode::ReadFailed;
  1501. // We intentionally do not increment bytes_read here, since these bytes will be read again
  1502. std::map<ReuseID, GroupEntry> reusable_groups;
  1503. std::mutex reusable_groups_mutex;
  1504. const auto set_up_compress_thread_state = [&](CompressThreadState* state) {
  1505. SetUpCompressor(&state->compressor, compression_type, compression_level, nullptr);
  1506. return ConversionResultCode::Success;
  1507. };
  1508. const auto process_and_compress = [&](CompressThreadState* state, CompressParameters parameters) {
  1509. const DataEntry& data_entry = *parameters.data_entry;
  1510. const FileSystem* file_system = data_entry.is_partition ?
  1511. partition_file_systems[data_entry.index] :
  1512. non_partition_file_system;
  1513. const bool compression = compression_type != WIARVZCompressionType::None;
  1514. return ProcessAndCompress(state, std::move(parameters), partition_entries, data_entries,
  1515. file_system, &reusable_groups, &reusable_groups_mutex,
  1516. chunks_per_wii_group, exception_lists_per_chunk,
  1517. compressed_exception_lists, compression);
  1518. };
  1519. const auto output = [&](OutputParameters parameters) {
  1520. const ConversionResultCode result =
  1521. Output(&parameters.entries, outfile, &reusable_groups, &reusable_groups_mutex,
  1522. &group_entries[parameters.group_index], &bytes_written);
  1523. if (result != ConversionResultCode::Success)
  1524. return result;
  1525. return RunCallback(parameters.group_index + parameters.entries.size(), parameters.bytes_read,
  1526. bytes_written, total_groups, iso_size, callback);
  1527. };
  1528. MultithreadedCompressor<CompressThreadState, CompressParameters, OutputParameters> mt_compressor(
  1529. set_up_compress_thread_state, process_and_compress, output);
  1530. for (const DataEntry& data_entry : data_entries)
  1531. {
  1532. u32 first_group;
  1533. u32 last_group;
  1534. u64 data_offset;
  1535. u64 data_size;
  1536. u64 data_offset_in_partition;
  1537. if (data_entry.is_partition)
  1538. {
  1539. const PartitionEntry& partition_entry = partition_entries[data_entry.index];
  1540. const PartitionDataEntry& partition_data_entry =
  1541. partition_entry.data_entries[data_entry.partition_data_index];
  1542. first_group = Common::swap32(partition_data_entry.group_index);
  1543. last_group = first_group + Common::swap32(partition_data_entry.number_of_groups);
  1544. const u32 first_sector = Common::swap32(partition_data_entry.first_sector);
  1545. data_offset = first_sector * VolumeWii::BLOCK_TOTAL_SIZE;
  1546. data_size =
  1547. Common::swap32(partition_data_entry.number_of_sectors) * VolumeWii::BLOCK_TOTAL_SIZE;
  1548. const u32 block_in_partition =
  1549. first_sector - Common::swap32(partition_entry.data_entries[0].first_sector);
  1550. data_offset_in_partition = block_in_partition * VolumeWii::BLOCK_DATA_SIZE;
  1551. }
  1552. else
  1553. {
  1554. const RawDataEntry& raw_data_entry = raw_data_entries[data_entry.index];
  1555. first_group = Common::swap32(raw_data_entry.group_index);
  1556. last_group = first_group + Common::swap32(raw_data_entry.number_of_groups);
  1557. data_offset = Common::swap64(raw_data_entry.data_offset);
  1558. data_size = Common::swap64(raw_data_entry.data_size);
  1559. const u64 skipped_data = data_offset % VolumeWii::BLOCK_TOTAL_SIZE;
  1560. data_offset -= skipped_data;
  1561. data_size += skipped_data;
  1562. data_offset_in_partition = data_offset;
  1563. }
  1564. ASSERT(groups_processed == first_group);
  1565. ASSERT(bytes_read == data_offset);
  1566. while (groups_processed < last_group)
  1567. {
  1568. const ConversionResultCode status = mt_compressor.GetStatus();
  1569. if (status != ConversionResultCode::Success)
  1570. return status;
  1571. u64 bytes_to_read = chunk_size;
  1572. if (data_entry.is_partition)
  1573. bytes_to_read = std::max<u64>(bytes_to_read, VolumeWii::GROUP_TOTAL_SIZE);
  1574. bytes_to_read = std::min<u64>(bytes_to_read, data_offset + data_size - bytes_read);
  1575. buffer.resize(bytes_to_read);
  1576. if (!infile->Read(bytes_read, bytes_to_read, buffer.data()))
  1577. return ConversionResultCode::ReadFailed;
  1578. bytes_read += bytes_to_read;
  1579. mt_compressor.CompressAndWrite(CompressParameters{
  1580. buffer, &data_entry, data_offset_in_partition, bytes_read, groups_processed});
  1581. data_offset += bytes_to_read;
  1582. data_size -= bytes_to_read;
  1583. if (data_entry.is_partition)
  1584. {
  1585. data_offset_in_partition +=
  1586. bytes_to_read / VolumeWii::BLOCK_TOTAL_SIZE * VolumeWii::BLOCK_DATA_SIZE;
  1587. }
  1588. else
  1589. {
  1590. data_offset_in_partition += bytes_to_read;
  1591. }
  1592. groups_processed += Common::AlignUp(bytes_to_read, chunk_size) / chunk_size;
  1593. }
  1594. ASSERT(data_size == 0);
  1595. }
  1596. ASSERT(groups_processed == total_groups);
  1597. ASSERT(bytes_read == iso_size);
  1598. mt_compressor.Shutdown();
  1599. const ConversionResultCode status = mt_compressor.GetStatus();
  1600. if (status != ConversionResultCode::Success)
  1601. return status;
  1602. std::unique_ptr<Compressor> compressor;
  1603. SetUpCompressor(&compressor, compression_type, compression_level, &header_2);
  1604. const std::optional<std::vector<u8>> compressed_raw_data_entries = Compress(
  1605. compressor.get(), reinterpret_cast<u8*>(raw_data_entries.data()), raw_data_entries_size);
  1606. if (!compressed_raw_data_entries)
  1607. return ConversionResultCode::InternalError;
  1608. const std::optional<std::vector<u8>> compressed_group_entries =
  1609. Compress(compressor.get(), reinterpret_cast<u8*>(group_entries.data()), group_entries_size);
  1610. if (!compressed_group_entries)
  1611. return ConversionResultCode::InternalError;
  1612. bytes_written = sizeof(WIAHeader1) + sizeof(WIAHeader2);
  1613. if (!outfile->Seek(sizeof(WIAHeader1) + sizeof(WIAHeader2), File::SeekOrigin::Begin))
  1614. return ConversionResultCode::WriteFailed;
  1615. u64 partition_entries_offset;
  1616. if (!WriteHeader(outfile, reinterpret_cast<u8*>(partition_entries.data()), partition_entries_size,
  1617. headers_size_upper_bound, &bytes_written, &partition_entries_offset))
  1618. {
  1619. return ConversionResultCode::WriteFailed;
  1620. }
  1621. u64 raw_data_entries_offset;
  1622. if (!WriteHeader(outfile, compressed_raw_data_entries->data(),
  1623. compressed_raw_data_entries->size(), headers_size_upper_bound, &bytes_written,
  1624. &raw_data_entries_offset))
  1625. {
  1626. return ConversionResultCode::WriteFailed;
  1627. }
  1628. u64 group_entries_offset;
  1629. if (!WriteHeader(outfile, compressed_group_entries->data(), compressed_group_entries->size(),
  1630. headers_size_upper_bound, &bytes_written, &group_entries_offset))
  1631. {
  1632. return ConversionResultCode::WriteFailed;
  1633. }
  1634. u32 disc_type = 0;
  1635. if (infile_volume)
  1636. {
  1637. if (infile_volume->GetVolumeType() == Platform::GameCubeDisc)
  1638. disc_type = 1;
  1639. else if (infile_volume->GetVolumeType() == Platform::WiiDisc)
  1640. disc_type = 2;
  1641. }
  1642. header_2.disc_type = Common::swap32(disc_type);
  1643. header_2.compression_type = Common::swap32(static_cast<u32>(compression_type));
  1644. header_2.compression_level =
  1645. static_cast<s32>(Common::swap32(static_cast<u32>(compression_level)));
  1646. header_2.chunk_size = Common::swap32(static_cast<u32>(chunk_size));
  1647. header_2.number_of_partition_entries = Common::swap32(static_cast<u32>(partition_entries.size()));
  1648. header_2.partition_entry_size = Common::swap32(sizeof(PartitionEntry));
  1649. header_2.partition_entries_offset = Common::swap64(partition_entries_offset);
  1650. header_2.partition_entries_hash = Common::SHA1::CalculateDigest(partition_entries);
  1651. header_2.number_of_raw_data_entries = Common::swap32(static_cast<u32>(raw_data_entries.size()));
  1652. header_2.raw_data_entries_offset = Common::swap64(raw_data_entries_offset);
  1653. header_2.raw_data_entries_size =
  1654. Common::swap32(static_cast<u32>(compressed_raw_data_entries->size()));
  1655. header_2.number_of_group_entries = Common::swap32(static_cast<u32>(group_entries.size()));
  1656. header_2.group_entries_offset = Common::swap64(group_entries_offset);
  1657. header_2.group_entries_size = Common::swap32(static_cast<u32>(compressed_group_entries->size()));
  1658. header_1.magic = RVZ ? RVZ_MAGIC : WIA_MAGIC;
  1659. header_1.version = Common::swap32(RVZ ? RVZ_VERSION : WIA_VERSION);
  1660. header_1.version_compatible =
  1661. Common::swap32(RVZ ? RVZ_VERSION_WRITE_COMPATIBLE : WIA_VERSION_WRITE_COMPATIBLE);
  1662. header_1.header_2_size = Common::swap32(sizeof(WIAHeader2));
  1663. header_1.header_2_hash =
  1664. Common::SHA1::CalculateDigest(reinterpret_cast<const u8*>(&header_2), sizeof(header_2));
  1665. header_1.iso_file_size = Common::swap64(infile->GetDataSize());
  1666. header_1.wia_file_size = Common::swap64(outfile->GetSize());
  1667. header_1.header_1_hash = Common::SHA1::CalculateDigest(reinterpret_cast<const u8*>(&header_1),
  1668. offsetof(WIAHeader1, header_1_hash));
  1669. if (!outfile->Seek(0, File::SeekOrigin::Begin))
  1670. return ConversionResultCode::WriteFailed;
  1671. if (!outfile->WriteArray(&header_1, 1))
  1672. return ConversionResultCode::WriteFailed;
  1673. if (!outfile->WriteArray(&header_2, 1))
  1674. return ConversionResultCode::WriteFailed;
  1675. return ConversionResultCode::Success;
  1676. }
  1677. bool ConvertToWIAOrRVZ(BlobReader* infile, const std::string& infile_path,
  1678. const std::string& outfile_path, bool rvz,
  1679. WIARVZCompressionType compression_type, int compression_level,
  1680. int chunk_size, CompressCB callback)
  1681. {
  1682. File::IOFile outfile(outfile_path, "wb");
  1683. if (!outfile)
  1684. {
  1685. PanicAlertFmtT(
  1686. "Failed to open the output file \"{0}\".\n"
  1687. "Check that you have permissions to write the target folder and that the media can "
  1688. "be written.",
  1689. outfile_path);
  1690. return false;
  1691. }
  1692. std::unique_ptr<VolumeDisc> infile_volume = CreateDisc(infile_path);
  1693. const auto convert = rvz ? RVZFileReader::Convert : WIAFileReader::Convert;
  1694. const ConversionResultCode result =
  1695. convert(infile, infile_volume.get(), &outfile, compression_type, compression_level,
  1696. chunk_size, callback);
  1697. if (result == ConversionResultCode::ReadFailed)
  1698. PanicAlertFmtT("Failed to read from the input file \"{0}\".", infile_path);
  1699. if (result == ConversionResultCode::WriteFailed)
  1700. {
  1701. PanicAlertFmtT("Failed to write the output file \"{0}\".\n"
  1702. "Check that you have enough space available on the target drive.",
  1703. outfile_path);
  1704. }
  1705. if (result != ConversionResultCode::Success)
  1706. {
  1707. // Remove the incomplete output file
  1708. outfile.Close();
  1709. File::Delete(outfile_path);
  1710. }
  1711. return result == ConversionResultCode::Success;
  1712. }
  1713. template class WIARVZFileReader<false>;
  1714. template class WIARVZFileReader<true>;
  1715. } // namespace DiscIO