mxaccessgw/src/ZB.MOM.WW.MxGateway.Tests/Security/Authentication/CachingApiKeyVerifierTests.cs

using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.Time.Testing;
using ZB.MOM.WW.Auth.Abstractions.ApiKeys;
using ZB.MOM.WW.MxGateway.Server.Security.Authentication;

using LibApiKeyIdentity = ZB.MOM.WW.Auth.Abstractions.ApiKeys.ApiKeyIdentity;

namespace ZB.MOM.WW.MxGateway.Tests.Security.Authentication;

/// <summary>
/// Hot-path decorators. Covers all three mechanisms: <see cref="CachingApiKeyVerifier"/>
/// (read/verification coalescing plus revoke/rotate invalidation),
/// <see cref="CoalescingMarkApiKeyStore"/> (the <c>last_used</c> write coalescing that keeps the
/// per-RPC database write off the throughput ceiling), and the constraint-blob cache inside
/// <see cref="GatewayApiKeyIdentityMapper"/> (which keeps the per-RPC constraints JSON parse off
/// the authenticated path).
/// </summary>
public sealed class CachingApiKeyVerifierTests
{
    private const string Header = "Bearer mxgw_operator01_super-secret";

    /// <summary>A cache hit within the TTL returns the cached result and never calls the inner verifier.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task VerifyAsync_RepeatedWithinTtl_CallsInnerOnce()
    {
        FakeVerifier inner = new(Success("operator01"));
        using MemoryCache cache = NewCache();
        CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(15));

        ApiKeyVerification first = await verifier.VerifyAsync(Header, CancellationToken.None);
        ApiKeyVerification second = await verifier.VerifyAsync(Header, CancellationToken.None);

