Meum/Meum.Core/QuicStreamWrapper.cs

using System.Buffers;
using System.Net.Quic;
using System.Runtime.InteropServices;
using Meum.Core.Messages;

namespace Meum.Core;

public class QuicStreamWrapper : IAsyncDisposable
{
    private QuicStream _stream;

    public QuicStreamWrapper(QuicStream stream)
    {
        _stream = stream;
    }
    
    public async ValueTask<T> ReadStructAsync<T>(CancellationToken ct = default)
        where T : struct
    {
        byte[] buffer = ArrayPool<byte>.Shared.Rent(Marshal.SizeOf<T>());
        var handle = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        try
        {
            await _stream.ReadExactlyAsync(buffer, ct);
            return (T) Marshal.PtrToStructure(handle.AddrOfPinnedObject(), typeof(T))!;
        }
        finally
        {
            handle.Free();
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }
    
    public ValueTask WriteStructAsync<T>(T msg_struct, CancellationToken ct = default)
        where T : struct
    {
        byte[] buffer = ArrayPool<byte>.Shared.Rent(Marshal.SizeOf<T>());
        var handle = GCHandle.Alloc(buffer, GCHandleType.Pinned);
        try
        {
            Marshal.StructureToPtr(msg_struct, handle.AddrOfPinnedObject(), false);
            return _stream.WriteAsync(buffer, ct);
        }
        finally
        {
            handle.Free();
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    public ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken ct = default) 
        => _stream.ReadAsync(buffer, ct);
    
    public ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken ct = default) 
        => _stream.WriteAsync(buffer, ct);

    public ValueTask WriteCodeMessageAsync(MessageTypeCode messageTypeCode, CancellationToken ct = default)
    {
        CodeMessage m = new CodeMessage(messageTypeCode);
        return WriteStructAsync(m, ct);
    }

    public async ValueTask<MessageTypeCode> ReadCodeMessageAsync(CancellationToken ct = default)
    {
        CodeMessage m = await ReadStructAsync<CodeMessage>(ct);
        m.ThrowIfInvalid();
        return m.type_code;
    }
    
    public async ValueTask<DataMessageHeader> ReadDataMessageHeaderAsync(CancellationToken ct = default)
    {
        var m = await ReadStructAsync<DataMessageHeader>(ct);
        m.ThrowIfInvalid();
        return m;
    }
    
    
    public async Task ReceivePingSendPong()
    {
        var messageCode = await ReadCodeMessageAsync();
        if(messageCode != MessageTypeCode.Ping)
            throw new Exception($"Failed to test application protocol: expected Ping, got {messageCode}");
        await WriteCodeMessageAsync(MessageTypeCode.Pong);
    }
    
    
    public async Task SendPingReceivePong()
    {
        await WriteCodeMessageAsync(MessageTypeCode.Ping);
        var messageCode = await ReadCodeMessageAsync();
        if(messageCode != MessageTypeCode.Pong)
            throw new Exception($"Failed to test application protocol: expected Pong, got {messageCode}");
    }
    
    public async ValueTask DisposeAsync()
    {
        _stream.Close();
        await _stream.DisposeAsync();
    }
}