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

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

namespace NATS.Server.Tests.LeafNodes;

/// <summary>
/// Tests for solicited (outbound) leaf node connections with retry logic,
/// exponential backoff, JetStream domain propagation, and cancellation.
/// Go reference: leafnode.go — connectSolicited, solicitLeafNode.
/// </summary>
public class LeafSolicitedConnectionTests
{
    // Go: TestLeafNodeBasicAuthSingleton server/leafnode_test.go:602
    [Fact]
    public async Task ConnectSolicited_ValidUrl_EstablishesConnection()
    {
        // Start a hub server with leaf node listener
        var hubOptions = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            LeafNode = new LeafNodeOptions
            {
                Host = "127.0.0.1",
                Port = 0,
            },
        };

        var hub = new NatsServer(hubOptions, NullLoggerFactory.Instance);
        var hubCts = new CancellationTokenSource();
        _ = hub.StartAsync(hubCts.Token);
        await hub.WaitForReadyAsync();

        // Create a spoke server that connects to the hub
        var spokeOptions = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            LeafNode = new LeafNodeOptions
            {
                Host = "127.0.0.1",
                Port = 0,
                Remotes = [hub.LeafListen!],
            },
        };

        var spoke = new NatsServer(spokeOptions, NullLoggerFactory.Instance);
        var spokeCts = new CancellationTokenSource();
        _ = spoke.StartAsync(spokeCts.Token);
        await spoke.WaitForReadyAsync();

        // Wait for leaf connections to establish
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        while (!timeout.IsCancellationRequested && (hub.Stats.Leafs == 0 || spoke.Stats.Leafs == 0))
            await Task.Delay(50, timeout.Token).ContinueWith(_ => { }, TaskScheduler.Default);

        hub.Stats.Leafs.ShouldBeGreaterThan(0);
        spoke.Stats.Leafs.ShouldBeGreaterThan(0);

        await spokeCts.CancelAsync();
        await hubCts.CancelAsync();
        spoke.Dispose();
        hub.Dispose();
        spokeCts.Dispose();
        hubCts.Dispose();
    }

    // Go: leafnode.go — reconnect with backoff on connection failure
    [Fact]
    public async Task ConnectSolicited_InvalidUrl_RetriesWithBackoff()
    {
        // Create a leaf node manager targeting a non-existent endpoint
        var options = new LeafNodeOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            Remotes = ["127.0.0.1:19999"], // Nothing listening here
        };

        var stats = new ServerStats();
        var manager = new LeafNodeManager(
            options, stats, "test-server",
            _ => { }, _ => { },
            NullLogger<LeafNodeManager>.Instance);

        // Start the manager — it will try to connect to 127.0.0.1:19999 and fail
        using var cts = new CancellationTokenSource();
        await manager.StartAsync(cts.Token);

        // Give it some time to attempt connections
        await Task.Delay(500);

        // No connections should have succeeded
        stats.Leafs.ShouldBe(0);

        await cts.CancelAsync();
        await manager.DisposeAsync();
    }

    // Go: leafnode.go — backoff caps at 60 seconds
    [Fact]
    public void ConnectSolicited_MaxBackoff_CapsAt60Seconds()
    {
        // Verify the backoff calculation caps at 60 seconds
        LeafNodeManager.ComputeBackoff(0).ShouldBe(TimeSpan.FromSeconds(1));
        LeafNodeManager.ComputeBackoff(1).ShouldBe(TimeSpan.FromSeconds(2));
        LeafNodeManager.ComputeBackoff(2).ShouldBe(TimeSpan.FromSeconds(4));
        LeafNodeManager.ComputeBackoff(3).ShouldBe(TimeSpan.FromSeconds(8));
        LeafNodeManager.ComputeBackoff(4).ShouldBe(TimeSpan.FromSeconds(16));
        LeafNodeManager.ComputeBackoff(5).ShouldBe(TimeSpan.FromSeconds(32));
        LeafNodeManager.ComputeBackoff(6).ShouldBe(TimeSpan.FromSeconds(60)); // Capped
        LeafNodeManager.ComputeBackoff(7).ShouldBe(TimeSpan.FromSeconds(60)); // Still capped
        LeafNodeManager.ComputeBackoff(100).ShouldBe(TimeSpan.FromSeconds(60)); // Still capped
    }

    // Go: leafnode.go — JsDomain in leafInfo propagated during handshake
    [Fact]
    public async Task JetStreamDomain_PropagatedInHandshake()
    {
        // Start a hub with JetStream domain
        var hubOptions = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            LeafNode = new LeafNodeOptions
            {
                Host = "127.0.0.1",
                Port = 0,
                JetStreamDomain = "hub-domain",
            },
        };

        var hub = new NatsServer(hubOptions, NullLoggerFactory.Instance);
        var hubCts = new CancellationTokenSource();
        _ = hub.StartAsync(hubCts.Token);
        await hub.WaitForReadyAsync();

        // Create a raw socket connection to verify the handshake includes domain
        using var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        var leafEndpoint = hub.LeafListen!.Split(':');
        await client.ConnectAsync(IPAddress.Parse(leafEndpoint[0]), int.Parse(leafEndpoint[1]));

        using var stream = new NetworkStream(client, ownsSocket: false);

        // Send our LEAF handshake with a domain
        var outMsg = Encoding.ASCII.GetBytes("LEAF test-spoke domain=spoke-domain\r\n");
        await stream.WriteAsync(outMsg);
        await stream.FlushAsync();

        // Read the hub's handshake response
        var response = await ReadLineAsync(stream);

        // The hub's handshake should include the JetStream domain
        response.ShouldStartWith("LEAF ");
        response.ShouldContain("domain=hub-domain");

        await hubCts.CancelAsync();
        hub.Dispose();
        hubCts.Dispose();
    }

    // Go: leafnode.go — cancellation stops reconnect loop
    [Fact]
    public async Task Retry_CancellationToken_StopsRetrying()
    {
        var options = new LeafNodeOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            Remotes = ["127.0.0.1:19998"], // Nothing listening
        };

        var stats = new ServerStats();
        var manager = new LeafNodeManager(
            options, stats, "test-server",
            _ => { }, _ => { },
            NullLogger<LeafNodeManager>.Instance);

        using var cts = new CancellationTokenSource();
        await manager.StartAsync(cts.Token);

        // Let it attempt at least one retry
        await Task.Delay(200);

        // Cancel — the retry loop should stop promptly
        await cts.CancelAsync();
        await manager.DisposeAsync();

        // No connections should have been established
        stats.Leafs.ShouldBe(0);
    }

    // Go: leafnode.go — verify backoff delay sequence
    [Fact]
    public void ExponentialBackoff_CalculatesCorrectDelays()
    {
        var delays = new List<TimeSpan>();
        for (var i = 0; i < 10; i++)
            delays.Add(LeafNodeManager.ComputeBackoff(i));

        // Verify the sequence: 1, 2, 4, 8, 16, 32, 60, 60, 60, 60
        delays[0].ShouldBe(TimeSpan.FromSeconds(1));
        delays[1].ShouldBe(TimeSpan.FromSeconds(2));
        delays[2].ShouldBe(TimeSpan.FromSeconds(4));
        delays[3].ShouldBe(TimeSpan.FromSeconds(8));
        delays[4].ShouldBe(TimeSpan.FromSeconds(16));
        delays[5].ShouldBe(TimeSpan.FromSeconds(32));

        // After attempt 5, all should be capped at 60s
        for (var i = 6; i < 10; i++)
            delays[i].ShouldBe(TimeSpan.FromSeconds(60));

        // Negative attempt should be treated as 0
        LeafNodeManager.ComputeBackoff(-1).ShouldBe(TimeSpan.FromSeconds(1));
    }

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

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