refactor: use threads!

src/jobs.hpp

@@ -0,0 +1,170 @@
+#pragma once
+#include <windows.h>
+
+#include <filesystem>
+#include <fstream>
+#include <mutex>
+#include <queue>
+#include <string>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include "qmc2/qmc2.h"
+
+class KggTask {
+   public:
+    explicit KggTask(std::filesystem::path kgg_path, std::filesystem::path out_dir)
+        : kgg_path_(std::move(kgg_path)), out_dir_(std::move(out_dir)) {}
+
+    void warning(const wchar_t* msg) const { fwprintf(stderr, L"[WARN] %s (%s)\n", msg, kgg_path_.filename().c_str()); }
+    void warning(const std::wstring& msg) const { warning(msg.c_str()); }
+
+    void error(const wchar_t* msg) const { fwprintf(stderr, L"[ERR ] %s (%s)\n", msg, kgg_path_.filename().c_str()); }
+    void error(const std::wstring& msg) const { error(msg.c_str()); }
+
+    void info(const wchar_t* msg) const { fwprintf(stderr, L"[INFO] %s (%s)\n", msg, kgg_path_.filename().c_str()); }
+    void info(const std::wstring& msg) const { info(msg.c_str()); }
+
+    void Execute(const Infra::kgm_ekey_db_t& ekey_db) {
+        constexpr static std::array<uint8_t, 16> kMagicHeader{0x7C, 0xD5, 0x32, 0xEB, 0x86, 0x02, 0x7F, 0x4B,
+                                                              0xA8, 0xAF, 0xA6, 0x8E, 0x0F, 0xFF, 0x99, 0x14};
+
+        std::ifstream kgg_stream_in(kgg_path_, std::ios::binary);
+        char header[0x100]{};
+        kgg_stream_in.read(header, sizeof(kgg_stream_in));
+        if (std::equal(kMagicHeader.cbegin(), kMagicHeader.cend(), header)) {
+            warning(L"invalid kgg header");
+            return;
+        }
+        uint32_t offset_to_audio = *reinterpret_cast<uint32_t*>(&header[0x10]);
+        uint32_t encrypt_mode = *reinterpret_cast<uint32_t*>(&header[0x14]);
+        if (encrypt_mode != 5) {
+            warning(std::format(L"unsupported enc_version (expect=0x05, got 0x{:02x})", encrypt_mode));
+            return;
+        }
+        uint32_t audio_hash_len = *reinterpret_cast<uint32_t*>(&header[0x44]);
+        if (audio_hash_len != 0x20) {
+            warning(std::format(L"audio hash length invalid (expect=0x20, got 0x{:02x})", audio_hash_len));
+            return;
+        }
+        std::string audio_hash(&header[0x48], &header[0x48 + audio_hash_len]);
+        std::string ekey{};
+        if (auto it = ekey_db.find(audio_hash); it != ekey_db.end()) {
+            ekey = it->second;
+        } else {
+            warning(L"ekey not found");
+            return;
+        }
+
+        auto qmc2 = QMC2::Create(ekey);
+
+        std::string magic(4, 0);
+        kgg_stream_in.seekg(offset_to_audio, std::ios::beg);
+        kgg_stream_in.read(magic.data(), 4);
+        qmc2->Decrypt(std::span(reinterpret_cast<uint8_t*>(magic.data()), 4), 0);
+        auto real_ext = DetectRealExt(magic);
+        auto out_path = out_dir_ / std::format(L"{}_kgg-dec.{}", kgg_path_.stem().wstring(), real_ext);
+
+        if (exists(out_path)) {
+            warning(std::format(L"output file already exists: {}", out_path.filename().wstring()));
+            return;
+        }
+
+        kgg_stream_in.seekg(offset_to_audio, std::ios::beg);
+        std::ofstream ofs_decrypted(out_path, std::ios::binary);
+        if (!ofs_decrypted.is_open()) {
+            error(L"failed to open output file");
+            return;
+        }
+
+        size_t offset{0};
+        thread_local std::array<uint8_t, 1024 * 1024> temp_buffer{};
+        auto buf_signed = std::span(reinterpret_cast<char*>(temp_buffer.data()), temp_buffer.size());
+        auto buf_unsigned = std::span(temp_buffer);
+
+        while (!kgg_stream_in.eof()) {
+            kgg_stream_in.read(buf_signed.data(), buf_signed.size());
+            const auto n = static_cast<size_t>(kgg_stream_in.gcount());
+            qmc2->Decrypt(buf_unsigned.subspan(0, n), offset);
+            ofs_decrypted.write(buf_signed.data(), n);
+            offset += n;
+        }
+
+        info(std::format(L"** OK **  -> {}", out_path.filename().wstring()));
+    }
+
+   private:
+    std::filesystem::path kgg_path_;
+    std::filesystem::path out_dir_;
+
+    static const wchar_t* DetectRealExt(std::string_view magic) {
+        if (magic == "fLaC") {
+            return L"flac";
+        }
+        if (magic == "OggS") {
+            return L"ogg";
+        }
+        return L"mp3";
+    }
+};
+
+class KggTaskQueue {
+   public:
+    explicit KggTaskQueue(Infra::kgm_ekey_db_t ekey_db) : ekey_db_(std::move(ekey_db)) {}
+
+    void Push(std::unique_ptr<KggTask> task) {
+        std::lock_guard lock(mutex_);
+        tasks_.push(std::move(task));
+        signal_.notify_one();
+    }
+
+    std::unique_ptr<KggTask> Pop() {
+        std::unique_lock lock(mutex_);
+        signal_.wait(lock, [this] { return !tasks_.empty() || thread_end_; });
+        if (tasks_.empty()) {
+            return {};
+        }
+
+        auto task = std::move(tasks_.front());
+        tasks_.pop();
+        return task;
+    }
+
+    [[nodiscard]] bool Finished() {
+        std::lock_guard lock(mutex_);
+        return tasks_.empty();
+    }
+
+    void AddWorkerThread() { threads_.emplace(&KggTaskQueue::WorkerThreadBody, this); }
+
+    void Join() {
+        thread_end_ = true;
+        signal_.notify_all();
+
+        for (int i = 1; !threads_.empty(); i++) {
+            threads_.front().join();
+            threads_.pop();
+#ifndef NDEBUG
+            fprintf(stderr, "[INFO] thread %d joined\n", i);
+#endif
+        }
+    }
+
+   private:
+    bool thread_end_{false};
+    Infra::kgm_ekey_db_t ekey_db_;
+    void WorkerThreadBody() {
+        SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_BELOW_NORMAL);
+
+        std::unique_ptr<KggTask> task{nullptr};
+        while ((task = Pop())) {
+            task->Execute(ekey_db_);
+        }
+    }
+
+    std::mutex mutex_{};
+    std::condition_variable signal_;
+    std::queue<std::unique_ptr<KggTask>> tasks_{};
+    std::queue<std::thread> threads_{};
+};

