natsdotnet/src/NATS.Server/LeafNodes/LeafNodeManager.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using Microsoft.Extensions.Logging;
using NATS.Server.Configuration;
using NATS.Server.Gateways;
using NATS.Server.Subscriptions;

namespace NATS.Server.LeafNodes;

/// <summary>
/// Manages leaf node connections — both inbound (accepted) and outbound (solicited).
/// Outbound connections use exponential backoff retry: 1s, 2s, 4s, ..., capped at 60s.
/// Subject filtering via DenyExports (hub→leaf) and DenyImports (leaf→hub) is applied
/// to both message forwarding and subscription propagation.
/// Go reference: leafnode.go.
/// </summary>
public sealed class LeafNodeManager : IAsyncDisposable
{
    public static readonly TimeSpan LeafNodeReconnectDelayAfterLoopDetected = TimeSpan.FromSeconds(30);
    public static readonly TimeSpan LeafNodeReconnectAfterPermViolation = TimeSpan.FromSeconds(30);
    public static readonly TimeSpan LeafNodeReconnectDelayAfterClusterNameSame = TimeSpan.FromSeconds(30);
    public static readonly TimeSpan LeafNodeWaitBeforeClose = TimeSpan.FromSeconds(5);

    private readonly LeafNodeOptions _options;
    private readonly ServerStats _stats;
    private readonly string _serverId;
    private readonly Action<RemoteSubscription> _remoteSubSink;
    private readonly Action<LeafMessage> _messageSink;
    private readonly ILogger<LeafNodeManager> _logger;
    private readonly ConcurrentDictionary<string, LeafConnection> _connections = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, bool> _disabledRemotes = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, LeafClusterInfo> _leafClusters = new(StringComparer.Ordinal);
    private readonly LeafHubSpokeMapper _subjectFilter;

    private CancellationTokenSource? _cts;
    private Socket? _listener;
    private Task? _acceptLoopTask;

    /// <summary>
    /// Invoked each time a connection is successfully registered.
    /// Exposed for test synchronization only.
    /// </summary>
    internal Action<string>? OnConnectionRegistered;

    /// <summary>
    /// Initial retry delay for solicited connections (1 second).
    /// Go reference: leafnode.go — DEFAULT_LEAF_NODE_RECONNECT constant.
    /// </summary>
    internal static readonly TimeSpan InitialRetryDelay = TimeSpan.FromSeconds(1);

    /// <summary>
    /// Maximum retry delay for solicited connections (60 seconds).
    /// </summary>
    internal static readonly TimeSpan MaxRetryDelay = TimeSpan.FromSeconds(60);

    public string ListenEndpoint => $"{_options.Host}:{_options.Port}";

    /// <summary>
    /// Current TLS certificate path, or null if TLS is not configured.
    /// Updated by <see cref="UpdateTlsConfig"/>.
    /// Go reference: leafnode.go — TLSConfig / tlsCertFile on server options.
    /// </summary>
    public string? CurrentCertPath { get; private set; }

    /// <summary>
    /// Current TLS private key path, or null if TLS is not configured.
    /// Updated by <see cref="UpdateTlsConfig"/>.
    /// </summary>
    public string? CurrentKeyPath { get; private set; }

    /// <summary>
    /// True when a TLS certificate is currently configured for leaf node connections.
    /// Go reference: leafnode.go — tls_required / tlsTimeout on LeafNodeOpts.
    /// </summary>
    public bool IsTlsEnabled => CurrentCertPath is not null;

    /// <summary>
    /// When true, all outbound leaf connections are disabled regardless of per-remote settings.
    /// Go reference: leafnode.go isLeafConnectDisabled — global disable flag.
    /// </summary>
    public bool IsGloballyDisabled { get; private set; }

    /// <summary>
    /// Returns the number of remotes that have been individually disabled via
    /// <see cref="DisableLeafConnect"/>.
    /// </summary>
    public int DisabledRemoteCount => _disabledRemotes.Count;

    /// <summary>
    /// Returns true when connections to <paramref name="remoteUrl"/> are currently disabled,
    /// either because it was individually disabled or because all leaf connections are globally
    /// disabled.
    /// Go reference: leafnode.go isLeafConnectDisabled.
    /// </summary>
    public bool IsLeafConnectDisabled(string remoteUrl)
        => IsGloballyDisabled || _disabledRemotes.ContainsKey(remoteUrl);

