Sparse.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;
using System.IO;

namespace FastbootCS
{
    public class Sparse : IDisposable
    {
        public const uint SPARSE_HEADER_MAGIC = 0xED26FF3A;
        public const ushort CHUNK_TYPE_RAW = 0xCAC1;
        public const ushort CHUNK_TYPE_FILL = 0xCAC2;
        public const ushort CHUNK_TYPE_DONT_CARE = 0xCAC3;
        public const ushort CHUNK_TYPE_CRC32 = 0xCAC4;
        public SparseHeader Header { get; private set; }
        public List<SparseChunk> chunks { get; private set; }
        private BinaryReader sparsebr { get; set; }
        private bool disposed;

        public Sparse(BinaryReader br)
        {
            sparsebr = br;
            chunks = new List<SparseChunk>();
            SparseHeader header = SparseReader.ReadStruct<SparseHeader>(sparsebr);
            if (header.Magic != SPARSE_HEADER_MAGIC)
            {
                throw new InvalidDataException("Not a valid sparse image!");
            }
            Header = header;
            // 读取所有chunks的头并记录数据在文件中的位置（按需加载数据以节省内存）
            for (int i = 0; i < Header.TotalChunks; i++)
            {
                long chunkHeaderPos = sparsebr.BaseStream.Position;
                var chunkHeader = SparseReader.ReadStruct<ChunkHeader>(sparsebr);

                long dataStartPos = sparsebr.BaseStream.Position;
                int dataSize = (int)(chunkHeader.TotalSize - Marshal.SizeOf<ChunkHeader>());

                // 跳过数据区域但不读取到内存
                if (dataSize > 0)
                {
                    // Ensure the stream can seek
                    if (!sparsebr.BaseStream.CanSeek)
                    {
                        // If not seekable, we must read and discard
                        sparsebr.ReadBytes(dataSize);
                    }
                    else
                    {
                        sparsebr.BaseStream.Seek(dataSize, SeekOrigin.Current);
                    }
                }

                SparseChunk chunk = new SparseChunk
                {
                    Header = chunkHeader,
                    StartPosition = chunkHeaderPos,
                    DataStartPosition = dataStartPos,
                    DataSize = dataSize,
                    Data = null // 按需加载
                };

                chunks.Add(chunk);
            }
        }
        public static bool IsSparseImage(Stream fileStream)
        {
            using (var br = new BinaryReader(fileStream, Encoding.Default, leaveOpen: true))
            {
                uint magic = br.ReadUInt32();
                fileStream.Position = 0; // 重置流位置
                if (magic == SPARSE_HEADER_MAGIC)
                {
                    return true;
                }
                else
                {
                    return false;
                }
            }
            
        }
        public static bool IsSparseImage(string filepath)
        {
            using var fs = File.OpenRead(filepath);
            return IsSparseImage(fs);

        }

