natsdotnet/src/NATS.Server/NatsClient.cs

using System.Buffers;
using System.IO.Pipelines;
using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using Microsoft.Extensions.Logging;
using NATS.Server.Auth;
using NATS.Server.Protocol;
using NATS.Server.Subscriptions;
using NATS.Server.Tls;

namespace NATS.Server;

public interface IMessageRouter
{
    void ProcessMessage(string subject, string? replyTo, ReadOnlyMemory<byte> headers,
        ReadOnlyMemory<byte> payload, NatsClient sender);
    void RemoveClient(NatsClient client);
}

public interface ISubListAccess
{
    SubList SubList { get; }
}

public sealed class NatsClient : IDisposable
{
    private readonly Socket _socket;
    private readonly Stream _stream;
    private readonly NatsOptions _options;
    private readonly ServerInfo _serverInfo;
    private readonly AuthService _authService;
    private readonly byte[]? _nonce;
    private readonly NatsParser _parser;
    private readonly SemaphoreSlim _writeLock = new(1, 1);
    private CancellationTokenSource? _clientCts;
    private readonly Dictionary<string, Subscription> _subs = new();
    private readonly ILogger _logger;
    private ClientPermissions? _permissions;
    private readonly ServerStats _serverStats;

    public ulong Id { get; }
    public ClientOptions? ClientOpts { get; private set; }
    public IMessageRouter? Router { get; set; }
    public Account? Account { get; private set; }

    // Thread-safe: read from auth timeout task on threadpool, written from command pipeline
    private int _connectReceived;
    public bool ConnectReceived => Volatile.Read(ref _connectReceived) != 0;

    public DateTime StartTime { get; }
    private long _lastActivityTicks;
    public DateTime LastActivity => new(Interlocked.Read(ref _lastActivityTicks), DateTimeKind.Utc);
    public string? RemoteIp { get; }
    public int RemotePort { get; }

    // Stats
    public long InMsgs;
    public long OutMsgs;
    public long InBytes;
    public long OutBytes;

    // PING keepalive state
    private int _pingsOut;
    private long _lastIn;

    public TlsConnectionState? TlsState { get; set; }
    public bool InfoAlreadySent { get; set; }

    public IReadOnlyDictionary<string, Subscription> Subscriptions => _subs;

    public NatsClient(ulong id, Stream stream, Socket socket, NatsOptions options, ServerInfo serverInfo,
        AuthService authService, byte[]? nonce, ILogger logger, ServerStats serverStats)
    {
        Id = id;
        _socket = socket;
        _stream = stream;
        _options = options;
        _serverInfo = serverInfo;
        _authService = authService;
        _nonce = nonce;
        _logger = logger;
        _serverStats = serverStats;
        _parser = new NatsParser(options.MaxPayload);
        StartTime = DateTime.UtcNow;
        _lastActivityTicks = StartTime.Ticks;
        if (socket.RemoteEndPoint is IPEndPoint ep)
        {
            RemoteIp = ep.Address.ToString();
            RemotePort = ep.Port;
        }
    }