    /// <summary>
    /// Returns true when the remote URL is still configured and not disabled.
    /// Go reference: leafnode.go remoteLeafNodeStillValid.
    /// </summary>
    internal bool RemoteLeafNodeStillValid(string remoteUrl)
    {
        if (IsLeafConnectDisabled(remoteUrl))
            return false;

        if (_options.Remotes.Any(r => string.Equals(r, remoteUrl, StringComparison.OrdinalIgnoreCase)))
            return true;

        foreach (var remote in _options.RemoteLeaves)
        {
            if (remote.Urls.Any(u => string.Equals(u, remoteUrl, StringComparison.OrdinalIgnoreCase)))
                return true;
        }

        return false;
    }

    /// <summary>
    /// Disables outbound leaf connections to the specified remote URL.
    /// Has no effect if the remote is already disabled.
    /// Go reference: leafnode.go isLeafConnectDisabled — per-remote disable tracking.
    /// </summary>
    public void DisableLeafConnect(string remoteUrl, string? reason = null)
    {
        _disabledRemotes.TryAdd(remoteUrl, true);
        _logger.LogInformation(
            "Leaf connect disabled for remote {RemoteUrl} (reason={Reason})",
            remoteUrl, reason ?? "unspecified");
    }

    /// <summary>
    /// Re-enables outbound leaf connections to the specified remote URL.
    /// Has no effect if the remote was not disabled.
    /// </summary>
    public void EnableLeafConnect(string remoteUrl)
    {
        _disabledRemotes.TryRemove(remoteUrl, out _);
        _logger.LogInformation("Leaf connect re-enabled for remote {RemoteUrl}", remoteUrl);
    }

    /// <summary>
    /// Disables all outbound leaf connections by setting the global disable flag.
    /// Per-remote disable state is preserved.
    /// Go reference: leafnode.go isLeafConnectDisabled — global flag.
    /// </summary>
    public void DisableAllLeafConnections(string? reason = null)
    {
        IsGloballyDisabled = true;
        _logger.LogInformation("All leaf connections globally disabled (reason={Reason})", reason ?? "unspecified");
    }

    /// <summary>
    /// Clears the global disable flag so outbound leaf connections may resume.
    /// Per-remote disable state is unchanged.
    /// </summary>
    public void EnableAllLeafConnections()
    {
        IsGloballyDisabled = false;
        _logger.LogInformation("All leaf connections globally re-enabled");
    }

    /// <summary>
    /// Returns a snapshot of the remote URLs that have been individually disabled via
    /// <see cref="DisableLeafConnect"/>.
    /// </summary>
    public IReadOnlyList<string> GetDisabledRemotes() => [.. _disabledRemotes.Keys];

    /// <summary>
    /// Incremented each time <see cref="UpdateTlsConfig"/> detects a change and applies it.
    /// Useful for testing and observability.
    /// </summary>
    public int TlsReloadCount { get; private set; }

    /// <summary>
    /// Compares <paramref name="newCertPath"/> and <paramref name="newKeyPath"/> against the
    /// currently active values. When either differs the paths are updated, the reload counter
    /// is incremented, and a <see cref="LeafTlsReloadResult"/> with <c>Changed = true</c> is
    /// returned. When both are identical a result with <c>Changed = false</c> is returned and
    /// no state is mutated.
    /// Go reference: leafnode.go — reloadTLSConfig hot-reload path.
    /// </summary>
    public LeafTlsReloadResult UpdateTlsConfig(string? newCertPath, string? newKeyPath)
    {
        var previousCert = CurrentCertPath;

        if (string.Equals(CurrentCertPath, newCertPath, StringComparison.Ordinal)
            && string.Equals(CurrentKeyPath, newKeyPath, StringComparison.Ordinal))
        {
            return new LeafTlsReloadResult(Changed: false, PreviousCertPath: previousCert, NewCertPath: newCertPath, Error: null);
        }

        CurrentCertPath = newCertPath;
        CurrentKeyPath = newKeyPath;
        TlsReloadCount++;

        _logger.LogInformation(
            "Leaf node TLS config updated (cert={CertPath}, key={KeyPath}, reloads={Count})",
            newCertPath, newKeyPath, TlsReloadCount);

        return new LeafTlsReloadResult(Changed: true, PreviousCertPath: previousCert, NewCertPath: newCertPath, Error: null);
    }

    public LeafNodeManager(
        LeafNodeOptions options,
        ServerStats stats,
        string serverId,
        Action<RemoteSubscription> remoteSubSink,
        Action<LeafMessage> messageSink,
        ILogger<LeafNodeManager> logger)
    {
        _options = options;
        _stats = stats;
        _serverId = serverId;
        _remoteSubSink = remoteSubSink;
        _messageSink = messageSink;
        _logger = logger;
        _subjectFilter = new LeafHubSpokeMapper(
            new Dictionary<string, string>(),
            options.DenyExports,
            options.DenyImports,
            options.ExportSubjects,
            options.ImportSubjects);
    }

