Rendezvous/src/FinalFactory.Rendezvous.Server/Abuse/HttpAbuseProtectionMiddleware.cs

using FinalFactory.Rendezvous.Contracts;
using Microsoft.AspNetCore.Http.Features;

namespace FinalFactory.Rendezvous.Server.Abuse;

internal sealed class HttpAbuseProtectionMiddleware(
    RequestDelegate next,
    AbuseProtectionService protection)
{
    public async Task InvokeAsync(HttpContext context)
    {
        IHttpMaxRequestBodySizeFeature? bodySize =
            context.Features.Get<IHttpMaxRequestBodySizeFeature>();
        if (bodySize is { IsReadOnly: false })
        {
            bodySize.MaxRequestBodySize = ContractLimits.HttpRequestMaxBytes;
        }

        string operation = context.GetEndpoint()?.Metadata.GetMetadata<IEndpointNameMetadata>()
            ?.EndpointName ?? "Unmatched";
        bool healthEndpoint = operation is "GetLiveness" or "GetReadiness";
        bool operatorEndpoint = operation is
            "GetOperatorStatus"
            or "RevokeOperatorListing"
            or "RevokeOperatorPrincipal"
            or "RevokeOperatorSigningKey"
            or "BeginOperatorDrain";
        if (operatorEndpoint
            && !protection.IsOperatorSourceAllowed(context.Connection.RemoteIpAddress))
        {
            bool deniedSourceAdmitted = protection.TryAcquireHttpIngress(
                context.Connection.RemoteIpAddress,
                "Unmatched",
                out AbuseProtectionService.AbuseLease? deniedSourceLease,
                out int deniedRetryAfterSeconds);
            using (deniedSourceLease)
            {
                if (!deniedSourceAdmitted)
                {
                    context.Response.Headers.RetryAfter = deniedRetryAfterSeconds.ToString(
                        System.Globalization.CultureInfo.InvariantCulture);
                    await WriteErrorAsync(
                        context,
                        StatusCodes.Status429TooManyRequests,
                        RendezvousErrorCode.RateLimited,
                        "The request rate limit was exceeded.",
                        deniedRetryAfterSeconds).ConfigureAwait(false);
                    return;
                }

                await WriteErrorAsync(
                    context,
                    StatusCodes.Status404NotFound,
                    RendezvousErrorCode.NotFound,
                    "The requested resource was not found.").ConfigureAwait(false);
            }

            return;
        }

        AbuseProtectionService.AbuseLease? lease;
        int retryAfterSeconds;
        bool acquired;
        if (healthEndpoint)
        {
            acquired = protection.TryAcquireHealthIngress(
                context.Connection.RemoteIpAddress,
                out lease,
                out retryAfterSeconds);
        }
        else if (operatorEndpoint)
        {
            acquired = protection.TryAcquireOperatorIngress(
                context.Connection.RemoteIpAddress,
                out lease,
                out retryAfterSeconds);
        }
        else
        {
            acquired = protection.TryAcquireHttpIngress(
                context.Connection.RemoteIpAddress,
                operation,
                out lease,
                out retryAfterSeconds);
        }

        if (!acquired)
        {
            context.Response.Headers.RetryAfter = retryAfterSeconds.ToString(
                System.Globalization.CultureInfo.InvariantCulture);
            await WriteErrorAsync(
                context,
                StatusCodes.Status429TooManyRequests,
                RendezvousErrorCode.RateLimited,
                "The request rate limit was exceeded.",
                retryAfterSeconds).ConfigureAwait(false);
            return;
        }

        using (lease)
        {
            if (context.Request.ContentLength > ContractLimits.HttpRequestMaxBytes)
            {
                await WriteErrorAsync(
                    context,
                    StatusCodes.Status413PayloadTooLarge,
                    RendezvousErrorCode.InvalidRequest,
                    "The request body exceeds the supported size.").ConfigureAwait(false);
                return;
            }

            await next(context).ConfigureAwait(false);
        }
    }

    private static Task WriteErrorAsync(
        HttpContext context,
        int status,
        RendezvousErrorCode code,
        string message,
        int? retryAfterSeconds = null)
    {
        context.Response.StatusCode = status;
        return context.Response.WriteAsJsonAsync(
            new ApiError
            {
                Code = code,
                Message = message,
                RetryAfterSeconds = retryAfterSeconds,
            },
            ContractJson.Options,
            contentType: "application/json",
            cancellationToken: context.RequestAborted);
    }
}