Fixes to CIA building

1. Optimizes ticket finding
2. Fixes progress reporting overshoot
3. Use u64 for size in general
4. Various other fixes and cleanups

src/common/CMakeLists.txt

@@ -10,6 +10,7 @@ add_library(common STATIC
   logging/log.cpp
   logging/log.h
   misc.cpp
+  progress_callback.cpp
   progress_callback.h
   scope_exit.h
   string_util.cpp

src/common/progress_callback.cpp

@@ -0,0 +1,32 @@
+// Copyright 2021 threeSD Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include "common/progress_callback.h"
+
+namespace Common {
+
+ProgressCallback ProgressCallbackWrapper::Wrap(const ProgressCallback& callback) {
+    current_done_size += current_pending_size; // Last content finished
+    return [this, callback](std::size_t current, std::size_t total) {
+        current_pending_size = total;
+        callback(current + current_done_size, total_size);
+    };
+}
+
+ProgressCallback ProgressCallbackWrapper::Wrap(
+    const std::function<void(std::size_t, std::size_t, std::size_t)>& callback) {
+
+    current_done_size += current_pending_size; // Last content finished
+    return [this, callback](std::size_t current, std::size_t total) {
+        current_pending_size = total;
+        callback(current, current + current_done_size, total_size);
+    };
+}
+
+void ProgressCallbackWrapper::SetCurrent(u64 current) {
+    current_done_size = current;
+    current_pending_size = 0;
+}
+
+} // namespace Common

src/common/progress_callback.h

@@ -5,10 +5,26 @@
 #pragma once
 
 #include <functional>
+#include "common/common_types.h"
 
 namespace Common {
 
 // (current_size, total_size)
-using ProgressCallback = std::function<void(std::size_t, std::size_t)>;
+using ProgressCallback = std::function<void(u64, u64)>;
+
+class ProgressCallbackWrapper {
+public:
+    // (total imported size, total size)
+    ProgressCallback Wrap(const ProgressCallback& callback);
+
+    // (current content imported size, total imported size, total size)
+    ProgressCallback Wrap(const std::function<void(u64, u64, u64)>& callback);
+
+    void SetCurrent(u64 current);
+
+    std::size_t total_size{};
+    std::size_t current_done_size{};
+    std::size_t current_pending_size{};
+};
 
 } // namespace Common

src/core/cia_builder.cpp

@@ -53,12 +53,28 @@ private:
     bool hash_enabled{};
 };
 
