missions/tests/Azaion.Missions.JwksMock/Services/KeyStore.cs

using System.Security.Cryptography;

namespace Azaion.Missions.JwksMock.Services;

/// <summary>
/// Holds the active ECDSA P-256 keypair used to sign test JWTs, plus an
/// optional retired keypair retained for <c>OldKeyGraceSeconds</c> after a
/// rotation so consumers can still validate in-flight tokens minted under the
/// previous kid (NFT-RES-07 / NFT-SEC-11).
/// </summary>
/// <remarks>
/// Singleton, thread-safe. The private key never leaves the container — only
/// public-half exports go out via the JWKS endpoint.
/// </remarks>
public sealed class KeyStore : IDisposable
{
    private readonly TimeSpan _graceWindow;
    private readonly TimeProvider _clock;
    private readonly Lock _gate = new();

    private KeypairEntry _active;
    private KeypairEntry? _retired;

    public KeyStore(TimeSpan graceWindow, TimeProvider clock)
    {
        _graceWindow = graceWindow;
        _clock = clock;
        _active = KeypairEntry.New();
    }

    public KeypairView Active
    {
        get
        {
            lock (_gate) return _active.View();
        }
    }

    public IReadOnlyList<KeypairView> PublishedKeys()
    {
        lock (_gate)
        {
            EvictExpiredRetired();
            if (_retired is null)
                return [_active.View()];
            return [_active.View(), _retired.View()];
        }
    }

    /// <summary>
    /// Rotate the active keypair. The previous active key is retained as the
    /// retired key (overwriting any older retired entry) until
    /// <c>OldKeyGraceSeconds</c> elapses.
    /// </summary>
    public KeypairView Rotate()
    {
        lock (_gate)
        {
            _retired?.Dispose();
            _retired = _active.WithRetiredAt(_clock.GetUtcNow().Add(_graceWindow));
            _active = KeypairEntry.New();
            return _active.View();
        }
    }

    public void Dispose()
    {
        lock (_gate)
        {
            _active.Dispose();
            _retired?.Dispose();
            _retired = null;
        }
    }

    private void EvictExpiredRetired()
    {
        if (_retired is null) return;
        if (_retired.RetiredAtUtc is { } retiredAt && _clock.GetUtcNow() > retiredAt)
        {
            _retired.Dispose();
            _retired = null;
        }
    }

    private sealed class KeypairEntry : IDisposable
    {
        public ECDsa Ec { get; }
        public string Kid { get; }
        public DateTimeOffset? RetiredAtUtc { get; }

        private KeypairEntry(ECDsa ec, string kid, DateTimeOffset? retiredAt)
        {
            Ec = ec;
            Kid = kid;
            RetiredAtUtc = retiredAt;
        }

        public static KeypairEntry New()
        {
            var ec = ECDsa.Create(ECCurve.NamedCurves.nistP256);
            // kid: SHA-256 of the public key parameters, base64url-truncated to 16 bytes.
            var pub = ec.ExportSubjectPublicKeyInfo();
            var hash = SHA256.HashData(pub);
            var kid = Base64Url.Encode(hash.AsSpan(0, 16));
            return new KeypairEntry(ec, kid, retiredAt: null);
        }

        public KeypairEntry WithRetiredAt(DateTimeOffset retiredAtUtc)
            => new(Ec, Kid, retiredAtUtc);

        public KeypairView View() => new(Kid, Ec, RetiredAtUtc);

        public void Dispose() => Ec.Dispose();
    }
}

public readonly record struct KeypairView(string Kid, ECDsa Ec, DateTimeOffset? RetiredAtUtc);