        Assert.True(first.Succeeded);
        Assert.True(second.Succeeded);
        Assert.Equal(1, inner.CallCount);
    }

    /// <summary>Different presented secrets are cached under distinct keys (no cross-secret aliasing).</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task VerifyAsync_DifferentTokens_NotAliased()
    {
        FakeVerifier inner = new(Success("operator01"));
        using MemoryCache cache = NewCache();
        CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(15));

        await verifier.VerifyAsync("Bearer mxgw_operator01_secret-a", CancellationToken.None);
        await verifier.VerifyAsync("Bearer mxgw_operator01_secret-b", CancellationToken.None);

        Assert.Equal(2, inner.CallCount);
    }

    /// <summary>Failed verifications are never cached; every attempt reaches the inner verifier.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task VerifyAsync_FailedVerification_NotCached()
    {
        FakeVerifier inner = new(Failure(ApiKeyFailure.SecretMismatch));
        using MemoryCache cache = NewCache();
        CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(15));

        await verifier.VerifyAsync(Header, CancellationToken.None);
        await verifier.VerifyAsync(Header, CancellationToken.None);

        Assert.Equal(2, inner.CallCount);
    }

    /// <summary>A TTL of zero disables caching: the inner verifier is called on every request.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task VerifyAsync_ZeroTtl_DisablesCache()
    {
        FakeVerifier inner = new(Success("operator01"));
        using MemoryCache cache = NewCache();
        CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.Zero);

        await verifier.VerifyAsync(Header, CancellationToken.None);
        await verifier.VerifyAsync(Header, CancellationToken.None);

        Assert.Equal(2, inner.CallCount);
    }

    /// <summary>Invalidating a key id (revoke/rotate) drops its cached verification, forcing a re-verify.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Invalidate_DropsCachedEntry_ForcesReverify()
    {
        FakeVerifier inner = new(Success("operator01"));
        using MemoryCache cache = NewCache();
        CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(30));

        await verifier.VerifyAsync(Header, CancellationToken.None);
        Assert.Equal(1, inner.CallCount);

        ((IApiKeyCacheInvalidator)verifier).Invalidate("operator01");

        await verifier.VerifyAsync(Header, CancellationToken.None);
        Assert.Equal(2, inner.CallCount);
    }

    /// <summary>
    /// SEC-34 window 3: an <see cref="IApiKeyCacheInvalidator.Invalidate"/> that lands while a
    /// verification is in flight in the inner library must discard that verification's repopulation,
    /// so the very next request re-verifies (revoke takes effect immediately, not after the TTL).
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Invalidate_DuringInFlightVerification_DiscardsStaleRepopulation()
    {
        GatedVerifier inner = new(Success("operator01"));
        using MemoryCache cache = NewCache();
        CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(30));

        // Begin a verification and wait until it is parked inside the inner verifier.
        Task<ApiKeyVerification> inFlight = verifier.VerifyAsync(Header, CancellationToken.None);
        await inner.Entered;

        // Revoke while the verify is still awaiting the inner library.
        ((IApiKeyCacheInvalidator)verifier).Invalidate("operator01");

        // Release the inner verifier; the in-flight call completes but must NOT cache its result.
        inner.Release();
        ApiKeyVerification result = await inFlight;
        Assert.True(result.Succeeded);

        // The follow-up request finds no cached entry and reaches the inner verifier again.
        await verifier.VerifyAsync(Header, CancellationToken.None);
        Assert.Equal(2, inner.CallCount);
    }

    /// <summary>
    /// The store decorator coalesces repeated <c>MarkUsed</c> writes for the same key inside the
    /// window down to a single forwarded write — the ≤1/min guarantee for <c>last_used_utc</c>.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CoalescingStore_RepeatedMarksWithinWindow_ForwardsOnce()
    {
        FakeStore inner = new();
        FakeTimeProvider clock = new(DateTimeOffset.UtcNow);
        CoalescingMarkApiKeyStore store = new(inner, TimeSpan.FromMinutes(1), clock);

        // Ten rapid authenticated calls in the same minute.
        for (int i = 0; i < 10; i++)
        {
            await store.MarkUsedAsync("operator01", clock.GetUtcNow(), CancellationToken.None);
            clock.Advance(TimeSpan.FromSeconds(5));
        }

        Assert.Equal(1, inner.MarkUsedCount);
    }

    /// <summary>After the window elapses the next mark is forwarded again (staleness is bounded, not frozen).</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CoalescingStore_AfterWindow_ForwardsAgain()
    {
        FakeStore inner = new();
        FakeTimeProvider clock = new(DateTimeOffset.UtcNow);
        CoalescingMarkApiKeyStore store = new(inner, TimeSpan.FromMinutes(1), clock);

        await store.MarkUsedAsync("operator01", clock.GetUtcNow(), CancellationToken.None);
        clock.Advance(TimeSpan.FromSeconds(61));
        await store.MarkUsedAsync("operator01", clock.GetUtcNow(), CancellationToken.None);

        Assert.Equal(2, inner.MarkUsedCount);
    }

    /// <summary>Distinct keys are coalesced independently.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CoalescingStore_DistinctKeys_TrackedSeparately()
    {
        FakeStore inner = new();
        FakeTimeProvider clock = new(DateTimeOffset.UtcNow);
        CoalescingMarkApiKeyStore store = new(inner, TimeSpan.FromMinutes(1), clock);

        await store.MarkUsedAsync("key-a", clock.GetUtcNow(), CancellationToken.None);
        await store.MarkUsedAsync("key-b", clock.GetUtcNow(), CancellationToken.None);
        await store.MarkUsedAsync("key-a", clock.GetUtcNow(), CancellationToken.None);

        Assert.Equal(2, inner.MarkUsedCount);
    }

    /// <summary>A zero window disables coalescing: every mark is forwarded.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CoalescingStore_ZeroWindow_ForwardsEveryMark()
    {
        FakeStore inner = new();
        FakeTimeProvider clock = new(DateTimeOffset.UtcNow);
        CoalescingMarkApiKeyStore store = new(inner, TimeSpan.Zero, clock);

        await store.MarkUsedAsync("operator01", clock.GetUtcNow(), CancellationToken.None);
        await store.MarkUsedAsync("operator01", clock.GetUtcNow(), CancellationToken.None);

        Assert.Equal(2, inner.MarkUsedCount);
    }

    /// <summary>
    /// Pins the header parse that arms the revoke-vs-in-flight generation guard. It runs on every
    /// cache miss and is written to allocate only the returned key id; these cases hold it to the
    /// rules the original <c>Split('_')</c> form applied. Note this parse is deliberately laxer than
    /// the interceptor's partition parse — it applies no key-id length cap and does not require a
    /// non-empty third segment — because a wrong <em>id</em> here would disarm the guard, whereas
    /// an over-long one merely fails to match any generation.
    /// </summary>
    /// <param name="authorizationHeader">The presented header value.</param>
    /// <param name="expected">The key id the parse must yield, or <see langword="null"/>.</param>
    [Theory]
    [InlineData(null, null)]
    [InlineData("", null)]
    [InlineData("   ", null)]
    [InlineData("Bearer", null)]
    [InlineData("Bearer ", null)]
    [InlineData("Bearer mxgw_operator01_super-secret", "operator01")]
    [InlineData("bearer mxgw_operator01_super-secret", "operator01")]
    [InlineData("  Bearer   mxgw_operator01_super-secret  ", "operator01")]
    [InlineData("mxgw_operator01_super-secret", "operator01")]
    [InlineData("Bearer mxgw_abc_sec_ret", "abc")]
    [InlineData("Bearer mxgw_a_b_c", "a")]

    // Laxer than the interceptor: an empty or absent third segment still yields the key id.
    [InlineData("Bearer mxgw_abc_", "abc")]
    [InlineData("Bearer mxgw_abc__secret", "abc")]
    [InlineData("Bearer mxgw_abc", null)]
    [InlineData("Bearer mxgwabcsecret", null)]
    [InlineData("Bearer mxgw__secret", null)]
    [InlineData("Bearer _mxgw_abc_secret", null)]
    [InlineData("Bearer MXGW_abc_secret", null)]
    [InlineData("Bearer xmxgw_abc_secret", null)]
    [InlineData("Bearer mxgw", null)]
    [InlineData("Bearer mxgw_", null)]
    [InlineData("Bearer ___", null)]
    public void TryParseKeyId_MatchesTokenShapeRules(string? authorizationHeader, string? expected)
    {
        Assert.Equal(expected, CachingApiKeyVerifier.TryParseKeyId(authorizationHeader));
    }

    /// <summary>The guard parse applies no key-id length cap: an over-long id is still returned whole.</summary>
    [Fact]
    public void TryParseKeyId_LongKeyId_ReturnedWhole()
    {
        string keyId = new('a', 65);

        Assert.Equal(keyId, CachingApiKeyVerifier.TryParseKeyId($"Bearer mxgw_{keyId}_secret"));
    }

    /// <summary>
    ///     The mapper's constraint-blob cache is bounded by eviction, not by a hard stop at the cap:
    ///     once it is full the oldest entry is dropped so a newly-seen blob is still cached. A cache
    ///     that merely stopped accepting entries would re-parse every blob beyond the cap on every
    ///     single RPC, forever. Asserted behaviourally through instance identity — a cached blob maps
    ///     to the same <see cref="ApiKeyConstraints"/> instance, a re-parsed one does not.
    /// </summary>
    [Fact]
    public void ToGatewayIdentity_ConstraintCacheOverCapacity_EvictsOldestAndKeepsCaching()
    {
        string firstJson = ConstraintsJson("Area_FifoProbe");
        ApiKeyConstraints first = MapConstraints(firstJson);
        Assert.Same(first, MapConstraints(firstJson));

        // Push strictly more than the cap through the cache after the probe blob, so FIFO eviction
        // is guaranteed to have reached it however full the (process-wide) cache already was.
        for (int i = 0; i < GatewayApiKeyIdentityMapper.MaxCachedConstraintBlobs + 8; i++)
        {
            MapConstraints(ConstraintsJson($"Area_FifoFlood_{i}"));
        }

        Assert.True(
            GatewayApiKeyIdentityMapper.CurrentCacheSize <= GatewayApiKeyIdentityMapper.MaxCachedConstraintBlobs,
            $"cache grew to {GatewayApiKeyIdentityMapper.CurrentCacheSize} entries, past the {GatewayApiKeyIdentityMapper.MaxCachedConstraintBlobs} cap");

        // Evicted, so the probe blob is parsed afresh...
        ApiKeyConstraints reparsed = MapConstraints(firstJson);
        Assert.NotSame(first, reparsed);
        Assert.Equal(first.ReadSubtrees, reparsed.ReadSubtrees);

        // ...and re-cached, rather than re-parsed on every later call.
        Assert.Same(reparsed, MapConstraints(firstJson));
    }

    /// <summary>
    ///     Stress: the constraint cache's bound is enforced by the inserting thread itself
    ///     (<c>GetOrAdd</c>, then enqueue, then evict), so concurrent inserters can each land an
    ///     entry before any of them reaches the eviction step. That overshoot is real but must be
    ///     transient and proportional to the in-flight inserters — not unbounded growth — and once
    ///     the churn stops the cache must be back at or under the cap. Hammered on both the
    ///     converging path (every iteration maps the same blob, so all but one <c>GetOrAdd</c>
    ///     loses) and the growth path (a distinct blob per iteration, which is what forces
    ///     eviction).
    /// </summary>
    [Fact]
    public void ToGatewayIdentity_ConcurrentBlobs_OvershootIsTransientAndCacheSettlesUnderCap()
    {
        const int cap = GatewayApiKeyIdentityMapper.MaxCachedConstraintBlobs;

        // Overshoot is bounded by how many inserters can sit between their GetOrAdd and their own
        // EvictIfOverCapacity, so scale the allowance with the available parallelism rather than
        // pinning a magic number. Generous on purpose: the assertion under test is "bounded", not
        // "bounded by exactly this".
        int transientAllowance = (Environment.ProcessorCount * 8) + 64;
        string sharedJson = ConstraintsJson("Area_StressShared");
        int peak = 0;

        Parallel.For(0, (cap * 2) + 64, index =>
        {
            MapConstraints(sharedJson);
            MapConstraints(ConstraintsJson($"Area_Stress_{index}"));

            int size = GatewayApiKeyIdentityMapper.CurrentCacheSize;
            int seen = Volatile.Read(ref peak);
            while (size > seen && Interlocked.CompareExchange(ref peak, size, seen) != seen)
            {
                seen = Volatile.Read(ref peak);
            }
        });

        Assert.True(
            peak <= cap + transientAllowance,
            $"cache peaked at {peak} entries, past the {cap} cap plus the {transientAllowance} transient allowance");

        // The cache is process-wide static and other test classes in this assembly map identities
        // too, so poll rather than asserting on the instant the loop returns.
        Assert.True(
            SpinWait.SpinUntil(
                () => GatewayApiKeyIdentityMapper.CurrentCacheSize <= cap,
                TimeSpan.FromSeconds(5)),
            $"cache settled at {GatewayApiKeyIdentityMapper.CurrentCacheSize} entries, past the {cap} cap");
    }

    private static ApiKeyConstraints MapConstraints(string constraintsJson) =>
        GatewayApiKeyIdentityMapper.ToGatewayIdentity(new LibApiKeyIdentity(
            KeyId: "operator01",
            DisplayName: "Operator Key",
            Scopes: new HashSet<string>(StringComparer.Ordinal),
            Constraints: constraintsJson)).EffectiveConstraints;

    private static string ConstraintsJson(string readSubtree) =>
        ApiKeyConstraintSerializer.Serialize(ApiKeyConstraints.Empty with { ReadSubtrees = [readSubtree] })!;

    private static MemoryCache NewCache() => new(new MemoryCacheOptions());

    private static ApiKeyVerification Success(string keyId) => new(
        Succeeded: true,
        Identity: new LibApiKeyIdentity(
            KeyId: keyId,
            DisplayName: "Operator Key",
            Scopes: new HashSet<string>(StringComparer.Ordinal),
            Constraints: null),
        Failure: null);

    private static ApiKeyVerification Failure(ApiKeyFailure failure) =>
        new(Succeeded: false, Identity: null, Failure: failure);

    private sealed class FakeVerifier(ApiKeyVerification result) : IApiKeyVerifier
    {
        /// <summary>Gets the number of times <see cref="VerifyAsync"/> has been called.</summary>
        public int CallCount { get; private set; }

        /// <summary>Records the call and returns the fixed <paramref name="result"/> supplied at construction.</summary>
        /// <param name="authorizationHeader">The authorization header presented by the caller.</param>
        /// <param name="ct">A token to observe for cancellation.</param>
        /// <returns>The fixed verification result.</returns>
        public Task<ApiKeyVerification> VerifyAsync(string authorizationHeader, CancellationToken ct)
        {
            CallCount++;
            return Task.FromResult(result);
        }
    }

    // A verifier that parks inside VerifyAsync until Release() is called, so a test can interleave an
    // Invalidate with an in-flight verification.
    private sealed class GatedVerifier(ApiKeyVerification result) : IApiKeyVerifier
    {
        private readonly TaskCompletionSource _entered =
            new(TaskCreationOptions.RunContinuationsAsynchronously);

        private readonly TaskCompletionSource _release =
            new(TaskCreationOptions.RunContinuationsAsynchronously);

        /// <summary>Gets the number of times <see cref="VerifyAsync"/> has been called.</summary>
        public int CallCount { get; private set; }

        /// <summary>Completes once the first call has entered and parked in the inner verifier.</summary>
        public Task Entered => _entered.Task;

        /// <summary>Unparks the first (gated) call.</summary>
        public void Release() => _release.TrySetResult();

        /// <summary>Records the call; the first call parks on the gate, later calls return immediately.</summary>
        /// <param name="authorizationHeader">The authorization header presented by the caller.</param>
        /// <param name="ct">A token to observe for cancellation.</param>
        /// <returns>The fixed verification result.</returns>
        public async Task<ApiKeyVerification> VerifyAsync(string authorizationHeader, CancellationToken ct)
        {
            bool first = CallCount == 0;
            CallCount++;
            if (first)
            {
                _entered.TrySetResult();
                await _release.Task.ConfigureAwait(false);
            }

            return result;
        }
    }

    private sealed class FakeStore : IApiKeyStore
    {
        /// <summary>Gets the number of times <see cref="MarkUsedAsync"/> has been called.</summary>
        public int MarkUsedCount { get; private set; }

        /// <summary>Always returns <see langword="null"/>; not exercised by these tests.</summary>
        /// <param name="keyId">The key id to look up.</param>
        /// <param name="ct">A token to observe for cancellation.</param>
        /// <returns><see langword="null"/>.</returns>
        public Task<ApiKeyRecord?> FindByKeyIdAsync(string keyId, CancellationToken ct)
            => Task.FromResult<ApiKeyRecord?>(null);

        /// <summary>Always returns <see langword="null"/>; not exercised by these tests.</summary>
        /// <param name="keyId">The key id to look up.</param>
        /// <param name="ct">A token to observe for cancellation.</param>
        /// <returns><see langword="null"/>.</returns>
        public Task<ApiKeyRecord?> FindActiveByKeyIdAsync(string keyId, CancellationToken ct)
            => Task.FromResult<ApiKeyRecord?>(null);

        /// <summary>Records the call by incrementing <see cref="MarkUsedCount"/>.</summary>
        /// <param name="keyId">The key id that was used.</param>
        /// <param name="whenUtc">The UTC timestamp of use.</param>
        /// <param name="ct">A token to observe for cancellation.</param>
        /// <returns>A task that represents the asynchronous operation.</returns>
        public Task MarkUsedAsync(string keyId, DateTimeOffset whenUtc, CancellationToken ct)
        {
            MarkUsedCount++;
            return Task.CompletedTask;
        }
    }
}