    public Task StartAsync(CancellationToken ct)
    {
        _cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        _listener = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        _listener.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
        _listener.Bind(new IPEndPoint(IPAddress.Parse(_options.Host), _options.Port));
        _listener.Listen(128);

        if (_options.Port == 0)
            _options.Port = ((IPEndPoint)_listener.LocalEndPoint!).Port;

        _acceptLoopTask = Task.Run(() => AcceptLoopAsync(_cts.Token));
        foreach (var remote in _options.Remotes.Distinct(StringComparer.OrdinalIgnoreCase))
            _ = Task.Run(() => ConnectSolicitedWithRetryAsync(remote, _options.JetStreamDomain, _cts.Token));

        // Also start solicited connections for remotes parsed from the config file (RemoteLeaves).
        // RemoteLeaves are populated by the config parser from leafnodes.remotes[] blocks;
        // _options.Remotes is the simple programmatic list only.
        // Go reference: leafnode.go — createLeafNode starts solicited connections via connectToRemoteLeaf.
        foreach (var remoteLeaf in _options.RemoteLeaves)
        {
            foreach (var url in remoteLeaf.Urls.Distinct(StringComparer.OrdinalIgnoreCase))
                _ = Task.Run(() => ConnectSolicitedWithRetryAsync(url, _options.JetStreamDomain, _cts.Token));
        }

        _logger.LogDebug("Leaf manager started (listen={Host}:{Port})", _options.Host, _options.Port);
        return Task.CompletedTask;
    }

    /// <summary>
    /// Establishes a single solicited (outbound) leaf connection to the specified URL.
    /// Performs socket connection and LEAF handshake. If a JetStream domain is specified,
    /// it is propagated during the handshake.
    /// Go reference: leafnode.go — connectSolicited.
    /// </summary>
    public async Task<LeafConnection> ConnectSolicitedAsync(string url, string? account, CancellationToken ct)
    {
        var endPoint = ParseEndpoint(url);
        var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        try
        {
            await socket.ConnectAsync(endPoint.Address, endPoint.Port, ct);
            var connection = new LeafConnection(socket)
            {
                JetStreamDomain = _options.JetStreamDomain,
                IsSolicited = true,
                IsSpoke = true,
            };
            await connection.PerformOutboundHandshakeAsync(_serverId, ct);
            Register(connection);
            _logger.LogDebug("Solicited leaf connection established to {Url} (account={Account})", url, account ?? "$G");
            return connection;
        }
        catch
        {
            socket.Dispose();
            throw;
        }
    }

    public async Task ForwardMessageAsync(string account, string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        // Apply subject filtering: outbound direction is hub→leaf (DenyExports).
        // The subject may be loop-marked ($LDS.{serverId}.{realSubject}), so we
        // strip the marker before checking the filter against the logical subject.
        // Go reference: leafnode.go:475-478 (DenyExports → Publish deny list).
        var filterSubject = LeafLoopDetector.TryUnmark(subject, out var unmarked) ? unmarked : subject;
        if (!_subjectFilter.IsSubjectAllowed(filterSubject, LeafMapDirection.Outbound))
        {
            _logger.LogDebug("Leaf outbound message denied for subject {Subject} (DenyExports)", filterSubject);
            return;
        }

        foreach (var connection in _connections.Values)
            await connection.SendMessageAsync(account, subject, replyTo, payload, ct);
    }

    public void PropagateLocalSubscription(string account, string subject, string? queue)
        => PropagateLocalSubscription(account, subject, queue, queueWeight: 0);

    public void PropagateLocalSubscription(string account, string subject, string? queue, int queueWeight)
    {
        // Subscription propagation is also subject to export filtering:
        // we don't propagate subscriptions for subjects that are denied.
        if (!_subjectFilter.IsSubjectAllowed(subject, LeafMapDirection.Outbound))
        {
            _logger.LogDebug("Leaf subscription propagation denied for subject {Subject} (DenyExports)", subject);
            return;
        }

        foreach (var connection in _connections.Values)
        {
            if (!CanSpokeSendSubscription(connection, subject))
            {
                _logger.LogDebug(
                    "Leaf subscription propagation denied for spoke connection {RemoteId} and subject {Subject} (subscribe permissions)",
                    connection.RemoteId ?? "<unknown>",
                    subject);
                continue;
            }

            _ = connection.SendLsPlusAsync(account, subject, queue, queueWeight, _cts?.Token ?? CancellationToken.None);
        }
    }

