LumaTunnel

using System.Net.Sockets;
using LumaTunnel.Client.Core.Abstractions;

namespace LumaTunnel.Client.Core.Proxy;

internal static class ProxyRelay
{
    public static async Task RunAsync(
        TcpClient client,
        Stream remote,
        Action<int> upload,
        Action<int> download,
        CancellationToken cancellationToken)
    {
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        var local = client.GetStream();
        var upstream = PumpAsync(local, remote, upload, async () =>
        {
            if (remote is IHalfCloseable halfCloseable)
                await halfCloseable.HalfCloseWriteAsync(linked.Token).ConfigureAwait(false);
        }, linked.Token);
        var downstream = PumpAsync(remote, local, download, () =>
        {
            try { client.Client.Shutdown(SocketShutdown.Send); } catch (SocketException) { }
            return ValueTask.CompletedTask;
        }, linked.Token);
        var first = await Task.WhenAny(upstream, downstream).ConfigureAwait(false);
        if (first.IsFaulted || first.IsCanceled)
            linked.Cancel();
        await Task.WhenAll(upstream, downstream).ConfigureAwait(false);
    }

    private static async Task PumpAsync(Stream source, Stream destination, Action<int> countBytes, Func<ValueTask> completed, CancellationToken cancellationToken)
    {
        var buffer = new byte[32 * 1024];
        while (true)
        {
            var count = await source.ReadAsync(buffer, cancellationToken).ConfigureAwait(false);
            if (count == 0)
                break;
            await destination.WriteAsync(buffer.AsMemory(0, count), cancellationToken).ConfigureAwait(false);
            countBytes(count);
        }
        await completed().ConfigureAwait(false);
    }
}