Chunked archive readers for large files

CUE4Parse/FileProvider/AbstractFileProvider.cs

@@ -549,6 +549,10 @@ public virtual async Task<IPackage> LoadPackageAsync(GameFile file)
             Files.TryGetValue(file.PathWithoutExtension + ".uexp", out var uexpFile);
             Files.TryGetValue(file.PathWithoutExtension + ".ubulk", out var ubulkFile);
             Files.TryGetValue(file.PathWithoutExtension + ".uptnl", out var uptnlFile);
+            if (ubulkFile != null && !Globals.AllowLargeFiles && ubulkFile.Size >= Globals.LargeFileLimit)
+                ubulkFile = null;
+            if (uptnlFile != null && !Globals.AllowLargeFiles && uptnlFile.Size >= Globals.LargeFileLimit)
+                uptnlFile = null;
             var uassetTask = file.CreateReaderAsync();
             var uexpTask = uexpFile?.CreateReaderAsync();
             var ubulkTask = ubulkFile?.CreateReaderAsync();
@@ -589,6 +593,10 @@ public virtual async Task<IPackage> LoadPackageAsync(GameFile file)
             Files.TryGetValue(file.PathWithoutExtension + ".uexp", out var uexpFile);
             Files.TryGetValue(file.PathWithoutExtension + ".ubulk", out var ubulkFile);
             Files.TryGetValue(file.PathWithoutExtension + ".uptnl", out var uptnlFile);
+            if (ubulkFile != null && !Globals.AllowLargeFiles && ubulkFile.Size >= Globals.LargeFileLimit)
+                ubulkFile = null;
+            if (uptnlFile != null && !Globals.AllowLargeFiles && uptnlFile.Size >= Globals.LargeFileLimit)
+                uptnlFile = null;
             var uassetTask = file.TryCreateReaderAsync().ConfigureAwait(false);
             var uexpTask = uexpFile?.TryCreateReaderAsync().ConfigureAwait(false);
             var lazyUbulk = ubulkFile != null ? new Lazy<FArchive?>(() => ubulkFile.TryCreateReader(out var reader) ? reader : null) : null;
@@ -653,6 +661,10 @@ public virtual async Task<IReadOnlyDictionary<string, byte[]>> SavePackageAsync(
             Files.TryGetValue(file.PathWithoutExtension + ".uexp", out var uexpFile);
             Files.TryGetValue(file.PathWithoutExtension + ".ubulk", out var ubulkFile);
             Files.TryGetValue(file.PathWithoutExtension + ".uptnl", out var uptnlFile);
+            if (ubulkFile != null && !Globals.AllowLargeFiles && ubulkFile.Size >= Globals.LargeFileLimit)
+                ubulkFile = null;
+            if (uptnlFile != null && !Globals.AllowLargeFiles && uptnlFile.Size >= Globals.LargeFileLimit)
+                uptnlFile = null;
             var uassetTask = file.ReadAsync();
             var uexpTask = uexpFile?.ReadAsync();
             var ubulkTask = ubulkFile?.ReadAsync();
@@ -688,6 +700,10 @@ public virtual async Task<IReadOnlyDictionary<string, byte[]>> SavePackageAsync(
             Files.TryGetValue(file.PathWithoutExtension + ".uexp", out var uexpFile);
             Files.TryGetValue(file.PathWithoutExtension + ".ubulk", out var ubulkFile);
             Files.TryGetValue(file.PathWithoutExtension + ".uptnl", out var uptnlFile);
+            if (ubulkFile != null && !Globals.AllowLargeFiles && ubulkFile.Size >= Globals.LargeFileLimit)
+                ubulkFile = null;
+            if (uptnlFile != null && !Globals.AllowLargeFiles && uptnlFile.Size >= Globals.LargeFileLimit)
+                uptnlFile = null;
             var uassetTask = file.TryReadAsync().ConfigureAwait(false);
             var uexpTask = uexpFile?.TryReadAsync().ConfigureAwait(false);
             var ubulkTask = ubulkFile?.TryReadAsync().ConfigureAwait(false);

CUE4Parse/Globals.cs

@@ -9,5 +9,8 @@ public static class Globals
         public static bool LogVfsMounts = true;
         public static bool FatalObjectSerializationErrors = false;
         public static bool WarnMissingImportPackage = true;
+        public static bool AlwaysUseChunkedReader = false;
+        public static bool AllowLargeFiles = false;
+        public static long LargeFileLimit = int.MaxValue;
     }
-}
+}

