From a9632c8d0016b01fef0ac64b37ff1cd4d118012e Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Sat, 10 Oct 2020 17:19:40 -0700 Subject: [PATCH] Rewrite the key file parser --- build/CodeGen/IncludedKeys.txt | 4 + build/CodeGen/Stage2/KeysCodeGen.cs | 6 +- src/LibHac/Common/Keys/ExternalKeyReader.cs | 380 +++++++++++++++----- src/LibHac/Common/Keys/KeyInfo.cs | 16 +- src/LibHac/Common/Keys/KeySet.cs | 6 + src/LibHac/Common/StringUtils.cs | 6 + src/LibHac/FsSrv/ExternalKeySet.cs | 5 + 7 files changed, 328 insertions(+), 95 deletions(-) diff --git a/build/CodeGen/IncludedKeys.txt b/build/CodeGen/IncludedKeys.txt index 3441482d..94ee69b7 100644 --- a/build/CodeGen/IncludedKeys.txt +++ b/build/CodeGen/IncludedKeys.txt @@ -18,6 +18,10 @@ mariko_master_kek_source_08 = 5C24E3B8B4F700C23CFD0ACE13C3DC23 mariko_master_kek_source_09 = 8669F00987C805AEB57B4874DE62A613 mariko_master_kek_source_0a = 0E440CEDB436C03FAA1DAEBF62B10982 +mariko_master_kek_source_dev_06 = CC974C462A0CB0A6C9C0B7BE302EC368 +mariko_master_kek_source_dev_07 = 86BD1D7650DF6DFA2C7D3322ABF18218 +mariko_master_kek_source_dev_08 = A3B1E0A958A2267F40BF5BBB87330B66 +mariko_master_kek_source_dev_09 = 82729165403B9D6660D01B3D4DA570E1 mariko_master_kek_source_dev_0a = F937CF9ABD86BBA99C9E03C4FCBC3BCE master_key_source = D8A2410AC6C59001C61D6A267C513F3C diff --git a/build/CodeGen/Stage2/KeysCodeGen.cs b/build/CodeGen/Stage2/KeysCodeGen.cs index f969b4a3..89bf7489 100644 --- a/build/CodeGen/Stage2/KeysCodeGen.cs +++ b/build/CodeGen/Stage2/KeysCodeGen.cs @@ -91,10 +91,10 @@ namespace LibHacBuild.CodeGen.Stage2 var keySet = new KeySet(); // Populate the key set with all the keys in IncludedKeys.txt - using (var reader = new StreamReader(Common.GetResource(InputMainKeyFileName))) + using (Stream keyFile = Common.GetResource(InputMainKeyFileName)) { - List list = ExternalKeyReader.CreateKeyList(); - ExternalKeyReader.ReadMainKeys(keySet, reader, list); + List list = KeySet.CreateKeyInfoList(); + ExternalKeyReader.ReadMainKeys(keySet, keyFile, list); } // Recover all the RSA key parameters and write the key to the key set diff --git a/src/LibHac/Common/Keys/ExternalKeyReader.cs b/src/LibHac/Common/Keys/ExternalKeyReader.cs index 10fe6a1d..dff25718 100644 --- a/src/LibHac/Common/Keys/ExternalKeyReader.cs +++ b/src/LibHac/Common/Keys/ExternalKeyReader.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.Diagnostics; using System.IO; +using System.Runtime.CompilerServices; using LibHac.Fs; using LibHac.Spl; @@ -45,20 +46,20 @@ namespace LibHac.Common.Keys if (filename != null) { - using var reader = new StreamReader(new FileStream(filename, FileMode.Open, FileAccess.Read)); - ReadMainKeys(keySet, reader, keyInfos, logger); + using var storage = new FileStream(filename, FileMode.Open, FileAccess.Read); + ReadMainKeys(keySet, storage, keyInfos, logger); } if (consoleKeysFilename != null) { - using var reader = new StreamReader(new FileStream(consoleKeysFilename, FileMode.Open, FileAccess.Read)); - ReadMainKeys(keySet, reader, keyInfos, logger); + using var storage = new FileStream(consoleKeysFilename, FileMode.Open, FileAccess.Read); + ReadMainKeys(keySet, storage, keyInfos, logger); } if (titleKeysFilename != null) { - using var reader = new StreamReader(new FileStream(titleKeysFilename, FileMode.Open, FileAccess.Read)); - ReadTitleKeys(keySet, reader, logger); + using var storage = new FileStream(titleKeysFilename, FileMode.Open, FileAccess.Read); + ReadTitleKeys(keySet, storage, logger); } keySet.DeriveKeys(logger); @@ -98,57 +99,65 @@ namespace LibHac.Common.Keys /// Missing keys will not be derived. /// /// The where the loaded keys will be placed. - /// A containing the keys to load. + /// A containing the keys to load. /// A list of all the keys that will be loaded into the key set. /// will create a list containing all loadable keys. /// An optional logger that key-parsing errors will be written to. - public static void ReadMainKeys(KeySet keySet, TextReader keyFileReader, List keyList, - IProgressReport logger = null) + public static void ReadMainKeys(KeySet keySet, Stream reader, List keyList, + IProgressReport logger = null) { - if (keyFileReader == null) return; + if (reader == null) return; - // Todo: Improve key parsing - string line; - while ((line = keyFileReader.ReadLine()) != null) + using var streamReader = new StreamReader(reader); + Span buffer = stackalloc char[1024]; + var ctx = new KvPairReaderContext(streamReader, buffer); + + while (true) { - string[] a = line.Split(',', '='); - if (a.Length != 2) continue; + ReaderStatus status = GetKeyValuePair(ref ctx); - string keyName = a[0].Trim(); - string valueStr = a[1].Trim(); - - if (!TryGetKeyInfo(out SpecificKeyInfo info, keyList, keyName)) + if (status == ReaderStatus.Error) { - logger?.LogMessage($"Failed to match key {keyName}"); - continue; + logger?.LogMessage($"Invalid line in key data: \"{ctx.CurrentKey.ToString()}\""); } - - Span key; - - // Get the dev key in the key set if needed. - if (info.IsDev && keySet.CurrentMode == KeySet.Mode.Prod) + else if (status == ReaderStatus.ReadKey) { - keySet.SetMode(KeySet.Mode.Dev); - key = info.Key.Getter(keySet, info.Index); - keySet.SetMode(KeySet.Mode.Prod); + if (!TryGetKeyInfo(out SpecificKeyInfo info, keyList, ctx.CurrentKey)) + { + logger?.LogMessage($"Failed to match key {ctx.CurrentKey.ToString()}"); + continue; + } + + Span key; + + // Get the dev key in the key set if needed. + if (info.IsDev && keySet.CurrentMode == KeySet.Mode.Prod) + { + keySet.SetMode(KeySet.Mode.Dev); + key = info.Key.Getter(keySet, info.Index); + keySet.SetMode(KeySet.Mode.Prod); + } + else + { + key = info.Key.Getter(keySet, info.Index); + } + + if (ctx.CurrentValue.Length != key.Length * 2) + { + logger?.LogMessage($"Key {ctx.CurrentKey.ToString()} has incorrect size {ctx.CurrentValue.Length}. Must be {key.Length * 2} hex digits."); + continue; + } + + if (!Utilities.TryToBytes(ctx.CurrentValue, key)) + { + key.Clear(); + + logger?.LogMessage($"Key {ctx.CurrentKey.ToString()} has an invalid value. Must be {key.Length * 2} hex digits."); + } } - else + else if (status == ReaderStatus.Finished) { - key = info.Key.Getter(keySet, info.Index); - } - - if (valueStr.Length != key.Length * 2) - { - logger?.LogMessage( - $"Key {keyName} had incorrect size {valueStr.Length}. Must be {key.Length * 2} hex digits."); - continue; - } - - if (!Utilities.TryToBytes(valueStr, key)) - { - key.Clear(); - - logger?.LogMessage($"Key {keyName} had an invalid value. Must be {key.Length * 2} hex digits."); + break; } } } @@ -157,59 +166,262 @@ namespace LibHac.Common.Keys /// Loads title keys from a into an existing . /// /// The where the loaded keys will be placed. - /// A containing the keys to load. + /// A containing the keys to load. /// An optional logger that key-parsing errors will be written to. - public static void ReadTitleKeys(KeySet keySet, TextReader keyFileReader, IProgressReport logger = null) + public static void ReadTitleKeys(KeySet keySet, Stream reader, IProgressReport logger = null) { - if (keyFileReader == null) return; + if (reader == null) return; - // Todo: Improve key parsing - string line; - while ((line = keyFileReader.ReadLine()) != null) + using var streamReader = new StreamReader(reader); + Span buffer = stackalloc char[1024]; + var ctx = new KvPairReaderContext(streamReader, buffer); + + keySet.ExternalKeySet.EnsureCapacity((int)reader.Length / 67); + + while (true) { - string[] splitLine; + ReaderStatus status = GetKeyValuePair(ref ctx); - // Some people use pipes as delimiters - if (line.Contains('|')) + if (status == ReaderStatus.Error) { - splitLine = line.Split('|'); + logger?.LogMessage($"Invalid line in key data: \"{ctx.CurrentKey.ToString()}\""); + Debugger.Break(); } - else + else if (status == ReaderStatus.ReadKey) { - splitLine = line.Split(',', '='); + if (ctx.CurrentKey.Length != TitleKeySize * 2) + { + logger?.LogMessage($"Rights ID {ctx.CurrentKey.ToString()} has incorrect size {ctx.CurrentKey.Length}. (Expected {TitleKeySize * 2})"); + continue; + } + + if (ctx.CurrentValue.Length != TitleKeySize * 2) + { + logger?.LogMessage($"Title key {ctx.CurrentValue.ToString()} has incorrect size {ctx.CurrentValue.Length}. (Expected {TitleKeySize * 2})"); + continue; + } + + var rightsId = new RightsId(); + var titleKey = new AccessKey(); + + if (!Utilities.TryToBytes(ctx.CurrentKey, SpanHelpers.AsByteSpan(ref rightsId))) + { + logger?.LogMessage($"Invalid rights ID \"{ctx.CurrentKey.ToString()}\" in title key file"); + continue; + } + + if (!Utilities.TryToBytes(ctx.CurrentValue, SpanHelpers.AsByteSpan(ref titleKey))) + { + logger?.LogMessage($"Invalid title key \"{ctx.CurrentValue.ToString()}\" in title key file"); + continue; + } + + keySet.ExternalKeySet.Add(rightsId, titleKey).ThrowIfFailure(); } - - if (splitLine.Length < 2) continue; - - if (!splitLine[0].Trim().TryToBytes(out byte[] rightsId)) + else if (status == ReaderStatus.Finished) { - logger?.LogMessage($"Invalid rights ID \"{splitLine[0].Trim()}\" in title key file"); - continue; + break; } - - if (!splitLine[1].Trim().TryToBytes(out byte[] titleKey)) - { - logger?.LogMessage($"Invalid title key \"{splitLine[1].Trim()}\" in title key file"); - continue; - } - - if (rightsId.Length != TitleKeySize) - { - logger?.LogMessage($"Rights ID {rightsId.ToHexString()} had incorrect size {rightsId.Length}. (Expected {TitleKeySize})"); - continue; - } - - if (titleKey.Length != TitleKeySize) - { - logger?.LogMessage($"Title key {titleKey.ToHexString()} had incorrect size {titleKey.Length}. (Expected {TitleKeySize})"); - continue; - } - - keySet.ExternalKeySet.Add(new RightsId(rightsId), new AccessKey(titleKey)).ThrowIfFailure(); } } - private static bool TryGetKeyInfo(out SpecificKeyInfo info, List keyList, string keyName) + private ref struct KvPairReaderContext + { + public TextReader Reader; + public Span Buffer; + public Span CurrentKey; + public Span CurrentValue; + public int BufferPos; + public bool NeedFillBuffer; + + public KvPairReaderContext(TextReader reader, Span buffer) + { + Reader = reader; + Buffer = buffer; + CurrentKey = default; + CurrentValue = default; + BufferPos = buffer.Length; + NeedFillBuffer = true; + } + } + + private enum ReaderStatus + { + ReadKey, + NoKeyRead, + Finished, + Error + } + + private enum ReaderState + { + Initial, + Key, + WhiteSpace1, + Delimiter, + Value, + WhiteSpace2, + End, + } + + private static ReaderStatus GetKeyValuePair(ref KvPairReaderContext reader) + { + Span buffer = reader.Buffer; + + if (reader.NeedFillBuffer) + { + // Move unread text to the front of the buffer + buffer.Slice(reader.BufferPos).CopyTo(buffer); + + int charsRead = reader.Reader.ReadBlock(buffer.Slice(buffer.Length - reader.BufferPos)); + + if (charsRead == 0) + { + return ReaderStatus.Finished; + } + + buffer = buffer.Slice(0, buffer.Length - reader.BufferPos + charsRead); + + reader.NeedFillBuffer = false; + reader.BufferPos = 0; + } + + // Skip any empty lines + while (reader.BufferPos < buffer.Length && IsEndOfLine(buffer[reader.BufferPos])) + { + reader.BufferPos++; + } + + var state = ReaderState.Initial; + int keyOffset = -1; + int keyLength = -1; + int valueOffset = -1; + int valueLength = -1; + int i; + + for (i = reader.BufferPos; i < buffer.Length; i++) + { + char c = buffer[i]; + + switch (state) + { + case ReaderState.Initial when IsWhiteSpace(c): + continue; + case ReaderState.Initial when IsValidNameChar(c): + state = ReaderState.Key; + ToLower(ref buffer[i]); + keyOffset = i; + continue; + case ReaderState.Key when IsValidNameChar(c): + ToLower(ref buffer[i]); + continue; + case ReaderState.Key when IsWhiteSpace(c): + state = ReaderState.WhiteSpace1; + keyLength = i - keyOffset; + continue; + case ReaderState.Key when IsDelimiter(c): + state = ReaderState.Delimiter; + keyLength = i - keyOffset; + continue; + case ReaderState.WhiteSpace1 when IsWhiteSpace(c): + continue; + case ReaderState.WhiteSpace1 when IsDelimiter(c): + state = ReaderState.Delimiter; + continue; + case ReaderState.Delimiter when IsWhiteSpace(c): + continue; + case ReaderState.Delimiter when StringUtils.IsHexDigit((byte)c): + state = ReaderState.Value; + valueOffset = i; + continue; + case ReaderState.Value when IsEndOfLine(c): + state = ReaderState.End; + valueLength = i - valueOffset; + continue; + case ReaderState.Value when IsWhiteSpace(c): + state = ReaderState.WhiteSpace2; + valueLength = i - valueOffset; + continue; + case ReaderState.Value when StringUtils.IsHexDigit((byte)c): + continue; + case ReaderState.WhiteSpace2 when IsWhiteSpace(c): + continue; + case ReaderState.WhiteSpace2 when IsEndOfLine(c): + state = ReaderState.End; + continue; + case ReaderState.End when IsEndOfLine(c): + continue; + case ReaderState.End when !IsEndOfLine(c): + break; + } + + // We've exited the state machine for one reason or another + break; + } + + // If we successfully read both the key and value + if (state == ReaderState.End || state == ReaderState.WhiteSpace2) + { + reader.CurrentKey = reader.Buffer.Slice(keyOffset, keyLength); + reader.CurrentValue = reader.Buffer.Slice(valueOffset, valueLength); + reader.BufferPos = i; + + return ReaderStatus.ReadKey; + } + + // We either ran out of buffer or hit an error reading the key-value pair. + // Advance to the end of the line if possible. + while (i < buffer.Length && !IsEndOfLine(buffer[i])) + { + i++; + } + + // We don't have a complete line. Return that the buffer needs to be refilled. + if (i == buffer.Length) + { + reader.NeedFillBuffer = true; + return ReaderStatus.NoKeyRead; + } + + // If we hit a line with an error, it'll be returned as "CurrentKey" in the reader context + reader.CurrentKey = buffer.Slice(reader.BufferPos, i - reader.BufferPos); + reader.BufferPos = i; + + return ReaderStatus.Error; + + static bool IsWhiteSpace(char c) + { + return c == ' ' || c == '\t'; + } + + static bool IsDelimiter(char c) + { + return c == '=' || c == ','; + } + + static bool IsEndOfLine(char c) + { + return c == '\0' || c == '\r' || c == '\n'; + } + + static void ToLower(ref char c) + { + // The only characters we need to worry about are underscores and alphanumerics + // Both lowercase and numbers have bit 5 set, so they're both treated the same + if (c != '_') + { + c |= (char)0b100000; + } + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static bool IsValidNameChar(char c) + { + return (c | 0x20u) - 'a' <= 'z' - 'a' || (uint)(c - '0') <= 9 || c == '_'; + } + + private static bool TryGetKeyInfo(out SpecificKeyInfo info, List keyList, ReadOnlySpan keyName) { for (int i = 0; i < keyList.Count; i++) { diff --git a/src/LibHac/Common/Keys/KeyInfo.cs b/src/LibHac/Common/Keys/KeyInfo.cs index 7b984392..09d2287f 100644 --- a/src/LibHac/Common/Keys/KeyInfo.cs +++ b/src/LibHac/Common/Keys/KeyInfo.cs @@ -71,7 +71,7 @@ namespace LibHac.Common.Keys Getter = retrieveFunc; } - public bool Matches(string keyName, out int keyIndex, out bool isDev) + public bool Matches(ReadOnlySpan keyName, out int keyIndex, out bool isDev) { keyIndex = default; isDev = default; @@ -84,7 +84,7 @@ namespace LibHac.Common.Keys }; } - private bool MatchesSingle(string keyName, out bool isDev) + private bool MatchesSingle(ReadOnlySpan keyName, out bool isDev) { Assert.Equal((int)KeyRangeType.Single, (int)RangeType); @@ -93,7 +93,7 @@ namespace LibHac.Common.Keys if (keyName.Length == NameLength + 4) { // Might be a dev key. Check if "_dev" comes after the base key name - if (!keyName.AsSpan(Name.Length, 4).SequenceEqual("_dev")) + if (!keyName.Slice(Name.Length, 4).SequenceEqual("_dev")) return false; isDev = true; @@ -104,13 +104,13 @@ namespace LibHac.Common.Keys } // Check if the base name matches - if (!keyName.AsSpan(0, Name.Length).SequenceEqual(Name)) + if (!keyName.Slice(0, Name.Length).SequenceEqual(Name)) return false; return true; } - private bool MatchesRangedKey(string keyName, ref int keyIndex, out bool isDev) + private bool MatchesRangedKey(ReadOnlySpan keyName, ref int keyIndex, out bool isDev) { Assert.Equal((int)KeyRangeType.Range, (int)RangeType); @@ -120,7 +120,7 @@ namespace LibHac.Common.Keys if (keyName.Length == Name.Length + 7) { // Check if "_dev" comes after the base key name - if (!keyName.AsSpan(Name.Length, 4).SequenceEqual("_dev")) + if (!keyName.Slice(Name.Length, 4).SequenceEqual("_dev")) return false; isDev = true; @@ -130,7 +130,7 @@ namespace LibHac.Common.Keys return false; // Check if the name before the "_XX" index matches - if (!keyName.AsSpan(0, Name.Length).SequenceEqual(Name)) + if (!keyName.Slice(0, Name.Length).SequenceEqual(Name)) return false; // The name should have an underscore before the index value @@ -140,7 +140,7 @@ namespace LibHac.Common.Keys byte index = default; // Try to get the index of the key name - if (!keyName.AsSpan(keyName.Length - 2, 2).TryToBytes(SpanHelpers.AsSpan(ref index))) + if (!keyName.Slice(keyName.Length - 2, 2).TryToBytes(SpanHelpers.AsSpan(ref index))) return false; // Check if the index is in this key's range diff --git a/src/LibHac/Common/Keys/KeySet.cs b/src/LibHac/Common/Keys/KeySet.cs index a984670c..490809fb 100644 --- a/src/LibHac/Common/Keys/KeySet.cs +++ b/src/LibHac/Common/Keys/KeySet.cs @@ -1,4 +1,5 @@ using System; +using System.Collections.Generic; using System.Runtime.InteropServices; using System.Security.Cryptography; using LibHac.Boot; @@ -201,6 +202,11 @@ namespace LibHac.Common.Keys return DefaultKeySet.CreateDefaultKeySet(); } + public static List CreateKeyInfoList() + { + return DefaultKeySet.CreateKeyList(); + } + public void DeriveKeys(IProgressReport logger = null) => KeyDerivation.DeriveAllKeys(this, logger); public void DeriveSdCardKeys() => KeyDerivation.DeriveSdCardKeys(this); diff --git a/src/LibHac/Common/StringUtils.cs b/src/LibHac/Common/StringUtils.cs index bc2173c5..4832e46f 100644 --- a/src/LibHac/Common/StringUtils.cs +++ b/src/LibHac/Common/StringUtils.cs @@ -178,5 +178,11 @@ namespace LibHac.Common { return (uint)(c - (byte)'0') <= 9; } + + public static bool IsHexDigit(byte c) + { + return (uint)(c - (byte)'0') <= 9 || + (c | 0x20u) - (byte)'a' <= 'f' - 'a'; + } } } diff --git a/src/LibHac/FsSrv/ExternalKeySet.cs b/src/LibHac/FsSrv/ExternalKeySet.cs index e5e32497..877f10c2 100644 --- a/src/LibHac/FsSrv/ExternalKeySet.cs +++ b/src/LibHac/FsSrv/ExternalKeySet.cs @@ -93,5 +93,10 @@ namespace LibHac.FsSrv ExternalKeys.TrimExcess(newCapacity); } } + + public void EnsureCapacity(int capacity) + { + ExternalKeys.EnsureCapacity(capacity); + } } }