    public void PropagateLocalUnsubscription(string account, string subject, string? queue)
    {
        foreach (var connection in _connections.Values)
            _ = connection.SendLsMinusAsync(account, subject, queue, _cts?.Token ?? CancellationToken.None);
    }

    /// <summary>
    /// Sends permission and account information to the specified leaf connection.
    /// Finds the connection by ID, sets the account name, and applies the publish/subscribe
    /// allow lists.  Returns a result describing whether the connection was found and its
    /// post-sync state.
    /// Go reference: leafnode.go — sendPermsAndAccountInfo.
    /// </summary>
    public LeafPermSyncResult SendPermsAndAccountInfo(
        string connectionId,
        string? account,
        IEnumerable<string>? pubAllow,
        IEnumerable<string>? subAllow)
    {
        if (!_connections.TryGetValue(connectionId, out var connection))
            return new LeafPermSyncResult(Found: false, PermsSynced: false, AccountName: null, PublishAllowCount: 0, SubscribeAllowCount: 0);

        connection.AccountName = account;
        connection.SetPermissions(pubAllow, subAllow);

        _logger.LogDebug(
            "Leaf perms synced for connection {ConnectionId} (account={Account}, pubAllow={PubCount}, subAllow={SubCount})",
            connectionId, account ?? "$G",
            connection.AllowedPublishSubjects.Count,
            connection.AllowedSubscribeSubjects.Count);

        return new LeafPermSyncResult(
            Found: true,
            PermsSynced: connection.PermsSynced,
            AccountName: connection.AccountName,
            PublishAllowCount: connection.AllowedPublishSubjects.Count,
            SubscribeAllowCount: connection.AllowedSubscribeSubjects.Count);
    }

    /// <summary>
    /// Initialises the subject map for a leaf connection from the supplied set of subjects
    /// and returns the count of subjects sent.  This mirrors the Go server's subject-map
    /// (smap) seeding that happens when a leaf first connects so the remote side knows
    /// which subjects have local interest.
    /// Go reference: leafnode.go — initLeafNodeSmapAndSendSubs.
    /// </summary>
    public int InitLeafNodeSmapAndSendSubs(string connectionId, IEnumerable<string> subjects)
    {
        if (!_connections.TryGetValue(connectionId, out var connection))
            return 0;

        var ct = _cts?.Token ?? CancellationToken.None;
        var count = 0;

        foreach (var subject in subjects)
        {
            _ = connection.SendLsPlusAsync("$G", subject, null, ct);
            count++;
        }

        _logger.LogDebug("Leaf smap initialised for connection {ConnectionId}: {Count} subjects sent", connectionId, count);
        return count;
    }

    /// <summary>
    /// Returns the current permission-sync status for the specified connection.
    /// Go reference: leafnode.go — sendPermsAndAccountInfo (read path).
    /// </summary>
    public LeafPermSyncResult GetPermSyncStatus(string connectionId)
    {
        if (!_connections.TryGetValue(connectionId, out var connection))
            return new LeafPermSyncResult(Found: false, PermsSynced: false, AccountName: null, PublishAllowCount: 0, SubscribeAllowCount: 0);

        return new LeafPermSyncResult(
            Found: true,
            PermsSynced: connection.PermsSynced,
            AccountName: connection.AccountName,
            PublishAllowCount: connection.AllowedPublishSubjects.Count,
            SubscribeAllowCount: connection.AllowedSubscribeSubjects.Count);
    }

    /// <summary>
    /// Validates whether a leaf connection can migrate its JetStream domain to a proposed value.
    /// Clearing the domain (null/empty proposedDomain) is always valid.
    /// If the proposed domain matches the current domain, no migration is needed.
    /// If another connection already uses the proposed domain, a conflict is reported.
    /// Go reference: leafnode.go checkJetStreamMigrate.
    /// </summary>
    public JetStreamMigrationResult CheckJetStreamMigrate(string connectionId, string? proposedDomain)
    {
        if (!_connections.TryGetValue(connectionId, out var connection))
            return new JetStreamMigrationResult(false, JetStreamMigrationStatus.ConnectionNotFound, $"Connection '{connectionId}' not found");

        // Clearing domain is always valid.
        if (string.IsNullOrEmpty(proposedDomain))
            return new JetStreamMigrationResult(true, JetStreamMigrationStatus.Valid, null);

        // If current domain already matches, no migration needed.
        if (string.Equals(connection.JetStreamDomain, proposedDomain, StringComparison.Ordinal))
            return new JetStreamMigrationResult(true, JetStreamMigrationStatus.NoChangeNeeded, null);

        // Check for domain conflict with other connections.
        foreach (var (key, conn) in _connections)
        {
            if (string.Equals(key, connectionId, StringComparison.Ordinal))
                continue;
            if (string.Equals(conn.JetStreamDomain, proposedDomain, StringComparison.Ordinal))
                return new JetStreamMigrationResult(false, JetStreamMigrationStatus.DomainConflict,
                    $"Domain '{proposedDomain}' is already in use by another connection");
        }

        return new JetStreamMigrationResult(true, JetStreamMigrationStatus.Valid, null);
    }

