SophonPatchAsset.Download.cs

using Hi3Helper.Sophon.Helper;
using Hi3Helper.Sophon.Infos;
using Hi3Helper.Sophon.Structs;
using System;
using System.Buffers;
using System.IO;
using System.Net.Http;
using System.Threading;
using System.Threading.Tasks;
// ReSharper disable IdentifierTypo
// ReSharper disable CommentTypo
// ReSharper disable AccessToModifiedClosure

using TaskExtensions = Hi3Helper.Sophon.Helper.TaskExtensions;
using ZstdStream = ZstdNet.DecompressionStream;

namespace Hi3Helper.Sophon
{
    public enum SophonPatchMethod
    {
        CopyOver,
        DownloadOver,
        Patch,
        Remove
    }

    public partial class SophonPatchAsset : SophonIdentifiableProperty
    {
        internal const int BufferSize = 64 << 10;

        public   SophonAsset       MainAssetInfo     { get; set; }
        public   SophonChunksInfo  PatchInfo         { get; set; }
        public   SophonPatchMethod PatchMethod       { get; set; }
        public   string            PatchNameSource   { get; set; }
        public   string            PatchHash         { get; set; }
        public   long              PatchOffset       { get; set; }
        public   long              PatchSize         { get; set; }
        public   long              PatchChunkLength  { get; set; }
        public   string            OriginalFilePath  { get; set; }
        public   string            OriginalFileHash  { get; set; }
        public   long              OriginalFileSize  { get; set; }
        public   string            TargetFilePath    { get; set; }
        public   string            TargetFileHash    { get; set; }
        public   long              TargetFileSize    { get; set; }

        public override string ToString() => TargetFilePath;

    #nullable enable
        public async Task<bool> DownloadPatchAsync(HttpClient                  client,
                                                   string                      inputDir,
                                                   string                      patchOutputDir,
                                                   bool                        forceVerification    = false,
                                                   Action<long>?               downloadReadDelegate = null,
                                                   SophonDownloadSpeedLimiter? downloadSpeedLimiter = null,
                                                   CancellationToken           token                = default)
        {
            // Ignore SophonPatchMethod.Remove and SophonPatchMethod.DownloadOver assets
            if (PatchMethod is SophonPatchMethod.Remove or
                               SophonPatchMethod.DownloadOver)
            {
                return true;
            }

            FileInfo patchFilePathHashedFileInfo = this.GetLegacyOrHoyoPlayPatchChunkPath(patchOutputDir);
            if (!PatchNameSource.TryGetChunkXxh64Hash(out byte[] patchHash))
            {
                patchHash = Extension.HexToBytes(PatchHash.AsSpan());
            }

            SophonChunk patchAsChunk = new()
            {
                ChunkHashDecompressed  = patchHash,
                ChunkName              = PatchNameSource,
                ChunkOffset            = 0,
                ChunkOldOffset         = 0,
                ChunkSize              = PatchSize,
                ChunkSizeDecompressed  = PatchSize,
                IsSkipHashCheckOnWrite = true
            };

            bool isPatchFileExist = patchFilePathHashedFileInfo.Exists;
            FileStream? fileStream = isPatchFileExist
                ? patchFilePathHashedFileInfo
                   .Open(new FileStreamOptions
                    {
                        // Use pre-check for file mode to prevent overhead over OS API.
                        Mode = FileMode.Open,
                        Access = FileAccess.Read,
                        Share = FileShare.Read, // Only share and open for read access to enable OS-level file caching.
                        BufferSize = PatchSize.GetFileStreamBufferSize()
                    })
                : null;

            try
            {
                bool isPatchMatched = fileStream != null && fileStream.Length == PatchSize;
                if (forceVerification && fileStream != null)
                {
                    isPatchMatched = patchHash.Length > 8
                        ? await patchAsChunk.CheckChunkMd5HashAsync(fileStream,
                                                                    true,
                                                                    token)
                        : await patchAsChunk.CheckChunkXxh64HashAsync(fileStream,
                                                                      patchHash,
                                                                      true,
                                                                      token);
                }

                if (isPatchMatched)
                {
#if DEBUG
                    this.PushLogDebug($"Skipping patch {PatchNameSource} for: {TargetFilePath}");
#endif
                    downloadReadDelegate?.Invoke(PatchSize);

                    return true;
                }

                if (fileStream != null)
                {
#if NET6_0_OR_GREATER
                    await fileStream.DisposeAsync();
#else
                    fileStream?.Dispose();
#endif
                }

                fileStream = patchFilePathHashedFileInfo.Open(new FileStreamOptions
                {
                    Mode       = FileMode.Create,
                    Access     = FileAccess.ReadWrite,
                    Share      = FileShare.ReadWrite,
                    BufferSize = PatchSize.GetFileStreamBufferSize()
                });

                await InnerWriteChunkCopyAsync(client,
                                               fileStream,
                                               patchAsChunk,
                                               PatchInfo,
                                               PatchInfo,
                                               writeInfoDelegate: null,
                                               downloadInfoDelegate: (_, y) =>
                                               {
                                                   downloadReadDelegate?.Invoke(y);
                                               },
                                               downloadSpeedLimiter: downloadSpeedLimiter,
                                               token: token);

                return true;
            }
            finally
            {
                if (fileStream != null)
                {
#if NET6_0_OR_GREATER
                    await fileStream.DisposeAsync();
#else
                    fileStream?.Dispose();
#endif
                }
            }
        }

