git.stella-ops.org/src/__Libraries/StellaOps.Router.Transport.Tcp/FrameProtocol.cs

using System.Buffers.Binary;
using StellaOps.Router.Common.Enums;
using StellaOps.Router.Common.Models;

namespace StellaOps.Router.Transport.Tcp;

/// <summary>
/// Handles reading and writing length-prefixed frames over a stream.
/// Frame format: [4-byte big-endian length][payload]
/// Payload format: [1-byte frame type][16-byte correlation GUID][remaining data]
/// </summary>
public static class FrameProtocol
{
    private const int LengthPrefixSize = 4;
    private const int FrameTypeSize = 1;
    private const int CorrelationIdSize = 16;
    private const int HeaderSize = FrameTypeSize + CorrelationIdSize;

    /// <summary>
    /// Reads a complete frame from the stream.
    /// </summary>
    /// <param name="stream">The stream to read from.</param>
    /// <param name="maxFrameSize">The maximum frame size allowed.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>The frame read, or null if the stream is closed.</returns>
    public static async Task<Frame?> ReadFrameAsync(
        Stream stream,
        int maxFrameSize,
        CancellationToken cancellationToken)
    {
        // Read length prefix (4 bytes, big-endian)
        var lengthBuffer = new byte[LengthPrefixSize];
        var bytesRead = await ReadExactAsync(stream, lengthBuffer, cancellationToken);
        if (bytesRead == 0)
        {
            return null; // Connection closed
        }

        if (bytesRead < LengthPrefixSize)
        {
            throw new InvalidOperationException("Incomplete length prefix received");
        }

        var payloadLength = BinaryPrimitives.ReadInt32BigEndian(lengthBuffer);
        if (payloadLength < HeaderSize)
        {
            throw new InvalidOperationException($"Invalid payload length: {payloadLength}");
        }

        if (payloadLength > maxFrameSize)
        {
            throw new InvalidOperationException(
                $"Frame size {payloadLength} exceeds maximum {maxFrameSize}");
        }

        // Read payload
        var payload = new byte[payloadLength];
        bytesRead = await ReadExactAsync(stream, payload, cancellationToken);
        if (bytesRead < payloadLength)
        {
            throw new InvalidOperationException(
                $"Incomplete payload: expected {payloadLength}, got {bytesRead}");
        }

        // Parse frame
        var frameType = (FrameType)payload[0];
        var correlationId = new Guid(payload.AsSpan(FrameTypeSize, CorrelationIdSize));
        var data = payload.AsMemory(HeaderSize);

        return new Frame
        {
            Type = frameType,
            CorrelationId = correlationId.ToString("N"),
            Payload = data
        };
    }

    /// <summary>
    /// Writes a frame to the stream.
    /// </summary>
    /// <param name="stream">The stream to write to.</param>
    /// <param name="frame">The frame to write.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    public static async Task WriteFrameAsync(
        Stream stream,
        Frame frame,
        CancellationToken cancellationToken)
    {
        // Parse or generate correlation ID
        var correlationGuid = frame.CorrelationId is not null &&
                              Guid.TryParse(frame.CorrelationId, out var parsed)
            ? parsed
            : Guid.NewGuid();

        var dataLength = frame.Payload.Length;
        var payloadLength = HeaderSize + dataLength;

        // Create buffer for the complete message
        var buffer = new byte[LengthPrefixSize + payloadLength];

        // Write length prefix (big-endian)
        BinaryPrimitives.WriteInt32BigEndian(buffer.AsSpan(0, LengthPrefixSize), payloadLength);

        // Write frame type
        buffer[LengthPrefixSize] = (byte)frame.Type;

        // Write correlation ID
        correlationGuid.TryWriteBytes(buffer.AsSpan(LengthPrefixSize + FrameTypeSize, CorrelationIdSize));

        // Write data
        if (dataLength > 0)
        {
            frame.Payload.Span.CopyTo(buffer.AsSpan(LengthPrefixSize + HeaderSize));
        }

        await stream.WriteAsync(buffer, cancellationToken);
    }

    /// <summary>
    /// Reads exactly the specified number of bytes from the stream.
    /// </summary>
    private static async Task<int> ReadExactAsync(
        Stream stream,
        Memory<byte> buffer,
        CancellationToken cancellationToken)
    {
        var totalRead = 0;
        while (totalRead < buffer.Length)
        {
            var read = await stream.ReadAsync(
                buffer[totalRead..],
                cancellationToken);

            if (read == 0)
            {
                return totalRead; // EOF
            }

            totalRead += read;
        }

        return totalRead;
    }
}