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

using System.Buffers;
using System.Net;
using System.Security.Cryptography;
using System.Text;
using FinalFactory.Rendezvous.Server.Observability;
using Microsoft.Extensions.Options;

namespace FinalFactory.Rendezvous.Server.Abuse;

internal sealed class AbuseProtectionService
{
    private readonly AbuseProtectionOptions _options;
    private readonly TimeProvider _timeProvider;
    private readonly TrackerState _httpTracker;
    private readonly TrackerState _udpTracker;
    private readonly RendezvousTelemetry? _telemetry;
    private readonly HashSet<string> _operatorAllowedAddresses;

    public AbuseProtectionService(
        IOptions<AbuseProtectionOptions> options,
        TimeProvider? timeProvider = null,
        RendezvousTelemetry? telemetry = null)
    {
        _options = options.Value;
        _timeProvider = timeProvider ?? TimeProvider.System;
        _telemetry = telemetry;
        _operatorAllowedAddresses = options.Value.OperatorAllowedAddresses
            .Select(static value => IPAddress.TryParse(value, out IPAddress? address)
                ? NormalizeAddress(address).ToString()
                : string.Empty)
            .Where(static value => value.Length > 0)
            .ToHashSet(StringComparer.Ordinal);
        DateTimeOffset now = _timeProvider.GetUtcNow();
        _httpTracker = new(now);
        _udpTracker = new(now);
    }

    public bool TryAcquireHttpIngress(
        IPAddress? remoteAddress,
        string operation,
        out AbuseLease? lease,
        out int retryAfterSeconds)
    {
        string prefix = GetNetworkPrefix(remoteAddress);
        List<RateDimension> rates =
        [
            new("http:rate:global", _options.HttpGlobalRequestsPerWindow),
            new($"http:rate:ip:{prefix}", _options.HttpIpPrefixRequestsPerWindow),
            new($"http:rate:operation:{operation}", _options.HttpOperationRequestsPerWindow),
        ];
        List<RateDimension> concurrency =
        [
            new("http:concurrency:global", _options.HttpGlobalConcurrency),
            new($"http:concurrency:ip:{prefix}", _options.HttpIpPrefixConcurrency),
            new($"http:concurrency:operation:{operation}", _options.HttpOperationConcurrency),
        ];
        if (!IsLeaseCriticalOperation(operation))
        {
            rates.Add(new("http:rate:optional", _options.HttpOptionalRequestsPerWindow));
            rates.Add(new($"http:rate:optional-ip:{prefix}",
                _options.HttpOptionalIpPrefixRequestsPerWindow));
            concurrency.Add(new("http:concurrency:optional", _options.HttpOptionalConcurrency));
            concurrency.Add(new($"http:concurrency:optional-ip:{prefix}",
                _options.HttpOptionalIpPrefixConcurrency));
        }

        return TryAcquire(
            [.. rates],
            [.. concurrency],
            TrackerDomain.Http,
            IsLeaseCriticalOperation(operation),
            out lease,
            out retryAfterSeconds);
    }

    public bool TryAcquireHealthIngress(
        IPAddress? remoteAddress,
        out AbuseLease? lease,
        out int retryAfterSeconds)
    {
        string prefix = GetNetworkPrefix(remoteAddress);
        RateDimension[] rates =
        [
            new("health:rate:global", _options.HealthGlobalRequestsPerWindow),
            new($"health:rate:ip:{prefix}", _options.HealthIpPrefixRequestsPerWindow),
        ];
        RateDimension[] concurrency =
        [
            new("health:concurrency:global", _options.HealthGlobalConcurrency),
            new($"health:concurrency:ip:{prefix}", _options.HealthIpPrefixConcurrency),
        ];
        return TryAcquire(
            rates,
            concurrency,
            TrackerDomain.Http,
            true,
            out lease,
            out retryAfterSeconds);
    }

    public bool IsOperatorSourceAllowed(IPAddress? remoteAddress) =>
        remoteAddress is not null
        && _operatorAllowedAddresses.Contains(NormalizeAddress(remoteAddress).ToString());

