PalladiumWallet/src/Core/Net/ElectrumClient.cs

using System.Buffers;
using System.Collections.Concurrent;
using System.IO.Pipelines;
using System.Net.Security;
using System.Net.Sockets;
using System.Text.Json;
using System.Threading.Channels;

namespace PalladiumWallet.Core.Net;

/// <summary>
/// Client for the indexing server protocol (blueprint §10):
/// newline-delimited JSON-RPC 2.0 over TCP, optionally TLS with TOFU pinning.
/// Notifications (subscriptions) arrive on the <see cref="NotificationReceived"/> event.
/// </summary>
public sealed class ElectrumClient : IAsyncDisposable
{
    public const string ClientName = "PalladiumWallet";
    public const string ProtocolVersion = "1.4";

    private readonly TcpClient _tcp;
    private readonly Stream _stream;
    private readonly ConcurrentDictionary<long, TaskCompletionSource<JsonElement>> _pending = new();
    private readonly CancellationTokenSource _cts = new();
    private readonly Task _readLoop;
    private readonly Task _writeLoop;

    // Single-reader channel: tasks write the payloads without a lock;
    // the write loop drains everything in a single WriteAsync+FlushAsync —
    // identical to Electrum's asyncio buffered writer.
    private readonly Channel<byte[]> _outgoing = Channel.CreateUnbounded<byte[]>(
        new UnboundedChannelOptions { SingleReader = true, AllowSynchronousContinuations = false });

    // Cap on in-flight requests (not in the write queue, but awaiting a
    // response from the server). The write loop already batches the send into a
    // single segment; this gate avoids flooding the server with thousands of
    // simultaneous requests on large wallets → no bursts of -101/-102 nor
    // connection drops. Writes still stay pipelined up to this degree.
    private const int MaxInFlight = 32;
    private readonly SemaphoreSlim _inFlight = new(MaxInFlight, MaxInFlight);

    private long _nextId;

    public string Host { get; }
    public int Port { get; }
    public bool UseSsl { get; }
    public bool IsConnected => _tcp.Connected && !_cts.IsCancellationRequested;

    public event Action<string, JsonElement>? NotificationReceived;
    public event Action<Exception?>? Disconnected;

    private ElectrumClient(TcpClient tcp, Stream stream, string host, int port, bool useSsl)
    {
        _tcp = tcp;
        _stream = stream;
        Host = host;
        Port = port;
        UseSsl = useSsl;
        _readLoop  = Task.Run(ReadLoopAsync);
        _writeLoop = Task.Run(WriteLoopAsync);
    }

    public static async Task<ElectrumClient> ConnectAsync(string host, int port, bool useSsl,
        CertificatePinStore? pins = null, CancellationToken ct = default)
    {
        var tcp = new TcpClient { NoDelay = true };
        try
        {
            await tcp.ConnectAsync(host, port, ct);
            Stream stream = tcp.GetStream();
            if (useSsl)
            {
                var pinOk = true;
                var ssl = new SslStream(stream, leaveInnerStreamOpen: false,
                    (_, cert, _, _) =>
                    {
                        if (cert is null) return false;
                        pinOk = pins is null || pins.VerifyOrPin(host, port, cert);
                        return pinOk;
                    });
                try
                {
                    await ssl.AuthenticateAsClientAsync(
                        new SslClientAuthenticationOptions { TargetHost = host }, ct);
                }
                catch (System.Security.Authentication.AuthenticationException) when (!pinOk)
                {
                    throw new CertificatePinMismatchException(host, port);
                }
                stream = ssl;
            }

            var client = new ElectrumClient(tcp, stream, host, port, useSsl);
            await client.RequestAsync("server.version", ct, ClientName, ProtocolVersion);
            return client;
        }
        catch
        {
            tcp.Dispose();
            throw;
        }
    }

    public async Task<JsonElement> RequestAsync(string method, CancellationToken ct = default,
        params object?[] parameters)
    {
        // Gate before going in-flight: beyond MaxInFlight pending requests
        // we wait for a response to free a slot, instead of flooding the server.
        await _inFlight.WaitAsync(ct);
        try
        {
            var id = Interlocked.Increment(ref _nextId);
            var tcs = new TaskCompletionSource<JsonElement>(TaskCreationOptions.RunContinuationsAsynchronously);
            _pending[id] = tcs;

            var payload = JsonSerializer.SerializeToUtf8Bytes(new
            {
                jsonrpc = "2.0",
                id,
                method,
                @params = parameters,
            });

            _outgoing.Writer.TryWrite(payload);

            await using var registration = ct.Register(() =>
            {
                _pending.TryRemove(id, out _);
                tcs.TrySetCanceled(ct);
            });
            return await tcs.Task;
        }
        finally
        {
            _inFlight.Release();
        }
    }

