MoShiro/backend/src/MoShiro.Backend.Application/Auth/LoginRateLimiter.cs

using System.Collections.Concurrent;

namespace MoShiro.Backend.Application;

/// <summary>登录限流：每个来源在滑动的一分钟内最多尝试若干次。</summary>
public sealed class LoginRateLimiter
{
    private readonly ConcurrentDictionary<string, Queue<long>> _attempts = new();
    private readonly BackendOptions _options;
    private readonly TimeProvider _time;

    public LoginRateLimiter(BackendOptions options, TimeProvider time)
    {
        _options = options;
        _time = time;
    }

    /// <summary>登记一次尝试；超限时返回 false 与建议的重试等待秒数。<paramref name="limit"/> 缺省是登录失败的额度（每分钟 <see cref="BackendOptions.LoginAttemptsPerMinute"/> 次）。</summary>
    public bool TryAcquire(string source, out int retryAfterSeconds, int? limit = null)
    {
        long now = _time.GetUtcNow().ToUnixTimeMilliseconds();
        Queue<long> queue = _attempts.GetOrAdd(source, _ => new Queue<long>());
        lock (queue)
        {
            while (queue.Count > 0 && now - queue.Peek() >= 60_000)
            {
                queue.Dequeue();
            }

            if (queue.Count >= (limit ?? _options.LoginAttemptsPerMinute))
            {
                retryAfterSeconds = (int)Math.Max(1, (60_000 - (now - queue.Peek()) + 999) / 1000);
                return false;
            }

            queue.Enqueue(now);
            retryAfterSeconds = 0;
            return true;
        }
    }

    /// <summary>清理很久没有尝试的来源，避免字典无限增长。</summary>
    public void Prune()
    {
        long now = _time.GetUtcNow().ToUnixTimeMilliseconds();
        foreach ((string key, Queue<long> queue) in _attempts)
        {
            lock (queue)
            {
                if (queue.Count == 0 || now - queue.Peek() >= 60_000)
                {
                    _attempts.TryRemove(key, out _);
                }
            }
        }
    }
}