        /// <summary>
        /// 将稀疏镜像分割为多个片段。每个片段会通过 <paramref name="createOutputStream"/> 创建一个可写的 Stream。
        /// 写入完毕后会将流 Position 重置为 0，然后调用 <paramref name="onSegmentReady"/>(index, stream) 由上层进行处理（例如上传并刷写）。
        /// onSegmentReady 返回后本方法会关闭并释放该 Stream。
        /// </summary>
        /// <param name="maxSize">单个片段的最大字节数（包括 sparse header 和 chunk headers + chunk data）。</param>
        /// <param name="createOutputStream">回调，传入当前片段索引，返回一个可写的 Stream 来接收片段数据。</param>
        /// <param name="onSegmentReady">当片段写入完成后被调用（stream.Position 已被重置为 0），返回 Task。</param>
        public async Task SplitSparseAsync(int maxSize, Func<int, Stream> createOutputStream, Func<int, Stream, Task> onSegmentReady)
        {
            if (createOutputStream == null) throw new ArgumentNullException(nameof(createOutputStream));
            if (onSegmentReady == null) throw new ArgumentNullException(nameof(onSegmentReady));

            // 第二步：分析块的逻辑布局，计算每个块的起始块位置
            List<ChunkLayoutInfo> layout = AnalyzeChunkLayout();

            // 第三步：分割chunks到多个输出流
            int fileIndex = 0;
            int currentChunkIndex = 0;

            while (currentChunkIndex < chunks.Count)
            {
                using (var outStream = createOutputStream(fileIndex))
                {
                    if (outStream == null) throw new InvalidOperationException("createOutputStream returned null.");
                    if (!outStream.CanWrite) throw new InvalidOperationException("createOutputStream must return a writable stream.");

                    using (var bw = new BinaryWriter(outStream, Encoding.UTF8, leaveOpen: true))
                    {
                        List<SparseChunk> currentFileChunks = new List<SparseChunk>();
                        int currentFileSize = Marshal.SizeOf<SparseHeader>(); // 文件头大小
                        uint lastBlockInFile = 0; // 跟踪文件中最后一个块的结束位置

                        // 收集当前文件的chunks
                        for (int i = currentChunkIndex; i < chunks.Count; i++)
                        {
                            var chunk = chunks[i];
                            var chunkLayout = layout[i];
                            int chunkSize = Marshal.SizeOf<ChunkHeader>() + chunk.DataSize;

                            // 检查CRC chunk
                            if (chunk.Header.ChunkType == CHUNK_TYPE_CRC32)
                            {
                                // CRC chunk 标记，但仍能写入
                            }

                            // 更准确的间隙检测
                            if (chunkLayout.StartBlock > lastBlockInFile)
                            {
                                uint gapBlocks = chunkLayout.StartBlock - lastBlockInFile;
                                var skipChunk = CreateSkipChunk(gapBlocks);
                                int skipChunkSize = Marshal.SizeOf<ChunkHeader>() + skipChunk.DataSize;

                                if (currentFileSize + skipChunkSize + chunkSize <= maxSize)
                                {
                                    currentFileChunks.Add(skipChunk);
                                    currentFileSize += skipChunkSize;
                                    lastBlockInFile = chunkLayout.StartBlock;
                                }
                                else
                                {
                                    // 间隙块都放不下，结束当前片段
                                    break;
                                }
                            }

                            // 检查是否超过最大大小
                            if (currentFileSize + chunkSize > maxSize)
                            {
                                int remainingSpace = maxSize - currentFileSize;

                                // 检查剩余空间是否足够并且为 RAW chunk 时尝试分割
                                if (remainingSpace > maxSize / 8 && chunk.Header.ChunkType == CHUNK_TYPE_RAW)
                                {
                                    // 不将整个chunk加载到内存，而是创建引用原始文件偏移的分割块
                                    var splitResult = SplitRawChunk(chunk, remainingSpace - Marshal.SizeOf<ChunkHeader>(), Header.BlockSize);

                                    if (splitResult.FirstPart.Header.ChunkSize > 0)
                                    {
                                        currentFileChunks.Add(splitResult.FirstPart);
                                        currentFileSize += Marshal.SizeOf<ChunkHeader>() + splitResult.FirstPart.DataSize;
                                        lastBlockInFile = chunkLayout.StartBlock + splitResult.FirstPart.Header.ChunkSize;

                                        // 更新原chunk为第二部分（仍为文件偏移引用）
                                        chunks[i] = splitResult.SecondPart;

                                        // 更新布局信息
                                        layout[i] = new ChunkLayoutInfo
                                        {
                                            StartBlock = lastBlockInFile,
                                            EndBlock = chunkLayout.EndBlock
                                        };

                                        // 分割后继续处理同一个chunk（现在是第二部分）
                                        i--; // 下次循环继续处理这个chunk
                                    }
                                    else
                                    {
                                        break;
                                    }
                                }
                                else
                                {
                                    // 当前chunk放在下一个片段
                                    break;
                                }
                            }
                            else
                            {
                                currentFileChunks.Add(chunk);
                                currentFileSize += chunkSize;
                                lastBlockInFile = chunkLayout.EndBlock;
                                currentChunkIndex = i + 1; // 移动到下一个chunk
                            }

                            // 如果添加了CRC chunk，这是最后一个chunk
                            if (chunk.Header.ChunkType == CHUNK_TYPE_CRC32)
                            {
                                break;
                            }
                        }

                        // 添加片段末尾的 DONT_CARE 块（如果需要）
                        uint totalBlocksInFile = lastBlockInFile;
                        uint originalTotalBlocks = Header.TotalBlocks;

                        if (totalBlocksInFile < originalTotalBlocks)
                        {
                            uint endGapBlocks = originalTotalBlocks - totalBlocksInFile;
                            var endSkipChunk = CreateSkipChunk(endGapBlocks);
                            int endSkipChunkSize = Marshal.SizeOf<ChunkHeader>() + endSkipChunk.DataSize;

                            if (currentFileSize + endSkipChunkSize <= maxSize)
                            {
                                currentFileChunks.Add(endSkipChunk);
                                currentFileSize += endSkipChunkSize;
                            }
                            else
                            {
                                throw new InvalidOperationException($"Unable to add gap block at end of segment, current segment size has reached maximum limit!");
                            }
                        }

                        // 写入当前片段的 sparse 头
                        var newHeader = CreateNewHeader((uint)currentFileChunks.Count);
                        WriteStruct(bw, newHeader);

                        // 写入所有 chunks 到提供的 Stream
                        foreach (var chunk in currentFileChunks)
                        {
                            WriteStruct(bw, chunk.Header);
                            if (chunk.DataSize > 0)
                            {
                                if (chunk.Data != null)
                                {
                                    // chunk.Data 已在内存中，直接写入（同步写）
                                    bw.Write(chunk.Data);
                                }
                                else
                                {
                                    // 按需从原始文件流异步拷贝数据到输出流，避免把所有数据加载到内存
                                    if (!sparsebr.BaseStream.CanSeek) throw new NotSupportedException("Underlying stream must be seekable for on-demand copy.");

                                    long originalPos = sparsebr.BaseStream.Position;
                                    try
                                    {
                                        sparsebr.BaseStream.Seek(chunk.DataStartPosition, SeekOrigin.Begin);

                                        const int bufferSize = 81920;
                                        byte[] buffer = new byte[bufferSize];
                                        int bytesRemaining = chunk.DataSize;
                                        while (bytesRemaining > 0)
                                        {
                                            int toRead = Math.Min(bufferSize, bytesRemaining);
                                            int read = await sparsebr.BaseStream.ReadAsync(buffer, 0, toRead);
                                            if (read <= 0) throw new EndOfStreamException("Unexpected end of stream while copying chunk data.");
                                            await outStream.WriteAsync(buffer, 0, read);
                                            bytesRemaining -= read;
                                        }
                                    }
                                    finally
                                    {
                                        if (originalPos >= 0)
                                        {
                                            sparsebr.BaseStream.Seek(originalPos, SeekOrigin.Begin);
                                        }
                                    }
                                }
                            }
                        }

                        // Ensure data flushed if stream supports flushing
                        try
                        {
                            outStream.Flush();
                        }
                        catch { }
                    } // bw disposed but leaveOpen:true 保证 outStream 仍可用

                    // 将流位置复位到开始，交给上层回调处理（如上传并刷写）
                    try
                    {
                        if (outStream.CanSeek) outStream.Position = 0;
                    }
                    catch { }

                    // 调用上层回调处理该 segment（上层应在回调返回后认为流已处理完毕）
                    await onSegmentReady(fileIndex, outStream);

                    // outStream 在 using 块结束时会被关闭/释放
                } // outStream disposed here

                fileIndex++;
            }
        }

