git.stella-ops.org/src/ReleaseOrchestrator/__Agents/StellaOps.Agent.WinRM/WinRmConnectionPool.cs

using Microsoft.Extensions.Logging;
using System.Collections.Concurrent;

namespace StellaOps.Agent.WinRM;

/// <summary>
/// Connection pool for WinRM sessions.
/// </summary>
public sealed class WinRmConnectionPool : IAsyncDisposable
{
    private readonly ConcurrentDictionary<string, PooledSession> _sessions = new();
    private readonly IHttpClientFactory _httpClientFactory;
    private readonly ILogger _logger;
    private readonly TimeSpan _idleTimeout;
    private readonly Timer _cleanupTimer;
    private bool _disposed;

    /// <summary>
    /// Creates a new WinRM connection pool.
    /// </summary>
    public WinRmConnectionPool(
        IHttpClientFactory httpClientFactory,
        ILogger<WinRmConnectionPool> logger,
        TimeSpan? idleTimeout = null)
    {
        _httpClientFactory = httpClientFactory;
        _logger = logger;
        _idleTimeout = idleTimeout ?? TimeSpan.FromMinutes(5);
        _cleanupTimer = new Timer(CleanupIdleSessions, null, _idleTimeout, _idleTimeout);
    }

    /// <summary>
    /// Gets or creates a WinRM session for the given connection info.
    /// </summary>
    public async Task<WinRmSession> GetSessionAsync(WinRmConnectionInfo connectionInfo, CancellationToken ct = default)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        var key = connectionInfo.GetConnectionKey();

        if (_sessions.TryGetValue(key, out var pooled) && !pooled.IsExpired(_idleTimeout))
        {
            pooled.Touch();
            return pooled.Session;
        }

        // Create new session
        var httpClient = CreateHttpClient(connectionInfo);
        var session = new WinRmSession(connectionInfo, httpClient, _logger);

        await session.ConnectAsync(ct);

        var newPooled = new PooledSession(session, httpClient);
        _sessions[key] = newPooled;

        _logger.LogDebug("Created new WinRM session for {Key}", key);

        return session;
    }

    /// <summary>
    /// Removes a session from the pool.
    /// </summary>
    public async Task RemoveSessionAsync(WinRmConnectionInfo connectionInfo, CancellationToken ct = default)
    {
        var key = connectionInfo.GetConnectionKey();

        if (_sessions.TryRemove(key, out var pooled))
        {
            await pooled.DisposeAsync();
            _logger.LogDebug("Removed WinRM session for {Key}", key);
        }
    }

    private HttpClient CreateHttpClient(WinRmConnectionInfo connectionInfo)
    {
        var client = _httpClientFactory.CreateClient("WinRM");
        client.Timeout = connectionInfo.Timeout;

        // Set up authentication based on mechanism
        var credentials = CreateCredentials(connectionInfo);
        if (credentials != null)
        {
            // Note: In production, use HttpClientHandler with credentials
            // For Basic auth, set Authorization header directly
            if (connectionInfo.AuthMechanism == WinRmAuthMechanism.Basic)
            {
                var authValue = Convert.ToBase64String(
                    System.Text.Encoding.UTF8.GetBytes(
                        $"{connectionInfo.Username}:{connectionInfo.Password}"));
                client.DefaultRequestHeaders.Authorization =
                    new System.Net.Http.Headers.AuthenticationHeaderValue("Basic", authValue);
            }
        }

        return client;
    }

    private static System.Net.NetworkCredential? CreateCredentials(WinRmConnectionInfo connectionInfo)
    {
        if (string.IsNullOrEmpty(connectionInfo.Password))
            return null;

        return new System.Net.NetworkCredential(
            connectionInfo.Username,
            connectionInfo.Password,
            connectionInfo.Domain ?? string.Empty);
    }

    private void CleanupIdleSessions(object? state)
    {
        if (_disposed)
            return;

        foreach (var kvp in _sessions)
        {
            if (kvp.Value.IsExpired(_idleTimeout))
            {
                if (_sessions.TryRemove(kvp.Key, out var pooled))
                {
                    _ = pooled.DisposeAsync();
                    _logger.LogDebug("Cleaned up idle WinRM session for {Key}", kvp.Key);
                }
            }
        }
    }

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

        _disposed = true;

        await _cleanupTimer.DisposeAsync();

        foreach (var kvp in _sessions)
        {
            await kvp.Value.DisposeAsync();
        }

        _sessions.Clear();
    }

    private sealed class PooledSession : IAsyncDisposable
    {
        private readonly HttpClient _httpClient;
        private DateTimeOffset _lastUsed;

        public WinRmSession Session { get; }

        public PooledSession(WinRmSession session, HttpClient httpClient)
        {
            Session = session;
            _httpClient = httpClient;
            _lastUsed = DateTimeOffset.UtcNow;
        }

        public void Touch() => _lastUsed = DateTimeOffset.UtcNow;

        public bool IsExpired(TimeSpan idleTimeout) =>
            DateTimeOffset.UtcNow - _lastUsed > idleTimeout;

        public async ValueTask DisposeAsync()
        {
            await Session.CloseAsync();
            Session.Dispose();
            _httpClient.Dispose();
        }
    }
}