AyCode.Core/AyCode.Services/SignalRs/AyCodeBinaryHubProtocol.cs

using System;
using System.Buffers;
using AyCode.Core.Serializers.Binaries;
using Microsoft.AspNetCore.SignalR.Protocol;
using Microsoft.Extensions.Logging;

namespace AyCode.Services.SignalRs;

/// <summary>
/// Project-specific binary protocol.
///
/// Adds a per-message wire header (via <see cref="WriteHeader"/>/<see cref="ReadHeader"/> hooks)
/// that expresses how the data argument (args[^1] by convention) should be handled by the client:
/// <list type="bullet">
///   <item><c>HasData</c>: the data arg is not null.</item>
///   <item><c>Streamed</c>: the data arg comes via CHUNK_DATA chunks (inline placeholder on the wire).</item>
///   <item><c>ConsumerDeserialize</c>: the client returns raw <c>byte[]</c> to the consumer (IsRawBytesData flow — e.g. DataSource PopulateMerge).</item>
///   <item><c>HasType</c>: the concrete AQN of the data arg follows (for typed deserialization).</item>
/// </list>
///
/// With this header the client no longer needs to inspect <c>SignalParams</c> in order to decide
/// how to treat the data arg — so it works even when <c>SignalParams</c> is itself streamed
/// (which happens when args[^1] is <c>byte[]</c> or <c>null</c>).
/// </summary>
public class AyCodeBinaryHubProtocol : AcBinaryHubProtocol
{
    public AyCodeBinaryHubProtocol() : this(AcBinarySerializerOptions.Default) { }
    public AyCodeBinaryHubProtocol(AcBinarySerializerOptions options, BinaryProtocolMode protocolMode = BinaryProtocolMode.Bytes, ILogger? logger = null) : base(options, protocolMode, logger) { }

    #region Wire header (per-message)

    [Flags]
    private enum DataFlags : byte
    {
        None                = 0,
        HasData             = 1 << 0,  // the data arg is not null
        Streamed            = 1 << 1,  // the data arg is delivered via CHUNK_DATA chunks
        ConsumerDeserialize = 1 << 2,  // client returns raw byte[] to the consumer (IsRawBytesData flow)
        HasType             = 1 << 3,  // a type AQN string follows (for typed deserialization)
    }

    /// <summary>
    /// Opaque context produced by <see cref="ReadHeader"/> and stashed in
    /// <see cref="AcBinaryHubProtocol._currentHeaderContext"/>. Consumed by
    /// <see cref="ReadSingleArgument"/> and <see cref="ResolveStreamedArgType"/>.
    /// </summary>
    private sealed class HeaderContext
    {
        public DataFlags Flags { get; }
        public Type? Type { get; }
        public HeaderContext(DataFlags flags, Type? type) { Flags = flags; Type = type; }
    }

    /// <summary>
    /// Writes the per-message header. See <see cref="DataFlags"/> for the semantics of each bit.
    /// </summary>
    protected override void WriteHeader(ref BufferWriterBinaryOutput bw, HubMessage message, object? streamedArg)
    {
        var dataArg = GetDataArg(message);
        var flags = DataFlags.None;
        string? typeName = null;

        if (dataArg != null)
        {
            flags |= DataFlags.HasData;

            // Streamed: chunked mode active AND streamedArg is the data arg
            if (streamedArg != null && ReferenceEquals(streamedArg, dataArg))
                flags |= DataFlags.Streamed;

            // ConsumerDeserialize: the client requested byte[] via IsRawBytesData
            // and the data arg is (or was pre-serialized to) byte[].
            // Detected via the sibling SignalParams arg (project convention).
            if (dataArg is byte[] && GetClientIsRawBytesData(message))
                flags |= DataFlags.ConsumerDeserialize;

            // HasType: for typed args we write the concrete AQN so the client can deserialize
            // directly (either inline or after streamed chunks are gathered).
            // byte[] needs no type info — either base byte[] fast-path or ConsumerDeserialize handles it.
            if (dataArg is not byte[])
            {
                flags |= DataFlags.HasType;
                typeName = dataArg.GetType().AssemblyQualifiedName;
            }
        }

        bw.WriteByte((byte)flags);
        if ((flags & DataFlags.HasType) != 0)
            WriteNullableString(ref bw, typeName);
    }

    /// <summary>
    /// Reads the per-message header and returns a <see cref="HeaderContext"/>.
    /// </summary>
    protected override object? ReadHeader(ref SequenceReader<byte> r)
    {
        r.TryRead(out byte flagsByte);
        var flags = (DataFlags)flagsByte;

        Type? resolvedType = null;
        if ((flags & DataFlags.HasType) != 0)
        {
            var typeName = ReadNullableString(ref r);
            if (typeName != null)
                resolvedType = Type.GetType(typeName);
        }

        return new HeaderContext(flags, resolvedType);
    }

