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