    /// <summary>
    /// Returns the distinct set of JetStream domains across all active connections.
    /// Connections without a domain (null/empty JetStreamDomain) are excluded.
    /// Go reference: leafnode.go — per-connection domain tracking.
    /// </summary>
    public IReadOnlyList<string> GetActiveJetStreamDomains()
    {
        var domains = new HashSet<string>(StringComparer.Ordinal);
        foreach (var conn in _connections.Values)
        {
            if (!string.IsNullOrEmpty(conn.JetStreamDomain))
                domains.Add(conn.JetStreamDomain);
        }
        return [.. domains];
    }

    /// <summary>
    /// Returns true if any currently active connection is associated with the specified JetStream domain.
    /// Go reference: leafnode.go — checkJetStreamMigrate domain-in-use check.
    /// </summary>
    public bool IsJetStreamDomainInUse(string domain)
    {
        foreach (var conn in _connections.Values)
        {
            if (string.Equals(conn.JetStreamDomain, domain, StringComparison.Ordinal))
                return true;
        }
        return false;
    }

    /// <summary>
    /// Count of connections that have a non-null, non-empty JetStream domain assigned.
    /// Go reference: leafnode.go — per-connection jsClusterDomain field.
    /// </summary>
    public int JetStreamEnabledConnectionCount
    {
        get
        {
            var count = 0;
            foreach (var conn in _connections.Values)
            {
                if (!string.IsNullOrEmpty(conn.JetStreamDomain))
                    count++;
            }
            return count;
        }
    }

    /// <summary>
    /// Registers a leaf cluster topology entry.
    /// Returns false if a cluster with the same name is already registered.
    /// Go reference: leafnode.go registerLeafNodeCluster.
    /// </summary>
    public bool RegisterLeafNodeCluster(string clusterName, string gatewayUrl, int connectionCount)
    {
        var info = new LeafClusterInfo
        {
            ClusterName = clusterName,
            GatewayUrl = gatewayUrl,
            ConnectionCount = connectionCount,
        };
        return _leafClusters.TryAdd(clusterName, info);
    }

    /// <summary>
    /// Removes a leaf cluster entry by name.
    /// Returns false if no entry with that name exists.
    /// Go reference: leafnode.go — leaf cluster topology removal.
    /// </summary>
    public bool UnregisterLeafNodeCluster(string clusterName) =>
        _leafClusters.TryRemove(clusterName, out _);

    /// <summary>
    /// Returns true if a leaf cluster with the given name is currently registered.
    /// Go reference: leafnode.go — leaf cluster topology lookup.
    /// </summary>
    public bool HasLeafNodeCluster(string clusterName) =>
        _leafClusters.ContainsKey(clusterName);

    /// <summary>
    /// Returns the <see cref="LeafClusterInfo"/> for the named cluster, or null if not registered.
    /// Go reference: leafnode.go — leaf cluster topology lookup.
    /// </summary>
    public LeafClusterInfo? GetLeafNodeCluster(string clusterName) =>
        _leafClusters.TryGetValue(clusterName, out var info) ? info : null;

    /// <summary>
    /// Returns all registered leaf cluster entries as a read-only list.
    /// Go reference: leafnode.go — leaf cluster topology enumeration.
    /// </summary>
    public IReadOnlyList<LeafClusterInfo> GetAllLeafClusters() =>
        [.. _leafClusters.Values];

    /// <summary>
    /// Count of registered leaf cluster topology entries.
    /// Go reference: leafnode.go — leaf cluster topology count.
    /// </summary>
    public int LeafClusterCount => _leafClusters.Count;

