From 824f10357d89fe103bcc92230f5f19b3f062e5a0 Mon Sep 17 00:00:00 2001 From: Gaoyang Date: Tue, 28 Jul 2026 16:38:25 +0800 Subject: [PATCH 1/4] fix: stream disk image save/load to avoid 2GB MemoryStream cap DiskImageSerializer.Save buffered the entire node region into a MemoryStream before compressing/encrypting it, which throws "Stream was too long" once a disk's actual content approaches the ~2GB int32 length limit -- including on the final save during unmount/app-exit/OS-shutdown, silently losing unsaved changes. Unencrypted saves/loads now stream directly to/from the file instead of buffering. Encrypted saves/loads switch to chunked AES-256-GCM (new image format version 4) so no single AES-GCM call or buffer needs to hold the whole node region; version 3's whole-blob encrypted format is still read for backward compatibility. --- .../Persistence/DiskImageSerializer.cs | 525 ++++++++++++++---- .../DiskImageSerializerTests.cs | 144 +++++ 2 files changed, 565 insertions(+), 104 deletions(-) diff --git a/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs b/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs index e01fbbf..04b72ee 100644 --- a/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs +++ b/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs @@ -20,7 +20,7 @@ namespace ManagedDrive.Core.Persistence; /// Image format (little-endian binary): /// /// 4-byte magic "MDRD" -/// Int32 version (currently 3) +/// Int32 version (currently 4) /// Byte holding an value (version 2+ only; absent in version 1, which is always uncompressed) /// Byte IsEncrypted (version 3+ only; absent/false in earlier versions) /// UInt64 capacity in bytes (always plaintext, so callers can preview it without a password) @@ -31,10 +31,23 @@ namespace ManagedDrive.Core.Persistence; /// wraps this randomly generated CEK; actual data is always encrypted with the CEK, so /// changing the password never requires re-encrypting existing data. /// -/// When encrypted: data nonce (12), data tag (16) for the node region below. -/// The node region, gzip-compressed whenever the level is not , then AES-256-GCM encrypted with the CEK when the image is encrypted: -/// Int32 node count -/// For each node: path, metadata, security descriptor bytes, file data bytes +/// +/// Version 3 (legacy, still readable): data nonce (12), data tag (16), then the entire +/// remaining file is one AES-256-GCM ciphertext blob wrapping the whole gzip-compressed +/// node region. Requires materializing the whole node region as a single byte array on +/// both save and load, which is capped by 's single-shot API and by +/// managed-array limits at roughly 2 GB — kept only so old images keep loading. +/// +/// +/// Version 4 (current) when encrypted: a random 12-byte base nonce, then a sequence of +/// chunks, each independently AES-256-GCM encrypted so no single buffer needs to hold the +/// whole node region. Each chunk is [Int32 ciphertext length][16-byte tag][ciphertext +/// bytes], terminated by a zero-length chunk. Per-chunk nonces are derived from the base +/// nonce by XOR-ing its last 4 bytes with the big-endian chunk index, guaranteeing a unique +/// nonce per chunk under the same key/base nonce (see ). +/// +/// When not encrypted (any version): the node region follows directly, gzip-compressed whenever the level is not , streamed straight from/to the file rather than buffered. +/// Node region contents: Int32 node count, then for each node: path, metadata, security descriptor bytes, file data bytes /// /// public static class DiskImageSerializer @@ -44,9 +57,25 @@ public static class DiskImageSerializer private const int Pbkdf2Iterations = 210_000; private const int SaltSize = 16; private const int TagSize = 16; - private const int Version = 3; + private const int Version = 4; private static readonly byte[] Magic = "MDRD"u8.ToArray(); + /// + /// Size of each independently AES-GCM-encrypted chunk when writing a version-4 encrypted + /// node region. Kept well under 2 GB so no single chunk buffer approaches managed-array or + /// single-shot limits. Overridable by tests via + /// to exercise the multi-chunk path without allocating a real 64 MB buffer. + /// + private const int DefaultChunkSize = 64 * 1024 * 1024; + + /// + /// Test-only override for ; means use the + /// production default. Set via InternalsVisibleTo("ManagedDrive.Tests"). + /// + internal static int? TestChunkSizeOverride; + + private static int ChunkSize => TestChunkSizeOverride ?? DefaultChunkSize; + /// /// Generates a fresh random 256-bit content-encryption key for use when encryption is first /// enabled on a disk. @@ -96,7 +125,7 @@ public static FileNodeMap Load( return version <= 2 ? LoadLegacy(stream, reader, version, level, out capacityBytes, out volumeLabel) - : LoadV3(stream, reader, level, isEncrypted, password, out capacityBytes, out volumeLabel, out cek); + : LoadCurrent(stream, reader, version, level, isEncrypted, password, out capacityBytes, out volumeLabel, out cek); } /// @@ -136,7 +165,7 @@ public static void PeekHeader( } else { - // Version 3 layout: capacity/label are always plaintext header fields. + // Version 3+ layout: capacity/label are always plaintext header fields. capacityBytes = reader.ReadUInt64(); volumeLabel = reader.ReadString(); } @@ -157,8 +186,8 @@ public static void PeekHeader( /// /// /// Optional progress reporter, updated with a fraction in [0, 1] as each node is written. - /// The subsequent gzip compression and AES-GCM encryption steps operate on the whole - /// serialized buffer as a single unit and are not individually reported. + /// The subsequent gzip compression and (when encrypting) AES-256-GCM chunk encryption happen + /// as nodes stream through and are not individually reported. /// public static void Save( FileNodeMap nodeMap, @@ -176,41 +205,6 @@ public static void Save( Directory.CreateDirectory(directory); } - // Serialize the node region into memory first: it needs to be a single contiguous byte - // array to gzip-compress and then AES-GCM-encrypt as one unit (the node map is already - // fully materialized in memory anyway, so this adds no meaningful overhead). - byte[] nodeRegionBytes; - using (var nodeRegionStream = new MemoryStream()) - { - using (var payloadWriter = new BinaryWriter( - compress ? new GZipStream(nodeRegionStream, ToCompressionLevel(level), leaveOpen: true) : nodeRegionStream, - System.Text.Encoding.UTF8, - leaveOpen: true)) - { - var nodes = nodeMap.GetAllNodes(); - payloadWriter.Write(nodes.Count); - - if (nodes.Count == 0) - { - progress?.Report(1.0); - } - else - { - var written = 0; - foreach (var kvp in nodes) - { - WriteNode(payloadWriter, kvp.Key, kvp.Value); - written++; - progress?.Report((double)written / nodes.Count); - } - } - - payloadWriter.Flush(); - } - - nodeRegionBytes = nodeRegionStream.ToArray(); - } - // Write to a sibling temp file and flush it to disk before atomically replacing the // real image path, so a process kill mid-write (e.g. during a Windows shutdown) never // leaves the actual image truncated — worst case is a stray .tmp file. @@ -240,24 +234,21 @@ public static void Save( writer.Write(wrapTag); writer.Write(wrappedCek); - var dataNonce = RandomNumberGenerator.GetBytes(NonceSize); - var ciphertext = new byte[nodeRegionBytes.Length]; - var dataTag = new byte[TagSize]; - using (var aesGcm = new AesGcm(enc.Cek, TagSize)) - { - aesGcm.Encrypt(dataNonce, nodeRegionBytes, ciphertext, dataTag); - } - - writer.Write(dataNonce); - writer.Write(dataTag); - writer.Write(ciphertext); + var baseNonce = RandomNumberGenerator.GetBytes(NonceSize); + writer.Write(baseNonce); + writer.Flush(); + + // Node data streams straight into chunked AES-GCM encryption below — + // never buffered whole, so there is no ~2 GB ceiling on disk content. + using var chunkedStream = new ChunkedGcmWriteStream(stream, enc.Cek, baseNonce, ChunkSize); + WriteNodeRegion(chunkedStream, compress, level, nodeMap, progress); + chunkedStream.Complete(); } else { - writer.Write(nodeRegionBytes); + writer.Flush(); + WriteNodeRegion(stream, compress, level, nodeMap, progress); } - - writer.Flush(); } stream.Flush(flushToDisk: true); @@ -278,9 +269,59 @@ public static void Save( throw; } + } + + /// + /// Writes the node-count-prefixed node region for directly into + /// , gzip-compressing on the fly when is + /// set. Never materializes the whole region as a single in-memory buffer, so disk content of + /// any size can be saved regardless of the ~2 GB limit on /managed + /// arrays. The (when used) is explicitly disposed here — rather than + /// relying on 's own disposal with leaveOpen: true, which + /// would skip it — so the deflate stream's final block/trailer is always flushed before + /// is used for anything else. + /// + private static void WriteNodeRegion( + Stream target, + bool compress, + ImageCompressionLevel level, + FileNodeMap nodeMap, + IProgress? progress) + { + var payloadStream = compress + ? new GZipStream(target, ToCompressionLevel(level), leaveOpen: true) + : target; + + try + { + using var payloadWriter = new BinaryWriter(payloadStream, System.Text.Encoding.UTF8, leaveOpen: true); + + var nodes = nodeMap.GetAllNodes(); + payloadWriter.Write(nodes.Count); + + if (nodes.Count == 0) + { + progress?.Report(1.0); + } + else + { + var written = 0; + foreach (var kvp in nodes) + { + WriteNode(payloadWriter, kvp.Key, kvp.Value); + written++; + progress?.Report((double)written / nodes.Count); + } + } + + payloadWriter.Flush(); + } finally { - CryptographicOperations.ZeroMemory(nodeRegionBytes); + if (compress) + { + payloadStream.Dispose(); + } } } @@ -309,13 +350,17 @@ private static FileNodeMap LoadLegacy( } /// - /// Reads a version 3 image: capacity/label are always plaintext header fields; the node - /// region (from node count onward) is compressed and, when encrypted, additionally wrapped - /// in AES-256-GCM using the content-encryption key unwrapped from the password. + /// Reads a version 3 or 4 image: capacity/label are always plaintext header fields; the node + /// region (from node count onward) is compressed and, when encrypted, additionally wrapped in + /// AES-256-GCM using the content-encryption key unwrapped from the password. Version 3 wraps + /// the whole node region as one ciphertext blob (legacy, kept only for backward compatibility); + /// version 4 uses independently encrypted chunks so no single buffer needs to hold the entire + /// node region — see the class remarks and . /// - private static FileNodeMap LoadV3( + private static FileNodeMap LoadCurrent( FileStream stream, BinaryReader reader, + int version, ImageCompressionLevel level, bool isEncrypted, string? password, @@ -326,65 +371,113 @@ private static FileNodeMap LoadV3( capacityBytes = reader.ReadUInt64(); volumeLabel = reader.ReadString(); cek = null; + var compressed = level != ImageCompressionLevel.None; - byte[] nodeRegionBytes; + if (!isEncrypted) + { + // The node region is the last thing in the file for an unencrypted image, so + // decompressing straight off the file stream (rather than buffering it) is safe — + // GZipStream simply reads until end of file. + return ReadNodeRegion(stream, compressed); + } - if (isEncrypted) + if (password is null) { - if (password is null) - { - throw new ImagePasswordRequiredException(); - } + throw new ImagePasswordRequiredException(); + } - var salt = reader.ReadBytes(SaltSize); - var iterations = reader.ReadInt32(); - var wrapNonce = reader.ReadBytes(NonceSize); - var wrapTag = reader.ReadBytes(TagSize); - var wrappedCek = reader.ReadBytes(CekSize); + var salt = reader.ReadBytes(SaltSize); + var iterations = reader.ReadInt32(); + var wrapNonce = reader.ReadBytes(NonceSize); + var wrapTag = reader.ReadBytes(TagSize); + var wrappedCek = reader.ReadBytes(CekSize); - var resolvedCek = UnwrapCek(wrappedCek, password, salt, iterations, wrapNonce, wrapTag); - cek = resolvedCek; + var resolvedCek = UnwrapCek(wrappedCek, password, salt, iterations, wrapNonce, wrapTag); + cek = resolvedCek; - var dataNonce = reader.ReadBytes(NonceSize); - var dataTag = reader.ReadBytes(TagSize); - var ciphertext = reader.ReadBytes((int)(stream.Length - stream.Position)); + return version == 3 + ? LoadLegacyEncryptedBlob(stream, reader, resolvedCek, compressed) + : LoadChunkedEncrypted(stream, reader, resolvedCek, compressed); + } - var plaintext = new byte[ciphertext.Length]; - try - { - using var aesGcm = new AesGcm(resolvedCek, TagSize); - aesGcm.Decrypt(dataNonce, ciphertext, dataTag, plaintext); - } - catch (CryptographicException) - { - throw new ImagePasswordIncorrectException(); - } + /// + /// Version 3's whole-blob encrypted node region: a single AES-256-GCM ciphertext covering the + /// entire (already gzip-compressed) node region. Requires materializing the whole region as + /// one byte array, which is what version 4 exists to avoid — kept only so pre-existing images + /// keep loading. + /// + private static FileNodeMap LoadLegacyEncryptedBlob( + FileStream stream, + BinaryReader reader, + byte[] cek, + bool compressed) + { + var dataNonce = reader.ReadBytes(NonceSize); + var dataTag = reader.ReadBytes(TagSize); + var ciphertext = reader.ReadBytes((int)(stream.Length - stream.Position)); - nodeRegionBytes = plaintext; + var plaintext = new byte[ciphertext.Length]; + try + { + using var aesGcm = new AesGcm(cek, TagSize); + aesGcm.Decrypt(dataNonce, ciphertext, dataTag, plaintext); } - else + catch (CryptographicException) { - nodeRegionBytes = reader.ReadBytes((int)(stream.Length - stream.Position)); + throw new ImagePasswordIncorrectException(); } try { - var compressed = level != ImageCompressionLevel.None; - using var nodeRegionStream = new MemoryStream(nodeRegionBytes); - using var payloadReader = new BinaryReader( - compressed - ? new GZipStream(nodeRegionStream, CompressionMode.Decompress, leaveOpen: true) - : nodeRegionStream, - System.Text.Encoding.UTF8); - - return ReadNodes(payloadReader); + using var nodeRegionStream = new MemoryStream(plaintext, writable: false); + return ReadNodeRegion(nodeRegionStream, compressed); } finally { - CryptographicOperations.ZeroMemory(nodeRegionBytes); + CryptographicOperations.ZeroMemory(plaintext); } } + /// + /// Version 4's chunked encrypted node region: each chunk was independently AES-256-GCM + /// encrypted on save, so decryption streams chunk-by-chunk via + /// rather than requiring the whole region in memory at once. + /// + private static FileNodeMap LoadChunkedEncrypted( + FileStream stream, + BinaryReader reader, + byte[] cek, + bool compressed) + { + var baseNonce = reader.ReadBytes(NonceSize); + + try + { + using var chunkedStream = new ChunkedGcmReadStream(stream, cek, baseNonce); + return ReadNodeRegion(chunkedStream, compressed); + } + catch (CryptographicException) + { + throw new ImagePasswordIncorrectException(); + } + } + + /// + /// Reads the node-count-prefixed node region from , transparently + /// gzip-decompressing when is set. Mirrors . + /// + private static FileNodeMap ReadNodeRegion(Stream source, bool compressed) + { + using var payloadReader = new BinaryReader( + compressed + ? new GZipStream(source, CompressionMode.Decompress, leaveOpen: true) + : source, + System.Text.Encoding.UTF8, + leaveOpen: true); + + return ReadNodes(payloadReader); + } + private static void ReadHeader( BinaryReader reader, out int version, @@ -398,7 +491,7 @@ private static void ReadHeader( } version = reader.ReadInt32(); - if (version is not (1 or 2 or 3)) + if (version is not (1 or 2 or 3 or 4)) { throw new InvalidDataException($"Unsupported image version: {version}."); } @@ -556,4 +649,228 @@ private static void WriteNode(BinaryWriter writer, string path, FileNode node) writer.Write(0L); } } + + /// + /// Derives a unique nonce for chunk from a random per-save + /// by XOR-ing its last 4 bytes with the big-endian chunk index. + /// This is a standard segmented-AEAD nonce derivation: as long as + /// is freshly random per save and chunk indices are never reused within that save (both true + /// here — increments a private counter once per chunk), + /// every chunk gets a distinct nonce under the same key, which is AES-GCM's only requirement. + /// + private static byte[] DeriveChunkNonce(byte[] baseNonce, int chunkIndex) + { + var nonce = (byte[])baseNonce.Clone(); + Span indexBytes = stackalloc byte[4]; + System.Buffers.Binary.BinaryPrimitives.WriteUInt32BigEndian(indexBytes, (uint)chunkIndex); + + for (var i = 0; i < indexBytes.Length; i++) + { + nonce[NonceSize - indexBytes.Length + i] ^= indexBytes[i]; + } + + return nonce; + } + + /// + /// Write-only that buffers up to (or the + /// test override) bytes at a time and, on each full buffer plus once more on , + /// AES-256-GCM-encrypts that chunk with a nonce derived via and + /// writes it to the underlying stream as [Int32 ciphertext length][16-byte tag][ciphertext]. + /// This lets encrypt node regions of any size without ever holding the whole + /// region (or a single AES-GCM call's input) in one buffer. + /// + private sealed class ChunkedGcmWriteStream(Stream output, byte[] key, byte[] baseNonce, int chunkSize) : Stream + { + private readonly byte[] _buffer = new byte[chunkSize]; + private int _bufferLength; + private int _chunkIndex; + private bool _completed; + + public override bool CanRead => false; + public override bool CanSeek => false; + public override bool CanWrite => true; + public override long Length => throw new NotSupportedException(); + + public override long Position + { + get => throw new NotSupportedException(); + set => throw new NotSupportedException(); + } + + public override void Write(byte[] buffer, int offset, int count) + { + while (count > 0) + { + var toCopy = Math.Min(count, _buffer.Length - _bufferLength); + Array.Copy(buffer, offset, _buffer, _bufferLength, toCopy); + _bufferLength += toCopy; + offset += toCopy; + count -= toCopy; + + if (_bufferLength == _buffer.Length) + { + FlushChunk(); + } + } + } + + public override void Flush() + { + } + + /// + /// Flushes any partially filled chunk, then writes a final zero-length chunk as an + /// explicit end-of-stream marker so the reader knows not to expect another chunk header. + /// Must be called exactly once after all plaintext has been written, before disposing. + /// + public void Complete() + { + if (_completed) + { + return; + } + + if (_bufferLength > 0) + { + FlushChunk(); + } + + WriteChunk(ReadOnlySpan.Empty); + _completed = true; + } + + private void FlushChunk() + { + WriteChunk(_buffer.AsSpan(0, _bufferLength)); + CryptographicOperations.ZeroMemory(_buffer.AsSpan(0, _bufferLength)); + _bufferLength = 0; + } + + private void WriteChunk(ReadOnlySpan plaintext) + { + var nonce = DeriveChunkNonce(baseNonce, _chunkIndex); + var ciphertext = plaintext.Length == 0 ? [] : new byte[plaintext.Length]; + var tag = new byte[TagSize]; + + using (var aesGcm = new AesGcm(key, TagSize)) + { + aesGcm.Encrypt(nonce, plaintext, ciphertext, tag); + } + + Span lengthBytes = stackalloc byte[4]; + System.Buffers.Binary.BinaryPrimitives.WriteInt32LittleEndian(lengthBytes, ciphertext.Length); + output.Write(lengthBytes); + output.Write(tag); + if (ciphertext.Length > 0) + { + output.Write(ciphertext); + } + + _chunkIndex++; + } + + public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + + public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); + + public override void SetLength(long value) => throw new NotSupportedException(); + } + + /// + /// Read-only counterpart to : reads + /// the [length][tag][ciphertext] chunk sequence from the underlying stream, decrypting + /// each chunk with the matching derived nonce and exposing the concatenated plaintext as a + /// normal readable stream (typically wrapped by a decompressing ). + /// Throws (translated by the caller into + /// ) if any chunk's tag fails to authenticate. + /// + private sealed class ChunkedGcmReadStream(Stream source, byte[] key, byte[] baseNonce) : Stream + { + private byte[] _currentChunk = []; + private int _currentChunkLength; + private int _positionInChunk; + private int _chunkIndex; + private bool _endOfStream; + + public override bool CanRead => true; + public override bool CanSeek => false; + public override bool CanWrite => false; + public override long Length => throw new NotSupportedException(); + + public override long Position + { + get => throw new NotSupportedException(); + set => throw new NotSupportedException(); + } + + public override int Read(byte[] buffer, int offset, int count) + { + var totalRead = 0; + + while (count > 0) + { + if (_positionInChunk == _currentChunkLength) + { + if (_endOfStream || !TryReadNextChunk()) + { + break; + } + } + + var toCopy = Math.Min(count, _currentChunkLength - _positionInChunk); + Array.Copy(_currentChunk, _positionInChunk, buffer, offset, toCopy); + _positionInChunk += toCopy; + offset += toCopy; + count -= toCopy; + totalRead += toCopy; + } + + return totalRead; + } + + private bool TryReadNextChunk() + { + Span lengthBytes = stackalloc byte[4]; + source.ReadExactly(lengthBytes); + var ciphertextLength = System.Buffers.Binary.BinaryPrimitives.ReadInt32LittleEndian(lengthBytes); + + var tag = new byte[TagSize]; + source.ReadExactly(tag); + + var ciphertext = ciphertextLength == 0 ? [] : new byte[ciphertextLength]; + if (ciphertextLength > 0) + { + source.ReadExactly(ciphertext); + } + + var nonce = DeriveChunkNonce(baseNonce, _chunkIndex); + var plaintext = ciphertextLength == 0 ? [] : new byte[ciphertextLength]; + using (var aesGcm = new AesGcm(key, TagSize)) + { + aesGcm.Decrypt(nonce, ciphertext, tag, plaintext); + } + + _chunkIndex++; + + if (ciphertextLength == 0) + { + _endOfStream = true; + return false; + } + + _currentChunk = plaintext; + _currentChunkLength = plaintext.Length; + _positionInChunk = 0; + return true; + } + + public override void Flush() => throw new NotSupportedException(); + + public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + + public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); + + public override void SetLength(long value) => throw new NotSupportedException(); + } } \ No newline at end of file diff --git a/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs b/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs index 76e9dff..c46f052 100644 --- a/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs +++ b/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs @@ -1,3 +1,6 @@ +using System.IO.Compression; +using System.Security.Cryptography; + namespace ManagedDrive.Tests; public sealed class DiskImageSerializerTests @@ -251,6 +254,147 @@ public void SaveThenLoad_WithPassword_RoundTrips(ImageCompressionLevel level) } } + [Theory] + [InlineData(ImageCompressionLevel.None)] + [InlineData(ImageCompressionLevel.Fastest)] + public void SaveThenLoad_WithPasswordAcrossMultipleChunks_RoundTrips(ImageCompressionLevel level) + { + // Force a tiny chunk size so a handful of KB of node data spans several chunks, + // exercising the version-4 chunked AES-GCM path without allocating a real 64 MB buffer. + DiskImageSerializer.TestChunkSizeOverride = 64; + var path = Path.Combine(Path.GetTempPath(), $"{Guid.NewGuid()}.mdr"); + try + { + var map = new FileNodeMap(); + map.Add("\\", MakeDir()); + for (var i = 0; i < 20; i++) + { + map.Add($"\\File{i}.txt", MakeFile(System.Text.Encoding.UTF8.GetBytes($"content for file number {i}"))); + } + + var cek = DiskImageSerializer.GenerateCek(); + DiskImageSerializer.Save(map, capacityBytes: 1024 * 1024, "ChunkedLabel", path, level, new ImageEncryptionInfo("s3cret", cek)); + + var loaded = DiskImageSerializer.Load(path, out var capacityBytes, out var volumeLabel, "s3cret", out var loadedCek); + + Assert.Equal(1024UL * 1024, capacityBytes); + Assert.Equal("ChunkedLabel", volumeLabel); + Assert.Equal(21, loaded.Count); + for (var i = 0; i < 20; i++) + { + Assert.True(loaded.TryGet($"\\File{i}.txt", out var node)); + var expected = System.Text.Encoding.UTF8.GetBytes($"content for file number {i}"); + Assert.Equal(expected, node!.FileData!.ToArray(expected.Length)); + } + + Assert.Equal(cek, loadedCek); + } + finally + { + DiskImageSerializer.TestChunkSizeOverride = null; + File.Delete(path); + } + } + + [Fact] + public void SaveThenLoad_WithWrongPasswordAcrossMultipleChunks_ThrowsPasswordIncorrect() + { + DiskImageSerializer.TestChunkSizeOverride = 64; + var path = Path.Combine(Path.GetTempPath(), $"{Guid.NewGuid()}.mdr"); + try + { + var map = new FileNodeMap(); + map.Add("\\", MakeDir()); + for (var i = 0; i < 20; i++) + { + map.Add($"\\File{i}.txt", MakeFile("hello world"u8.ToArray())); + } + + DiskImageSerializer.Save(map, capacityBytes: 1024 * 1024, "MyLabel", path, ImageCompressionLevel.Fastest, + new ImageEncryptionInfo("s3cret", DiskImageSerializer.GenerateCek())); + + Assert.Throws(() => + DiskImageSerializer.Load(path, out _, out _, "wrong-password", out _)); + } + finally + { + DiskImageSerializer.TestChunkSizeOverride = null; + File.Delete(path); + } + } + + [Fact] + public void Load_LegacyVersion3EncryptedWholeBlobImage_StillLoads() + { + var path = Path.Combine(Path.GetTempPath(), $"{Guid.NewGuid()}.mdr"); + try + { + const string password = "s3cret"; + const int iterations = 210_000; + var cek = DiskImageSerializer.GenerateCek(); + + byte[] nodeRegion; + using (var nodeRegionStream = new MemoryStream()) + { + using (var gzip = new GZipStream(nodeRegionStream, CompressionLevel.Fastest, leaveOpen: true)) + using (var payloadWriter = new BinaryWriter(gzip, System.Text.Encoding.UTF8, leaveOpen: true)) + { + payloadWriter.Write(0); // node count + } + + nodeRegion = nodeRegionStream.ToArray(); + } + + var salt = RandomNumberGenerator.GetBytes(16); + var kek = Rfc2898DeriveBytes.Pbkdf2(password, salt, iterations, HashAlgorithmName.SHA256, 32); + var wrapNonce = RandomNumberGenerator.GetBytes(12); + var wrappedCek = new byte[cek.Length]; + var wrapTag = new byte[16]; + using (var aesGcm = new AesGcm(kek, 16)) + { + aesGcm.Encrypt(wrapNonce, cek, wrappedCek, wrapTag); + } + + var dataNonce = RandomNumberGenerator.GetBytes(12); + var ciphertext = new byte[nodeRegion.Length]; + var dataTag = new byte[16]; + using (var aesGcm = new AesGcm(cek, 16)) + { + aesGcm.Encrypt(dataNonce, nodeRegion, ciphertext, dataTag); + } + + using (var stream = new FileStream(path, FileMode.Create, FileAccess.Write)) + using (var writer = new BinaryWriter(stream, System.Text.Encoding.UTF8)) + { + writer.Write("MDRD"u8.ToArray()); + writer.Write(3); // legacy whole-blob encrypted version + writer.Write((byte)ImageCompressionLevel.Fastest); + writer.Write((byte)1); // isEncrypted + writer.Write(2048UL); + writer.Write("LegacyEncryptedLabel"); + writer.Write(salt); + writer.Write(iterations); + writer.Write(wrapNonce); + writer.Write(wrapTag); + writer.Write(wrappedCek); + writer.Write(dataNonce); + writer.Write(dataTag); + writer.Write(ciphertext); + } + + var loaded = DiskImageSerializer.Load(path, out var capacityBytes, out var volumeLabel, password, out var loadedCek); + + Assert.Equal(2048UL, capacityBytes); + Assert.Equal("LegacyEncryptedLabel", volumeLabel); + Assert.Equal(0, loaded.Count); + Assert.Equal(cek, loadedCek); + } + finally + { + File.Delete(path); + } + } + private static FileNode MakeDir() => new() { FileInfo = { FileAttributes = (uint)FileAttributes.Directory }, From e69c5fde6833c331cbd806f5de00dbebf6e72fdf Mon Sep 17 00:00:00 2001 From: Gaoyang Date: Tue, 28 Jul 2026 16:48:42 +0800 Subject: [PATCH 2/4] fix: stream SnapshotStore blob read/write to avoid whole-buffer caps Individual snapshot file blobs had the same class of bug fixed for DiskImageSerializer: EnsureBlobWritten/ReadBlob materialized the full compressed/encrypted payload as a single buffer, and WriteNode truncated file size to int before that -- both breaking for single files approaching/exceeding ~2GB. Extracted the chunked AES-256-GCM stream pair used by DiskImageSerializer into a shared ChunkedGcm helper and reused it here so encrypted blobs are written/read chunk-by-chunk. Unencrypted blobs now stream straight through FileContent/GZipStream instead of via MemoryStream. Old whole-blob encrypted blobs are still read via a legacy branch for backward compatibility. --- .../Persistence/ChunkedGcm.cs | 257 ++++++++++++++++++ .../Persistence/DiskImageSerializer.cs | 250 +---------------- .../Snapshots/SnapshotStore.cs | 198 ++++++++------ .../DiskImageSerializerTests.cs | 8 +- .../SnapshotManagerTests.cs | 99 +++++++ 5 files changed, 482 insertions(+), 330 deletions(-) create mode 100644 src/ManagedDrive.Core/Persistence/ChunkedGcm.cs diff --git a/src/ManagedDrive.Core/Persistence/ChunkedGcm.cs b/src/ManagedDrive.Core/Persistence/ChunkedGcm.cs new file mode 100644 index 0000000..6903549 --- /dev/null +++ b/src/ManagedDrive.Core/Persistence/ChunkedGcm.cs @@ -0,0 +1,257 @@ +using System.Buffers.Binary; +using System.Security.Cryptography; + +namespace ManagedDrive.Core.Persistence; + +/// +/// Shared chunked AES-256-GCM read/write stream pair used to encrypt/decrypt data of any size +/// without ever holding the whole plaintext or ciphertext in a single buffer — +/// only exposes single-shot Encrypt/Decrypt over one buffer, and a managed array is +/// itself practically capped near 2 GB, so this splits the data into independently +/// encrypted/authenticated chunks instead. Used by both (the +/// node region of a .mdr image) and (individual +/// content-addressed file blobs). +/// +internal static class ChunkedGcm +{ + internal const int NonceSize = 12; + internal const int TagSize = 16; + + /// + /// Size of each independently AES-GCM-encrypted chunk. Kept well under 2 GB so no single + /// chunk buffer approaches managed-array or single-shot limits. + /// Overridable by tests via to exercise the multi-chunk + /// path without allocating a real 64 MB buffer. + /// + private const int DefaultChunkSize = 64 * 1024 * 1024; + + /// + /// Test-only override for ; means use the + /// production default. Set via InternalsVisibleTo("ManagedDrive.Tests"). + /// + internal static int? TestChunkSizeOverride; + + internal static int ChunkSize => TestChunkSizeOverride ?? DefaultChunkSize; + + /// + /// Derives a unique nonce for chunk from a random per-save + /// by XOR-ing its last 4 bytes with the big-endian chunk index. + /// This is a standard segmented-AEAD nonce derivation: as long as + /// is freshly random per save and chunk indices are never reused within that save (both true + /// here — increments a private counter once per chunk), every chunk + /// gets a distinct nonce under the same key, which is AES-GCM's only requirement. + /// + internal static byte[] DeriveChunkNonce(byte[] baseNonce, int chunkIndex) + { + var nonce = (byte[])baseNonce.Clone(); + Span indexBytes = stackalloc byte[4]; + BinaryPrimitives.WriteUInt32BigEndian(indexBytes, (uint)chunkIndex); + + for (var i = 0; i < indexBytes.Length; i++) + { + nonce[NonceSize - indexBytes.Length + i] ^= indexBytes[i]; + } + + return nonce; + } + + /// + /// Write-only that buffers up to bytes at a + /// time and, on each full buffer plus once more on , AES-256-GCM-encrypts + /// that chunk with a nonce derived via and writes it to the + /// underlying stream as [Int32 ciphertext length][16-byte tag][ciphertext]. + /// + internal sealed class WriteStream(Stream output, byte[] key, byte[] baseNonce, int chunkSize) : Stream + { + private readonly byte[] _buffer = new byte[chunkSize]; + private int _bufferLength; + private int _chunkIndex; + private bool _completed; + + public override bool CanRead => false; + public override bool CanSeek => false; + public override bool CanWrite => true; + public override long Length => throw new NotSupportedException(); + + public override long Position + { + get => throw new NotSupportedException(); + set => throw new NotSupportedException(); + } + + public override void Write(byte[] buffer, int offset, int count) + { + while (count > 0) + { + var toCopy = Math.Min(count, _buffer.Length - _bufferLength); + Array.Copy(buffer, offset, _buffer, _bufferLength, toCopy); + _bufferLength += toCopy; + offset += toCopy; + count -= toCopy; + + if (_bufferLength == _buffer.Length) + { + FlushChunk(); + } + } + } + + public override void Flush() + { + } + + /// + /// Flushes any partially filled chunk, then writes a final zero-length chunk as an + /// explicit end-of-stream marker so the reader knows not to expect another chunk header. + /// Must be called exactly once after all plaintext has been written, before disposing. + /// + public void Complete() + { + if (_completed) + { + return; + } + + if (_bufferLength > 0) + { + FlushChunk(); + } + + WriteChunk(ReadOnlySpan.Empty); + _completed = true; + } + + private void FlushChunk() + { + WriteChunk(_buffer.AsSpan(0, _bufferLength)); + CryptographicOperations.ZeroMemory(_buffer.AsSpan(0, _bufferLength)); + _bufferLength = 0; + } + + private void WriteChunk(ReadOnlySpan plaintext) + { + var nonce = DeriveChunkNonce(baseNonce, _chunkIndex); + var ciphertext = plaintext.Length == 0 ? [] : new byte[plaintext.Length]; + var tag = new byte[TagSize]; + + using (var aesGcm = new AesGcm(key, TagSize)) + { + aesGcm.Encrypt(nonce, plaintext, ciphertext, tag); + } + + Span lengthBytes = stackalloc byte[4]; + BinaryPrimitives.WriteInt32LittleEndian(lengthBytes, ciphertext.Length); + output.Write(lengthBytes); + output.Write(tag); + if (ciphertext.Length > 0) + { + output.Write(ciphertext); + } + + _chunkIndex++; + } + + public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + + public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); + + public override void SetLength(long value) => throw new NotSupportedException(); + } + + /// + /// Read-only counterpart to : reads the + /// [length][tag][ciphertext] chunk sequence from the underlying stream, decrypting + /// each chunk with the matching derived nonce and exposing the concatenated plaintext as a + /// normal readable stream (typically wrapped by a decompressing ). + /// Throws if any chunk's tag fails to authenticate — + /// callers should translate that into their own password-incorrect exception type. + /// + internal sealed class ReadStream(Stream source, byte[] key, byte[] baseNonce) : Stream + { + private byte[] _currentChunk = []; + private int _currentChunkLength; + private int _positionInChunk; + private int _chunkIndex; + private bool _endOfStream; + + public override bool CanRead => true; + public override bool CanSeek => false; + public override bool CanWrite => false; + public override long Length => throw new NotSupportedException(); + + public override long Position + { + get => throw new NotSupportedException(); + set => throw new NotSupportedException(); + } + + public override int Read(byte[] buffer, int offset, int count) + { + var totalRead = 0; + + while (count > 0) + { + if (_positionInChunk == _currentChunkLength) + { + if (_endOfStream || !TryReadNextChunk()) + { + break; + } + } + + var toCopy = Math.Min(count, _currentChunkLength - _positionInChunk); + Array.Copy(_currentChunk, _positionInChunk, buffer, offset, toCopy); + _positionInChunk += toCopy; + offset += toCopy; + count -= toCopy; + totalRead += toCopy; + } + + return totalRead; + } + + private bool TryReadNextChunk() + { + Span lengthBytes = stackalloc byte[4]; + source.ReadExactly(lengthBytes); + var ciphertextLength = BinaryPrimitives.ReadInt32LittleEndian(lengthBytes); + + var tag = new byte[TagSize]; + source.ReadExactly(tag); + + var ciphertext = ciphertextLength == 0 ? [] : new byte[ciphertextLength]; + if (ciphertextLength > 0) + { + source.ReadExactly(ciphertext); + } + + var nonce = DeriveChunkNonce(baseNonce, _chunkIndex); + var plaintext = ciphertextLength == 0 ? [] : new byte[ciphertextLength]; + using (var aesGcm = new AesGcm(key, TagSize)) + { + aesGcm.Decrypt(nonce, ciphertext, tag, plaintext); + } + + _chunkIndex++; + + if (ciphertextLength == 0) + { + _endOfStream = true; + return false; + } + + _currentChunk = plaintext; + _currentChunkLength = plaintext.Length; + _positionInChunk = 0; + return true; + } + + public override void Flush() => throw new NotSupportedException(); + + public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + + public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); + + public override void SetLength(long value) => throw new NotSupportedException(); + } +} diff --git a/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs b/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs index 04b72ee..b2b9714 100644 --- a/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs +++ b/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs @@ -44,7 +44,7 @@ namespace ManagedDrive.Core.Persistence; /// whole node region. Each chunk is [Int32 ciphertext length][16-byte tag][ciphertext /// bytes], terminated by a zero-length chunk. Per-chunk nonces are derived from the base /// nonce by XOR-ing its last 4 bytes with the big-endian chunk index, guaranteeing a unique -/// nonce per chunk under the same key/base nonce (see ). +/// nonce per chunk under the same key/base nonce (see ). /// /// When not encrypted (any version): the node region follows directly, gzip-compressed whenever the level is not , streamed straight from/to the file rather than buffered. /// Node region contents: Int32 node count, then for each node: path, metadata, security descriptor bytes, file data bytes @@ -60,22 +60,6 @@ public static class DiskImageSerializer private const int Version = 4; private static readonly byte[] Magic = "MDRD"u8.ToArray(); - /// - /// Size of each independently AES-GCM-encrypted chunk when writing a version-4 encrypted - /// node region. Kept well under 2 GB so no single chunk buffer approaches managed-array or - /// single-shot limits. Overridable by tests via - /// to exercise the multi-chunk path without allocating a real 64 MB buffer. - /// - private const int DefaultChunkSize = 64 * 1024 * 1024; - - /// - /// Test-only override for ; means use the - /// production default. Set via InternalsVisibleTo("ManagedDrive.Tests"). - /// - internal static int? TestChunkSizeOverride; - - private static int ChunkSize => TestChunkSizeOverride ?? DefaultChunkSize; - /// /// Generates a fresh random 256-bit content-encryption key for use when encryption is first /// enabled on a disk. @@ -240,7 +224,7 @@ public static void Save( // Node data streams straight into chunked AES-GCM encryption below — // never buffered whole, so there is no ~2 GB ceiling on disk content. - using var chunkedStream = new ChunkedGcmWriteStream(stream, enc.Cek, baseNonce, ChunkSize); + using var chunkedStream = new ChunkedGcm.WriteStream(stream, enc.Cek, baseNonce, ChunkedGcm.ChunkSize); WriteNodeRegion(chunkedStream, compress, level, nodeMap, progress); chunkedStream.Complete(); } @@ -355,7 +339,7 @@ private static FileNodeMap LoadLegacy( /// AES-256-GCM using the content-encryption key unwrapped from the password. Version 3 wraps /// the whole node region as one ciphertext blob (legacy, kept only for backward compatibility); /// version 4 uses independently encrypted chunks so no single buffer needs to hold the entire - /// node region — see the class remarks and . + /// node region — see the class remarks and . /// private static FileNodeMap LoadCurrent( FileStream stream, @@ -440,7 +424,7 @@ private static FileNodeMap LoadLegacyEncryptedBlob( /// /// Version 4's chunked encrypted node region: each chunk was independently AES-256-GCM - /// encrypted on save, so decryption streams chunk-by-chunk via + /// encrypted on save, so decryption streams chunk-by-chunk via /// rather than requiring the whole region in memory at once. /// private static FileNodeMap LoadChunkedEncrypted( @@ -453,7 +437,7 @@ private static FileNodeMap LoadChunkedEncrypted( try { - using var chunkedStream = new ChunkedGcmReadStream(stream, cek, baseNonce); + using var chunkedStream = new ChunkedGcm.ReadStream(stream, cek, baseNonce); return ReadNodeRegion(chunkedStream, compressed); } catch (CryptographicException) @@ -649,228 +633,4 @@ private static void WriteNode(BinaryWriter writer, string path, FileNode node) writer.Write(0L); } } - - /// - /// Derives a unique nonce for chunk from a random per-save - /// by XOR-ing its last 4 bytes with the big-endian chunk index. - /// This is a standard segmented-AEAD nonce derivation: as long as - /// is freshly random per save and chunk indices are never reused within that save (both true - /// here — increments a private counter once per chunk), - /// every chunk gets a distinct nonce under the same key, which is AES-GCM's only requirement. - /// - private static byte[] DeriveChunkNonce(byte[] baseNonce, int chunkIndex) - { - var nonce = (byte[])baseNonce.Clone(); - Span indexBytes = stackalloc byte[4]; - System.Buffers.Binary.BinaryPrimitives.WriteUInt32BigEndian(indexBytes, (uint)chunkIndex); - - for (var i = 0; i < indexBytes.Length; i++) - { - nonce[NonceSize - indexBytes.Length + i] ^= indexBytes[i]; - } - - return nonce; - } - - /// - /// Write-only that buffers up to (or the - /// test override) bytes at a time and, on each full buffer plus once more on , - /// AES-256-GCM-encrypts that chunk with a nonce derived via and - /// writes it to the underlying stream as [Int32 ciphertext length][16-byte tag][ciphertext]. - /// This lets encrypt node regions of any size without ever holding the whole - /// region (or a single AES-GCM call's input) in one buffer. - /// - private sealed class ChunkedGcmWriteStream(Stream output, byte[] key, byte[] baseNonce, int chunkSize) : Stream - { - private readonly byte[] _buffer = new byte[chunkSize]; - private int _bufferLength; - private int _chunkIndex; - private bool _completed; - - public override bool CanRead => false; - public override bool CanSeek => false; - public override bool CanWrite => true; - public override long Length => throw new NotSupportedException(); - - public override long Position - { - get => throw new NotSupportedException(); - set => throw new NotSupportedException(); - } - - public override void Write(byte[] buffer, int offset, int count) - { - while (count > 0) - { - var toCopy = Math.Min(count, _buffer.Length - _bufferLength); - Array.Copy(buffer, offset, _buffer, _bufferLength, toCopy); - _bufferLength += toCopy; - offset += toCopy; - count -= toCopy; - - if (_bufferLength == _buffer.Length) - { - FlushChunk(); - } - } - } - - public override void Flush() - { - } - - /// - /// Flushes any partially filled chunk, then writes a final zero-length chunk as an - /// explicit end-of-stream marker so the reader knows not to expect another chunk header. - /// Must be called exactly once after all plaintext has been written, before disposing. - /// - public void Complete() - { - if (_completed) - { - return; - } - - if (_bufferLength > 0) - { - FlushChunk(); - } - - WriteChunk(ReadOnlySpan.Empty); - _completed = true; - } - - private void FlushChunk() - { - WriteChunk(_buffer.AsSpan(0, _bufferLength)); - CryptographicOperations.ZeroMemory(_buffer.AsSpan(0, _bufferLength)); - _bufferLength = 0; - } - - private void WriteChunk(ReadOnlySpan plaintext) - { - var nonce = DeriveChunkNonce(baseNonce, _chunkIndex); - var ciphertext = plaintext.Length == 0 ? [] : new byte[plaintext.Length]; - var tag = new byte[TagSize]; - - using (var aesGcm = new AesGcm(key, TagSize)) - { - aesGcm.Encrypt(nonce, plaintext, ciphertext, tag); - } - - Span lengthBytes = stackalloc byte[4]; - System.Buffers.Binary.BinaryPrimitives.WriteInt32LittleEndian(lengthBytes, ciphertext.Length); - output.Write(lengthBytes); - output.Write(tag); - if (ciphertext.Length > 0) - { - output.Write(ciphertext); - } - - _chunkIndex++; - } - - public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException(); - - public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); - - public override void SetLength(long value) => throw new NotSupportedException(); - } - - /// - /// Read-only counterpart to : reads - /// the [length][tag][ciphertext] chunk sequence from the underlying stream, decrypting - /// each chunk with the matching derived nonce and exposing the concatenated plaintext as a - /// normal readable stream (typically wrapped by a decompressing ). - /// Throws (translated by the caller into - /// ) if any chunk's tag fails to authenticate. - /// - private sealed class ChunkedGcmReadStream(Stream source, byte[] key, byte[] baseNonce) : Stream - { - private byte[] _currentChunk = []; - private int _currentChunkLength; - private int _positionInChunk; - private int _chunkIndex; - private bool _endOfStream; - - public override bool CanRead => true; - public override bool CanSeek => false; - public override bool CanWrite => false; - public override long Length => throw new NotSupportedException(); - - public override long Position - { - get => throw new NotSupportedException(); - set => throw new NotSupportedException(); - } - - public override int Read(byte[] buffer, int offset, int count) - { - var totalRead = 0; - - while (count > 0) - { - if (_positionInChunk == _currentChunkLength) - { - if (_endOfStream || !TryReadNextChunk()) - { - break; - } - } - - var toCopy = Math.Min(count, _currentChunkLength - _positionInChunk); - Array.Copy(_currentChunk, _positionInChunk, buffer, offset, toCopy); - _positionInChunk += toCopy; - offset += toCopy; - count -= toCopy; - totalRead += toCopy; - } - - return totalRead; - } - - private bool TryReadNextChunk() - { - Span lengthBytes = stackalloc byte[4]; - source.ReadExactly(lengthBytes); - var ciphertextLength = System.Buffers.Binary.BinaryPrimitives.ReadInt32LittleEndian(lengthBytes); - - var tag = new byte[TagSize]; - source.ReadExactly(tag); - - var ciphertext = ciphertextLength == 0 ? [] : new byte[ciphertextLength]; - if (ciphertextLength > 0) - { - source.ReadExactly(ciphertext); - } - - var nonce = DeriveChunkNonce(baseNonce, _chunkIndex); - var plaintext = ciphertextLength == 0 ? [] : new byte[ciphertextLength]; - using (var aesGcm = new AesGcm(key, TagSize)) - { - aesGcm.Decrypt(nonce, ciphertext, tag, plaintext); - } - - _chunkIndex++; - - if (ciphertextLength == 0) - { - _endOfStream = true; - return false; - } - - _currentChunk = plaintext; - _currentChunkLength = plaintext.Length; - _positionInChunk = 0; - return true; - } - - public override void Flush() => throw new NotSupportedException(); - - public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException(); - - public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); - - public override void SetLength(long value) => throw new NotSupportedException(); - } } \ No newline at end of file diff --git a/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs b/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs index 525b51b..d7fa833 100644 --- a/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs +++ b/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs @@ -13,8 +13,17 @@ namespace ManagedDrive.Core.Snapshots; /// internal static class SnapshotStore { - private const int BlobFlagCompressed = 0b01; - private const int BlobFlagEncrypted = 0b10; + private const int BlobFlagCompressed = 0b001; + private const int BlobFlagEncrypted = 0b010; + + /// + /// Marks a blob's encrypted payload as chunked AES-256-GCM (see ) + /// rather than the legacy whole-blob single-shot layout. Only meaningful when + /// is also set. Never set for blobs written before this + /// flag existed, so old blobs keep loading via the legacy branch in . + /// + private const int BlobFlagChunked = 0b100; + private const int BlobNonceSize = 12; private const int BlobTagSize = 16; private const int Version = 1; @@ -241,7 +250,16 @@ internal static void Write( } } - private static void EnsureBlobWritten(string blobDirectory, byte[] hash, FileContent data, int length, ImageCompressionLevel level, byte[]? cek) + /// + /// Writes the blob for if it doesn't already exist. Streams + /// straight from through gzip compression + /// and (when is set) chunked AES-256-GCM encryption directly into the + /// destination file — no whole-file buffer is ever materialized, so a single blob's size is + /// not limited by 's ~2 GB cap. New encrypted blobs always use the + /// chunked layout (, flagged via ); see + /// for the legacy whole-blob layout this format replaces. + /// + private static void EnsureBlobWritten(string blobDirectory, byte[] hash, FileContent data, long length, ImageCompressionLevel level, byte[]? cek) { var blobPath = HashToBlobPath(blobDirectory, hash); if (File.Exists(blobPath)) @@ -252,38 +270,7 @@ private static void EnsureBlobWritten(string blobDirectory, byte[] hash, FileCon Directory.CreateDirectory(Path.GetDirectoryName(blobPath)!); var compress = level != ImageCompressionLevel.None; - - byte[] payload; - if (compress) - { - using var ms = new MemoryStream(); - using (var gzip = new GZipStream(ms, ToCompressionLevel(level), leaveOpen: true)) - { - data.CopyTo(gzip, length); - } - - payload = ms.ToArray(); - } - else - { - payload = data.ToArray(length); - } - - var flag = compress ? BlobFlagCompressed : 0; - byte[]? nonce = null; - byte[]? tag = null; - - if (cek is not null) - { - flag |= BlobFlagEncrypted; - nonce = RandomNumberGenerator.GetBytes(BlobNonceSize); - var ciphertext = new byte[payload.Length]; - var localTag = new byte[BlobTagSize]; - using var aesGcm = new AesGcm(cek, BlobTagSize); - aesGcm.Encrypt(nonce, payload, ciphertext, localTag); - tag = localTag; - payload = ciphertext; - } + var flag = (compress ? BlobFlagCompressed : 0) | (cek is not null ? BlobFlagEncrypted | BlobFlagChunked : 0); var tempPath = blobPath + "." + Guid.NewGuid().ToString("N") + ".tmp"; @@ -293,13 +280,37 @@ private static void EnsureBlobWritten(string blobDirectory, byte[] hash, FileCon { stream.WriteByte((byte)flag); - if (nonce is not null && tag is not null) + Stream target = stream; + ChunkedGcm.WriteStream? chunkedStream = null; + if (cek is not null) { - stream.Write(nonce); - stream.Write(tag); + var baseNonce = RandomNumberGenerator.GetBytes(BlobNonceSize); + stream.Write(baseNonce); + chunkedStream = new ChunkedGcm.WriteStream(stream, cek, baseNonce, ChunkedGcm.ChunkSize); + target = chunkedStream; } - stream.Write(payload); + if (compress) + { + var gzip = new GZipStream(target, ToCompressionLevel(level), leaveOpen: true); + try + { + data.CopyTo(gzip, length); + } + finally + { + // Explicitly disposed (rather than relying on leaveOpen semantics further + // up the chain) so the deflate stream's final block is flushed before the + // chunked encryption below is completed. + gzip.Dispose(); + } + } + else + { + data.CopyTo(target, length); + } + + chunkedStream?.Complete(); stream.Flush(flushToDisk: true); } @@ -320,6 +331,12 @@ private static void EnsureBlobWritten(string blobDirectory, byte[] hash, FileCon } } + /// + /// Reads the blob for straight into a via + /// , decrypting (chunked or legacy whole-blob, see + /// ) and decompressing on the fly rather than materializing the + /// ciphertext, plaintext, and decompressed bytes as three separate whole-file buffers. + /// private static FileContent ReadBlob(string blobDirectory, byte[] hash, string nodePath, ulong fileSize, ulong allocationSize, byte[]? cek) { var blobPath = HashToBlobPath(blobDirectory, hash); @@ -333,67 +350,86 @@ private static FileContent ReadBlob(string blobDirectory, byte[] hash, string no var flag = stream.ReadByte(); var compressed = (flag & BlobFlagCompressed) != 0; var encrypted = (flag & BlobFlagEncrypted) != 0; + var chunked = (flag & BlobFlagChunked) != 0; + + Stream plaintextStream; + byte[]? legacyPlaintext = null; - byte[] payload; - if (encrypted) + if (!encrypted) + { + plaintextStream = stream; + } + else { if (cek is null) { throw new ImagePasswordRequiredException(); } - var nonce = new byte[BlobNonceSize]; - stream.ReadExactly(nonce); - var tag = new byte[BlobTagSize]; - stream.ReadExactly(tag); - - using var cipherStream = new MemoryStream(); - stream.CopyTo(cipherStream); - var ciphertext = cipherStream.ToArray(); - - var plaintext = new byte[ciphertext.Length]; - try + if (chunked) { - using var aesGcm = new AesGcm(cek, BlobTagSize); - aesGcm.Decrypt(nonce, ciphertext, tag, plaintext); + var baseNonce = new byte[BlobNonceSize]; + stream.ReadExactly(baseNonce); + plaintextStream = new ChunkedGcm.ReadStream(stream, cek, baseNonce); } - catch (CryptographicException) + else { - throw new ImagePasswordIncorrectException(); - } + // Legacy whole-blob layout: a single AES-256-GCM ciphertext covering the entire + // (already gzip-compressed) payload. Kept only so pre-existing blobs keep loading. + var nonce = new byte[BlobNonceSize]; + stream.ReadExactly(nonce); + var tag = new byte[BlobTagSize]; + stream.ReadExactly(tag); + + using var cipherStream = new MemoryStream(); + stream.CopyTo(cipherStream); + var ciphertext = cipherStream.ToArray(); + + var plaintext = new byte[ciphertext.Length]; + try + { + using var aesGcm = new AesGcm(cek, BlobTagSize); + aesGcm.Decrypt(nonce, ciphertext, tag, plaintext); + } + catch (CryptographicException) + { + throw new ImagePasswordIncorrectException(); + } - payload = plaintext; - } - else - { - using var raw = new MemoryStream(); - stream.CopyTo(raw); - payload = raw.ToArray(); + legacyPlaintext = plaintext; + plaintextStream = new MemoryStream(plaintext, writable: false); + } } - byte[] content; - if (compressed) + var sourceStream = compressed + ? new GZipStream(plaintextStream, CompressionMode.Decompress) + : plaintextStream; + + var aligned = FileNode.AlignToAllocationUnit(allocationSize); + var content = FileContent.CreateZeroed(aligned); + + try { - using var payloadStream = new MemoryStream(payload); - using var gzip = new GZipStream(payloadStream, CompressionMode.Decompress); - using var decompressed = new MemoryStream(); - gzip.CopyTo(decompressed); - content = decompressed.ToArray(); + content.FillFromStream(sourceStream, (long)fileSize); } - else + catch (CryptographicException) { - content = payload; + throw new ImagePasswordIncorrectException(); } - - if ((ulong)content.Length != fileSize) + finally { - throw new InvalidDataException( - $"Snapshot blob for '{nodePath}' (hash {Convert.ToHexStringLower(hash)}) has unexpected length " + - $"{content.Length} bytes; expected {fileSize}. The snapshot may be corrupted."); + if (compressed) + { + sourceStream.Dispose(); + } + + if (legacyPlaintext is not null) + { + CryptographicOperations.ZeroMemory(legacyPlaintext); + } } - var aligned = FileNode.AlignToAllocationUnit(allocationSize); - return FileContent.FromSpan(content, aligned); + return content; } private static void ReadHeader(BinaryReader reader) @@ -499,7 +535,7 @@ private static void WriteNode(BinaryWriter writer, string path, FileNode node, s return; } - var fileSize = (int)Math.Min(node.FileInfo.FileSize, (ulong)node.FileData.Length); + var fileSize = (long)Math.Min(node.FileInfo.FileSize, (ulong)node.FileData.Length); byte[] hash; using (var incrementalHash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256)) diff --git a/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs b/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs index c46f052..e8c30ab 100644 --- a/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs +++ b/tests/ManagedDrive.Tests/DiskImageSerializerTests.cs @@ -261,7 +261,7 @@ public void SaveThenLoad_WithPasswordAcrossMultipleChunks_RoundTrips(ImageCompre { // Force a tiny chunk size so a handful of KB of node data spans several chunks, // exercising the version-4 chunked AES-GCM path without allocating a real 64 MB buffer. - DiskImageSerializer.TestChunkSizeOverride = 64; + ChunkedGcm.TestChunkSizeOverride = 64; var path = Path.Combine(Path.GetTempPath(), $"{Guid.NewGuid()}.mdr"); try { @@ -291,7 +291,7 @@ public void SaveThenLoad_WithPasswordAcrossMultipleChunks_RoundTrips(ImageCompre } finally { - DiskImageSerializer.TestChunkSizeOverride = null; + ChunkedGcm.TestChunkSizeOverride = null; File.Delete(path); } } @@ -299,7 +299,7 @@ public void SaveThenLoad_WithPasswordAcrossMultipleChunks_RoundTrips(ImageCompre [Fact] public void SaveThenLoad_WithWrongPasswordAcrossMultipleChunks_ThrowsPasswordIncorrect() { - DiskImageSerializer.TestChunkSizeOverride = 64; + ChunkedGcm.TestChunkSizeOverride = 64; var path = Path.Combine(Path.GetTempPath(), $"{Guid.NewGuid()}.mdr"); try { @@ -318,7 +318,7 @@ public void SaveThenLoad_WithWrongPasswordAcrossMultipleChunks_ThrowsPasswordInc } finally { - DiskImageSerializer.TestChunkSizeOverride = null; + ChunkedGcm.TestChunkSizeOverride = null; File.Delete(path); } } diff --git a/tests/ManagedDrive.Tests/SnapshotManagerTests.cs b/tests/ManagedDrive.Tests/SnapshotManagerTests.cs index a5c15a6..03c2b90 100644 --- a/tests/ManagedDrive.Tests/SnapshotManagerTests.cs +++ b/tests/ManagedDrive.Tests/SnapshotManagerTests.cs @@ -1,3 +1,5 @@ +using System.Security.Cryptography; + namespace ManagedDrive.Tests; public sealed class SnapshotManagerTests : IDisposable @@ -315,6 +317,103 @@ public void LoadSnapshot_EncryptedBlobWithWrongCek_ThrowsPasswordIncorrect() SnapshotManager.LoadSnapshot(snapshot.Path, out _, out _, wrongCek)); } + [Fact] + public void WriteSnapshot_EncryptedAcrossMultipleChunks_RoundTrips() + { + // Force a tiny chunk size so a small file's blob spans several chunks, exercising the + // chunked AES-GCM blob path without allocating a real 64 MB buffer. + ChunkedGcm.TestChunkSizeOverride = 64; + try + { + var cek = DiskImageSerializer.GenerateCek(); + var content = System.Text.Encoding.UTF8.GetBytes(string.Concat(Enumerable.Range(0, 20).Select(i => $"line {i} of content;"))); + + var nodeMap = new FileNodeMap(); + nodeMap.Add("\\", MakeDir()); + nodeMap.Add("\\a.txt", MakeFile(content)); + SnapshotManager.WriteSnapshot(nodeMap, 1024, "Label", _mainImagePath, + new(2026, 1, 1, 0, 0, 0, TimeSpan.Zero), ImageCompressionLevel.Fastest, cek); + + var snapshot = Assert.Single(SnapshotManager.ListSnapshots(_mainImagePath)); + var loaded = SnapshotManager.LoadSnapshot(snapshot.Path, out _, out _, cek); + + Assert.True(loaded.TryGet("\\a.txt", out var node)); + Assert.Equal(content, node!.FileData!.ToArray(content.Length)); + } + finally + { + ChunkedGcm.TestChunkSizeOverride = null; + } + } + + [Fact] + public void LoadSnapshot_EncryptedAcrossMultipleChunksWithWrongCek_ThrowsPasswordIncorrect() + { + ChunkedGcm.TestChunkSizeOverride = 64; + try + { + var cek = DiskImageSerializer.GenerateCek(); + var wrongCek = DiskImageSerializer.GenerateCek(); + var content = System.Text.Encoding.UTF8.GetBytes(string.Concat(Enumerable.Range(0, 20).Select(i => $"line {i} of content;"))); + + var nodeMap = new FileNodeMap(); + nodeMap.Add("\\", MakeDir()); + nodeMap.Add("\\a.txt", MakeFile(content)); + SnapshotManager.WriteSnapshot(nodeMap, 1024, "Label", _mainImagePath, + new(2026, 1, 1, 0, 0, 0, TimeSpan.Zero), ImageCompressionLevel.Fastest, cek); + + var snapshot = Assert.Single(SnapshotManager.ListSnapshots(_mainImagePath)); + + Assert.Throws(() => + SnapshotManager.LoadSnapshot(snapshot.Path, out _, out _, wrongCek)); + } + finally + { + ChunkedGcm.TestChunkSizeOverride = null; + } + } + + [Fact] + public void LoadSnapshot_LegacyWholeBlobEncryptedBlob_StillLoads() + { + var cek = DiskImageSerializer.GenerateCek(); + var content = new byte[] { 7, 7, 7 }; + + // Write a normal encrypted snapshot first so the index and blob directory exist with the + // expected layout, then hand-overwrite the blob it produced in the pre-chunking format + // (flag without the chunked bit; nonce+tag+whole ciphertext) to simulate a blob written + // before BlobFlagChunked existed. + var nodeMap = new FileNodeMap(); + nodeMap.Add("\\", MakeDir()); + nodeMap.Add("\\a.txt", MakeFile(content)); + SnapshotManager.WriteSnapshot(nodeMap, 1024, "Label", _mainImagePath, + new(2026, 1, 1, 0, 0, 0, TimeSpan.Zero), ImageCompressionLevel.None, cek); + + var blobPath = Directory.EnumerateFiles(BlobDirectory, "*.blob", SearchOption.AllDirectories).Single(); + + var nonce = RandomNumberGenerator.GetBytes(12); + var ciphertext = new byte[content.Length]; + var tag = new byte[16]; + using (var aesGcm = new AesGcm(cek, 16)) + { + aesGcm.Encrypt(nonce, content, ciphertext, tag); + } + + using (var stream = new FileStream(blobPath, FileMode.Create, FileAccess.Write)) + { + stream.WriteByte(0b010); // Encrypted, not Compressed, not Chunked + stream.Write(nonce); + stream.Write(tag); + stream.Write(ciphertext); + } + + var snapshot = Assert.Single(SnapshotManager.ListSnapshots(_mainImagePath)); + var loaded = SnapshotManager.LoadSnapshot(snapshot.Path, out _, out _, cek); + + Assert.True(loaded.TryGet("\\a.txt", out var node)); + Assert.Equal(content, node!.FileData!.ToArray(content.Length)); + } + [Fact] public void LoadSnapshot_MissingBlob_ReturnsClearError() { From 514da3c019a49439adf9101da3f985171df2b879 Mon Sep 17 00:00:00 2001 From: Gaoyang Date: Tue, 28 Jul 2026 16:51:09 +0800 Subject: [PATCH 3/4] feat: show per-disk read/write throughput and a speed history chart Tracks cumulative bytes read/written per mounted disk (MemoryFileSystem, RamDisk), derives an instantaneous rate via ThroughputTracker, and surfaces it on the tray icon tooltip and a new SpeedHistoryChart control on the disk card. --- .../Controls/SpeedHistoryChart.xaml | 37 +++++ .../Controls/SpeedHistoryChart.xaml.cs | 119 ++++++++++++++++ src/ManagedDrive.App/GlobalUsings.cs | 1 + .../Localization/Strings.en-US.xaml | 2 + .../Localization/Strings.zh-CN.xaml | 2 + src/ManagedDrive.App/MainWindow.xaml | 31 +++++ src/ManagedDrive.App/MainWindow.xaml.cs | 64 +++++++++ .../Services/DiskNotificationService.cs | 40 ++++++ .../Services/TrayIconController.cs | 130 ++++++++++++++++-- .../Themes/AppTheme.Colors.Dark.xaml | 1 + .../Themes/AppTheme.Colors.Light.xaml | 1 + src/ManagedDrive.App/Themes/AppTheme.xaml | 22 +++ .../ViewModels/DiskViewModel.cs | 63 +++++++++ src/ManagedDrive.Cli.Core/ByteFormatter.cs | 14 ++ .../FileSystem/MemoryFileSystem.cs | 16 +++ src/ManagedDrive.Core/Mounting/RamDisk.cs | 12 ++ .../Mounting/ThroughputTracker.cs | 47 +++++++ .../ManagedDrive.Tests/ByteFormatterTests.cs | 20 +++ .../ManagedDrive.Tests.csproj | 1 + .../SpeedHistoryChartTests.cs | 47 +++++++ .../ThroughputTrackerTests.cs | 72 ++++++++++ 21 files changed, 731 insertions(+), 11 deletions(-) create mode 100644 src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml create mode 100644 src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml.cs create mode 100644 src/ManagedDrive.Core/Mounting/ThroughputTracker.cs create mode 100644 tests/ManagedDrive.Tests/ByteFormatterTests.cs create mode 100644 tests/ManagedDrive.Tests/SpeedHistoryChartTests.cs create mode 100644 tests/ManagedDrive.Tests/ThroughputTrackerTests.cs diff --git a/src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml b/src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml new file mode 100644 index 0000000..62cd61d --- /dev/null +++ b/src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml.cs b/src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml.cs new file mode 100644 index 0000000..a198b7c --- /dev/null +++ b/src/ManagedDrive.App/Controls/SpeedHistoryChart.xaml.cs @@ -0,0 +1,119 @@ +using System.Windows.Media; +using ManagedDrive.Cli.Core; + +namespace ManagedDrive.App.Controls; + +/// +/// Renders a fixed-size read/write speed history as two overlaid polylines. Intended for use +/// inside a hover-triggered — it is not +/// rendered while off-screen, so no visibility/throttling logic lives here. +/// +public partial class SpeedHistoryChart +{ + /// Identifies the dependency property. + public static readonly DependencyProperty ReadHistoryProperty = DependencyProperty.Register( + nameof(ReadHistory), typeof(IReadOnlyList), typeof(SpeedHistoryChart), + new PropertyMetadata(null, OnHistoryChanged)); + + /// Identifies the dependency property. + public static readonly DependencyProperty WriteHistoryProperty = DependencyProperty.Register( + nameof(WriteHistory), typeof(IReadOnlyList), typeof(SpeedHistoryChart), + new PropertyMetadata(null, OnHistoryChanged)); + + /// + /// Initializes a new instance of . + /// + public SpeedHistoryChart() + { + InitializeComponent(); + } + + /// + /// Gets or sets the oldest-first read-speed history (bytes/sec) to plot. + /// + public IReadOnlyList? ReadHistory + { + get => (IReadOnlyList?)GetValue(ReadHistoryProperty); + set => SetValue(ReadHistoryProperty, value); + } + + /// + /// Gets or sets the oldest-first write-speed history (bytes/sec) to plot. + /// + public IReadOnlyList? WriteHistory + { + get => (IReadOnlyList?)GetValue(WriteHistoryProperty); + set => SetValue(WriteHistoryProperty, value); + } + + /// + /// Maps a sequence of values onto polyline points that fit within + /// x , scaled by the largest value across + /// both series () so the read and write lines share one scale. + /// A flat, near-baseline line is returned when every value is zero, rather than an empty + /// point list, so the chart never looks unloaded. + /// + public static PointCollection NormalizePoints(IReadOnlyList values, double width, double height, double scaleMax) + { + var points = new PointCollection(); + if (values.Count == 0) + { + return points; + } + + var stepX = values.Count > 1 ? width / (values.Count - 1) : 0; + for (var i = 0; i < values.Count; i++) + { + var x = i * stepX; + var y = scaleMax > 0 + ? height - (values[i] / scaleMax * height) + : height; + points.Add(new(x, y)); + } + + return points; + } + + private static void OnHistoryChanged(DependencyObject d, DependencyPropertyChangedEventArgs e) => + ((SpeedHistoryChart)d).Redraw(); + + private void ChartArea_SizeChanged(object sender, SizeChangedEventArgs e) => Redraw(); + + private void Redraw() + { + var width = ChartArea.ActualWidth; + var height = ChartArea.ActualHeight; + if (width <= 0 || height <= 0) + { + return; + } + + var read = ReadHistory ?? Array.Empty(); + var write = WriteHistory ?? Array.Empty(); + var scaleMax = Math.Max( + read.Count > 0 ? read.Max() : 0.0, + write.Count > 0 ? write.Max() : 0.0); + + ReadLine.Points = NormalizePoints(read, width, height, scaleMax); + WriteLine.Points = NormalizePoints(write, width, height, scaleMax); + + TopGridLine.X1 = 0; + TopGridLine.Y1 = 0; + TopGridLine.X2 = width; + TopGridLine.Y2 = 0; + + MidGridLine.X1 = 0; + MidGridLine.Y1 = height / 2; + MidGridLine.X2 = width; + MidGridLine.Y2 = height / 2; + + BottomGridLine.X1 = 0; + BottomGridLine.Y1 = height; + BottomGridLine.X2 = width; + BottomGridLine.Y2 = height; + + AxisMaxLabel.Text = scaleMax > 0 ? ByteFormatter.FormatRate(scaleMax) : string.Empty; + AxisMidLabel.Text = scaleMax > 0 ? ByteFormatter.FormatRate(scaleMax / 2) : string.Empty; + AxisZeroLabel.Text = "0 B/s"; + } +} diff --git a/src/ManagedDrive.App/GlobalUsings.cs b/src/ManagedDrive.App/GlobalUsings.cs index 18380e2..23aa094 100644 --- a/src/ManagedDrive.App/GlobalUsings.cs +++ b/src/ManagedDrive.App/GlobalUsings.cs @@ -1,4 +1,5 @@ global using ManagedDrive.App.Cli; +global using ManagedDrive.App.Controls; global using ManagedDrive.App.Infrastructure; global using ManagedDrive.App.Localization; global using ManagedDrive.App.Models; diff --git a/src/ManagedDrive.App/Localization/Strings.en-US.xaml b/src/ManagedDrive.App/Localization/Strings.en-US.xaml index 788e22c..3d204cc 100644 --- a/src/ManagedDrive.App/Localization/Strings.en-US.xaml +++ b/src/ManagedDrive.App/Localization/Strings.en-US.xaml @@ -89,6 +89,8 @@ Not yet modified Last image save: {0} Image not yet saved + Read speed + Write speed Just now {0} min ago {0} hours ago diff --git a/src/ManagedDrive.App/Localization/Strings.zh-CN.xaml b/src/ManagedDrive.App/Localization/Strings.zh-CN.xaml index 8cfa94a..66435f5 100644 --- a/src/ManagedDrive.App/Localization/Strings.zh-CN.xaml +++ b/src/ManagedDrive.App/Localization/Strings.zh-CN.xaml @@ -89,6 +89,8 @@ 尚未写入 上次保存镜像：{0} 镜像尚未保存 + 读取速度 + 写入速度 刚才 {0} 分钟前 {0} 小时前 diff --git a/src/ManagedDrive.App/MainWindow.xaml b/src/ManagedDrive.App/MainWindow.xaml index 51ca860..de384f0 100644 --- a/src/ManagedDrive.App/MainWindow.xaml +++ b/src/ManagedDrive.App/MainWindow.xaml @@ -5,6 +5,7 @@ xmlns:mc="http://schemas.openxmlformats.org/markup-compatibility/2006" xmlns:vm="clr-namespace:ManagedDrive.App.ViewModels" xmlns:infra="clr-namespace:ManagedDrive.App.Infrastructure" + xmlns:controls="clr-namespace:ManagedDrive.App.Controls" mc:Ignorable="d" Style="{StaticResource AppWindow}" WindowStyle="None" @@ -400,6 +401,36 @@ Foreground="{DynamicResource AppForegroundLight}" Visibility="{Binding ShowLastAutoSave, Converter={StaticResource BoolToVis}}"/> + + + + + + + + + + + + + diff --git a/src/ManagedDrive.App/MainWindow.xaml.cs b/src/ManagedDrive.App/MainWindow.xaml.cs index dec517b..5830fca 100644 --- a/src/ManagedDrive.App/MainWindow.xaml.cs +++ b/src/ManagedDrive.App/MainWindow.xaml.cs @@ -1,3 +1,7 @@ +using System.Windows.Controls; +using System.Windows.Controls.Primitives; +using System.Windows.Input; + namespace ManagedDrive.App; /// @@ -5,6 +9,11 @@ namespace ManagedDrive.App; /// public partial class MainWindow { + private static readonly TimeSpan SpeedPopupCloseDelay = TimeSpan.FromMilliseconds(150); + + private DispatcherTimer? _speedPopupCloseTimer; + private Popup? _pendingCloseSpeedPopup; + /// /// Initializes the main window and binds the supplied view model. /// @@ -33,4 +42,59 @@ private void OpenAttachedContextMenu(object sender) btn.ContextMenu.IsOpen = true; } } + + /// + /// Opens the speed-history popup when the mouse enters the speed row or the popup itself, + /// cancelling any pending close scheduled for a different popup. + /// + private void SpeedRow_MouseEnter(object sender, MouseEventArgs e) + { + var popup = FindSpeedPopup(sender); + if (popup is null) + { + return; + } + + if (_pendingCloseSpeedPopup is not null && _pendingCloseSpeedPopup != popup) + { + CloseSpeedPopup(_pendingCloseSpeedPopup); + } + + _speedPopupCloseTimer?.Stop(); + _pendingCloseSpeedPopup = null; + popup.IsOpen = true; + } + + /// + /// Schedules the speed-history popup to close shortly after the mouse leaves the speed row + /// or the popup itself, so briefly crossing the gap between them doesn't flicker it shut. + /// + private void SpeedRow_MouseLeave(object sender, MouseEventArgs e) + { + var popup = FindSpeedPopup(sender); + if (popup is null) + { + return; + } + + _speedPopupCloseTimer?.Stop(); + _pendingCloseSpeedPopup = popup; + _speedPopupCloseTimer = new DispatcherTimer { Interval = SpeedPopupCloseDelay }; + _speedPopupCloseTimer.Tick += (_, _) => + { + _speedPopupCloseTimer?.Stop(); + CloseSpeedPopup(popup); + _pendingCloseSpeedPopup = null; + }; + _speedPopupCloseTimer.Start(); + } + + private static void CloseSpeedPopup(Popup popup) => popup.IsOpen = false; + + private static Popup? FindSpeedPopup(object sender) => sender switch + { + StackPanel speedRow => speedRow.Children.OfType().FirstOrDefault(), + FrameworkElement { Parent: Popup popup } => popup, + _ => null, + }; } \ No newline at end of file diff --git a/src/ManagedDrive.App/Services/DiskNotificationService.cs b/src/ManagedDrive.App/Services/DiskNotificationService.cs index 8b83ad6..d4b194e 100644 --- a/src/ManagedDrive.App/Services/DiskNotificationService.cs +++ b/src/ManagedDrive.App/Services/DiskNotificationService.cs @@ -7,6 +7,7 @@ namespace ManagedDrive.App.Services; /// public sealed class DiskNotificationService { + private readonly HashSet _highUsageDisks = []; private readonly Func _isMainWindowVisible; private readonly MainViewModel _mainViewModel; private readonly TrayIconController _trayIconController; @@ -33,12 +34,19 @@ public DiskNotificationService(MainViewModel mainViewModel, TrayIconController t vm.HighUsageWarning += OnDiskHighUsageWarning; vm.SaveFailed += OnDiskSaveFailed; vm.ActivityObserved += OnDiskActivityObserved; + vm.PropertyChanged += OnDiskPropertyChanged; vm.SetActivityTrackingEnabled(_isMainWindowVisible()); if (vm is { CapacityAdjustedOnLoad: true, Disk.Options.SourceArchivePath: null }) { OnDiskCapacityAdjusted(vm); } + + if (vm.IsHighUsage) + { + _highUsageDisks.Add(vm); + _trayIconController.SetHighUsageWarningActive(true); + } } } @@ -49,6 +57,12 @@ public DiskNotificationService(MainViewModel mainViewModel, TrayIconController t vm.HighUsageWarning -= OnDiskHighUsageWarning; vm.SaveFailed -= OnDiskSaveFailed; vm.ActivityObserved -= OnDiskActivityObserved; + vm.PropertyChanged -= OnDiskPropertyChanged; + + if (_highUsageDisks.Remove(vm)) + { + _trayIconController.SetHighUsageWarningActive(_highUsageDisks.Count > 0); + } } } }; @@ -76,6 +90,32 @@ private void OnDiskCapacityAdjusted(DiskViewModel vm) _mainViewModel.StatusText = Loc.Format("Status.CapacityAdjusted", vm.MountPoint, originalMb, newMb); } + /// + /// Tracks every disk currently in the state (which + /// fires on both the rising and falling edge, unlike the one-shot + /// event used for the balloon tip below) and reduces it to a single tray blink call: the tray + /// icon has no per-disk concept, so it only needs to know whether any disk is currently over + /// its threshold. + /// + private void OnDiskPropertyChanged(object? sender, PropertyChangedEventArgs e) + { + if (e.PropertyName != nameof(DiskViewModel.IsHighUsage) || sender is not DiskViewModel vm) + { + return; + } + + if (vm.IsHighUsage) + { + _highUsageDisks.Add(vm); + } + else + { + _highUsageDisks.Remove(vm); + } + + _trayIconController.SetHighUsageWarningActive(_highUsageDisks.Count > 0); + } + private void OnDiskHighUsageWarning(object? sender, EventArgs e) { if (sender is not DiskViewModel vm) diff --git a/src/ManagedDrive.App/Services/TrayIconController.cs b/src/ManagedDrive.App/Services/TrayIconController.cs index 2486ac3..f879457 100644 --- a/src/ManagedDrive.App/Services/TrayIconController.cs +++ b/src/ManagedDrive.App/Services/TrayIconController.cs @@ -1,3 +1,4 @@ +using System.Drawing.Drawing2D; using System.Runtime.InteropServices; namespace ManagedDrive.App.Services; @@ -16,21 +17,33 @@ public sealed class TrayIconController : IDisposable /// private static readonly TimeSpan ActivityFlashDuration = TimeSpan.FromMilliseconds(300); + /// + /// Toggle interval for the high-usage warning blink, driven by . + /// Deliberately slower than so the two remain visually distinct. + /// + private static readonly TimeSpan HighUsageBlinkInterval = TimeSpan.FromMilliseconds(800); + private readonly Dispatcher _dispatcher; private readonly DispatcherTimer _timerActivityFlash = new() { Interval = ActivityFlashDuration }; - private readonly IntPtr[] _trayActivityIconHandles = new IntPtr[3]; + private readonly DispatcherTimer _timerHighUsageBlink = new() + { + Interval = HighUsageBlinkInterval + }; + private readonly IntPtr[] _trayActivityIconHandles = new IntPtr[4]; /// - /// [0] normal, [1] read indicator, [2] write indicator. Generated once at startup from the - /// base tray icon. + /// [0] normal, [1] read indicator, [2] write indicator, [3] high-usage warning indicator. + /// Generated once at startup from the base tray icon. /// - private readonly Icon?[] _trayActivityIcons = new Icon?[3]; + private readonly Icon?[] _trayActivityIcons = new Icon?[4]; private readonly System.Windows.Forms.NotifyIcon _trayIcon; private readonly Icon _trayIconNormal; + private bool _blinkOn; + private bool _isHighUsageActive; /// /// Builds the tray icon, its context menu, and the activity-flash timer. @@ -84,10 +97,22 @@ public TrayIconController( _timerActivityFlash.Tick += (_, _) => { - SetTrayIcon(_trayActivityIcons[0]); + SetTrayIcon(_isHighUsageActive && _blinkOn ? _trayActivityIcons[3] : _trayActivityIcons[0]); _timerActivityFlash.Stop(); }; + _timerHighUsageBlink.Tick += (_, _) => + { + _blinkOn = !_blinkOn; + + // Let an in-progress read/write flash keep showing; its own Tick will pick up the + // right idle/warning icon (see above) once it reverts. + if (!_timerActivityFlash.IsEnabled) + { + SetTrayIcon(_trayActivityIcons[_blinkOn ? 3 : 0]); + } + }; + LanguageManager.Instance.LanguageChanged += (_, _) => UpdateTrayMenuHeaders(); ApplyTrayMenuTheme(); ThemeManager.Instance.ThemeChanged += (_, _) => dispatcher.Invoke(ApplyTrayMenuTheme); @@ -109,13 +134,14 @@ public bool Visible } /// - /// Stops the activity blink timer and releases the tray icon and all three generated - /// activity-indicator variants, including the HICONs backing - /// which alone would leak. + /// Stops the activity and high-usage blink timers and releases the tray icon and all four + /// generated activity-indicator variants, including the HICONs backing + /// which alone would leak. /// public void Dispose() { _timerActivityFlash.Stop(); + _timerHighUsageBlink.Stop(); _trayIcon.Dispose(); _trayIconNormal.Dispose(); @@ -157,6 +183,44 @@ public void OnActivityDetected(bool isWrite) }); } + /// + /// Starts or stops the sustained high-usage warning blink: while active, the tray icon + /// repeatedly toggles between the idle icon and the warning indicator every + /// until turned off again. Unlike + /// 's one-shot flash, this stays on for as long as the + /// caller reports the condition is active (see , + /// which aggregates every disk's IsHighUsage state into a single call here). + /// + /// Whether any disk currently exceeds its high-usage threshold. + public void SetHighUsageWarningActive(bool active) + { + _dispatcher.BeginInvoke(() => + { + if (active == _isHighUsageActive || _trayActivityIcons[3] == null) + { + return; + } + + _isHighUsageActive = active; + + if (active) + { + _blinkOn = true; + SetTrayIcon(_trayActivityIcons[3]); + _timerHighUsageBlink.Start(); + } + else + { + _timerHighUsageBlink.Stop(); + + if (!_timerActivityFlash.IsEnabled) + { + SetTrayIcon(_trayActivityIcons[0]); + } + } + }); + } + /// /// Shows a balloon tip from the tray icon. /// @@ -192,8 +256,11 @@ private void ApplyTrayMenuTheme() /// /// Generates the tray icon variants indexed by (0 = normal, - /// 1 = read indicator, 2 = write indicator) by overlaying a colored dot in the top-right - /// corner of : green for reads, orange for writes. Each generated + /// 1 = read indicator, 2 = write indicator, 3 = high-usage warning indicator) by overlaying an + /// indicator in the top-right corner of : a green dot for reads, an + /// orange dot for writes, and a red exclamation-mark badge (distinct shape, not just a + /// differently-colored dot) for the warning state — the shape difference keeps the warning + /// blink from being confused with a write flash at 16x16 tray size. Each generated /// 's backing HICON is recorded in so /// it can be released via on shutdown, since /// alone does not release a handle obtained from . @@ -203,11 +270,14 @@ private void BuildTrayActivityIcons(Icon baseIcon) var size = baseIcon.Size; var dotDiameter = Math.Max(4, size.Width / 3); var dotRect = new Rectangle(size.Width - dotDiameter, 0, dotDiameter, dotDiameter); + var badgeDiameter = Math.Max(dotDiameter, size.Height / 2); + var badgeRect = new Rectangle(size.Width - badgeDiameter, 0, badgeDiameter, badgeDiameter); Brush?[] overlayBrushes = [ null, new SolidBrush(Color.FromArgb(255, 0, 230, 118)), new SolidBrush(Color.FromArgb(255, 255, 50, 0)), + null, ]; for (var state = 0; state < _trayActivityIcons.Length; state++) @@ -216,7 +286,11 @@ private void BuildTrayActivityIcons(Icon baseIcon) using (var graphics = Graphics.FromImage(bitmap)) { graphics.DrawIcon(baseIcon, new Rectangle(0, 0, size.Width, size.Height)); - if (overlayBrushes[state] is { } brush) + if (state == 3) + { + DrawWarningBadge(graphics, badgeRect); + } + else if (overlayBrushes[state] is { } brush) { graphics.FillEllipse(brush, dotRect); } @@ -228,6 +302,40 @@ private void BuildTrayActivityIcons(Icon baseIcon) } } + /// + /// Draws a red exclamation-mark warning badge (filled triangle with a white "!" glyph built + /// from two primitives rather than rendered text, so it stays legible at tray-icon scale + /// without depending on font hinting/DPI) inside . + /// + private static void DrawWarningBadge(Graphics graphics, Rectangle bounds) + { + using var fillBrush = new SolidBrush(Color.FromArgb(255, 230, 0, 0)); + using var outlinePen = new Pen(Color.FromArgb(255, 90, 0, 0), Math.Max(1f, bounds.Width / 8f)); + using var glyphBrush = new SolidBrush(Color.White); + + graphics.SmoothingMode = SmoothingMode.AntiAlias; + + System.Drawing.Point[] triangle = + [ + new(bounds.Left + bounds.Width / 2, bounds.Top), + new(bounds.Right, bounds.Bottom), + new(bounds.Left, bounds.Bottom), + ]; + graphics.FillPolygon(fillBrush, triangle); + graphics.DrawPolygon(outlinePen, triangle); + + var stemWidth = Math.Max(1f, bounds.Width / 6f); + var stemHeight = bounds.Height * 0.40f; + var stemLeft = bounds.Left + (bounds.Width - stemWidth) / 2f; + var stemTop = bounds.Top + bounds.Height * 0.38f; + graphics.FillRectangle(glyphBrush, stemLeft, stemTop, stemWidth, stemHeight); + + var dotDiameter = stemWidth; + var dotLeft = bounds.Left + (bounds.Width - dotDiameter) / 2f; + var dotTop = bounds.Top + bounds.Height * 0.82f; + graphics.FillEllipse(glyphBrush, dotLeft, dotTop, dotDiameter, dotDiameter); + } + /// /// Assigns to the tray icon only if it differs from the current one, /// avoiding redundant reassignment that diff --git a/src/ManagedDrive.App/Themes/AppTheme.Colors.Dark.xaml b/src/ManagedDrive.App/Themes/AppTheme.Colors.Dark.xaml index 5171e90..c2a4ec5 100644 --- a/src/ManagedDrive.App/Themes/AppTheme.Colors.Dark.xaml +++ b/src/ManagedDrive.App/Themes/AppTheme.Colors.Dark.xaml @@ -16,5 +16,6 @@ + diff --git a/src/ManagedDrive.App/Themes/AppTheme.Colors.Light.xaml b/src/ManagedDrive.App/Themes/AppTheme.Colors.Light.xaml index eaf5fc8..20ef4f7 100644 --- a/src/ManagedDrive.App/Themes/AppTheme.Colors.Light.xaml +++ b/src/ManagedDrive.App/Themes/AppTheme.Colors.Light.xaml @@ -16,5 +16,6 @@ + diff --git a/src/ManagedDrive.App/Themes/AppTheme.xaml b/src/ManagedDrive.App/Themes/AppTheme.xaml index 4232017..3f5b634 100644 --- a/src/ManagedDrive.App/Themes/AppTheme.xaml +++ b/src/ManagedDrive.App/Themes/AppTheme.xaml @@ -677,6 +677,28 @@ + + + + + + + + + + + + + + + + + diff --git a/src/ManagedDrive.App/ViewModels/DiskViewModel.cs b/src/ManagedDrive.App/ViewModels/DiskViewModel.cs index 830491a..1fa981e 100644 --- a/src/ManagedDrive.App/ViewModels/DiskViewModel.cs +++ b/src/ManagedDrive.App/ViewModels/DiskViewModel.cs @@ -22,14 +22,29 @@ public sealed class DiskViewModel : INotifyPropertyChanged, IDisposable /// private static readonly TimeSpan ActivityThrottleWindow = TimeSpan.FromMilliseconds(300); + /// + /// Number of speed samples retained for the read/write history popup: 30 minutes at the + /// 2-second cadence. Sized generously since the buffer is only + /// rendered on hover — sampling itself is just two array writes per tick, independent of + /// whether anyone is looking at it. + /// + private const int SpeedHistoryLength = 900; + private readonly DispatcherTimer _activityThrottleTimer; private readonly DispatcherTimer _refreshTimer; + private readonly double[] _readSpeedHistory = new double[SpeedHistoryLength]; + private readonly ThroughputTracker _readThroughput = new(); + private readonly double[] _writeSpeedHistory = new double[SpeedHistoryLength]; + private readonly ThroughputTracker _writeThroughput = new(); private bool _activityTrackingEnabled; private ulong _freeBytes; private bool _isCurrentTempDir; private DiskActivityEventArgs? _pendingActivity; + private double _readBytesPerSecond; + private int _speedHistoryHead; private ulong _usedBytes; + private double _writeBytesPerSecond; /// /// Initializes a new view model for . @@ -275,6 +290,29 @@ public RelayCommand OpenInExplorerCommand /// public string? SourcePath => Disk.Options.SourceArchivePath ?? Disk.Options.PersistImagePath; + /// + /// Gets the most recently sampled read throughput, formatted as e.g. "1.2 MB/s". + /// + public string ReadSpeedFormatted => ByteFormatter.FormatRate(_readBytesPerSecond); + + /// + /// Gets the last 30 minutes of sampled read-speed history (bytes/sec), oldest first, for + /// display in the hover-triggered history chart. Only meaningfully consumed while that + /// popup is open; sampling itself runs unconditionally in . + /// + public IReadOnlyList ReadSpeedHistory => SnapshotHistory(_readSpeedHistory); + + /// + /// Gets the most recently sampled write throughput, formatted as e.g. "1.2 MB/s". + /// + public string WriteSpeedFormatted => ByteFormatter.FormatRate(_writeBytesPerSecond); + + /// + /// Gets the last 30 minutes of sampled write-speed history (bytes/sec), oldest first. See + /// . + /// + public IReadOnlyList WriteSpeedHistory => SnapshotHistory(_writeSpeedHistory); + /// /// Gets the amount of used space formatted as a human-readable string. /// @@ -302,6 +340,20 @@ public void Dispose() _activityThrottleTimer.Tick -= OnActivityThrottleTick; SetActivityTrackingEnabled(false); Disk.SaveFailed -= OnDiskSaveFailed; + _readThroughput.Reset(); + _writeThroughput.Reset(); + } + + /// + /// Reorders a fixed-length ring buffer written via into + /// oldest-first order for display. + /// + private double[] SnapshotHistory(double[] buffer) + { + var result = new double[buffer.Length]; + Array.Copy(buffer, _speedHistoryHead, result, 0, buffer.Length - _speedHistoryHead); + Array.Copy(buffer, 0, result, buffer.Length - _speedHistoryHead, _speedHistoryHead); + return result; } /// @@ -317,6 +369,13 @@ public void Refresh() _usedBytes = Disk.UsedBytes; _freeBytes = Disk.FreeBytes; + var now = DateTimeOffset.UtcNow; + _readBytesPerSecond = _readThroughput.Sample(Disk.TotalBytesRead, now); + _writeBytesPerSecond = _writeThroughput.Sample(Disk.TotalBytesWritten, now); + _readSpeedHistory[_speedHistoryHead] = _readBytesPerSecond; + _writeSpeedHistory[_speedHistoryHead] = _writeBytesPerSecond; + _speedHistoryHead = (_speedHistoryHead + 1) % SpeedHistoryLength; + var isUiVisible = Application.Current?.MainWindow is { IsVisible: true }; if (isUiVisible) { @@ -333,6 +392,10 @@ public void Refresh() OnPropertyChanged(nameof(SourcePath)); OnPropertyChanged(nameof(HasImagePath)); OnPropertyChanged(nameof(IsPasswordProtected)); + OnPropertyChanged(nameof(ReadSpeedFormatted)); + OnPropertyChanged(nameof(ReadSpeedHistory)); + OnPropertyChanged(nameof(WriteSpeedFormatted)); + OnPropertyChanged(nameof(WriteSpeedHistory)); IsCurrentTempDir = CheckIsCurrentTempDir(); } diff --git a/src/ManagedDrive.Cli.Core/ByteFormatter.cs b/src/ManagedDrive.Cli.Core/ByteFormatter.cs index a5f81f4..32609be 100644 --- a/src/ManagedDrive.Cli.Core/ByteFormatter.cs +++ b/src/ManagedDrive.Cli.Core/ByteFormatter.cs @@ -28,4 +28,18 @@ public static string Format(ulong bytes) return $"{bytes} B"; } + + /// + /// Formats a byte-per-second rate using the same unit thresholds as , + /// with an "/s" suffix (e.g. "1.2 MB/s"). Values below 1 B/s are reported as "0 B/s". + /// + public static string FormatRate(double bytesPerSecond) + { + if (bytesPerSecond < 1.0) + { + return "0 B/s"; + } + + return $"{Format((ulong)Math.Round(bytesPerSecond))}/s"; + } } \ No newline at end of file diff --git a/src/ManagedDrive.Core/FileSystem/MemoryFileSystem.cs b/src/ManagedDrive.Core/FileSystem/MemoryFileSystem.cs index 948341c..81b8ee2 100644 --- a/src/ManagedDrive.Core/FileSystem/MemoryFileSystem.cs +++ b/src/ManagedDrive.Core/FileSystem/MemoryFileSystem.cs @@ -20,6 +20,8 @@ public sealed class MemoryFileSystem : FileSystemBase private long _lastContentReadTicks; private ContentAccessInfo? _lastContentWriteAccess; private long _lastContentWriteTicks; + private long _totalBytesRead; + private long _totalBytesWritten; private ulong _maxCapacity; private string _volumeLabel; @@ -107,6 +109,18 @@ internal DateTimeOffset? LastContentWriteTimeUtc } } + /// + /// Gets the cumulative number of bytes read from file content since mount. Never resets; + /// consumers derive a rate by sampling the delta between two reads of this value over time. + /// + internal long TotalBytesRead => Interlocked.Read(ref _totalBytesRead); + + /// + /// Gets the cumulative number of bytes written to file content since mount. Never resets; + /// consumers derive a rate by sampling the delta between two reads of this value over time. + /// + internal long TotalBytesWritten => Interlocked.Read(ref _totalBytesWritten); + /// /// Exposes the underlying node map for serialization and capacity queries. /// @@ -516,6 +530,7 @@ public override int Read( { node.FileData.ReadTo(offset, buffer, toRead); bytesTransferred = toRead; + Interlocked.Add(ref _totalBytesRead, toRead); var readNow = DateTimeOffset.UtcNow; Interlocked.Exchange(ref _lastContentReadTicks, readNow.UtcTicks); Interlocked.Exchange(ref _lastContentReadAccess, new(readNow, node.FilePath)); @@ -773,6 +788,7 @@ public override int Write( } bytesTransferred = length; + Interlocked.Add(ref _totalBytesWritten, length); var now = FileTimeNow(); node.FileInfo.LastAccessTime = now; diff --git a/src/ManagedDrive.Core/Mounting/RamDisk.cs b/src/ManagedDrive.Core/Mounting/RamDisk.cs index 595a41a..829ce9b 100644 --- a/src/ManagedDrive.Core/Mounting/RamDisk.cs +++ b/src/ManagedDrive.Core/Mounting/RamDisk.cs @@ -93,6 +93,18 @@ private RamDisk(MemoryFileSystem fs, FileSystemHost host, DiskOptions options) /// public DateTimeOffset? LastContentWriteTime => _fs.LastContentWriteTimeUtc; + /// + /// Gets the cumulative number of bytes read from file content since mount. Never resets; + /// consumers derive a rate by sampling the delta over time (see ). + /// + public long TotalBytesRead => _fs.TotalBytesRead; + + /// + /// Gets the cumulative number of bytes written to file content since mount. Never resets; + /// consumers derive a rate by sampling the delta over time (see ). + /// + public long TotalBytesWritten => _fs.TotalBytesWritten; + /// /// Gets the UTC timestamp of the most recent successful image save (auto-save, final /// save on unmount, or manual save via ). null if no diff --git a/src/ManagedDrive.Core/Mounting/ThroughputTracker.cs b/src/ManagedDrive.Core/Mounting/ThroughputTracker.cs new file mode 100644 index 0000000..88363cd --- /dev/null +++ b/src/ManagedDrive.Core/Mounting/ThroughputTracker.cs @@ -0,0 +1,47 @@ +namespace ManagedDrive.Core.Mounting; + +/// +/// Converts a cumulative byte counter sampled at arbitrary intervals into an instantaneous +/// bytes/sec rate. Pure logic, independent of WinFsp — the counter and clock are supplied by +/// the caller on every call. +/// +public sealed class ThroughputTracker +{ + private long? _lastBytes; + private DateTimeOffset? _lastSampleTime; + + /// + /// Feeds a new cumulative byte total and returns the instantaneous rate in bytes/sec since + /// the previous call. The first call for a given instance has no baseline and returns 0. + /// + public double Sample(long cumulativeBytes, DateTimeOffset now) + { + if (_lastBytes is not { } lastBytes || _lastSampleTime is not { } lastTime) + { + _lastBytes = cumulativeBytes; + _lastSampleTime = now; + return 0.0; + } + + var elapsed = (now - lastTime).TotalSeconds; + _lastBytes = cumulativeBytes; + _lastSampleTime = now; + + if (elapsed <= 0) + { + return 0.0; + } + + var delta = cumulativeBytes - lastBytes; + return delta <= 0 ? 0.0 : delta / elapsed; + } + + /// + /// Discards the current baseline so the next call starts fresh. + /// + public void Reset() + { + _lastBytes = null; + _lastSampleTime = null; + } +} diff --git a/tests/ManagedDrive.Tests/ByteFormatterTests.cs b/tests/ManagedDrive.Tests/ByteFormatterTests.cs new file mode 100644 index 0000000..1c76019 --- /dev/null +++ b/tests/ManagedDrive.Tests/ByteFormatterTests.cs @@ -0,0 +1,20 @@ +using ManagedDrive.Cli.Core; + +namespace ManagedDrive.Tests; + +public sealed class ByteFormatterTests +{ + [Theory] + [InlineData(0.0, "0 B/s")] + [InlineData(0.5, "0 B/s")] + [InlineData(512.0, "512 B/s")] + [InlineData(1024.0 * 5, "5.0 KB/s")] + [InlineData(1024.0 * 1024 * 3, "3.0 MB/s")] + [InlineData(1024.0 * 1024 * 1024 * 2, "2.0 GB/s")] + public void FormatRate_UsesSameUnitThresholdsAsFormat(double bytesPerSecond, string expected) + { + var formatted = ByteFormatter.FormatRate(bytesPerSecond); + + Assert.Equal(expected, formatted); + } +} diff --git a/tests/ManagedDrive.Tests/ManagedDrive.Tests.csproj b/tests/ManagedDrive.Tests/ManagedDrive.Tests.csproj index 6201e0a..447290c 100644 --- a/tests/ManagedDrive.Tests/ManagedDrive.Tests.csproj +++ b/tests/ManagedDrive.Tests/ManagedDrive.Tests.csproj @@ -22,6 +22,7 @@ + diff --git a/tests/ManagedDrive.Tests/SpeedHistoryChartTests.cs b/tests/ManagedDrive.Tests/SpeedHistoryChartTests.cs new file mode 100644 index 0000000..ff69d5d --- /dev/null +++ b/tests/ManagedDrive.Tests/SpeedHistoryChartTests.cs @@ -0,0 +1,47 @@ +namespace ManagedDrive.Tests; + +public sealed class SpeedHistoryChartTests +{ + [Fact] + public void NormalizePoints_EmptyValues_ReturnsEmptyCollection() + { + var points = SpeedHistoryChart.NormalizePoints(Array.Empty(), width: 100, height: 50, scaleMax: 0); + + Assert.Empty(points); + } + + [Fact] + public void NormalizePoints_AllZero_ReturnsFlatLineAtBaseline() + { + var points = SpeedHistoryChart.NormalizePoints([0, 0, 0], width: 100, height: 50, scaleMax: 0); + + Assert.All(points, p => Assert.Equal(50, p.Y)); + } + + [Fact] + public void NormalizePoints_MaxValue_TouchesTop() + { + var points = SpeedHistoryChart.NormalizePoints([0, 100], width: 100, height: 50, scaleMax: 100); + + Assert.Equal(50, points[0].Y); + Assert.Equal(0, points[1].Y); + } + + [Fact] + public void NormalizePoints_SinglePoint_DoesNotDivideByZeroWidth() + { + var points = SpeedHistoryChart.NormalizePoints([42], width: 100, height: 50, scaleMax: 42); + + Assert.Single(points); + Assert.Equal(0, points[0].X); + } + + [Fact] + public void NormalizePoints_SpansFullWidth() + { + var points = SpeedHistoryChart.NormalizePoints([0, 0, 0, 0], width: 90, height: 50, scaleMax: 0); + + Assert.Equal(0, points[0].X); + Assert.Equal(90, points[3].X); + } +} diff --git a/tests/ManagedDrive.Tests/ThroughputTrackerTests.cs b/tests/ManagedDrive.Tests/ThroughputTrackerTests.cs new file mode 100644 index 0000000..c2eb4d5 --- /dev/null +++ b/tests/ManagedDrive.Tests/ThroughputTrackerTests.cs @@ -0,0 +1,72 @@ +namespace ManagedDrive.Tests; + +public sealed class ThroughputTrackerTests +{ + private static readonly DateTimeOffset BaseTime = new(2026, 1, 1, 0, 0, 0, TimeSpan.Zero); + + [Fact] + public void Sample_FirstCall_ReturnsZeroAndEstablishesBaseline() + { + var tracker = new ThroughputTracker(); + + var rate = tracker.Sample(1000, BaseTime); + + Assert.Equal(0.0, rate); + } + + [Fact] + public void Sample_SteadyGrowth_ReturnsExpectedBytesPerSecond() + { + var tracker = new ThroughputTracker(); + tracker.Sample(0, BaseTime); + + var rate = tracker.Sample(2000, BaseTime.AddSeconds(2)); + + Assert.Equal(1000.0, rate); + } + + [Fact] + public void Sample_NoGrowth_ReturnsZero() + { + var tracker = new ThroughputTracker(); + tracker.Sample(1000, BaseTime); + + var rate = tracker.Sample(1000, BaseTime.AddSeconds(2)); + + Assert.Equal(0.0, rate); + } + + [Fact] + public void Sample_NonMonotonicDecrease_ReturnsZeroNotNegative() + { + var tracker = new ThroughputTracker(); + tracker.Sample(1000, BaseTime); + + var rate = tracker.Sample(500, BaseTime.AddSeconds(2)); + + Assert.Equal(0.0, rate); + } + + [Fact] + public void Sample_ZeroElapsedTime_ReturnsZeroNotInfinity() + { + var tracker = new ThroughputTracker(); + tracker.Sample(1000, BaseTime); + + var rate = tracker.Sample(2000, BaseTime); + + Assert.Equal(0.0, rate); + } + + [Fact] + public void Reset_ClearsBaseline_NextSampleReturnsZero() + { + var tracker = new ThroughputTracker(); + tracker.Sample(1000, BaseTime); + tracker.Reset(); + + var rate = tracker.Sample(5000, BaseTime.AddSeconds(2)); + + Assert.Equal(0.0, rate); + } +} From 1adfe6ef74585bc780a8b435d307f33fb4e92083 Mon Sep 17 00:00:00 2001 From: Gaoyang Date: Tue, 28 Jul 2026 16:59:27 +0800 Subject: [PATCH 4/4] fix: address code-review findings on snapshot/image/tray changes - SnapshotStore.ReadBlob: restore the corrupted-length check that was dropped when the blob read path was rewritten to stream via FillFromStream; FillFromStream now returns the actual byte count read so callers can detect truncated/corrupt sources. - DiskImageSerializer.LoadCurrent: dispatch encrypted-region loading on an explicit version switch (3/4) instead of a boolean version == 3 check, so an unexpected version value fails loudly instead of silently falling into the chunked path. - TrayIconController.SetHighUsageWarningActive: record the active state before bailing out on missing tray icons, so a call that arrives before icon generation finishes doesn't get silently dropped and desync future calls from the real state. - MainWindow: resolve the speed-history popup via an explicit Tag binding instead of FrameworkElement.Parent, which is not a reliable way to reach a Popup from its hosted content. --- src/ManagedDrive.App/MainWindow.xaml | 1 + src/ManagedDrive.App/MainWindow.xaml.cs | 11 ++++++++++- .../Services/TrayIconController.cs | 11 ++++++++++- src/ManagedDrive.Core/FileSystem/FileContent.cs | 13 ++++++++++--- .../Persistence/DiskImageSerializer.cs | 9 ++++++--- src/ManagedDrive.Core/Snapshots/SnapshotStore.cs | 10 +++++++++- tests/ManagedDrive.Tests/SnapshotManagerTests.cs | 16 ++++++++++++++++ 7 files changed, 62 insertions(+), 9 deletions(-) diff --git a/src/ManagedDrive.App/MainWindow.xaml b/src/ManagedDrive.App/MainWindow.xaml index de384f0..fa7bd9a 100644 --- a/src/ManagedDrive.App/MainWindow.xaml +++ b/src/ManagedDrive.App/MainWindow.xaml @@ -425,6 +425,7 @@ diff --git a/src/ManagedDrive.App/MainWindow.xaml.cs b/src/ManagedDrive.App/MainWindow.xaml.cs index 5830fca..ad44ec3 100644 --- a/src/ManagedDrive.App/MainWindow.xaml.cs +++ b/src/ManagedDrive.App/MainWindow.xaml.cs @@ -91,10 +91,19 @@ private void SpeedRow_MouseLeave(object sender, MouseEventArgs e) private static void CloseSpeedPopup(Popup popup) => popup.IsOpen = false; + /// + /// Resolves the speed-history for a mouse event raised either by the + /// speed row itself or by the popup's own content border. The border's Tag is bound to + /// the popup by name in XAML (Tag="{Binding ElementName=SpeedPopup}") rather than + /// resolved here via FrameworkElement.Parent — a 's child is hosted + /// in a separate visual root, so relying on logical-tree parentage to hold at arbitrary + /// MouseEnter/MouseLeave timing is fragile; the explicit binding is resolved once when the + /// template loads, before any such event can fire. + /// private static Popup? FindSpeedPopup(object sender) => sender switch { StackPanel speedRow => speedRow.Children.OfType().FirstOrDefault(), - FrameworkElement { Parent: Popup popup } => popup, + FrameworkElement { Tag: Popup popup } => popup, _ => null, }; } \ No newline at end of file diff --git a/src/ManagedDrive.App/Services/TrayIconController.cs b/src/ManagedDrive.App/Services/TrayIconController.cs index f879457..0e6ce9b 100644 --- a/src/ManagedDrive.App/Services/TrayIconController.cs +++ b/src/ManagedDrive.App/Services/TrayIconController.cs @@ -196,13 +196,22 @@ public void SetHighUsageWarningActive(bool active) { _dispatcher.BeginInvoke(() => { - if (active == _isHighUsageActive || _trayActivityIcons[3] == null) + if (active == _isHighUsageActive) { return; } + // Record the state change before checking icon readiness below, so a call that + // arrives before icon generation finishes doesn't get treated as a no-op — otherwise + // a later, genuine call with the same target value would be dropped by the + // short-circuit above, permanently losing the warning blink. _isHighUsageActive = active; + if (_trayActivityIcons[3] == null) + { + return; + } + if (active) { _blinkOn = true; diff --git a/src/ManagedDrive.Core/FileSystem/FileContent.cs b/src/ManagedDrive.Core/FileSystem/FileContent.cs index 88bca24..46851d3 100644 --- a/src/ManagedDrive.Core/FileSystem/FileContent.cs +++ b/src/ManagedDrive.Core/FileSystem/FileContent.cs @@ -265,14 +265,18 @@ public void CopyTo(Stream destination, long count) /// Reads up to bytes from into the content /// starting at offset 0. The content must already be at least bytes. /// If the stream ends early, the remaining bytes keep their current (zero) value — mirroring a - /// bounded copy rather than throwing. + /// bounded copy rather than throwing. Returns the number of bytes actually read, so callers + /// that need to detect a truncated/corrupt source can compare it against + /// themselves. /// /// The stream to read from. /// Maximum number of bytes to read. - public void FillFromStream(Stream source, long count) + /// The number of bytes actually filled, which is less than if ended early. + public long FillFromStream(Stream source, long count) { var remaining = count; var chunkIndex = 0; + var totalFilled = 0L; while (remaining > 0) { @@ -284,15 +288,18 @@ public void FillFromStream(Stream source, long count) var read = source.Read(_chunks[chunkIndex], filled, segment - filled); if (read == 0) { - return; + return totalFilled + filled; } filled += read; } + totalFilled += filled; remaining -= segment; chunkIndex++; } + + return totalFilled; } /// diff --git a/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs b/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs index b2b9714..95392b1 100644 --- a/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs +++ b/src/ManagedDrive.Core/Persistence/DiskImageSerializer.cs @@ -379,9 +379,12 @@ private static FileNodeMap LoadCurrent( var resolvedCek = UnwrapCek(wrappedCek, password, salt, iterations, wrapNonce, wrapTag); cek = resolvedCek; - return version == 3 - ? LoadLegacyEncryptedBlob(stream, reader, resolvedCek, compressed) - : LoadChunkedEncrypted(stream, reader, resolvedCek, compressed); + return version switch + { + 3 => LoadLegacyEncryptedBlob(stream, reader, resolvedCek, compressed), + 4 => LoadChunkedEncrypted(stream, reader, resolvedCek, compressed), + _ => throw new InvalidDataException($"Unsupported image version: {version}."), + }; } /// diff --git a/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs b/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs index d7fa833..1fc0994 100644 --- a/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs +++ b/src/ManagedDrive.Core/Snapshots/SnapshotStore.cs @@ -407,10 +407,11 @@ private static FileContent ReadBlob(string blobDirectory, byte[] hash, string no var aligned = FileNode.AlignToAllocationUnit(allocationSize); var content = FileContent.CreateZeroed(aligned); + long filled; try { - content.FillFromStream(sourceStream, (long)fileSize); + filled = content.FillFromStream(sourceStream, (long)fileSize); } catch (CryptographicException) { @@ -429,6 +430,13 @@ private static FileContent ReadBlob(string blobDirectory, byte[] hash, string no } } + if ((ulong)filled != fileSize) + { + throw new InvalidDataException( + $"Snapshot blob for '{nodePath}' (hash {Convert.ToHexStringLower(hash)}) has unexpected length " + + $"{filled} bytes; expected {fileSize}. The snapshot may be corrupted."); + } + return content; } diff --git a/tests/ManagedDrive.Tests/SnapshotManagerTests.cs b/tests/ManagedDrive.Tests/SnapshotManagerTests.cs index 03c2b90..7054924 100644 --- a/tests/ManagedDrive.Tests/SnapshotManagerTests.cs +++ b/tests/ManagedDrive.Tests/SnapshotManagerTests.cs @@ -414,6 +414,22 @@ public void LoadSnapshot_LegacyWholeBlobEncryptedBlob_StillLoads() Assert.Equal(content, node!.FileData!.ToArray(content.Length)); } + [Fact] + public void LoadSnapshot_TruncatedBlob_ThrowsInvalidData() + { + WriteSnapshotWithFile(new(2026, 1, 1, 0, 0, 0, TimeSpan.Zero), "\\a.txt", [1, 2, 3, 4, 5]); + + var snapshot = Assert.Single(SnapshotManager.ListSnapshots(_mainImagePath)); + var blobFile = Directory.EnumerateFiles(BlobDirectory, "*.blob", SearchOption.AllDirectories).Single(); + + // Truncate the (uncompressed, unencrypted) blob so its content is shorter than the file + // size recorded in the index, simulating a corrupted/incompletely-written blob. + var bytes = File.ReadAllBytes(blobFile); + File.WriteAllBytes(blobFile, bytes[..^2]); + + Assert.Throws(() => SnapshotManager.LoadSnapshot(snapshot.Path, out _, out _)); + } + [Fact] public void LoadSnapshot_MissingBlob_ReturnsClearError() {