-CIABuilder::CIABuilder() = default;
+CIABuilder::CIABuilder(const Config& config) {
+    if (!config.ticket_db_path.empty()) {
+        ticket_db = std::make_unique<TicketDB>(config.ticket_db_path);
+    }
+    if (!ticket_db || !ticket_db->IsGood()) {
+        LOG_WARNING(Core, "ticket.db not present or is invalid");
+        ticket_db.reset();
+    }
+
+    if (!config.enc_title_keys_bin_path.empty()) {
+        enc_title_keys_bin = std::make_unique<EncTitleKeysBin>();
+        if (!LoadTitleKeysBin(*enc_title_keys_bin, config.enc_title_keys_bin_path)) {
+            LOG_WARNING(Core, "encTitleKeys.bin invalid");
+            enc_title_keys_bin.reset();
+        }
+    }
+}
+
 CIABuilder::~CIABuilder() = default;
 
 bool CIABuilder::Init(CIABuildType type_, const std::string& destination, TitleMetadata tmd_,
-                      const Config& config, std::size_t total_size_,
-                      const Common::ProgressCallback& callback_) {
+                      std::size_t total_size_, const Common::ProgressCallback& callback_) {
 
     type = type_;
     header = {};
@@ -91,14 +107,14 @@ bool CIABuilder::Init(CIABuildType type_, const std::string& destination, TitleM
     // Cert
     cert_offset = Common::AlignUp(header.header_size, CIA_ALIGNMENT);
     header.cert_size = CIA_CERT_SIZE;
-    if (!WriteCert(config.certs_db_path)) {
+    if (!WriteCert()) {
         LOG_ERROR(Core, "Could not write cert to file {}", destination);
         return false;
     }
 
     // Ticket
     ticket_offset = Common::AlignUp(cert_offset + header.cert_size, CIA_ALIGNMENT);
-    if (!WriteTicket(config.ticket_db_path, config.enc_title_keys_bin_path)) {
+    if (!WriteTicket()) {
         return false;
     }
 
@@ -112,15 +128,23 @@ bool CIABuilder::Init(CIABuildType type_, const std::string& destination, TitleM
     // Meta will be written in Finalize
     header.meta_size = 0;
 
+    // Initialize variables
     written = content_offset;
     total_size = total_size_;
+
     callback = callback_;
+    wrapper.total_size = total_size;
+    wrapper.SetCurrent(written);
 
     callback(written, total_size);
     return true;
 }
 
-bool CIABuilder::WriteCert(const std::string& certs_db_path) {
+void CIABuilder::Cleanup() {
+    file.reset();
+}
+
+bool CIABuilder::WriteCert() {
     if (!Certs::IsLoaded()) {
         return false;
     }
@@ -135,10 +159,9 @@ bool CIABuilder::WriteCert(const std::string& certs_db_path) {
     return true;
 }
 
-static bool FindLegitTicket(Ticket& ticket, u64 title_id, const std::string& ticket_db_path) {
-    TicketDB ticket_db(ticket_db_path);
-    if (ticket_db.IsGood() && ticket_db.tickets.count(title_id)) {
-        ticket = ticket_db.tickets.at(title_id);
+bool CIABuilder::FindLegitTicket(Ticket& ticket, u64 title_id) const {
+    if (ticket_db && ticket_db->tickets.count(title_id)) {
+        ticket = ticket_db->tickets.at(title_id);
         if (!ticket.ValidateSignature()) {
             LOG_ERROR(Core, "Ticket in ticket.db for {:016x} is not legit", title_id);
             return false;
@@ -150,24 +173,17 @@ static bool FindLegitTicket(Ticket& ticket, u64 title_id, const std::string& tic
     return false;
 }
 
-static Ticket BuildStandardTicket(u64 title_id, const std::string& ticket_db_path,
-                                  const std::string& enc_title_keys_bin_path) {
+Ticket CIABuilder::BuildStandardTicket(u64 title_id) const {
     Ticket ticket = BuildFakeTicket(title_id);
 
     // Fill in common_key_index and title_key from either ticket.db (installed tickets)
     // or GM9 support files (encTitleKeys.bin) found on the SD card
-    if (TicketDB ticket_db(ticket_db_path);
-        ticket_db.IsGood() && ticket_db.tickets.count(title_id)) { // ticket.db
-
-        const auto& legit_ticket = ticket_db.tickets.at(title_id);
+    if (ticket_db && ticket_db->tickets.count(title_id)) { // ticket.db
+        const auto& legit_ticket = ticket_db->tickets.at(title_id);
         ticket.body.common_key_index = legit_ticket.body.common_key_index;
         ticket.body.title_key = legit_ticket.body.title_key;
-
-    } else if (TitleKeysMap enc_title_keys;
-               LoadTitleKeysBin(enc_title_keys, enc_title_keys_bin_path) &&
-               enc_title_keys.count(title_id)) { // support files
-
-        const auto& entry = enc_title_keys.at(title_id);
+    } else if (enc_title_keys_bin && enc_title_keys_bin->count(title_id)) { // support files
+        const auto& entry = enc_title_keys_bin->at(title_id);
         ticket.body.common_key_index = entry.common_key_index;
         ticket.body.title_key = entry.title_key;
     } else {
@@ -195,18 +211,16 @@ static Key::AESKey GetTitleKey(const Ticket& ticket) {
     return title_key;
 }
 
-bool CIABuilder::WriteTicket(const std::string& ticket_db_path,
-                             const std::string& enc_title_keys_bin_path) {
-
+bool CIABuilder::WriteTicket() {
     const auto title_id = tmd.GetTitleID();
 
     Ticket ticket;
     if (type == CIABuildType::Legit) {
-        if (!FindLegitTicket(ticket, title_id, ticket_db_path)) {
+        if (!FindLegitTicket(ticket, title_id)) {
             return false;
         }
     } else {
-        ticket = BuildStandardTicket(title_id, ticket_db_path, enc_title_keys_bin_path);
+        ticket = BuildStandardTicket(title_id);
     }
     title_key = GetTitleKey(ticket);
 
@@ -248,11 +262,10 @@ bool CIABuilder::AddContent(u16 content_id, NCCHContainer& ncch) {
     }
 
     file->Seek(written, SEEK_SET); // To enforce alignment
+    wrapper.SetCurrent(written);
 
     auto& tmd_chunk = tmd.GetContentChunkByID(content_id);
-    const auto progress_callback = [this](std::size_t current, std::size_t total) {
-        callback(written + current, total_size);
-    };
+
     if (type == CIABuildType::Standard) {
         // Decrypt the NCCH. We created a HashedFile to transparently calculate the hash as there
         // is no easy way to get decrypted NCCH content otherwise.
@@ -261,7 +274,7 @@ bool CIABuilder::AddContent(u16 content_id, NCCHContainer& ncch) {
             std::lock_guard lock{abort_ncch_mutex};
             abort_ncch = &ncch;
         }
-        const auto ret = ncch.DecryptToFile(file, progress_callback);
+        const auto ret = ncch.DecryptToFile(file, wrapper.Wrap(callback));
         {
             std::lock_guard lock{abort_ncch_mutex};
             abort_ncch = nullptr;
@@ -292,7 +305,7 @@ bool CIABuilder::AddContent(u16 content_id, NCCHContainer& ncch) {
         }
         decryptor.SetCrypto(crypto);
         if (!decryptor.CryptAndWriteFile(ncch.file, ncch.file->GetSize(), file,
-                                         progress_callback)) {
+                                         wrapper.Wrap(callback))) {
 
             return false;
         }

src/core/cia_builder.h

@@ -25,11 +25,14 @@ constexpr std::size_t CIA_CERT_SIZE = 0xA00;
 constexpr std::size_t CIA_METADATA_SIZE = 0x3AC0;
 
 struct Config;
+class EncTitleKeysBin;
 class HashedFile;
+class Ticket;
+class TicketDB;
 
 class CIABuilder {
 public:
-    explicit CIABuilder();
+    explicit CIABuilder(const Config& config);
     ~CIABuilder();
 
     /**
@@ -37,8 +40,9 @@ public:
      * @return true on success, false otherwise
      */
     bool Init(CIABuildType type, const std::string& destination, TitleMetadata tmd,
-              const Config& config, std::size_t total_size,
-              const Common::ProgressCallback& callback);
+              std::size_t total_size, const Common::ProgressCallback& callback);
+
+    void Cleanup();
 
     /**
      * Adds an NCCH content to the CIA.
@@ -93,9 +97,17 @@ private:
 
     static_assert(sizeof(Metadata) == CIA_METADATA_SIZE, "CIA Metadata structure size is wrong");
 
-    bool WriteCert(const std::string& certs_db_path);
-    bool WriteTicket(const std::string& ticket_db_path, const std::string& enc_title_keys_bin_path);
+    bool WriteCert();
 
+    bool FindLegitTicket(Ticket& ticket, u64 title_id) const;
+    Ticket BuildStandardTicket(u64 title_id) const;
+    bool WriteTicket();
+
+    // Persistent state
+    std::unique_ptr<TicketDB> ticket_db;
+    std::unique_ptr<EncTitleKeysBin> enc_title_keys_bin;
+
+    // State of a single task
     CIABuildType type;
 
     Header header{};
@@ -113,6 +125,7 @@ private:
     std::size_t written{}; // size written (with alignment)
     std::size_t total_size{};
     Common::ProgressCallback callback;
+    Common::ProgressCallbackWrapper wrapper;
 
     // The NCCH to abort on
     std::mutex abort_ncch_mutex;

src/core/db/seed_db.cpp

@@ -116,6 +116,11 @@ void Load(const std::string& path) {
     g_seeddb_loaded = g_seeddb.Load(path);
 }
 
+void Clear() {
+    g_seeddb.Clear();
+    g_seeddb_loaded = false;
+}
+
 std::optional<Seed::Data> GetSeed(u64 title_id) {
     if (!g_seeddb_loaded) {
         Load(fmt::format("{}/seeddb.bin", FileUtil::GetUserPath(FileUtil::UserPath::SysDataDir)));

src/core/db/seed_db.h

@@ -33,6 +33,10 @@ public:
     std::size_t Size() const;
     std::optional<Seed::Data> Get(u64 title_id) const;
 
+    void Clear() {
+        seeds.clear();
+    }
+
     auto begin() {
         return seeds.begin();
     }
@@ -56,6 +60,7 @@ private:
 namespace Seeds {
 
 void Load(const std::string& path);
+void Clear();
 std::optional<Seed::Data> GetSeed(u64 title_id);
 
 } // namespace Seeds

src/core/db/title_keys_bin.h

@@ -27,6 +27,7 @@ struct TitleKeysBinEntry {
 static_assert(sizeof(TitleKeysBinEntry) == 32);
 
 using TitleKeysMap = std::unordered_map<u64, TitleKeysBinEntry>;
+class EncTitleKeysBin : public TitleKeysMap {};
 
 // GM9 support files encTitleKeys.bin and decTitleKeys.bin.
 bool LoadTitleKeysBin(TitleKeysMap& out, const std::string& path);

src/core/file_decryptor.cpp

@@ -50,7 +50,7 @@ bool FileDecryptor::CryptAndWriteFile(std::shared_ptr<FileUtil::IOFile> source_,
     destination = std::move(destination_);
     callback = callback_;
 
-    current_total_size = size;
+    total_size = size;
 
     is_good = is_running = true;
 
@@ -88,7 +88,7 @@ void FileDecryptor::DataReadLoop() {
         return;
     }
 
-    std::size_t file_size = current_total_size;
+    std::size_t file_size = total_size;
 
     while (is_running && file_size > 0) {
         if (is_first_run) {
@@ -114,7 +114,7 @@ void FileDecryptor::DataReadLoop() {
 
 void FileDecryptor::DataDecryptLoop() {
     std::size_t current_buffer = 0;
-    std::size_t file_size = current_total_size;
+    std::size_t file_size = total_size;
 
     while (is_running && file_size > 0) {
         data_read_event[current_buffer].Wait();
@@ -138,7 +138,8 @@ void FileDecryptor::DataWriteLoop() {
         return;
     }
 
-    std::size_t file_size = current_total_size;
+    std::size_t file_size = total_size;
+    std::size_t imported_size = 0;
     std::size_t iteration = 0;
     /// The number of iterations each progress report covers. 32 * 16K = 512K
     constexpr std::size_t ProgressReportFreq = 32;
@@ -180,11 +181,6 @@ void FileDecryptor::Abort() {
     }
 }
 
-void FileDecryptor::Reset(std::size_t total_size_) {
-    total_size = total_size_;
-    imported_size = 0;
-}
-
 CryptoFunc::~CryptoFunc() = default;
 
 class CryptoFunc_AES_CTR final : public CryptoFunc {

src/core/file_decryptor.h

@@ -44,7 +44,7 @@ public:
     bool CryptAndWriteFile(
         std::shared_ptr<FileUtil::IOFile> source, std::size_t size,
         std::shared_ptr<FileUtil::IOFile> destination,
-        const Common::ProgressCallback& callback = [](std::size_t, std::size_t) {});
+        const Common::ProgressCallback& callback = [](u64, u64) {});
 
     void DataReadLoop();
     void DataDecryptLoop();
@@ -52,9 +52,6 @@ public:
 
     void Abort();
 
-    /// Reset the imported_size counter for this content and set a new total_size.
-    void Reset(std::size_t total_size);
-
 private:
     static constexpr std::size_t BufferSize = 16 * 1024; // 16 KB
 
@@ -62,12 +59,7 @@ private:
     std::shared_ptr<FileUtil::IOFile> destination;
     std::shared_ptr<CryptoFunc> crypto;
 
-    // Total size of this content, may consist of multiple files
     std::size_t total_size{};
-    // Total size of the current file to process
-    std::size_t current_total_size{};
-    // Total imported size for this content
-    std::size_t imported_size{};
 
     std::array<std::array<u8, BufferSize>, 3> buffers;
     std::array<Common::Event, 3> data_read_event;

src/core/file_sys/certificate.cpp

@@ -97,7 +97,7 @@ static std::unordered_map<std::string, Certificate> g_certs;
 static bool g_is_loaded = false;
 
 bool Load(const std::string& path) {
-    g_certs.clear();
+    Clear();
 
     FileUtil::IOFile file(path, "rb");
     DataContainer container(file.GetData());
@@ -151,6 +151,11 @@ bool IsLoaded() {
     return g_is_loaded;
 }
 
+void Clear() {
+    g_is_loaded = false;
+    g_certs.clear();
+}
+
 const Certificate& Get(const std::string& name) {
     return g_certs.at(name);
 }

src/core/file_sys/certificate.h

@@ -61,6 +61,7 @@ namespace Certs {
 
 bool Load(const std::string& path);
 bool IsLoaded();
+void Clear();
 const Certificate& Get(const std::string& name);
 bool Exists(const std::string& name);
 

src/core/file_sys/ncch_container.cpp

@@ -447,19 +447,12 @@ bool NCCHContainer::DecryptToFile(std::shared_ptr<FileUtil::IOFile> dest_file,
 
         const auto size = file->GetSize();
 
-        decryptor.Reset(size);
         decryptor.SetCrypto(nullptr);
         return decryptor.CryptAndWriteFile(file, size, dest_file, callback);
     }
 
     const auto total_size = file->GetSize();
-    decryptor.Reset(total_size); // This is inaccurate but doesn't really matter as we don't use it
-
     std::size_t written{};
-    const auto decryptor_callback = [total_size, &written, &callback](std::size_t current,
-                                                                      std::size_t /*total*/) {
-        callback(written + current, total_size);
-    };
 
     // Write NCCH header
     NCCH_Header modified_header = ncch_header;
@@ -487,6 +480,7 @@ bool NCCHContainer::DecryptToFile(std::shared_ptr<FileUtil::IOFile> dest_file,
         written += sizeof(ExHeader_Header);
     }
 
+    Common::ProgressCallbackWrapper wrapper(total_size);
     const auto Write = [&](std::string_view name, std::size_t offset, std::size_t size,
                            bool decrypt = false, const Key::AESKey& key = {},
                            const Key::AESKey& ctr = {}, std::size_t aes_seek_pos = 0) {
@@ -518,9 +512,10 @@ bool NCCHContainer::DecryptToFile(std::shared_ptr<FileUtil::IOFile> dest_file,
         }
 
         written = offset;
+        wrapper.SetCurrent(written);
 
         decryptor.SetCrypto(decrypt ? CreateCTRCrypto(key, ctr, aes_seek_pos) : nullptr);
-        if (!decryptor.CryptAndWriteFile(file, size, dest_file, decryptor_callback)) {
+        if (!decryptor.CryptAndWriteFile(file, size, dest_file, wrapper.Wrap(callback))) {
             LOG_ERROR(Core, "Could not write {}", name);
             return false;
         }

src/core/file_sys/ncch_container.h

@@ -271,7 +271,7 @@ public:
      */
     bool DecryptToFile(
         std::shared_ptr<FileUtil::IOFile> dest_file,
-        const Common::ProgressCallback& callback = [](std::size_t, std::size_t) {});
+        const Common::ProgressCallback& callback = [](u64, u64) {});
 
     /**
      * Aborts DecryptToFile. Simply aborts the decryptor.

src/core/importer.cpp

@@ -28,7 +28,11 @@ SDMCImporter::SDMCImporter(const Config& config_) : config(config_) {
     is_good = Init();
 }
 
-SDMCImporter::~SDMCImporter() = default;
+SDMCImporter::~SDMCImporter() {
+    // Unload global DBs
+    Certs::Clear();
+    Seeds::Clear();
+}
 
 bool SDMCImporter::Init() {
     ASSERT_MSG(!config.sdmc_path.empty() && !config.user_path.empty() &&
@@ -53,6 +57,7 @@ bool SDMCImporter::Init() {
         return false;
     }
 
+    // Load global DBs
     if (!config.seed_db_path.empty()) {
         Seeds::Load(config.seed_db_path);
     }
@@ -60,7 +65,9 @@ bool SDMCImporter::Init() {
         Certs::Load(config.certs_db_path);
     }
 
+    // Create children
     sdmc_decryptor = std::make_unique<SDMCDecryptor>(config.sdmc_path);
+    cia_builder = std::make_unique<CIABuilder>(config);
 
     // Load SDMC Title DB
     {
@@ -138,16 +145,16 @@ bool SDMCImporter::ImportContentImpl(const ContentSpecifier& specifier,
 
 namespace {
 
-template <typename Dec>
-bool ImportTitleGeneric(Dec& decryptor, const std::string& base_path,
-                        const ContentSpecifier& specifier,
-                        const std::function<bool(const std::string&)>& decryption_func) {
+using DecryptionFunc = std::function<bool(const std::string&, const Common::ProgressCallback&)>;
+bool ImportTitleGeneric(const std::string& base_path, const ContentSpecifier& specifier,
+                        const Common::ProgressCallback& callback,
+                        const DecryptionFunc& decryption_func) {
 
-    decryptor.Reset(specifier.maximum_size);
+    Common::ProgressCallbackWrapper wrapper{specifier.maximum_size};
     const FileUtil::DirectoryEntryCallable DirectoryEntryCallback =
-        [size = base_path.size(), &DirectoryEntryCallback,
-         &decryption_func](u64* /*num_entries_out*/, const std::string& directory,
-                           const std::string& virtual_name) {
+        [size = base_path.size(), &DirectoryEntryCallback, &callback, &decryption_func,
+         &wrapper](u64* /*num_entries_out*/, const std::string& directory,
+                   const std::string& virtual_name) {
             if (FileUtil::IsDirectory(directory + virtual_name + "/")) {
                 if (virtual_name == "cmd") {
                     return true; // Skip cmd (not used in Citra)
@@ -157,7 +164,7 @@ bool ImportTitleGeneric(Dec& decryptor, const std::string& base_path,
                                                        DirectoryEntryCallback);
             }
             const auto filepath = (directory + virtual_name).substr(size - 1);
-            return decryption_func(filepath);
+            return decryption_func(filepath, wrapper.Wrap(callback));
         };
     const auto path = fmt::format("title/{:08x}/{:08x}/content/", (specifier.id >> 32),
                                   (specifier.id & 0xFFFFFFFF));
@@ -169,15 +176,15 @@ bool ImportTitleGeneric(Dec& decryptor, const std::string& base_path,
 bool SDMCImporter::ImportTitle(const ContentSpecifier& specifier,
                                const Common::ProgressCallback& callback) {
     return ImportTitleGeneric(
-        *sdmc_decryptor, config.sdmc_path, specifier,
-        [this, &callback](const std::string& filepath) {
+        config.sdmc_path, specifier, callback,
+        [this](const std::string& filepath, const Common::ProgressCallback& wrapped_callback) {
             return sdmc_decryptor->DecryptAndWriteFile(
                 filepath,
                 FileUtil::GetUserPath(FileUtil::UserPath::SDMCDir) +
                     "Nintendo "
                     "3DS/00000000000000000000000000000000/00000000000000000000000000000000" +
                     filepath,
-                callback);
+                wrapped_callback);
         });
 }
 
@@ -188,8 +195,9 @@ bool SDMCImporter::ImportNandTitle(const ContentSpecifier& specifier,
         config.system_titles_path.substr(0, config.system_titles_path.size() - 6);
     FileDecryptor decryptor;
     return ImportTitleGeneric(
-        decryptor, base_path, specifier,
-        [&base_path, &decryptor, &callback](const std::string& filepath) {
+        base_path, specifier, callback,
+        [&base_path, &decryptor](const std::string& filepath,
+                                 const Common::ProgressCallback& wrapped_callback) {
             const auto physical_path = base_path + filepath.substr(1);
             const auto citra_path = FileUtil::GetUserPath(FileUtil::UserPath::NANDDir) +
                                     "00000000000000000000000000000000" + filepath;
@@ -201,7 +209,7 @@ bool SDMCImporter::ImportNandTitle(const ContentSpecifier& specifier,
             return decryptor.CryptAndWriteFile(
                 std::make_shared<FileUtil::IOFile>(physical_path, "rb"),
                 FileUtil::GetSize(physical_path),
-                std::make_shared<FileUtil::IOFile>(citra_path, "wb"), callback);
+                std::make_shared<FileUtil::IOFile>(citra_path, "wb"), wrapped_callback);
         });
 }
 
@@ -716,18 +724,9 @@ bool SDMCImporter::BuildCIA(CIABuildType build_type, const ContentSpecifier& spe
         }
     }
 
-    {
-        std::lock_guard lock{cia_builder_mutex};
-        cia_builder = std::make_unique<CIABuilder>();
-    }
-
-    bool ret = cia_builder->Init(build_type, destination, tmd, config,
-                                 FileUtil::GetDirectoryTreeSize(physical_path), callback);
+    bool ret = cia_builder->Init(build_type, destination, tmd, specifier.maximum_size, callback);
     SCOPE_EXIT({
-        {
-            std::lock_guard lock{cia_builder_mutex};
-            cia_builder.reset(); // To release file handles, etc
-        }
+        cia_builder->Cleanup();
         if (!ret) { // Remove borked file
             FileUtil::Delete(destination);
         }
@@ -782,12 +781,12 @@ bool SDMCImporter::BuildCIA(CIABuildType build_type, const ContentSpecifier& spe
 }
 
 void SDMCImporter::AbortBuildCIA() {
-    std::lock_guard lock{cia_builder_mutex};
-    if (cia_builder) {
-        cia_builder->Abort();
-    }
+    cia_builder->Abort();
 }
 
+// Add a certain amount to the titles' maximum sizes, so that they are always larger than CIA sizes
+constexpr u64 TitleSizeAllowance = 0xA000;
+
 void SDMCImporter::ListTitle(std::vector<ContentSpecifier>& out) const {
     const auto ProcessDirectory = [this, &out, &sdmc_path = config.sdmc_path](ContentType type,
                                                                               u64 high_id) {
@@ -836,10 +835,11 @@ void SDMCImporter::ListTitle(std::vector<ContentSpecifier>& out) const {
 
                         const auto& [name, extdata_id, encryption, seed_crypto, icon] =
                             LoadTitleData(ncch);
-                        out.push_back(
-                            {type, id, FileUtil::Exists(citra_path + "content/"),
-                             FileUtil::GetDirectoryTreeSize(directory + virtual_name + "/content/"),
-                             name, extdata_id, encryption, seed_crypto, icon});
+                        const auto size =
+                            FileUtil::GetDirectoryTreeSize(directory + virtual_name + "/content/") +
+                            TitleSizeAllowance;
+                        out.push_back({type, id, FileUtil::Exists(citra_path + "content/"), size,
+                                       name, extdata_id, encryption, seed_crypto, icon});
                     } while (false);
                 }
 
@@ -907,9 +907,6 @@ void SDMCImporter::ListNandTitle(std::vector<ContentSpecifier>& out) const {
                             std::make_shared<FileUtil::IOFile>(boot_content_path, "rb"));
                         if (!ncch.Load()) {
                             LOG_WARNING(Core, "Could not load NCCH {}", boot_content_path);
-                            out.push_back({ContentType::SystemTitle, id,
-                                           FileUtil::Exists(citra_path + "content/"),
-                                           FileUtil::GetDirectoryTreeSize(content_path)});
                             break;
                         }
 
@@ -917,10 +914,11 @@ void SDMCImporter::ListNandTitle(std::vector<ContentSpecifier>& out) const {
                             LoadTitleData(ncch);
                         const auto type = (id >> 32) == 0x00040030 ? ContentType::SystemApplet
                                                                    : ContentType::SystemTitle;
-                        out.push_back(
-                            {type, id, FileUtil::Exists(citra_path + "content/"),
-                             FileUtil::GetDirectoryTreeSize(directory + virtual_name + "/content/"),
-                             name, extdata_id, encryption, seed_crypto, icon});
+                        const auto size =
+                            FileUtil::GetDirectoryTreeSize(directory + virtual_name + "/content/") +
+                            TitleSizeAllowance;
+                        out.push_back({type, id, FileUtil::Exists(citra_path + "content/"), size,
+                                       name, extdata_id, encryption, seed_crypto, icon});
                     } while (false);
                 }
                 return true;

src/core/importer.h

@@ -6,7 +6,6 @@
 
 #include <functional>
 #include <memory>
-#include <mutex>
 #include <string>
 #include <vector>
 #include "common/common_types.h"
@@ -17,6 +16,7 @@ namespace Core {
 
 class CIABuilder;
 class SDMCDecryptor;
+class TicketDB;
 class TitleDB;
 class TitleMetadata;
 
@@ -120,7 +120,7 @@ public:
      */
     bool ImportContent(
         const ContentSpecifier& specifier,
-        const Common::ProgressCallback& callback = [](std::size_t, std::size_t) {});
+        const Common::ProgressCallback& callback = [](u64, u64) {});
 
     /**
      * Aborts current importing.
@@ -175,7 +175,7 @@ private:
     // Impl of ImportContent without deleting mechanism.
     bool ImportContentImpl(
         const ContentSpecifier& specifier,
-        const Common::ProgressCallback& callback = [](std::size_t, std::size_t) {});
+        const Common::ProgressCallback& callback = [](u64, u64) {});
     bool ImportTitle(const ContentSpecifier& specifier, const Common::ProgressCallback& callback);
     bool ImportNandTitle(const ContentSpecifier& specifier,
                          const Common::ProgressCallback& callback);
@@ -207,8 +207,8 @@ private:
     Config config;
     std::unique_ptr<SDMCDecryptor> sdmc_decryptor;
 
+    // Used for CIA building
     std::unique_ptr<CIABuilder> cia_builder;
-    std::mutex cia_builder_mutex;
 
     // The NCCH used to dump CXIs.
     std::unique_ptr<NCCHContainer> dump_cxi_ncch;

src/core/sdmc_decryptor.cpp

@@ -48,10 +48,6 @@ std::array<u8, 16> GetFileCTR(const std::string& path) {
 }
 } // namespace
 
-void SDMCDecryptor::Reset(std::size_t total_size) {
-    file_decryptor.Reset(total_size);
-}
-
 bool SDMCDecryptor::DecryptAndWriteFile(const std::string& source, const std::string& destination,
                                         const Common::ProgressCallback& callback) {
     if (!FileUtil::CreateFullPath(destination)) {

src/core/sdmc_decryptor.h

@@ -33,7 +33,7 @@ public:
      */
     bool DecryptAndWriteFile(
         const std::string& source, const std::string& destination,
-        const Common::ProgressCallback& callback = [](std::size_t, std::size_t) {});
+        const Common::ProgressCallback& callback = [](u64, u64) {});
 
     void Abort();
 
@@ -43,14 +43,6 @@ public:
      */
     std::vector<u8> DecryptFile(const std::string& source) const;
 
-    /**
-     * Marks the beginning of a new content, resetting imported_size counter, and setting an new
-     * total_size for the next content.
-     * This doesn't affect at all how the contents will be imported, but will make sure the callback
-     * is properly invoked.
-     */
-    void Reset(std::size_t total_size);
-
 private:
     std::string root_folder;
     FileDecryptor file_decryptor;

src/frontend/helpers/multi_job.cpp

@@ -19,35 +19,35 @@ void MultiJob::run() {
         total_size += content.maximum_size;
     }
 
-    u64 size_imported = 0, count = 0;
+    std::size_t count = 0;
     int eta = -1;
 
     const auto initial_time = std::chrono::steady_clock::now();
     const auto UpdateETA = [total_size, &eta, initial_time](u64 size_imported) {
-        if (size_imported >= 10 * 1024 * 1024) { // 10M Threshold
-            using namespace std::chrono;
-            const u64 time_elapsed =
-                duration_cast<milliseconds>(steady_clock::now() - initial_time).count();
-            eta = static_cast<int>(time_elapsed * (total_size - size_imported) / (size_imported) /
-                                   1000);
+        if (size_imported < 10 * 1024 * 1024) { // 10M Threshold
+            return;
         }
+        using namespace std::chrono;
+        const u64 time_elapsed =
+            duration_cast<milliseconds>(steady_clock::now() - initial_time).count();
+        eta =
+            static_cast<int>(time_elapsed * (total_size - size_imported) / (size_imported) / 1000);
+    };
+    const auto Callback = [this, &eta, &UpdateETA](u64 current_imported_size,
+                                                   u64 total_imported_size, u64 /*total_size*/) {
+        UpdateETA(total_imported_size);
+        emit ProgressUpdated(current_imported_size, total_imported_size, eta);
     };
 
+    Common::ProgressCallbackWrapper wrapper(total_size);
     for (const auto& content : contents) {
-        emit NextContent(size_imported, count + 1, content, eta);
-        const auto callback = [this, size_imported, &eta, &UpdateETA](std::size_t current_size,
-                                                                      std::size_t /*total_size*/) {
-            UpdateETA(size_imported + current_size);
-            emit ProgressUpdated(size_imported + current_size, current_size, eta);
-        };
-        if (!execute_func(importer, content, callback)) {
+        emit NextContent(count + 1, content, eta);
+        if (!execute_func(importer, content, wrapper.Wrap(Callback))) {
             if (!cancelled) {
                 failed_contents.emplace_back(content);
             }
         }
         count++;
-        size_imported += content.maximum_size;
-        UpdateETA(size_imported);
 
         if (cancelled) {
             break;

src/frontend/helpers/multi_job.h

@@ -31,16 +31,15 @@ public:
 signals:
     /**
      * Called when progress is updated on the current content.
-     * @param total_size_imported Total imported size taking all previous contents into
+     * @param current_imported_size Imported size of the current content.
+     * @param total_imported_size Total imported size taking all previous contents into
      * consideration.
-     * @param current_size_imported Imported size of the current content.
      * @param eta ETA in seconds, 0 when not determined.
      */
-    void ProgressUpdated(u64 total_size_imported, u64 current_size_imported, int eta);
+    void ProgressUpdated(u64 current_imported_size, u64 total_imported_size, int eta);
 
     /// Dumping of a content has been finished, go on to the next. Called at start as well.
-    void NextContent(u64 size_imported, u64 count, const Core::ContentSpecifier& next_content,
-                     int eta);
+    void NextContent(std::size_t count, const Core::ContentSpecifier& next_content, int eta);
 
     void Completed();
 

src/frontend/helpers/simple_job.cpp

@@ -10,8 +10,8 @@ SimpleJob::SimpleJob(QObject* parent, ExecuteFunc execute_, AbortFunc abort_)
 SimpleJob::~SimpleJob() = default;
 
 void SimpleJob::run() {
-    const bool ret = execute(
-        [this](std::size_t current, std::size_t total) { emit ProgressUpdated(current, total); });
+    const bool ret =
+        execute([this](u64 current, u64 total) { emit ProgressUpdated(current, total); });
 
     if (ret || canceled) {
         emit Completed();

src/frontend/import_dialog.cpp

@@ -93,11 +93,11 @@ static QPixmap GetContentIcon(const Core::ContentSpecifier& specifier,
                                      QImage::Format::Format_RGB16));
 }
 
-ImportDialog::ImportDialog(QWidget* parent, const Core::Config& config)
-    : QDialog(parent), ui(std::make_unique<Ui::ImportDialog>()), user_path(config.user_path),
-      has_cert_db(!config.certs_db_path.empty()), importer(config) {
+ImportDialog::ImportDialog(QWidget* parent, const Core::Config& config_)
+    : QDialog(parent), ui(std::make_unique<Ui::ImportDialog>()), config(config_) {
 
     qRegisterMetaType<u64>("u64");
+    qRegisterMetaType<std::size_t>("std::size_t");
     qRegisterMetaType<Core::ContentSpecifier>();
 
     ui->setupUi(this);
@@ -105,12 +105,8 @@ ImportDialog::ImportDialog(QWidget* parent, const Core::Config& config)
     const double scale = qApp->desktop()->logicalDpiX() / 96.0;
     resize(static_cast<int>(width() * scale), static_cast<int>(height() * scale));
 
-    if (!importer.IsGood()) {
-        QMessageBox::critical(
-            this, tr("Importer Error"),
-            tr("Failed to initalize the importer.\nRefer to the log for details."));
-        reject();
-    }
+    RelistContent();
+    UpdateSizeDisplay();
 
     ui->title_view_button->setChecked(true);
 
@@ -128,9 +124,6 @@ ImportDialog::ImportDialog(QWidget* parent, const Core::Config& config)
     connect(ui->title_view_button, &QRadioButton::toggled, this, &ImportDialog::RepopulateContent);
     connect(ui->advanced_button, &QPushButton::clicked, this, &ImportDialog::ShowAdvancedMenu);
 
-    RelistContent();
-    UpdateSizeDisplay();
-
     // Set up column widths.
     // These values are tweaked with regard to the default dialog size.
     ui->main->setColumnWidth(0, width() * 0.11);
@@ -155,12 +148,25 @@ void ImportDialog::RelistContent() {
     auto* future_watcher = new FutureWatcher(this);
     connect(future_watcher, &FutureWatcher::finished, this, [this, dialog] {
         dialog->hide();
-        RepopulateContent();
+        if (importer->IsGood()) {
+            RepopulateContent();
+        } else {
+            QMessageBox::critical(
+                this, tr("Importer Error"),
+                tr("Failed to initalize the importer.\nRefer to the log for details."));
+            reject();
+        }
     });
 
-    auto future = QtConcurrent::run([&contents = this->contents, &importer = this->importer] {
-        contents = importer.ListContent();
-    });
+    auto future = QtConcurrent::run(
+        [&importer = this->importer, &config = this->config, &contents = this->contents] {
+            if (!importer) {
+                importer = std::make_unique<Core::SDMCImporter>(config);
+            }
+            if (importer->IsGood()) {
+                contents = importer->ListContent();
+            }
+        });
     future_watcher->setFuture(future);
 }
 
@@ -323,6 +329,12 @@ void ImportDialog::InsertSecondLevelItem(std::size_t row, const Core::ContentSpe
 }
 
 void ImportDialog::RepopulateContent() {
+    if (contents.empty()) { // why???
+        QMessageBox::warning(this, tr("threeSD"), tr("Sorry, there are no contents available."));
+        reject();
+        return;
+    }
+
     total_selected_size = 0;
     ui->main->clear();
     ui->main->setSortingEnabled(false);
@@ -418,13 +430,14 @@ void ImportDialog::RepopulateContent() {
 }
 
 void ImportDialog::UpdateSizeDisplay() {
-    QStorageInfo storage(QString::fromStdString(user_path));
+    QStorageInfo storage(QString::fromStdString(config.user_path));
     if (!storage.isValid() || !storage.isReady()) {
-        LOG_ERROR(Frontend, "Storage {} is not good", user_path);
+        LOG_ERROR(Frontend, "Storage {} is not good", config.user_path);
         QMessageBox::critical(
             this, tr("Bad Storage"),
             tr("An error occured while trying to get available space for the storage."));
         reject();
+        return;
     }
 
     ui->availableSpace->setText(
@@ -602,9 +615,8 @@ void ImportDialog::RunMultiJob(MultiJob* job, std::size_t total_count, u64 total
     dialog->setMinimumDuration(0);
 
     connect(job, &MultiJob::NextContent, this,
-            [this, bar, dialog, multiplier, total_count](
-                u64 size_imported, u64 count, const Core::ContentSpecifier& next_content, int eta) {
-                bar->setValue(static_cast<int>(size_imported / multiplier));
+            [this, bar, dialog, multiplier,
+             total_count](std::size_t count, const Core::ContentSpecifier& next_content, int eta) {
                 dialog->setLabelText(
                     tr("<p>(%1/%2) %3 (%4)</p><p>&nbsp;</p><p align=\"right\">%5</p>")
                         .arg(count)
@@ -616,16 +628,16 @@ void ImportDialog::RunMultiJob(MultiJob* job, std::size_t total_count, u64 total
                 current_count = count;
             });
     connect(job, &MultiJob::ProgressUpdated, this,
-            [this, bar, dialog, multiplier, total_count](u64 total_size_imported,
-                                                         u64 current_size_imported, int eta) {
-                bar->setValue(static_cast<int>(total_size_imported / multiplier));
+            [this, bar, dialog, multiplier, total_count](u64 current_imported_size,
+                                                         u64 total_imported_size, int eta) {
+                bar->setValue(static_cast<int>(total_imported_size / multiplier));
                 dialog->setLabelText(tr("<p>(%1/%2) %3 (%4)</p><p align=\"center\">%5 "
                                         "/ %6</p><p align=\"right\">%7</p>")
                                          .arg(current_count)
                                          .arg(total_count)
                                          .arg(GetContentName(current_content))
                                          .arg(GetContentTypeName(current_content.type))
-                                         .arg(ReadableByteSize(current_size_imported))
+                                         .arg(ReadableByteSize(current_imported_size))
                                          .arg(ReadableByteSize(current_content.maximum_size))
                                          .arg(FormatETA(eta)));
             });
@@ -709,7 +721,7 @@ void ImportDialog::StartImporting() {
     const std::size_t total_count = to_import.size();
 
     auto* job =
-        new MultiJob(this, importer, std::move(to_import), &Core::SDMCImporter::ImportContent,
+        new MultiJob(this, *importer, std::move(to_import), &Core::SDMCImporter::ImportContent,
                      &Core::SDMCImporter::AbortImporting);
 
     RunMultiJob(job, total_count, total_selected_size);
@@ -728,9 +740,9 @@ void ImportDialog::StartDumpingCXISingle(const Core::ContentSpecifier& specifier
     auto* job = new SimpleJob(
         this,
         [this, specifier, path](const Common::ProgressCallback& callback) {
-            return importer.DumpCXI(specifier, path.toStdString(), callback);
+            return importer->DumpCXI(specifier, path.toStdString(), callback);
         },
-        [this] { importer.AbortDumpCXI(); });
+        [this] { importer->AbortDumpCXI(); });
     RunSimpleJob(job);
 }
 
@@ -771,13 +783,12 @@ void ImportDialog::StartBatchDumpingCXI() {
     }
 
     const auto total_count = to_import.size();
-    const auto total_size =
-        std::accumulate(to_import.begin(), to_import.end(), std::size_t{0},
-                        [](std::size_t sum, const Core::ContentSpecifier& specifier) {
-                            return sum + specifier.maximum_size;
-                        });
+    const auto total_size = std::accumulate(to_import.begin(), to_import.end(), u64{0},
+                                            [](u64 sum, const Core::ContentSpecifier& specifier) {
+                                                return sum + specifier.maximum_size;
+                                            });
     auto* job = new MultiJob(
-        this, importer, std::move(to_import),
+        this, *importer, std::move(to_import),
         [path](Core::SDMCImporter& importer, const Core::ContentSpecifier& specifier,
                const Common::ProgressCallback& callback) {
             return importer.DumpCXI(specifier, path.toStdString(), callback, true);
@@ -792,7 +803,7 @@ void ImportDialog::StartBuildingCIASingle(const Core::ContentSpecifier& specifie
     CIABuildDialog dialog(this,
                           /*is_dir*/ false,
                           /*is_nand*/ specifier.type == Core::ContentType::SystemTitle,
-                          /*enable_legit*/ importer.CanBuildLegitCIA(specifier),
+                          /*enable_legit*/ importer->CanBuildLegitCIA(specifier),
                           last_build_cia_path);
     if (dialog.exec() != QDialog::Accepted) {
         return;
@@ -803,9 +814,9 @@ void ImportDialog::StartBuildingCIASingle(const Core::ContentSpecifier& specifie
     auto* job = new SimpleJob(
         this,
         [this, specifier, path = path, type = type](const Common::ProgressCallback& callback) {
-            return importer.BuildCIA(type, specifier, path.toStdString(), callback);
+            return importer->BuildCIA(type, specifier, path.toStdString(), callback);
         },
-        [this] { importer.AbortBuildCIA(); });
+        [this] { importer->AbortBuildCIA(); });
     RunSimpleJob(job);
 }
 
@@ -845,7 +856,7 @@ void ImportDialog::StartBatchBuildingCIA() {
                                      });
     const bool enable_legit = std::all_of(to_import.begin(), to_import.end(),
                                           [this](const Core::ContentSpecifier& specifier) {
-                                              return importer.CanBuildLegitCIA(specifier);
+                                              return importer->CanBuildLegitCIA(specifier);
                                           });
     CIABuildDialog dialog(this, /*is_dir*/ true, is_nand, enable_legit, last_batch_build_cia_path);
     if (dialog.exec() != QDialog::Accepted) {
@@ -858,13 +869,12 @@ void ImportDialog::StartBatchBuildingCIA() {
     }
 
     const auto total_count = to_import.size();
-    const auto total_size =
-        std::accumulate(to_import.begin(), to_import.end(), std::size_t{0},
-                        [](std::size_t sum, const Core::ContentSpecifier& specifier) {
-                            return sum + specifier.maximum_size;
-                        });
+    const auto total_size = std::accumulate(to_import.begin(), to_import.end(), u64{0},
+                                            [](u64 sum, const Core::ContentSpecifier& specifier) {
+                                                return sum + specifier.maximum_size;
+                                            });
     auto* job = new MultiJob(
-        this, importer, std::move(to_import),
+        this, *importer, std::move(to_import),
         [path = path, type = type](Core::SDMCImporter& importer,
                                    const Core::ContentSpecifier& specifier,
                                    const Common::ProgressCallback& callback) {

src/frontend/import_dialog.h

@@ -63,9 +63,9 @@ private:
 
     std::unique_ptr<Ui::ImportDialog> ui;
 
-    std::string user_path;
-    bool has_cert_db = false;
-    Core::SDMCImporter importer;
+    std::unique_ptr<Core::SDMCImporter> importer;
+    const Core::Config config;
+
     std::vector<Core::ContentSpecifier> contents;
     u64 total_selected_size = 0;
 
@@ -84,5 +84,5 @@ private:
 
     // TODO: Why this won't work as locals?
     Core::ContentSpecifier current_content = {};
-    u64 current_count = 0;
+    std::size_t current_count = 0;
 };