mxaccessgw/src/ZB.MOM.WW.MxGateway.Server/Security/Authorization/ApiKeyFailureLimiter.cs

using System.Collections.Concurrent;
using ZB.MOM.WW.MxGateway.Server.Configuration;

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

/// <summary>
/// Cheap, in-process per-peer sliding-window failure counter for the gRPC auth path (SEC-11). It is
/// checked BEFORE the API-key verification store read and short-circuits a peer that has exceeded
/// <see cref="SecurityOptions.ApiKeyFailureLimit"/> failed attempts within
/// <see cref="SecurityOptions.ApiKeyFailureWindowSeconds"/>; a successful verification resets the
/// peer's counter.
/// </summary>
/// <remarks>
/// <para>
/// The peer key is the API key id when the presented token parses, falling back to the transport
/// peer address otherwise. Keying on key id (per the NAT caveat in <c>glauth.md</c>) means a single
/// abusive credential behind a shared NAT is throttled without locking out unrelated clients on the
/// same address.
/// </para>
/// <para>
/// The tracked-peer set is a bounded LRU (<see cref="SecurityOptions.ApiKeyFailureTrackedPeers"/>) so
/// a spray of unique peer keys cannot grow memory without limit. Successful peers are removed on
/// reset, so in steady state the map holds only peers with recent failures — the common success path
/// is a lock-free dictionary miss.
/// </para>
/// </remarks>
public sealed class ApiKeyFailureLimiter
{
    private readonly int _limit;
    private readonly long _windowTicks;
    private readonly int _maxPeers;
    private readonly TimeProvider _clock;

    private readonly ConcurrentDictionary<string, PeerState> _peers = new(StringComparer.Ordinal);

    /// <summary>Initializes a new instance of the <see cref="ApiKeyFailureLimiter"/> class.</summary>
    /// <param name="security">Security options carrying the failure-limit knobs.</param>
    /// <param name="clock">The time provider.</param>
    public ApiKeyFailureLimiter(SecurityOptions security, TimeProvider clock)
        : this(
            (security ?? throw new ArgumentNullException(nameof(security))).ApiKeyFailureLimit,
            TimeSpan.FromSeconds(security.ApiKeyFailureWindowSeconds),
            security.ApiKeyFailureTrackedPeers,
            clock)
    {
    }

    // Test/explicit seam.
    internal ApiKeyFailureLimiter(int limit, TimeSpan window, int maxPeers, TimeProvider clock)
    {
        ArgumentNullException.ThrowIfNull(clock);
        _limit = limit;
        _windowTicks = window.Ticks;
        _maxPeers = maxPeers;
        _clock = clock;
    }

    /// <summary>Returns whether the peer has reached the failure limit within the current window.</summary>
    /// <param name="peer">The peer key (key id or peer address).</param>
    /// <returns><see langword="true"/> when the peer should be short-circuited.</returns>
    public bool IsBlocked(string peer)
    {
        ArgumentNullException.ThrowIfNull(peer);
        if (_limit <= 0)
        {
            return false;
        }

        if (!_peers.TryGetValue(peer, out PeerState? state))
        {
            return false;
        }

        long now = _clock.GetUtcNow().UtcTicks;
        lock (state)
        {
            Prune(state, now);
            return state.FailureTicks.Count >= _limit;
        }
    }

    /// <summary>Records a failed verification attempt for the peer.</summary>
    /// <param name="peer">The peer key (key id or peer address).</param>
    public void RecordFailure(string peer)
    {
        ArgumentNullException.ThrowIfNull(peer);
        if (_limit <= 0)
        {
            return;
        }

        long now = _clock.GetUtcNow().UtcTicks;
        PeerState state = _peers.GetOrAdd(peer, static _ => new PeerState());
        lock (state)
        {
            Prune(state, now);
            state.FailureTicks.Enqueue(now);
            state.LastActivityTicks = now;
        }

        EvictIfOverCapacity();
    }

    /// <summary>Clears the peer's failure count after a successful verification.</summary>
    /// <param name="peer">The peer key (key id or peer address).</param>
    public void Reset(string peer)
    {
        ArgumentNullException.ThrowIfNull(peer);
        _peers.TryRemove(peer, out _);
    }

    private void Prune(PeerState state, long now)
    {
        while (state.FailureTicks.Count > 0 && now - state.FailureTicks.Peek() >= _windowTicks)
        {
            state.FailureTicks.Dequeue();
        }
    }

    private void EvictIfOverCapacity()
    {
        // Best-effort eviction: only runs when the map exceeds the cap (rare, since only peers with
        // recent failures are tracked). Removes the least-recently-active peer. Racy under
        // concurrency, which is acceptable for a bound rather than an exact policy.
        while (_peers.Count > _maxPeers)
        {
            string? oldest = null;
            long oldestTicks = long.MaxValue;
            foreach (KeyValuePair<string, PeerState> entry in _peers)
            {
                long activity = Volatile.Read(ref entry.Value.LastActivityTicks);
                if (activity < oldestTicks)
                {
                    oldestTicks = activity;
                    oldest = entry.Key;
                }
            }

            if (oldest is null || !_peers.TryRemove(oldest, out _))
            {
                break;
            }
        }
    }

    private sealed class PeerState
    {
        public Queue<long> FailureTicks { get; } = new();

        public long LastActivityTicks;
    }
}