CUE4Parse/UE4/IO/IoStoreReader.cs

@@ -199,7 +199,7 @@ public virtual byte[] Read(FIoChunkId chunkId)
             throw new KeyNotFoundException($"Couldn't find chunk {chunkId} in IoStore {Name}");
         }
 
-        private byte[] Read(long offset, long length)
+        public byte[] Read(long offset, long length)
         {
             var compressionBlockSize = TocResource.Header.CompressionBlockSize;
             var dst = new byte[length];

CUE4Parse/UE4/IO/Objects/FIoStoreEntry.cs

@@ -1,4 +1,5 @@
-using System.Runtime.CompilerServices;
+using System;
+using System.Runtime.CompilerServices;
 using CUE4Parse.Compression;
 using CUE4Parse.UE4.Readers;
 using CUE4Parse.UE4.VirtualFileSystem;
@@ -39,6 +40,6 @@ public IoStoreReader IoStoreReader
         public override byte[] Read() => Vfs.Extract(this);
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public override FArchive CreateReader() => new FByteArchive(Path, Read(), Vfs.Versions);
+        public override FArchive CreateReader() => Globals.AlwaysUseChunkedReader || Size > Math.Min(Globals.LargeFileLimit, int.MaxValue) ? new FIoChunkArchive(Path, this, Vfs.Versions) :  new FByteArchive(Path, Read(), Vfs.Versions);
     }
 }

CUE4Parse/UE4/Pak/Objects/FPakEntry.cs

@@ -328,6 +328,6 @@ public PakFileReader PakFileReader
         public override byte[] Read() => Vfs.Extract(this);
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public override FArchive CreateReader() => new FByteArchive(Path, Read(), Vfs.Versions);
+        public override FArchive CreateReader() => Globals.AlwaysUseChunkedReader || Size > Math.Min(Globals.LargeFileLimit, int.MaxValue) ? new FPakChunkArchive(Path, this, Vfs.Versions) :  new FByteArchive(Path, Read(), Vfs.Versions);
     }
 }

CUE4Parse/UE4/Pak/PakFileReader.cs

