natsdotnet/tests/NATS.Server.LeafNodes.Tests/LeafNodes/LeafNodeManagerParityBatch5Tests.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.LeafNodes.Tests.LeafNodes;

public class LeafNodeManagerParityBatch5Tests
{
    [Fact]
    [SlopwatchSuppress("SW004", "Delay verifies a blocked subject is NOT forwarded; absence of a frame cannot be observed via synchronization primitives")]
    public async Task PropagateLocalSubscription_enforces_spoke_subscribe_permissions_and_keeps_queue_weight()
    {
        await using var ctx = await CreateManagerWithInboundConnectionAsync();
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        var conn = ctx.Manager.GetConnectionByRemoteId("SPOKE1");
        conn.ShouldNotBeNull();
        conn!.IsSpoke = true;

        var sync = ctx.Manager.SendPermsAndAccountInfo(
            ctx.ConnectionId,
            "$G",
            pubAllow: null,
            subAllow: ["allowed.>"]);
        sync.Found.ShouldBeTrue();
        sync.PermsSynced.ShouldBeTrue();

        ctx.Manager.PropagateLocalSubscription("$G", "blocked.data", null);
        ctx.Manager.PropagateLocalSubscription("$G", "allowed.data", "workers", queueWeight: 4);

        // Only the allowed subject should appear on the wire; the blocked one is filtered synchronously.
        var line = await ReadLineAsync(ctx.RemoteSocket, timeout.Token);
        line.ShouldBe("LS+ $G allowed.data workers 4");
    }

    [Fact]
    public async Task PropagateLocalSubscription_allows_loop_and_gateway_reply_prefixes_for_spoke()
    {
        await using var ctx = await CreateManagerWithInboundConnectionAsync();
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        var conn = ctx.Manager.GetConnectionByRemoteId("SPOKE1");
        conn.ShouldNotBeNull();
        conn!.IsSpoke = true;

        ctx.Manager.SendPermsAndAccountInfo(
            ctx.ConnectionId,
            "$G",
            pubAllow: null,
            subAllow: ["allowed.>"]);

        ctx.Manager.PropagateLocalSubscription("$G", "$LDS.HUB.loop", null);
        (await ReadLineAsync(ctx.RemoteSocket, timeout.Token)).ShouldBe("LS+ $G $LDS.HUB.loop");

        ctx.Manager.PropagateLocalSubscription("$G", "_GR_.A.reply", null);
        (await ReadLineAsync(ctx.RemoteSocket, timeout.Token)).ShouldBe("LS+ $G _GR_.A.reply");
    }

    private sealed class ManagerContext : IAsyncDisposable
    {
        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 Socket RemoteSocket { get; }

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

    private static async Task<ManagerContext> CreateManagerWithInboundConnectionAsync()
    {
        var options = new LeafNodeOptions { Host = "127.0.0.1", Port = 0 };
        var manager = new LeafNodeManager(
            options,
            new ServerStats(),
            "HUB",
            _ => { },
            _ => { },
            NullLogger<LeafNodeManager>.Instance);

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

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

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var registered = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
        manager.OnConnectionRegistered = id => registered.TrySetResult(id);
        timeout.Token.Register(() => registered.TrySetCanceled(timeout.Token));

        await WriteLineAsync(remoteSocket, "LEAF SPOKE1", timeout.Token);
        (await ReadLineAsync(remoteSocket, timeout.Token)).ShouldStartWith("LEAF ");

        var connectionId = await registered.Task;
        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)
                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();
}