    /// <summary>
    /// Drain loop: empties the channel into a single buffer → one WriteAsync+FlushAsync
    /// for all queued messages. When N requests are queued, they are
    /// transmitted in a single TCP segment instead of N serial flushes.
    /// </summary>
    private async Task WriteLoopAsync()
    {
        try
        {
            while (await _outgoing.Reader.WaitToReadAsync(_cts.Token))
            {
                using var ms = new MemoryStream(512);
                while (_outgoing.Reader.TryRead(out var data))
                {
                    ms.Write(data);
                    ms.WriteByte((byte)'\n');
                }
                await _stream.WriteAsync(ms.GetBuffer().AsMemory(0, (int)ms.Length), _cts.Token);
                await _stream.FlushAsync(_cts.Token);
            }
        }
        catch (OperationCanceledException) { }
        catch (Exception ex)
        {
            foreach (var (_, pending) in _pending)
                pending.TrySetException(ex);
        }
    }

    /// <summary>
    /// Read loop with PipeReader: pooled buffers, zero per-response allocations
    /// (no intermediate string), JSON parsing directly from a byte span.
    /// </summary>
    private async Task ReadLoopAsync()
    {
        Exception? failure = null;
        var pipe = PipeReader.Create(_stream, new StreamPipeReaderOptions(leaveOpen: true));
        try
        {
            while (true)
            {
                ReadResult result;
                try { result = await pipe.ReadAsync(_cts.Token); }
                catch (OperationCanceledException) { break; }

                var buffer = result.Buffer;
                try
                {
                    while (TrySliceLine(ref buffer, out var line))
                        if (!line.IsEmpty)
                            DispatchLine(line);
                }
                finally
                {
                    // consumed = everything we have consumed (up to the last \n)
                    // examined = everything we have looked at (up to the end of the buffer)
                    pipe.AdvanceTo(buffer.Start, buffer.End);
                }

                if (result.IsCompleted) break;
            }
        }
        catch (Exception ex)
        {
            failure = ex;
        }
        finally
        {
            await pipe.CompleteAsync();
            foreach (var (_, tcs) in _pending)
                tcs.TrySetException(failure ?? new IOException("Connection to the server closed."));
            _pending.Clear();
            Disconnected?.Invoke(failure);
        }
    }

    private static bool TrySliceLine(ref ReadOnlySequence<byte> buffer,
        out ReadOnlySequence<byte> line)
    {
        var pos = buffer.PositionOf((byte)'\n');
        if (pos is null) { line = default; return false; }
        line   = buffer.Slice(0, pos.Value);
        buffer = buffer.Slice(buffer.GetPosition(1, pos.Value));
        return true;
    }

    private void DispatchLine(ReadOnlySequence<byte> line)
    {
        if (line.IsSingleSegment)
        {
            DispatchSpan(line.FirstSpan);
            return;
        }
        // Multi-segment (very long response): copy to ArrayPool then parse.
        var len = (int)line.Length;
        var buf = ArrayPool<byte>.Shared.Rent(len);
        try
        {
            line.CopyTo(buf);
            DispatchSpan(buf.AsSpan(0, len));
        }
        finally { ArrayPool<byte>.Shared.Return(buf); }
    }

    private void DispatchSpan(ReadOnlySpan<byte> utf8)
    {
        // JsonDocument.Parse via Utf8JsonReader: no intermediate string,
        // parsing directly from the pooled span.
        var reader = new Utf8JsonReader(utf8);
        using var doc = JsonDocument.ParseValue(ref reader);
        var root = doc.RootElement;

        if (root.TryGetProperty("id", out var idEl) && idEl.ValueKind == JsonValueKind.Number)
        {
            if (!_pending.TryRemove(idEl.GetInt64(), out var tcs)) return;
            if (root.TryGetProperty("error", out var err) && err.ValueKind != JsonValueKind.Null)
                tcs.TrySetException(new ElectrumServerException(err.ToString()));
            else
                tcs.TrySetResult(root.GetProperty("result").Clone());
        }
        else if (root.TryGetProperty("method", out var methodEl))
        {
            NotificationReceived?.Invoke(
                methodEl.GetString()!,
                root.GetProperty("params").Clone());
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _cts.CancelAsync();
        _outgoing.Writer.Complete();
        _tcp.Close();
        try { await _readLoop;  } catch { }
        try { await _writeLoop; } catch { }
        _cts.Dispose();
        _inFlight.Dispose();
    }
}

/// <summary>Error returned by the server ("error" field of the JSON-RPC response).</summary>
public sealed class ElectrumServerException(string error) : Exception(error);

/// <summary>
/// The server certificate changed with respect to the saved one (TOFU, §9).
/// It is unlocked with an explicit reset of the certificates.
/// </summary>
public sealed class CertificatePinMismatchException(string host, int port) : Exception(
    $"Il certificato TLS di {host}:{port} è cambiato rispetto a quello salvato. " +
    "Se il server ha rinnovato il certificato, esegui il reset dei certificati SSL.");