    public bool TryAcquireOperatorIngress(
        IPAddress? remoteAddress,
        out AbuseLease? lease,
        out int retryAfterSeconds)
    {
        string prefix = GetNetworkPrefix(remoteAddress);
        RateDimension[] rates =
        [
            new("operator:rate:global", _options.OperatorGlobalRequestsPerWindow),
            new($"operator:rate:ip:{prefix}", _options.OperatorIpPrefixRequestsPerWindow),
        ];
        RateDimension[] concurrency =
        [
            new("operator:concurrency:global", _options.OperatorGlobalConcurrency),
            new($"operator:concurrency:ip:{prefix}", _options.OperatorIpPrefixConcurrency),
        ];
        return TryAcquire(
            rates,
            concurrency,
            TrackerDomain.Http,
            true,
            out lease,
            out retryAfterSeconds);
    }

    public bool TryAcquireHttpIdentity(
        string operation,
        string? tenant,
        string? principal,
        string? resource,
        out AbuseLease? lease,
        out int retryAfterSeconds) => TryAcquireHttpIdentity(
        operation,
        null,
        tenant,
        principal,
        resource,
        out lease,
        out retryAfterSeconds);

    public bool TryAcquireHttpIdentity(
        string operation,
        IPAddress? remoteAddress,
        string? tenant,
        string? principal,
        string? resource,
        out AbuseLease? lease,
        out int retryAfterSeconds)
    {
        string sourcePrefix = GetNetworkPrefix(remoteAddress);
        List<RateDimension> rates = [];
        List<RateDimension> concurrency = [];
        AddDimension(rates, concurrency, "tenant", tenant,
            _options.HttpTenantRequestsPerWindow, _options.HttpTenantConcurrency);
        AddDimension(rates, concurrency, "principal", principal,
            _options.HttpPrincipalRequestsPerWindow, _options.HttpPrincipalConcurrency);
        AddDimension(rates, concurrency, "resource", resource,
            _options.HttpResourceRequestsPerWindow, _options.HttpResourceConcurrency);
        return TryAcquire(
            [.. rates],
            [.. concurrency],
            TrackerDomain.Http,
            IsLeaseCriticalOperation(operation),
            out lease,
            out retryAfterSeconds);

        void AddDimension(
            List<RateDimension> rateDimensions,
            List<RateDimension> concurrencyDimensions,
            string kind,
            string? value,
            int rateLimit,
            int concurrencyLimit)
        {
            if (string.IsNullOrEmpty(value))
            {
                return;
            }

            if (kind == "resource")
            {
                string validationKey =
                    $"http:source-resource:{operation}:{sourcePrefix}:{value}";
                rateDimensions.Add(new($"{validationKey}:rate", rateLimit));
                concurrencyDimensions.Add(new($"{validationKey}:concurrency", concurrencyLimit));
            }

            string key = kind == "resource"
                ? $"http:resource-scoped:{operation}:{tenant ?? string.Empty}|{principal ?? string.Empty}:{value}"
                : $"http:{kind}:{operation}:{value}";
            rateDimensions.Add(new($"{key}:rate", rateLimit));
            concurrencyDimensions.Add(new($"{key}:concurrency", concurrencyLimit));
        }
    }

    public bool TryAcceptUdpIngress(IPAddress? remoteAddress, string operation)
    {
        string prefix = GetNetworkPrefix(remoteAddress);
        RateDimension[] rates =
        [
            new("udp:rate:global", _options.UdpGlobalDatagramsPerWindow),
            new($"udp:rate:ip:{prefix}", _options.UdpIpPrefixDatagramsPerWindow),
            new($"udp:rate:operation:{operation}", _options.UdpOperationDatagramsPerWindow),
        ];
        return TryAcquire(
                   rates,
                   [],
                   TrackerDomain.Udp,
                   false,
                   out AbuseLease? lease,
                   out _)
               && DisposeAccepted(lease);
    }

    public bool TryAcceptUdpIdentity(
        string operation,
        string capability,
        string resource) => TryAcceptUdpIdentity(
        operation,
        null,
        capability,
        resource);

