1 files changed, 35 insertions, 33 deletions

diff --git a/src/core/crypto/key_manager.h b/src/core/crypto/key_manager.h
index 0c62d4421..7ca3e6cbc 100644
--- a/src/core/crypto/key_manager.h
+++ b/src/core/crypto/key_manager.h
@@ -6,11 +6,13 @@
 
 #include <array>
 #include <string>
+#include <string_view>
 #include <type_traits>
-#include <unordered_map>
 #include <vector>
+#include <boost/container/flat_map.hpp>
 #include <fmt/format.h>
 #include "common/common_types.h"
+#include "core/loader/loader.h"
 
 namespace Core::Crypto {
 
@@ -22,9 +24,8 @@ static_assert(sizeof(Key128) == 16, "Key128 must be 128 bytes big.");
 static_assert(sizeof(Key256) == 32, "Key128 must be 128 bytes big.");
 
 enum class S256KeyType : u64 {
-    Header, //
-    SDSave, //
-    SDNCA,  //
+    Header,      //
+    SDKeySource, // f1=SDKeyType
 };
 
 enum class S128KeyType : u64 {
@@ -36,6 +37,7 @@ enum class S128KeyType : u64 {
     KeyArea,       // f1=crypto revision f2=type {app, ocean, system}
     SDSeed,        //
     Titlekey,      // f1=rights id LSB f2=rights id MSB
+    Source,        // f1=source type, f2= sub id
 };
 
 enum class KeyAreaKeyType : u8 {
@@ -44,6 +46,17 @@ enum class KeyAreaKeyType : u8 {
     System,
 };
 
+enum class SourceKeyType : u8 {
+    SDKEK,
+    AESKEKGeneration,
+    AESKeyGeneration,
+};
+
+enum class SDKeyType : u8 {
+    Save,
+    NCA,
+};
+
 template <typename KeyType>
 struct KeyIndex {
     KeyType type;
@@ -59,34 +72,12 @@ struct KeyIndex {
     }
 };
 
-// The following two (== and hash) are so KeyIndex can be a key in unordered_map
-
+// boost flat_map requires operator< for O(log(n)) lookups.
 template <typename KeyType>
-bool operator==(const KeyIndex<KeyType>& lhs, const KeyIndex<KeyType>& rhs) {
-    return std::tie(lhs.type, lhs.field1, lhs.field2) == std::tie(rhs.type, rhs.field1, rhs.field2);
+bool operator<(const KeyIndex<KeyType>& lhs, const KeyIndex<KeyType>& rhs) {
+    return std::tie(lhs.type, lhs.field1, lhs.field2) < std::tie(rhs.type, rhs.field1, rhs.field2);
 }
 
-template <typename KeyType>
-bool operator!=(const KeyIndex<KeyType>& lhs, const KeyIndex<KeyType>& rhs) {
-    return !operator==(lhs, rhs);
-}
-
-} // namespace Core::Crypto
-
-namespace std {
-template <typename KeyType>
-struct hash<Core::Crypto::KeyIndex<KeyType>> {
-    size_t operator()(const Core::Crypto::KeyIndex<KeyType>& k) const {
-        using std::hash;
-
-        return ((hash<u64>()(static_cast<u64>(k.type)) ^ (hash<u64>()(k.field1) << 1)) >> 1) ^
-               (hash<u64>()(k.field2) << 1);
-    }
-};
-} // namespace std
-
-namespace Core::Crypto {
-
 class KeyManager {
 public:
     KeyManager();
@@ -102,16 +93,27 @@ public:
 
     static bool KeyFileExists(bool title);
 
+    // Call before using the sd seed to attempt to derive it if it dosen't exist. Needs system save
+    // 8*43 and the private file to exist.
+    void DeriveSDSeedLazy();
+
 private:
-    std::unordered_map<KeyIndex<S128KeyType>, Key128> s128_keys;
-    std::unordered_map<KeyIndex<S256KeyType>, Key256> s256_keys;
+    boost::container::flat_map<KeyIndex<S128KeyType>, Key128> s128_keys;
+    boost::container::flat_map<KeyIndex<S256KeyType>, Key256> s256_keys;
 
     bool dev_mode;
     void LoadFromFile(const std::string& filename, bool is_title_keys);
     void AttemptLoadKeyFile(const std::string& dir1, const std::string& dir2,
                             const std::string& filename, bool title);
+    template <size_t Size>
+    void WriteKeyToFile(bool title_key, std::string_view keyname, const std::array<u8, Size>& key);
 
-    static const std::unordered_map<std::string, KeyIndex<S128KeyType>> s128_file_id;
-    static const std::unordered_map<std::string, KeyIndex<S256KeyType>> s256_file_id;
+    static const boost::container::flat_map<std::string, KeyIndex<S128KeyType>> s128_file_id;
+    static const boost::container::flat_map<std::string, KeyIndex<S256KeyType>> s256_file_id;
 };
+
+Key128 GenerateKeyEncryptionKey(Key128 source, Key128 master, Key128 kek_seed, Key128 key_seed);
+boost::optional<Key128> DeriveSDSeed();
+Loader::ResultStatus DeriveSDKeys(std::array<Key256, 2>& sd_keys, const KeyManager& keys);
+
 } // namespace Core::Crypto