    /// <summary>
    /// Updates the connection count for the named leaf cluster.
    /// No-op if the cluster is not registered.
    /// Go reference: leafnode.go — leaf cluster connection count update.
    /// </summary>
    public void UpdateLeafClusterConnectionCount(string clusterName, int newCount)
    {
        if (_leafClusters.TryGetValue(clusterName, out var info))
            info.ConnectionCount = newCount;
    }

    /// <summary>
    /// Returns all current connection IDs. Useful for tests and monitoring.
    /// </summary>
    internal IReadOnlyCollection<string> GetConnectionIds() => _connections.Keys.ToArray();

    /// <summary>
    /// Injects a <see cref="LeafConnection"/> directly into the tracked connections.
    /// For testing only — bypasses the normal handshake and registration path.
    /// </summary>
    internal void InjectConnectionForTesting(LeafConnection connection)
    {
        var key = $"{connection.RemoteId}:{connection.RemoteEndpoint}:{Guid.NewGuid():N}";
        _connections.TryAdd(key, connection);
    }

    public async ValueTask DisposeAsync()
    {
        if (_cts == null)
            return;

        await _cts.CancelAsync();
        _listener?.Dispose();
        if (_acceptLoopTask != null)
            await _acceptLoopTask.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);

        foreach (var connection in _connections.Values)
            await connection.DisposeAsync();
        _connections.Clear();
        Interlocked.Exchange(ref _stats.Leafs, 0);
        _cts.Dispose();
        _cts = null;
        _logger.LogDebug("Leaf manager stopped");
    }

    /// <summary>
    /// Computes the next backoff delay using exponential backoff with a cap.
    /// Delay sequence: 1s, 2s, 4s, 8s, 16s, 32s, 60s, 60s, ...
    /// </summary>
    internal static TimeSpan ComputeBackoff(int attempt)
    {
        if (attempt < 0) attempt = 0;
        var seconds = Math.Min(InitialRetryDelay.TotalSeconds * Math.Pow(2, attempt), MaxRetryDelay.TotalSeconds);
        return TimeSpan.FromSeconds(seconds);
    }

    private async Task AcceptLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            Socket socket;
            try
            {
                socket = await _listener!.AcceptAsync(ct);
            }
            catch
            {
                break;
            }

            _ = Task.Run(() => HandleInboundAsync(socket, ct), ct);
        }
    }

    private async Task HandleInboundAsync(Socket socket, CancellationToken ct)
    {
        var connection = new LeafConnection(socket)
        {
            JetStreamDomain = _options.JetStreamDomain,
        };
        try
        {
            await connection.PerformInboundHandshakeAsync(_serverId, ct);
            Register(connection);
        }
        catch
        {
            await connection.DisposeAsync();
        }
    }

    private async Task ConnectSolicitedWithRetryAsync(string remote, string? jetStreamDomain, CancellationToken ct)
    {
        var attempt = 0;
        while (!ct.IsCancellationRequested)
        {
            if (!RemoteLeafNodeStillValid(remote))
                return;

            try
            {
                var endPoint = ParseEndpoint(remote);
                var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
                try
                {
                    await socket.ConnectAsync(endPoint.Address, endPoint.Port, ct);
                    var connection = new LeafConnection(socket)
                    {
                        JetStreamDomain = jetStreamDomain,
                        IsSolicited = true,
                        IsSpoke = true,
                    };
                    await connection.PerformOutboundHandshakeAsync(_serverId, ct);
                    Register(connection);
                    _logger.LogDebug("Solicited leaf connection established to {Remote}", remote);
                    return;
                }
                catch
                {
                    socket.Dispose();
                    throw;
                }
            }
            catch (OperationCanceledException)
            {
                return;
            }
            catch (Exception ex)
            {
                _logger.LogDebug(ex, "Leaf connect retry for {Remote} (attempt {Attempt})", remote, attempt);
            }

            var delay = ComputeBackoff(attempt);
            attempt++;
            try
            {
                await Task.Delay(delay, ct);
            }
            catch (OperationCanceledException)
            {
                return;
            }
        }
    }

    private void Register(LeafConnection connection)
    {
        var key = $"{connection.RemoteId}:{connection.RemoteEndpoint}:{Guid.NewGuid():N}";
        if (!_connections.TryAdd(key, connection))
        {
            _ = connection.DisposeAsync();
            return;
        }

        OnConnectionRegistered?.Invoke(key);

        connection.RemoteSubscriptionReceived = sub =>
        {
            _remoteSubSink(sub);
            return Task.CompletedTask;
        };
        connection.MessageReceived = msg =>
        {
            // Apply inbound filtering: DenyImports restricts leaf→hub messages.
            // The subject may be loop-marked ($LDS.{serverId}.{realSubject}), so we
            // strip the marker before checking the filter against the logical subject.
            // Go reference: leafnode.go:480-481 (DenyImports → Subscribe deny list).
            var filterSubject = LeafLoopDetector.TryUnmark(msg.Subject, out var unmarked)
                ? unmarked
                : msg.Subject;
            if (!_subjectFilter.IsSubjectAllowed(filterSubject, LeafMapDirection.Inbound))
            {
                _logger.LogDebug("Leaf inbound message denied for subject {Subject} (DenyImports)", filterSubject);
                return Task.CompletedTask;
            }

            _messageSink(msg);
            return Task.CompletedTask;
        };
        connection.StartLoop(_cts!.Token);
        Interlocked.Increment(ref _stats.Leafs);
        _ = Task.Run(() => WatchConnectionAsync(key, connection, _cts!.Token));
    }

    private async Task WatchConnectionAsync(string key, LeafConnection connection, CancellationToken ct)
    {
        try
        {
            await connection.WaitUntilClosedAsync(ct);
        }
        catch
        {
        }
        finally
        {
            if (_connections.TryRemove(key, out _))
                Interlocked.Decrement(ref _stats.Leafs);
            await connection.DisposeAsync();
        }
    }

    /// <summary>
    /// Validates a reconnecting leaf node against the current server state.
    /// Checks for self-connect, duplicate connections, and JetStream domain conflicts.
    /// Go reference: leafnode.go addLeafNodeConnection — duplicate and domain checks.
    /// </summary>
    public LeafValidationResult ValidateRemoteLeafNode(string remoteId, string? account, string? jsDomain)
    {
        if (IsSelfConnect(remoteId))
            return new LeafValidationResult(false, $"Self-connect detected: remoteId '{remoteId}' matches local server ID", LeafValidationError.SelfConnect);

        var existing = GetConnectionByRemoteId(remoteId);
        if (existing != null)
            return new LeafValidationResult(false, $"Duplicate connection: a connection from '{remoteId}' already exists", LeafValidationError.DuplicateConnection);

        if (!string.IsNullOrEmpty(jsDomain))
        {
            foreach (var conn in _connections.Values)
            {
                if (!string.IsNullOrEmpty(conn.JetStreamDomain) &&
                    !string.Equals(conn.JetStreamDomain, jsDomain, StringComparison.Ordinal))
                    return new LeafValidationResult(false, $"JetStream domain conflict: incoming domain '{jsDomain}' conflicts with existing domain '{conn.JetStreamDomain}'", LeafValidationError.JetStreamDomainConflict);
            }
        }

        return new LeafValidationResult(true, null, LeafValidationError.None);
    }

    /// <summary>
    /// Returns true if the given remoteId matches this server's own ID (self-connect detection).
    /// Go reference: leafnode.go loop detection via server ID comparison.
    /// </summary>
    public bool IsSelfConnect(string remoteId) => string.Equals(remoteId, _serverId, StringComparison.Ordinal);

    /// <summary>
    /// Returns true if any currently registered connection has the specified remote server ID.
    /// </summary>
    public bool HasConnection(string remoteId) => GetConnectionByRemoteId(remoteId) != null;

    /// <summary>
    /// Returns the first registered connection whose RemoteId matches the given value, or null if none.
    /// </summary>
    public LeafConnection? GetConnectionByRemoteId(string remoteId)
    {
        foreach (var conn in _connections.Values)
        {
            if (string.Equals(conn.RemoteId, remoteId, StringComparison.Ordinal))
                return conn;
        }
        return null;
    }

    private static bool CanSpokeSendSubscription(LeafConnection connection, string subject)
    {
        if (!connection.IsSpokeLeafNode())
            return true;

        if (ShouldBypassSpokeSubscribePermission(subject))
            return true;

        if (!connection.PermsSynced || connection.AllowedSubscribeSubjects.Count == 0)
            return true;

        for (var i = 0; i < connection.AllowedSubscribeSubjects.Count; i++)
        {
            if (SubjectMatch.MatchLiteral(subject, connection.AllowedSubscribeSubjects[i]))
                return true;
        }

        return false;
    }

    private static bool ShouldBypassSpokeSubscribePermission(string subject)
    {
        if (string.IsNullOrEmpty(subject))
            return false;

        if (subject[0] != '$' && subject[0] != '_')
            return false;

        return subject.StartsWith("$LDS.", StringComparison.Ordinal)
            || subject.StartsWith(ReplyMapper.GatewayReplyPrefix, StringComparison.Ordinal)
            || subject.StartsWith(ReplyMapper.OldGatewayReplyPrefix, StringComparison.Ordinal);
    }

    private static IPEndPoint ParseEndpoint(string endpoint)
    {
        // Handle full URLs with a scheme (e.g. "nats-leaf://127.0.0.1:5222").
        // Uri.TryCreate handles both schemed URLs and bare "host:port" strings.
        if (Uri.TryCreate(endpoint, UriKind.Absolute, out var uri))
            return new IPEndPoint(IPAddress.Parse(uri.Host), uri.Port);

        // Fall back to bare "host:port" splitting for plain strings without a scheme.
        var parts = endpoint.Split(':', 2, StringSplitOptions.TrimEntries | StringSplitOptions.RemoveEmptyEntries);
        if (parts.Length != 2)
            throw new FormatException($"Invalid endpoint: {endpoint}");

        return new IPEndPoint(IPAddress.Parse(parts[0]), int.Parse(parts[1]));
    }
}

