From 0c06d9e0b3585db0ae7ced39a38f995ec3a46c00 Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Mon, 18 May 2020 22:15:18 -0700 Subject: [PATCH 1/8] Implement most of Nintendo's bucket tree code --- build/CodeGen/results.csv | 11 + src/LibHac/Fs/ResultFs.cs | 21 + src/LibHac/FsSystem/BucketTree2.cs | 730 +++++++++++++++++++++++++++++ src/LibHac/Util.cs | 22 + 4 files changed, 784 insertions(+) create mode 100644 src/LibHac/FsSystem/BucketTree2.cs diff --git a/build/CodeGen/results.csv b/build/CodeGen/results.csv index dd634f7f..1b9bd8e7 100644 --- a/build/CodeGen/results.csv +++ b/build/CodeGen/results.csv @@ -52,6 +52,16 @@ Module,DescriptionStart,DescriptionEnd,Name,Summary 2,4001,4299,RomCorrupted, 2,4023,,InvalidIndirectStorageSource, +2,4031,4039,BucketTreeCorrupted, +2,4032,,InvalidBucketTreeSignature, +2,4033,,InvalidBucketTreeEntryCount, +2,4034,,InvalidBucketTreeNodeEntryCount, +2,4035,,InvalidBucketTreeNodeOffset, +2,4036,,InvalidBucketTreeEntryOffset, +2,4037,,InvalidBucketTreeEntrySetOffset, +2,4038,,InvalidBucketTreeNodeIndex, +2,4039,,BucketTreeEntryNotFound, + 2,4241,4259,RomHostFileSystemCorrupted, 2,4242,,RomHostEntryCorrupted, 2,4243,,RomHostFileDataCorrupted, @@ -220,6 +230,7 @@ Module,DescriptionStart,DescriptionEnd,Name,Summary 2,6606,,TargetProgramIndexNotFound,Specified program index is not found 2,6700,6799,OutOfResource, +2,6705,,BufferAllocationFailed, 2,6706,,MappingTableFull, 2,6707,,AllocationTableInsufficientFreeBlocks, 2,6709,,OpenCountLimit, diff --git a/src/LibHac/Fs/ResultFs.cs b/src/LibHac/Fs/ResultFs.cs index f2af9707..179f181a 100644 --- a/src/LibHac/Fs/ResultFs.cs +++ b/src/LibHac/Fs/ResultFs.cs @@ -117,6 +117,25 @@ namespace LibHac.Fs /// Error code: 2002-4023; Inner value: 0x1f6e02 public static Result.Base InvalidIndirectStorageSource => new Result.Base(ModuleFs, 4023); + /// 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); } + /// Error code: 2002-4032; Inner value: 0x1f8002 + public static Result.Base InvalidBucketTreeSignature => new Result.Base(ModuleFs, 4032); + /// Error code: 2002-4033; Inner value: 0x1f8202 + public static Result.Base InvalidBucketTreeEntryCount => new Result.Base(ModuleFs, 4033); + /// Error code: 2002-4034; Inner value: 0x1f8402 + public static Result.Base InvalidBucketTreeNodeEntryCount => new Result.Base(ModuleFs, 4034); + /// Error code: 2002-4035; Inner value: 0x1f8602 + public static Result.Base InvalidBucketTreeNodeOffset => new Result.Base(ModuleFs, 4035); + /// Error code: 2002-4036; Inner value: 0x1f8802 + public static Result.Base InvalidBucketTreeEntryOffset => new Result.Base(ModuleFs, 4036); + /// Error code: 2002-4037; Inner value: 0x1f8a02 + public static Result.Base InvalidBucketTreeEntrySetOffset => new Result.Base(ModuleFs, 4037); + /// Error code: 2002-4038; Inner value: 0x1f8c02 + public static Result.Base InvalidBucketTreeNodeIndex => new Result.Base(ModuleFs, 4038); + /// Error code: 2002-4039; Inner value: 0x1f8e02 + public static Result.Base BucketTreeEntryNotFound => new Result.Base(ModuleFs, 4039); + /// Error code: 2002-4241; Range: 4241-4259; Inner value: 0x212202 public static Result.Base RomHostFileSystemCorrupted { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => new Result.Base(ModuleFs, 4241, 4259); } /// Error code: 2002-4242; Inner value: 0x212402 @@ -426,6 +445,8 @@ namespace LibHac.Fs /// Error code: 2002-6700; Range: 6700-6799; Inner value: 0x345802 public static Result.Base OutOfResource { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => new Result.Base(ModuleFs, 6700, 6799); } + /// Error code: 2002-6705; Inner value: 0x346202 + public static Result.Base BufferAllocationFailed => new Result.Base(ModuleFs, 6705); /// Error code: 2002-6706; Inner value: 0x346402 public static Result.Base MappingTableFull => new Result.Base(ModuleFs, 6706); /// Error code: 2002-6707; Inner value: 0x346602 diff --git a/src/LibHac/FsSystem/BucketTree2.cs b/src/LibHac/FsSystem/BucketTree2.cs new file mode 100644 index 00000000..ace5ef7a --- /dev/null +++ b/src/LibHac/FsSystem/BucketTree2.cs @@ -0,0 +1,730 @@ +using System; +using System.Buffers.Binary; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using LibHac.Common; +using LibHac.Diag; +using LibHac.Fs; + +namespace LibHac.FsSystem +{ + public class BucketTree2 + { + private const uint ExpectedMagic = 0x52544B42; // BKTR + private const int MaxVersion = 1; + + private const int NodeSizeMin = 1024; + private const int NodeSizeMax = 1024 * 512; + + private static int NodeHeaderSize => Unsafe.SizeOf(); + + private SubStorage2 NodeStorage { get; set; } + private SubStorage2 EntryStorage { get; set; } + + private NodeBuffer NodeL1 { get; } = new NodeBuffer(); + + private long NodeSize { get; set; } + private long EntrySize { get; set; } + private int EntryCount { get; set; } + private int OffsetCount { get; set; } + private int EntrySetCount { get; set; } + private long StartOffset { get; set; } + private long EndOffset { get; set; } + + public Result Initialize(SubStorage2 nodeStorage, SubStorage2 entryStorage, int nodeSize, int entrySize, + int entryCount) + { + Assert.AssertTrue(entrySize >= sizeof(long)); + Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); + Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); + Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + Assert.AssertTrue(!IsInitialized()); + + // Ensure valid entry count. + if (entryCount <= 0) + return ResultFs.InvalidArgument.Log(); + + // Allocate node. + if (!NodeL1.Allocate(nodeSize)) + return ResultFs.BufferAllocationFailed.Log(); + + bool needFree = true; + try + { + // Read node. + Result rc = nodeStorage.Read(0, NodeL1.GetBuffer()); + if (rc.IsFailure()) return rc; + + // Verify node. + rc = NodeL1.GetHeader().Verify(0, nodeSize, sizeof(long)); + if (rc.IsFailure()) return rc; + + // Validate offsets. + int offsetCount = GetOffsetCount(nodeSize); + int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); + BucketTreeNode node = NodeL1.GetNode(); + + long startOffset; + if (offsetCount < entrySetCount && node.GetCount() < offsetCount) + { + startOffset = node.GetL2BeginOffset(); + } + else + { + startOffset = node.GetBeginOffset(); + } + + long endOffset = node.GetEndOffset(); + + if (startOffset < 0 || startOffset > node.GetBeginOffset() || startOffset >= endOffset) + return ResultFs.InvalidBucketTreeEntryOffset.Log(); + + NodeStorage = nodeStorage; + EntryStorage = entryStorage; + NodeSize = nodeSize; + EntrySize = entrySize; + EntryCount = entryCount; + OffsetCount = offsetCount; + EntrySetCount = entrySetCount; + StartOffset = startOffset; + EndOffset = endOffset; + + needFree = false; + + return Result.Success; + } + finally + { + if (needFree) + NodeL1.Free(); + } + } + + public bool IsInitialized() => NodeSize > 0; + public bool IsEmpty() => EntrySize == 0; + + public Result Find(ref Visitor visitor, long virtualAddress) + { + Assert.AssertTrue(IsInitialized()); + + if (virtualAddress < 0) + return ResultFs.InvalidOffset.Log(); + + if (IsEmpty()) + return ResultFs.OutOfRange.Log(); + + Result rc = visitor.Initialize(this); + if (rc.IsFailure()) return rc; + + return visitor.Find(virtualAddress); + } + + public static int QueryHeaderStorageSize() => Unsafe.SizeOf
(); + + public static long QueryNodeStorageSize(long nodeSize, long entrySize, int entryCount) + { + Assert.AssertTrue(entrySize >= sizeof(long)); + Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); + Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); + Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + Assert.AssertTrue(entryCount >= 0); + + if (entryCount <= 0) + return 0; + + return (1 + GetNodeL2Count(nodeSize, entrySize, entryCount)) * nodeSize; + } + + public static long QueryEntryStorageSize(long nodeSize, long entrySize, int entryCount) + { + Assert.AssertTrue(entrySize >= sizeof(long)); + Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); + Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); + Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + Assert.AssertTrue(entryCount >= 0); + + if (entryCount <= 0) + return 0; + + return GetEntrySetCount(nodeSize, entrySize, entryCount) * nodeSize; + } + + private static int GetEntryCount(long nodeSize, long entrySize) + { + return (int)((nodeSize - Unsafe.SizeOf()) / entrySize); + } + + private static int GetOffsetCount(long nodeSize) + { + return (int)((nodeSize - Unsafe.SizeOf()) / sizeof(long)); + } + + private static int GetEntrySetCount(long nodeSize, long entrySize, int entryCount) + { + int entryCountPerNode = GetEntryCount(nodeSize, entrySize); + return Util.DivideByRoundUp(entryCount, entryCountPerNode); + } + + private static int GetNodeL2Count(long nodeSize, long entrySize, int entryCount) + { + int offsetCountPerNode = GetOffsetCount(nodeSize); + int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); + + if (entrySetCount <= offsetCountPerNode) + return 0; + + int nodeL2Count = Util.DivideByRoundUp(entrySetCount, offsetCountPerNode); + Abort.DoAbortUnless(nodeL2Count <= offsetCountPerNode); + + return Util.DivideByRoundUp(entrySetCount - (offsetCountPerNode - (nodeL2Count - 1)), offsetCountPerNode); + } + + private static long GetBucketTreeEntryOffset(long entrySetOffset, long entrySize, int entryIndex) + { + return entrySetOffset + Unsafe.SizeOf() + entryIndex * entrySize; + } + + private static long GetBucketTreeEntryOffset(int entrySetIndex, long nodeSize, long entrySize, int entryIndex) + { + return GetBucketTreeEntryOffset(entrySetIndex * nodeSize, entrySize, entryIndex); + } + + private bool IsExistL2() => OffsetCount < EntrySetCount; + private bool IsExistOffsetL2OnL1() => IsExistL2() && NodeL1.GetHeader().Count < OffsetCount; + + private long GetEntrySetIndex(int nodeIndex, int offsetIndex) + { + return (OffsetCount - NodeL1.GetHeader().Count) + (OffsetCount * nodeIndex) + offsetIndex; + } + + public struct Header + { + public uint Magic; + public uint Version; + public int EntryCount; + private int _reserved; + + public void Format(int entryCount) + { + Magic = ExpectedMagic; + Version = MaxVersion; + EntryCount = entryCount; + _reserved = 0; + } + + public Result Verify() + { + if (Magic != ExpectedMagic) + return ResultFs.InvalidBucketTreeSignature.Log(); + + if (EntryCount < 0) + return ResultFs.InvalidBucketTreeEntryCount.Log(); + + if (Version > MaxVersion) + return ResultFs.UnsupportedVersion.Log(); + + return Result.Success; + } + } + + public struct NodeHeader + { + public int Index; + public int Count; + public long Offset; + + public Result Verify(int nodeIndex, long nodeSize, long entrySize) + { + if (Index != nodeIndex) + return ResultFs.InvalidBucketTreeNodeIndex.Log(); + + if (entrySize == 0 || nodeSize < entrySize + NodeHeaderSize) + return ResultFs.InvalidSize.Log(); + + long maxEntryCount = (nodeSize - NodeHeaderSize) / entrySize; + + if (Count <= 0 || maxEntryCount < Count) + return ResultFs.InvalidBucketTreeNodeEntryCount.Log(); + + if (Offset < 0) + return ResultFs.InvalidBucketTreeNodeOffset.Log(); + + return Result.Success; + } + } + + private class NodeBuffer + { + // Use long to ensure alignment + private long[] _header; + + public bool Allocate(int nodeSize) + { + Assert.AssertTrue(_header == null); + + _header = new long[nodeSize / sizeof(long)]; + + return _header != null; + } + + public void Free() + { + _header = null; + } + + public void FillZero() + { + if (_header != null) + { + Array.Fill(_header, 0); + } + } + + public ref NodeHeader GetHeader() + { + Assert.AssertTrue(_header.Length / sizeof(long) >= Unsafe.SizeOf()); + + return ref Unsafe.As(ref _header[0]); + } + + public Span GetBuffer() + { + return MemoryMarshal.AsBytes(_header.AsSpan()); + } + + public BucketTreeNode GetNode() where TEntry : unmanaged + { + return new BucketTreeNode(GetBuffer()); + } + + public ref T Get() where T : unmanaged + { + if (Unsafe.SizeOf() != Unsafe.SizeOf()) + { + throw new InvalidOperationException(); + } + + Assert.AssertTrue(_header.Length / sizeof(long) >= Unsafe.SizeOf()); + + return ref Unsafe.As(ref _header[0]); + } + } + + public readonly ref struct BucketTreeNode where TEntry : unmanaged + { + private readonly ReadOnlySpan _buffer; + + public BucketTreeNode(ReadOnlySpan buffer) + { + _buffer = buffer; + + Assert.AssertTrue(_buffer.Length >= Unsafe.SizeOf()); + Assert.AssertTrue(_buffer.Length >= Unsafe.SizeOf() + GetHeader().Count * Unsafe.SizeOf()); + } + + public int GetCount() => GetHeader().Count; + + public ReadOnlySpan GetArray() => GetArray(); + + public long GetBeginOffset() => GetArray()[0]; + public long GetEndOffset() => GetHeader().Offset; + public long GetL2BeginOffset() => GetArray()[GetCount()]; + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ReadOnlySpan GetArray() where TElement : unmanaged + { + return MemoryMarshal.Cast(_buffer.Slice(Unsafe.SizeOf())); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private ref NodeHeader GetHeader() + { + return ref Unsafe.As(ref MemoryMarshal.GetReference(_buffer)); + } + } + + public ref struct Visitor + { + private BucketTree2 Tree { get; set; } + private byte[] Entry { get; set; } + private int EntryIndex { get; set; } + private int EntrySetCount { get; set; } + private EntrySetHeader _entrySet; + + [StructLayout(LayoutKind.Explicit)] + private struct EntrySetHeader + { + [FieldOffset(0)] public NodeHeader Header; + [FieldOffset(0)] public EntrySetInfo Info; + + [StructLayout(LayoutKind.Sequential)] + public struct EntrySetInfo + { + public int Index; + public int Count; + public long End; + public long Start; + } + } + + public Result Initialize(BucketTree2 tree) + { + Assert.AssertTrue(tree != null); + Assert.AssertTrue(Tree == null || tree == Tree); + + if (Entry == null) + { + Entry = new byte[tree.EntrySize]; + Tree = tree; + EntryIndex = -1; + } + + return Result.Success; + } + + public bool IsValid() => EntryIndex >= 0; + + public bool CanMoveNext() + { + return IsValid() && (EntryIndex + 1 < _entrySet.Info.Count || _entrySet.Info.Index + 1 < EntrySetCount); + } + + public bool CanMovePrevious() + { + return IsValid() && (EntryIndex > 0 || _entrySet.Info.Index > 0); + } + + public ref T Get() where T : unmanaged + { + return ref MemoryMarshal.Cast(Entry)[0]; + } + + public Result MoveNext() + { + Result rc; + + if (!IsValid()) + return ResultFs.OutOfRange.Log(); + + int entryIndex = EntryIndex + 1; + + // Invalidate our index, and read the header for the next index. + if (entryIndex == _entrySet.Info.Count) + { + int entrySetIndex = _entrySet.Info.Index + 1; + if (entrySetIndex >= EntrySetCount) + return ResultFs.OutOfRange.Log(); + + EntryIndex = -1; + + long end = _entrySet.Info.End; + + long entrySetSize = Tree.NodeSize; + long entrySetOffset = entrySetIndex * entrySetSize; + + rc = Tree.EntryStorage.Read(entrySetOffset, SpanHelpers.AsByteSpan(ref _entrySet)); + if (rc.IsFailure()) return rc; + + rc = _entrySet.Header.Verify(entrySetIndex, entrySetSize, Tree.EntrySize); + if (rc.IsFailure()) return rc; + + if (_entrySet.Info.Start != end || _entrySet.Info.Start >= _entrySet.Info.End) + return ResultFs.InvalidBucketTreeEntrySetOffset.Log(); + + entryIndex = 0; + } + else + { + EntryIndex = 1; + } + + // Read the new entry + long entrySize = Tree.EntrySize; + long entryOffset = GetBucketTreeEntryOffset(_entrySet.Info.Index, Tree.NodeSize, entrySize, entryIndex); + + rc = Tree.EntryStorage.Read(entryOffset, Entry); + if (rc.IsFailure()) return rc; + + // Note that we changed index. + EntryIndex = entryIndex; + return Result.Success; + } + + public Result MovePrevious() + { + Result rc; + + if (!IsValid()) + return ResultFs.OutOfRange.Log(); + + int entryIndex = EntryIndex; + + if (entryIndex == 0) + { + if (_entrySet.Info.Index <= 0) + return ResultFs.OutOfRange.Log(); + + EntryIndex = -1; + + long start = _entrySet.Info.Start; + + long entrySetSize = Tree.NodeSize; + int entrySetIndex = _entrySet.Info.Index - 1; + long entrySetOffset = entrySetIndex * entrySetSize; + + rc = Tree.EntryStorage.Read(entrySetOffset, SpanHelpers.AsByteSpan(ref _entrySet)); + if (rc.IsFailure()) return rc; + + rc = _entrySet.Header.Verify(entrySetIndex, entrySetSize, Tree.EntrySize); + if (rc.IsFailure()) return rc; + + if (_entrySet.Info.End != start || _entrySet.Info.Start >= _entrySet.Info.End) + return ResultFs.InvalidBucketTreeEntrySetOffset.Log(); + + entryIndex = _entrySet.Info.Count; + } + else + { + EntryIndex = -1; + } + + // Read the new entry + long entrySize = Tree.EntrySize; + long entryOffset = GetBucketTreeEntryOffset(_entrySet.Info.Index, Tree.NodeSize, entrySize, entryIndex); + + rc = Tree.EntryStorage.Read(entryOffset, Entry); + if (rc.IsFailure()) return rc; + + // Note that we changed index. + EntryIndex = entryIndex; + return Result.Success; + } + + public Result Find(long virtualAddress) + { + Result rc; + + // Get the node. + BucketTreeNode node = Tree.NodeL1.GetNode(); + + if (virtualAddress >= node.GetEndOffset()) + return ResultFs.OutOfRange.Log(); + + int entrySetIndex; + + if (Tree.IsExistOffsetL2OnL1() && virtualAddress < node.GetBeginOffset()) + { + // The portion of the L2 offsets containing our target offset is stored in the L1 node + ReadOnlySpan offsets = node.GetArray().Slice(node.GetCount()); + int index = offsets.BinarySearch(virtualAddress); + if (index < 0) index = (~index) - 1; + + if (index < 0) + return ResultFs.OutOfRange.Log(); + + entrySetIndex = index; + } + else + { + ReadOnlySpan offsets = node.GetArray().Slice(0, node.GetCount()); + int index = offsets.BinarySearch(virtualAddress); + if (index < 0) index = (~index) - 1; + + if (index < 0) + return ResultFs.OutOfRange.Log(); + + if (Tree.IsExistL2()) + { + if (index >= Tree.OffsetCount) + return ResultFs.InvalidBucketTreeNodeOffset.Log(); + + rc = FindEntrySet(out entrySetIndex, virtualAddress, index); + if (rc.IsFailure()) return rc; + } + else + { + entrySetIndex = index; + } + } + + // Validate the entry set index. + if (entrySetIndex < 0 || entrySetIndex >= Tree.EntrySetCount) + return ResultFs.InvalidBucketTreeNodeOffset.Log(); + + // Find the entry. + rc = FindEntry(virtualAddress, entrySetIndex); + if (rc.IsFailure()) return rc; + + // Set count. + EntrySetCount = Tree.EntrySetCount; + return Result.Success; + } + + private Result FindEntrySet(out int entrySetIndex, long virtualAddress, int nodeIndex) + { + long nodeSize = Tree.NodeSize; + + using (var rented = new RentedArray((int)nodeSize)) + { + return FindEntrySetWithBuffer(out entrySetIndex, virtualAddress, nodeIndex, rented.Span); + } + } + + private Result FindEntrySetWithBuffer(out int outIndex, long virtualAddress, int nodeIndex, + Span buffer) + { + outIndex = default; + + // Calculate node extents. + long nodeSize = Tree.NodeSize; + long nodeOffset = (nodeIndex + 1) * nodeSize; + SubStorage2 storage = Tree.NodeStorage; + + // Read the node. + Result rc = storage.Read(nodeOffset, buffer.Slice(0, (int)nodeSize)); + if (rc.IsFailure()) return rc; + + // Validate the header. + NodeHeader header = MemoryMarshal.Cast(buffer)[0]; + rc = header.Verify(nodeIndex, nodeSize, sizeof(long)); + if (rc.IsFailure()) return rc; + + // Create the node and find. + var node = new StorageNode(sizeof(long), header.Count); + node.Find(buffer, virtualAddress); + + if (node.GetIndex() < 0) + return ResultFs.BucketTreeEntryNotFound.Log(); + + // Return the index. + outIndex = (int)Tree.GetEntrySetIndex(header.Index, node.GetIndex()); + return Result.Success; + } + + private Result FindEntry(long virtualAddress, int entrySetIndex) + { + long entrySetSize = Tree.NodeSize; + + using (var rented = new RentedArray((int)entrySetSize)) + { + return FindEntryWithBuffer(virtualAddress, entrySetIndex, rented.Span); + } + } + + private Result FindEntryWithBuffer(long virtualAddress, int entrySetIndex, Span buffer) + { + // Calculate entry set extents. + long entrySize = Tree.EntrySize; + long entrySetSize = Tree.NodeSize; + long entrySetOffset = entrySetIndex * entrySetSize; + SubStorage2 storage = Tree.EntryStorage; + + // Read the entry set. + Result rc = storage.Read(entrySetOffset, buffer.Slice(0, (int)entrySetSize)); + if (rc.IsFailure()) return rc; + + // Validate the entry set. + EntrySetHeader entrySet = MemoryMarshal.Cast(buffer)[0]; + rc = entrySet.Header.Verify(entrySetIndex, entrySetSize, entrySize); + if (rc.IsFailure()) return rc; + + // Create the node, and find. + var node = new StorageNode(entrySize, entrySet.Info.Count); + node.Find(buffer, virtualAddress); + + if (node.GetIndex() < 0) + return ResultFs.BucketTreeEntryNotFound.Log(); + + // Copy the data into entry. + int entryIndex = node.GetIndex(); + long entryOffset = GetBucketTreeEntryOffset(0, entrySize, entryIndex); + buffer.Slice((int)entryOffset, (int)entrySize).CopyTo(Entry); + + // Set our entry set/index. + _entrySet = entrySet; + EntryIndex = entryIndex; + + return Result.Success; + } + + private struct StorageNode + { + private Offset _start; + private int _count; + private int _index; + + public StorageNode(long size, int count) + { + _start = new Offset(NodeHeaderSize, (int)size); + _count = count; + _index = -1; + } + + public StorageNode(long offset, long size, int count) + { + _start = new Offset(offset + NodeHeaderSize, (int)size); + _count = count; + _index = -1; + } + + public int GetIndex() => _index; + + public void Find(ReadOnlySpan buffer, long virtualAddress) + { + int end = _count; + Offset pos = _start; + + while (end > 0) + { + int half = end / 2; + Offset mid = pos + half; + + long offset = BinaryPrimitives.ReadInt64LittleEndian(buffer.Slice((int)mid.Get())); + + if (offset <= virtualAddress) + { + pos = mid + 1; + end -= half + 1; + } + else + { + end = half; + } + } + + _index = (int)(pos - _start) - 1; + } + + private readonly struct Offset + { + private readonly long _offset; + private readonly int _stride; + + public Offset(long offset, int stride) + { + _offset = offset; + _stride = stride; + } + + public long Get() => _offset; + + public static Offset operator ++(Offset left) => left + 1; + public static Offset operator --(Offset left) => left - 1; + + public static Offset operator +(Offset left, long right) => new Offset(left._offset + right * left._stride, left._stride); + public static Offset operator -(Offset left, long right) => new Offset(left._offset - right * left._stride, left._stride); + + public static long operator -(Offset left, Offset right) => + (left._offset - right._offset) / left._stride; + + public static bool operator ==(Offset left, Offset right) => left._offset == right._offset; + public static bool operator !=(Offset left, Offset right) => left._offset != right._offset; + + public bool Equals(Offset other) => _offset == other._offset; + public override bool Equals(object obj) => obj is Offset other && Equals(other); + public override int GetHashCode() => _offset.GetHashCode(); + } + } + } + } +} diff --git a/src/LibHac/Util.cs b/src/LibHac/Util.cs index 7c226b5e..31c7ce7b 100644 --- a/src/LibHac/Util.cs +++ b/src/LibHac/Util.cs @@ -491,5 +491,27 @@ namespace LibHac return keyGeneration - 1; } + + public static bool IsPowerOfTwo(int value) + { + return value > 0 && ResetLeastSignificantOneBit(value) == 0; + } + + public static bool IsPowerOfTwo(long value) + { + return value > 0 && ResetLeastSignificantOneBit(value) == 0; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static int ResetLeastSignificantOneBit(int value) + { + return value & (value - 1); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static long ResetLeastSignificantOneBit(long value) + { + return value & (value - 1); + } } } From 9589f681a6a9874329864633c0f087257f0e4f2f Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Wed, 17 Jun 2020 20:56:40 -0700 Subject: [PATCH 2/8] Add a bucket tree builder --- build/CodeGen/results.csv | 2 +- src/LibHac/Common/SpanHelpers.cs | 29 +++ src/LibHac/Diag/Assert.cs | 9 + src/LibHac/Fs/ResultFs.cs | 2 +- src/LibHac/FsSystem/BucketTree2.cs | 27 ++- src/LibHac/FsSystem/BucketTreeBuilder.cs | 286 +++++++++++++++++++++++ 6 files changed, 345 insertions(+), 10 deletions(-) create mode 100644 src/LibHac/FsSystem/BucketTreeBuilder.cs diff --git a/build/CodeGen/results.csv b/build/CodeGen/results.csv index 1b9bd8e7..087dccac 100644 --- a/build/CodeGen/results.csv +++ b/build/CodeGen/results.csv @@ -60,7 +60,7 @@ Module,DescriptionStart,DescriptionEnd,Name,Summary 2,4036,,InvalidBucketTreeEntryOffset, 2,4037,,InvalidBucketTreeEntrySetOffset, 2,4038,,InvalidBucketTreeNodeIndex, -2,4039,,BucketTreeEntryNotFound, +2,4039,,InvalidBucketTreeVirtualOffset, 2,4241,4259,RomHostFileSystemCorrupted, 2,4242,,RomHostEntryCorrupted, diff --git a/src/LibHac/Common/SpanHelpers.cs b/src/LibHac/Common/SpanHelpers.cs index f63835d4..77ddbeac 100644 --- a/src/LibHac/Common/SpanHelpers.cs +++ b/src/LibHac/Common/SpanHelpers.cs @@ -57,5 +57,34 @@ namespace LibHac.Common { return CreateReadOnlySpan(ref Unsafe.As(ref reference), Unsafe.SizeOf()); } + + // All AsStruct methods do bounds checks on the input + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ref T AsStruct(Span span) where T : unmanaged + { + return ref MemoryMarshal.Cast(span)[0]; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ref readonly T AsReadOnlyStruct(ReadOnlySpan span) where T : unmanaged + { + return ref MemoryMarshal.Cast(span)[0]; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ref TTo AsStruct(Span span) + where TFrom : unmanaged + where TTo : unmanaged + { + return ref MemoryMarshal.Cast(span)[0]; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static ref readonly TTo AsStruct(ReadOnlySpan span) + where TFrom : unmanaged + where TTo : unmanaged + { + return ref MemoryMarshal.Cast(span)[0]; + } } } diff --git a/src/LibHac/Diag/Assert.cs b/src/LibHac/Diag/Assert.cs index 14aab339..84c98323 100644 --- a/src/LibHac/Diag/Assert.cs +++ b/src/LibHac/Diag/Assert.cs @@ -18,5 +18,14 @@ namespace LibHac.Diag throw new LibHacException($"Assertion failed: {message}"); } + + [Conditional("DEBUG")] + public static void NotNull([NotNull] T item) where T : class + { + if (!(item is null)) + { + throw new LibHacException("Not-null assertion failed."); + } + } } } diff --git a/src/LibHac/Fs/ResultFs.cs b/src/LibHac/Fs/ResultFs.cs index 179f181a..dd8ca748 100644 --- a/src/LibHac/Fs/ResultFs.cs +++ b/src/LibHac/Fs/ResultFs.cs @@ -134,7 +134,7 @@ namespace LibHac.Fs /// Error code: 2002-4038; Inner value: 0x1f8c02 public static Result.Base InvalidBucketTreeNodeIndex => new Result.Base(ModuleFs, 4038); /// Error code: 2002-4039; Inner value: 0x1f8e02 - public static Result.Base BucketTreeEntryNotFound => new Result.Base(ModuleFs, 4039); + public static Result.Base InvalidBucketTreeVirtualOffset => new Result.Base(ModuleFs, 4039); /// Error code: 2002-4241; Range: 4241-4259; Inner value: 0x212202 public static Result.Base RomHostFileSystemCorrupted { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => new Result.Base(ModuleFs, 4241, 4259); } diff --git a/src/LibHac/FsSystem/BucketTree2.cs b/src/LibHac/FsSystem/BucketTree2.cs index ace5ef7a..8e2ab8c2 100644 --- a/src/LibHac/FsSystem/BucketTree2.cs +++ b/src/LibHac/FsSystem/BucketTree2.cs @@ -8,7 +8,7 @@ using LibHac.Fs; namespace LibHac.FsSystem { - public class BucketTree2 + public partial class BucketTree2 { private const uint ExpectedMagic = 0x52544B42; // BKTR private const int MaxVersion = 1; @@ -103,6 +103,10 @@ namespace LibHac.FsSystem public bool IsInitialized() => NodeSize > 0; public bool IsEmpty() => EntrySize == 0; + public long GetStart() => StartOffset; + public long GetEnd() => EndOffset; + public long GetSize() => EndOffset - StartOffset; + public Result Find(ref Visitor visitor, long virtualAddress) { Assert.AssertTrue(IsInitialized()); @@ -165,7 +169,7 @@ namespace LibHac.FsSystem return Util.DivideByRoundUp(entryCount, entryCountPerNode); } - private static int GetNodeL2Count(long nodeSize, long entrySize, int entryCount) + public static int GetNodeL2Count(long nodeSize, long entrySize, int entryCount) { int offsetCountPerNode = GetOffsetCount(nodeSize); int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); @@ -312,9 +316,9 @@ namespace LibHac.FsSystem public readonly ref struct BucketTreeNode where TEntry : unmanaged { - private readonly ReadOnlySpan _buffer; + private readonly Span _buffer; - public BucketTreeNode(ReadOnlySpan buffer) + public BucketTreeNode(Span buffer) { _buffer = buffer; @@ -324,7 +328,8 @@ namespace LibHac.FsSystem public int GetCount() => GetHeader().Count; - public ReadOnlySpan GetArray() => GetArray(); + public ReadOnlySpan GetArray() => GetWritableArray(); + internal Span GetWritableArray() => GetWritableArray(); public long GetBeginOffset() => GetArray()[0]; public long GetEndOffset() => GetHeader().Offset; @@ -332,12 +337,18 @@ namespace LibHac.FsSystem [MethodImpl(MethodImplOptions.AggressiveInlining)] public ReadOnlySpan GetArray() where TElement : unmanaged + { + return GetWritableArray(); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private Span GetWritableArray() where TElement : unmanaged { return MemoryMarshal.Cast(_buffer.Slice(Unsafe.SizeOf())); } [MethodImpl(MethodImplOptions.AggressiveInlining)] - private ref NodeHeader GetHeader() + internal ref NodeHeader GetHeader() { return ref Unsafe.As(ref MemoryMarshal.GetReference(_buffer)); } @@ -594,7 +605,7 @@ namespace LibHac.FsSystem node.Find(buffer, virtualAddress); if (node.GetIndex() < 0) - return ResultFs.BucketTreeEntryNotFound.Log(); + return ResultFs.InvalidBucketTreeVirtualOffset.Log(); // Return the index. outIndex = (int)Tree.GetEntrySetIndex(header.Index, node.GetIndex()); @@ -633,7 +644,7 @@ namespace LibHac.FsSystem node.Find(buffer, virtualAddress); if (node.GetIndex() < 0) - return ResultFs.BucketTreeEntryNotFound.Log(); + return ResultFs.InvalidBucketTreeVirtualOffset.Log(); // Copy the data into entry. int entryIndex = node.GetIndex(); diff --git a/src/LibHac/FsSystem/BucketTreeBuilder.cs b/src/LibHac/FsSystem/BucketTreeBuilder.cs new file mode 100644 index 00000000..4b810fb4 --- /dev/null +++ b/src/LibHac/FsSystem/BucketTreeBuilder.cs @@ -0,0 +1,286 @@ +using System; +using System.Buffers.Binary; +using System.Runtime.CompilerServices; +using LibHac.Common; +using LibHac.Diag; +using LibHac.Fs; + +namespace LibHac.FsSystem +{ + public partial class BucketTree2 + { + public ref struct Builder + { + private Span NodeBuffer { get; set; } + private Span EntryBuffer { get; set; } + + private int NodeSize { get; set; } + private int EntrySize { get; set; } + private int EntryCount { get; set; } + private int EntriesPerEntrySet { get; set; } + private int OffsetsPerNode { get; set; } + + private int CurrentL2OffsetIndex { get; set; } + private int CurrentEntryIndex { get; set; } + private long CurrentOffset { get; set; } + + /// + /// Initializes the bucket tree builder. + /// + /// The buffer for the tree's header. Must be at least the size in bytes returned by . + /// The buffer for the tree's nodes. Must be at least the size in bytes returned by . + /// The buffer for the tree's entries. Must be at least the size in bytes returned by . + /// The size of each node in the bucket tree. + /// The size of each entry that will be stored in the bucket tree. + /// The exact number of entries that will be added to the bucket tree. + /// The of the operation. + public Result Initialize(Span headerBuffer, Span nodeBuffer, Span entryBuffer, + int nodeSize, int entrySize, int entryCount) + { + Assert.AssertTrue(entrySize >= sizeof(long)); + Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); + Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); + Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + + NodeSize = nodeSize; + EntrySize = entrySize; + EntryCount = entryCount; + + EntriesPerEntrySet = GetEntryCount(nodeSize, entrySize); + OffsetsPerNode = GetOffsetCount(nodeSize); + CurrentL2OffsetIndex = GetNodeL2Count(nodeSize, entrySize, entryCount); + + // Verify the provided buffers are large enough + int nodeStorageSize = (int)QueryNodeStorageSize(nodeSize, entrySize, entryCount); + int entryStorageSize = (int)QueryEntryStorageSize(nodeSize, entrySize, entryCount); + + if (headerBuffer.Length < QueryHeaderStorageSize() || + nodeBuffer.Length < nodeStorageSize || + entryBuffer.Length < entryStorageSize) + { + return ResultFs.InvalidSize.Log(); + } + + // Set and clear the buffers + NodeBuffer = nodeBuffer.Slice(0, nodeStorageSize); + EntryBuffer = entryBuffer.Slice(0, entryStorageSize); + + nodeBuffer.Clear(); + entryBuffer.Clear(); + + // Format the tree header + ref Header header = ref SpanHelpers.AsStruct
(headerBuffer); + header.Format(entryCount); + + // Set the initial position + CurrentEntryIndex = 0; + CurrentOffset = -1; + + return Result.Success; + } + + /// + /// Adds a new entry to the bucket tree. + /// + /// The type of the entry to add. Added entries should all be the same type. + /// The entry to add. + /// The of the operation. + public Result Add(ref T entry) where T : unmanaged + { + Assert.AssertTrue(Unsafe.SizeOf() == EntrySize); + + if (CurrentEntryIndex >= EntryCount) + return ResultFs.OutOfRange.Log(); + + long entryOffset = BinaryPrimitives.ReadInt64LittleEndian(SpanHelpers.AsByteSpan(ref entry)); + + if (entryOffset <= CurrentOffset) + return ResultFs.InvalidOffset.Log(); + + FinalizePreviousEntrySet(entryOffset); + AddEntryOffset(entryOffset); + + int entrySetIndex = CurrentEntryIndex / EntriesPerEntrySet; + int indexInEntrySet = CurrentEntryIndex % EntriesPerEntrySet; + + GetEntrySet(entrySetIndex).GetWritableArray()[indexInEntrySet] = entry; + + CurrentOffset = entryOffset; + CurrentEntryIndex++; + + return Result.Success; + } + + /// + /// Checks if a new entry set is being started. If so, sets the end offset of the previous entry set. + /// + /// The end offset of the previous entry. + private void FinalizePreviousEntrySet(long endOffset) + { + int prevEntrySetIndex = CurrentEntryIndex / EntriesPerEntrySet - 1; + int indexInEntrySet = CurrentEntryIndex % EntriesPerEntrySet; + + // If the previous Add finished an entry set + if (CurrentEntryIndex > 0 && indexInEntrySet == 0) + { + // Set the end offset of that entry set + BucketTreeNode prevEntrySet = GetEntrySet(prevEntrySetIndex); + + prevEntrySet.GetHeader().Index = prevEntrySetIndex; + prevEntrySet.GetHeader().Count = EntriesPerEntrySet; + prevEntrySet.GetHeader().Offset = endOffset; + + // Check if we're writing in L2 nodes + if (CurrentL2OffsetIndex > OffsetsPerNode) + { + int prevL2NodeIndex = CurrentL2OffsetIndex / OffsetsPerNode - 2; + int indexInL2Node = CurrentL2OffsetIndex % OffsetsPerNode; + + // If the previous Add finished an L2 node + if (indexInL2Node == 0) + { + // Set the end offset of that node + BucketTreeNode prevL2Node = GetL2Node(prevL2NodeIndex); + + prevL2Node.GetHeader().Index = prevL2NodeIndex; + prevL2Node.GetHeader().Count = OffsetsPerNode; + prevL2Node.GetHeader().Offset = endOffset; + } + } + } + } + + /// + /// If needed, adds a new entry set's start offset to the L1 or L2 nodes. + /// + /// The start offset of the entry being added. + private void AddEntryOffset(long entryOffset) + { + int entrySetIndex = CurrentEntryIndex / EntriesPerEntrySet; + int indexInEntrySet = CurrentEntryIndex % EntriesPerEntrySet; + + // If we're starting a new entry set we need to add its start offset to the L1/L2 nodes + if (indexInEntrySet == 0) + { + BucketTreeNode l1Node = GetL1Node(); + + if (CurrentL2OffsetIndex == 0) + { + // There are no L2 nodes. Write the entry set end offset directly to L1 + l1Node.GetWritableArray()[entrySetIndex] = entryOffset; + } + else + { + if (CurrentL2OffsetIndex < OffsetsPerNode) + { + // The current L2 offset is stored in the L1 node + l1Node.GetWritableArray()[CurrentL2OffsetIndex] = entryOffset; + } + else + { + // Write the entry set offset to the current L2 node + int l2NodeIndex = CurrentL2OffsetIndex / OffsetsPerNode; + int indexInL2Node = CurrentL2OffsetIndex % OffsetsPerNode; + + BucketTreeNode l2Node = GetL2Node(l2NodeIndex - 1); + + l2Node.GetWritableArray()[indexInL2Node] = entryOffset; + + // If we're starting a new L2 node we need to add its start offset to the L1 node + if (indexInL2Node == 0) + { + l1Node.GetWritableArray()[l2NodeIndex - 1] = entryOffset; + } + } + + CurrentL2OffsetIndex++; + } + } + } + + /// + /// Finalizes the bucket tree. Must be called after all entries are added. + /// + /// The end offset of the bucket tree. + /// The of the operation. + public Result Finalize(long endOffset) + { + // Finalize must only be called after all entries are added + if (EntryCount != CurrentEntryIndex) + return ResultFs.OutOfRange.Log(); + + if (endOffset <= CurrentOffset) + return ResultFs.InvalidOffset.Log(); + + FinalizePreviousEntrySet(endOffset); + + int entrySetIndex = CurrentEntryIndex / EntriesPerEntrySet; + int indexInEntrySet = CurrentEntryIndex % EntriesPerEntrySet; + + // Finalize the current entry set if needed + if (indexInEntrySet != 0) + { + ref NodeHeader entrySetHeader = ref GetEntrySetHeader(entrySetIndex); + + entrySetHeader.Index = entrySetIndex; + entrySetHeader.Count = indexInEntrySet; + entrySetHeader.Offset = endOffset; + } + + int l2NodeIndex = Util.DivideByRoundUp(CurrentL2OffsetIndex, OffsetsPerNode) - 2; + int indexInL2Node = CurrentL2OffsetIndex % OffsetsPerNode; + + // Finalize the current L2 node if needed + if (CurrentL2OffsetIndex > OffsetsPerNode && (indexInEntrySet != 0 || indexInL2Node != 0)) + { + ref NodeHeader l2NodeHeader = ref GetL2Node(l2NodeIndex).GetHeader(); + l2NodeHeader.Index = l2NodeIndex; + l2NodeHeader.Count = indexInL2Node != 0 ? indexInL2Node : OffsetsPerNode; + l2NodeHeader.Offset = endOffset; + } + + // Finalize the L1 node + ref NodeHeader l1Header = ref GetL1Node().GetHeader(); + + l1Header.Index = 0; + l1Header.Offset = endOffset; + + // L1 count depends on the existence or absence of L2 nodes + if (CurrentL2OffsetIndex == 0) + { + l1Header.Count = Util.DivideByRoundUp(CurrentEntryIndex, EntriesPerEntrySet); + } + else + { + l1Header.Count = l2NodeIndex + 1; + } + + return Result.Success; + } + + private ref NodeHeader GetEntrySetHeader(int index) + { + BucketTreeNode entrySetNode = GetEntrySet(index); + return ref entrySetNode.GetHeader(); + } + + private BucketTreeNode GetEntrySet(int index) where T : unmanaged + { + Span entrySetBuffer = EntryBuffer.Slice(NodeSize * index, NodeSize); + return new BucketTreeNode(entrySetBuffer); + } + + private BucketTreeNode GetL1Node() + { + Span l1NodeBuffer = NodeBuffer.Slice(0, NodeSize); + return new BucketTreeNode(l1NodeBuffer); + } + + private BucketTreeNode GetL2Node(int index) + { + Span l2NodeBuffer = NodeBuffer.Slice(NodeSize * (index + 1), NodeSize); + return new BucketTreeNode(l2NodeBuffer); + } + } + } +} From 4b4b354a7ee073584e2fb990820100faf30232ce Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Thu, 18 Jun 2020 15:26:32 -0700 Subject: [PATCH 3/8] Use IStorage for the bucket tree builder instead of Spans --- src/LibHac/FsSystem/BucketTree2.cs | 2 +- src/LibHac/FsSystem/BucketTreeBuilder.cs | 181 +++++++++++++---------- 2 files changed, 103 insertions(+), 80 deletions(-) diff --git a/src/LibHac/FsSystem/BucketTree2.cs b/src/LibHac/FsSystem/BucketTree2.cs index 8e2ab8c2..87337c64 100644 --- a/src/LibHac/FsSystem/BucketTree2.cs +++ b/src/LibHac/FsSystem/BucketTree2.cs @@ -257,7 +257,7 @@ namespace LibHac.FsSystem } } - private class NodeBuffer + private struct NodeBuffer { // Use long to ensure alignment private long[] _header; diff --git a/src/LibHac/FsSystem/BucketTreeBuilder.cs b/src/LibHac/FsSystem/BucketTreeBuilder.cs index 4b810fb4..309b08e4 100644 --- a/src/LibHac/FsSystem/BucketTreeBuilder.cs +++ b/src/LibHac/FsSystem/BucketTreeBuilder.cs @@ -9,10 +9,14 @@ namespace LibHac.FsSystem { public partial class BucketTree2 { - public ref struct Builder + public class Builder { - private Span NodeBuffer { get; set; } - private Span EntryBuffer { get; set; } + private SubStorage2 NodeStorage { get; set; } + private SubStorage2 EntryStorage { get; set; } + + private NodeBuffer _l1Node = new NodeBuffer(); + private NodeBuffer _l2Node = new NodeBuffer(); + private NodeBuffer _entrySet = new NodeBuffer(); private int NodeSize { get; set; } private int EntrySize { get; set; } @@ -22,19 +26,19 @@ namespace LibHac.FsSystem private int CurrentL2OffsetIndex { get; set; } private int CurrentEntryIndex { get; set; } - private long CurrentOffset { get; set; } + private long CurrentOffset { get; set; } = -1; /// /// Initializes the bucket tree builder. /// - /// The buffer for the tree's header. Must be at least the size in bytes returned by . - /// The buffer for the tree's nodes. Must be at least the size in bytes returned by . - /// The buffer for the tree's entries. Must be at least the size in bytes returned by . - /// The size of each node in the bucket tree. + /// The the tree's header will be written to.Must be at least the size in bytes returned by . + /// The the tree's nodes will be written to. Must be at least the size in bytes returned by . + /// The the tree's entries will be written to. Must be at least the size in bytes returned by . + /// The size of each node in the bucket tree. Must be a power of 2. /// The size of each entry that will be stored in the bucket tree. /// The exact number of entries that will be added to the bucket tree. /// The of the operation. - public Result Initialize(Span headerBuffer, Span nodeBuffer, Span entryBuffer, + public Result Initialize(SubStorage2 headerStorage, SubStorage2 nodeStorage, SubStorage2 entryStorage, int nodeSize, int entrySize, int entryCount) { Assert.AssertTrue(entrySize >= sizeof(long)); @@ -42,6 +46,10 @@ namespace LibHac.FsSystem Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + if (headerStorage is null || nodeStorage is null || entryStorage is null) + return ResultFs.NullptrArgument.Log(); + + // Set the builder parameters NodeSize = nodeSize; EntrySize = entrySize; EntryCount = entryCount; @@ -50,27 +58,30 @@ namespace LibHac.FsSystem OffsetsPerNode = GetOffsetCount(nodeSize); CurrentL2OffsetIndex = GetNodeL2Count(nodeSize, entrySize, entryCount); - // Verify the provided buffers are large enough - int nodeStorageSize = (int)QueryNodeStorageSize(nodeSize, entrySize, entryCount); - int entryStorageSize = (int)QueryEntryStorageSize(nodeSize, entrySize, entryCount); + // Create and write the header + var header = new Header(); + header.Format(entryCount); + Result rc = headerStorage.Write(0, SpanHelpers.AsByteSpan(ref header)); + if (rc.IsFailure()) return rc; - if (headerBuffer.Length < QueryHeaderStorageSize() || - nodeBuffer.Length < nodeStorageSize || - entryBuffer.Length < entryStorageSize) + // Allocate buffers for the L1 node and entry sets + _l1Node.Allocate(nodeSize); + _entrySet.Allocate(nodeSize); + + int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); + + // Allocate an L2 node buffer if there are more entry sets than will fit in the L1 node + if (OffsetsPerNode < entrySetCount) { - return ResultFs.InvalidSize.Log(); + _l2Node.Allocate(nodeSize); } - // Set and clear the buffers - NodeBuffer = nodeBuffer.Slice(0, nodeStorageSize); - EntryBuffer = entryBuffer.Slice(0, entryStorageSize); + _l1Node.FillZero(); + _l2Node.FillZero(); + _entrySet.FillZero(); - nodeBuffer.Clear(); - entryBuffer.Clear(); - - // Format the tree header - ref Header header = ref SpanHelpers.AsStruct
(headerBuffer); - header.Format(entryCount); + NodeStorage = nodeStorage; + EntryStorage = entryStorage; // Set the initial position CurrentEntryIndex = 0; @@ -92,18 +103,20 @@ namespace LibHac.FsSystem if (CurrentEntryIndex >= EntryCount) return ResultFs.OutOfRange.Log(); + // The entry offset must always be the first 8 bytes of the struct long entryOffset = BinaryPrimitives.ReadInt64LittleEndian(SpanHelpers.AsByteSpan(ref entry)); if (entryOffset <= CurrentOffset) return ResultFs.InvalidOffset.Log(); - FinalizePreviousEntrySet(entryOffset); + Result rc = FinalizePreviousEntrySet(entryOffset); + if (rc.IsFailure()) return rc; + AddEntryOffset(entryOffset); - int entrySetIndex = CurrentEntryIndex / EntriesPerEntrySet; + // Write the new entry int indexInEntrySet = CurrentEntryIndex % EntriesPerEntrySet; - - GetEntrySet(entrySetIndex).GetWritableArray()[indexInEntrySet] = entry; + _entrySet.GetNode().GetWritableArray()[indexInEntrySet] = entry; CurrentOffset = entryOffset; CurrentEntryIndex++; @@ -112,10 +125,12 @@ namespace LibHac.FsSystem } /// - /// Checks if a new entry set is being started. If so, sets the end offset of the previous entry set. + /// Checks if a new entry set is being started. If so, sets the end offset of the previous + /// entry set and writes it to the output storage. /// /// The end offset of the previous entry. - private void FinalizePreviousEntrySet(long endOffset) + /// The of the operation. + private Result FinalizePreviousEntrySet(long endOffset) { int prevEntrySetIndex = CurrentEntryIndex / EntriesPerEntrySet - 1; int indexInEntrySet = CurrentEntryIndex % EntriesPerEntrySet; @@ -124,11 +139,19 @@ namespace LibHac.FsSystem if (CurrentEntryIndex > 0 && indexInEntrySet == 0) { // Set the end offset of that entry set - BucketTreeNode prevEntrySet = GetEntrySet(prevEntrySetIndex); + ref NodeHeader entrySetHeader = ref _entrySet.GetHeader(); - prevEntrySet.GetHeader().Index = prevEntrySetIndex; - prevEntrySet.GetHeader().Count = EntriesPerEntrySet; - prevEntrySet.GetHeader().Offset = endOffset; + entrySetHeader.Index = prevEntrySetIndex; + entrySetHeader.Count = EntriesPerEntrySet; + entrySetHeader.Offset = endOffset; + + // Write the entry set to the entry storage + long storageOffset = (long)NodeSize * prevEntrySetIndex; + Result rc = EntryStorage.Write(storageOffset, _entrySet.GetBuffer()); + if (rc.IsFailure()) return rc; + + // Clear the entry set buffer to begin the new entry set + _entrySet.FillZero(); // Check if we're writing in L2 nodes if (CurrentL2OffsetIndex > OffsetsPerNode) @@ -140,14 +163,24 @@ namespace LibHac.FsSystem if (indexInL2Node == 0) { // Set the end offset of that node - BucketTreeNode prevL2Node = GetL2Node(prevL2NodeIndex); + ref NodeHeader l2NodeHeader = ref _l2Node.GetHeader(); - prevL2Node.GetHeader().Index = prevL2NodeIndex; - prevL2Node.GetHeader().Count = OffsetsPerNode; - prevL2Node.GetHeader().Offset = endOffset; + l2NodeHeader.Index = prevL2NodeIndex; + l2NodeHeader.Count = OffsetsPerNode; + l2NodeHeader.Offset = endOffset; + + // Write the L2 node to the node storage + long nodeOffset = (long)NodeSize * (prevL2NodeIndex + 1); + rc = NodeStorage.Write(nodeOffset, _l2Node.GetBuffer()); + if (rc.IsFailure()) return rc; + + // Clear the L2 node buffer to begin the new node + _l2Node.FillZero(); } } } + + return Result.Success; } /// @@ -162,19 +195,19 @@ namespace LibHac.FsSystem // If we're starting a new entry set we need to add its start offset to the L1/L2 nodes if (indexInEntrySet == 0) { - BucketTreeNode l1Node = GetL1Node(); + Span l1Data = _l1Node.GetNode().GetWritableArray(); if (CurrentL2OffsetIndex == 0) { // There are no L2 nodes. Write the entry set end offset directly to L1 - l1Node.GetWritableArray()[entrySetIndex] = entryOffset; + l1Data[entrySetIndex] = entryOffset; } else { if (CurrentL2OffsetIndex < OffsetsPerNode) { // The current L2 offset is stored in the L1 node - l1Node.GetWritableArray()[CurrentL2OffsetIndex] = entryOffset; + l1Data[CurrentL2OffsetIndex] = entryOffset; } else { @@ -182,14 +215,13 @@ namespace LibHac.FsSystem int l2NodeIndex = CurrentL2OffsetIndex / OffsetsPerNode; int indexInL2Node = CurrentL2OffsetIndex % OffsetsPerNode; - BucketTreeNode l2Node = GetL2Node(l2NodeIndex - 1); - - l2Node.GetWritableArray()[indexInL2Node] = entryOffset; + Span l2Data = _l2Node.GetNode().GetWritableArray(); + l2Data[indexInL2Node] = entryOffset; // If we're starting a new L2 node we need to add its start offset to the L1 node if (indexInL2Node == 0) { - l1Node.GetWritableArray()[l2NodeIndex - 1] = entryOffset; + l1Data[l2NodeIndex - 1] = entryOffset; } } @@ -212,7 +244,11 @@ namespace LibHac.FsSystem if (endOffset <= CurrentOffset) return ResultFs.InvalidOffset.Log(); - FinalizePreviousEntrySet(endOffset); + if (CurrentOffset == -1) + return Result.Success; + + Result rc = FinalizePreviousEntrySet(endOffset); + if (rc.IsFailure()) return rc; int entrySetIndex = CurrentEntryIndex / EntriesPerEntrySet; int indexInEntrySet = CurrentEntryIndex % EntriesPerEntrySet; @@ -220,11 +256,15 @@ namespace LibHac.FsSystem // Finalize the current entry set if needed if (indexInEntrySet != 0) { - ref NodeHeader entrySetHeader = ref GetEntrySetHeader(entrySetIndex); + ref NodeHeader entrySetHeader = ref _entrySet.GetHeader(); entrySetHeader.Index = entrySetIndex; entrySetHeader.Count = indexInEntrySet; entrySetHeader.Offset = endOffset; + + long entryStorageOffset = (long)NodeSize * entrySetIndex; + rc = EntryStorage.Write(entryStorageOffset, _entrySet.GetBuffer()); + if (rc.IsFailure()) return rc; } int l2NodeIndex = Util.DivideByRoundUp(CurrentL2OffsetIndex, OffsetsPerNode) - 2; @@ -233,54 +273,37 @@ namespace LibHac.FsSystem // Finalize the current L2 node if needed if (CurrentL2OffsetIndex > OffsetsPerNode && (indexInEntrySet != 0 || indexInL2Node != 0)) { - ref NodeHeader l2NodeHeader = ref GetL2Node(l2NodeIndex).GetHeader(); + ref NodeHeader l2NodeHeader = ref _l2Node.GetHeader(); l2NodeHeader.Index = l2NodeIndex; l2NodeHeader.Count = indexInL2Node != 0 ? indexInL2Node : OffsetsPerNode; l2NodeHeader.Offset = endOffset; + + long l2NodeStorageOffset = NodeSize * (l2NodeIndex + 1); + rc = NodeStorage.Write(l2NodeStorageOffset, _l2Node.GetBuffer()); + if (rc.IsFailure()) return rc; } // Finalize the L1 node - ref NodeHeader l1Header = ref GetL1Node().GetHeader(); - - l1Header.Index = 0; - l1Header.Offset = endOffset; + ref NodeHeader l1NodeHeader = ref _l1Node.GetHeader(); + l1NodeHeader.Index = 0; + l1NodeHeader.Offset = endOffset; // L1 count depends on the existence or absence of L2 nodes if (CurrentL2OffsetIndex == 0) { - l1Header.Count = Util.DivideByRoundUp(CurrentEntryIndex, EntriesPerEntrySet); + l1NodeHeader.Count = Util.DivideByRoundUp(CurrentEntryIndex, EntriesPerEntrySet); } else { - l1Header.Count = l2NodeIndex + 1; + l1NodeHeader.Count = l2NodeIndex + 1; } + rc = NodeStorage.Write(0, _l1Node.GetBuffer()); + if (rc.IsFailure()) return rc; + + CurrentOffset = long.MaxValue; return Result.Success; } - - private ref NodeHeader GetEntrySetHeader(int index) - { - BucketTreeNode entrySetNode = GetEntrySet(index); - return ref entrySetNode.GetHeader(); - } - - private BucketTreeNode GetEntrySet(int index) where T : unmanaged - { - Span entrySetBuffer = EntryBuffer.Slice(NodeSize * index, NodeSize); - return new BucketTreeNode(entrySetBuffer); - } - - private BucketTreeNode GetL1Node() - { - Span l1NodeBuffer = NodeBuffer.Slice(0, NodeSize); - return new BucketTreeNode(l1NodeBuffer); - } - - private BucketTreeNode GetL2Node(int index) - { - Span l2NodeBuffer = NodeBuffer.Slice(NodeSize * (index + 1), NodeSize); - return new BucketTreeNode(l2NodeBuffer); - } } } } From c2247e583f4e9d8e25d7ac6f6ac21ba2c663b234 Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Thu, 18 Jun 2020 15:42:11 -0700 Subject: [PATCH 4/8] Address compiler warnings --- src/LibHac/FsSystem/BucketTree2.cs | 42 +++++++++--------------------- 1 file changed, 12 insertions(+), 30 deletions(-) diff --git a/src/LibHac/FsSystem/BucketTree2.cs b/src/LibHac/FsSystem/BucketTree2.cs index 87337c64..e65bcff8 100644 --- a/src/LibHac/FsSystem/BucketTree2.cs +++ b/src/LibHac/FsSystem/BucketTree2.cs @@ -21,11 +21,10 @@ namespace LibHac.FsSystem private SubStorage2 NodeStorage { get; set; } private SubStorage2 EntryStorage { get; set; } - private NodeBuffer NodeL1 { get; } = new NodeBuffer(); + private NodeBuffer _nodeL1 = new NodeBuffer(); private long NodeSize { get; set; } private long EntrySize { get; set; } - private int EntryCount { get; set; } private int OffsetCount { get; set; } private int EntrySetCount { get; set; } private long StartOffset { get; set; } @@ -45,24 +44,24 @@ namespace LibHac.FsSystem return ResultFs.InvalidArgument.Log(); // Allocate node. - if (!NodeL1.Allocate(nodeSize)) + if (!_nodeL1.Allocate(nodeSize)) return ResultFs.BufferAllocationFailed.Log(); bool needFree = true; try { // Read node. - Result rc = nodeStorage.Read(0, NodeL1.GetBuffer()); + Result rc = nodeStorage.Read(0, _nodeL1.GetBuffer()); if (rc.IsFailure()) return rc; // Verify node. - rc = NodeL1.GetHeader().Verify(0, nodeSize, sizeof(long)); + rc = _nodeL1.GetHeader().Verify(0, nodeSize, sizeof(long)); if (rc.IsFailure()) return rc; // Validate offsets. int offsetCount = GetOffsetCount(nodeSize); int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); - BucketTreeNode node = NodeL1.GetNode(); + BucketTreeNode node = _nodeL1.GetNode(); long startOffset; if (offsetCount < entrySetCount && node.GetCount() < offsetCount) @@ -83,7 +82,6 @@ namespace LibHac.FsSystem EntryStorage = entryStorage; NodeSize = nodeSize; EntrySize = entrySize; - EntryCount = entryCount; OffsetCount = offsetCount; EntrySetCount = entrySetCount; StartOffset = startOffset; @@ -96,7 +94,7 @@ namespace LibHac.FsSystem finally { if (needFree) - NodeL1.Free(); + _nodeL1.Free(); } } @@ -194,11 +192,11 @@ namespace LibHac.FsSystem } private bool IsExistL2() => OffsetCount < EntrySetCount; - private bool IsExistOffsetL2OnL1() => IsExistL2() && NodeL1.GetHeader().Count < OffsetCount; + private bool IsExistOffsetL2OnL1() => IsExistL2() && _nodeL1.GetHeader().Count < OffsetCount; private long GetEntrySetIndex(int nodeIndex, int offsetIndex) { - return (OffsetCount - NodeL1.GetHeader().Count) + (OffsetCount * nodeIndex) + offsetIndex; + return (OffsetCount - _nodeL1.GetHeader().Count) + (OffsetCount * nodeIndex) + offsetIndex; } public struct Header @@ -206,7 +204,9 @@ namespace LibHac.FsSystem public uint Magic; public uint Version; public int EntryCount; +#pragma warning disable 414 private int _reserved; +#pragma warning restore 414 public void Format(int entryCount) { @@ -300,18 +300,6 @@ namespace LibHac.FsSystem { return new BucketTreeNode(GetBuffer()); } - - public ref T Get() where T : unmanaged - { - if (Unsafe.SizeOf() != Unsafe.SizeOf()) - { - throw new InvalidOperationException(); - } - - Assert.AssertTrue(_header.Length / sizeof(long) >= Unsafe.SizeOf()); - - return ref Unsafe.As(ref _header[0]); - } } public readonly ref struct BucketTreeNode where TEntry : unmanaged @@ -365,6 +353,7 @@ namespace LibHac.FsSystem [StructLayout(LayoutKind.Explicit)] private struct EntrySetHeader { + // ReSharper disable once MemberHidesStaticFromOuterClass [FieldOffset(0)] public NodeHeader Header; [FieldOffset(0)] public EntrySetInfo Info; @@ -516,7 +505,7 @@ namespace LibHac.FsSystem Result rc; // Get the node. - BucketTreeNode node = Tree.NodeL1.GetNode(); + BucketTreeNode node = Tree._nodeL1.GetNode(); if (virtualAddress >= node.GetEndOffset()) return ResultFs.OutOfRange.Log(); @@ -671,13 +660,6 @@ namespace LibHac.FsSystem _index = -1; } - public StorageNode(long offset, long size, int count) - { - _start = new Offset(offset + NodeHeaderSize, (int)size); - _count = count; - _index = -1; - } - public int GetIndex() => _index; public void Find(ReadOnlySpan buffer, long virtualAddress) From 33af34cefce165606bfd6101931434892e9fb438 Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Thu, 25 Jun 2020 14:01:45 -0700 Subject: [PATCH 5/8] 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) From d3abdeacc42692844ee809082c9b2eda511548d4 Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Thu, 25 Jun 2020 19:03:51 -0700 Subject: [PATCH 6/8] Use the new BucketTree in Aes128CtrExStorage --- src/LibHac/FsSystem/Aes128CtrExStorage.cs | 74 +++++++++++++---------- src/LibHac/FsSystem/NcaUtils/Nca.cs | 17 +++++- 2 files changed, 56 insertions(+), 35 deletions(-) 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); } From 33b414a15c046d88595f144eb6b6f6da396783ae Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Thu, 25 Jun 2020 20:50:54 -0700 Subject: [PATCH 7/8] Remove old BucketTree --- src/LibHac/FsSystem/Aes128CtrExStorage.cs | 4 +- src/LibHac/FsSystem/BucketTree.cs | 783 +++++++++++++++++++--- src/LibHac/FsSystem/BucketTree2.cs | 738 -------------------- src/LibHac/FsSystem/BucketTreeBuilder.cs | 2 +- src/LibHac/FsSystem/IndirectStorage.cs | 16 +- src/LibHac/FsSystem/NcaUtils/Nca.cs | 4 +- 6 files changed, 719 insertions(+), 828 deletions(-) delete mode 100644 src/LibHac/FsSystem/BucketTree2.cs diff --git a/src/LibHac/FsSystem/Aes128CtrExStorage.cs b/src/LibHac/FsSystem/Aes128CtrExStorage.cs index 42446e43..79f305b1 100644 --- a/src/LibHac/FsSystem/Aes128CtrExStorage.cs +++ b/src/LibHac/FsSystem/Aes128CtrExStorage.cs @@ -9,7 +9,7 @@ namespace LibHac.FsSystem { public static readonly int NodeSize = 1024 * 16; - private BucketTree2 Table { get; } = new BucketTree2(); + private BucketTree Table { get; } = new BucketTree(); private readonly object _locker = new object(); @@ -34,7 +34,7 @@ namespace LibHac.FsSystem if (destination.Length == 0) return Result.Success; - var visitor = new BucketTree2.Visitor(); + var visitor = new BucketTree.Visitor(); Result rc = Table.Find(ref visitor, offset); if (rc.IsFailure()) return rc; diff --git a/src/LibHac/FsSystem/BucketTree.cs b/src/LibHac/FsSystem/BucketTree.cs index 85f6608e..a9cb125b 100644 --- a/src/LibHac/FsSystem/BucketTree.cs +++ b/src/LibHac/FsSystem/BucketTree.cs @@ -1,109 +1,738 @@ -using System.Collections.Generic; -using System.IO; -using System.Linq; +using System; +using System.Buffers.Binary; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using LibHac.Common; +using LibHac.Diag; using LibHac.Fs; namespace LibHac.FsSystem { - public class BucketTree where T : BucketTreeEntry, new() + public partial class BucketTree { - private const int BucketAlignment = 0x4000; - public BucketTreeBucket BucketOffsets { get; } - public BucketTreeBucket[] Buckets { get; } + private const uint ExpectedMagic = 0x52544B42; // BKTR + private const int MaxVersion = 1; - public BucketTree(IStorage data) + private const int NodeSizeMin = 1024; + private const int NodeSizeMax = 1024 * 512; + + private static int NodeHeaderSize => Unsafe.SizeOf(); + + private SubStorage2 NodeStorage { get; set; } + private SubStorage2 EntryStorage { get; set; } + + private NodeBuffer _nodeL1 = new NodeBuffer(); + + private long NodeSize { get; set; } + private long EntrySize { get; set; } + private int OffsetCount { get; set; } + private int EntrySetCount { get; set; } + private long StartOffset { get; set; } + private long EndOffset { get; set; } + + public Result Initialize(SubStorage2 nodeStorage, SubStorage2 entryStorage, int nodeSize, int entrySize, + int entryCount) { - var reader = new BinaryReader(data.AsStream()); + Assert.AssertTrue(entrySize >= sizeof(long)); + Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); + Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); + Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + Assert.AssertTrue(!IsInitialized()); - BucketOffsets = new BucketTreeBucket(reader); + // Ensure valid entry count. + if (entryCount <= 0) + return ResultFs.InvalidArgument.Log(); - Buckets = new BucketTreeBucket[BucketOffsets.EntryCount]; + // Allocate node. + if (!_nodeL1.Allocate(nodeSize)) + return ResultFs.BufferAllocationFailed.Log(); - for (int i = 0; i < BucketOffsets.EntryCount; i++) + bool needFree = true; + try { - reader.BaseStream.Position = (i + 1) * BucketAlignment; - Buckets[i] = new BucketTreeBucket(reader); + // Read node. + Result rc = nodeStorage.Read(0, _nodeL1.GetBuffer()); + if (rc.IsFailure()) return rc; + + // Verify node. + rc = _nodeL1.GetHeader().Verify(0, nodeSize, sizeof(long)); + if (rc.IsFailure()) return rc; + + // Validate offsets. + int offsetCount = GetOffsetCount(nodeSize); + int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); + BucketTreeNode node = _nodeL1.GetNode(); + + long startOffset; + if (offsetCount < entrySetCount && node.GetCount() < offsetCount) + { + startOffset = node.GetL2BeginOffset(); + } + else + { + startOffset = node.GetBeginOffset(); + } + + long endOffset = node.GetEndOffset(); + + if (startOffset < 0 || startOffset > node.GetBeginOffset() || startOffset >= endOffset) + return ResultFs.InvalidBucketTreeEntryOffset.Log(); + + NodeStorage = nodeStorage; + EntryStorage = entryStorage; + NodeSize = nodeSize; + EntrySize = entrySize; + OffsetCount = offsetCount; + EntrySetCount = entrySetCount; + StartOffset = startOffset; + EndOffset = endOffset; + + needFree = false; + + return Result.Success; + } + finally + { + if (needFree) + _nodeL1.Free(); } } - public List GetEntryList() - { - List list = Buckets.SelectMany(x => x.Entries).ToList(); + public bool IsInitialized() => NodeSize > 0; + public bool IsEmpty() => EntrySize == 0; - for (int i = 0; i < list.Count - 1; i++) + public long GetStart() => StartOffset; + 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()); + + if (virtualAddress < 0) + return ResultFs.InvalidOffset.Log(); + + if (IsEmpty()) + return ResultFs.OutOfRange.Log(); + + Result rc = visitor.Initialize(this); + if (rc.IsFailure()) return rc; + + return visitor.Find(virtualAddress); + } + + public static int QueryHeaderStorageSize() => Unsafe.SizeOf
(); + + public static long QueryNodeStorageSize(long nodeSize, long entrySize, int entryCount) + { + Assert.AssertTrue(entrySize >= sizeof(long)); + Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); + Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); + Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + Assert.AssertTrue(entryCount >= 0); + + if (entryCount <= 0) + return 0; + + return (1 + GetNodeL2Count(nodeSize, entrySize, entryCount)) * nodeSize; + } + + public static long QueryEntryStorageSize(long nodeSize, long entrySize, int entryCount) + { + Assert.AssertTrue(entrySize >= sizeof(long)); + Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); + Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); + Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); + Assert.AssertTrue(entryCount >= 0); + + if (entryCount <= 0) + return 0; + + return GetEntrySetCount(nodeSize, entrySize, entryCount) * nodeSize; + } + + private static int GetEntryCount(long nodeSize, long entrySize) + { + return (int)((nodeSize - Unsafe.SizeOf()) / entrySize); + } + + private static int GetOffsetCount(long nodeSize) + { + return (int)((nodeSize - Unsafe.SizeOf()) / sizeof(long)); + } + + private static int GetEntrySetCount(long nodeSize, long entrySize, int entryCount) + { + int entryCountPerNode = GetEntryCount(nodeSize, entrySize); + return Util.DivideByRoundUp(entryCount, entryCountPerNode); + } + + public static int GetNodeL2Count(long nodeSize, long entrySize, int entryCount) + { + int offsetCountPerNode = GetOffsetCount(nodeSize); + int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); + + if (entrySetCount <= offsetCountPerNode) + return 0; + + int nodeL2Count = Util.DivideByRoundUp(entrySetCount, offsetCountPerNode); + Abort.DoAbortUnless(nodeL2Count <= offsetCountPerNode); + + return Util.DivideByRoundUp(entrySetCount - (offsetCountPerNode - (nodeL2Count - 1)), offsetCountPerNode); + } + + private static long GetBucketTreeEntryOffset(long entrySetOffset, long entrySize, int entryIndex) + { + return entrySetOffset + Unsafe.SizeOf() + entryIndex * entrySize; + } + + private static long GetBucketTreeEntryOffset(int entrySetIndex, long nodeSize, long entrySize, int entryIndex) + { + return GetBucketTreeEntryOffset(entrySetIndex * nodeSize, entrySize, entryIndex); + } + + private bool IsExistL2() => OffsetCount < EntrySetCount; + private bool IsExistOffsetL2OnL1() => IsExistL2() && _nodeL1.GetHeader().Count < OffsetCount; + + private long GetEntrySetIndex(int nodeIndex, int offsetIndex) + { + return (OffsetCount - _nodeL1.GetHeader().Count) + (OffsetCount * nodeIndex) + offsetIndex; + } + + public struct Header + { + public uint Magic; + public uint Version; + public int EntryCount; +#pragma warning disable 414 + private int _reserved; +#pragma warning restore 414 + + public void Format(int entryCount) { - list[i].Next = list[i + 1]; - list[i].OffsetEnd = list[i + 1].Offset; + Magic = ExpectedMagic; + Version = MaxVersion; + EntryCount = entryCount; + _reserved = 0; } - list[list.Count - 1].OffsetEnd = BucketOffsets.OffsetEnd; - - return list; - } - } - - public class BucketTreeBucket where T : BucketTreeEntry, new() - { - public int Index; - public int EntryCount; - public long OffsetEnd; - public T[] Entries; - - public BucketTreeBucket(BinaryReader reader) - { - Index = reader.ReadInt32(); - EntryCount = reader.ReadInt32(); - OffsetEnd = reader.ReadInt64(); - Entries = new T[EntryCount]; - - for (int i = 0; i < EntryCount; i++) + public Result Verify() { - Entries[i] = new T().Read(reader); + if (Magic != ExpectedMagic) + return ResultFs.InvalidBucketTreeSignature.Log(); + + if (EntryCount < 0) + return ResultFs.InvalidBucketTreeEntryCount.Log(); + + if (Version > MaxVersion) + return ResultFs.UnsupportedVersion.Log(); + + return Result.Success; } } - } - public abstract class BucketTreeEntry where T : BucketTreeEntry - { - public long Offset { get; set; } - public long OffsetEnd { get; set; } - public T Next { get; set; } - - protected abstract void ReadSpecific(BinaryReader reader); - internal T Read(BinaryReader reader) + public struct NodeHeader { - Offset = reader.ReadInt64(); - ReadSpecific(reader); - return (T)this; + public int Index; + public int Count; + public long Offset; + + public Result Verify(int nodeIndex, long nodeSize, long entrySize) + { + if (Index != nodeIndex) + return ResultFs.InvalidBucketTreeNodeIndex.Log(); + + if (entrySize == 0 || nodeSize < entrySize + NodeHeaderSize) + return ResultFs.InvalidSize.Log(); + + long maxEntryCount = (nodeSize - NodeHeaderSize) / entrySize; + + if (Count <= 0 || maxEntryCount < Count) + return ResultFs.InvalidBucketTreeNodeEntryCount.Log(); + + if (Offset < 0) + return ResultFs.InvalidBucketTreeNodeOffset.Log(); + + return Result.Success; + } } - } - public class OffsetEntry : BucketTreeEntry - { - protected override void ReadSpecific(BinaryReader reader) { } - } - - public class AesSubsectionEntry : BucketTreeEntry - { - public uint Field8 { get; set; } - public uint Counter { get; set; } - - protected override void ReadSpecific(BinaryReader reader) + private struct NodeBuffer { - Field8 = reader.ReadUInt32(); - Counter = reader.ReadUInt32(); + // Use long to ensure alignment + private long[] _header; + + public bool Allocate(int nodeSize) + { + Assert.AssertTrue(_header == null); + + _header = new long[nodeSize / sizeof(long)]; + + return _header != null; + } + + public void Free() + { + _header = null; + } + + public void FillZero() + { + if (_header != null) + { + Array.Fill(_header, 0); + } + } + + public ref NodeHeader GetHeader() + { + Assert.AssertTrue(_header.Length / sizeof(long) >= Unsafe.SizeOf()); + + return ref Unsafe.As(ref _header[0]); + } + + public Span GetBuffer() + { + return MemoryMarshal.AsBytes(_header.AsSpan()); + } + + public BucketTreeNode GetNode() where TEntry : unmanaged + { + return new BucketTreeNode(GetBuffer()); + } } - } - public class RelocationEntry : BucketTreeEntry - { - public long SourceOffset { get; set; } - public int SourceIndex { get; set; } - - protected override void ReadSpecific(BinaryReader reader) + public readonly ref struct BucketTreeNode where TEntry : unmanaged { - SourceOffset = reader.ReadInt64(); - SourceIndex = reader.ReadInt32(); + private readonly Span _buffer; + + public BucketTreeNode(Span buffer) + { + _buffer = buffer; + + Assert.AssertTrue(_buffer.Length >= Unsafe.SizeOf()); + Assert.AssertTrue(_buffer.Length >= Unsafe.SizeOf() + GetHeader().Count * Unsafe.SizeOf()); + } + + public int GetCount() => GetHeader().Count; + + public ReadOnlySpan GetArray() => GetWritableArray(); + internal Span GetWritableArray() => GetWritableArray(); + + public long GetBeginOffset() => GetArray()[0]; + public long GetEndOffset() => GetHeader().Offset; + public long GetL2BeginOffset() => GetArray()[GetCount()]; + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public ReadOnlySpan GetArray() where TElement : unmanaged + { + return GetWritableArray(); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private Span GetWritableArray() where TElement : unmanaged + { + return MemoryMarshal.Cast(_buffer.Slice(Unsafe.SizeOf())); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal ref NodeHeader GetHeader() + { + return ref Unsafe.As(ref MemoryMarshal.GetReference(_buffer)); + } + } + + public ref struct Visitor + { + private BucketTree Tree { get; set; } + private byte[] Entry { get; set; } + private int EntryIndex { get; set; } + private int EntrySetCount { get; set; } + private EntrySetHeader _entrySet; + + [StructLayout(LayoutKind.Explicit)] + private struct EntrySetHeader + { + // ReSharper disable once MemberHidesStaticFromOuterClass + [FieldOffset(0)] public NodeHeader Header; + [FieldOffset(0)] public EntrySetInfo Info; + + [StructLayout(LayoutKind.Sequential)] + public struct EntrySetInfo + { + public int Index; + public int Count; + public long End; + public long Start; + } + } + + public Result Initialize(BucketTree tree) + { + Assert.AssertTrue(tree != null); + Assert.AssertTrue(Tree == null || tree == Tree); + + if (Entry == null) + { + Entry = new byte[tree.EntrySize]; + Tree = tree; + EntryIndex = -1; + } + + return Result.Success; + } + + public void Dispose() + { + // todo: try using shared arrays + } + + public bool IsValid() => EntryIndex >= 0; + + public bool CanMoveNext() + { + return IsValid() && (EntryIndex + 1 < _entrySet.Info.Count || _entrySet.Info.Index + 1 < EntrySetCount); + } + + public bool CanMovePrevious() + { + return IsValid() && (EntryIndex > 0 || _entrySet.Info.Index > 0); + } + + public ref T Get() where T : unmanaged + { + return ref MemoryMarshal.Cast(Entry)[0]; + } + + public Result MoveNext() + { + Result rc; + + if (!IsValid()) + return ResultFs.OutOfRange.Log(); + + int entryIndex = EntryIndex + 1; + + // Invalidate our index, and read the header for the next index. + if (entryIndex == _entrySet.Info.Count) + { + int entrySetIndex = _entrySet.Info.Index + 1; + if (entrySetIndex >= EntrySetCount) + return ResultFs.OutOfRange.Log(); + + EntryIndex = -1; + + long end = _entrySet.Info.End; + + long entrySetSize = Tree.NodeSize; + long entrySetOffset = entrySetIndex * entrySetSize; + + rc = Tree.EntryStorage.Read(entrySetOffset, SpanHelpers.AsByteSpan(ref _entrySet)); + if (rc.IsFailure()) return rc; + + rc = _entrySet.Header.Verify(entrySetIndex, entrySetSize, Tree.EntrySize); + if (rc.IsFailure()) return rc; + + if (_entrySet.Info.Start != end || _entrySet.Info.Start >= _entrySet.Info.End) + return ResultFs.InvalidBucketTreeEntrySetOffset.Log(); + + entryIndex = 0; + } + else + { + EntryIndex = 1; + } + + // Read the new entry + long entrySize = Tree.EntrySize; + long entryOffset = GetBucketTreeEntryOffset(_entrySet.Info.Index, Tree.NodeSize, entrySize, entryIndex); + + rc = Tree.EntryStorage.Read(entryOffset, Entry); + if (rc.IsFailure()) return rc; + + // Note that we changed index. + EntryIndex = entryIndex; + return Result.Success; + } + + public Result MovePrevious() + { + Result rc; + + if (!IsValid()) + return ResultFs.OutOfRange.Log(); + + int entryIndex = EntryIndex; + + if (entryIndex == 0) + { + if (_entrySet.Info.Index <= 0) + return ResultFs.OutOfRange.Log(); + + EntryIndex = -1; + + long start = _entrySet.Info.Start; + + long entrySetSize = Tree.NodeSize; + int entrySetIndex = _entrySet.Info.Index - 1; + long entrySetOffset = entrySetIndex * entrySetSize; + + rc = Tree.EntryStorage.Read(entrySetOffset, SpanHelpers.AsByteSpan(ref _entrySet)); + if (rc.IsFailure()) return rc; + + rc = _entrySet.Header.Verify(entrySetIndex, entrySetSize, Tree.EntrySize); + if (rc.IsFailure()) return rc; + + if (_entrySet.Info.End != start || _entrySet.Info.Start >= _entrySet.Info.End) + return ResultFs.InvalidBucketTreeEntrySetOffset.Log(); + + entryIndex = _entrySet.Info.Count; + } + else + { + EntryIndex = -1; + } + + // Read the new entry + long entrySize = Tree.EntrySize; + long entryOffset = GetBucketTreeEntryOffset(_entrySet.Info.Index, Tree.NodeSize, entrySize, entryIndex); + + rc = Tree.EntryStorage.Read(entryOffset, Entry); + if (rc.IsFailure()) return rc; + + // Note that we changed index. + EntryIndex = entryIndex; + return Result.Success; + } + + public Result Find(long virtualAddress) + { + Result rc; + + // Get the node. + BucketTreeNode node = Tree._nodeL1.GetNode(); + + if (virtualAddress >= node.GetEndOffset()) + return ResultFs.OutOfRange.Log(); + + int entrySetIndex; + + if (Tree.IsExistOffsetL2OnL1() && virtualAddress < node.GetBeginOffset()) + { + // The portion of the L2 offsets containing our target offset is stored in the L1 node + ReadOnlySpan offsets = node.GetArray().Slice(node.GetCount()); + int index = offsets.BinarySearch(virtualAddress); + if (index < 0) index = (~index) - 1; + + if (index < 0) + return ResultFs.OutOfRange.Log(); + + entrySetIndex = index; + } + else + { + ReadOnlySpan offsets = node.GetArray().Slice(0, node.GetCount()); + int index = offsets.BinarySearch(virtualAddress); + if (index < 0) index = (~index) - 1; + + if (index < 0) + return ResultFs.OutOfRange.Log(); + + if (Tree.IsExistL2()) + { + if (index >= Tree.OffsetCount) + return ResultFs.InvalidBucketTreeNodeOffset.Log(); + + rc = FindEntrySet(out entrySetIndex, virtualAddress, index); + if (rc.IsFailure()) return rc; + } + else + { + entrySetIndex = index; + } + } + + // Validate the entry set index. + if (entrySetIndex < 0 || entrySetIndex >= Tree.EntrySetCount) + return ResultFs.InvalidBucketTreeNodeOffset.Log(); + + // Find the entry. + rc = FindEntry(virtualAddress, entrySetIndex); + if (rc.IsFailure()) return rc; + + // Set count. + EntrySetCount = Tree.EntrySetCount; + return Result.Success; + } + + private Result FindEntrySet(out int entrySetIndex, long virtualAddress, int nodeIndex) + { + long nodeSize = Tree.NodeSize; + + using (var rented = new RentedArray((int)nodeSize)) + { + return FindEntrySetWithBuffer(out entrySetIndex, virtualAddress, nodeIndex, rented.Span); + } + } + + private Result FindEntrySetWithBuffer(out int outIndex, long virtualAddress, int nodeIndex, + Span buffer) + { + outIndex = default; + + // Calculate node extents. + long nodeSize = Tree.NodeSize; + long nodeOffset = (nodeIndex + 1) * nodeSize; + SubStorage2 storage = Tree.NodeStorage; + + // Read the node. + Result rc = storage.Read(nodeOffset, buffer.Slice(0, (int)nodeSize)); + if (rc.IsFailure()) return rc; + + // Validate the header. + NodeHeader header = MemoryMarshal.Cast(buffer)[0]; + rc = header.Verify(nodeIndex, nodeSize, sizeof(long)); + if (rc.IsFailure()) return rc; + + // Create the node and find. + var node = new StorageNode(sizeof(long), header.Count); + node.Find(buffer, virtualAddress); + + if (node.GetIndex() < 0) + return ResultFs.InvalidBucketTreeVirtualOffset.Log(); + + // Return the index. + outIndex = (int)Tree.GetEntrySetIndex(header.Index, node.GetIndex()); + return Result.Success; + } + + private Result FindEntry(long virtualAddress, int entrySetIndex) + { + long entrySetSize = Tree.NodeSize; + + using (var rented = new RentedArray((int)entrySetSize)) + { + return FindEntryWithBuffer(virtualAddress, entrySetIndex, rented.Span); + } + } + + private Result FindEntryWithBuffer(long virtualAddress, int entrySetIndex, Span buffer) + { + // Calculate entry set extents. + long entrySize = Tree.EntrySize; + long entrySetSize = Tree.NodeSize; + long entrySetOffset = entrySetIndex * entrySetSize; + SubStorage2 storage = Tree.EntryStorage; + + // Read the entry set. + Result rc = storage.Read(entrySetOffset, buffer.Slice(0, (int)entrySetSize)); + if (rc.IsFailure()) return rc; + + // Validate the entry set. + EntrySetHeader entrySet = MemoryMarshal.Cast(buffer)[0]; + rc = entrySet.Header.Verify(entrySetIndex, entrySetSize, entrySize); + if (rc.IsFailure()) return rc; + + // Create the node, and find. + var node = new StorageNode(entrySize, entrySet.Info.Count); + node.Find(buffer, virtualAddress); + + if (node.GetIndex() < 0) + return ResultFs.InvalidBucketTreeVirtualOffset.Log(); + + // Copy the data into entry. + int entryIndex = node.GetIndex(); + long entryOffset = GetBucketTreeEntryOffset(0, entrySize, entryIndex); + buffer.Slice((int)entryOffset, (int)entrySize).CopyTo(Entry); + + // Set our entry set/index. + _entrySet = entrySet; + EntryIndex = entryIndex; + + return Result.Success; + } + + private struct StorageNode + { + private Offset _start; + private int _count; + private int _index; + + public StorageNode(long size, int count) + { + _start = new Offset(NodeHeaderSize, (int)size); + _count = count; + _index = -1; + } + + public int GetIndex() => _index; + + public void Find(ReadOnlySpan buffer, long virtualAddress) + { + int end = _count; + Offset pos = _start; + + while (end > 0) + { + int half = end / 2; + Offset mid = pos + half; + + long offset = BinaryPrimitives.ReadInt64LittleEndian(buffer.Slice((int)mid.Get())); + + if (offset <= virtualAddress) + { + pos = mid + 1; + end -= half + 1; + } + else + { + end = half; + } + } + + _index = (int)(pos - _start) - 1; + } + + private readonly struct Offset + { + private readonly long _offset; + private readonly int _stride; + + public Offset(long offset, int stride) + { + _offset = offset; + _stride = stride; + } + + public long Get() => _offset; + + public static Offset operator ++(Offset left) => left + 1; + public static Offset operator --(Offset left) => left - 1; + + public static Offset operator +(Offset left, long right) => new Offset(left._offset + right * left._stride, left._stride); + public static Offset operator -(Offset left, long right) => new Offset(left._offset - right * left._stride, left._stride); + + public static long operator -(Offset left, Offset right) => + (left._offset - right._offset) / left._stride; + + public static bool operator ==(Offset left, Offset right) => left._offset == right._offset; + public static bool operator !=(Offset left, Offset right) => left._offset != right._offset; + + public bool Equals(Offset other) => _offset == other._offset; + public override bool Equals(object obj) => obj is Offset other && Equals(other); + public override int GetHashCode() => _offset.GetHashCode(); + } + } } } } diff --git a/src/LibHac/FsSystem/BucketTree2.cs b/src/LibHac/FsSystem/BucketTree2.cs deleted file mode 100644 index 97598b30..00000000 --- a/src/LibHac/FsSystem/BucketTree2.cs +++ /dev/null @@ -1,738 +0,0 @@ -using System; -using System.Buffers.Binary; -using System.Runtime.CompilerServices; -using System.Runtime.InteropServices; -using LibHac.Common; -using LibHac.Diag; -using LibHac.Fs; - -namespace LibHac.FsSystem -{ - public partial class BucketTree2 - { - private const uint ExpectedMagic = 0x52544B42; // BKTR - private const int MaxVersion = 1; - - private const int NodeSizeMin = 1024; - private const int NodeSizeMax = 1024 * 512; - - private static int NodeHeaderSize => Unsafe.SizeOf(); - - private SubStorage2 NodeStorage { get; set; } - private SubStorage2 EntryStorage { get; set; } - - private NodeBuffer _nodeL1 = new NodeBuffer(); - - private long NodeSize { get; set; } - private long EntrySize { get; set; } - private int OffsetCount { get; set; } - private int EntrySetCount { get; set; } - private long StartOffset { get; set; } - private long EndOffset { get; set; } - - public Result Initialize(SubStorage2 nodeStorage, SubStorage2 entryStorage, int nodeSize, int entrySize, - int entryCount) - { - Assert.AssertTrue(entrySize >= sizeof(long)); - Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); - Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); - Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); - Assert.AssertTrue(!IsInitialized()); - - // Ensure valid entry count. - if (entryCount <= 0) - return ResultFs.InvalidArgument.Log(); - - // Allocate node. - if (!_nodeL1.Allocate(nodeSize)) - return ResultFs.BufferAllocationFailed.Log(); - - bool needFree = true; - try - { - // Read node. - Result rc = nodeStorage.Read(0, _nodeL1.GetBuffer()); - if (rc.IsFailure()) return rc; - - // Verify node. - rc = _nodeL1.GetHeader().Verify(0, nodeSize, sizeof(long)); - if (rc.IsFailure()) return rc; - - // Validate offsets. - int offsetCount = GetOffsetCount(nodeSize); - int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); - BucketTreeNode node = _nodeL1.GetNode(); - - long startOffset; - if (offsetCount < entrySetCount && node.GetCount() < offsetCount) - { - startOffset = node.GetL2BeginOffset(); - } - else - { - startOffset = node.GetBeginOffset(); - } - - long endOffset = node.GetEndOffset(); - - if (startOffset < 0 || startOffset > node.GetBeginOffset() || startOffset >= endOffset) - return ResultFs.InvalidBucketTreeEntryOffset.Log(); - - NodeStorage = nodeStorage; - EntryStorage = entryStorage; - NodeSize = nodeSize; - EntrySize = entrySize; - OffsetCount = offsetCount; - EntrySetCount = entrySetCount; - StartOffset = startOffset; - EndOffset = endOffset; - - needFree = false; - - return Result.Success; - } - finally - { - if (needFree) - _nodeL1.Free(); - } - } - - public bool IsInitialized() => NodeSize > 0; - public bool IsEmpty() => EntrySize == 0; - - public long GetStart() => StartOffset; - 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()); - - if (virtualAddress < 0) - return ResultFs.InvalidOffset.Log(); - - if (IsEmpty()) - return ResultFs.OutOfRange.Log(); - - Result rc = visitor.Initialize(this); - if (rc.IsFailure()) return rc; - - return visitor.Find(virtualAddress); - } - - public static int QueryHeaderStorageSize() => Unsafe.SizeOf
(); - - public static long QueryNodeStorageSize(long nodeSize, long entrySize, int entryCount) - { - Assert.AssertTrue(entrySize >= sizeof(long)); - Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); - Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); - Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); - Assert.AssertTrue(entryCount >= 0); - - if (entryCount <= 0) - return 0; - - return (1 + GetNodeL2Count(nodeSize, entrySize, entryCount)) * nodeSize; - } - - public static long QueryEntryStorageSize(long nodeSize, long entrySize, int entryCount) - { - Assert.AssertTrue(entrySize >= sizeof(long)); - Assert.AssertTrue(nodeSize >= entrySize + Unsafe.SizeOf()); - Assert.AssertTrue(NodeSizeMin <= nodeSize && nodeSize <= NodeSizeMax); - Assert.AssertTrue(Util.IsPowerOfTwo(nodeSize)); - Assert.AssertTrue(entryCount >= 0); - - if (entryCount <= 0) - return 0; - - return GetEntrySetCount(nodeSize, entrySize, entryCount) * nodeSize; - } - - private static int GetEntryCount(long nodeSize, long entrySize) - { - return (int)((nodeSize - Unsafe.SizeOf()) / entrySize); - } - - private static int GetOffsetCount(long nodeSize) - { - return (int)((nodeSize - Unsafe.SizeOf()) / sizeof(long)); - } - - private static int GetEntrySetCount(long nodeSize, long entrySize, int entryCount) - { - int entryCountPerNode = GetEntryCount(nodeSize, entrySize); - return Util.DivideByRoundUp(entryCount, entryCountPerNode); - } - - public static int GetNodeL2Count(long nodeSize, long entrySize, int entryCount) - { - int offsetCountPerNode = GetOffsetCount(nodeSize); - int entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount); - - if (entrySetCount <= offsetCountPerNode) - return 0; - - int nodeL2Count = Util.DivideByRoundUp(entrySetCount, offsetCountPerNode); - Abort.DoAbortUnless(nodeL2Count <= offsetCountPerNode); - - return Util.DivideByRoundUp(entrySetCount - (offsetCountPerNode - (nodeL2Count - 1)), offsetCountPerNode); - } - - private static long GetBucketTreeEntryOffset(long entrySetOffset, long entrySize, int entryIndex) - { - return entrySetOffset + Unsafe.SizeOf() + entryIndex * entrySize; - } - - private static long GetBucketTreeEntryOffset(int entrySetIndex, long nodeSize, long entrySize, int entryIndex) - { - return GetBucketTreeEntryOffset(entrySetIndex * nodeSize, entrySize, entryIndex); - } - - private bool IsExistL2() => OffsetCount < EntrySetCount; - private bool IsExistOffsetL2OnL1() => IsExistL2() && _nodeL1.GetHeader().Count < OffsetCount; - - private long GetEntrySetIndex(int nodeIndex, int offsetIndex) - { - return (OffsetCount - _nodeL1.GetHeader().Count) + (OffsetCount * nodeIndex) + offsetIndex; - } - - public struct Header - { - public uint Magic; - public uint Version; - public int EntryCount; -#pragma warning disable 414 - private int _reserved; -#pragma warning restore 414 - - public void Format(int entryCount) - { - Magic = ExpectedMagic; - Version = MaxVersion; - EntryCount = entryCount; - _reserved = 0; - } - - public Result Verify() - { - if (Magic != ExpectedMagic) - return ResultFs.InvalidBucketTreeSignature.Log(); - - if (EntryCount < 0) - return ResultFs.InvalidBucketTreeEntryCount.Log(); - - if (Version > MaxVersion) - return ResultFs.UnsupportedVersion.Log(); - - return Result.Success; - } - } - - public struct NodeHeader - { - public int Index; - public int Count; - public long Offset; - - public Result Verify(int nodeIndex, long nodeSize, long entrySize) - { - if (Index != nodeIndex) - return ResultFs.InvalidBucketTreeNodeIndex.Log(); - - if (entrySize == 0 || nodeSize < entrySize + NodeHeaderSize) - return ResultFs.InvalidSize.Log(); - - long maxEntryCount = (nodeSize - NodeHeaderSize) / entrySize; - - if (Count <= 0 || maxEntryCount < Count) - return ResultFs.InvalidBucketTreeNodeEntryCount.Log(); - - if (Offset < 0) - return ResultFs.InvalidBucketTreeNodeOffset.Log(); - - return Result.Success; - } - } - - private struct NodeBuffer - { - // Use long to ensure alignment - private long[] _header; - - public bool Allocate(int nodeSize) - { - Assert.AssertTrue(_header == null); - - _header = new long[nodeSize / sizeof(long)]; - - return _header != null; - } - - public void Free() - { - _header = null; - } - - public void FillZero() - { - if (_header != null) - { - Array.Fill(_header, 0); - } - } - - public ref NodeHeader GetHeader() - { - Assert.AssertTrue(_header.Length / sizeof(long) >= Unsafe.SizeOf()); - - return ref Unsafe.As(ref _header[0]); - } - - public Span GetBuffer() - { - return MemoryMarshal.AsBytes(_header.AsSpan()); - } - - public BucketTreeNode GetNode() where TEntry : unmanaged - { - return new BucketTreeNode(GetBuffer()); - } - } - - public readonly ref struct BucketTreeNode where TEntry : unmanaged - { - private readonly Span _buffer; - - public BucketTreeNode(Span buffer) - { - _buffer = buffer; - - Assert.AssertTrue(_buffer.Length >= Unsafe.SizeOf()); - Assert.AssertTrue(_buffer.Length >= Unsafe.SizeOf() + GetHeader().Count * Unsafe.SizeOf()); - } - - public int GetCount() => GetHeader().Count; - - public ReadOnlySpan GetArray() => GetWritableArray(); - internal Span GetWritableArray() => GetWritableArray(); - - public long GetBeginOffset() => GetArray()[0]; - public long GetEndOffset() => GetHeader().Offset; - public long GetL2BeginOffset() => GetArray()[GetCount()]; - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ReadOnlySpan GetArray() where TElement : unmanaged - { - return GetWritableArray(); - } - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - private Span GetWritableArray() where TElement : unmanaged - { - return MemoryMarshal.Cast(_buffer.Slice(Unsafe.SizeOf())); - } - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal ref NodeHeader GetHeader() - { - return ref Unsafe.As(ref MemoryMarshal.GetReference(_buffer)); - } - } - - public ref struct Visitor - { - private BucketTree2 Tree { get; set; } - private byte[] Entry { get; set; } - private int EntryIndex { get; set; } - private int EntrySetCount { get; set; } - private EntrySetHeader _entrySet; - - [StructLayout(LayoutKind.Explicit)] - private struct EntrySetHeader - { - // ReSharper disable once MemberHidesStaticFromOuterClass - [FieldOffset(0)] public NodeHeader Header; - [FieldOffset(0)] public EntrySetInfo Info; - - [StructLayout(LayoutKind.Sequential)] - public struct EntrySetInfo - { - public int Index; - public int Count; - public long End; - public long Start; - } - } - - public Result Initialize(BucketTree2 tree) - { - Assert.AssertTrue(tree != null); - Assert.AssertTrue(Tree == null || tree == Tree); - - if (Entry == null) - { - Entry = new byte[tree.EntrySize]; - Tree = tree; - EntryIndex = -1; - } - - return Result.Success; - } - - public void Dispose() - { - // todo: try using shared arrays - } - - public bool IsValid() => EntryIndex >= 0; - - public bool CanMoveNext() - { - return IsValid() && (EntryIndex + 1 < _entrySet.Info.Count || _entrySet.Info.Index + 1 < EntrySetCount); - } - - public bool CanMovePrevious() - { - return IsValid() && (EntryIndex > 0 || _entrySet.Info.Index > 0); - } - - public ref T Get() where T : unmanaged - { - return ref MemoryMarshal.Cast(Entry)[0]; - } - - public Result MoveNext() - { - Result rc; - - if (!IsValid()) - return ResultFs.OutOfRange.Log(); - - int entryIndex = EntryIndex + 1; - - // Invalidate our index, and read the header for the next index. - if (entryIndex == _entrySet.Info.Count) - { - int entrySetIndex = _entrySet.Info.Index + 1; - if (entrySetIndex >= EntrySetCount) - return ResultFs.OutOfRange.Log(); - - EntryIndex = -1; - - long end = _entrySet.Info.End; - - long entrySetSize = Tree.NodeSize; - long entrySetOffset = entrySetIndex * entrySetSize; - - rc = Tree.EntryStorage.Read(entrySetOffset, SpanHelpers.AsByteSpan(ref _entrySet)); - if (rc.IsFailure()) return rc; - - rc = _entrySet.Header.Verify(entrySetIndex, entrySetSize, Tree.EntrySize); - if (rc.IsFailure()) return rc; - - if (_entrySet.Info.Start != end || _entrySet.Info.Start >= _entrySet.Info.End) - return ResultFs.InvalidBucketTreeEntrySetOffset.Log(); - - entryIndex = 0; - } - else - { - EntryIndex = 1; - } - - // Read the new entry - long entrySize = Tree.EntrySize; - long entryOffset = GetBucketTreeEntryOffset(_entrySet.Info.Index, Tree.NodeSize, entrySize, entryIndex); - - rc = Tree.EntryStorage.Read(entryOffset, Entry); - if (rc.IsFailure()) return rc; - - // Note that we changed index. - EntryIndex = entryIndex; - return Result.Success; - } - - public Result MovePrevious() - { - Result rc; - - if (!IsValid()) - return ResultFs.OutOfRange.Log(); - - int entryIndex = EntryIndex; - - if (entryIndex == 0) - { - if (_entrySet.Info.Index <= 0) - return ResultFs.OutOfRange.Log(); - - EntryIndex = -1; - - long start = _entrySet.Info.Start; - - long entrySetSize = Tree.NodeSize; - int entrySetIndex = _entrySet.Info.Index - 1; - long entrySetOffset = entrySetIndex * entrySetSize; - - rc = Tree.EntryStorage.Read(entrySetOffset, SpanHelpers.AsByteSpan(ref _entrySet)); - if (rc.IsFailure()) return rc; - - rc = _entrySet.Header.Verify(entrySetIndex, entrySetSize, Tree.EntrySize); - if (rc.IsFailure()) return rc; - - if (_entrySet.Info.End != start || _entrySet.Info.Start >= _entrySet.Info.End) - return ResultFs.InvalidBucketTreeEntrySetOffset.Log(); - - entryIndex = _entrySet.Info.Count; - } - else - { - EntryIndex = -1; - } - - // Read the new entry - long entrySize = Tree.EntrySize; - long entryOffset = GetBucketTreeEntryOffset(_entrySet.Info.Index, Tree.NodeSize, entrySize, entryIndex); - - rc = Tree.EntryStorage.Read(entryOffset, Entry); - if (rc.IsFailure()) return rc; - - // Note that we changed index. - EntryIndex = entryIndex; - return Result.Success; - } - - public Result Find(long virtualAddress) - { - Result rc; - - // Get the node. - BucketTreeNode node = Tree._nodeL1.GetNode(); - - if (virtualAddress >= node.GetEndOffset()) - return ResultFs.OutOfRange.Log(); - - int entrySetIndex; - - if (Tree.IsExistOffsetL2OnL1() && virtualAddress < node.GetBeginOffset()) - { - // The portion of the L2 offsets containing our target offset is stored in the L1 node - ReadOnlySpan offsets = node.GetArray().Slice(node.GetCount()); - int index = offsets.BinarySearch(virtualAddress); - if (index < 0) index = (~index) - 1; - - if (index < 0) - return ResultFs.OutOfRange.Log(); - - entrySetIndex = index; - } - else - { - ReadOnlySpan offsets = node.GetArray().Slice(0, node.GetCount()); - int index = offsets.BinarySearch(virtualAddress); - if (index < 0) index = (~index) - 1; - - if (index < 0) - return ResultFs.OutOfRange.Log(); - - if (Tree.IsExistL2()) - { - if (index >= Tree.OffsetCount) - return ResultFs.InvalidBucketTreeNodeOffset.Log(); - - rc = FindEntrySet(out entrySetIndex, virtualAddress, index); - if (rc.IsFailure()) return rc; - } - else - { - entrySetIndex = index; - } - } - - // Validate the entry set index. - if (entrySetIndex < 0 || entrySetIndex >= Tree.EntrySetCount) - return ResultFs.InvalidBucketTreeNodeOffset.Log(); - - // Find the entry. - rc = FindEntry(virtualAddress, entrySetIndex); - if (rc.IsFailure()) return rc; - - // Set count. - EntrySetCount = Tree.EntrySetCount; - return Result.Success; - } - - private Result FindEntrySet(out int entrySetIndex, long virtualAddress, int nodeIndex) - { - long nodeSize = Tree.NodeSize; - - using (var rented = new RentedArray((int)nodeSize)) - { - return FindEntrySetWithBuffer(out entrySetIndex, virtualAddress, nodeIndex, rented.Span); - } - } - - private Result FindEntrySetWithBuffer(out int outIndex, long virtualAddress, int nodeIndex, - Span buffer) - { - outIndex = default; - - // Calculate node extents. - long nodeSize = Tree.NodeSize; - long nodeOffset = (nodeIndex + 1) * nodeSize; - SubStorage2 storage = Tree.NodeStorage; - - // Read the node. - Result rc = storage.Read(nodeOffset, buffer.Slice(0, (int)nodeSize)); - if (rc.IsFailure()) return rc; - - // Validate the header. - NodeHeader header = MemoryMarshal.Cast(buffer)[0]; - rc = header.Verify(nodeIndex, nodeSize, sizeof(long)); - if (rc.IsFailure()) return rc; - - // Create the node and find. - var node = new StorageNode(sizeof(long), header.Count); - node.Find(buffer, virtualAddress); - - if (node.GetIndex() < 0) - return ResultFs.InvalidBucketTreeVirtualOffset.Log(); - - // Return the index. - outIndex = (int)Tree.GetEntrySetIndex(header.Index, node.GetIndex()); - return Result.Success; - } - - private Result FindEntry(long virtualAddress, int entrySetIndex) - { - long entrySetSize = Tree.NodeSize; - - using (var rented = new RentedArray((int)entrySetSize)) - { - return FindEntryWithBuffer(virtualAddress, entrySetIndex, rented.Span); - } - } - - private Result FindEntryWithBuffer(long virtualAddress, int entrySetIndex, Span buffer) - { - // Calculate entry set extents. - long entrySize = Tree.EntrySize; - long entrySetSize = Tree.NodeSize; - long entrySetOffset = entrySetIndex * entrySetSize; - SubStorage2 storage = Tree.EntryStorage; - - // Read the entry set. - Result rc = storage.Read(entrySetOffset, buffer.Slice(0, (int)entrySetSize)); - if (rc.IsFailure()) return rc; - - // Validate the entry set. - EntrySetHeader entrySet = MemoryMarshal.Cast(buffer)[0]; - rc = entrySet.Header.Verify(entrySetIndex, entrySetSize, entrySize); - if (rc.IsFailure()) return rc; - - // Create the node, and find. - var node = new StorageNode(entrySize, entrySet.Info.Count); - node.Find(buffer, virtualAddress); - - if (node.GetIndex() < 0) - return ResultFs.InvalidBucketTreeVirtualOffset.Log(); - - // Copy the data into entry. - int entryIndex = node.GetIndex(); - long entryOffset = GetBucketTreeEntryOffset(0, entrySize, entryIndex); - buffer.Slice((int)entryOffset, (int)entrySize).CopyTo(Entry); - - // Set our entry set/index. - _entrySet = entrySet; - EntryIndex = entryIndex; - - return Result.Success; - } - - private struct StorageNode - { - private Offset _start; - private int _count; - private int _index; - - public StorageNode(long size, int count) - { - _start = new Offset(NodeHeaderSize, (int)size); - _count = count; - _index = -1; - } - - public int GetIndex() => _index; - - public void Find(ReadOnlySpan buffer, long virtualAddress) - { - int end = _count; - Offset pos = _start; - - while (end > 0) - { - int half = end / 2; - Offset mid = pos + half; - - long offset = BinaryPrimitives.ReadInt64LittleEndian(buffer.Slice((int)mid.Get())); - - if (offset <= virtualAddress) - { - pos = mid + 1; - end -= half + 1; - } - else - { - end = half; - } - } - - _index = (int)(pos - _start) - 1; - } - - private readonly struct Offset - { - private readonly long _offset; - private readonly int _stride; - - public Offset(long offset, int stride) - { - _offset = offset; - _stride = stride; - } - - public long Get() => _offset; - - public static Offset operator ++(Offset left) => left + 1; - public static Offset operator --(Offset left) => left - 1; - - public static Offset operator +(Offset left, long right) => new Offset(left._offset + right * left._stride, left._stride); - public static Offset operator -(Offset left, long right) => new Offset(left._offset - right * left._stride, left._stride); - - public static long operator -(Offset left, Offset right) => - (left._offset - right._offset) / left._stride; - - public static bool operator ==(Offset left, Offset right) => left._offset == right._offset; - public static bool operator !=(Offset left, Offset right) => left._offset != right._offset; - - public bool Equals(Offset other) => _offset == other._offset; - public override bool Equals(object obj) => obj is Offset other && Equals(other); - public override int GetHashCode() => _offset.GetHashCode(); - } - } - } - } -} diff --git a/src/LibHac/FsSystem/BucketTreeBuilder.cs b/src/LibHac/FsSystem/BucketTreeBuilder.cs index 309b08e4..01d82ebb 100644 --- a/src/LibHac/FsSystem/BucketTreeBuilder.cs +++ b/src/LibHac/FsSystem/BucketTreeBuilder.cs @@ -7,7 +7,7 @@ using LibHac.Fs; namespace LibHac.FsSystem { - public partial class BucketTree2 + public partial class BucketTree { public class Builder { diff --git a/src/LibHac/FsSystem/IndirectStorage.cs b/src/LibHac/FsSystem/IndirectStorage.cs index d1df2015..65664251 100644 --- a/src/LibHac/FsSystem/IndirectStorage.cs +++ b/src/LibHac/FsSystem/IndirectStorage.cs @@ -12,10 +12,10 @@ namespace LibHac.FsSystem public static readonly int StorageCount = 2; public static readonly int NodeSize = 1024 * 16; - private BucketTree2 Table { get; } = new BucketTree2(); + private BucketTree Table { get; } = new BucketTree(); private SubStorage2[] DataStorage { get; } = new SubStorage2[StorageCount]; - [StructLayout(LayoutKind.Sequential, Size = 0x14)] + [StructLayout(LayoutKind.Sequential, Size = 0x14, Pack = 4)] public struct Entry { private long VirtualOffset; @@ -29,13 +29,13 @@ namespace LibHac.FsSystem public long GetPhysicalOffset() => PhysicalOffset; } - public static long QueryHeaderStorageSize() => BucketTree2.QueryHeaderStorageSize(); + public static long QueryHeaderStorageSize() => BucketTree.QueryHeaderStorageSize(); public static long QueryNodeStorageSize(int entryCount) => - BucketTree2.QueryNodeStorageSize(NodeSize, Unsafe.SizeOf(), entryCount); + BucketTree.QueryNodeStorageSize(NodeSize, Unsafe.SizeOf(), entryCount); public static long QueryEntryStorageSize(int entryCount) => - BucketTree2.QueryEntryStorageSize(NodeSize, Unsafe.SizeOf(), entryCount); + BucketTree.QueryEntryStorageSize(NodeSize, Unsafe.SizeOf(), entryCount); public bool IsInitialized() => Table.IsInitialized(); @@ -43,7 +43,7 @@ namespace LibHac.FsSystem { // Read and verify the bucket tree header. // note: skip init - var header = new BucketTree2.Header(); + var header = new BucketTree.Header(); Result rc = tableStorage.Read(0, SpanHelpers.AsByteSpan(ref header)); if (rc.IsFailure()) return rc; @@ -100,7 +100,7 @@ namespace LibHac.FsSystem return ResultFs.OutOfRange.Log(); // Find the offset in our tree - var visitor = new BucketTree2.Visitor(); + var visitor = new BucketTree.Visitor(); Result rc = Table.Find(ref visitor, offset); if (rc.IsFailure()) return rc; @@ -214,7 +214,7 @@ namespace LibHac.FsSystem return ResultFs.OutOfRange.Log(); // Find the offset in our tree - var visitor = new BucketTree2.Visitor(); + var visitor = new BucketTree.Visitor(); Result rc = Table.Find(ref visitor, offset); if (rc.IsFailure()) return rc; diff --git a/src/LibHac/FsSystem/NcaUtils/Nca.cs b/src/LibHac/FsSystem/NcaUtils/Nca.cs index 1681ef08..672dcfa6 100644 --- a/src/LibHac/FsSystem/NcaUtils/Nca.cs +++ b/src/LibHac/FsSystem/NcaUtils/Nca.cs @@ -199,7 +199,7 @@ namespace LibHac.FsSystem.NcaUtils var cachedBucketTreeData = new CachedStorage(encryptionBucketTreeData, IndirectStorage.NodeSize, 6, true); - var treeHeader = new BucketTree2.Header(); + var treeHeader = new BucketTree.Header(); info.EncryptionTreeHeader.CopyTo(SpanHelpers.AsByteSpan(ref treeHeader)); long nodeStorageSize = IndirectStorage.QueryNodeStorageSize(treeHeader.EntryCount); long entryStorageSize = IndirectStorage.QueryEntryStorageSize(treeHeader.EntryCount); @@ -237,7 +237,7 @@ namespace LibHac.FsSystem.NcaUtils return patchStorage; } - var treeHeader = new BucketTree2.Header(); + var treeHeader = new BucketTree.Header(); patchInfo.RelocationTreeHeader.CopyTo(SpanHelpers.AsByteSpan(ref treeHeader)); long nodeStorageSize = IndirectStorage.QueryNodeStorageSize(treeHeader.EntryCount); long entryStorageSize = IndirectStorage.QueryEntryStorageSize(treeHeader.EntryCount); From 83dc874df14f2af6b5e444a494348622d9cc6341 Mon Sep 17 00:00:00 2001 From: Alex Barney Date: Fri, 26 Jun 2020 09:11:06 -0700 Subject: [PATCH 8/8] Rent arrays in BucketTree.Visitor --- src/LibHac/FsSystem/Aes128CtrExStorage.cs | 74 ++++---- src/LibHac/FsSystem/BucketTree.cs | 9 +- src/LibHac/FsSystem/IndirectStorage.cs | 215 +++++++++++----------- 3 files changed, 158 insertions(+), 140 deletions(-) diff --git a/src/LibHac/FsSystem/Aes128CtrExStorage.cs b/src/LibHac/FsSystem/Aes128CtrExStorage.cs index 79f305b1..46d4a3d6 100644 --- a/src/LibHac/FsSystem/Aes128CtrExStorage.cs +++ b/src/LibHac/FsSystem/Aes128CtrExStorage.cs @@ -36,47 +36,51 @@ namespace LibHac.FsSystem var visitor = new BucketTree.Visitor(); - Result rc = Table.Find(ref visitor, offset); - if (rc.IsFailure()) return rc; - - long inPos = offset; - int outPos = 0; - int remaining = destination.Length; - - while (remaining > 0) + try { - var currentEntry = visitor.Get(); + Result rc = Table.Find(ref visitor, offset); + if (rc.IsFailure()) return rc; - // Get and validate the next entry offset - long nextEntryOffset; - if (visitor.CanMoveNext()) + long inPos = offset; + int outPos = 0; + int remaining = destination.Length; + + while (remaining > 0) { - rc = visitor.MoveNext(); - if (rc.IsFailure()) return rc; + var currentEntry = visitor.Get(); - nextEntryOffset = visitor.Get().Offset; - if (!Table.Includes(nextEntryOffset)) - return ResultFs.InvalidIndirectEntryOffset.Log(); + // 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((uint)currentEntry.Generation); + + rc = base.DoRead(inPos, destination.Slice(outPos, bytesToRead)); + if (rc.IsFailure()) return rc; + } + + outPos += bytesToRead; + inPos += bytesToRead; + remaining -= bytesToRead; } - else - { - nextEntryOffset = Table.GetEnd(); - } - - int bytesToRead = (int)Math.Min(nextEntryOffset - inPos, remaining); - - lock (_locker) - { - UpdateCounterSubsection((uint)currentEntry.Generation); - - rc = base.DoRead(inPos, destination.Slice(outPos, bytesToRead)); - if (rc.IsFailure()) return rc; - } - - outPos += bytesToRead; - inPos += bytesToRead; - remaining -= bytesToRead; } + finally { visitor.Dispose(); } return Result.Success; } diff --git a/src/LibHac/FsSystem/BucketTree.cs b/src/LibHac/FsSystem/BucketTree.cs index a9cb125b..bede8cfc 100644 --- a/src/LibHac/FsSystem/BucketTree.cs +++ b/src/LibHac/FsSystem/BucketTree.cs @@ -1,4 +1,5 @@ using System; +using System.Buffers; using System.Buffers.Binary; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; @@ -384,7 +385,7 @@ namespace LibHac.FsSystem if (Entry == null) { - Entry = new byte[tree.EntrySize]; + Entry = ArrayPool.Shared.Rent((int)tree.EntrySize); Tree = tree; EntryIndex = -1; } @@ -394,7 +395,11 @@ namespace LibHac.FsSystem public void Dispose() { - // todo: try using shared arrays + if (Entry != null) + { + ArrayPool.Shared.Return(Entry); + Entry = null; + } } public bool IsValid() => EntryIndex >= 0; diff --git a/src/LibHac/FsSystem/IndirectStorage.cs b/src/LibHac/FsSystem/IndirectStorage.cs index 65664251..ba073e53 100644 --- a/src/LibHac/FsSystem/IndirectStorage.cs +++ b/src/LibHac/FsSystem/IndirectStorage.cs @@ -101,48 +101,53 @@ namespace LibHac.FsSystem // Find the offset in our tree var visitor = new BucketTree.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) + try { - // Try to write the entry to the out list - if (entryBuffer.Length != 0) + 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) { - if (count >= entryBuffer.Length) + // Try to write the entry to the out list + if (entryBuffer.Length != 0) + { + 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; - - 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; } - - // Write the entry count - outputEntryCount = count; - return Result.Success; + finally { visitor.Dispose(); } } protected override unsafe Result DoRead(long offset, Span destination) @@ -216,84 +221,88 @@ namespace LibHac.FsSystem // Find the offset in our tree var visitor = new BucketTree.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) + try { - // Get the current entry - var currentEntry = visitor.Get(); + Result rc = Table.Find(ref visitor, offset); + if (rc.IsFailure()) return rc; - // 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) + long entryOffset = visitor.Get().GetVirtualOffset(); + if (entryOffset < 0 || !Table.Includes(entryOffset)) return ResultFs.InvalidIndirectEntryStorageIndex.Log(); - // todo: Implement continuous reading + // Prepare to operate in chunks + long currentOffset = offset; + long endOffset = offset + size; - // Get and validate the next entry offset - long nextEntryOffset; - if (visitor.CanMoveNext()) + while (currentOffset < endOffset) { - rc = visitor.MoveNext(); - if (rc.IsFailure()) return rc; + // Get the current entry + var currentEntry = visitor.Get(); - nextEntryOffset = visitor.Get().GetVirtualOffset(); - if (!Table.Includes(nextEntryOffset)) + // 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; } - 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; } + finally { visitor.Dispose(); } return Result.Success; }