        internal List<ChunkLayoutInfo> AnalyzeChunkLayout()
        {
            List<ChunkLayoutInfo> layout = new List<ChunkLayoutInfo>();
            uint currentBlock = 0;

            foreach (var chunk in chunks)
            {
                uint startBlock = currentBlock;
                uint endBlock = startBlock + chunk.Header.ChunkSize;

                // 对于非DONT_CARE块，需要检查是否有间隙
                if (chunk.Header.ChunkType != CHUNK_TYPE_DONT_CARE)
                {
                    startBlock = currentBlock;
                    endBlock = startBlock + chunk.Header.ChunkSize;
                    currentBlock = endBlock;
                }
                else
                {
                    // DONT_CARE块本身表示间隙
                    startBlock = currentBlock;
                    endBlock = startBlock + chunk.Header.ChunkSize;
                    currentBlock = endBlock;
                }

                layout.Add(new ChunkLayoutInfo
                {
                    StartBlock = startBlock,
                    EndBlock = endBlock
                });
            }

            return layout;
        }

        internal SparseChunk CreateSkipChunk(uint blockCount)
        {
            return new SparseChunk
            {
                Header = new ChunkHeader
                {
                    ChunkType = CHUNK_TYPE_DONT_CARE,
                    Reserved1 = 0,
                    ChunkSize = blockCount,
                    TotalSize = (uint)Marshal.SizeOf<ChunkHeader>() // DONT_CARE 块没有数据
                },
                DataSize = 0,
                Data = Array.Empty<byte>(),
                DataStartPosition = -1
            };
        }

