fix(flash): Improve sparse flashing, download handling & tests

Add tests and mocks for multipart downloads and resparse/merge scenarios, including a new SparseFileResparseTests and MultiPartDownloadCaptureTransport. Improve robustness of download/flash logic: tighten DATA-response handling (Bytes returns a Fail response; Stream throws on protocol mismatch), add debug logs (FastbootDebug.Log) for response handling and steps, and refine FlashSparseFile to choose part sizes via FindMaxBlockCountWithinLimit (binary search) while preserving original total blocks and providing a fallback for tiny limits. Also adjust NotifyCurrentStep/NotifyProgress to print default console output when no handlers are attached and add related unit tests. Minor comment and using cleanup in FastbootDriver.

Author: 13584452567

FirmwareKit.Comm.Fastboot.Tests/FastbootProtocolTests.cs

@@ -1,5 +1,6 @@
 using FirmwareKit.Comm.Fastboot.Usb;
 using FirmwareKit.Sparse.Core;
+using FirmwareKit.Sparse.Models;
 using System.Globalization;
 using System.Text;
 
@@ -193,6 +194,78 @@ public long Write(byte[] data, int length)
             public void Dispose() { }
         }
 
+        private sealed class MultiPartDownloadCaptureTransport : IFastbootTransport
+        {
+            private readonly Dictionary<string, string> _responses;
+            private readonly Queue<byte[]> _readQueue = new();
+            private MemoryStream? _currentPayload;
+            private int _pendingDownloadBytes;
+
+            public List<string> Commands { get; } = new();
+            public List<byte[]> DownloadPayloads { get; } = new();
+
+            public MultiPartDownloadCaptureTransport(Dictionary<string, string>? responses = null)
+            {
+                _responses = responses ?? new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
+            }
+
+            public byte[] Read(int length)
+            {
+                if (_readQueue.Count == 0)
+                {
+                    return Encoding.UTF8.GetBytes("OKAY");
+                }
+
+                return _readQueue.Dequeue();
+            }
+
+            public long Write(byte[] data, int length)
+            {
+                if (_pendingDownloadBytes > 0)
+                {
+                    _currentPayload ??= new MemoryStream();
+                    _currentPayload.Write(data, 0, length);
+                    _pendingDownloadBytes -= length;
+                    if (_pendingDownloadBytes <= 0)
+                    {
+                        DownloadPayloads.Add(_currentPayload.ToArray());
+                        _currentPayload.Dispose();
+                        _currentPayload = null;
+                        _readQueue.Enqueue(Encoding.UTF8.GetBytes("OKAY"));
+                    }
+
+                    return length;
+                }
+
+                string command = Encoding.UTF8.GetString(data, 0, length);
+                Commands.Add(command);
+
+                if (command.StartsWith("download:", StringComparison.OrdinalIgnoreCase))
+                {
+                    string hex = command.Substring("download:".Length);
+                    int size = int.Parse(hex, NumberStyles.HexNumber, CultureInfo.InvariantCulture);
+                    _pendingDownloadBytes = size;
+                    _currentPayload = new MemoryStream(size);
+                    _readQueue.Enqueue(Encoding.UTF8.GetBytes($"DATA{size:x8}"));
+                    return length;
+                }
+
+                if (_responses.TryGetValue(command, out string? response))
+                {
+                    _readQueue.Enqueue(Encoding.UTF8.GetBytes(response));
+                    return length;
+                }
+
+                _readQueue.Enqueue(Encoding.UTF8.GetBytes("OKAY"));
+                return length;
+            }
+
+            public void Dispose()
+            {
+                _currentPayload?.Dispose();
+            }
+        }
+
         private sealed class NonSeekableStream : Stream
         {
             private readonly Stream _inner;
@@ -408,6 +481,23 @@ public void DownloadDataStream_Retry_UsesInitialStreamPosition()
             Assert.Equal(new byte[] { 0x01, 0x02, 0x03, 0x04 }, transport.DownloadPayload.ToArray());
         }
 
