mxaccessgw/src/ZB.MOM.WW.MxGateway.Server/Security/Authentication/CachingApiKeyVerifier.cs

using System.Collections.Concurrent;
using System.Security.Cryptography;
using System.Text;
using Microsoft.Extensions.Caching.Memory;
using ZB.MOM.WW.Auth.Abstractions.ApiKeys;
using ZB.MOM.WW.MxGateway.Server.Configuration;

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

/// <summary>
/// Lets gateway-side admin call sites drop a cached API-key verification when they revoke, rotate
/// or delete a key, so a mutation applied through the running gateway process takes effect
/// immediately rather than after the cache TTL elapses.
/// </summary>
public interface IApiKeyCacheInvalidator
{
    /// <summary>Removes every cached verification for the given key id.</summary>
    /// <param name="keyId">The key id whose cached verifications should be dropped.</param>
    void Invalidate(string keyId);
}

/// <summary>
/// An <see cref="IApiKeyVerifier"/> decorator that caches successful verifications for a short,
/// configurable TTL (<see cref="SecurityOptions.ApiKeyVerificationCacheSeconds"/>), keyed on a
/// SHA-256 hash of the presented token. A cache hit returns the cached result without calling the
/// inner (library) verifier, so it avoids both the per-call SQLite read and the <c>last_used</c>
/// write the library verifier couples into <see cref="IApiKeyVerifier.VerifyAsync"/>.
/// </summary>
/// <remarks>
/// <para>
/// Only successful verifications are cached. Failures and unparseable headers always fall through
/// to the inner verifier — caching a failure risks pinning a transiently-wrong negative, and the
/// per-peer failure counter (not this cache) is what bounds brute-force cost.
/// </para>
/// <para>
/// The cache key is the hex SHA-256 of the presented token (which embeds both the key id and the
/// secret). Hashing means the plaintext secret is never stored in memory. The hash is a cache index
/// only; the authentic constant-time secret comparison remains inside the inner verifier, reached
/// on every cache miss.
/// </para>
/// <para>
/// Gateway-initiated revoke/rotate/delete call <see cref="Invalidate"/> directly (see
/// <c>DashboardApiKeyManagementService</c>), which bumps a per-key generation counter <em>before</em>
/// evicting so an in-flight verification that started under the old generation discards its own
/// repopulation (set-then-recheck in <c>VerifyAsync</c>) — making the "gateway-initiated mutations
/// take effect immediately" contract true even against a verify that was already in the inner library
/// when the revoke landed.
/// </para>
/// <para>
/// The short TTL is the backstop for two remaining bounded-staleness windows. (1) Out-of-band
/// mutations — a direct DB edit, or a revoke issued by the separate <c>apikey</c> CLI process whose
/// in-memory cache is not this process's cache — take effect only after the TTL elapses. (2) A key
/// whose expiry passes while cached keeps authenticating for up to the TTL: expiry is enforced by the
/// inner library verifier, which a cache hit never reaches, and the verification identity the library
/// returns (<c>ZB.MOM.WW.Auth.Abstractions.ApiKeys.ApiKeyIdentity</c>) carries no expiry timestamp, so
/// the cache cannot cap an entry at the key's <c>ExpiresUtc</c>. Capping it requires the donor library
/// to surface expiry on the verification identity (donor-library ask); until then the TTL bounds this
/// window and it is intentionally kept short (<see cref="SecurityOptions.ApiKeyVerificationCacheSeconds"/>,
/// default 15 s).
/// </para>
/// <para>
/// Misses are single-flight: concurrent misses for the same token and the same key generation share
/// one inner verification, so a TTL expiry under load costs one SQLite read rather than one per
/// in-flight RPC. A miss that starts after an <see cref="Invalidate"/> never joins a flight that
/// started before it (the flight is keyed on the generation too), so the revoke guarantee holds for
/// joiners as well. The shared inner call is not bound to any one caller's cancellation token; each
/// caller stops waiting on its own token without cancelling the others.
/// </para>
/// <para>
/// The key-id index is pruned by a post-eviction callback whenever the cache drops an entry on its
/// own (expiry, compaction), not only by <see cref="Invalidate"/>, so it stays proportional to the
/// live cache rather than to every token ever verified.
/// </para>
/// </remarks>
public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalidator
{
    private const string CacheKeyPrefix = "mxgw:apikeyverif:";
    private const string TokenPrefix = "mxgw";

    private readonly IApiKeyVerifier _inner;
    private readonly IMemoryCache _cache;
    private readonly TimeSpan _ttl;

    // keyId -> set of cache keys currently held for that id, so a revoke/rotate can evict every
    // secret variant (e.g. an old secret still within TTL after a rotate). Guarded by _indexSync,
    // which also covers the cache Set that an index entry describes: the post-eviction callback
    // runs asynchronously, so Set + index and check-cache + unindex must each be atomic with
    // respect to the other, or a late callback could unindex a freshly re-cached entry and leave
    // it out of reach of a later Invalidate. Every operation under the lock is O(1) and the lock is
    // only taken on misses, evictions and invalidations — never on a cache hit.
    private readonly object _indexSync = new();
    private readonly Dictionary<string, HashSet<string>> _keyIdIndex = new(StringComparer.Ordinal);

    // (cache key, key generation) -> the in-flight inner verification for that pair (single-flight).
    private readonly ConcurrentDictionary<(string CacheKey, long Generation), Lazy<Task<ApiKeyVerification>>> _inflight =
        new();

    private readonly PostEvictionDelegate _onEvicted;

    // keyId -> monotonic generation, bumped by Invalidate before it evicts. A VerifyAsync snapshots
    // the generation before calling the inner verifier and only caches if it is unchanged after the
    // inner call and again after the Set — closing the revoke-vs-in-flight-repopulation race.
    private readonly ConcurrentDictionary<string, long> _generations = new(StringComparer.Ordinal);

    /// <summary>Initializes a new instance of the <see cref="CachingApiKeyVerifier"/> class.</summary>
    /// <param name="inner">The wrapped verifier (the library verifier) reached on a cache miss.</param>
    /// <param name="cache">The shared memory cache.</param>
    /// <param name="security">Security options carrying the verification-cache TTL.</param>
    public CachingApiKeyVerifier(IApiKeyVerifier inner, IMemoryCache cache, SecurityOptions security)
        : this(
            inner,
            cache,
            TimeSpan.FromSeconds((security ?? throw new ArgumentNullException(nameof(security)))
                .ApiKeyVerificationCacheSeconds))
    {
    }

    /// <summary>Initializes a new instance of the <see cref="CachingApiKeyVerifier"/> class with an explicit TTL.</summary>
    /// <param name="inner">The wrapped verifier (the library verifier) reached on a cache miss.</param>
    /// <param name="cache">The shared memory cache.</param>
    /// <param name="ttl">The verification-cache TTL; a zero or negative value disables caching.</param>
    /// <remarks>Test/explicit-TTL seam.</remarks>
    internal CachingApiKeyVerifier(IApiKeyVerifier inner, IMemoryCache cache, TimeSpan ttl)
    {
        ArgumentNullException.ThrowIfNull(inner);
        ArgumentNullException.ThrowIfNull(cache);
        _inner = inner;
        _cache = cache;
        _ttl = ttl;
        _onEvicted = OnCacheEntryEvicted;
    }

    /// <summary>Gets the number of cache keys currently held in the key-id index (test seam).</summary>
    internal int IndexedCacheKeyCount
    {
        get
        {
            lock (_indexSync)
            {
                int count = 0;
                foreach (HashSet<string> set in _keyIdIndex.Values)
                {
                    count += set.Count;
                }

                return count;
            }
        }
    }

    /// <summary>Gets the number of key ids currently present in the key-id index (test seam).</summary>
    internal int IndexedKeyIdCount
    {
        get
        {
            lock (_indexSync)
            {
                return _keyIdIndex.Count;
            }
        }
    }

    /// <summary>
    /// Verifies the given authorization header, returning a cached successful result when one is
    /// still within TTL, or falling through to the inner verifier on a cache miss or non-cacheable
    /// header.
    /// </summary>
    /// <param name="authorizationHeader">The raw <c>Authorization</c> header value to verify.</param>
    /// <param name="ct">Token to cancel the asynchronous operation.</param>
    /// <returns>The verification result, cached or freshly computed.</returns>
    public async Task<ApiKeyVerification> VerifyAsync(string authorizationHeader, CancellationToken ct)
    {
        if (_ttl <= TimeSpan.Zero || !TryComputeCacheKey(authorizationHeader, out string cacheKey))
        {
            return await _inner.VerifyAsync(authorizationHeader, ct).ConfigureAwait(false);
        }

        if (_cache.TryGetValue(cacheKey, out ApiKeyVerification? cached) && cached is not null)
        {
            return cached;
        }

        // Snapshot the key's generation BEFORE the inner verify so a concurrent Invalidate that runs
        // while we are in the library is detected and its repopulation discarded. The snapshot is
        // also half of the single-flight key, so a miss that starts after an Invalidate cannot join
        // a flight that started before it.
        string? tokenKeyId = TryParseKeyId(authorizationHeader);
        long generationAtStart = tokenKeyId is null ? 0 : ReadGeneration(tokenKeyId);
        (string CacheKey, long Generation) flightKey = (cacheKey, generationAtStart);

        Lazy<Task<ApiKeyVerification>> created = new(
            () => VerifyMissAsync(authorizationHeader, cacheKey, tokenKeyId, generationAtStart));
        Lazy<Task<ApiKeyVerification>> flight = _inflight.GetOrAdd(flightKey, created);
        Task<ApiKeyVerification> verification = flight.Value;
        if (ReferenceEquals(flight, created))
        {
            // The flight's creator retires it once it completes (success, failure or fault), so the
            // next miss for this token starts fresh. Removing by (key, value) pair never retires a
            // newer flight that replaced this one.
            _ = verification.ContinueWith(
                static (_, state) =>
                {
                    var (owner, key, completed) =
                        ((CachingApiKeyVerifier, (string, long), Lazy<Task<ApiKeyVerification>>))state!;
                    owner._inflight.TryRemove(
                        new KeyValuePair<(string CacheKey, long Generation), Lazy<Task<ApiKeyVerification>>>(key, completed));
                },
                (this, flightKey, flight),
                CancellationToken.None,
                TaskContinuationOptions.ExecuteSynchronously,
                TaskScheduler.Default);
        }

        return await verification.WaitAsync(ct).ConfigureAwait(false);
    }

    // The shared body of one single-flight miss. Not bound to any caller's token: a caller that
    // gives up must not cancel the verification other callers are waiting on.
    private async Task<ApiKeyVerification> VerifyMissAsync(
        string authorizationHeader,
        string cacheKey,
        string? tokenKeyId,
        long generationAtStart)
    {
        ApiKeyVerification result = await _inner
            .VerifyAsync(authorizationHeader, CancellationToken.None)
            .ConfigureAwait(false);

        if (result.Succeeded && result.Identity is not null)
        {
            TryCacheSuccess(cacheKey, result, tokenKeyId, generationAtStart);
        }

        return result;
    }

    // Caches a successful verification unless the key was invalidated while the inner verify was in
    // flight (detected via the generation snapshot). The entry lifetime stays the
    // TTL: the library verification identity carries no expiry, so the cache cannot cap at the key's
    // ExpiresUtc — that window remains TTL-bounded and documented in the class remarks.
    private void TryCacheSuccess(
        string cacheKey,
        ApiKeyVerification result,
        string? tokenKeyId,
        long generationAtStart)
    {
        // If Invalidate bumped the generation while we were verifying, the identity we hold may be
        // stale — do not cache it.
        if (tokenKeyId is not null && ReadGeneration(tokenKeyId) != generationAtStart)
        {
            return;
        }

        string keyId = result.Identity!.KeyId;
        MemoryCacheEntryOptions entryOptions = new()
        {
            AbsoluteExpirationRelativeToNow = _ttl,
        };
        entryOptions.PostEvictionCallbacks.Add(new PostEvictionCallbackRegistration
        {
            EvictionCallback = _onEvicted,
            State = keyId,
        });

        lock (_indexSync)
        {
            _cache.Set(cacheKey, result, entryOptions);
            IndexCacheKey(keyId, cacheKey);
        }

        // Set-then-recheck: an Invalidate that landed between the pre-Set check and IndexCacheKey
        // would have missed this entry (it was not yet indexed). If the generation has moved, evict
        // the just-written entry so the revoke still takes effect immediately.
        if (tokenKeyId is not null && ReadGeneration(tokenKeyId) != generationAtStart)
        {
            _cache.Remove(cacheKey);
            lock (_indexSync)
            {
                UnindexCacheKey(keyId, cacheKey);
            }
        }
    }

    // Post-eviction callback: drops the evicted entry's key-id index slot so the index tracks the
    // live cache. MemoryCache raises it asynchronously, so by the time it runs the same cache key
    // may already hold a newer entry (re-verified after the expiry); that entry owns the slot, and
    // the check below runs under the lock that also covers Set + index, so it cannot interleave
    // with a re-cache.
    private void OnCacheEntryEvicted(object key, object? value, EvictionReason reason, object? state)
    {
        if (reason == EvictionReason.Replaced || key is not string cacheKey || state is not string keyId)
        {
            // Replaced: a newer entry under the same key was Set in this entry's place and is
            // indexed under the same slot already.
            return;
        }

        lock (_indexSync)
        {
            if (_cache.TryGetValue(cacheKey, out _))
            {
                return;
            }

            UnindexCacheKey(keyId, cacheKey);
        }
    }

    /// <inheritdoc />
    public void Invalidate(string keyId)
    {
        if (string.IsNullOrEmpty(keyId))
        {
            return;
        }

        // Bump the generation BEFORE evicting so an in-flight VerifyAsync that snapshotted the old
        // generation refuses to cache (or self-evicts) its now-stale repopulation.
        _generations.AddOrUpdate(keyId, 1, static (_, current) => current + 1);

        string[] cacheKeys;
        lock (_indexSync)
        {
            if (!_keyIdIndex.Remove(keyId, out HashSet<string>? indexed))
            {
                return;
            }

            cacheKeys = [.. indexed];
        }

        // Evicted outside the lock: the eviction callbacks these raise take it themselves (and find
        // the slot already gone).
        foreach (string cacheKey in cacheKeys)
        {
            _cache.Remove(cacheKey);
        }
    }

    private long ReadGeneration(string keyId) => _generations.TryGetValue(keyId, out long generation)
        ? generation
        : 0;

    /// <summary>
    /// Parses the key id out of a <c>Bearer mxgw_&lt;keyId&gt;_&lt;secret&gt;</c> header without any store access,
    /// using the same split the authorization interceptor does.
    /// </summary>
    /// <param name="authorizationHeader">The raw authorization header value, or <see langword="null"/>.</param>
    /// <returns>The full key id, or <see langword="null"/> for a header this cache cannot attribute to a key id (the generation race-guard is then not applied).</returns>
    // Implementation note:
    // Correctness of the generation guard rests on parts[1] being the FULL key id: the '_'
    // split would truncate a key id that itself contained '_', silently disarming the guard for that
    // key. This is safe because '_' is the token's field delimiter and both — and the only — key
    // creation paths in the gateway forbid it: ApiKeyAdminCommandLineParser.IsValidKeyId and
    // DashboardApiKeyManagementService.ValidateKeyId each restrict a key id to
    // char.IsAsciiLetterOrDigit || '.' || '-'. Key ids are never library-generated, so no path can
    // mint one containing '_'.
    //
    // Internal rather than private so the parse rules can be pinned directly by test: the guard's
    // correctness depends on this returning the full key id.
    internal static string? TryParseKeyId(string? authorizationHeader)
    {
        if (string.IsNullOrEmpty(authorizationHeader))
        {
            return null;
        }

        ReadOnlySpan<char> header = authorizationHeader.AsSpan().Trim();
        const string bearer = "Bearer ";
        ReadOnlySpan<char> token = header.StartsWith(bearer, StringComparison.OrdinalIgnoreCase)
            ? header[bearer.Length..].Trim()
            : header;

        // Scanned rather than split, for the same reason as the interceptor's copy: Split would
        // allocate a token copy, an array and a string per segment on every cache miss to produce
        // one key id. Two IndexOf scans allocate only that key id.
        //
        // The '_' checked immediately after the prefix is what makes "mxgw" the whole first segment
        // (so "mxgwabc_..." is still rejected), and the second separator must exist because the
        // split form required three segments — a token with no secret delimiter is not a key token.
        if (!token.StartsWith(TokenPrefix, StringComparison.Ordinal)
            || token.Length <= TokenPrefix.Length
            || token[TokenPrefix.Length] != '_')
        {
            return null;
        }

        ReadOnlySpan<char> afterPrefix = token[(TokenPrefix.Length + 1)..];
        int separator = afterPrefix.IndexOf('_');
        if (separator <= 0)
        {
            // -1 is a token with no second separator; 0 is an empty key id.
            return null;
        }

        return new string(afterPrefix[..separator]);
    }

    // Caller holds _indexSync.
    private void IndexCacheKey(string keyId, string cacheKey)
    {
        if (!_keyIdIndex.TryGetValue(keyId, out HashSet<string>? set))
        {
            set = new HashSet<string>(StringComparer.Ordinal);
            _keyIdIndex[keyId] = set;
        }

        set.Add(cacheKey);
    }

    // Caller holds _indexSync. Drops the key id's slot entirely once its last cache key goes, so
    // the index does not accumulate empty sets for keys that are no longer cached.
    private void UnindexCacheKey(string keyId, string cacheKey)
    {
        if (_keyIdIndex.TryGetValue(keyId, out HashSet<string>? set)
            && set.Remove(cacheKey)
            && set.Count == 0)
        {
            _keyIdIndex.Remove(keyId);
        }
    }

    private static bool TryComputeCacheKey(string? authorizationHeader, out string cacheKey)
    {
        cacheKey = string.Empty;
        if (string.IsNullOrEmpty(authorizationHeader))
        {
            return false;
        }

        ReadOnlySpan<char> header = authorizationHeader.AsSpan().Trim();
        const string bearer = "Bearer ";
        ReadOnlySpan<char> token = header.StartsWith(bearer, StringComparison.OrdinalIgnoreCase)
            ? header[bearer.Length..].Trim()
            : header;

        if (token.IsEmpty)
        {
            return false;
        }

        byte[] hash = SHA256.HashData(Encoding.UTF8.GetBytes(token.ToString()));
        cacheKey = CacheKeyPrefix + Convert.ToHexString(hash);
        return true;
    }
}