git.stella-ops.org/src/__Libraries/StellaOps.Cryptography.Kms/AwsKmsClient.cs

using System.Collections.Concurrent;
using System.Collections.Immutable;
using System.Security.Cryptography;
using Microsoft.IdentityModel.Tokens;

namespace StellaOps.Cryptography.Kms;

/// <summary>
/// AWS KMS implementation of <see cref="IKmsClient"/>.
/// </summary>
public sealed class AwsKmsClient : IKmsClient, IDisposable
{
    private readonly IAwsKmsFacade _facade;
    private readonly TimeProvider _timeProvider;
    private readonly TimeSpan _metadataCacheDuration;
    private readonly TimeSpan _publicKeyCacheDuration;

    private readonly ConcurrentDictionary<string, CachedMetadata> _metadataCache = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, CachedPublicKey> _publicKeyCache = new(StringComparer.Ordinal);
    private bool _disposed;

    public AwsKmsClient(IAwsKmsFacade facade, AwsKmsOptions options, TimeProvider? timeProvider = null)
    {
        _facade = facade ?? throw new ArgumentNullException(nameof(facade));
        ArgumentNullException.ThrowIfNull(options);

        _timeProvider = timeProvider ?? TimeProvider.System;
        _metadataCacheDuration = options.MetadataCacheDuration;
        _publicKeyCacheDuration = options.PublicKeyCacheDuration;
    }

    public async Task<KmsSignResult> SignAsync(
        string keyId,
        string? keyVersion,
        ReadOnlyMemory<byte> data,
        CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);
        if (data.IsEmpty)
        {
            throw new ArgumentException("Signing payload cannot be empty.", nameof(data));
        }