    /// <summary>
    /// The data arg by project convention — the last argument for Invocation messages,
    /// <see cref="StreamItemMessage.Item"/> for stream items, <see cref="CompletionMessage.Result"/> for completions.
    /// </summary>
    private static object? GetDataArg(HubMessage message) => message switch
    {
        InvocationMessage m when m.Arguments.Length > 0 => m.Arguments[m.Arguments.Length - 1],
        StreamInvocationMessage m when m.Arguments.Length > 0 => m.Arguments[m.Arguments.Length - 1],
        StreamItemMessage m => m.Item,
        CompletionMessage m => m.HasResult ? m.Result : null,
        _ => null
    };

    /// <summary>
    /// Extracts <see cref="SignalParams.IsRawBytesData"/> from the message, assuming the project
    /// convention of <c>OnReceiveMessage(int, int?, SignalParams, object)</c> — arg[2] is SignalParams.
    /// Returns false for messages that don't follow this shape.
    /// </summary>
    private static bool GetClientIsRawBytesData(HubMessage message)
    {
        if (message is InvocationMessage im && im.Arguments.Length >= 3
            && im.Arguments[2] is SignalParams sp)
            return sp.IsRawBytesData;
        return false;
    }

    #endregion

    /// <summary>
    /// For the chunked streaming path: resolve the type the background Task will deserialize into.
    /// Prefers the concrete type from the wire header (set by <see cref="ReadHeader"/>) when present,
    /// otherwise falls back to the binder-provided type (base behavior).
    /// </summary>
    protected override Type ResolveStreamedArgType(Type binderType)
    {
        if (_currentHeaderContext is HeaderContext hctx && hctx.Type != null)
            return hctx.Type;
        return base.ResolveStreamedArgType(binderType);
    }

    /// <summary>
    /// Read a single argument, using the per-message header to decide how to treat <c>object</c>-typed args.
    /// Decision order:
    /// <list type="number">
    ///   <item>Base byte[] fast-path — tag <c>0x44</c> present (file/image/raw bytes from a typed <c>byte[]</c> param).</item>
    ///   <item>Header <see cref="DataFlags.ConsumerDeserialize"/> — return raw bytes (consumer handles deserialization later).</item>
    ///   <item>Header <see cref="DataFlags.HasType"/> — resolve target type from the header and deserialize.</item>
    ///   <item>Fall through to base typed deserialization against the binder-provided target type.</item>
    /// </list>
    /// </summary>
    protected override object? ReadSingleArgument(ref SequenceReader<byte> r, Type targetType)
    {
        r.TryReadLittleEndian(out int argLength);
        if (argLength == 0)
            return null;

        // AsyncSegment: streamed arg marker (INT32 -1) → placeholder for chunked deserialization
        if (argLength == -1)
            return StreamedArgPlaceholder;

        if (argLength == 1)
        {
            r.TryPeek(out byte marker);
            if (marker == 0) { r.Advance(1); return null; }
        }

        var argSlice = r.UnreadSequence.Slice(0, argLength);
        r.Advance(argLength);

        // 1. Base byte[] fast-path: [0x44 tag][raw bytes] — strip tag, return byte[]
        var argReader = new SequenceReader<byte>(argSlice);
        if (argReader.TryPeek(out byte tag) && tag == BinaryTypeCode.ByteArray)
            return SequenceToByteArray(argSlice.Slice(1));

        var hctx = _currentHeaderContext as HeaderContext;

        // 2. Header ConsumerDeserialize: no tag on wire (isAcBinary path on server),
        //    consumer wants raw byte[] — return as-is without deserialization.
        //    Applies only to the data arg (convention: targetType == typeof(object));
        //    typed args (Int32, SignalParams, etc.) are unaffected.
        if (targetType == typeof(object)
            && hctx != null && (hctx.Flags & DataFlags.ConsumerDeserialize) != 0)
            return SequenceToByteArray(argSlice);

        // 3. Type resolution: prefer concrete type from header over binder type (which is often typeof(object))
        if (targetType == typeof(object) && hctx?.Type != null)
            targetType = hctx.Type;

        // 4. Deserialize — Bytes mode linearizes, Segment/AsyncSegment uses the sequence directly
        if (_protocolMode == BinaryProtocolMode.Bytes)
        {
            var bytes = SequenceToByteArray(argSlice);
            return AcBinaryDeserializer.Deserialize(bytes, targetType, Options);
        }

        return DeserializeFromSequence(argSlice, targetType, Options);
    }
}