natsdotnet/tests/NATS.Server.Clustering.Tests/InterServerAccountProtocolTests.cs

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

namespace NATS.Server.Clustering.Tests;

public class InterServerAccountProtocolTests
{
    [Fact]
    public async Task Aplus_Aminus_frames_include_account_scope_and_do_not_leak_interest_across_accounts()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();

        var port = ((IPEndPoint)listener.LocalEndpoint).Port;
        using var remoteSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await remoteSocket.ConnectAsync(IPAddress.Loopback, port);
        using var gatewaySocket = await listener.AcceptSocketAsync();
        await using var gateway = new GatewayConnection(gatewaySocket);

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var handshakeTask = gateway.PerformOutboundHandshakeAsync("LOCAL", timeout.Token);
        (await ReadLineAsync(remoteSocket, timeout.Token)).ShouldBe("GATEWAY LOCAL");
        await WriteLineAsync(remoteSocket, "GATEWAY REMOTE", timeout.Token);
        await handshakeTask;

        var received = new TaskCompletionSource<RemoteSubscription>(TaskCreationOptions.RunContinuationsAsynchronously);
        gateway.RemoteSubscriptionReceived = sub =>
        {
            received.TrySetResult(sub);
            return Task.CompletedTask;
        };
        gateway.StartLoop(timeout.Token);

        await WriteLineAsync(remoteSocket, "A+ A orders.*", timeout.Token);
        var aPlus = await received.Task.WaitAsync(timeout.Token);
        aPlus.Account.ShouldBe("A");
        aPlus.Subject.ShouldBe("orders.*");
        aPlus.IsRemoval.ShouldBeFalse();

        var subList = new SubList();
        subList.ApplyRemoteSub(aPlus);
        subList.HasRemoteInterest("A", "orders.created").ShouldBeTrue();
        subList.HasRemoteInterest("B", "orders.created").ShouldBeFalse();

        var removedTcs = new TaskCompletionSource<RemoteSubscription>(TaskCreationOptions.RunContinuationsAsynchronously);
        gateway.RemoteSubscriptionReceived = sub =>
        {
            removedTcs.TrySetResult(sub);
            return Task.CompletedTask;
        };

        await WriteLineAsync(remoteSocket, "A- A orders.*", timeout.Token);
        var aMinus = await removedTcs.Task.WaitAsync(timeout.Token);
        aMinus.Account.ShouldBe("A");
        aMinus.IsRemoval.ShouldBeTrue();

        subList.ApplyRemoteSub(aMinus);
        subList.HasRemoteInterest("A", "orders.created").ShouldBeFalse();
    }

    private static async Task<string> ReadLineAsync(Socket socket, CancellationToken ct)
    {
        var bytes = new List<byte>(64);
        var single = new byte[1];
        while (true)
        {
            var read = await socket.ReceiveAsync(single, SocketFlags.None, ct);
            if (read == 0)
                break;
            if (single[0] == (byte)'\n')
                break;
            if (single[0] != (byte)'\r')
                bytes.Add(single[0]);
        }

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

    private static Task WriteLineAsync(Socket socket, string line, CancellationToken ct)
        => socket.SendAsync(Encoding.ASCII.GetBytes($"{line}\r\n"), SocketFlags.None, ct).AsTask();
}