        var digest = ComputeSha256(data);
        try
        {
            var resource = ResolveResource(keyId, keyVersion);
            var result = await _facade.SignAsync(resource, digest, cancellationToken).ConfigureAwait(false);

            return new KmsSignResult(
                keyId,
                string.IsNullOrWhiteSpace(result.VersionId) ? resource : result.VersionId,
                KmsAlgorithms.Es256,
                result.Signature);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(digest.AsSpan());
        }
    }

    public async Task<bool> VerifyAsync(
        string keyId,
        string? keyVersion,
        ReadOnlyMemory<byte> data,
        ReadOnlyMemory<byte> signature,
        CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);
        if (data.IsEmpty || signature.IsEmpty)
        {
            return false;
        }

        var digest = ComputeSha256(data);
        try
        {
            var resource = ResolveResource(keyId, keyVersion);
            return await _facade.VerifyAsync(resource, digest, signature, cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(digest.AsSpan());
        }
    }

    public async Task<KmsKeyMetadata> GetMetadataAsync(string keyId, CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);

        var metadata = await GetCachedMetadataAsync(keyId, cancellationToken).ConfigureAwait(false);
        var publicKey = await GetCachedPublicKeyAsync(metadata.KeyId, cancellationToken).ConfigureAwait(false);

        var versionState = MapState(metadata.Status);
        var versionMetadata = ImmutableArray.Create(
            new KmsKeyVersionMetadata(
                publicKey.VersionId,
                versionState,
                metadata.CreatedAt,
                null,
                Convert.ToBase64String(publicKey.SubjectPublicKeyInfo),
                ResolveCurveName(publicKey.Curve)));

        return new KmsKeyMetadata(
            metadata.KeyId,
            KmsAlgorithms.Es256,
            versionState,
            metadata.CreatedAt,
            versionMetadata);
    }

    public async Task<KmsKeyMaterial> ExportAsync(
        string keyId,
        string? keyVersion,
        CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
        ArgumentException.ThrowIfNullOrWhiteSpace(keyId);

        var metadata = await GetCachedMetadataAsync(keyId, cancellationToken).ConfigureAwait(false);
        var resource = ResolveResource(metadata.KeyId, keyVersion);
        var publicKey = await GetCachedPublicKeyAsync(resource, cancellationToken).ConfigureAwait(false);

        using var ecdsa = ECDsa.Create();
        ecdsa.ImportSubjectPublicKeyInfo(publicKey.SubjectPublicKeyInfo, out _);
        var parameters = ecdsa.ExportParameters(false);

        return new KmsKeyMaterial(
            metadata.KeyId,
            publicKey.VersionId,
            KmsAlgorithms.Es256,
            ResolveCurveName(publicKey.Curve),
            Array.Empty<byte>(),
            parameters.Q.X ?? throw new InvalidOperationException("Public key missing X coordinate."),
            parameters.Q.Y ?? throw new InvalidOperationException("Public key missing Y coordinate."),
            metadata.CreatedAt);
    }

    public Task<KmsKeyMetadata> RotateAsync(string keyId, CancellationToken cancellationToken = default)
        => throw new NotSupportedException("AWS KMS rotation must be orchestrated via AWS KMS policies or schedules.");

    public Task RevokeAsync(string keyId, CancellationToken cancellationToken = default)
        => throw new NotSupportedException("AWS KMS key revocation must be managed through AWS KMS APIs or console.");

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;
        _facade.Dispose();
    }

    private async Task<AwsKeyMetadata> GetCachedMetadataAsync(string keyId, CancellationToken cancellationToken)
    {
        var now = _timeProvider.GetUtcNow();
        if (_metadataCache.TryGetValue(keyId, out var cached) && cached.ExpiresAt > now)
        {
            return cached.Metadata;
        }

        var metadata = await _facade.GetMetadataAsync(keyId, cancellationToken).ConfigureAwait(false);
        var entry = new CachedMetadata(metadata, now.Add(_metadataCacheDuration));
        _metadataCache[keyId] = entry;
        return metadata;
    }

    private async Task<AwsPublicKeyMaterial> GetCachedPublicKeyAsync(string resource, CancellationToken cancellationToken)
    {
        var now = _timeProvider.GetUtcNow();
        if (_publicKeyCache.TryGetValue(resource, out var cached) && cached.ExpiresAt > now)
        {
            return cached.Material;
        }

        var material = await _facade.GetPublicKeyAsync(resource, cancellationToken).ConfigureAwait(false);
        var entry = new CachedPublicKey(material, now.Add(_publicKeyCacheDuration));
        _publicKeyCache[resource] = entry;
        return material;
    }

    private static byte[] ComputeSha256(ReadOnlyMemory<byte> data)
    {
        var digest = new byte[32];
        if (!SHA256.TryHashData(data.Span, digest, out _))
        {
            throw new InvalidOperationException("Failed to hash payload with SHA-256.");
        }

        return digest;
    }

    private static string ResolveResource(string keyId, string? version)
        => string.IsNullOrWhiteSpace(version) ? keyId : version;

    private static string ResolveCurveName(string curve)
    {
        if (string.Equals(curve, "ECC_NIST_P256", StringComparison.OrdinalIgnoreCase) ||
            string.Equals(curve, "P-256", StringComparison.OrdinalIgnoreCase))
        {
            return JsonWebKeyECTypes.P256;
        }

        if (string.Equals(curve, "ECC_NIST_P384", StringComparison.OrdinalIgnoreCase) ||
            string.Equals(curve, "P-384", StringComparison.OrdinalIgnoreCase))
        {
            return JsonWebKeyECTypes.P384;
        }

        if (string.Equals(curve, "ECC_NIST_P521", StringComparison.OrdinalIgnoreCase) ||
            string.Equals(curve, "P-521", StringComparison.OrdinalIgnoreCase))
        {
            return JsonWebKeyECTypes.P521;
        }

        if (string.Equals(curve, "SECP256K1", StringComparison.OrdinalIgnoreCase) ||
            string.Equals(curve, "ECC_SECG_P256K1", StringComparison.OrdinalIgnoreCase))
        {
            return "secp256k1";
        }

        return curve;
    }

    private static KmsKeyState MapState(AwsKeyStatus status)
        => status switch
        {
            AwsKeyStatus.Enabled => KmsKeyState.Active,
            AwsKeyStatus.PendingImport or AwsKeyStatus.PendingUpdate => KmsKeyState.PendingRotation,
            AwsKeyStatus.Disabled or AwsKeyStatus.PendingDeletion or AwsKeyStatus.Unavailable => KmsKeyState.Revoked,
            _ => KmsKeyState.Active,
        };

    private void ThrowIfDisposed()
    {
        if (_disposed)
        {
            throw new ObjectDisposedException(nameof(AwsKmsClient));
        }
    }

    private sealed record CachedMetadata(AwsKeyMetadata Metadata, DateTimeOffset ExpiresAt);

    private sealed record CachedPublicKey(AwsPublicKeyMaterial Material, DateTimeOffset ExpiresAt);
}