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>
/// </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).
    private readonly ConcurrentDictionary<string, ConcurrentDictionary<string, byte>> _keyIdIndex =
        new(StringComparer.Ordinal);

    // 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;
    }

    /// <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 below and its repopulation discarded.
        string? tokenKeyId = TryParseKeyId(authorizationHeader);
        long generationAtStart = tokenKeyId is null ? 0 : ReadGeneration(tokenKeyId);

        ApiKeyVerification result = await _inner.VerifyAsync(authorizationHeader, ct).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;
        _cache.Set(cacheKey, result, new MemoryCacheEntryOptions
        {
            AbsoluteExpirationRelativeToNow = _ttl,
        });
        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);
            if (_keyIdIndex.TryGetValue(keyId, out ConcurrentDictionary<string, byte>? set))
            {
                set.TryRemove(cacheKey, out _);
            }
        }
    }

    /// <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);

        if (_keyIdIndex.TryRemove(keyId, out ConcurrentDictionary<string, byte>? cacheKeys))
        {
            foreach (string cacheKey in cacheKeys.Keys)
            {
                _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]);
    }

    private void IndexCacheKey(string keyId, string cacheKey)
    {
        ConcurrentDictionary<string, byte> set = _keyIdIndex.GetOrAdd(
            keyId,
            static _ => new ConcurrentDictionary<string, byte>(StringComparer.Ordinal));
        set.TryAdd(cacheKey, 0);
    }

    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;
    }
}