Add key id based m2p rotation support

src/PackLib/M2PackRuntimeKeyProvider.cpp

@@ -17,6 +17,7 @@ constexpr const char* M2PACK_DEFAULT_MAP_NAME = "Local\\M2PackSharedKeys";
 constexpr const char* M2PACK_ENV_MASTER_KEY = "M2PACK_MASTER_KEY_HEX";
 constexpr const char* M2PACK_ENV_PUBLIC_KEY = "M2PACK_SIGN_PUBKEY_HEX";
 constexpr const char* M2PACK_ENV_MAP_NAME = "M2PACK_KEY_MAP";
+constexpr const char* M2PACK_ENV_KEY_ID = "M2PACK_KEY_ID";
 
 #pragma pack(push, 1)
 struct M2PackSharedKeys
@@ -24,6 +25,7 @@ struct M2PackSharedKeys
 	char magic[8];
 	uint32_t version;
 	uint32_t flags;
+	uint32_t key_id;
 	uint8_t master_key[M2PACK_KEY_SIZE];
 	uint8_t sign_public_key[M2PACK_PUBLIC_KEY_SIZE];
 };
@@ -31,8 +33,9 @@ struct M2PackSharedKeys
 
 struct RuntimeKeyState
 {
+	uint32_t master_key_id = 0;
 	std::array<uint8_t, M2PACK_KEY_SIZE> master_key {};
-	std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE> public_key = M2PACK_SIGN_PUBLIC_KEY;
+	std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE> public_key {};
 	bool runtime_master_key = false;
 	bool runtime_public_key = false;
 	bool initialized = false;
@@ -82,6 +85,32 @@ std::string GetEnvString(const char* name)
 	return std::string(buffer, buffer + len);
 }
 
+uint32_t ParseUInt32(const std::string& value)
+{
+	if (value.empty())
+		return 0;
+
+	try
+	{
+		return static_cast<uint32_t>(std::stoul(value));
+	}
+	catch (...)
+	{
+		return 0;
+	}
+}
+
+const std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE>* FindCompiledPublicKey(uint32_t keyId)
+{
+	for (std::size_t i = 0; i < M2PACK_SIGN_KEY_IDS.size(); ++i)
+	{
+		if (M2PACK_SIGN_KEY_IDS[i] == keyId)
+			return &M2PACK_SIGN_PUBLIC_KEYS[i];
+	}
+
+	return nullptr;
+}
+
 bool LoadFromSharedMapping(const std::string& mappingName)
 {
 	const HANDLE mapping = OpenFileMappingA(FILE_MAP_READ, FALSE, mappingName.c_str());
@@ -108,6 +137,7 @@ bool LoadFromSharedMapping(const std::string& mappingName)
 	}
 
 	memcpy(g_state.master_key.data(), blob.master_key, g_state.master_key.size());
+	g_state.master_key_id = blob.key_id;
 	memcpy(g_state.public_key.data(), blob.sign_public_key, g_state.public_key.size());
 	g_state.runtime_master_key = true;
 	g_state.runtime_public_key = true;
@@ -129,6 +159,20 @@ void ApplyCommandLineOption(const std::string& key, const std::string& value)
 		return;
 	}
 
+	if (key == "--m2pack-key-id")
+	{
+		const auto parsed = ParseUInt32(value);
+		if (parsed != 0)
+		{
+			g_state.master_key_id = parsed;
+		}
+		else
+		{
+			TraceError("Invalid value for --m2pack-key-id");
+		}
+		return;
+	}
+
 	if (key == "--m2pack-pubkey-hex")
 	{
 		if (ParseHexInto(value, g_state.public_key))
@@ -177,6 +221,16 @@ bool InitializeM2PackRuntimeKeyProvider(const char* commandLine)
 			TraceError("Invalid M2PACK_MASTER_KEY_HEX value");
 	}
 
+	const std::string envKeyId = GetEnvString(M2PACK_ENV_KEY_ID);
+	if (!envKeyId.empty())
+	{
+		const auto parsed = ParseUInt32(envKeyId);
+		if (parsed != 0)
+			g_state.master_key_id = parsed;
+		else
+			TraceError("Invalid M2PACK_KEY_ID value");
+	}
+
 	const std::string envPublic = GetEnvString(M2PACK_ENV_PUBLIC_KEY);
 	if (!envPublic.empty())
 	{
@@ -193,7 +247,7 @@ bool InitializeM2PackRuntimeKeyProvider(const char* commandLine)
 		for (int i = 0; i < argc; ++i)
 		{
 			const std::string key = argv[i];
-			if ((key == "--m2pack-key-hex" || key == "--m2pack-pubkey-hex" || key == "--m2pack-key-map") && i + 1 < argc)
+			if ((key == "--m2pack-key-hex" || key == "--m2pack-pubkey-hex" || key == "--m2pack-key-map" || key == "--m2pack-key-id") && i + 1 < argc)
 			{
 				ApplyCommandLineOption(key, argv[i + 1]);
 				++i;
@@ -223,9 +277,19 @@ const std::array<uint8_t, M2PACK_KEY_SIZE>& GetM2PackActiveMasterKey()
 	return g_state.master_key;
 }
 
-const std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE>& GetM2PackActivePublicKey()
+const std::array<uint8_t, M2PACK_PUBLIC_KEY_SIZE>* GetM2PackPublicKeyForKeyId(uint32_t keyId)
 {
-	return g_state.public_key;
+	if (g_state.runtime_public_key && g_state.master_key_id == keyId)
+	{
+		return &g_state.public_key;
+	}
+
+	return FindCompiledPublicKey(keyId);
+}
+
+uint32_t GetM2PackActiveMasterKeyId()
+{
+	return g_state.master_key_id;
 }
 
 bool HasM2PackRuntimeMasterKey()
@@ -233,9 +297,9 @@ bool HasM2PackRuntimeMasterKey()
 	return g_state.runtime_master_key;
 }
 
-bool HasM2PackRuntimeKeysForArchiveLoad()
+bool HasM2PackRuntimeKeysForArchiveLoad(uint32_t keyId)
 {
-	return g_state.runtime_master_key;
+	return g_state.runtime_master_key && g_state.master_key_id == keyId;
 }
 
 bool IsM2PackUsingRuntimeMasterKey()