src/main.cpp

@@ -1,3 +1,8 @@
+#include "infra/infra.h"
+#include "infra/sqlite_error.h"
+#include "jobs.hpp"
+#include "qmc2/qmc2.h"
+
 // clang-format off
 #include <windows.h>
 #include <shlobj.h>
@@ -5,102 +10,41 @@
 #include <knownfolders.h>
 // clang-format on
 
+#include <algorithm>
 #include <cstdint>
 #include <cstdio>
-#include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <span>
-
+#include <utility>
-#include "infra/infra.h"
-#include "infra/sqlite_error.h"
-#include "qmc2/qmc2.h"
 
 using Infra::kgm_ekey_db_t;
 
-constexpr std::array<uint8_t, 16> kMagicHeader{0x7C, 0xD5, 0x32, 0xEB, 0x86, 0x02, 0x7F, 0x4B,
+void WalkFileOrDir(KggTaskQueue& queue, const std::filesystem::path& input_path, bool scan_all_exts) {
-                                               0xA8, 0xAF, 0xA6, 0x8E, 0x0F, 0xFF, 0x99, 0x14};
+    std::queue<std::filesystem::path> file_queue;
+    file_queue.push(absolute(input_path));
 
-void DecryptKGG(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& kgg_file) {
+    while (!file_queue.empty()) {
-    std::ifstream ifs_kgg(kgg_file, std::ios::binary);
+        auto target_path = std::move(file_queue.front());
-    char header[0x100]{};
+        file_queue.pop();
-    ifs_kgg.read(header, sizeof(ifs_kgg));
-    if (std::equal(kMagicHeader.cbegin(), kMagicHeader.cend(), header)) {
-        fprintf(stderr, "[WARN] ! invalid kgg header, skip.\n");
-        return;
-    }
-    uint32_t offset_to_audio = *reinterpret_cast<uint32_t*>(&header[0x10]);
-    uint32_t encrypt_mode = *reinterpret_cast<uint32_t*>(&header[0x14]);
-    if (encrypt_mode != 5) {
-        fprintf(stderr, "[WARN] ! invalid enc_version (expect=0x05, got 0x%02x), skip.\n", encrypt_mode);
-        return;
-    }
-    uint32_t audio_hash_len = *reinterpret_cast<uint32_t*>(&header[0x44]);
-    if (audio_hash_len != 0x20) {
-        fprintf(stderr, "[WARN] ! audio hash length invalid (expect=0x20, got 0x%02x), skip.\n", audio_hash_len);
-        return;
-    }
-    std::string audio_hash(&header[0x48], &header[0x48 + audio_hash_len]);
-    std::string ekey{};
-    if (auto it = ekey_db.find(audio_hash); it != ekey_db.end()) {
-        ekey = it->second;
-    } else {
-        fprintf(stderr, "[WARN] ! ekey not found, skip.\n");
-        return;
-    }
 
-    auto qmc2 = QMC2::Create(ekey);
+        if (is_regular_file(target_path)) {
-    if (!qmc2) {
+            if (!scan_all_exts && target_path.extension() != L".kgg") {
-        fprintf(stderr, "[WARN] ! create qmc2 failed, skip.\n");
+                continue;
-        return;
+            }
-    }
 
-    ifs_kgg.seekg(offset_to_audio, std::ios::beg);
+            queue.Push(std::make_unique<KggTask>(target_path, target_path.parent_path()));
-    std::vector<uint8_t> buffer(1024 * 1024);
+            continue;
-    ifs_kgg.read(reinterpret_cast<char*>(buffer.data()), 4);
-    auto n = static_cast<size_t>(ifs_kgg.gcount());
-    qmc2->Decrypt(std::span(buffer.begin(), n), 0);
-
-    auto decrypted_path = kgg_file;
-    auto magic = std::span(buffer.cbegin(), 4);
-    if (std::equal(magic.begin(), magic.end(), "fLaC")) {
-        decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.flac");
-    } else if (std::equal(magic.begin(), magic.end(), "OggS")) {
-        decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.ogg");
-    } else {
-        decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.mp3");
-    }
-    if (exists(decrypted_path)) {
-        fprintf(stderr, "[WARN] ! output file '%s' exists, skip.\n", decrypted_path.filename().string().c_str());
-        return;
-    }
-
-    std::ofstream ofs_decrypted(decrypted_path, std::ios::binary);
-    ofs_decrypted.write(reinterpret_cast<char*>(buffer.data()), n);
-    size_t offset{n};
-    while (!ifs_kgg.eof()) {
-        ifs_kgg.read(reinterpret_cast<char*>(buffer.data()), buffer.size());
-        n = static_cast<size_t>(ifs_kgg.gcount());
-        qmc2->Decrypt(std::span(buffer.begin(), n), offset);
-        ofs_decrypted.write(reinterpret_cast<char*>(buffer.data()), n);
-        offset += n;
-    }
-    qmc2.reset();
-
-    auto kgg_fname = kgg_file.filename();
-    auto decrypted_fname = decrypted_path.filename();
-    fwprintf(stderr, L"[INFO] * OK! %s --> %s\n", kgg_fname.c_str(), decrypted_fname.c_str());
-}
-
-void WalkFileOrDir(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& input_path) {
-    if (is_directory(input_path)) {
-        for (auto const& dir_entry : std::filesystem::directory_iterator{input_path}) {
-            DecryptKGG(ekey_db, dir_entry.path());
         }
-    } else if (is_regular_file(input_path)) {
+
-        DecryptKGG(ekey_db, input_path);
+        if (is_directory(target_path)) {
-    } else {
+            for (auto const& dir_entry : std::filesystem::directory_iterator{target_path}) {
-        fwprintf(stderr, L"[WARN] invalid path: %s\n", input_path.c_str());
+                file_queue.push(dir_entry.path());
+            }
+            continue;
+        }
+
+        fwprintf(stderr, L"[WARN] invalid path: %s\n", target_path.c_str());
     }
 }
 
@@ -144,6 +88,7 @@ int main() {
     fputs(
         "Usage: kgg-dec "
         "[--infra-dll infra.dll] "
+        "[--scan-all-file-ext 0] "
         "[--db /path/to/KGMusicV3.db] "
         "[--] [kgg-dir... = '.']\n\n",
         stderr);
@@ -164,6 +109,11 @@ int main() {
         return 1;
     }
 
+    bool scan_all_exts{false};
+    if (auto it = named_args.find(L"scan-all-file-ext"); it != named_args.end()) {
+        scan_all_exts = it->second == L"1";
+    }
+
     if (auto it = named_args.find(L"db"); it != named_args.end()) {
         kgm_db_path = std::filesystem::path{it->second};
     } else {
@@ -190,16 +140,29 @@ int main() {
     }
     db.Close();
 
-#ifdef _DEBUG
+#ifndef NDEBUG
-    printf("ekey_db:\n");
+    fprintf(stderr, "ekey_db:\n");
     for (auto& [a, b] : ekey_db) {
-        printf("%s --> %s\n", a.c_str(), b.c_str());
+        fprintf(stderr, "%s --> %s\n", a.c_str(), b.c_str());
     }
 #endif
 
-    for (auto& positional_arg : positional_args) {
+    KggTaskQueue queue(ekey_db);
-        WalkFileOrDir(ekey_db, positional_arg);
+    auto thread_count =
+#ifndef NDEBUG
+        1;
+#else
+        std::max(static_cast<int>(std::thread::hardware_concurrency()) - 2, 2);
+#endif
+
+    for (int i = 0; i < thread_count; i++) {
+        queue.AddWorkerThread();
     }
 
+    for (auto& positional_arg : positional_args) {
+        WalkFileOrDir(queue, positional_arg, scan_all_exts);
+    }
+    queue.Join();
+
     return 0;
 }

src/qmc2/qmc2.h

@@ -11,10 +11,10 @@ namespace QMC2 {
 
 class QMC2_Base {
    public:
-    QMC2_Base(){};
+    QMC2_Base() = default;
     virtual ~QMC2_Base() = default;
 
-    virtual void Decrypt(std::span<uint8_t> data, size_t offset) = 0;
+    virtual void Decrypt(std::span<uint8_t> data, size_t offset) const = 0;
 };
 
 constexpr size_t kMapIndexOffset = 71214;
@@ -24,7 +24,7 @@ class QMC2_MAP : public QMC2_Base {
    public:
     explicit QMC2_MAP(std::span<uint8_t> key);
 
-    void Decrypt(std::span<uint8_t> data, size_t offset) override;
+    void Decrypt(std::span<uint8_t> data, size_t offset) const override;
 
    private:
     std::array<uint8_t, kMapKeySize> key_{};
@@ -38,15 +38,15 @@ class QMC2_RC4 : public QMC2_Base {
    public:
     explicit QMC2_RC4(std::span<uint8_t> key);
 
-    void Decrypt(std::span<uint8_t> data, size_t offset) override;
+    void Decrypt(std::span<uint8_t> data, size_t offset) const override;
 
    private:
     std::vector<uint8_t> key_{};
     double hash_{0};
     std::array<uint8_t, kRC4StreamSize> key_stream_{};
 
-    size_t DecryptFirstSegment(std::span<uint8_t> data, size_t offset);
+    [[nodiscard]] size_t DecryptFirstSegment(std::span<uint8_t> data, size_t offset) const;
-    size_t DecryptOtherSegment(std::span<uint8_t> data, size_t offset);
+    [[nodiscard]] size_t DecryptOtherSegment(std::span<uint8_t> data, size_t offset) const;
 };
 
 std::unique_ptr<QMC2_Base> Create(std::string_view ekey);

src/qmc2/qmc2_map.cpp

@@ -13,7 +13,7 @@ QMC2_MAP::QMC2_MAP(std::span<uint8_t> key) {
     }
 }
 
-void QMC2_MAP::Decrypt(std::span<uint8_t> data, size_t offset) {
+void QMC2_MAP::Decrypt(std::span<uint8_t> data, size_t offset) const {
     for (auto& it : data) {
         size_t idx = (offset <= kMapOffsetBoundary) ? offset : (offset % kMapOffsetBoundary);
         it ^= key_[idx % key_.size()];

src/qmc2/qmc2_rc4.cpp

@@ -14,7 +14,7 @@ inline double hash(const uint8_t* key, size_t len) {
         }
 
         // Overflow check.
-        uint32_t next_hash = hash * (uint32_t)(*key);
+        uint32_t next_hash = hash * static_cast<uint32_t>(*key);
         if (next_hash <= hash) {
             break;
         }
@@ -47,7 +47,7 @@ class RC4 {
 
         for (auto& it : buffer) {
             i = (i + 1) % n;
-            j = (j + (size_t)s[i]) % n;
+            j = (j + s[i]) % n;
             std::swap(s[i], s[j]);
 
             const size_t final_idx = (s[i] + s[j]) % n;
@@ -70,8 +70,8 @@ inline size_t get_segment_key(double key_hash, size_t segment_id, uint8_t seed)
         return 0;
     }
 
-    double result = key_hash / (double)(seed * (segment_id + 1)) * 100.0;
+    const double result = key_hash / static_cast<double>(seed * (segment_id + 1)) * 100.0;
-    return (size_t)result;
+    return static_cast<size_t>(result);
 }
 
 QMC2_RC4::QMC2_RC4(std::span<uint8_t> key) {
@@ -82,7 +82,7 @@ QMC2_RC4::QMC2_RC4(std::span<uint8_t> key) {
     rc4.Derive(std::span<uint8_t>(key_stream_));
 }
 
-void QMC2_RC4::Decrypt(std::span<uint8_t> data, size_t offset) {
+void QMC2_RC4::Decrypt(std::span<uint8_t> data, size_t offset) const {
     if (offset < kFirstSegmentSize) {
         const auto n = DecryptFirstSegment(data, offset);
         offset += n;
@@ -96,7 +96,7 @@ void QMC2_RC4::Decrypt(std::span<uint8_t> data, size_t offset) {
     }
 }
 
-size_t QMC2_RC4::DecryptFirstSegment(std::span<uint8_t> data, size_t offset) {
+size_t QMC2_RC4::DecryptFirstSegment(std::span<uint8_t> data, size_t offset) const {
     const uint8_t* key = key_.data();
     const size_t n = this->key_.size();
 
@@ -109,7 +109,7 @@ size_t QMC2_RC4::DecryptFirstSegment(std::span<uint8_t> data, size_t offset) {
     return process_len;
 }
 
-size_t QMC2_RC4::DecryptOtherSegment(std::span<uint8_t> data, size_t offset) {
+size_t QMC2_RC4::DecryptOtherSegment(std::span<uint8_t> data, size_t offset) const {
     const size_t n = this->key_.size();
 
     size_t segment_idx = offset / kOtherSegmentSize;