@@ -55,40 +55,7 @@ public PakFileReader(string filePath, Stream stream, VersionContainer? versions
         public override byte[] Extract(VfsEntry entry)
         {
             if (entry is not FPakEntry pakEntry || entry.Vfs != this) throw new ArgumentException($"Wrong pak file reader, required {entry.Vfs.Name}, this is {Name}");
-            // If this reader is used as a concurrent reader create a clone of the main reader to provide thread safety
-            var reader = IsConcurrent ? (FArchive) Ar.Clone() : Ar;
-            if (pakEntry.IsCompressed)
-            {
-#if DEBUG
-                Log.Debug($"{pakEntry.Name} is compressed with {pakEntry.CompressionMethod}");
-#endif
-                var uncompressed = new byte[(int) pakEntry.UncompressedSize];
-                var uncompressedOff = 0;
-                foreach (var block in pakEntry.CompressionBlocks)
-                {
-                    reader.Position = block.CompressedStart;
-                    var blockSize = (int) block.Size;
-                    var srcSize = blockSize.Align(pakEntry.IsEncrypted ? Aes.ALIGN : 1);
-                    // Read the compressed block
-                    byte[] compressed = ReadAndDecrypt(srcSize, reader, pakEntry.IsEncrypted);
-                    // Calculate the uncompressed size,
-                    // its either just the compression block size
-                    // or if its the last block its the remaining data size
-                    var uncompressedSize = (int) Math.Min(pakEntry.CompressionBlockSize, pakEntry.UncompressedSize - uncompressedOff);
-                    Decompress(compressed, 0, blockSize, uncompressed, uncompressedOff, uncompressedSize, pakEntry.CompressionMethod);
-                    uncompressedOff += (int) pakEntry.CompressionBlockSize;
-                }
-
-                return uncompressed;
-            }
-
-            // Pak Entry is written before the file data,
-            // but its the same as the one from the index, just without a name
-            // We don't need to serialize that again so + file.StructSize
-            reader.Position = pakEntry.Offset + pakEntry.StructSize; // Doesn't seem to be the case with older pak versions
-            var size = (int) pakEntry.UncompressedSize.Align(pakEntry.IsEncrypted ? Aes.ALIGN : 1);
-            var data = ReadAndDecrypt(size, reader, pakEntry.IsEncrypted);
-            return size != pakEntry.UncompressedSize ? data.SubByteArray((int) pakEntry.UncompressedSize) : data;
+            return Read(pakEntry, 0, pakEntry.UncompressedSize).ToArray();
         }
 
         public override IReadOnlyDictionary<string, GameFile> Mount(bool caseInsensitive = false)
@@ -304,5 +271,65 @@ public override void Dispose()
         {
             Ar.Dispose();
         }
+
+        public Span<byte> Read(FPakEntry pakEntry, long offset, long size)
+        {
+            // If this reader is used as a concurrent reader create a clone of the main reader to provide thread safety
+            var reader = IsConcurrent ? (FArchive)Ar.Clone() : Ar;
+            if (pakEntry.IsCompressed)
+            {
+#if DEBUG
+                Log.Debug($"{pakEntry.Name} is compressed with {pakEntry.CompressionMethod}");
+#endif
+                var uncompressed = new byte[(int)size + pakEntry.CompressionBlockSize];
+                var uncompressedOff = 0;
+                var shift = 0;
+                foreach (var block in pakEntry.CompressionBlocks)
+                {
+                    if (offset >= block.Size)
+                    {
+                        offset -= block.Size;
+                        continue;
+                    }
+
+                    if (offset != 0)
+                        shift = (int)offset;
+
+                    offset = 0;
+                    reader.Position = block.CompressedStart;
+                    var blockSize = (int)block.Size;
+                    var srcSize = blockSize.Align(pakEntry.IsEncrypted ? Aes.ALIGN : 1);
+                    // Read the compressed block
+                    byte[] compressed = ReadAndDecrypt(srcSize, reader, pakEntry.IsEncrypted);
+                    // Calculate the uncompressed size,
+                    // its either just the compression block size
+                    // or if its the last block its the remaining data size
+                    var uncompressedSize = (int)Math.Min(pakEntry.CompressionBlockSize, pakEntry.UncompressedSize - uncompressedOff);
+                    Decompress(compressed, 0, blockSize, uncompressed, uncompressedOff, uncompressedSize, pakEntry.CompressionMethod);
+                    uncompressedOff += (int)pakEntry.CompressionBlockSize;
+
+                    if (uncompressedOff >= size)
+                        break;
+                }
+
+                return uncompressed.AsSpan(shift, (int)size);
+            }
+
+            // Pak Entry is written before the file data,
+            // but its the same as the one from the index, just without a name
+            // We don't need to serialize that again so + file.StructSize
+            var readSize = (int)size.Align(pakEntry.IsEncrypted ? Aes.ALIGN : 1);
+            var readShift = 0;
+            if (pakEntry.IsEncrypted && offset.Align(Aes.ALIGN) != offset)
+            {
+                var tmp = offset.Align(Aes.ALIGN) - Aes.ALIGN;
+                readShift = (int)(offset - tmp);
+                offset = tmp;
+                readSize += Aes.ALIGN;
+            }
+
+            reader.Position = pakEntry.Offset + offset + pakEntry.StructSize; // Doesn't seem to be the case with older pak versions
+            return ReadAndDecrypt(readSize, reader, pakEntry.IsEncrypted).AsSpan(readShift, (int)size);
+        }
     }
 }

CUE4Parse/UE4/Readers/FChunkedArchive.cs