        private async
            Task
            InnerWriteChunkCopyAsync(HttpClient                  client,
                                     Stream                      outStream,
                                     SophonChunk                 chunk,
                                     SophonChunksInfo            currentSophonChunkInfo,
                                     SophonChunksInfo            altSophonChunkInfo,
                                     DelegateWriteStreamInfo?    writeInfoDelegate,
                                     DelegateWriteDownloadInfo?  downloadInfoDelegate,
                                     SophonDownloadSpeedLimiter? downloadSpeedLimiter,
                                     CancellationToken           token)
        {
            const int retryCount = TaskExtensions.DefaultRetryAttempt;
            int currentRetry = 0;

            long currentWriteOffset = 0;

#if !NOSTREAMLOCK
            if (outStream is FileStream fs)
            {
                fs.Lock(chunk.ChunkOffset, chunk.ChunkSizeDecompressed);
                this.PushLogDebug($"Locked data stream from pos: 0x{chunk.ChunkOffset:x8}" +
                $" -> L: 0x{chunk.ChunkSizeDecompressed:x8} for chunk: {chunk.ChunkName}" +
                $" by asset: {TargetFilePath}");
            }
#endif

            while (true)
            {
                bool allowDispose = false;
                HttpResponseMessage? httpResponseMessage = null;
                Stream? httpResponseStream = null;
                Stream? sourceStream = null;

                byte[] buffer = ArrayPool<byte>.Shared.Rent(BufferSize);

                {
                    try
                    {
                        CancellationTokenSource innerTimeoutToken =
                            new(TimeSpan.FromSeconds(TaskExtensions.DefaultTimeoutSec)
                            #if NET8_0_OR_GREATER
                              , TimeProvider.System
                            #endif
                               );
                        CancellationTokenSource cooperatedToken =
                            CancellationTokenSource.CreateLinkedTokenSource(token, innerTimeoutToken.Token);

#if DEBUG
                        this.PushLogDebug($"Init. by offset: 0x{chunk.ChunkOffset:x8}" +
                            $" -> L: 0x{chunk.ChunkSizeDecompressed:x8} for chunk: {chunk.ChunkName}");

#endif
                        outStream.Position = 0;
                        httpResponseMessage = await client
                            .GetChunkAndIfAltAsync(chunk.ChunkName,
                                                   currentSophonChunkInfo,
                                                   altSophonChunkInfo,
                                                   cooperatedToken.Token);
                        httpResponseStream = await httpResponseMessage
                                                  .EnsureSuccessStatusCode()
                                                  .Content
                                                  .ReadAsStreamAsync(
#if NET6_0_OR_GREATER
                                                                     cooperatedToken.Token
#endif
                                                                    );

                        sourceStream = httpResponseStream;
#if DEBUG
                        this.PushLogDebug($"[Complete init.] by offset: 0x{chunk.ChunkOffset:x8}" +
                            $" -> L: 0x{chunk.ChunkSizeDecompressed:x8} for chunk: {chunk.ChunkName}");
#endif

                        int read;
                        while ((read = await sourceStream.ReadAsync(
#if NET6_0_OR_GREATER
                                                                    buffer
#else
                                                                    buffer, 0, buffer.Length
#endif
                                                                  , cooperatedToken.Token)
                                                         .ConfigureAwait(false)) > 0)
                        {
                            await (downloadSpeedLimiter?.AddBytesOrWaitAsync(read, token) ?? ValueTask.CompletedTask);

#if NET6_0_OR_GREATER
                            await outStream.WriteAsync(buffer.AsMemory(0, read), cooperatedToken.Token).ConfigureAwait(false);
#else
                            await outStream.WriteAsync(buffer, 0, read, cooperatedToken.Token).ConfigureAwait(false);
#endif

                            currentWriteOffset += read;
                            writeInfoDelegate?.Invoke(read);
                            downloadInfoDelegate?.Invoke(read, read);

                            currentRetry = 0;
                            innerTimeoutToken.Dispose();
                            cooperatedToken.Dispose();

                            innerTimeoutToken =
                                new CancellationTokenSource(TimeSpan.FromSeconds(TaskExtensions.DefaultTimeoutSec)
#if NET8_0_OR_GREATER
                                                          , TimeProvider.System
#endif
                                                           );
                            cooperatedToken =
                                CancellationTokenSource.CreateLinkedTokenSource(token, innerTimeoutToken.Token);
                        }

                        if (chunk.IsSkipHashCheckOnWrite)
                        {
                            return;
                        }

                        outStream.Position = 0;
                        Stream checkHashStream = outStream;

                        bool isHashVerified;
                        if (chunk.ChunkName.TryGetChunkXxh64Hash(out byte[] outHash))
                        {
                            isHashVerified =
                                await chunk.CheckChunkXxh64HashAsync(checkHashStream,
                                                                     outHash,
                                                                     true,
                                                                     cooperatedToken.Token);
                        }
                        else
                        {
                            if (PatchInfo.IsUseCompression)
                            {
                                checkHashStream = new ZstdStream(checkHashStream);
                            }

                            isHashVerified =
                                await chunk.CheckChunkMd5HashAsync(checkHashStream,
                                                                   true,
                                                                   cooperatedToken.Token);
                        }

                        if (!isHashVerified)
                        {
                            writeInfoDelegate?.Invoke(-chunk.ChunkSizeDecompressed);
                            downloadInfoDelegate?.Invoke(-chunk.ChunkSizeDecompressed, 0);
                            this.PushLogWarning("Output data seems to be corrupted at transport.\r\n" +
                                $"Restarting download for chunk: {chunk.ChunkName} | 0x{chunk.ChunkOffset:x8}" +
                                $" -> L: 0x{chunk.ChunkSizeDecompressed:x8} for: {TargetFilePath}");
                            continue;
                        }

#if DEBUG
                        this.PushLogDebug($"Download completed! Chunk: {chunk.ChunkName}" +
                            $" | 0x{chunk.ChunkOffset:x8} -> L: 0x{chunk.ChunkSizeDecompressed:x8}" +
                            $" for: {TargetFilePath}");
#endif
                        return;
                    }
                    catch (OperationCanceledException) when (token.IsCancellationRequested)
                    {
                        allowDispose = true;
                        throw;
                    }
                    catch (Exception ex)
                    {
                        if (currentRetry < retryCount)
                        {
                            writeInfoDelegate?.Invoke(-currentWriteOffset);
                            downloadInfoDelegate?.Invoke(-currentWriteOffset, 0);
                            currentWriteOffset = 0;
                            currentRetry++;

                            this.PushLogWarning("An error has occurred while" +
                                $" downloading chunk: {chunk.ChunkName}" +
                                $" | 0x{chunk.ChunkOffset:x8} -> L: 0x{chunk.ChunkSizeDecompressed:x8}" +
                                $" for: {TargetFilePath}\r\n{ex}");
                            await Task.Delay(TimeSpan.FromSeconds(1), token);
                            continue;
                        }

                        allowDispose = true;
                        this.PushLogError($"An unhandled error has occurred while downloading chunk:" +
                            $" {chunk.ChunkName} | 0x{chunk.ChunkOffset:x8} ->" +
                            $" L: 0x{chunk.ChunkSizeDecompressed:x8} for: {TargetFilePath}\r\n{ex}");
                        throw;
                    }
                    finally
                    {
                        if (allowDispose)
                        {
                            httpResponseMessage?.Dispose();
#if NET6_0_OR_GREATER
                            if (httpResponseStream != null)
                            {
                                await httpResponseStream.DisposeAsync();
                            }

                            if (sourceStream != null)
                            {
                                await sourceStream.DisposeAsync();
                            }
#else
                            sourceStream?.Dispose();
                            httpResponseStream?.Dispose();
#endif
                        }

                        ArrayPool<byte>.Shared.Return(buffer);
                    }
                }
            }
        }
    }
}