    public bool TryAcceptUdpIdentity(
        string operation,
        IPAddress? remoteAddress,
        string capability,
        string resource)
    {
        string sourcePrefix = GetNetworkPrefix(remoteAddress);
        string capabilityFingerprint = FingerprintSecret(capability);
        RateDimension[] rates =
        [
            new($"udp:rate:capability:{operation}:{capabilityFingerprint}",
                _options.UdpCapabilityDatagramsPerWindow),
            new($"udp:rate:source-resource:{operation}:{sourcePrefix}:{resource}",
                _options.UdpResourceDatagramsPerWindow),
            new($"udp:rate:resource:{operation}:{capabilityFingerprint}:{resource}",
                _options.UdpResourceDatagramsPerWindow),
        ];
        return TryAcquire(
                   rates,
                   [],
                   TrackerDomain.Udp,
                   false,
                   out AbuseLease? lease,
                   out _)
               && DisposeAccepted(lease);
    }

    public static string FingerprintSecret(string secret)
    {
        int byteCount = Encoding.UTF8.GetByteCount(secret);
        byte[]? rented = null;
        Span<byte> encoded = byteCount <= 1_024
            ? stackalloc byte[byteCount]
            : (rented = ArrayPool<byte>.Shared.Rent(byteCount)).AsSpan(0, byteCount);
        Span<byte> digest = stackalloc byte[32];
        try
        {
            _ = Encoding.UTF8.GetBytes(secret, encoded);
            _ = SHA256.HashData(encoded, digest);
            return Convert.ToHexString(digest[..12]);
        }
        finally
        {
            CryptographicOperations.ZeroMemory(encoded);
            CryptographicOperations.ZeroMemory(digest);
            if (rented is not null)
            {
                ArrayPool<byte>.Shared.Return(rented);
            }
        }
    }

    internal int TrackedKeyCount
    {
        get
        {
            int http;
            int udp;
            lock (_httpTracker.Gate)
            {
                http = _httpTracker.WindowCounts.Count + _httpTracker.ConcurrencyCounts.Count;
            }

            lock (_udpTracker.Gate)
            {
                udp = _udpTracker.WindowCounts.Count + _udpTracker.ConcurrencyCounts.Count;
            }

            return http + udp;
        }
    }

    private bool TryAcquire(
        ReadOnlySpan<RateDimension> rates,
        ReadOnlySpan<RateDimension> concurrency,
        TrackerDomain domain,
        bool canUseCriticalReserve,
        out AbuseLease? lease,
        out int retryAfterSeconds)
    {
        TrackerState tracker = domain == TrackerDomain.Udp ? _udpTracker : _httpTracker;
        bool accepted;
        lock (tracker.Gate)
        {
            accepted = TryAcquireLocked(
                tracker,
                rates,
                concurrency,
                domain,
                canUseCriticalReserve,
                out lease,
                out retryAfterSeconds);
        }

        if (!accepted)
        {
            _telemetry?.RecordLimiterDrop(
                domain == TrackerDomain.Udp ? "udp" : "http",
                "rate-or-concurrency");
        }

        return accepted;
    }

    private bool TryAcquireLocked(
        TrackerState tracker,
        ReadOnlySpan<RateDimension> rates,
        ReadOnlySpan<RateDimension> concurrency,
        TrackerDomain domain,
        bool canUseCriticalReserve,
        out AbuseLease? lease,
        out int retryAfterSeconds)
    {
        DateTimeOffset now = _timeProvider.GetUtcNow();
        TimeSpan window = TimeSpan.FromSeconds(_options.WindowSeconds);
        if (now - tracker.WindowStartedAt >= window || now < tracker.WindowStartedAt)
        {
            tracker.WindowCounts.Clear();
            tracker.WindowStartedAt = now;
        }

        retryAfterSeconds = Math.Max(
            1,
            (int)Math.Ceiling((window - (now - tracker.WindowStartedAt)).TotalSeconds));
        int stagedNewKeys = 0;
        int partitionLimit = domain == TrackerDomain.Udp
            ? _options.UdpTrackedKeyLimit
            : _options.MaxTrackedKeys - _options.UdpTrackedKeyLimit;
        int maxTrackedKeys = domain == TrackerDomain.Udp || canUseCriticalReserve
            ? partitionLimit
            : partitionLimit - _options.CriticalTrackedKeyReserve;
        if (!CanAcquireAll(
                tracker,
                tracker.WindowCounts,
                rates,
                maxTrackedKeys,
                ref stagedNewKeys)
            || !CanAcquireAll(
                tracker,
                tracker.ConcurrencyCounts,
                concurrency,
                maxTrackedKeys,
                ref stagedNewKeys))
        {
            lease = null;
            return false;
        }

        foreach (RateDimension dimension in rates)
        {
            tracker.WindowCounts[dimension.Key] =
                tracker.WindowCounts.GetValueOrDefault(dimension.Key) + 1;
        }

        if (concurrency.IsEmpty)
        {
            lease = null;
            return true;
        }

        string[] acquiredConcurrency = new string[concurrency.Length];
        for (int index = 0; index < concurrency.Length; index++)
        {
            RateDimension dimension = concurrency[index];
            tracker.ConcurrencyCounts[dimension.Key] =
                tracker.ConcurrencyCounts.GetValueOrDefault(dimension.Key) + 1;
            acquiredConcurrency[index] = dimension.Key;
        }

        lease = new AbuseLease(this, tracker, acquiredConcurrency);
        return true;
    }