/// <summary>
/// Describes the outcome of a <see cref="LeafNodeManager.SendPermsAndAccountInfo"/> or
/// <see cref="LeafNodeManager.GetPermSyncStatus"/> call.
/// </summary>
/// <param name="Found">True when the connection ID was found in the manager.</param>
/// <param name="PermsSynced">True when <see cref="LeafConnection.SetPermissions"/> has been called at least once.</param>
/// <param name="AccountName">The account name associated with the connection, or null if not set.</param>
/// <param name="PublishAllowCount">Number of allowed publish subjects currently configured.</param>
/// <param name="SubscribeAllowCount">Number of allowed subscribe subjects currently configured.</param>
public sealed record LeafPermSyncResult(
    bool Found,
    bool PermsSynced,
    string? AccountName,
    int PublishAllowCount,
    int SubscribeAllowCount);

/// <summary>
/// Describes the outcome of a <see cref="LeafNodeManager.UpdateTlsConfig"/> call.
/// </summary>
/// <param name="Changed">True when the cert or key path differed from the previously active values.</param>
/// <param name="PreviousCertPath">The cert path that was active before this call.</param>
/// <param name="NewCertPath">The cert path supplied to this call.</param>
/// <param name="Error">Non-null when an error prevented the reload (reserved for future use).</param>
public sealed record LeafTlsReloadResult(
    bool Changed,
    string? PreviousCertPath,
    string? NewCertPath,
    string? Error);

