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); + } + } + } +}