    private static bool CanAcquireAll(
        TrackerState tracker,
        Dictionary<string, int> counts,
        ReadOnlySpan<RateDimension> dimensions,
        int maxTrackedKeys,
        ref int stagedNewKeys)
    {
        foreach (RateDimension dimension in dimensions)
        {
            if (counts.TryGetValue(dimension.Key, out int current))
            {
                if (current >= dimension.Limit)
                {
                    return false;
                }

                continue;
            }

            stagedNewKeys++;
            if (tracker.WindowCounts.Count + tracker.ConcurrencyCounts.Count + stagedNewKeys
                > maxTrackedKeys)
            {
                return false;
            }
        }

        return true;
    }

    private static void Release(TrackerState tracker, string[] keys)
    {
        lock (tracker.Gate)
        {
            foreach (string key in keys)
            {
                if (!tracker.ConcurrencyCounts.TryGetValue(key, out int current))
                {
                    continue;
                }

                if (current <= 1)
                {
                    tracker.ConcurrencyCounts.Remove(key);
                }
                else
                {
                    tracker.ConcurrencyCounts[key] = current - 1;
                }
            }
        }
    }

    private static bool DisposeAccepted(AbuseLease? lease)
    {
        lease?.Dispose();
        return true;
    }

    private static bool IsLeaseCriticalOperation(string operation) => operation is
        "RenewSessionLease" or "UpdateSession" or "DeleteSession";

    private static string GetNetworkPrefix(IPAddress? address)
    {
        if (address is null)
        {
            return "unknown";
        }

        IPAddress normalized = address.IsIPv4MappedToIPv6 ? address.MapToIPv4() : address;
        byte[] bytes = normalized.GetAddressBytes();
        if (bytes.Length == 4)
        {
            bytes[3] = 0;
            return $"4:{Convert.ToHexString(bytes)}:24";
        }

        if (bytes.Length == 16)
        {
            Array.Clear(bytes, 7, 9);
            return $"6:{Convert.ToHexString(bytes)}:56";
        }

        return "unknown";
    }

    private static IPAddress NormalizeAddress(IPAddress address) =>
        address.IsIPv4MappedToIPv6 ? address.MapToIPv4() : address;

    private readonly record struct RateDimension(string Key, int Limit);

    private enum TrackerDomain
    {
        Http,
        Udp,
    }

    internal sealed class TrackerState(DateTimeOffset windowStartedAt)
    {
        public object Gate { get; } = new();
        public Dictionary<string, int> WindowCounts { get; } = new(StringComparer.Ordinal);
        public Dictionary<string, int> ConcurrencyCounts { get; } = new(StringComparer.Ordinal);
        public DateTimeOffset WindowStartedAt { get; set; } = windowStartedAt;
    }

    internal sealed class AbuseLease : IDisposable
    {
        private AbuseProtectionService? _owner;
        private readonly TrackerState _tracker;
        private readonly string[] _keys;

        internal AbuseLease(
            AbuseProtectionService owner,
            TrackerState tracker,
            string[] keys)
        {
            _owner = owner;
            _tracker = tracker;
            _keys = keys;
        }

        public void Dispose()
        {
            if (Interlocked.Exchange(ref _owner, null) is not null)
            {
                Release(_tracker, _keys);
            }
        }
    }
}