mirror of
https://github.com/Thealexbarney/LibHac.git
synced 2024-11-14 10:49:41 +01:00
Use the new BucketTree in Aes128CtrExStorage
This commit is contained in:
parent
33af34cefc
commit
d3abdeacc4
2 changed files with 56 additions and 35 deletions
|
@ -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<AesSubsectionEntry> SubsectionEntries { get; }
|
||||
private List<long> SubsectionOffsets { get; }
|
||||
private BucketTree<AesSubsectionEntry> 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<AesSubsectionEntry>(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<AesSubsectionEntry>(bucketTreeData);
|
||||
|
||||
SubsectionEntries = BucketTree.GetEntryList();
|
||||
SubsectionOffsets = SubsectionEntries.Select(x => x.Offset).ToList();
|
||||
Result rc = Table.Initialize(nodeStorage, entryStorage, NodeSize, Unsafe.SizeOf<Entry>(), entryCount);
|
||||
rc.ThrowIfFailure();
|
||||
}
|
||||
|
||||
protected override Result DoRead(long offset, Span<byte> 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<Entry>();
|
||||
|
||||
// Get and validate the next entry offset
|
||||
long nextEntryOffset;
|
||||
if (visitor.CanMoveNext())
|
||||
{
|
||||
rc = visitor.MoveNext();
|
||||
if (rc.IsFailure()) return rc;
|
||||
|
||||
nextEntryOffset = visitor.Get<Entry>().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;
|
||||
|
|
|
@ -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);
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue