git.stella-ops.org/src/__Libraries/StellaOps.Router.Transport.Tls/TlsConnection.cs

using System.Net.Security;
using System.Net.Sockets;
using System.Security.Cryptography.X509Certificates;
using Microsoft.Extensions.Logging;
using StellaOps.Router.Common.Models;
using StellaOps.Router.Transport.Tcp;

namespace StellaOps.Router.Transport.Tls;

/// <summary>
/// Represents a TLS-secured connection to a microservice.
/// </summary>
public sealed class TlsConnection : IAsyncDisposable
{
    private readonly TcpClient _client;
    private readonly SslStream _sslStream;
    private readonly SemaphoreSlim _writeLock = new(1, 1);
    private readonly TlsTransportOptions _options;
    private readonly ILogger _logger;
    private readonly CancellationTokenSource _connectionCts = new();
    private bool _disposed;

    /// <summary>
    /// Gets the connection ID.
    /// </summary>
    public string ConnectionId { get; }

    /// <summary>
    /// Gets the remote endpoint as a string.
    /// </summary>
    public string RemoteEndpoint { get; }

    /// <summary>
    /// Gets a value indicating whether the connection is active.
    /// </summary>
    public bool IsConnected => _client.Connected && !_disposed;

    /// <summary>
    /// Gets the connection state.
    /// </summary>
    public ConnectionState? State { get; set; }

    /// <summary>
    /// Gets the cancellation token for this connection.
    /// </summary>
    public CancellationToken ConnectionToken => _connectionCts.Token;

    /// <summary>
    /// Gets the remote certificate (if mTLS).
    /// </summary>
    public X509Certificate? RemoteCertificate => _sslStream.RemoteCertificate;

    /// <summary>
    /// Gets the peer identity extracted from the certificate.
    /// </summary>
    public string? PeerIdentity { get; }

    /// <summary>
    /// Event raised when a frame is received.
    /// </summary>
    public event Action<TlsConnection, Frame>? OnFrameReceived;

    /// <summary>
    /// Event raised when the connection is closed.
    /// </summary>
    public event Action<TlsConnection, Exception?>? OnDisconnected;

    /// <summary>
    /// Initializes a new instance of the <see cref="TlsConnection"/> class.
    /// </summary>
    public TlsConnection(
        string connectionId,
        TcpClient client,
        SslStream sslStream,
        TlsTransportOptions options,
        ILogger logger)
    {
        ConnectionId = connectionId;
        _client = client;
        _sslStream = sslStream;
        _options = options;
        _logger = logger;
        RemoteEndpoint = client.Client.RemoteEndPoint?.ToString() ?? "unknown";

        // Extract peer identity from certificate
        if (_sslStream.RemoteCertificate is X509Certificate2 cert)
        {
            PeerIdentity = ExtractIdentityFromCertificate(cert);
        }

        // Configure socket options
        client.ReceiveBufferSize = options.ReceiveBufferSize;
        client.SendBufferSize = options.SendBufferSize;
        client.NoDelay = true;
    }

    /// <summary>
    /// Extracts identity from a certificate.
    /// </summary>
    private static string? ExtractIdentityFromCertificate(X509Certificate2 cert)
    {
        // Try to get Common Name (CN)
        var cn = cert.GetNameInfo(X509NameType.SimpleName, forIssuer: false);
        if (!string.IsNullOrEmpty(cn))
        {
            return cn;
        }

        // Fallback to subject
        return cert.Subject;
    }

    /// <summary>
    /// Starts the read loop to receive frames.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    public async Task ReadLoopAsync(CancellationToken cancellationToken)
    {
        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(
            cancellationToken, _connectionCts.Token);

        Exception? disconnectException = null;

        try
        {
            while (!linkedCts.Token.IsCancellationRequested)
            {
                var frame = await FrameProtocol.ReadFrameAsync(
                    _sslStream,
                    _options.MaxFrameSize,
                    linkedCts.Token);

                if (frame is null)
                {
                    _logger.LogDebug("TLS connection {ConnectionId} closed by remote", ConnectionId);
                    break;
                }

                OnFrameReceived?.Invoke(this, frame);
            }
        }
        catch (OperationCanceledException)
        {
            // Expected on shutdown
        }
        catch (IOException ex) when (ex.InnerException is SocketException)
        {
            disconnectException = ex;
            _logger.LogDebug(ex, "TLS connection {ConnectionId} socket error", ConnectionId);
        }
        catch (Exception ex)
        {
            disconnectException = ex;
            _logger.LogWarning(ex, "TLS connection {ConnectionId} read error", ConnectionId);
        }

        OnDisconnected?.Invoke(this, disconnectException);
    }

    /// <summary>
    /// Writes a frame to the connection.
    /// </summary>
    /// <param name="frame">The frame to write.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    public async Task WriteFrameAsync(Frame frame, CancellationToken cancellationToken = default)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        await _writeLock.WaitAsync(cancellationToken);
        try
        {
            await FrameProtocol.WriteFrameAsync(_sslStream, frame, cancellationToken);
            await _sslStream.FlushAsync(cancellationToken);
        }
        finally
        {
            _writeLock.Release();
        }
    }

    /// <summary>
    /// Closes the connection.
    /// </summary>
    public void Close()
    {
        if (_disposed) return;

        try
        {
            _connectionCts.Cancel();
            _sslStream.Close();
            _client.Close();
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "Error closing TLS connection {ConnectionId}", ConnectionId);
        }
    }

    /// <inheritdoc />
    public async ValueTask DisposeAsync()
    {
        if (_disposed) return;
        _disposed = true;

        try
        {
            await _connectionCts.CancelAsync();
        }
        catch
        {
            // Ignore
        }

        await _sslStream.DisposeAsync();
        _client.Dispose();
        _writeLock.Dispose();
        _connectionCts.Dispose();
    }
}