LengthPrefixedMessageFrameCodec.cs

이 파일의 문서화 페이지로 가기

using System.Buffers.Binary;
 
namespace Dreamine.Communication.Core.Framing;

public sealed class LengthPrefixedMessageFrameCodec : IMessageFrameCodec
{
    private const int HeaderSize = 4;

    private readonly int _maxFrameLength;

    public LengthPrefixedMessageFrameCodec()
        : this(1024 * 1024)
    {
    }

    public LengthPrefixedMessageFrameCodec(int maxFrameLength)
    {
        if (maxFrameLength <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(maxFrameLength));
        }
 
        _maxFrameLength = maxFrameLength;
    }

    public async Task WriteFrameAsync(
        Stream stream,
        byte[] payload,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(stream);
        ArgumentNullException.ThrowIfNull(payload);
 
        if (payload.Length > _maxFrameLength)
        {
            throw new InvalidDataException(
                $"Frame length exceeded. Length={payload.Length}, MaxFrameLength={_maxFrameLength}");
        }
 
        var header = new byte[HeaderSize];
        BinaryPrimitives.WriteInt32BigEndian(header, payload.Length);
 
        await stream.WriteAsync(header, cancellationToken).ConfigureAwait(false);
        await stream.WriteAsync(payload, cancellationToken).ConfigureAwait(false);
        await stream.FlushAsync(cancellationToken).ConfigureAwait(false);
    }

    public async Task<byte[]?> ReadFrameAsync(
        Stream stream,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(stream);
 
        var header = await ReadExactAsync(stream, HeaderSize, cancellationToken)
            .ConfigureAwait(false);
 
        if (header is null)
        {
            return null;
        }
 
        var length = BinaryPrimitives.ReadInt32BigEndian(header);
 
        if (length < 0 || length > _maxFrameLength)
        {
            throw new InvalidDataException(
                $"Invalid frame length. Length={length}, MaxFrameLength={_maxFrameLength}");
        }
 
        if (length == 0)
        {
            return Array.Empty<byte>();
        }
 
        return await ReadExactAsync(stream, length, cancellationToken)
            .ConfigureAwait(false);
    }

    private static async Task<byte[]?> ReadExactAsync(
        Stream stream,
        int length,
        CancellationToken cancellationToken)
    {
        var buffer = new byte[length];
        var offset = 0;
 
        while (offset < length)
        {
            var read = await stream.ReadAsync(
                    buffer.AsMemory(offset, length - offset),
                    cancellationToken)
                .ConfigureAwait(false);
 
            if (read == 0)
            {
                return null;
            }
 
            offset += read;
        }
 
        return buffer;
    }
}