natsdotnet/src/NATS.Server/LeafNodes/LeafConnection.cs

using System.Net.Sockets;
using System.Text;
using NATS.Server.Subscriptions;

namespace NATS.Server.LeafNodes;

public sealed class LeafConnection(Socket socket) : IAsyncDisposable
{
    private readonly NetworkStream _stream = new(socket, ownsSocket: true);
    private readonly SemaphoreSlim _writeGate = new(1, 1);
    private readonly CancellationTokenSource _closedCts = new();
    private Task? _loopTask;

    public string? RemoteId { get; private set; }
    public string RemoteEndpoint => socket.RemoteEndPoint?.ToString() ?? Guid.NewGuid().ToString("N");
    public Func<RemoteSubscription, Task>? RemoteSubscriptionReceived { get; set; }
    public Func<LeafMessage, Task>? MessageReceived { get; set; }

    public async Task PerformOutboundHandshakeAsync(string serverId, CancellationToken ct)
    {
        await WriteLineAsync($"LEAF {serverId}", ct);
        var line = await ReadLineAsync(ct);
        RemoteId = ParseHandshake(line);
    }

    public async Task PerformInboundHandshakeAsync(string serverId, CancellationToken ct)
    {
        var line = await ReadLineAsync(ct);
        RemoteId = ParseHandshake(line);
        await WriteLineAsync($"LEAF {serverId}", ct);
    }

    public void StartLoop(CancellationToken ct)
    {
        if (_loopTask != null)
            return;

        var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _closedCts.Token);
        _loopTask = Task.Run(() => ReadLoopAsync(linked.Token), linked.Token);
    }

    public Task WaitUntilClosedAsync(CancellationToken ct)
        => _loopTask?.WaitAsync(ct) ?? Task.CompletedTask;

    public Task SendLsPlusAsync(string account, string subject, string? queue, CancellationToken ct)
        => WriteLineAsync(queue is { Length: > 0 } ? $"LS+ {account} {subject} {queue}" : $"LS+ {account} {subject}", ct);

    public Task SendLsMinusAsync(string account, string subject, string? queue, CancellationToken ct)
        => WriteLineAsync(queue is { Length: > 0 } ? $"LS- {account} {subject} {queue}" : $"LS- {account} {subject}", ct);

    public async Task SendMessageAsync(string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        var reply = string.IsNullOrEmpty(replyTo) ? "-" : replyTo;
        await _writeGate.WaitAsync(ct);
        try
        {
            var control = Encoding.ASCII.GetBytes($"LMSG {subject} {reply} {payload.Length}\r\n");
            await _stream.WriteAsync(control, ct);
            if (!payload.IsEmpty)
                await _stream.WriteAsync(payload, ct);
            await _stream.WriteAsync("\r\n"u8.ToArray(), ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _closedCts.CancelAsync();
        if (_loopTask != null)
            await _loopTask.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
        _closedCts.Dispose();
        _writeGate.Dispose();
        await _stream.DisposeAsync();
    }

    private async Task ReadLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            string line;
            try
            {
                line = await ReadLineAsync(ct);
            }
            catch
            {
                break;
            }

            if (line.StartsWith("LS+ ", StringComparison.Ordinal))
            {
                var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
                if (RemoteSubscriptionReceived != null && TryParseAccountScopedInterest(parts, out var account, out var parsedSubject, out var queue))
                {
                    await RemoteSubscriptionReceived(new RemoteSubscription(parsedSubject, queue, RemoteId ?? string.Empty, account));
                }
                continue;
            }

            if (line.StartsWith("LS- ", StringComparison.Ordinal))
            {
                var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
                if (RemoteSubscriptionReceived != null && TryParseAccountScopedInterest(parts, out var account, out var parsedSubject, out var queue))
                {
                    await RemoteSubscriptionReceived(RemoteSubscription.Removal(parsedSubject, queue, RemoteId ?? string.Empty, account));
                }
                continue;
            }

            if (!line.StartsWith("LMSG ", StringComparison.Ordinal))
                continue;

            var args = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
            if (args.Length < 4 || !int.TryParse(args[3], out var size) || size < 0)
                continue;

            var payload = await ReadPayloadAsync(size, ct);
            if (MessageReceived != null)
                await MessageReceived(new LeafMessage(args[1], args[2] == "-" ? null : args[2], payload));
        }
    }

    private async Task<ReadOnlyMemory<byte>> ReadPayloadAsync(int size, CancellationToken ct)
    {
        var payload = new byte[size];
        var offset = 0;
        while (offset < size)
        {
            var read = await _stream.ReadAsync(payload.AsMemory(offset, size - offset), ct);
            if (read == 0)
                throw new IOException("Leaf payload read closed");
            offset += read;
        }

        var trailer = new byte[2];
        _ = await _stream.ReadAsync(trailer, ct);
        return payload;
    }

    private async Task WriteLineAsync(string line, CancellationToken ct)
    {
        await _writeGate.WaitAsync(ct);
        try
        {
            var bytes = Encoding.ASCII.GetBytes($"{line}\r\n");
            await _stream.WriteAsync(bytes, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    private async Task<string> ReadLineAsync(CancellationToken ct)
    {
        var bytes = new List<byte>(64);
        var single = new byte[1];
        while (true)
        {
            var read = await _stream.ReadAsync(single, ct);
            if (read == 0)
                throw new IOException("Leaf closed");
            if (single[0] == (byte)'\n')
                break;
            if (single[0] != (byte)'\r')
                bytes.Add(single[0]);
        }

        return Encoding.ASCII.GetString([.. bytes]);
    }

    private static string ParseHandshake(string line)
    {
        if (!line.StartsWith("LEAF ", StringComparison.OrdinalIgnoreCase))
            throw new InvalidOperationException("Invalid leaf handshake");

        var id = line[5..].Trim();
        if (id.Length == 0)
            throw new InvalidOperationException("Leaf handshake missing id");
        return id;
    }

    private static bool TryParseAccountScopedInterest(string[] parts, out string account, out string subject, out string? queue)
    {
        account = "$G";
        subject = string.Empty;
        queue = null;

        if (parts.Length < 2)
            return false;

        // New format: LS+ <account> <subject> [queue]
        // Legacy format: LS+ <subject> [queue]
        if (parts.Length >= 3 && !LooksLikeSubject(parts[1]))
        {
            account = parts[1];
            subject = parts[2];
            queue = parts.Length >= 4 ? parts[3] : null;
            return true;
        }

        subject = parts[1];
        queue = parts.Length >= 3 ? parts[2] : null;
        return true;
    }

    private static bool LooksLikeSubject(string token)
        => token.Contains('.', StringComparison.Ordinal)
            || token.Contains('*', StringComparison.Ordinal)
            || token.Contains('>', StringComparison.Ordinal);
}

public sealed record LeafMessage(string Subject, string? ReplyTo, ReadOnlyMemory<byte> Payload);