natsdotnet/tests/NATS.Server.Tests/GatewayProtocolTests.cs

using System.Net;
using System.Net.Sockets;
using System.Text;
using Microsoft.Extensions.Logging.Abstractions;
using NATS.Server.Configuration;

namespace NATS.Server.Tests;

public class GatewayProtocolTests
{
    [Fact]
    public async Task Gateway_link_establishes_and_forwards_interested_message()
    {
        await using var fx = await GatewayFixture.StartTwoClustersAsync();
        await fx.SubscribeRemoteClusterAsync("g.>");
        await fx.PublishLocalClusterAsync("g.test", "hello");
        (await fx.ReadRemoteClusterMessageAsync()).ShouldContain("hello");
    }
}

internal sealed class GatewayFixture : IAsyncDisposable
{
    private readonly NatsServer _local;
    private readonly NatsServer _remote;
    private readonly CancellationTokenSource _localCts;
    private readonly CancellationTokenSource _remoteCts;
    private Socket? _remoteSubscriber;
    private Socket? _localPublisher;

    private GatewayFixture(NatsServer local, NatsServer remote, CancellationTokenSource localCts, CancellationTokenSource remoteCts)
    {
        _local = local;
        _remote = remote;
        _localCts = localCts;
        _remoteCts = remoteCts;
    }

    public static async Task<GatewayFixture> StartTwoClustersAsync()
    {
        var localOptions = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            Gateway = new GatewayOptions
            {
                Name = "LOCAL",
                Host = "127.0.0.1",
                Port = 0,
            },
        };

        var local = new NatsServer(localOptions, NullLoggerFactory.Instance);
        var localCts = new CancellationTokenSource();
        _ = local.StartAsync(localCts.Token);
        await local.WaitForReadyAsync();

        var remoteOptions = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            Gateway = new GatewayOptions
            {
                Name = "REMOTE",
                Host = "127.0.0.1",
                Port = 0,
                Remotes = [local.GatewayListen!],
            },
        };

        var remote = new NatsServer(remoteOptions, NullLoggerFactory.Instance);
        var remoteCts = new CancellationTokenSource();
        _ = remote.StartAsync(remoteCts.Token);
        await remote.WaitForReadyAsync();

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        while (!timeout.IsCancellationRequested && (local.Stats.Gateways == 0 || remote.Stats.Gateways == 0))
            await Task.Delay(50, timeout.Token).ContinueWith(_ => { }, TaskScheduler.Default);

        return new GatewayFixture(local, remote, localCts, remoteCts);
    }

    public async Task SubscribeRemoteClusterAsync(string subject)
    {
        var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await sock.ConnectAsync(IPAddress.Loopback, _remote.Port);
        _remoteSubscriber = sock;

        _ = await ReadLineAsync(sock); // INFO
        await sock.SendAsync(Encoding.ASCII.GetBytes($"CONNECT {{}}\r\nSUB {subject} 1\r\nPING\r\n"));
        await ReadUntilAsync(sock, "PONG");
    }

    public async Task PublishLocalClusterAsync(string subject, string payload)
    {
        var sock = _localPublisher;
        if (sock == null)
        {
            sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            await sock.ConnectAsync(IPAddress.Loopback, _local.Port);
            _localPublisher = sock;
            _ = await ReadLineAsync(sock); // INFO
            await sock.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
            await ReadUntilAsync(sock, "PONG");
        }

        await sock.SendAsync(Encoding.ASCII.GetBytes($"PUB {subject} {payload.Length}\r\n{payload}\r\nPING\r\n"));
        await ReadUntilAsync(sock, "PONG");
    }

    public Task<string> ReadRemoteClusterMessageAsync()
    {
        if (_remoteSubscriber == null)
            throw new InvalidOperationException("Remote subscriber was not initialized.");

        return ReadUntilAsync(_remoteSubscriber, "MSG ");
    }

    public async ValueTask DisposeAsync()
    {
        _remoteSubscriber?.Dispose();
        _localPublisher?.Dispose();
        await _localCts.CancelAsync();
        await _remoteCts.CancelAsync();
        _local.Dispose();
        _remote.Dispose();
        _localCts.Dispose();
        _remoteCts.Dispose();
    }

    private static async Task<string> ReadLineAsync(Socket sock)
    {
        var buf = new byte[4096];
        var n = await sock.ReceiveAsync(buf, SocketFlags.None);
        return Encoding.ASCII.GetString(buf, 0, n);
    }

    private static async Task<string> ReadUntilAsync(Socket sock, string expected)
    {
        var sb = new StringBuilder();
        var buf = new byte[4096];
        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        while (!sb.ToString().Contains(expected, StringComparison.Ordinal))
        {
            var n = await sock.ReceiveAsync(buf, SocketFlags.None, cts.Token);
            if (n == 0)
                break;
            sb.Append(Encoding.ASCII.GetString(buf, 0, n));
        }

        return sb.ToString();
    }
}