TaskApi/Middleware/RateLimitingMiddleware.cs - master - Kovalenko/TODO

TaskApi/Middleware/RateLimitingMiddleware.cs
99 строк3 KB
Kovalenko1
Done
29 дек 2025, 11:50
29 дек 2025, 11:50f453b3b
using System.Collections.Concurrent;
using System.Globalization;
using System.Security.Claims;
using Microsoft.Extensions.Options;

namespace TaskApi.Middleware;

public class RateLimitOptions
{
    public int PermitLimit { get; init; } = 60;

    public int WindowSeconds { get; init; } = 60;
}

internal sealed class RateLimitState
{
    public object SyncRoot { get; } = new();

    public DateTimeOffset WindowStart { get; set; }

    public int Count { get; set; }
}

public class RateLimitingMiddleware
{
    private readonly RequestDelegate _next;
    private readonly RateLimitOptions _options;
    private readonly ConcurrentDictionary<string, RateLimitState> _states = new();

    public RateLimitingMiddleware(RequestDelegate next, IOptions<RateLimitOptions> options)
    {
        _next = next;
        _options = options.Value;
    }

    public async Task InvokeAsync(HttpContext context)
    {
        var key = GetClientKey(context);
        var now = DateTimeOffset.UtcNow;
        var window = TimeSpan.FromSeconds(Math.Max(1, _options.WindowSeconds));
        var limit = Math.Max(1, _options.PermitLimit);

        var state = _states.GetOrAdd(key, _ => new RateLimitState { WindowStart = now, Count = 0 });
        bool allowed;
        int remaining;
        int retryAfter;

        lock (state.SyncRoot)
        {
            if (now - state.WindowStart >= window)
            {
                state.WindowStart = now;
                state.Count = 0;
            }

            if (state.Count >= limit)
            {
                allowed = false;
                remaining = 0;
            }
            else
            {
                state.Count++;
                allowed = true;
                remaining = limit - state.Count;
            }

            retryAfter = Math.Max(0, (int)Math.Ceiling((window - (now - state.WindowStart)).TotalSeconds));
        }

        context.Response.OnStarting(() =>
        {
            context.Response.Headers["X-Limit-Remaining"] = remaining.ToString(CultureInfo.InvariantCulture);
            context.Response.Headers["Retry-After"] = retryAfter.ToString(CultureInfo.InvariantCulture);
            return Task.CompletedTask;
        });

        if (!allowed)
        {
            context.Response.StatusCode = StatusCodes.Status429TooManyRequests;
            await context.Response.WriteAsync("Rate limit exceeded.");
            return;
        }

        await _next(context);
    }

    private static string GetClientKey(HttpContext context)
    {
        var userId = context.User.FindFirst(ClaimTypes.NameIdentifier)?.Value;
        if (!string.IsNullOrWhiteSpace(userId))
        {
            return $"user:{userId}";
        }

        var ip = context.Connection.RemoteIpAddress?.ToString() ?? "unknown";
        return $"ip:{ip}";
    }
}