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

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

namespace NATS.Server.Tests.LeafNodes;

/// <summary>
/// Unit tests for leaf node permission and account syncing (Gap 12.2).
/// Verifies <see cref="LeafConnection.SetPermissions"/>, <see cref="LeafNodeManager.SendPermsAndAccountInfo"/>,
/// <see cref="LeafNodeManager.InitLeafNodeSmapAndSendSubs"/>, and
/// <see cref="LeafNodeManager.GetPermSyncStatus"/>.
/// Go reference: leafnode.go — sendPermsAndAccountInfo, initLeafNodeSmapAndSendSubs.
/// </summary>
public class LeafPermissionSyncTests
{
    // ── LeafConnection.SetPermissions ─────────────────────────────────────────

    // Go: leafnode.go — sendPermsAndAccountInfo sets client.perm.pub.allow
    [Fact]
    public async Task SetPermissions_SetsPublishAllow()
    {
        await using var leaf = await CreateConnectedLeafAsync();

        leaf.SetPermissions(["orders.*", "events.>"], null);

        leaf.AllowedPublishSubjects.Count.ShouldBe(2);
        leaf.AllowedPublishSubjects.ShouldContain("orders.*");
        leaf.AllowedPublishSubjects.ShouldContain("events.>");
    }

    // Go: leafnode.go — sendPermsAndAccountInfo sets client.perm.sub.allow
    [Fact]
    public async Task SetPermissions_SetsSubscribeAllow()
    {
        await using var leaf = await CreateConnectedLeafAsync();

        leaf.SetPermissions(null, ["metrics.cpu", "metrics.memory"]);

        leaf.AllowedSubscribeSubjects.Count.ShouldBe(2);
        leaf.AllowedSubscribeSubjects.ShouldContain("metrics.cpu");
        leaf.AllowedSubscribeSubjects.ShouldContain("metrics.memory");
    }

    // Go: leafnode.go — sendPermsAndAccountInfo marks perms as synced
    [Fact]
    public async Task SetPermissions_MarksPermsSynced()
    {
        await using var leaf = await CreateConnectedLeafAsync();
        leaf.PermsSynced.ShouldBeFalse();

        leaf.SetPermissions(["foo.*"], ["bar.*"]);

        leaf.PermsSynced.ShouldBeTrue();
    }

    // Go: leafnode.go — clearing perms via null resets the allow lists
    [Fact]
    public async Task SetPermissions_NullLists_ClearsPermissions()
    {
        await using var leaf = await CreateConnectedLeafAsync();
        leaf.SetPermissions(["orders.*"], ["events.*"]);
        leaf.AllowedPublishSubjects.Count.ShouldBe(1);
        leaf.AllowedSubscribeSubjects.Count.ShouldBe(1);

        // Passing null clears both lists but still marks PermsSynced
        leaf.SetPermissions(null, null);

        leaf.AllowedPublishSubjects.ShouldBeEmpty();
        leaf.AllowedSubscribeSubjects.ShouldBeEmpty();
        leaf.PermsSynced.ShouldBeTrue();
    }

    // ── LeafNodeManager.SendPermsAndAccountInfo ───────────────────────────────

    // Go: leafnode.go — sendPermsAndAccountInfo for known connection
    [Fact]
    public async Task SendPermsAndAccountInfo_ExistingConnection_SyncsPerms()
    {
        await using var ctx = await CreateManagerWithConnectionAsync();

        var result = ctx.Manager.SendPermsAndAccountInfo(
            ctx.ConnectionId,
            "myaccount",
            ["pub.>"],
            ["sub.>"]);

        result.Found.ShouldBeTrue();
        result.PermsSynced.ShouldBeTrue();
        result.PublishAllowCount.ShouldBe(1);
        result.SubscribeAllowCount.ShouldBe(1);
    }

    // Go: leafnode.go — sendPermsAndAccountInfo returns early when connection not found
    [Fact]
    public async Task SendPermsAndAccountInfo_NonExistent_ReturnsNotFound()
    {
        await using var ctx = await CreateManagerWithConnectionAsync();

        var result = ctx.Manager.SendPermsAndAccountInfo(
            "no-such-connection-id",
            "account",
            ["foo.*"],
            null);

        result.Found.ShouldBeFalse();
        result.PermsSynced.ShouldBeFalse();
        result.PublishAllowCount.ShouldBe(0);
        result.SubscribeAllowCount.ShouldBe(0);
    }

    // Go: leafnode.go — sendPermsAndAccountInfo sets client.acc (account name)
    [Fact]
    public async Task SendPermsAndAccountInfo_SetsAccountName()
    {
        await using var ctx = await CreateManagerWithConnectionAsync();

        var result = ctx.Manager.SendPermsAndAccountInfo(
            ctx.ConnectionId,
            "tenant-alpha",
            null,
            null);

        result.Found.ShouldBeTrue();
        result.AccountName.ShouldBe("tenant-alpha");
    }

    // ── LeafNodeManager.InitLeafNodeSmapAndSendSubs ───────────────────────────

