Allow key loading from %YUZU_DIR%/keys in addition to ~/.switch

5 files changed, 23 insertions, 7 deletions

diff --git a/src/common/common_paths.h b/src/common/common_paths.h
index 6799a357a..df2ce80b1 100644
--- a/src/common/common_paths.h
+++ b/src/common/common_paths.h
@@ -32,6 +32,7 @@
 #define SDMC_DIR "sdmc"
 #define NAND_DIR "nand"
 #define SYSDATA_DIR "sysdata"
+#define KEYS_DIR "keys"
 #define LOG_DIR "log"
 
 // Filenames
diff --git a/src/common/file_util.cpp b/src/common/file_util.cpp
index ed38c3ced..7aeda737f 100644
--- a/src/common/file_util.cpp
+++ b/src/common/file_util.cpp
@@ -706,6 +706,7 @@ const std::string& GetUserPath(UserPath path, const std::string& new_path) {
         paths.emplace(UserPath::SDMCDir, user_path + SDMC_DIR DIR_SEP);
         paths.emplace(UserPath::NANDDir, user_path + NAND_DIR DIR_SEP);
         paths.emplace(UserPath::SysDataDir, user_path + SYSDATA_DIR DIR_SEP);
+        paths.emplace(UserPath::KeysDir, user_path + KEYS_DIR DIR_SEP);
         // TODO: Put the logs in a better location for each OS
         paths.emplace(UserPath::LogDir, user_path + LOG_DIR DIR_SEP);
     }
diff --git a/src/common/file_util.h b/src/common/file_util.h
index d530d86c9..28697d527 100644
--- a/src/common/file_util.h
+++ b/src/common/file_util.h
@@ -23,6 +23,7 @@ namespace FileUtil {
 enum class UserPath {
     CacheDir,
     ConfigDir,
+    KeysDir,
     LogDir,
     NANDDir,
     RootDir,
diff --git a/src/core/crypto/key_manager.cpp b/src/core/crypto/key_manager.cpp
index ea20bc31b..dea092b5e 100644
--- a/src/core/crypto/key_manager.cpp
+++ b/src/core/crypto/key_manager.cpp
@@ -12,6 +12,7 @@
 #include "common/logging/log.h"
 #include "core/crypto/key_manager.h"
 #include "core/settings.h"
+#include "key_manager.h"
 
 namespace Core::Crypto {
 
@@ -50,18 +51,17 @@ std::array<u8, 32> operator""_array32(const char* str, size_t len) {
 
 KeyManager::KeyManager() {
     // Initialize keys
-    std::string keys_dir = FileUtil::GetHactoolConfigurationPath();
+    std::string hactool_keys_dir = FileUtil::GetHactoolConfigurationPath();
+    std::string yuzu_keys_dir = FileUtil::GetUserPath(FileUtil::UserPath::KeysDir);
     if (Settings::values.use_dev_keys) {
         dev_mode = true;
-        if (FileUtil::Exists(keys_dir + DIR_SEP + "dev.keys"))
-            LoadFromFile(keys_dir + DIR_SEP + "dev.keys", false);
+        AttemptLoadKeyFile(yuzu_keys_dir, hactool_keys_dir, "dev.keys", false);
     } else {
         dev_mode = false;
-        if (FileUtil::Exists(keys_dir + DIR_SEP + "prod.keys"))
-            LoadFromFile(keys_dir + DIR_SEP + "prod.keys", false);
+        AttemptLoadKeyFile(yuzu_keys_dir, hactool_keys_dir, "prod.keys", false);
     }
-    if (FileUtil::Exists(keys_dir + DIR_SEP + "title.keys"))
-        LoadFromFile(keys_dir + DIR_SEP + "title.keys", true);
+
+    AttemptLoadKeyFile(yuzu_keys_dir, hactool_keys_dir, "title.keys", true);
 }
 
 void KeyManager::LoadFromFile(std::string_view filename_, bool is_title_keys) {
@@ -104,6 +104,17 @@ void KeyManager::LoadFromFile(std::string_view filename_, bool is_title_keys) {
     }
 }
 
+void KeyManager::AttemptLoadKeyFile(std::string_view dir1_, std::string_view dir2_,
+                                    std::string_view filename_, bool title) {
+    std::string dir1(dir1_);
+    std::string dir2(dir2_);
+    std::string filename(filename_);
+    if (FileUtil::Exists(dir1 + DIR_SEP + filename))
+        LoadFromFile(dir1 + DIR_SEP + filename, title);
+    else if (FileUtil::Exists(dir2 + DIR_SEP + filename))
+        LoadFromFile(dir2 + DIR_SEP + filename, title);
+}
+
 bool KeyManager::HasKey(S128KeyType id, u64 field1, u64 field2) const {
     return s128_keys.find({id, field1, field2}) != s128_keys.end();
 }
diff --git a/src/core/crypto/key_manager.h b/src/core/crypto/key_manager.h
index e04f1d49f..a52ea4cb9 100644
--- a/src/core/crypto/key_manager.h
+++ b/src/core/crypto/key_manager.h
@@ -107,6 +107,8 @@ private:
 
     bool dev_mode;
     void LoadFromFile(std::string_view filename, bool is_title_keys);
+    void AttemptLoadKeyFile(std::string_view dir1, std::string_view dir2, std::string_view filename,
+                            bool title);
 
     static std::unordered_map<std::string, KeyIndex<S128KeyType>> s128_file_id;
     static std::unordered_map<std::string, KeyIndex<S256KeyType>> s256_file_id;