tiku-backend.net/Tiku.Api/Middleware/AuthRateLimitPartitionMiddleware.cs

using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using Microsoft.AspNetCore.RateLimiting;
using Microsoft.Extensions.Options;
using Tiku.Api.Options;
using Tiku.Application.Security;
using Tiku.Infrastructure.Security;

namespace Tiku.Api.Middleware;

public sealed class AuthRateLimitPartitionMiddleware(
    RequestDelegate next,
    IRedisSecurityStore redisSecurityStore,
    IOptions<AuthRateLimitOptions> options)
{
    public AuthRateLimitPartitionMiddleware(RequestDelegate next)
        : this(
            next,
            new NullRedisSecurityStore(),
            Microsoft.Extensions.Options.Options.Create(new AuthRateLimitOptions()))
    {
    }

    public async Task InvokeAsync(HttpContext context)
    {
        var policy = context.GetEndpoint()?
            .Metadata
            .GetMetadata<EnableRateLimitingAttribute>()?
            .PolicyName;
        var propertyName = policy switch
        {
            AuthRateLimitPolicies.Password => "identifier",
            AuthRateLimitPolicies.Sms => "phone",
            _ => null
        };

        if (HttpMethods.IsPost(context.Request.Method) && propertyName is not null)
        {
            await CaptureAccountHashAsync(context, propertyName);
            if (redisSecurityStore.IsConfigured)
            {
                await ConsumeDistributedLimitAsync(context, policy!);
                if (context.Response.HasStarted) return;
            }
        }

        await next(context);
    }

    private async Task ConsumeDistributedLimitAsync(HttpContext context, string policyName)
    {
        var ip = Hash(context.Connection.RemoteIpAddress?.ToString() ?? "unknown-ip");
        var account = context.Items.TryGetValue(AuthRateLimitPartitionKey.AccountHashItemKey, out var value)
            ? value as string ?? "unknown-account"
            : "unknown-account";
        var isPassword = policyName == AuthRateLimitPolicies.Password;
        var limit = isPassword ? options.Value.PasswordPermitLimit : options.Value.SmsPermitLimit;
        var window = TimeSpan.FromSeconds(isPassword
            ? options.Value.PasswordWindowSeconds
            : options.Value.SmsWindowSeconds);

        DistributedRateLimitResult result;
        try
        {
            result = await redisSecurityStore.ConsumeAsync(
            [
                new DistributedRateLimitBucket($"{policyName}:ip:{ip}", limit * 4, window),
                new DistributedRateLimitBucket($"{policyName}:ip-account:{ip}:{account}", limit, window)
            ], context.RequestAborted);
        }
        catch (RedisSecurityUnavailableException)
        {
            context.Response.StatusCode = StatusCodes.Status503ServiceUnavailable;
            await context.Response.WriteAsJsonAsync(new
            {
                title = "Authentication security service is unavailable.",
                status = StatusCodes.Status503ServiceUnavailable,
                code = "auth_security_unavailable",
                traceId = context.TraceIdentifier
            }, context.RequestAborted);
            return;
        }

        if (!result.Allowed)
        {
            if (result.RetryAfter is { } retryAfter)
                context.Response.Headers.RetryAfter =
                    Math.Max(1, (int)Math.Ceiling(retryAfter.TotalSeconds)).ToString();
            context.Response.StatusCode = StatusCodes.Status429TooManyRequests;
            await context.Response.WriteAsJsonAsync(new
            {
                title = "Too many requests.",
                status = StatusCodes.Status429TooManyRequests,
                code = "rate_limited",
                traceId = context.TraceIdentifier
            }, context.RequestAborted);
        }
    }

    private static async Task CaptureAccountHashAsync(HttpContext context, string propertyName)
    {
        context.Request.EnableBuffering(4096, 16_384);
        try
        {
            using var document = await JsonDocument.ParseAsync(
                context.Request.Body,
                cancellationToken: context.RequestAborted);
            var captured = TryGetStringProperty(document.RootElement, propertyName) ??
                           (propertyName == "identifier" ? TryGetStringProperty(document.RootElement, "phone") : null);
            if (captured is { } value &&
                !string.IsNullOrWhiteSpace(value))
                context.Items[AuthRateLimitPartitionKey.AccountHashItemKey] = Hash(value.Trim());
        }
        catch (JsonException)
        {
            // MVC will produce the canonical malformed JSON response.
        }
        catch (IOException)
        {
            // Oversized or unreadable bodies share the IP-only fallback partition.
        }
        finally
        {
            if (context.Request.Body.CanSeek) context.Request.Body.Position = 0;
        }
    }

    private static string? TryGetStringProperty(JsonElement element, string propertyName)
    {
        if (element.ValueKind != JsonValueKind.Object) return null;

        foreach (var property in element.EnumerateObject())
            if (string.Equals(property.Name, propertyName, StringComparison.OrdinalIgnoreCase) &&
                property.Value.ValueKind == JsonValueKind.String)
                return property.Value.GetString();

        return null;
    }

    private static string Hash(string value)
    {
        return Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(value)))
            .ToLowerInvariant();
    }
}

public static class AuthRateLimitPartitionKey
{
    internal const string AccountHashItemKey = "tiku.auth_rate_limit.account_hash";

    public static string Resolve(HttpContext context, string policyName)
    {
        var ipAddress = context.Connection.RemoteIpAddress?.ToString() ?? "unknown-ip";
        var accountHash = context.Items.TryGetValue(AccountHashItemKey, out var value) &&
                          value is string hash &&
                          !string.IsNullOrWhiteSpace(hash)
            ? hash
            : "unknown-account";
        return $"{policyName}:{ipAddress}:{accountHash}";
    }
}