From 33af34cefce165606bfd6101931434892e9fb438 Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Thu, 25 Jun 2020 14:01:45 -0700 Subject: [PATCH] Rewrite IndirectStorage --- build/CodeGen/results.csv | 9 +- src/LibHac/Diag/Assert.cs | 15 ++ src/LibHac/Fs/IStorage.cs | 2 +- src/LibHac/Fs/ResultFs.cs | 16 +- src/LibHac/FsSystem/BucketTree2.cs | 15 ++ src/LibHac/FsSystem/IndirectStorage.cs | 320 ++++++++++++++++++++----- src/LibHac/FsSystem/NcaUtils/Nca.cs | 23 +- src/LibHac/Util.cs | 15 +- 8 files changed, 332 insertions(+), 83 deletions(-) diff --git a/build/CodeGen/results.csv b/build/CodeGen/results.csv index 087dccac..3adccd8c 100644 --- a/build/CodeGen/results.csv +++ b/build/CodeGen/results.csv @@ -50,7 +50,14 @@ Module,DescriptionStart,DescriptionEnd,Name,Summary 2,4000,4999,DataCorrupted, 2,4001,4299,RomCorrupted, -2,4023,,InvalidIndirectStorageSource, + +2,4021,4029,IndirectStorageCorrupted, +2,4022,,InvalidIndirectEntryOffset, +2,4023,,InvalidIndirectEntryStorageIndex, +2,4024,,InvalidIndirectStorageSize, +2,4025,,InvalidIndirectVirtualOffset, +2,4026,,InvalidIndirectPhysicalOffset, +2,4027,,InvalidIndirectStorageIndex, 2,4031,4039,BucketTreeCorrupted, 2,4032,,InvalidBucketTreeSignature, diff --git a/src/LibHac/Diag/Assert.cs b/src/LibHac/Diag/Assert.cs index 84c98323..a269a4fb 100644 --- a/src/LibHac/Diag/Assert.cs +++ b/src/LibHac/Diag/Assert.cs @@ -27,5 +27,20 @@ namespace LibHac.Diag throw new LibHacException("Not-null assertion failed."); } } + + [Conditional("DEBUG")] + public static void InRange(int value, int lowerInclusive, int upperExclusive) + { + InRange((long)value, lowerInclusive, upperExclusive); + } + + [Conditional("DEBUG")] + public static void InRange(long value, long lowerInclusive, long upperExclusive) + { + if (value < lowerInclusive || value >= upperExclusive) + { + throw new LibHacException($"Value {value} is not in the range {lowerInclusive} to {upperExclusive}"); + } + } } } diff --git a/src/LibHac/Fs/IStorage.cs b/src/LibHac/Fs/IStorage.cs index 05121695..9c023322 100644 --- a/src/LibHac/Fs/IStorage.cs +++ b/src/LibHac/Fs/IStorage.cs @@ -132,8 +132,8 @@ namespace LibHac.Fs protected abstract Result DoRead(long offset, Span destination); protected abstract Result DoWrite(long offset, ReadOnlySpan source); protected abstract Result DoFlush(); - protected abstract Result DoGetSize(out long size); protected abstract Result DoSetSize(long size); + protected abstract Result DoGetSize(out long size); protected virtual Result DoOperateRange(Span outBuffer, OperationId operationId, long offset, long size, ReadOnlySpan inBuffer) diff --git a/src/LibHac/Fs/ResultFs.cs b/src/LibHac/Fs/ResultFs.cs index dd8ca748..a1bf0ae9 100644 --- a/src/LibHac/Fs/ResultFs.cs +++ b/src/LibHac/Fs/ResultFs.cs @@ -114,8 +114,20 @@ namespace LibHac.Fs public static Result.Base DataCorrupted { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => new Result.Base(ModuleFs, 4000, 4999); } /// Error code: 2002-4001; Range: 4001-4299; Inner value: 0x1f4202 public static Result.Base RomCorrupted { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => new Result.Base(ModuleFs, 4001, 4299); } - /// Error code: 2002-4023; Inner value: 0x1f6e02 - public static Result.Base InvalidIndirectStorageSource => new Result.Base(ModuleFs, 4023); + /// Error code: 2002-4021; Range: 4021-4029; Inner value: 0x1f6a02 + public static Result.Base IndirectStorageCorrupted { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => new Result.Base(ModuleFs, 4021, 4029); } + /// Error code: 2002-4022; Inner value: 0x1f6c02 + public static Result.Base InvalidIndirectEntryOffset => new Result.Base(ModuleFs, 4022); + /// Error code: 2002-4023; Inner value: 0x1f6e02 + public static Result.Base InvalidIndirectEntryStorageIndex => new Result.Base(ModuleFs, 4023); + /// Error code: 2002-4024; Inner value: 0x1f7002 + public static Result.Base InvalidIndirectStorageSize => new Result.Base(ModuleFs, 4024); + /// Error code: 2002-4025; Inner value: 0x1f7202 + public static Result.Base InvalidIndirectVirtualOffset => new Result.Base(ModuleFs, 4025); + /// Error code: 2002-4026; Inner value: 0x1f7402 + public static Result.Base InvalidIndirectPhysicalOffset => new Result.Base(ModuleFs, 4026); + /// Error code: 2002-4027; Inner value: 0x1f7602 + public static Result.Base InvalidIndirectStorageIndex => new Result.Base(ModuleFs, 4027); /// Error code: 2002-4031; Range: 4031-4039; Inner value: 0x1f7e02 public static Result.Base BucketTreeCorrupted { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => new Result.Base(ModuleFs, 4031, 4039); } diff --git a/src/LibHac/FsSystem/BucketTree2.cs b/src/LibHac/FsSystem/BucketTree2.cs index e65bcff8..97598b30 100644 --- a/src/LibHac/FsSystem/BucketTree2.cs +++ b/src/LibHac/FsSystem/BucketTree2.cs @@ -105,6 +105,16 @@ namespace LibHac.FsSystem public long GetEnd() => EndOffset; public long GetSize() => EndOffset - StartOffset; + public bool Includes(long offset) + { + return StartOffset <= offset && offset < EndOffset; + } + + public bool Includes(long offset, long size) + { + return size > 0 && StartOffset <= offset && size <= EndOffset - offset; + } + public Result Find(ref Visitor visitor, long virtualAddress) { Assert.AssertTrue(IsInitialized()); @@ -382,6 +392,11 @@ namespace LibHac.FsSystem return Result.Success; } + public void Dispose() + { + // todo: try using shared arrays + } + public bool IsValid() => EntryIndex >= 0; public bool CanMoveNext() diff --git a/src/LibHac/FsSystem/IndirectStorage.cs b/src/LibHac/FsSystem/IndirectStorage.cs index 2acb3ae0..d1df2015 100644 --- a/src/LibHac/FsSystem/IndirectStorage.cs +++ b/src/LibHac/FsSystem/IndirectStorage.cs @@ -1,72 +1,183 @@ using System; -using System.Collections.Generic; -using System.Linq; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using LibHac.Common; +using LibHac.Diag; using LibHac.Fs; namespace LibHac.FsSystem { public class IndirectStorage : IStorage { - private List RelocationEntries { get; } - private List RelocationOffsets { get; } + public static readonly int StorageCount = 2; + public static readonly int NodeSize = 1024 * 16; - private List Sources { get; } = new List(); - private BucketTree BucketTree { get; } - private long Length { get; } - private bool LeaveOpen { get; } + private BucketTree2 Table { get; } = new BucketTree2(); + private SubStorage2[] DataStorage { get; } = new SubStorage2[StorageCount]; - public IndirectStorage(IStorage bucketTreeData, bool leaveOpen, params IStorage[] sources) + [StructLayout(LayoutKind.Sequential, Size = 0x14)] + public struct Entry { - Sources.AddRange(sources); + private long VirtualOffset; + private long PhysicalOffset; + public int StorageIndex; - LeaveOpen = leaveOpen; + public void SetVirtualOffset(long offset) => VirtualOffset = offset; + public long GetVirtualOffset() => VirtualOffset; - BucketTree = new BucketTree(bucketTreeData); - - RelocationEntries = BucketTree.GetEntryList(); - RelocationOffsets = RelocationEntries.Select(x => x.Offset).ToList(); - - Length = BucketTree.BucketOffsets.OffsetEnd; + public void SetPhysicalOffset(long offset) => PhysicalOffset = offset; + public long GetPhysicalOffset() => PhysicalOffset; } - protected override Result DoRead(long offset, Span destination) + public static long QueryHeaderStorageSize() => BucketTree2.QueryHeaderStorageSize(); + + public static long QueryNodeStorageSize(int entryCount) => + BucketTree2.QueryNodeStorageSize(NodeSize, Unsafe.SizeOf(), entryCount); + + public static long QueryEntryStorageSize(int entryCount) => + BucketTree2.QueryEntryStorageSize(NodeSize, Unsafe.SizeOf(), entryCount); + + public bool IsInitialized() => Table.IsInitialized(); + + public Result Initialize(SubStorage2 tableStorage) { - RelocationEntry entry = GetRelocationEntry(offset); + // Read and verify the bucket tree header. + // note: skip init + var header = new BucketTree2.Header(); - if (entry.SourceIndex > Sources.Count) + Result rc = tableStorage.Read(0, SpanHelpers.AsByteSpan(ref header)); + if (rc.IsFailure()) return rc; + + rc = header.Verify(); + if (rc.IsFailure()) return rc; + + // Determine extents. + long nodeStorageSize = QueryNodeStorageSize(header.EntryCount); + long entryStorageSize = QueryEntryStorageSize(header.EntryCount); + long nodeStorageOffset = QueryHeaderStorageSize(); + long entryStorageOffset = nodeStorageOffset + nodeStorageSize; + + // Initialize. + var nodeStorage = new SubStorage2(tableStorage, nodeStorageOffset, nodeStorageSize); + var entryStorage = new SubStorage2(tableStorage, entryStorageOffset, entryStorageSize); + + return Initialize(nodeStorage, entryStorage, header.EntryCount); + } + + public Result Initialize(SubStorage2 nodeStorage, SubStorage2 entryStorage, int entryCount) + { + return Table.Initialize(nodeStorage, entryStorage, NodeSize, Unsafe.SizeOf(), entryCount); + } + + public void SetStorage(int index, SubStorage2 storage) + { + Assert.InRange(index, 0, StorageCount); + DataStorage[index] = storage; + } + + public void SetStorage(int index, IStorage storage, long offset, long size) + { + Assert.InRange(index, 0, StorageCount); + DataStorage[index] = new SubStorage2(storage, offset, size); + } + + public Result GetEntryList(Span entryBuffer, out int outputEntryCount, long offset, long size) + { + // Validate pre-conditions + Assert.AssertTrue(offset >= 0); + Assert.AssertTrue(size >= 0); + Assert.AssertTrue(IsInitialized()); + + // Clear the out count + outputEntryCount = 0; + + // Succeed if there's no range + if (size == 0) + return Result.Success; + + // Check that our range is valid + if (!Table.Includes(offset, size)) + return ResultFs.OutOfRange.Log(); + + // Find the offset in our tree + var visitor = new BucketTree2.Visitor(); + Result rc = Table.Find(ref visitor, offset); + if (rc.IsFailure()) return rc; + + long entryOffset = visitor.Get().GetVirtualOffset(); + if (entryOffset > 0 || !Table.Includes(entryOffset)) + return ResultFs.InvalidIndirectEntryOffset.Log(); + + // Prepare to loop over entries + long endOffset = offset + size; + int count = 0; + + ref Entry currentEntry = ref visitor.Get(); + while (currentEntry.GetVirtualOffset() < endOffset) { - return ResultFs.InvalidIndirectStorageSource.Log(); - } - - long inPos = offset; - int outPos = 0; - int remaining = destination.Length; - - while (remaining > 0) - { - long entryPos = inPos - entry.Offset; - - int bytesToRead = (int)Math.Min(entry.OffsetEnd - inPos, remaining); - - Result rc = Sources[entry.SourceIndex].Read(entry.SourceOffset + entryPos, destination.Slice(outPos, bytesToRead)); - if (rc.IsFailure()) return rc; - - outPos += bytesToRead; - inPos += bytesToRead; - remaining -= bytesToRead; - - if (inPos >= entry.OffsetEnd) + // Try to write the entry to the out list + if (entryBuffer.Length != 0) { - entry = entry.Next; + if (count >= entryBuffer.Length) + break; + + entryBuffer[count] = currentEntry; + } + + count++; + + // Advance + if (visitor.CanMoveNext()) + { + rc = visitor.MoveNext(); + if (rc.IsFailure()) return rc; + + currentEntry = ref visitor.Get(); + } + else + { + break; } } + // Write the entry count + outputEntryCount = count; return Result.Success; } + protected override unsafe Result DoRead(long offset, Span destination) + { + // Validate pre-conditions + Assert.AssertTrue(offset >= 0); + Assert.AssertTrue(IsInitialized()); + + // Succeed if there's nothing to read + if (destination.Length == 0) + return Result.Success; + + // Pin and recreate the span because C# can't use byref-like types in a closure + int bufferSize = destination.Length; + fixed (byte* pBuffer = destination) + { + // Copy the pointer to workaround CS1764. + // OperatePerEntry won't store the delegate anywhere, so it should be safe + byte* pBuffer2 = pBuffer; + + Result Operate(IStorage storage, long dataOffset, long currentOffset, long currentSize) + { + var buffer = new Span(pBuffer2, bufferSize); + + return storage.Read(dataOffset, + buffer.Slice((int)(currentOffset - offset), (int)currentSize)); + } + + return OperatePerEntry(offset, destination.Length, Operate); + } + } + protected override Result DoWrite(long offset, ReadOnlySpan source) { - return ResultFs.UnsupportedOperationInIndirectStorageSetSize.Log(); + return ResultFs.UnsupportedOperationInIndirectStorageWrite.Log(); } protected override Result DoFlush() @@ -74,36 +185,117 @@ namespace LibHac.FsSystem return Result.Success; } - protected override Result DoGetSize(out long size) - { - size = Length; - return Result.Success; - } - protected override Result DoSetSize(long size) { return ResultFs.UnsupportedOperationInIndirectStorageSetSize.Log(); } - protected override void Dispose(bool disposing) + protected override Result DoGetSize(out long size) { - if (disposing) - { - if (!LeaveOpen && Sources != null) - { - foreach (IStorage storage in Sources) - { - storage?.Dispose(); - } - } - } + size = Table.GetEnd(); + return Result.Success; } - private RelocationEntry GetRelocationEntry(long offset) + private delegate Result OperateFunc(IStorage storage, long dataOffset, long currentOffset, long currentSize); + + private Result OperatePerEntry(long offset, long size, OperateFunc func) { - int index = RelocationOffsets.BinarySearch(offset); - if (index < 0) index = ~index - 1; - return RelocationEntries[index]; + // Validate preconditions + Assert.AssertTrue(offset >= 0); + Assert.AssertTrue(size >= 0); + Assert.AssertTrue(IsInitialized()); + + // Succeed if there's nothing to operate on + if (size == 0) + return Result.Success; + + // Validate arguments + if (!Table.Includes(offset, size)) + return ResultFs.OutOfRange.Log(); + + // Find the offset in our tree + var visitor = new BucketTree2.Visitor(); + + Result rc = Table.Find(ref visitor, offset); + if (rc.IsFailure()) return rc; + + long entryOffset = visitor.Get().GetVirtualOffset(); + if (entryOffset < 0 || !Table.Includes(entryOffset)) + return ResultFs.InvalidIndirectEntryStorageIndex.Log(); + + // Prepare to operate in chunks + long currentOffset = offset; + long endOffset = offset + size; + + while (currentOffset < endOffset) + { + // Get the current entry + var currentEntry = visitor.Get(); + + // Get and validate the entry's offset + long currentEntryOffset = currentEntry.GetVirtualOffset(); + if (currentEntryOffset > currentOffset) + return ResultFs.InvalidIndirectEntryOffset.Log(); + + // Validate the storage index + if (currentEntry.StorageIndex < 0 || currentEntry.StorageIndex >= StorageCount) + return ResultFs.InvalidIndirectEntryStorageIndex.Log(); + + // todo: Implement continuous reading + + // Get and validate the next entry offset + long nextEntryOffset; + if (visitor.CanMoveNext()) + { + rc = visitor.MoveNext(); + if (rc.IsFailure()) return rc; + + nextEntryOffset = visitor.Get().GetVirtualOffset(); + if (!Table.Includes(nextEntryOffset)) + return ResultFs.InvalidIndirectEntryOffset.Log(); + } + else + { + nextEntryOffset = Table.GetEnd(); + } + + if (currentOffset >= nextEntryOffset) + return ResultFs.InvalidIndirectEntryOffset.Log(); + + // Get the offset of the entry in the data we read + long dataOffset = currentOffset - currentEntryOffset; + long dataSize = nextEntryOffset - currentEntryOffset - dataOffset; + Assert.AssertTrue(dataSize > 0); + + // Determine how much is left + long remainingSize = endOffset - currentOffset; + long currentSize = Math.Min(remainingSize, dataSize); + Assert.AssertTrue(currentSize <= size); + + { + SubStorage2 currentStorage = DataStorage[currentEntry.StorageIndex]; + + // Get the current data storage's size. + rc = currentStorage.GetSize(out long currentDataStorageSize); + if (rc.IsFailure()) return rc; + + // Ensure that we remain within range. + long currentEntryPhysicalOffset = currentEntry.GetPhysicalOffset(); + + if (currentEntryPhysicalOffset < 0 || currentEntryPhysicalOffset > currentDataStorageSize) + return ResultFs.IndirectStorageCorrupted.Log(); + + if (currentDataStorageSize < currentEntryPhysicalOffset + dataOffset + currentSize) + return ResultFs.IndirectStorageCorrupted.Log(); + + rc = func(currentStorage, currentEntryPhysicalOffset + dataOffset, currentOffset, currentSize); + if (rc.IsFailure()) return rc; + } + + currentOffset += currentSize; + } + + return Result.Success; } } } diff --git a/src/LibHac/FsSystem/NcaUtils/Nca.cs b/src/LibHac/FsSystem/NcaUtils/Nca.cs index 0626f84b..ccb11b87 100644 --- a/src/LibHac/FsSystem/NcaUtils/Nca.cs +++ b/src/LibHac/FsSystem/NcaUtils/Nca.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.Diagnostics; using System.IO; +using LibHac.Common; using LibHac.Fs; using LibHac.Fs.Fsa; using LibHac.FsSystem.RomFs; @@ -214,6 +215,9 @@ namespace LibHac.FsSystem.NcaUtils IStorage patchStorage = patchNca.OpenRawStorage(index); IStorage baseStorage = SectionExists(index) ? OpenRawStorage(index) : new NullStorage(); + patchStorage.GetSize(out long patchSize).ThrowIfFailure(); + baseStorage.GetSize(out long baseSize).ThrowIfFailure(); + NcaFsHeader header = patchNca.Header.GetFsHeader(index); NcaFsPatchInfo patchInfo = header.GetPatchInfo(); @@ -222,9 +226,24 @@ namespace LibHac.FsSystem.NcaUtils return patchStorage; } - IStorage relocationTableStorage = patchStorage.Slice(patchInfo.RelocationTreeOffset, patchInfo.RelocationTreeSize); + var treeHeader = new BucketTree2.Header(); + patchInfo.RelocationTreeHeader.CopyTo(SpanHelpers.AsByteSpan(ref treeHeader)); + long nodeStorageSize = IndirectStorage.QueryNodeStorageSize(treeHeader.EntryCount); + long entryStorageSize = IndirectStorage.QueryEntryStorageSize(treeHeader.EntryCount); - return new IndirectStorage(relocationTableStorage, true, baseStorage, patchStorage); + var relocationTableStorage = new SubStorage2(patchStorage, patchInfo.RelocationTreeOffset, patchInfo.RelocationTreeSize); + var cachedTableStorage = new CachedStorage(relocationTableStorage, IndirectStorage.NodeSize, 4, true); + + var tableNodeStorage = new SubStorage2(cachedTableStorage, 0, nodeStorageSize); + var tableEntryStorage = new SubStorage2(cachedTableStorage, nodeStorageSize, entryStorageSize); + + var storage = new IndirectStorage(); + storage.Initialize(tableNodeStorage, tableEntryStorage, treeHeader.EntryCount).ThrowIfFailure(); + + storage.SetStorage(0, baseStorage, 0, baseSize); + storage.SetStorage(1, patchStorage, 0, patchSize); + + return storage; } public IStorage OpenStorage(int index, IntegrityCheckLevel integrityCheckLevel) diff --git a/src/LibHac/Util.cs b/src/LibHac/Util.cs index 31c7ce7b..2bfd26bb 100644 --- a/src/LibHac/Util.cs +++ b/src/LibHac/Util.cs @@ -48,20 +48,9 @@ namespace LibHac return true; } - public static bool SpansEqual(Span a1, Span a2) + public static bool SpansEqual(Span a1, Span a2) where T : IEquatable { - if (a1 == a2) return true; - if (a1.Length != a2.Length) return false; - - for (int i = 0; i < a1.Length; i++) - { - if (!a1[i].Equals(a2[i])) - { - return false; - } - } - - return true; + return a1.SequenceEqual(a2); } public static ReadOnlySpan GetUtf8Bytes(string value)