        internal (SparseChunk FirstPart, SparseChunk SecondPart) SplitRawChunk(SparseChunk original, int maxDataSize, uint blockSize)
        {
            // 确保分割大小是块大小的倍数
            int alignedMaxSize = ((maxDataSize / (int)blockSize)) * (int)blockSize;

            if (alignedMaxSize <= 0)
            {
                alignedMaxSize = (int)blockSize; // 至少一个块
            }

            // 不再将整个原始数据加载到内存，而是通过调整 DataStartPosition 和 DataSize 来引用文件中的数据区域

            int originalDataSize = original.DataSize;

            // 第一部分
            var firstHeader = original.Header;
            firstHeader.ChunkSize = (uint)(alignedMaxSize / blockSize);
            firstHeader.TotalSize = (uint)(Marshal.SizeOf<ChunkHeader>() + alignedMaxSize);

            var firstChunk = new SparseChunk
            {
                Header = firstHeader,
                DataSize = Math.Min(alignedMaxSize, originalDataSize),
                Data = null, // 不在内存中
                DataStartPosition = original.DataStartPosition,
                StartPosition = original.StartPosition
            };

            // 第二部分
            int remainingSize = originalDataSize - firstChunk.DataSize;
            if (remainingSize < 0) remainingSize = 0;

            var secondHeader = original.Header;
            secondHeader.ChunkSize = (uint)(remainingSize / blockSize);
            secondHeader.TotalSize = (uint)(Marshal.SizeOf<ChunkHeader>() + remainingSize);

            var secondChunk = new SparseChunk
            {
                Header = secondHeader,
                DataSize = remainingSize,
                Data = null,
                DataStartPosition = (remainingSize > 0) ? original.DataStartPosition + firstChunk.DataSize : -1,
                StartPosition = original.StartPosition // note: header pos for split chunk is synthetic
            };

            return (firstChunk, secondChunk);
        }

        internal SparseHeader CreateNewHeader(uint chunkcount)
        {
            return new SparseHeader
            {
                Magic = Header.Magic,
                MajorVersion = Header.MajorVersion,
                MinorVersion = Header.MinorVersion,
                FileHeaderSize = Header.FileHeaderSize,
                ChunkHeaderSize = Header.ChunkHeaderSize,
                BlockSize = Header.BlockSize,
                TotalBlocks = Header.TotalBlocks,  // 保持与原始文件一致
                TotalChunks = chunkcount,          // 只重新计算chunk数量
                ImageChecksum = Header.ImageChecksum
            };
        }

