natsdotnet/tests/NATS.Server.LeafNodes.Tests/LeafNodes/LeafConnectionParityBatch3Tests.cs

using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Text.Json;
using NATS.Server.Configuration;
using NATS.Server.LeafNodes;

namespace NATS.Server.LeafNodes.Tests.LeafNodes;

public class LeafConnectionParityBatch3Tests
{
    [Fact]
    public async Task SendLeafConnect_writes_connect_json_payload_with_expected_fields()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var endpoint = (IPEndPoint)listener.LocalEndpoint;

        using var remoteSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await remoteSocket.ConnectAsync(IPAddress.Loopback, endpoint.Port);
        using var leafSocket = await listener.AcceptSocketAsync();
        await using var leaf = new LeafConnection(leafSocket);

        var info = new LeafConnectInfo
        {
            Jwt = "jwt-token",
            Nkey = "nkey",
            Sig = "sig",
            Hub = true,
            Cluster = "C1",
            Headers = true,
            JetStream = true,
            Compression = "s2_auto",
            RemoteAccount = "A",
            Proto = 1,
        };

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        await leaf.SendLeafConnectAsync(info, timeout.Token);

        var line = await ReadLineAsync(remoteSocket, timeout.Token);
        line.ShouldStartWith("CONNECT ");

        var payload = line["CONNECT ".Length..];
        using var json = JsonDocument.Parse(payload);
        var root = json.RootElement;

        root.GetProperty("jwt").GetString().ShouldBe("jwt-token");
        root.GetProperty("nkey").GetString().ShouldBe("nkey");
        root.GetProperty("sig").GetString().ShouldBe("sig");
        root.GetProperty("hub").GetBoolean().ShouldBeTrue();
        root.GetProperty("cluster").GetString().ShouldBe("C1");
        root.GetProperty("headers").GetBoolean().ShouldBeTrue();
        root.GetProperty("jetstream").GetBoolean().ShouldBeTrue();
        root.GetProperty("compression").GetString().ShouldBe("s2_auto");
        root.GetProperty("remote_account").GetString().ShouldBe("A");
        root.GetProperty("proto").GetInt32().ShouldBe(1);
    }

    [Fact]
    public async Task RemoteCluster_returns_cluster_from_leaf_handshake_attributes()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var endpoint = (IPEndPoint)listener.LocalEndpoint;

        using var remoteSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await remoteSocket.ConnectAsync(IPAddress.Loopback, endpoint.Port);
        using var leafSocket = await listener.AcceptSocketAsync();
        await using var leaf = new LeafConnection(leafSocket);

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var handshakeTask = leaf.PerformOutboundHandshakeAsync("LOCAL", timeout.Token);

        (await ReadLineAsync(remoteSocket, timeout.Token)).ShouldBe("LEAF LOCAL");
        await WriteLineAsync(remoteSocket, "LEAF REMOTE cluster=HUB-A domain=JS-A", timeout.Token);

        await handshakeTask;

        leaf.RemoteId.ShouldBe("REMOTE");
        leaf.RemoteCluster().ShouldBe("HUB-A");
        leaf.RemoteJetStreamDomain.ShouldBe("JS-A");
    }

    [Fact]
    public async Task SetLeafConnectDelayIfSoliciting_sets_delay_only_for_solicited_connections()
    {
        await using var solicited = await CreateConnectionAsync();
        solicited.IsSolicited = true;
        solicited.SetLeafConnectDelayIfSoliciting(TimeSpan.FromSeconds(30));
        solicited.GetConnectDelay().ShouldBe(TimeSpan.FromSeconds(30));

        await using var inbound = await CreateConnectionAsync();
        inbound.IsSolicited = false;
        inbound.SetLeafConnectDelayIfSoliciting(TimeSpan.FromSeconds(30));
        inbound.GetConnectDelay().ShouldBe(TimeSpan.Zero);
    }

    [Fact]
    public async Task LeafProcessErr_maps_known_errors_to_reconnect_delays_on_solicited_connections()
    {
        await using var leaf = await CreateConnectionAsync();
        leaf.IsSolicited = true;

        leaf.LeafProcessErr("Permissions Violation for Subscription to foo");
        leaf.GetConnectDelay().ShouldBe(LeafNodeManager.LeafNodeReconnectAfterPermViolation);

        leaf.LeafProcessErr("Loop detected");
        leaf.GetConnectDelay().ShouldBe(LeafNodeManager.LeafNodeReconnectDelayAfterLoopDetected);

        leaf.LeafProcessErr("Cluster name is same");
        leaf.GetConnectDelay().ShouldBe(LeafNodeManager.LeafNodeReconnectDelayAfterClusterNameSame);
    }

    [Fact]
    public async Task LeafSubPermViolation_and_LeafPermViolation_set_permission_delay()
    {
        await using var leaf = await CreateConnectionAsync();
        leaf.IsSolicited = true;

        leaf.LeafSubPermViolation("subj.A");
        leaf.GetConnectDelay().ShouldBe(LeafNodeManager.LeafNodeReconnectAfterPermViolation);

        leaf.SetLeafConnectDelayIfSoliciting(TimeSpan.Zero);
        leaf.LeafPermViolation(pub: true, subj: "subj.B");
        leaf.GetConnectDelay().ShouldBe(LeafNodeManager.LeafNodeReconnectAfterPermViolation);
    }

    [Fact]
    public async Task CancelMigrateTimer_stops_pending_timer_callback()
    {
        var remote = new RemoteLeafOptions();
        using var signal = new SemaphoreSlim(0, 1);

        remote.StartMigrateTimer(_ => signal.Release(), TimeSpan.FromMilliseconds(120));
        remote.CancelMigrateTimer();

        // The timer was disposed before its 120 ms deadline. We wait 300 ms via WaitAsync;
        // if the callback somehow fires it will release the semaphore and WaitAsync returns
        // true, which the assertion catches. No Task.Delay required.
        var fired = await signal.WaitAsync(TimeSpan.FromMilliseconds(300));
        fired.ShouldBeFalse();
    }

    private static async Task<LeafConnection> CreateConnectionAsync()
    {
        var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var endpoint = (IPEndPoint)listener.LocalEndpoint;

        var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client.ConnectAsync(IPAddress.Loopback, endpoint.Port);

        var server = await listener.AcceptSocketAsync();
        listener.Stop();
        client.Dispose();

        return new LeafConnection(server);
    }

    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)
                throw new IOException("Connection closed while reading line");
            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();
}