+        [Fact]
+        public void DownloadDataStream_Retry_WhenDownloadHandshakeIsNotData()
+        {
+            var transport = new MockTransport();
+            transport.EnqueueResponse("OKAYno");
+            transport.EnqueueResponse("OKAY");
+            transport.EnqueueResponse("DATA00000004");
+            transport.EnqueueResponse("OKAY");
+            var util = new FastbootDriver(transport);
+
+            using var stream = new MemoryStream(new byte[] { 0x01, 0x02, 0x03, 0x04 });
+            var response = util.DownloadData(stream, 4);
+
+            Assert.Equal(FastbootState.Success, response.Result);
+            Assert.True(transport.Commands.Count(c => c == "download:00000004") >= 2);
+        }
+
         [Fact]
         public void DownloadDataStream_NonSeekable_DoesNotRetry()
         {
@@ -617,6 +707,107 @@ public void FlashUnsparseImage_LogicalPartition_ResizesBeforeFlash()
             Assert.Contains("flash:system_b", transport.Commands);
         }
 
+        [Fact]
+        public void FlashUnsparseImage_JustUnderLimit_UsesRaw()
+        {
+            // size one byte less than max-download-size (0x40 from protocol stub)
+            byte[] imageBytes = new byte[0x3F];
+            var transport = new ProtocolDownloadCaptureTransport(new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase)
+            {
+                ["getvar:has-slot:boot"] = "OKAYno",
+                ["getvar:is-logical:boot"] = "OKAYno",
+                ["getvar:max-download-size"] = "OKAY0x40",
+                ["flash:boot"] = "OKAY"
+            });
+            var util = new FastbootDriver(transport);
+            using var stream = new MemoryStream(imageBytes);
+
+            var response = util.FlashUnsparseImage("boot", stream, stream.Length);
+            Assert.Equal(FastbootState.Success, response.Result);
+
+            // ensure downloaded payload matches original bytes and not a sparse header
+            var payload = transport.DownloadPayload.ToArray();
+            Assert.Equal(imageBytes.Length, payload.Length);
+            Assert.NotEqual(0xed, payload[0]); // first byte of sparse magic is 0x3a? actually magic is 0xed26ff3a; just ensure not starting with ED
+        }
+
+        [Fact]
+        public void FlashUnsparseImage_ExactLimit_CompressesIfNeeded()
+        {
+            // image equal to limit - will trigger sparse conversion because size > max
+            byte[] imageBytes = new byte[0x41];
+            // put some zeros to make sparse compressible
+            for (int i = 1; i < imageBytes.Length; i++) imageBytes[i] = 0;
+
+            var transport = new ProtocolDownloadCaptureTransport(new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase)
+            {
+                ["getvar:has-slot:boot"] = "OKAYno",
+                ["getvar:is-logical:boot"] = "OKAYno",
+                ["getvar:max-download-size"] = "OKAY0x40",
+                ["flash:boot"] = "OKAY"
+            });
+            var util = new FastbootDriver(transport);
+            using var stream = new MemoryStream(imageBytes);
+
+            var response = util.FlashUnsparseImage("boot", stream, stream.Length);
+            Assert.Equal(FastbootState.Success, response.Result);
+
+            var payload = transport.DownloadPayload.ToArray();
+            // sparse images always start with the 32‑bit magic value 0xED26FF3A
+            // the header is written little‑endian so first four bytes should be 3A FF 26 ED
+            Assert.True(payload.Length >= 4);
+            Assert.Equal(0x3a, payload[0]);
+            Assert.Equal(0xff, payload[1]);
+            Assert.Equal(0x26, payload[2]);
+            Assert.Equal(0xed, payload[3]);
+        }
+
+        [Fact]
+        public void NotifyCurrentStep_DefaultWritesToConsole()
+        {
+            var util = new FastbootDriver(new ProtocolDownloadCaptureTransport());
+            using var err = new StringWriter();
+            var orig = Console.Error;
+            Console.SetError(err);
+            util.NotifyCurrentStep("example step");
+            Console.SetError(orig);
+
+            string output = err.ToString();
+            Assert.Contains("example step", output);
+        }
+
+        [Fact]
+        public void NotifyProgress_DefaultWritesToConsole()
+        {
+            var util = new FastbootDriver(new ProtocolDownloadCaptureTransport());
+            using var err = new StringWriter();
+            var orig = Console.Error;
+            Console.SetError(err);
+            util.NotifyProgress(50, 200);
+            Console.SetError(orig);
+
+            string output = err.ToString();
+            Assert.Contains("50/200", output);
+            Assert.Contains("25%", output);
+        }
+
+        [Fact]
+        public void Progress_WithHandler_DoesNotDoublePrint()
+        {
+            var util = new FastbootDriver(new ProtocolDownloadCaptureTransport());
+            bool called = false;
+            util.DataTransferProgressChanged += (s, e) => called = true;
+
+            using var err = new StringWriter();
+            var orig = Console.Error;
+            Console.SetError(err);
+            util.NotifyProgress(1, 1);
+            Console.SetError(orig);
+
+            Assert.True(called);
+            Assert.Equal(string.Empty, err.ToString());
+        }
+
         [Fact]
         public void FlashSparseFile_TinyLimit_FallsBackToSingleSparseTransfer()
         {
@@ -635,5 +826,34 @@ public void FlashSparseFile_TinyLimit_FallsBackToSingleSparseTransfer()
             Assert.Contains(transport.Commands, c => c.StartsWith("download:", StringComparison.OrdinalIgnoreCase));
             Assert.Contains("flash:boot", transport.Commands);
         }
