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

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

namespace StellaOps.Cryptography.Kms;

/// <summary>
/// PKCS#11-backed implementation of <see cref="IKmsClient"/>.
/// </summary>
public sealed class Pkcs11KmsClient : IKmsClient
{
    private readonly IPkcs11Facade _facade;
    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 Pkcs11KmsClient(IPkcs11Facade facade, Pkcs11Options options)
    {
        _facade = facade ?? throw new ArgumentNullException(nameof(facade));
        ArgumentNullException.ThrowIfNull(options);

        _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 descriptor = await GetCachedMetadataAsync(keyId, cancellationToken).ConfigureAwait(false);
            var signature = await _facade.SignDigestAsync(digest, cancellationToken).ConfigureAwait(false);

            return new KmsSignResult(
                descriptor.Descriptor.KeyId,
                descriptor.Descriptor.KeyId,
                KmsAlgorithms.Es256,
                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 publicMaterial = await GetCachedPublicKeyAsync(keyId, cancellationToken).ConfigureAwait(false);

            using var ecdsa = ECDsa.Create(new ECParameters
            {
                Curve = ResolveCurve(publicMaterial.Material.Curve),
                Q =
                {
                    X = publicMaterial.Material.Qx,
                    Y = publicMaterial.Material.Qy,
                },
            });

            return ecdsa.VerifyHash(digest, signature.ToArray());
        }
        finally
        {
            CryptographicOperations.ZeroMemory(digest.AsSpan());
        }
    }

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

        var descriptor = await GetCachedMetadataAsync(keyId, cancellationToken).ConfigureAwait(false);
        var publicMaterial = await GetCachedPublicKeyAsync(keyId, cancellationToken).ConfigureAwait(false);

        using var ecdsa = ECDsa.Create(new ECParameters
        {
            Curve = ResolveCurve(publicMaterial.Material.Curve),
            Q =
            {
                X = publicMaterial.Material.Qx,
                Y = publicMaterial.Material.Qy,
            },
        });

        var subjectInfo = Convert.ToBase64String(ecdsa.ExportSubjectPublicKeyInfo());

        var version = new KmsKeyVersionMetadata(
            descriptor.Descriptor.KeyId,
            KmsKeyState.Active,
            descriptor.Descriptor.CreatedAt,
            null,
            subjectInfo,
            publicMaterial.Material.Curve);

        return new KmsKeyMetadata(
            descriptor.Descriptor.KeyId,
            KmsAlgorithms.Es256,
            KmsKeyState.Active,
            descriptor.Descriptor.CreatedAt,
            ImmutableArray.Create(version));
    }

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

        var descriptor = await GetCachedMetadataAsync(keyId, cancellationToken).ConfigureAwait(false);
        var publicMaterial = await GetCachedPublicKeyAsync(keyId, cancellationToken).ConfigureAwait(false);

        return new KmsKeyMaterial(
            descriptor.Descriptor.KeyId,
            descriptor.Descriptor.KeyId,
            KmsAlgorithms.Es256,
            publicMaterial.Material.Curve,
            Array.Empty<byte>(),
            publicMaterial.Material.Qx,
            publicMaterial.Material.Qy,
            descriptor.Descriptor.CreatedAt);
    }

    public Task<KmsKeyMetadata> RotateAsync(string keyId, CancellationToken cancellationToken = default)
        => throw new NotSupportedException("PKCS#11 rotation requires HSM administrative tooling.");

    public Task RevokeAsync(string keyId, CancellationToken cancellationToken = default)
        => throw new NotSupportedException("PKCS#11 revocation must be handled by HSM policies.");

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

        _disposed = true;
        _facade.Dispose();
    }

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

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

    private async Task<CachedPublicKey> GetCachedPublicKeyAsync(string keyId, CancellationToken cancellationToken)
    {
        var now = DateTimeOffset.UtcNow;
        if (_publicKeyCache.TryGetValue(keyId, out var cached) && cached.ExpiresAt > now)
        {
            return cached;
        }

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

    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 ECCurve ResolveCurve(string curve)
        => curve switch
        {
            JsonWebKeyECTypes.P256 => ECCurve.NamedCurves.nistP256,
            JsonWebKeyECTypes.P384 => ECCurve.NamedCurves.nistP384,
            JsonWebKeyECTypes.P521 => ECCurve.NamedCurves.nistP521,
            _ => throw new InvalidOperationException($"Unsupported EC curve '{curve}'."),
        };

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

    private sealed record CachedMetadata(Pkcs11KeyDescriptor Descriptor, DateTimeOffset ExpiresAt);

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