101 lines
3.2 KiB
C#
101 lines
3.2 KiB
C#
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;
|
|
}
|
|
}
|
|
}
|