    public async Task RunAsync(CancellationToken ct)
    {
        _clientCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        Interlocked.Exchange(ref _lastIn, Environment.TickCount64);
        var pipe = new Pipe();
        try
        {
            // Send INFO (skip if already sent during TLS negotiation)
            if (!InfoAlreadySent)
                await SendInfoAsync(_clientCts.Token);

            // Start auth timeout if auth is required
            Task? authTimeoutTask = null;
            if (_authService.IsAuthRequired)
            {
                authTimeoutTask = Task.Run(async () =>
                {
                    try
                    {
                        await Task.Delay(_options.AuthTimeout, _clientCts!.Token);
                        if (!ConnectReceived)
                        {
                            _logger.LogDebug("Client {ClientId} auth timeout", Id);
                            await SendErrAndCloseAsync(NatsProtocol.ErrAuthTimeout);
                        }
                    }
                    catch (OperationCanceledException)
                    {
                        // Normal -- client connected or was cancelled
                    }
                }, _clientCts.Token);
            }

            // Start read pump, command processing, and ping timer in parallel
            var fillTask = FillPipeAsync(pipe.Writer, _clientCts.Token);
            var processTask = ProcessCommandsAsync(pipe.Reader, _clientCts.Token);
            var pingTask = RunPingTimerAsync(_clientCts.Token);

            if (authTimeoutTask != null)
                await Task.WhenAny(fillTask, processTask, pingTask, authTimeoutTask);
            else
                await Task.WhenAny(fillTask, processTask, pingTask);
        }
        catch (OperationCanceledException)
        {
            _logger.LogDebug("Client {ClientId} operation cancelled", Id);
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "Client {ClientId} connection error", Id);
        }
        finally
        {
            try { _socket.Shutdown(SocketShutdown.Both); }
            catch (SocketException) { }
            catch (ObjectDisposedException) { }
            Router?.RemoveClient(this);
        }
    }

    private async Task FillPipeAsync(PipeWriter writer, CancellationToken ct)
    {
        try
        {
            while (!ct.IsCancellationRequested)
            {
                var memory = writer.GetMemory(4096);
                int bytesRead = await _stream.ReadAsync(memory, ct);
                if (bytesRead == 0)
                    break;

                writer.Advance(bytesRead);
                var result = await writer.FlushAsync(ct);
                if (result.IsCompleted)
                    break;
            }
        }
        finally
        {
            await writer.CompleteAsync();
        }
    }

    private async Task ProcessCommandsAsync(PipeReader reader, CancellationToken ct)
    {
        try
        {
            while (!ct.IsCancellationRequested)
            {
                var result = await reader.ReadAsync(ct);
                var buffer = result.Buffer;

                while (_parser.TryParse(ref buffer, out var cmd))
                {
                    Interlocked.Exchange(ref _lastIn, Environment.TickCount64);
                    await DispatchCommandAsync(cmd, ct);
                }

                reader.AdvanceTo(buffer.Start, buffer.End);

                if (result.IsCompleted)
                    break;
            }
        }
        finally
        {
            await reader.CompleteAsync();
        }
    }

    private async ValueTask DispatchCommandAsync(ParsedCommand cmd, CancellationToken ct)
    {
        Interlocked.Exchange(ref _lastActivityTicks, DateTime.UtcNow.Ticks);

        // If auth is required and CONNECT hasn't been received yet,
        // only allow CONNECT and PING commands
        if (_authService.IsAuthRequired && !ConnectReceived)
        {
            switch (cmd.Type)
            {
                case CommandType.Connect:
                    await ProcessConnectAsync(cmd);
                    return;
                case CommandType.Ping:
                    await WriteAsync(NatsProtocol.PongBytes, ct);
                    return;
                default:
                    // Ignore all other commands until authenticated
                    return;
            }
        }

        switch (cmd.Type)
        {
            case CommandType.Connect:
                await ProcessConnectAsync(cmd);
                break;

            case CommandType.Ping:
                await WriteAsync(NatsProtocol.PongBytes, ct);
                break;

            case CommandType.Pong:
                Interlocked.Exchange(ref _pingsOut, 0);
                break;

            case CommandType.Sub:
                await ProcessSubAsync(cmd);
                break;

            case CommandType.Unsub:
                ProcessUnsub(cmd);
                break;

            case CommandType.Pub:
            case CommandType.HPub:
                await ProcessPubAsync(cmd);
                break;
        }
    }

    private async ValueTask ProcessConnectAsync(ParsedCommand cmd)
    {
        ClientOpts = JsonSerializer.Deserialize<ClientOptions>(cmd.Payload.Span)
                     ?? new ClientOptions();

        // Authenticate if auth is required
        if (_authService.IsAuthRequired)
        {
            var context = new ClientAuthContext
            {
                Opts = ClientOpts,
                Nonce = _nonce ?? [],
            };

            var result = _authService.Authenticate(context);
            if (result == null)
            {
                _logger.LogWarning("Client {ClientId} authentication failed", Id);
                await SendErrAndCloseAsync(NatsProtocol.ErrAuthorizationViolation);
                return;
            }

            // Build permissions from auth result
            _permissions = ClientPermissions.Build(result.Permissions);

            // Resolve account
            if (Router is NatsServer server)
            {
                var accountName = result.AccountName ?? Account.GlobalAccountName;
                Account = server.GetOrCreateAccount(accountName);
                Account.AddClient(Id);
            }

            _logger.LogDebug("Client {ClientId} authenticated as {Identity}", Id, result.Identity);

            // Clear nonce after use -- defense-in-depth against memory dumps
            if (_nonce != null)
                CryptographicOperations.ZeroMemory(_nonce);
        }

        // If no account was assigned by auth, assign global account
        if (Account == null && Router is NatsServer server2)
        {
            Account = server2.GetOrCreateAccount(Account.GlobalAccountName);
            Account.AddClient(Id);
        }

        Volatile.Write(ref _connectReceived, 1);
        _logger.LogDebug("CONNECT received from client {ClientId}, name={ClientName}", Id, ClientOpts?.Name);
    }

    private async ValueTask ProcessSubAsync(ParsedCommand cmd)
    {
        // Permission check for subscribe
        if (_permissions != null && !_permissions.IsSubscribeAllowed(cmd.Subject!, cmd.Queue))
        {
            _logger.LogDebug("Client {ClientId} subscribe permission denied for {Subject}", Id, cmd.Subject);
            await SendErrAsync(NatsProtocol.ErrPermissionsSubscribe);
            return;
        }

        var sub = new Subscription
        {
            Subject = cmd.Subject!,
            Queue = cmd.Queue,
            Sid = cmd.Sid!,
        };

        _subs[cmd.Sid!] = sub;
        sub.Client = this;

        _logger.LogDebug("SUB {Subject} {Sid} from client {ClientId}", cmd.Subject, cmd.Sid, Id);

        Account?.SubList.Insert(sub);
    }

    private void ProcessUnsub(ParsedCommand cmd)
    {
        _logger.LogDebug("UNSUB {Sid} from client {ClientId}", cmd.Sid, Id);

        if (!_subs.TryGetValue(cmd.Sid!, out var sub))
            return;

        if (cmd.MaxMessages > 0)
        {
            sub.MaxMessages = cmd.MaxMessages;
            // Will be cleaned up when MessageCount reaches MaxMessages
            return;
        }

        _subs.Remove(cmd.Sid!);

        Account?.SubList.Remove(sub);
    }

    private async ValueTask ProcessPubAsync(ParsedCommand cmd)
    {
        Interlocked.Increment(ref InMsgs);
        Interlocked.Add(ref InBytes, cmd.Payload.Length);
        Interlocked.Increment(ref _serverStats.InMsgs);
        Interlocked.Add(ref _serverStats.InBytes, cmd.Payload.Length);

        // Max payload validation (always, hard close)
        if (cmd.Payload.Length > _options.MaxPayload)
        {
            _logger.LogWarning("Client {ClientId} exceeded max payload: {Size} > {MaxPayload}",
                Id, cmd.Payload.Length, _options.MaxPayload);
            await SendErrAndCloseAsync(NatsProtocol.ErrMaxPayloadViolation);
            return;
        }

        // Pedantic mode: validate publish subject
        if (ClientOpts?.Pedantic == true && !SubjectMatch.IsValidPublishSubject(cmd.Subject!))
        {
            _logger.LogDebug("Client {ClientId} invalid publish subject: {Subject}", Id, cmd.Subject);
            await SendErrAsync(NatsProtocol.ErrInvalidPublishSubject);
            return;
        }

        // Permission check for publish
        if (_permissions != null && !_permissions.IsPublishAllowed(cmd.Subject!))
        {
            _logger.LogDebug("Client {ClientId} publish permission denied for {Subject}", Id, cmd.Subject);
            await SendErrAsync(NatsProtocol.ErrPermissionsPublish);
            return;
        }

        ReadOnlyMemory<byte> headers = default;
        ReadOnlyMemory<byte> payload = cmd.Payload;

        if (cmd.Type == CommandType.HPub && cmd.HeaderSize > 0)
        {
            headers = cmd.Payload[..cmd.HeaderSize];
            payload = cmd.Payload[cmd.HeaderSize..];
        }

        Router?.ProcessMessage(cmd.Subject!, cmd.ReplyTo, headers, payload, this);
    }

    private async Task SendInfoAsync(CancellationToken ct)
    {
        var infoJson = JsonSerializer.Serialize(_serverInfo);
        var infoLine = Encoding.ASCII.GetBytes($"INFO {infoJson}\r\n");
        await WriteAsync(infoLine, ct);
    }

    public async Task SendMessageAsync(string subject, string sid, string? replyTo,
        ReadOnlyMemory<byte> headers, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        Interlocked.Increment(ref OutMsgs);
        Interlocked.Add(ref OutBytes, payload.Length + headers.Length);
        Interlocked.Increment(ref _serverStats.OutMsgs);
        Interlocked.Add(ref _serverStats.OutBytes, payload.Length + headers.Length);

        byte[] line;
        if (headers.Length > 0)
        {
            int totalSize = headers.Length + payload.Length;
            line = Encoding.ASCII.GetBytes($"HMSG {subject} {sid} {(replyTo != null ? replyTo + " " : "")}{headers.Length} {totalSize}\r\n");
        }
        else
        {
            line = Encoding.ASCII.GetBytes($"MSG {subject} {sid} {(replyTo != null ? replyTo + " " : "")}{payload.Length}\r\n");
        }

        await _writeLock.WaitAsync(ct);
        try
        {
            await _stream.WriteAsync(line, ct);
            if (headers.Length > 0)
                await _stream.WriteAsync(headers, ct);
            if (payload.Length > 0)
                await _stream.WriteAsync(payload, ct);
            await _stream.WriteAsync(NatsProtocol.CrLf, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeLock.Release();
        }
    }

    private async Task WriteAsync(byte[] data, CancellationToken ct)
    {
        await _writeLock.WaitAsync(ct);
        try
        {
            await _stream.WriteAsync(data, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeLock.Release();
        }
    }

    public async Task SendErrAsync(string message)
    {
        var errLine = Encoding.ASCII.GetBytes($"-ERR '{message}'\r\n");
        try
        {
            await WriteAsync(errLine, _clientCts?.Token ?? CancellationToken.None);
        }
        catch (OperationCanceledException)
        {
            // Expected during shutdown
        }
        catch (IOException ex)
        {
            _logger.LogDebug(ex, "Client {ClientId} failed to send -ERR", Id);
        }
        catch (ObjectDisposedException ex)
        {
            _logger.LogDebug(ex, "Client {ClientId} failed to send -ERR (disposed)", Id);
        }
    }

    public async Task SendErrAndCloseAsync(string message)
    {
        await SendErrAsync(message);
        if (_clientCts is { } cts)
            await cts.CancelAsync();
        else
            _socket.Close();
    }

    private async Task RunPingTimerAsync(CancellationToken ct)
    {
        using var timer = new PeriodicTimer(_options.PingInterval);
        try
        {
            while (await timer.WaitForNextTickAsync(ct))
            {
                var elapsed = Environment.TickCount64 - Interlocked.Read(ref _lastIn);
                if (elapsed < (long)_options.PingInterval.TotalMilliseconds)
                {
                    // Client was recently active, skip ping
                    Interlocked.Exchange(ref _pingsOut, 0);
                    continue;
                }

                if (Volatile.Read(ref _pingsOut) + 1 > _options.MaxPingsOut)
                {
                    _logger.LogDebug("Client {ClientId} stale connection -- closing", Id);
                    await SendErrAndCloseAsync(NatsProtocol.ErrStaleConnection);
                    return;
                }

                var currentPingsOut = Interlocked.Increment(ref _pingsOut);
                _logger.LogDebug("Client {ClientId} sending PING ({PingsOut}/{MaxPingsOut})",
                    Id, currentPingsOut, _options.MaxPingsOut);
                try
                {
                    await WriteAsync(NatsProtocol.PingBytes, ct);
                }
                catch (Exception ex)
                {
                    _logger.LogDebug(ex, "Client {ClientId} failed to send PING", Id);
                    return;
                }
            }
        }
        catch (OperationCanceledException)
        {
            // Normal shutdown
        }
    }

    public void RemoveAllSubscriptions(SubList subList)
    {
        foreach (var sub in _subs.Values)
            subList.Remove(sub);
        _subs.Clear();
    }

    public void Dispose()
    {
        _permissions?.Dispose();
        _clientCts?.Dispose();
        _stream.Dispose();
        _socket.Dispose();
        _writeLock.Dispose();
    }
}