/// <summary>
/// Result of validating a reconnecting leaf node.
/// Go reference: leafnode.go addLeafNodeConnection validation logic.
/// </summary>
public sealed record LeafValidationResult(
    bool Valid,
    string? Error,
    LeafValidationError ErrorCode);

/// <summary>
/// Error codes for leaf node validation failures.
/// </summary>
public enum LeafValidationError
{
    None,
    SelfConnect,
    DuplicateConnection,
    JetStreamDomainConflict
}

/// <summary>
/// Describes the outcome of a <see cref="LeafNodeManager.CheckJetStreamMigrate"/> call.
/// </summary>
/// <param name="Valid">True when migration to the proposed domain is allowed.</param>
/// <param name="Status">Detailed status code for the migration check.</param>
/// <param name="Error">Human-readable error message when <paramref name="Valid"/> is false, otherwise null.</param>
public sealed record JetStreamMigrationResult(
    bool Valid,
    JetStreamMigrationStatus Status,
    string? Error);

/// <summary>
/// Status codes for <see cref="JetStreamMigrationResult"/>.
/// Go reference: leafnode.go checkJetStreamMigrate return values.
/// </summary>
public enum JetStreamMigrationStatus
{
    /// <summary>Migration to the proposed domain is allowed.</summary>
    Valid,
    /// <summary>The specified connection ID was not found.</summary>
    ConnectionNotFound,
    /// <summary>The proposed domain is identical to the current domain — no migration required.</summary>
    NoChangeNeeded,
    /// <summary>Another connection already uses the proposed domain.</summary>
    DomainConflict
}

/// <summary>
/// Holds topology information for a registered leaf cluster entry.
/// Go reference: leafnode.go — leaf cluster registration / registerLeafNodeCluster.
/// </summary>
public sealed class LeafClusterInfo
{
    public required string ClusterName { get; init; }
    public required string GatewayUrl { get; init; }
    public int ConnectionCount { get; set; }
    public DateTime RegisteredAt { get; init; } = DateTime.UtcNow;
}