    // Go: leafnode.go — initLeafNodeSmapAndSendSubs returns subject count
    [Fact]
    public async Task InitLeafNodeSmapAndSendSubs_ReturnSubjectCount()
    {
        await using var ctx = await CreateManagerWithConnectionAsync();

        var count = ctx.Manager.InitLeafNodeSmapAndSendSubs(
            ctx.ConnectionId,
            ["orders.new", "orders.updated", "events.>"]);

        count.ShouldBe(3);
    }

    // ── LeafNodeManager.GetPermSyncStatus ─────────────────────────────────────

    // Go: leafnode.go — GetPermSyncStatus returns status for known connection
    [Fact]
    public async Task GetPermSyncStatus_FoundConnection_ReturnsStatus()
    {
        await using var ctx = await CreateManagerWithConnectionAsync();
        ctx.Manager.SendPermsAndAccountInfo(ctx.ConnectionId, "acct", ["pub.*"], ["sub.*"]);

        var status = ctx.Manager.GetPermSyncStatus(ctx.ConnectionId);

        status.Found.ShouldBeTrue();
        status.PermsSynced.ShouldBeTrue();
        status.AccountName.ShouldBe("acct");
        status.PublishAllowCount.ShouldBe(1);
        status.SubscribeAllowCount.ShouldBe(1);
    }

    // Go: leafnode.go — GetPermSyncStatus returns not-found for unknown ID
    [Fact]
    public async Task GetPermSyncStatus_NotFound_ReturnsNotFound()
    {
        await using var ctx = await CreateManagerWithConnectionAsync();

        var status = ctx.Manager.GetPermSyncStatus("unknown-id");

        status.Found.ShouldBeFalse();
        status.PermsSynced.ShouldBeFalse();
        status.AccountName.ShouldBeNull();
        status.PublishAllowCount.ShouldBe(0);
        status.SubscribeAllowCount.ShouldBe(0);
    }

    // ── Helpers ───────────────────────────────────────────────────────────────

    /// <summary>
    /// Creates a connected pair of sockets and returns the server-side LeafConnection
    /// after completing an outbound handshake. Callers must await-dispose the returned
    /// connection to avoid socket leaks.
    /// </summary>
    private static async Task<LeafConnection> CreateConnectedLeafAsync()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var port = ((IPEndPoint)listener.LocalEndpoint).Port;

        using var clientSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await clientSocket.ConnectAsync(IPAddress.Loopback, port);
        var serverSocket = await listener.AcceptSocketAsync();
        listener.Stop();

        var leaf = new LeafConnection(serverSocket);

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

        // Answer the handshake from the remote side
        var line = await ReadLineAsync(clientSocket, cts.Token);
        line.ShouldStartWith("LEAF ");
        await WriteLineAsync(clientSocket, "LEAF REMOTE", cts.Token);
        await handshakeTask;

        clientSocket.Close();
        return leaf;
    }

    /// <summary>
    /// Context returned by <see cref="CreateManagerWithConnectionAsync"/>.
    /// Disposes the manager and drains sockets on disposal.
    /// </summary>
    private sealed class ManagerContext : IAsyncDisposable
    {
        private readonly Socket _remoteSocket;
        private readonly CancellationTokenSource _cts;

        public ManagerContext(LeafNodeManager manager, string connectionId, Socket remoteSocket, CancellationTokenSource cts)
        {
            Manager = manager;
            ConnectionId = connectionId;
            _remoteSocket = remoteSocket;
            _cts = cts;
        }

        public LeafNodeManager Manager { get; }
        public string ConnectionId { get; }

        public async ValueTask DisposeAsync()
        {
            _remoteSocket.Close();
            await Manager.DisposeAsync();
            _cts.Dispose();
        }
    }

    /// <summary>
    /// Starts a <see cref="LeafNodeManager"/>, establishes one inbound leaf connection, waits
    /// for registration, and returns a context containing the manager plus the registered
    /// connection ID.
    /// </summary>
    private static async Task<ManagerContext> CreateManagerWithConnectionAsync()
    {
        var options = new LeafNodeOptions { Host = "127.0.0.1", Port = 0 };
        var manager = new LeafNodeManager(
            options,
            new ServerStats(),
            "HUB-SERVER",
            _ => { },
            _ => { },
            NullLogger<LeafNodeManager>.Instance);

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

        // Connect a raw socket acting as the spoke (inbound to the manager's listener)
        var remoteSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await remoteSocket.ConnectAsync(IPAddress.Loopback, options.Port);

        using var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        // Set up a TaskCompletionSource that fires as soon as the manager registers
        // the inbound connection — avoids any polling or timing-dependent delays.
        var registered = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
        manager.OnConnectionRegistered = id => registered.TrySetResult(id);
        timeoutCts.Token.Register(() => registered.TrySetCanceled(timeoutCts.Token));

        // For inbound connections the manager reads first, then writes
        await WriteLineAsync(remoteSocket, "LEAF SPOKE1", timeoutCts.Token);
        var response = await ReadLineAsync(remoteSocket, timeoutCts.Token);
        response.ShouldStartWith("LEAF ");

        var connectionId = await registered.Task;
        connectionId.ShouldNotBeNull("Manager should have registered the inbound connection");
        return new ManagerContext(manager, connectionId, remoteSocket, cts);
    }

    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();
}