+
+        [Fact]
+        public void FlashSparseFile_Multipart_PreservesOriginalTotalBlocksInEachSparseHeader()
+        {
+            var transport = new MultiPartDownloadCaptureTransport(new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase)
+            {
+                ["flash:system_a"] = "OKAY"
+            });
+            var util = new FastbootDriver(transport);
+
+            using var sparse = new SparseFile(4096, 3 * 1024 * 1024);
+            sparse.AddRawChunk(new byte[1536 * 1024]);
+            sparse.AddDontCareChunk(512 * 1024);
+            sparse.AddRawChunk(new byte[1024 * 1024]);
+
+            uint originalTotalBlocks = sparse.Header.TotalBlocks;
+
+            var response = util.FlashSparseFile("system_a", sparse, 1024 * 1024);
+
+            Assert.Equal(FastbootState.Success, response.Result);
+            Assert.True(transport.DownloadPayloads.Count > 1);
+
+            foreach (var payload in transport.DownloadPayloads)
+            {
+                Assert.True(payload.Length >= SparseFormat.SparseHeaderSize);
+                var header = SparseHeader.FromBytes(payload.AsSpan(0, SparseFormat.SparseHeaderSize));
+                Assert.Equal(originalTotalBlocks, header.TotalBlocks);
+            }
+        }
     }
 }

FirmwareKit.Comm.Fastboot.Tests/SparseFileResparseTests.cs

@@ -0,0 +1,95 @@
+using FirmwareKit.Sparse.Core;
+using FirmwareKit.Sparse.Models;
+
+namespace FirmwareKit.Comm.Fastboot.Tests
+{
+    public class SparseFileResparseTests
+    {
+        [Fact]
+        public async Task ResparseAndMerge_ShouldProduceIdenticalSparseImage()
+        {
+            // Arrange: 创建一个临时稀疏文件
+            var tempDir = Path.Combine(Path.GetTempPath(), Guid.NewGuid().ToString());
+            Directory.CreateDirectory(tempDir);
+            var originalPath = Path.Combine(tempDir, "original_sparse.img");
+            var mergedPath = Path.Combine(tempDir, "merged_sparse.img");
+            var blockSize = 4096u;
+            var totalSize = 10 * 1024 * 1024; // 10MB
+            using (var sparse = new SparseFile(blockSize, totalSize))
+            {
+                // 填充一些数据块
+                sparse.AddFillChunk(0xAABBCCDD, 2 * 1024 * 1024); // 2MB填充
+                sparse.AddDontCareChunk(1 * 1024 * 1024); // 1MB空洞
+                sparse.AddRawChunk(new byte[blockSize * 10]); // 10块原始数据
+                using (var fs = File.Create(originalPath))
+                {
+                    sparse.WriteToStream(fs, sparse: true);
+                }
+            }
+
+            // Act: 读取稀疏文件并重稀疏化再合并
+            using var original = SparseFile.FromImageFile(originalPath);
+            var parts = original.Resparse(2 * 1024 * 1024).ToArray(); // 2MB分块
+
+            // 合并所有分块为一个新的SparseFile对象，使用公开API重建chunk
+            var mergedSparse = new SparseFile(parts[0].Header.BlockSize, original.Header.TotalBlocks * parts[0].Header.BlockSize);
+            foreach (var part in parts)
+            {
+                foreach (var chunk in part.Chunks)
+                {
+                    switch ((ChunkType)chunk.Header.ChunkType)
+                    {
+                        case ChunkType.Raw:
+                            if (chunk.DataProvider != null)
+                            {
+                                var buffer = new byte[chunk.Header.ChunkSize * mergedSparse.Header.BlockSize];
+                                chunk.DataProvider.Read(0, buffer, 0, buffer.Length);
+                                mergedSparse.AddRawChunk(buffer);
+                            }
+                            break;
+                        case ChunkType.Fill:
+                            mergedSparse.AddFillChunk(chunk.FillValue, chunk.Header.ChunkSize * mergedSparse.Header.BlockSize);
+                            break;
+                        case ChunkType.DontCare:
+                            mergedSparse.AddDontCareChunk(chunk.Header.ChunkSize * mergedSparse.Header.BlockSize);
+                            break;
+                    }
+                }
+                part.Dispose();
+            }
+            // 合并后补齐DONT_CARE块，确保总块数一致
+            if (mergedSparse.CurrentBlock < mergedSparse.Header.TotalBlocks)
+            {
+                var remainBlocks = mergedSparse.Header.TotalBlocks - mergedSparse.CurrentBlock;
+                mergedSparse.AddDontCareChunk(remainBlocks * mergedSparse.Header.BlockSize);
+            }
+            using (var fs = File.Create(mergedPath))
+            {
+                mergedSparse.WriteToStream(fs, sparse: true);
+            }
+            mergedSparse.Dispose();
+
+            // Assert: 解包为raw image后比对内容
+            var originalRaw = Path.Combine(tempDir, "original.raw");
+            var mergedRaw = Path.Combine(tempDir, "merged.raw");
+            using (var origSparse = SparseFile.FromImageFile(originalPath))
+            using (var origRaw = File.Create(originalRaw))
+                origSparse.WriteRawToStream(origRaw);
+
+            using (var mergedSparseFile = SparseFile.FromImageFile(mergedPath))
+            using (var mergedRawStream = File.Create(mergedRaw))
+                mergedSparseFile.WriteRawToStream(mergedRawStream);
+
+            var originalRawBytes = await File.ReadAllBytesAsync(originalRaw);
+            var mergedRawBytes = await File.ReadAllBytesAsync(mergedRaw);
+            Assert.Equal(originalRawBytes, mergedRawBytes);
+
+            // 清理
+            File.Delete(originalPath);
+            File.Delete(mergedPath);
+            File.Delete(originalRaw);
+            File.Delete(mergedRaw);
+            Directory.Delete(tempDir);
+        }
+    }
+}