        public static void WriteStruct<T>(BinaryWriter writer, T structure) where T : unmanaged
        {
            int size = Marshal.SizeOf<T>();
            byte[] buffer = new byte[size];
            GCHandle handle = GCHandle.Alloc(buffer, GCHandleType.Pinned);
            try
            {
                Marshal.StructureToPtr(structure, handle.AddrOfPinnedObject(), false);
                writer.Write(buffer);
            }
            finally
            {
                handle.Free();
            }
        }


        /// <summary>
        /// 按需加载chunk的数据到内存（如果尚未加载）
        /// </summary>
        /// <param name="chunk"></param>
        internal void EnsureChunkDataLoaded(SparseChunk chunk)
        {
            if (chunk.Data != null) return;
            if (chunk.DataSize <= 0)
            {
                chunk.Data = Array.Empty<byte>();
                return;
            }

            if (!sparsebr.BaseStream.CanSeek)
            {
                throw new NotSupportedException("Underlying stream is not seekable; cannot load chunk data on demand.");
            }

            long originalPos = sparsebr.BaseStream.Position;
            try
            {
                sparsebr.BaseStream.Seek(chunk.DataStartPosition, SeekOrigin.Begin);
                byte[] buf = new byte[chunk.DataSize];
                int read = 0;
                while (read < buf.Length)
                {
                    int r = sparsebr.BaseStream.Read(buf, read, buf.Length - read);
                    if (r <= 0) throw new EndOfStreamException("Unexpected end of stream while reading chunk data.");
                    read += r;
                }
                chunk.Data = buf;
            }
            finally
            {
                sparsebr.BaseStream.Seek(originalPos, SeekOrigin.Begin);
            }
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (disposed) return;
            if (disposing)
            {
                sparsebr?.Dispose();
                sparsebr = null;
                // 清理已加载的数据以释放内存
                if (chunks != null)
                {
                    foreach (var c in chunks)
                    {
                        c.Data = null;
                    }
                    chunks.Clear();
                }
            }

            disposed = true;
        }
    }

    public class ChunkLayoutInfo
    {
        public uint StartBlock { get; set; }
        public uint EndBlock { get; set; }
    }

    public class SparseChunk
    {
        public ChunkHeader Header { get; set; }
        public long StartPosition { get; set; }
        public long DataStartPosition { get; set; }
        public int DataSize { get; set; }
        public byte[] Data { get; set; }
    }
    // 对齐方式与C保持一致
    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    public struct SparseHeader
    {
        public uint Magic;             /* 0xed26ff3a */
        public ushort MajorVersion;    /* (0x1) - reject images with higher major versions */
        public ushort MinorVersion;    /* (0x0) - allow images with higer minor versions */
        public ushort FileHeaderSize;  /* 28 bytes for first revision of the file format */
        public ushort ChunkHeaderSize; /* 12 bytes for first revision of the file format */
        public uint BlockSize;         /* block size in bytes, must be a multiple of 4 (4096) */
        public uint TotalBlocks;       /* total blocks in the non-sparse output image */
        public uint TotalChunks;       /* total chunks in the sparse input image */
        public uint ImageChecksum;     /* CRC32 checksum of the original data, counting "don't care" */
    }

    [StructLayout(LayoutKind.Sequential, Pack = 1)]
    public struct ChunkHeader
    {
        public ushort ChunkType;  // 0xCAC1 -> raw; 0xCAC2 -> fill; 0xCAC3 -> don't care
        public ushort Reserved1;
        public uint ChunkSize;    /* in blocks in output image */
        public uint TotalSize;    /* in bytes of chunk input file including chunk header and data */
    }
    public static class SparseReader
    {
        public static T ReadStruct<T>(BinaryReader reader) where T : unmanaged
        {
            int size = Unsafe.SizeOf<T>();
            Span<byte> buffer = stackalloc byte[size];
            int read = reader.Read(buffer);
            if (read != size) throw new EndOfStreamException($"Unable to read structure of size {size} bytes.");
            return MemoryMarshal.Read<T>(buffer);
        }
    }
}