diff --git a/src/LibHac/FsSystem/Aes128CtrExStorage.cs b/src/LibHac/FsSystem/Aes128CtrExStorage.cs index a6cae265..42446e43 100644 --- a/src/LibHac/FsSystem/Aes128CtrExStorage.cs +++ b/src/LibHac/FsSystem/Aes128CtrExStorage.cs @@ -1,39 +1,43 @@ using System; -using System.Collections.Generic; -using System.Linq; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; using LibHac.Fs; namespace LibHac.FsSystem { public class Aes128CtrExStorage : Aes128CtrStorage { - private List SubsectionEntries { get; } - private List SubsectionOffsets { get; } - private BucketTree BucketTree { get; } + public static readonly int NodeSize = 1024 * 16; + + private BucketTree2 Table { get; } = new BucketTree2(); private readonly object _locker = new object(); - public Aes128CtrExStorage(IStorage baseStorage, IStorage bucketTreeData, byte[] key, long counterOffset, byte[] ctrHi, bool leaveOpen) - : base(baseStorage, key, counterOffset, ctrHi, leaveOpen) + [StructLayout(LayoutKind.Sequential, Size = 0x10)] + public struct Entry { - BucketTree = new BucketTree(bucketTreeData); - - SubsectionEntries = BucketTree.GetEntryList(); - SubsectionOffsets = SubsectionEntries.Select(x => x.Offset).ToList(); + public long Offset; + public int Reserved; + public int Generation; } - public Aes128CtrExStorage(IStorage baseStorage, IStorage bucketTreeData, byte[] key, byte[] counter, bool leaveOpen) + public Aes128CtrExStorage(IStorage baseStorage, SubStorage2 nodeStorage, SubStorage2 entryStorage, + int entryCount, byte[] key, byte[] counter, bool leaveOpen) : base(baseStorage, key, counter, leaveOpen) { - BucketTree = new BucketTree(bucketTreeData); - - SubsectionEntries = BucketTree.GetEntryList(); - SubsectionOffsets = SubsectionEntries.Select(x => x.Offset).ToList(); + Result rc = Table.Initialize(nodeStorage, entryStorage, NodeSize, Unsafe.SizeOf(), entryCount); + rc.ThrowIfFailure(); } protected override Result DoRead(long offset, Span destination) { - AesSubsectionEntry entry = GetSubsectionEntry(offset); + if (destination.Length == 0) + return Result.Success; + + var visitor = new BucketTree2.Visitor(); + + Result rc = Table.Find(ref visitor, offset); + if (rc.IsFailure()) return rc; long inPos = offset; int outPos = 0; @@ -41,24 +45,37 @@ namespace LibHac.FsSystem while (remaining > 0) { - int bytesToRead = (int)Math.Min(entry.OffsetEnd - inPos, remaining); + var currentEntry = visitor.Get(); + + // Get and validate the next entry offset + long nextEntryOffset; + if (visitor.CanMoveNext()) + { + rc = visitor.MoveNext(); + if (rc.IsFailure()) return rc; + + nextEntryOffset = visitor.Get().Offset; + if (!Table.Includes(nextEntryOffset)) + return ResultFs.InvalidIndirectEntryOffset.Log(); + } + else + { + nextEntryOffset = Table.GetEnd(); + } + + int bytesToRead = (int)Math.Min(nextEntryOffset - inPos, remaining); lock (_locker) { - UpdateCounterSubsection(entry.Counter); + UpdateCounterSubsection((uint)currentEntry.Generation); - Result rc = base.DoRead(inPos, destination.Slice(outPos, bytesToRead)); + rc = base.DoRead(inPos, destination.Slice(outPos, bytesToRead)); if (rc.IsFailure()) return rc; } outPos += bytesToRead; inPos += bytesToRead; remaining -= bytesToRead; - - if (remaining != 0 && inPos >= entry.OffsetEnd) - { - entry = entry.Next; - } } return Result.Success; @@ -74,13 +91,6 @@ namespace LibHac.FsSystem return Result.Success; } - private AesSubsectionEntry GetSubsectionEntry(long offset) - { - int index = SubsectionOffsets.BinarySearch(offset); - if (index < 0) index = ~index - 1; - return SubsectionEntries[index]; - } - private void UpdateCounterSubsection(uint value) { Counter[7] = (byte)value; diff --git a/src/LibHac/FsSystem/NcaUtils/Nca.cs b/src/LibHac/FsSystem/NcaUtils/Nca.cs index ccb11b87..1681ef08 100644 --- a/src/LibHac/FsSystem/NcaUtils/Nca.cs +++ b/src/LibHac/FsSystem/NcaUtils/Nca.cs @@ -194,10 +194,21 @@ namespace LibHac.FsSystem.NcaUtils byte[] counterEx = Aes128CtrStorage.CreateCounter(fsHeader.Counter, sectionOffset); IStorage bucketTreeData = new CachedStorage(new Aes128CtrStorage(baseStorage.Slice(bktrOffset, bktrSize), key, counter, true), 4, true); + var encryptionBucketTreeData = new SubStorage2(bucketTreeData, + info.EncryptionTreeOffset - bktrOffset, sectionSize - info.EncryptionTreeOffset); - IStorage encryptionBucketTreeData = bucketTreeData.Slice(info.EncryptionTreeOffset - bktrOffset); - IStorage decStorage = new Aes128CtrExStorage(baseStorage.Slice(0, dataSize), encryptionBucketTreeData, key, counterEx, true); - decStorage = new CachedStorage(decStorage, 0x4000, 4, true); + var cachedBucketTreeData = new CachedStorage(encryptionBucketTreeData, IndirectStorage.NodeSize, 6, true); + + var treeHeader = new BucketTree2.Header(); + info.EncryptionTreeHeader.CopyTo(SpanHelpers.AsByteSpan(ref treeHeader)); + long nodeStorageSize = IndirectStorage.QueryNodeStorageSize(treeHeader.EntryCount); + long entryStorageSize = IndirectStorage.QueryEntryStorageSize(treeHeader.EntryCount); + + var tableNodeStorage = new SubStorage2(cachedBucketTreeData, 0, nodeStorageSize); + var tableEntryStorage = new SubStorage2(cachedBucketTreeData, nodeStorageSize, entryStorageSize); + + IStorage decStorage = new Aes128CtrExStorage(baseStorage.Slice(0, dataSize), tableNodeStorage, + tableEntryStorage, treeHeader.EntryCount, key, counterEx, true); return new ConcatenationStorage(new[] { decStorage, bucketTreeData }, true); }