FirmwareKit.Comm.Fastboot/Command/Flash.cs

@@ -42,13 +42,17 @@ static bool IsSparseHeader(Stream s)
             ResizeLogicalPartition(partition, imageSize);
         }
 
+        // log sizes so that callers can understand why we convert to sparse
+        FastbootDebug.Log($"FlashUnsparseImage: imageSize={imageSize}, maxDownloadSize={maxDownloadSize}, isSparse={isSparse}");
+
         if (isSparse)
         {
             NotifyCurrentStep($"Flashing sparse image to {partition}...");
             using var sparseImage = SparseFile.ImportAuto(stream, validateCrc: false, verbose: false);
             return FlashSparseFile(partition, sparseImage, maxDownloadSize);
         }
 
+        // Only send raw if the image will fit in a single transfer
         if (imageSize <= maxDownloadSize)
         {
             if (canSeek) stream.Seek(originalPosition, SeekOrigin.Begin);

FirmwareKit.Comm.Fastboot/Command/InternalDownloadDataBytes.cs

@@ -22,7 +22,14 @@ public FastbootResponse DownloadData(byte[] data)
 
         FastbootResponse response = RawCommand("download:" + data.Length.ToString("x8"));
         if (response.Result != FastbootState.Data)
-            return response;
+        {
+            return new FastbootResponse
+            {
+                Result = FastbootState.Fail,
+                Response = "protocol error: expected DATA response to download command, got " + response.Result +
+                    (string.IsNullOrEmpty(response.Response) ? string.Empty : " (" + response.Response + ")")
+            };
+        }
         if (response.DataSize != data.Length)
         {
             return new FastbootResponse

FirmwareKit.Comm.Fastboot/Command/InternalDownloadDataStream.cs

@@ -79,7 +79,10 @@ private FastbootResponse DownloadDataInternal(Stream stream, long length, bool o
         // AOSP uses %08" PRIx32 which is 8 chars hex with leading zeros
         FastbootResponse response = RawCommand("download:" + length.ToString("x8"));
         if (response.Result != FastbootState.Data)
-            return response;
+        {
+            throw new Exception("protocol error: expected DATA response to download command, got " + response.Result +
+                (string.IsNullOrEmpty(response.Response) ? string.Empty : " (" + response.Response + ")"));
+        }
         if (response.DataSize != length)
         {
             return new FastbootResponse