diff --git a/ValvePak/ValvePak.Test/MemoryMapTest.cs b/ValvePak/ValvePak.Test/MemoryMapTest.cs index 9a5c33f..fd85252 100644 --- a/ValvePak/ValvePak.Test/MemoryMapTest.cs +++ b/ValvePak/ValvePak.Test/MemoryMapTest.cs @@ -1,6 +1,5 @@ using System; using System.IO; -using System.IO.MemoryMappedFiles; using System.Security.Cryptography; using System.Threading.Tasks; @@ -20,7 +19,7 @@ public async Task ReturnsCorrectStreamsForSplitPackages() package.Read(path); using var stream = package.GetMemoryMappedStreamIfPossible(package.FindEntry("kitten.jpg")!); - await Assert.That(stream).IsAssignableTo(); + await Assert.That(stream).IsNotAssignableTo(); await VerifyKitten(stream); using var stream2 = package.GetMemoryMappedStreamIfPossible(package.FindEntry("steammessages_base.proto")!); @@ -37,7 +36,7 @@ public async Task ReturnsCorrectStreams() package.Read(path); using var stream = package.GetMemoryMappedStreamIfPossible(package.FindEntry("kitten.jpg")!); - await Assert.That(stream).IsAssignableTo(); + await Assert.That(stream).IsNotAssignableTo(); await VerifyKitten(stream); using var stream2 = package.GetMemoryMappedStreamIfPossible(package.FindEntry("steammessages_base.proto")!); @@ -56,7 +55,7 @@ public async Task ReturnsCorrectStreamsWhenUsingFileStream() package.Read(fileStream); using var stream = package.GetMemoryMappedStreamIfPossible(package.FindEntry("kitten.jpg")!); - await Assert.That(stream).IsAssignableTo(); + await Assert.That(stream).IsNotAssignableTo(); await VerifyKitten(stream); using var stream2 = package.GetMemoryMappedStreamIfPossible(package.FindEntry("steammessages_base.proto")!); diff --git a/ValvePak/ValvePak.Test/PackageTest.cs b/ValvePak/ValvePak.Test/PackageTest.cs index 59a0f9d..bdebb1f 100644 --- a/ValvePak/ValvePak.Test/PackageTest.cs +++ b/ValvePak/ValvePak.Test/PackageTest.cs @@ -373,6 +373,63 @@ public async Task ExtractIntoUserProvidedByteArray() await Assert.That(() => package.ReadEntry(entry, smallBuffer)).ThrowsExactly(); } + [Test] + public async Task ThrowsWhenEntryExtendsPastEndOfArchive() + { + var directory = Path.Combine(Path.GetTempPath(), "ValvePakTest_" + Path.GetRandomFileName()); + Directory.CreateDirectory(directory); + + try + { + // Cut off the last entry (steammessages_clientserver.proto, 39177 bytes) and the end of the one before it (steammessages_base.proto, 2563 bytes) + const int Cut = 39177 + 100; + + File.Copy(Path.Combine(AppContext.BaseDirectory, "Files", "test_dir.vpk"), Path.Combine(directory, "test_dir.vpk")); + var chunk = await File.ReadAllBytesAsync(Path.Combine(AppContext.BaseDirectory, "Files", "test_000.vpk")); + await File.WriteAllBytesAsync(Path.Combine(directory, "test_000.vpk"), chunk[..^Cut]); + + var single = await File.ReadAllBytesAsync(Path.Combine(AppContext.BaseDirectory, "Files", "test_single.vpk")); + await File.WriteAllBytesAsync(Path.Combine(directory, "test_single.vpk"), single[..^Cut]); + + foreach (var (name, fromMemory) in new[] { ("test_dir.vpk", false), ("test_single.vpk", false), ("test_single.vpk", true) }) + { + var path = Path.Combine(directory, name); + + using var package = new Package(); + + if (fromMemory) + { + package.SetFileName(path); + package.Read(new MemoryStream(await File.ReadAllBytesAsync(path))); + } + else + { + package.Read(path); + } + + var kitten = package.FindEntry("kitten.jpg")!; + await Assert.That(() => package.ReadEntry(kitten, out _, validateCrc: true)).ThrowsNothing(); + + foreach (var truncatedName in new[] { "steammessages_base.proto", "steammessages_clientserver.proto" }) + { + var truncated = package.FindEntry(truncatedName)!; + + using (Assert.Multiple()) + { + await Assert.That(() => package.ReadEntry(truncated, out _, validateCrc: false)).ThrowsExactly(); + await Assert.That(() => package.ReadEntry(truncated, out _, validateCrc: true)).ThrowsExactly(); + await Assert.That(() => package.ReadEntry(truncated, new byte[truncated.TotalLength], validateCrc: false)).ThrowsExactly(); + await Assert.That(() => package.GetMemoryMappedStreamIfPossible(truncated)).ThrowsExactly(); + } + } + } + } + finally + { + Directory.Delete(directory, recursive: true); + } + } + [Test] public async Task TestFileChecksums() { diff --git a/ValvePak/ValvePak/ArchiveFile.cs b/ValvePak/ValvePak/ArchiveFile.cs new file mode 100644 index 0000000..34f43fb --- /dev/null +++ b/ValvePak/ValvePak/ArchiveFile.cs @@ -0,0 +1,108 @@ +using System; +using System.IO; +using System.IO.MemoryMappedFiles; +using System.Threading; +using Microsoft.Win32.SafeHandles; + +namespace ValvePak +{ + /// + /// A file containing archive data of a package, kept open and memory mapped on first use. + /// + internal sealed class ArchiveFile : IDisposable + { + private readonly bool OwnsHandle; + private readonly Lock MappingLock = new(); + private bool IsMapped; + private MemoryMappedFile? Mapping; + + /// + /// Initializes a new instance of the class. + /// + /// Handle of the file, opened for reading. + /// Whether the handle is closed when this archive is disposed. + public ArchiveFile(SafeFileHandle handle, bool ownsHandle) + { + Handle = handle; + OwnsHandle = ownsHandle; + } + + /// + /// Gets the file handle. + /// + public SafeFileHandle Handle { get; } + + /// + /// Gets the length in bytes of the file at the time it was mapped. + /// + public long MappedLength { get; private set; } + + /// + /// Gets the view of the whole mapped file, only created in 64-bit processes. + /// + public MemoryMappedViewAccessor? View { get; private set; } + + /// + /// Gets the memory mapping of the file, creating it on first use. + /// + /// The mapping, or null if the file is empty. + public MemoryMappedFile? GetMapping() + { + if (!Volatile.Read(ref IsMapped)) + { + lock (MappingLock) + { + if (!IsMapped) + { + CreateMapping(); + Volatile.Write(ref IsMapped, true); + } + } + } + + return Mapping; + } + + private void CreateMapping() + { + var length = RandomAccess.GetLength(Handle); + + if (length == 0) + { + return; + } + + var mapping = MemoryMappedFile.CreateFromFile(Handle, null, 0, MemoryMappedFileAccess.Read, HandleInheritability.None, leaveOpen: true); + + if (Environment.Is64BitProcess) + { + try + { + View = mapping.CreateViewAccessor(0, length, MemoryMappedFileAccess.Read); + } + catch + { + mapping.Dispose(); + throw; + } + } + + Mapping = mapping; + MappedLength = length; + } + + /// + /// Releases the view and the mapping, and closes the handle if it is owned. + /// + public void Dispose() + { + View?.Dispose(); + Mapping?.Dispose(); + + if (OwnsHandle) + { + Handle.Dispose(); + } + } + } +} diff --git a/ValvePak/ValvePak/Package.Read.cs b/ValvePak/ValvePak/Package.Read.cs index 542d885..96e4369 100644 --- a/ValvePak/ValvePak/Package.Read.cs +++ b/ValvePak/ValvePak/Package.Read.cs @@ -1,4 +1,5 @@ using System; +using System.Buffers.Binary; using System.Collections.Generic; using System.Diagnostics; using System.IO; @@ -21,6 +22,8 @@ public partial class Package : IDisposable /// The file to open and read. public void Read(string filename) { + ObjectDisposedException.ThrowIf(IsDisposed, this); + SetFileName(filename); var fs = new FileStream(filename, FileMode.Open, FileAccess.Read, FileShare.ReadWrite); @@ -30,17 +33,25 @@ public void Read(string filename) /// /// Reads the given . + /// If this package was read before, everything read from the previous package is replaced. + /// + /// The stream must stay open while the package is used, entry data is read from it on demand. + /// It is disposed together with the package. /// /// The input to read from. public void Read(Stream input) { ArgumentNullException.ThrowIfNull(input); + ObjectDisposedException.ThrowIf(IsDisposed, this); if (FileName == null) { throw new InvalidOperationException("If you call Read() directly with a stream, you must call SetFileName() first."); } + ResetReadState(input); + + StreamOffset = input.Position; Reader = new BinaryReader(input); if (Reader.ReadUInt32() != MAGIC) @@ -73,7 +84,7 @@ public void Read(Stream input) Version = version; - HeaderSize = (uint)input.Position; + HeaderSize = (uint)(input.Position - StreamOffset); ReadEntries(); @@ -90,13 +101,16 @@ public void Read(Stream input) /// /// Reads the entry from the VPK package. + /// This method is safe to call from multiple threads. /// /// Package entry. /// Output buffer. /// If true, CRC32 will be calculated and verified for read data. + /// The archive ends before all bytes of the entry could be read. public void ReadEntry(PackageEntry entry, out byte[] output, bool validateCrc = true) { ArgumentNullException.ThrowIfNull(entry); + ObjectDisposedException.ThrowIf(IsDisposed, this); output = new byte[entry.TotalLength]; @@ -105,14 +119,17 @@ public void ReadEntry(PackageEntry entry, out byte[] output, bool validateCrc = /// /// Reads the entry from the VPK package into a user-provided output byte array. + /// This method is safe to call from multiple threads. /// /// Package entry. /// Output buffer, size of the buffer must be at least . /// If true, CRC32 will be calculated and verified for read data. + /// The archive ends before all bytes of the entry could be read. public void ReadEntry(PackageEntry entry, byte[] output, bool validateCrc = true) { ArgumentNullException.ThrowIfNull(entry); ArgumentNullException.ThrowIfNull(output); + ObjectDisposedException.ThrowIf(IsDisposed, this); var totalLength = (int)entry.TotalLength; @@ -128,29 +145,12 @@ public void ReadEntry(PackageEntry entry, byte[] output, bool validateCrc = true if (entry.Length > 0) { -#pragma warning disable CA2000 // Dispose objects before losing scope - var fs = GetFileStream(entry.ArchiveIndex); -#pragma warning restore CA2000 // Stream is base reader stream when reading from non-split vpk, it should not be disposed + var length = (int)entry.Length; + var bytesRead = ReadArchiveData(entry.ArchiveIndex, entry.Offset, output.AsSpan(entry.SmallData.Length, length)); - try - { - fs.Seek(entry.Offset, SeekOrigin.Current); - - int length = (int)entry.Length; - int readOffset = entry.SmallData.Length; - int bytesRead; - int totalRead = 0; - while ((bytesRead = fs.Read(output, readOffset + totalRead, length - totalRead)) != 0) - { - totalRead += bytesRead; - } - } - finally + if (bytesRead != length) { - if (entry.ArchiveIndex != 0x7FFF) - { - fs.Dispose(); - } + throw new EndOfStreamException($"Reached the end of archive {entry.ArchiveIndex} after reading {bytesRead} of {length} bytes for '{entry.GetFullPath()}'."); } } @@ -167,12 +167,153 @@ public void ReadEntry(PackageEntry entry, byte[] output, bool validateCrc = true } } + /// + /// Largest directory tree in bytes that is read into a single buffer, larger trees are parsed from the stream. + /// + private const uint MaxTreeBufferSize = 64 * 1024 * 1024; + private void ReadEntries() { Debug.Assert(Reader != null); - var stringComparer = Comparer == null ? null : StringComparer.FromComparison(Comparer.Comparison); - var typeEntries = new Dictionary>(stringComparer); + var input = Reader.BaseStream; + var treeStart = input.Position; + + if (TreeSize > 0 && TreeSize <= MaxTreeBufferSize && input.CanSeek && TreeSize <= input.Length - treeStart) + { + var tree = GC.AllocateUninitializedArray((int)TreeSize); + var treeLength = input.ReadAtLeast(tree, tree.Length, throwOnEndOfStream: false); + + if (TryReadEntries(tree.AsSpan(0, treeLength), out var typeEntries, out var treeSize)) + { + Entries = typeEntries; + TreeSize = (uint)treeSize; + input.Position = treeStart + treeSize; + return; + } + + input.Position = treeStart; + } + + ReadEntriesFromStream(); + } + + /// + /// Parses the directory tree from a buffer. + /// Returns false if the buffer ends before the tree does, in which case it has to be read from the stream. + /// + private bool TryReadEntries(ReadOnlySpan tree, out Dictionary> typeEntries, out int treeSize) + { + const int EntryMetadataSize = 18; + + typeEntries = CreateEntriesDictionary(); + treeSize = 0; + + var position = 0; + + // Types + while (true) + { + if (!TryReadNullTermUtf8String(tree, ref position, out var typeName)) + { + return false; + } + + if (typeName.Length == 0) + { + break; + } + + var entries = new List(); + + // Directories + while (true) + { + if (!TryReadNullTermUtf8String(tree, ref position, out var directoryName)) + { + return false; + } + + if (directoryName.Length == 0) + { + break; + } + + // Files + while (true) + { + if (!TryReadNullTermUtf8String(tree, ref position, out var fileName)) + { + return false; + } + + if (fileName.Length == 0) + { + break; + } + + if (tree.Length - position < EntryMetadataSize) + { + return false; + } + + var metadata = tree.Slice(position, EntryMetadataSize); + position += EntryMetadataSize; + + var terminator = BinaryPrimitives.ReadUInt16LittleEndian(metadata[16..]); + + if (terminator != 0xFFFF) + { + throw new FormatException($"Invalid terminator, was 0x{terminator:X} but expected 0x{0xFFFF:X}."); + } + + var entry = new PackageEntry + { + FileName = fileName, + DirectoryName = directoryName, + TypeName = typeName, + CRC32 = BinaryPrimitives.ReadUInt32LittleEndian(metadata), + ArchiveIndex = BinaryPrimitives.ReadUInt16LittleEndian(metadata[6..]), + Offset = BinaryPrimitives.ReadUInt32LittleEndian(metadata[8..]), + Length = BinaryPrimitives.ReadUInt32LittleEndian(metadata[12..]), + }; + + var smallDataSize = BinaryPrimitives.ReadUInt16LittleEndian(metadata[4..]); + + if (smallDataSize > 0) + { + if (tree.Length - position < smallDataSize) + { + return false; + } + + entry.SmallData = tree.Slice(position, smallDataSize).ToArray(); + position += smallDataSize; + } + + entries.Add(entry); + } + } + + if (Comparer != null) + { + // Sorting at the end is faster than doing BinarySearch+Insert + entries.Sort(Comparer); + } + + typeEntries.Add(typeName, entries); + } + + treeSize = position; + + return true; + } + + private void ReadEntriesFromStream() + { + Debug.Assert(Reader != null); + + var typeEntries = CreateEntriesDictionary(); using var ms = new MemoryStream(); // Types @@ -229,13 +370,7 @@ private void ReadEntries() if (smallDataSize > 0) { entry.SmallData = new byte[smallDataSize]; - - int bytesRead; - int totalRead = 0; - while ((bytesRead = Reader.Read(entry.SmallData, totalRead, entry.SmallData.Length - totalRead)) != 0) - { - totalRead += bytesRead; - } + Reader.BaseStream.ReadExactly(entry.SmallData); } entries.Add(entry); @@ -254,14 +389,14 @@ private void ReadEntries() Entries = typeEntries; // Set to real size that was read for hash verification, in case it was tampered with - TreeSize = (uint)Reader.BaseStream.Position - HeaderSize; + TreeSize = (uint)(Reader.BaseStream.Position - StreamOffset - HeaderSize); } private void ReadArchiveMD5Section() { Debug.Assert(Reader != null); - FileSizeBeforeArchiveMD5Entries = (uint)Reader.BaseStream.Position; + FileSizeBeforeArchiveMD5Entries = (uint)(Reader.BaseStream.Position - StreamOffset); if (ArchiveMD5SectionSize == 0) { @@ -269,29 +404,67 @@ private void ReadArchiveMD5Section() return; } + if (ArchiveMD5SectionSize % ChunkHashFraction.SectionEntrySize != 0) + { + throw new InvalidDataException($"Archive MD5 section size ({ArchiveMD5SectionSize}) is not a multiple of {ChunkHashFraction.SectionEntrySize}."); + } + var entries = (int)(ArchiveMD5SectionSize / ChunkHashFraction.SectionEntrySize); AccessPackFileHashes = new List(entries); + var input = Reader.BaseStream; + var sectionStart = input.Position; + var sectionLength = (long)entries * ChunkHashFraction.SectionEntrySize; + + if (sectionLength <= Array.MaxLength && input.CanSeek && sectionLength <= input.Length - sectionStart) + { + var section = GC.AllocateUninitializedArray((int)sectionLength); + + if (input.ReadAtLeast(section, section.Length, throwOnEndOfStream: false) == section.Length) + { + for (var i = 0; i < entries; i++) + { + var data = section.AsSpan(i * ChunkHashFraction.SectionEntrySize, ChunkHashFraction.SectionEntrySize); + + AddChunkHashFraction(new ChunkHashFraction + { + ArchiveIndex = BinaryPrimitives.ReadUInt16LittleEndian(data), + HashType = (EHashType)BinaryPrimitives.ReadUInt16LittleEndian(data[2..]), + Offset = BinaryPrimitives.ReadUInt32LittleEndian(data[4..]), + Length = BinaryPrimitives.ReadUInt32LittleEndian(data[8..]), + Checksum = data[12..].ToArray() + }); + } + + return; + } + + input.Position = sectionStart; + } + for (var i = 0; i < entries; i++) { - var hashFraction = new ChunkHashFraction + AddChunkHashFraction(new ChunkHashFraction { ArchiveIndex = Reader.ReadUInt16(), HashType = (EHashType)Reader.ReadUInt16(), Offset = Reader.ReadUInt32(), Length = Reader.ReadUInt32(), Checksum = Reader.ReadBytes(16) - }; - - if (hashFraction.ArchiveIndex == 0 && hashFraction.HashType == (EHashType)0x8000) - { - hashFraction.ArchiveIndex = 0x7FFF; - hashFraction.HashType = EHashType.MD5; - } + }); + } + } - AccessPackFileHashes.Add(hashFraction); + private void AddChunkHashFraction(ChunkHashFraction hashFraction) + { + if (hashFraction.ArchiveIndex == 0 && hashFraction.HashType == (EHashType)0x8000) + { + hashFraction.ArchiveIndex = 0x7FFF; + hashFraction.HashType = EHashType.MD5; } + + AccessPackFileHashes.Add(hashFraction); } private void ReadOtherMD5Section() @@ -305,7 +478,7 @@ private void ReadOtherMD5Section() TreeChecksum = Reader.ReadBytes(16); ArchiveMD5EntriesChecksum = Reader.ReadBytes(16); - FileSizeBeforeWholeFileHash = (uint)Reader.BaseStream.Position; + FileSizeBeforeWholeFileHash = (uint)(Reader.BaseStream.Position - StreamOffset); WholeFileChecksum = Reader.ReadBytes(16); } @@ -313,7 +486,7 @@ private void ReadSignatureSection() { Debug.Assert(Reader != null); - FileSizeBeforeSignature = (uint)Reader.BaseStream.Position; + FileSizeBeforeSignature = (uint)(Reader.BaseStream.Position - StreamOffset); if (SignatureSectionSize == 0) { @@ -329,105 +502,194 @@ private void ReadSignatureSection() var signatureSize2 = Reader.ReadInt32(); Reader.ReadInt32(); // reserved? - if (publicKeySize > 0) + if (publicKeySize != 0) { - PublicKey = Reader.ReadBytes(publicKeySize); + PublicKey = ReadSignatureBytes(publicKeySize, "public key"); } - if (signatureSize2 > 0) + if (signatureSize2 != 0) { - Signature = Reader.ReadBytes(signatureSize2); + Signature = ReadSignatureBytes(signatureSize2, "signature"); } return; } SignatureType = ESignatureType.FullFile; - PublicKey = Reader.ReadBytes(publicKeySize); + PublicKey = ReadSignatureBytes(publicKeySize, "public key"); var signatureSize = Reader.ReadInt32(); - Signature = Reader.ReadBytes(signatureSize); + Signature = ReadSignatureBytes(signatureSize, "signature"); } - private Stream GetFileStream(ushort archiveIndex) + private byte[] ReadSignatureBytes(int size, string name) { - Stream stream; + Debug.Assert(Reader != null); + + var input = Reader.BaseStream; - if (archiveIndex != 0x7FFF) + if (size < 0 || (input.CanSeek && size > input.Length - input.Position)) { - var fileName = GetArchiveIndexFullFilePath(FileName, archiveIndex); + throw new InvalidDataException($"Invalid {name} size in the signature section ({size})."); + } - stream = new FileStream(fileName, FileMode.Open, FileAccess.Read); + return Reader.ReadBytes(size); + } + + /// + /// Returns the archive for the given index, opening its file on first use. + /// The embedded archive (0x7FFF) shares the handle of the the package was read from, + /// and is null if the package was read from any other kind of stream. + /// + private ArchiveFile? GetArchive(ushort archiveIndex) + { + if (Archives.TryGetValue(archiveIndex, out var archive)) + { + return archive; } - else + + if (archiveIndex == 0x7FFF && Reader?.BaseStream is not FileStream) { - Debug.Assert(Reader != null); + ObjectDisposedException.ThrowIf(IsDisposed, this); - stream = Reader.BaseStream; - stream.Seek(HeaderSize + TreeSize, SeekOrigin.Begin); + return null; } - return stream; + lock (ArchivesLock) + { + ObjectDisposedException.ThrowIf(IsDisposed, this); + + if (Archives.TryGetValue(archiveIndex, out archive)) + { + return archive; + } + + if (archiveIndex == 0x7FFF) + { + if (Reader?.BaseStream is not FileStream fileStream) + { + return null; + } + + archive = new ArchiveFile(fileStream.SafeFileHandle, ownsHandle: false); + } + else + { + var fileName = GetArchiveIndexFullFilePath(FileName, archiveIndex); + var handle = File.OpenHandle(fileName, FileMode.Open, FileAccess.Read, FileShare.ReadWrite); + + archive = new ArchiveFile(handle, ownsHandle: true); + } + + Archives[archiveIndex] = archive; + + return archive; + } + } + + /// + /// Returns the offset in the archive file of data at the given offset in the archive data section. + /// The package itself may start at a non-zero position of the stream it was read from. + /// + private long GetArchiveDataOffset(ushort archiveIndex, long offset) + { + if (archiveIndex == 0x7FFF) + { + return StreamOffset + offset + HeaderSize + TreeSize; + } + + return offset; + } + + /// + /// Reads archive data at the given offset until the output is full or the end of the archive is reached. + /// This method is safe to call from multiple threads. + /// + /// The number of bytes read. + private int ReadArchiveData(ushort archiveIndex, long offset, Span output) + { + var fileOffset = GetArchiveDataOffset(archiveIndex, offset); + var archive = GetArchive(archiveIndex); + var totalRead = 0; + int bytesRead; + + if (archive != null) + { + while (totalRead < output.Length && (bytesRead = RandomAccess.Read(archive.Handle, output[totalRead..], fileOffset + totalRead)) != 0) + { + totalRead += bytesRead; + } + + return totalRead; + } + + var reader = Reader; + ObjectDisposedException.ThrowIf(reader == null, this); + + lock (reader) + { + var stream = reader.BaseStream; + stream.Position = fileOffset; + + while (totalRead < output.Length && (bytesRead = stream.Read(output[totalRead..])) != 0) + { + totalRead += bytesRead; + } + } + + return totalRead; } /// - /// Returns when possible, otherwise reads entry into a byte array and returns . + /// Returns a stream that reads the entry directly from a memory mapping of its archive when possible, + /// otherwise reads entry into a byte array and returns . /// This only works when entries have no preload bytes. - /// Files smaller or equal to 4096 bytes will be read into memory. + /// Files smaller or equal to 4096 bytes will be copied into a . /// /// When trying to map a file in a non-split packages (), /// this will only memory map if the package was read by providing or using . + /// + /// Each archive is mapped once and stays mapped until the package is disposed. + /// Returned streams can only be read while the package is open, they throw + /// once the package is disposed or read again. + /// This method is safe to call from multiple threads. /// /// Package entry. /// Stream for a given package entry contents. + /// The archive ends before all bytes of the entry could be read. public Stream GetMemoryMappedStreamIfPossible(PackageEntry entry) { ArgumentNullException.ThrowIfNull(entry); + ObjectDisposedException.ThrowIf(IsDisposed, this); - if (entry.Length <= 4096 || entry.SmallData.Length > 0) - { - ReadEntry(entry, out var output, false); - return new MemoryStream(output); - } - - if (!MemoryMappedPaks.TryGetValue(entry.ArchiveIndex, out var stream)) + if (entry.Length > 0 && entry.SmallData.Length == 0 && GetArchive(entry.ArchiveIndex) is { } archive) { - string path; + var offset = GetArchiveDataOffset(entry.ArchiveIndex, entry.Offset); + var mapping = archive.GetMapping(); + var isInRange = mapping != null && offset + entry.Length <= archive.MappedLength; - // If the package was opened by providing a file path, we can memory map non directory files - if (entry.ArchiveIndex == 0x7FFF) + if (isInRange && archive.View is { } view) { - Debug.Assert(Reader != null); - - if (Reader.BaseStream is FileStream fileStream) - { - path = fileStream.Name; - } - else + if (entry.Length > 4096) { - // This package was read by in a stream that is not FileStream, or not using Read(path). - // In this case fallback to reading the file into memory. We could rely on the what was passed into SetFileName, but why. - ReadEntry(entry, out var output, false); - return new MemoryStream(output); + return new UnmanagedMemoryStream(view.SafeMemoryMappedViewHandle, view.PointerOffset + offset, entry.Length); } + + var data = new byte[entry.Length]; + view.SafeMemoryMappedViewHandle.ReadSpan((ulong)(view.PointerOffset + offset), data.AsSpan()); + + return new MemoryStream(data); } - else + + if (isInRange && entry.Length > 4096) { - path = GetArchiveIndexFullFilePath(FileName, entry.ArchiveIndex); + return mapping!.CreateViewStream(offset, entry.Length, MemoryMappedFileAccess.Read); } - - stream = MemoryMappedFile.CreateFromFile(path, FileMode.Open, null, 0, MemoryMappedFileAccess.Read); - MemoryMappedPaks[entry.ArchiveIndex] = stream; } - var offset = entry.Offset; + ReadEntry(entry, out var output, false); - if (entry.ArchiveIndex == 0x7FFF) - { - offset += HeaderSize + TreeSize; - } - - return stream.CreateViewStream(offset, entry.Length, MemoryMappedFileAccess.Read); + return new MemoryStream(output); } [MethodImpl(MethodImplOptions.AggressiveInlining)] @@ -436,6 +698,22 @@ private static string GetArchiveIndexFullFilePath(string? baseFileName, ushort a return $"{baseFileName}_{archiveIndex:D3}.vpk"; } + private static bool TryReadNullTermUtf8String(ReadOnlySpan buffer, ref int position, out string value) + { + var length = buffer[position..].IndexOf((byte)0); + + if (length < 0) + { + value = string.Empty; + return false; + } + + value = Encoding.UTF8.GetString(buffer.Slice(position, length)); + position += length + 1; + + return true; + } + [MethodImpl(MethodImplOptions.AggressiveInlining)] private static string ReadNullTermUtf8String(BinaryReader reader, MemoryStream ms) { diff --git a/ValvePak/ValvePak/Package.Save.cs b/ValvePak/ValvePak/Package.Save.cs index 6de3384..444b64f 100644 --- a/ValvePak/ValvePak/Package.Save.cs +++ b/ValvePak/ValvePak/Package.Save.cs @@ -138,11 +138,7 @@ public PackageEntry AddFile(string filePath, byte[] fileData, bool multiChunk = ArchiveIndex = archiveIndex, }; - if (Entries == null) - { - var stringComparer = Comparer == null ? null : StringComparer.FromComparison(Comparer.Comparison); - Entries = new Dictionary>(stringComparer); - } + Entries ??= CreateEntriesDictionary(); if (!Entries.TryGetValue(extension, out var typeEntries)) { @@ -150,7 +146,15 @@ public PackageEntry AddFile(string filePath, byte[] fileData, bool multiChunk = Entries[extension] = typeEntries; } - typeEntries.Add(entry); + if (Comparer == null) + { + typeEntries.Add(entry); + } + else + { + var index = typeEntries.BinarySearch(entry, Comparer); + typeEntries.Insert(index < 0 ? ~index : index, entry); + } return entry; } @@ -325,7 +329,7 @@ private void Write(Stream stream, string? chunkFileBaseName) writer.Write(NullByte); - var fileTreeSize = stream.Position - headerSize; + var fileTreeSize = stream.Position - streamOffset - headerSize; var chunkHashFractions = new List(); // File data, one archive at a time so only a single chunk file is open and written sequentially @@ -354,12 +358,14 @@ private void Write(Stream stream, string? chunkFileBaseName) } var afterFileData = stream.Position; - var fileDataSize = afterFileData - fileTreeSize - headerSize; + var fileDataSize = afterFileData - streamOffset - headerSize - fileTreeSize; // Archive MD5 section, contains hashes of chunk files var archiveMd5SectionBytes = SerializeChunkHashFractions(chunkHashFractions); writer.Write(archiveMd5SectionBytes); + var hashedLength = stream.Position - streamOffset; + // Set tree size // TODO: It is possible to precalculate these sizes to remove seeking stream.Seek(streamOffset + (2 * sizeof(int)), SeekOrigin.Begin); @@ -369,6 +375,7 @@ private void Write(Stream stream, string? chunkFileBaseName) { // Version 1 has no checksums stream.Seek(afterFileData, SeekOrigin.Begin); + TruncateAfterPackage(stream); return; } @@ -386,33 +393,26 @@ private void Write(Stream stream, string? chunkFileBaseName) using var fileTreeMD5 = MD5.Create(); using var fullFileMD5 = MD5.Create(); - stream.ReadExactly(buffer, 0, headerSize); - fullFileMD5.TransformBlock(buffer, 0, headerSize, null, 0); + var treeEnd = headerSize + fileTreeSize; - int bytesRead; - var fileTreeRead = 0; - - // Calculate file tree size hash - while ((bytesRead = stream.Read(buffer, 0, buffer.Length)) > 0) + for (long position = 0; position < hashedLength;) { - fullFileMD5.TransformBlock(buffer, 0, bytesRead, null, 0); + var count = (int)Math.Min(buffer.Length, hashedLength - position); + stream.ReadExactly(buffer, 0, count); + fullFileMD5.TransformBlock(buffer, 0, count, null, 0); - fileTreeRead += bytesRead; + var treeFrom = Math.Max(position, headerSize); + var treeTo = Math.Min(position + count, treeEnd); - if (fileTreeRead >= fileTreeSize) + if (treeFrom < treeTo) { - fileTreeMD5.TransformFinalBlock(buffer, 0, (int)(fileTreeSize - (fileTreeRead - bytesRead))); - break; + fileTreeMD5.TransformBlock(buffer, (int)(treeFrom - position), (int)(treeTo - treeFrom), null, 0); } - fileTreeMD5.TransformBlock(buffer, 0, bytesRead, null, 0); + position += count; } - // Calculate remaining file data hash - while ((bytesRead = stream.Read(buffer, 0, buffer.Length)) > 0) - { - fullFileMD5.TransformBlock(buffer, 0, bytesRead, null, 0); - } + fileTreeMD5.TransformFinalBlock(buffer, 0, 0); // File tree hash var treeHash = fileTreeMD5.Hash; @@ -436,6 +436,27 @@ private void Write(Stream stream, string? chunkFileBaseName) { ArrayPool.Shared.Return(buffer); } + + TruncateAfterPackage(stream); + } + + /// + /// Removes data left in the stream after the written package, if the stream can be resized. + /// + private static void TruncateAfterPackage(Stream stream) + { + if (stream.Length <= stream.Position) + { + return; + } + + try + { + stream.SetLength(stream.Position); + } + catch (NotSupportedException) + { + } } /// diff --git a/ValvePak/ValvePak/Package.Verify.cs b/ValvePak/ValvePak/Package.Verify.cs index 463df40..543719e 100644 --- a/ValvePak/ValvePak/Package.Verify.cs +++ b/ValvePak/ValvePak/Package.Verify.cs @@ -1,5 +1,4 @@ using System; -using System.Diagnostics; using System.Diagnostics.CodeAnalysis; using System.IO; using System.Linq; @@ -51,6 +50,11 @@ public override long Position public override void Write(byte[] buffer, int offset, int count) => throw new NotImplementedException(); } + /// + /// Size in bytes of the buffer that chunk hash fractions are read into, which covers a whole fraction as created by Valve. + /// + private const uint ChunkHashBufferSize = 1024 * 1024; + /// /// Gets the size in bytes of the whole file before . /// @@ -69,36 +73,39 @@ public override long Position /// /// Verify MD5 hashes provided in the VPK. /// + /// The package has not been read. + /// The package has no checksums, or a checksum does not match. public void VerifyHashes() { + ObjectDisposedException.ThrowIf(IsDisposed, this); + + var reader = Reader ?? throw new InvalidOperationException("The package has not been read."); + if (Version != 2) { throw new InvalidDataException("Only version 2 is supported."); } - Debug.Assert(Reader != null); - Debug.Assert(TreeChecksum != null); - Debug.Assert(ArchiveMD5EntriesChecksum != null); - Debug.Assert(WholeFileChecksum != null); + if (TreeChecksum == null || ArchiveMD5EntriesChecksum == null || WholeFileChecksum == null) + { + throw new InvalidDataException("The package has no checksum section."); + } - using var subStream = new SubStream(Reader.BaseStream, HeaderSize, (int)TreeSize); - var hash = MD5.HashData(subStream); + var hash = HashPackageRange(reader, HeaderSize, TreeSize); if (!hash.SequenceEqual(TreeChecksum)) { throw new InvalidDataException($"File tree checksum mismatch ({BitConverter.ToString(hash)} != expected {BitConverter.ToString(TreeChecksum)})"); } - using var subStream2 = new SubStream(Reader.BaseStream, FileSizeBeforeArchiveMD5Entries, (int)ArchiveMD5SectionSize); - hash = MD5.HashData(subStream2); + hash = HashPackageRange(reader, FileSizeBeforeArchiveMD5Entries, ArchiveMD5SectionSize); if (!hash.SequenceEqual(ArchiveMD5EntriesChecksum)) { throw new InvalidDataException($"Archive MD5 entries checksum mismatch ({BitConverter.ToString(hash)} != expected {BitConverter.ToString(ArchiveMD5EntriesChecksum)})"); } - using var subStream3 = new SubStream(Reader.BaseStream, 0, FileSizeBeforeWholeFileHash); - hash = MD5.HashData(subStream3); + hash = HashPackageRange(reader, 0, FileSizeBeforeWholeFileHash); if (!hash.SequenceEqual(WholeFileChecksum)) { @@ -106,93 +113,112 @@ public void VerifyHashes() } } + /// + /// Computes the MD5 hash of a range of the package, at an offset from the start of the package, + /// while holding the reader lock that also takes. + /// + private byte[] HashPackageRange(BinaryReader reader, long offset, long length) + { + lock (reader) + { + using var subStream = new SubStream(reader.BaseStream, StreamOffset + offset, length); + + return MD5.HashData(subStream); + } + } + /// /// Verify hashes of individual chunk files provided in the VPK. /// /// If provided, will report a string with the current verification progress. public void VerifyChunkHashes(IProgress? progressReporter = null) { - Stream? stream = null; - var lastArchiveIndex = ushort.MaxValue; + ObjectDisposedException.ThrowIf(IsDisposed, this); // When created by Valve, entries are sorted, and are 1MB chunks var allEntries = AccessPackFileHashes - .OrderBy(file => file.Offset) - .GroupBy(file => file.ArchiveIndex) - .OrderBy(x => x.Key) - .SelectMany(x => x); + .OrderBy(file => file.ArchiveIndex) + .ThenBy(file => file.Offset) + .ToList(); + var maxLength = allEntries.Count > 0 ? allEntries.Max(file => file.Length) : 0; + var buffer = new byte[Math.Min(maxLength, ChunkHashBufferSize)]; Span hash = stackalloc byte[16]; - try + using var md5 = IncrementalHash.CreateHash(HashAlgorithmName.MD5); + + foreach (var entry in allEntries) { - foreach (var entry in allEntries) + var hashTypeName = entry.HashType switch { - var hashTypeName = entry.HashType switch - { - EHashType.Blake3 => "Blake3", - EHashType.MD5 => "MD5", - _ => string.Empty, - }; - - progressReporter?.Report($"Verifying {hashTypeName} hash at offset {entry.Offset} in archive {entry.ArchiveIndex}."); - - if (stream == null || lastArchiveIndex != entry.ArchiveIndex) - { - if (lastArchiveIndex != 0x7FFF) - { - stream?.Close(); - } + EHashType.Blake3 => "Blake3", + EHashType.MD5 => "MD5", + _ => string.Empty, + }; + + progressReporter?.Report($"Verifying {hashTypeName} hash at offset {entry.Offset} in archive {entry.ArchiveIndex}."); - stream = GetFileStream(entry.ArchiveIndex); - lastArchiveIndex = entry.ArchiveIndex; - } - else - { - var offset = entry.ArchiveIndex == 0x7FFF ? HeaderSize + TreeSize : 0; - stream.Seek(offset, SeekOrigin.Begin); - } - - using var subStream = new SubStream(stream, stream.Position + entry.Offset, entry.Length); - - switch (entry.HashType) - { - case EHashType.Blake3: - using (var hasher = Blake3.Hasher.New()) + GetArchive(entry.ArchiveIndex); + + long position = 0; + ReadOnlySpan data; + + switch (entry.HashType) + { + case EHashType.Blake3: + using (var hasher = Blake3.Hasher.New()) + { + while (ReadChunkHashFraction(entry, buffer, ref position, out data)) { - var buffer = new byte[8192]; // TODO: Fix this alloc - int bytesRead; - while ((bytesRead = subStream.Read(buffer, 0, buffer.Length)) > 0) - { - hasher.UpdateWithJoin(buffer.AsSpan(0, bytesRead)); - } - hasher.Finalize(hash); + hasher.Update(data); } - break; - case EHashType.MD5: - MD5.HashData(subStream, hash); - break; + hasher.Finalize(hash); + } + break; - default: - throw new InvalidDataException($"Unrecognized hash type: {entry.HashType} ({(ushort)entry.HashType})"); - } + case EHashType.MD5: + while (ReadChunkHashFraction(entry, buffer, ref position, out data)) + { + md5.AppendData(data); + } + + md5.GetHashAndReset(hash); + break; - if (!hash.SequenceEqual(entry.Checksum)) - { - throw new InvalidDataException($"Package checksum mismatch ({hashTypeName}) in archive {entry.ArchiveIndex} at {entry.Offset} ({Convert.ToHexString(hash)} != expected {Convert.ToHexString(entry.Checksum)})"); - } + default: + throw new InvalidDataException($"Unrecognized hash type: {entry.HashType} ({(ushort)entry.HashType})"); } - progressReporter?.Report("Successfully verified archive hashes."); - } - finally - { - if (lastArchiveIndex != 0x7FFF) + if (!hash.SequenceEqual(entry.Checksum)) { - stream?.Close(); + throw new InvalidDataException($"Package checksum mismatch ({hashTypeName}) in archive {entry.ArchiveIndex} at {entry.Offset} ({Convert.ToHexString(hash)} != expected {Convert.ToHexString(entry.Checksum)})"); } } + + progressReporter?.Report("Successfully verified archive hashes."); + } + + /// + /// Reads the next part of the data covered by a chunk hash fraction into the buffer. + /// + /// False once all data was read, or the end of the archive was reached. + private bool ReadChunkHashFraction(ChunkHashFraction fraction, byte[] buffer, scoped ref long position, out ReadOnlySpan data) + { + var length = (int)Math.Min(fraction.Length - position, buffer.Length); + + if (length <= 0) + { + data = default; + return false; + } + + var bytesRead = ReadArchiveData(fraction.ArchiveIndex, fraction.Offset + position, buffer.AsSpan(0, length)); + + data = buffer.AsSpan(0, bytesRead); + position = bytesRead < length ? fraction.Length : position + bytesRead; + + return bytesRead > 0; } /// @@ -201,6 +227,8 @@ public void VerifyChunkHashes(IProgress? progressReporter = null) /// If provided, will report a string with the current verification progress. public void VerifyFileChecksums(IProgress? progressReporter = null) { + ObjectDisposedException.ThrowIf(IsDisposed, this); + if (Entries == null) { return; @@ -208,16 +236,21 @@ public void VerifyFileChecksums(IProgress? progressReporter = null) var allEntries = Entries .SelectMany(file => file.Value) - .OrderBy(file => file.Offset) - .GroupBy(file => file.ArchiveIndex) - .OrderBy(x => x.Key) - .SelectMany(x => x); + .OrderBy(file => file.ArchiveIndex) + .ThenBy(file => file.Offset); + + var buffer = Array.Empty(); foreach (var entry in allEntries) { progressReporter?.Report($"Verifying CRC32 checksum for '{entry.GetFullPath()}' in archive {entry.ArchiveIndex}."); - ReadEntry(entry, out var _, validateCrc: true); + if (buffer.Length < entry.TotalLength) + { + buffer = new byte[entry.TotalLength]; + } + + ReadEntry(entry, buffer, validateCrc: true); } progressReporter?.Report("Successfully verified file CRC32 checksums."); @@ -234,7 +267,9 @@ public bool IsSignatureValid() return true; } - if (Reader == null) + var reader = Reader; + + if (reader == null) { return false; } @@ -249,8 +284,7 @@ public bool IsSignatureValid() return false; } - using var subStream3 = new SubStream(Reader.BaseStream, 0, FileSizeBeforeWholeFileHash); - var hash = MD5.HashData(subStream3); + var hash = HashPackageRange(reader, 0, FileSizeBeforeWholeFileHash); if (!hash.SequenceEqual(WholeFileChecksum)) { @@ -266,9 +300,12 @@ public bool IsSignatureValid() return false; } - using var subStream = new SubStream(Reader.BaseStream, 0, FileSizeBeforeSignature); + lock (reader) + { + using var subStream = new SubStream(reader.BaseStream, StreamOffset, FileSizeBeforeSignature); - return rsa.VerifyData(subStream, Signature, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1); + return rsa.VerifyData(subStream, Signature, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1); + } } } } diff --git a/ValvePak/ValvePak/Package.cs b/ValvePak/ValvePak/Package.cs index a0b6d0f..9bf70f8 100644 --- a/ValvePak/ValvePak/Package.cs +++ b/ValvePak/ValvePak/Package.cs @@ -1,7 +1,8 @@ using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.IO; -using System.IO.MemoryMappedFiles; +using System.Threading; namespace ValvePak { @@ -26,7 +27,14 @@ public partial class Package : IDisposable private const char WindowsDirectorySeparator = '\\'; private BinaryReader? Reader; - private readonly Dictionary MemoryMappedPaks = []; + + /// + /// Position in the read stream where the package starts. + /// + private long StreamOffset; + + private readonly ConcurrentDictionary Archives = new(); + private readonly Lock ArchivesLock = new(); /// /// Gets the file name. @@ -128,6 +136,7 @@ public uint Version public List AccessPackFileHashes { get; private set; } = []; private CaseInsensitivePackageEntryComparer? Comparer; + private bool IsDisposed; /// /// Releases binary reader. @@ -145,35 +154,98 @@ protected virtual void Dispose(bool disposing) { if (disposing) { + lock (ArchivesLock) + { + IsDisposed = true; + } + + ReleaseArchives(); + if (Reader != null) { Reader.Dispose(); Reader = null; } + } + } + + /// + /// Closes all archive files. + /// + private void ReleaseArchives() + { + lock (ArchivesLock) + { + foreach (var archive in Archives.Values) + { + archive.Dispose(); + } + + Archives.Clear(); + } + } - foreach (var stream in MemoryMappedPaks.Values) + /// + /// Releases the previous reader and archive files and clears everything that was read, + /// so that reading a package again behaves the same as reading it with a new instance. + /// + /// The stream that is about to be read, which is not disposed even if the previous reader uses it. + private void ResetReadState(Stream input) + { + ReleaseArchives(); + + if (Reader != null) + { + if (Reader.BaseStream != input) { - stream.Dispose(); + Reader.Dispose(); } - MemoryMappedPaks.Clear(); + Reader = null; } + + Version = 2; + StreamOffset = 0; + HeaderSize = 0; + TreeSize = 0; + FileDataSectionSize = 0; + ArchiveMD5SectionSize = 0; + OtherMD5SectionSize = 0; + SignatureSectionSize = 0; + TreeChecksum = null; + ArchiveMD5EntriesChecksum = null; + WholeFileChecksum = null; + PublicKey = null; + Signature = null; + SignatureType = ESignatureType.Unknown; + Entries = null; + AccessPackFileHashes = []; + FileSizeBeforeArchiveMD5Entries = 0; + FileSizeBeforeWholeFileHash = 0; + FileSizeBeforeSignature = 0; + CurrentChunkFileIndex = 0; + CurrentChunkFileSize = 0; + } + + private Dictionary> CreateEntriesDictionary() + { + var stringComparer = Comparer == null ? null : StringComparer.FromComparison(Comparer.Comparison); + + return new Dictionary>(stringComparer); } /// - /// Sets the file name. + /// Sets the file name. Archive files opened for the previous file name are closed. /// /// Filename. public void SetFileName(string fileName) { ArgumentNullException.ThrowIfNull(fileName); - FileName = StripDirVpkSuffixes(fileName, out var isDirVpk); + ReleaseArchives(); - if (isDirVpk) - { - IsDirVPK = true; - } + FileName = StripDirVpkSuffixes(fileName, out var isDirVpk); + IsDirVPK = isDirVpk; } ///