git.stella-ops.org/src/StellaOps.Signer/StellaOps.Signer.Infrastructure/Quotas/InMemoryQuotaService.cs

using System;
using System.Collections.Concurrent;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using StellaOps.Signer.Core;

namespace StellaOps.Signer.Infrastructure.Quotas;

public sealed class InMemoryQuotaService : ISignerQuotaService
{
    private readonly ConcurrentDictionary<string, QuotaWindow> _windows = new(StringComparer.Ordinal);
    private readonly TimeProvider _timeProvider;
    private readonly ILogger<InMemoryQuotaService> _logger;

    public InMemoryQuotaService(TimeProvider timeProvider, ILogger<InMemoryQuotaService> logger)
    {
        _timeProvider = timeProvider ?? TimeProvider.System;
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

    public ValueTask EnsureWithinLimitsAsync(
        SigningRequest request,
        ProofOfEntitlementResult entitlement,
        CallerContext caller,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);
        ArgumentNullException.ThrowIfNull(entitlement);
        ArgumentNullException.ThrowIfNull(caller);

        var payloadSize = EstimatePayloadSize(request);
        if (payloadSize > entitlement.MaxArtifactBytes)
        {
            throw new SignerQuotaException("artifact_too_large", $"Artifact size {payloadSize} exceeds plan cap ({entitlement.MaxArtifactBytes}).");
        }

        if (entitlement.QpsLimit <= 0)
        {
            return ValueTask.CompletedTask;
        }

        var window = _windows.GetOrAdd(caller.Tenant, static _ => new QuotaWindow());
        lock (window)
        {
            var now = _timeProvider.GetUtcNow();
            if (window.ResetAt <= now)
            {
                window.Reset(now, entitlement.QpsLimit);
            }

            if (window.Remaining <= 0)
            {
                _logger.LogWarning("Quota exceeded for tenant {Tenant}", caller.Tenant);
                throw new SignerQuotaException("plan_throttled", "Plan QPS limit exceeded.");
            }

            window.Remaining--;
            window.LastUpdated = now;
        }

        return ValueTask.CompletedTask;
    }

    private static int EstimatePayloadSize(SigningRequest request)
    {
        var predicateBytes = request.Predicate is null
            ? Array.Empty<byte>()
            : Encoding.UTF8.GetBytes(request.Predicate.RootElement.GetRawText());

        var subjectBytes = 0;
        foreach (var subject in request.Subjects)
        {
            subjectBytes += subject.Name.Length;
            foreach (var digest in subject.Digest)
            {
                subjectBytes += digest.Key.Length + digest.Value.Length;
            }
        }

        return predicateBytes.Length + subjectBytes;
    }

    private sealed class QuotaWindow
    {
        public DateTimeOffset ResetAt { get; private set; } = DateTimeOffset.MinValue;

        public int Remaining { get; set; }

        public DateTimeOffset LastUpdated { get; set; }

        public void Reset(DateTimeOffset now, int limit)
        {
            ResetAt = now.AddSeconds(1);
            Remaining = limit;
            LastUpdated = now;
        }
    }
}