Chunker.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
 
using Microsoft.Extensions.Logging;
 
using Newtonsoft.Json;
 
namespace Tgstation.Server.Host.Components.Interop
{
    abstract class Chunker
    {
        protected ILogger<Chunker> Logger { get; }
 
        protected uint NextPayloadId
        {
            get
            {
                // 0 is special, since BYOND doesn't use it we reserve it for if we have to send a chunked topic request immediately upon reattaching
                // Otherwise, jump ahead a bunch compared to what was last used/seen, so we don't accidentally clash
                lock (chunkSets)
                    return highestSeenPayloadId == 0 ? 0 : highestSeenPayloadId + 20;
            }
        }
 
        readonly Dictionary<uint, Tuple<ChunkSetInfo, string[]>> chunkSets;
 
        uint highestSeenPayloadId;
 
        protected Chunker(ILogger<Chunker> logger)
        {
            Logger = logger ?? throw new ArgumentNullException(nameof(logger));
            chunkSets = new Dictionary<uint, Tuple<ChunkSetInfo, string[]>>();
        }
 
        protected async ValueTask<TResponse?> ProcessChunk<TCommunication, TResponse>(
            Func<TCommunication, CancellationToken, ValueTask<TResponse?>> completionCallback,
            Func<string, TResponse?> chunkErrorCallback,
            ChunkData? chunk,
            CancellationToken cancellationToken)
            where TCommunication : class
            where TResponse : IMissingPayloadsCommunication, new()
        {
            if (chunk == null)
                return chunkErrorCallback("Missing chunk!");
 
            if (!chunk.PayloadId.HasValue)
                return chunkErrorCallback("Missing chunk payloadId!");
 
            if (!chunk.SequenceId.HasValue)
                return chunkErrorCallback("Missing chunk sequenceId!");
 
            if (chunk.Payload == null)
                return chunkErrorCallback("Missing chunk payload!");
 
            ChunkSetInfo requestInfo;
            string[] payloads;
            lock (chunkSets)
            {
                highestSeenPayloadId = Math.Max(chunk.PayloadId.Value, highestSeenPayloadId);
                if (!chunkSets.TryGetValue(chunk.PayloadId.Value, out var tuple))
                {
                    // first time seeing this payload
                    if (chunk.TotalChunks == 0)
                        return chunkErrorCallback("Receieved chunked request with 0 totalChunks!");
 
                    tuple = Tuple.Create<ChunkSetInfo, string[]>(chunk, new string[chunk.TotalChunks]);
                    chunkSets.Add(chunk.PayloadId.Value, tuple);
                }
 
                requestInfo = tuple.Item1;
                payloads = tuple.Item2;
 
                Logger.LogTrace("Received chunk P{payloadId}: {sequenceId}/{totalChunks}", requestInfo.PayloadId, chunk.SequenceId + 1, requestInfo.TotalChunks);
 
                if (chunk.TotalChunks != requestInfo.TotalChunks)
                {
                    chunkSets.Remove(requestInfo.PayloadId!.Value);
                    return chunkErrorCallback("Received differing total chunks for same payloadId! Invalidating payloadId!");
                }
 
                if (payloads[chunk.SequenceId.Value] != null && payloads[chunk.SequenceId.Value] != chunk.Payload)
                {
                    chunkSets.Remove(requestInfo.PayloadId!.Value);
                    return chunkErrorCallback("Received differing payload for same sequenceId! Invalidating payloadId!");
                }
 
                payloads[chunk.SequenceId.Value] = chunk.Payload;
                var missingPayloads = new List<uint>();
                for (var i = 0U; i < payloads.Length; ++i)
                    if (payloads[i] == null)
                        missingPayloads.Add(i);
 
                if (missingPayloads.Count > 0)
                    return new TResponse
                    {
                        MissingChunks = missingPayloads,
                    };
 
                Logger.LogTrace("Received all chunks for P{payloadId}, processing request...", requestInfo.PayloadId);
                chunkSets.Remove(requestInfo.PayloadId!.Value);
            }
 
            TCommunication? completedCommunication;
            var fullCommunicationJson = String.Concat(payloads);
            try
            {
                completedCommunication = JsonConvert.DeserializeObject<TCommunication>(fullCommunicationJson, DMApiConstants.SerializerSettings);
            }
            catch (Exception ex)
            {
                Logger.LogDebug(ex, "Bad chunked communication for payload {payloadId}!", requestInfo.PayloadId);
                completedCommunication = null;
            }
 
            if (completedCommunication == null)
                return chunkErrorCallback("Chunked request completed with bad JSON!");
 
            return await completionCallback(completedCommunication, cancellationToken);
        }
    }
}