@@ -0,0 +1,95 @@
+using System;
+using System.Buffers;
+using System.IO;
+using System.Runtime.CompilerServices;
+using CUE4Parse.UE4.Versions;
+using CUE4Parse.Utils;
+
+namespace CUE4Parse.UE4.Readers;
+
+public abstract class FChunkedArchive : FArchive
+{
+    private const int BUFFER_SIZE = 1024 * 1024 * 128; // 128 MiB.
+
+    protected FChunkedArchive(string path, VersionContainer versions) : base(versions)
+    {
+        Name = path;
+        Buffer = ArrayPool<byte>.Shared.Rent(BUFFER_SIZE);
+    }
+
+    protected abstract Span<byte> ReadChunks(long offset, long size);
+
+    public override int Read(byte[] buffer, int offset, int count)
+    {
+        if (Position < 0)
+            return 0;
+
+        var n = (int)(Length - Position);
+        if (n > count) n = count;
+        if (n <= 0)
+            return 0;
+
+        Span<byte> data;
+        if (n < BUFFER_SIZE)
+        {
+            var bufferRangeStart = BufferOffset;
+            var bufferRangeEnd = BufferOffset + BUFFER_SIZE;
+            if (!(bufferRangeStart <= Position && Position <= bufferRangeEnd))
+            {
+                BufferOffset = Position;
+                var blockSize = Math.Min(BUFFER_SIZE, Length - Position).Align(BUFFER_SIZE);
+                if (Position.Align(BUFFER_SIZE) != Position)
+                    BufferOffset = Position.Align(BUFFER_SIZE) - BUFFER_SIZE;
+
+                if ((int) (Position - BufferOffset) + n <= BUFFER_SIZE) //overflow check
+                    ReadChunks(BufferOffset, blockSize).CopyTo(Buffer);
+                else
+                    BufferOffset = -1; // reset buffer position because we didn't actually read
+            }
+
+            if (BufferOffset == -1 || (int) (Position - BufferOffset) + n > BUFFER_SIZE)
+                data = ReadChunks(Position, n);
+            else
+                data = Buffer.AsSpan().Slice((int) (Position - BufferOffset), n);
+        }
+        else
+        {
+            data = ReadChunks(Position, n);
+        }
+
+        data.CopyTo(buffer.AsSpan(offset));
+        Position += n;
+        return n;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public override long Seek(long offset, SeekOrigin origin)
+    {
+        Position = origin switch
+        {
+            SeekOrigin.Begin => offset,
+            SeekOrigin.Current => Position + offset,
+            SeekOrigin.End => Length + offset,
+            _ => throw new ArgumentOutOfRangeException()
+        };
+        return Position;
+    }
+
+    public override bool CanSeek => true;
+    public override long Position { get; set; }
+    public override string Name { get; }
+    private byte[] Buffer { get; }
+    private long BufferOffset { get; set; } = -BUFFER_SIZE - 1;
+
+    protected void ImportBuffer(FChunkedArchive other)
+    {
+        BufferOffset = other.BufferOffset;
+        other.Buffer.AsSpan().CopyTo(Buffer.AsSpan());
+    }
+
+    public override void Close()
+    {
+        base.Close();
+        ArrayPool<byte>.Shared.Return(Buffer);
+    }
+}

CUE4Parse/UE4/Readers/FIoChunkArchive.cs

@@ -0,0 +1,32 @@
+using System;
+using CUE4Parse.UE4.IO.Objects;
+using CUE4Parse.UE4.Versions;
+
+namespace CUE4Parse.UE4.Readers;
+
+public class FIoChunkArchive : FChunkedArchive
+{
+    public FIoChunkArchive(string path, FIoStoreEntry entry, VersionContainer versions) : base(path, versions)
+    {
+        Entry = entry;
+    }
+
+    public override long Length => Entry.Size;
+    public FIoStoreEntry Entry { get; }
+
+    public override object Clone()
+    {
+        var instance = new FIoChunkArchive(Name, Entry, Versions)
+        {
+            Position = Position,
+        };
+        instance.ImportBuffer(this);
+        return instance;
+    }
+
+    protected override Span<byte> ReadChunks(long offset, long size)
+    {
+        var remaining = Math.Min(size, Entry.Size - offset);
+        return remaining <= 0 ? Span<byte>.Empty : Entry.IoStoreReader.Read(offset + Entry.Offset, remaining);
+    }
+}

CUE4Parse/UE4/Readers/FPakChunkArchive.cs

@@ -0,0 +1,32 @@
+using System;
+using CUE4Parse.UE4.Pak.Objects;
+using CUE4Parse.UE4.Versions;
+
+namespace CUE4Parse.UE4.Readers;
+
+public class FPakChunkArchive : FChunkedArchive
+{
+    public FPakChunkArchive(string path, FPakEntry entry, VersionContainer versions) : base(path, versions)
+    {
+        Entry = entry;
+    }
+
+    public override long Length => Entry.UncompressedSize;
+    public FPakEntry Entry { get; }
+
+    public override object Clone()
+    {
+        var instance = new FPakChunkArchive(Name, Entry, Versions)
+        {
+            Position = Position,
+        };
+        instance.ImportBuffer(this);
+        return instance;
+    }
+
+    protected override Span<byte> ReadChunks(long offset, long size)
+    {
+        var remaining = Math.Min(size, Entry.UncompressedSize - offset);
+        return remaining <= 0 ? Span<byte>.Empty : Entry.PakFileReader.Read(Entry, offset, remaining);
+    }
+}