|
| 1 | +// Copyright (c) EasyAppDev. All rights reserved. |
| 2 | +// Licensed under the MIT License. |
| 3 | + |
| 4 | +namespace EasyAppDev.Blazor.AutoComplete.AI.RateLimiting; |
| 5 | + |
| 6 | +/// <summary> |
| 7 | +/// Thread-safe rate limiter using token bucket algorithm. |
| 8 | +/// Prevents cost amplification attacks by limiting the rate of API calls. |
| 9 | +/// </summary> |
| 10 | +/// <remarks> |
| 11 | +/// <para> |
| 12 | +/// The token bucket algorithm maintains a fixed-size bucket of tokens. |
| 13 | +/// Each API call consumes one token. Tokens are refilled at a constant rate. |
| 14 | +/// When the bucket is empty, requests are rejected or delayed. |
| 15 | +/// </para> |
| 16 | +/// <para> |
| 17 | +/// <b>Security Features:</b> |
| 18 | +/// <list type="bullet"> |
| 19 | +/// <item>Thread-safe for concurrent access using SemaphoreSlim</item> |
| 20 | +/// <item>Prevents API quota exhaustion</item> |
| 21 | +/// <item>Protects against cost amplification attacks</item> |
| 22 | +/// <item>Configurable per-component or per-user limits</item> |
| 23 | +/// </list> |
| 24 | +/// </para> |
| 25 | +/// </remarks> |
| 26 | +public class RateLimiter : IDisposable |
| 27 | +{ |
| 28 | + private readonly int _maxTokens; |
| 29 | + private readonly TimeSpan _refillInterval; |
| 30 | + private readonly SemaphoreSlim _lock = new(1, 1); |
| 31 | + private int _currentTokens; |
| 32 | + private DateTime _lastRefill; |
| 33 | + private bool _disposed; |
| 34 | + |
| 35 | + /// <summary> |
| 36 | + /// Initializes a new rate limiter. |
| 37 | + /// </summary> |
| 38 | + /// <param name="maxRequestsPerMinute">Maximum requests allowed per minute. Must be greater than 0.</param> |
| 39 | + /// <exception cref="ArgumentOutOfRangeException">Thrown when maxRequestsPerMinute is less than or equal to 0.</exception> |
| 40 | + /// <example> |
| 41 | + /// <code> |
| 42 | + /// // Allow 60 requests per minute (1 per second) |
| 43 | + /// var limiter = new RateLimiter(maxRequestsPerMinute: 60); |
| 44 | + /// |
| 45 | + /// // Allow 10 requests per minute (conservative for multi-user apps) |
| 46 | + /// var limiter = new RateLimiter(maxRequestsPerMinute: 10); |
| 47 | + /// </code> |
| 48 | + /// </example> |
| 49 | + public RateLimiter(int maxRequestsPerMinute) |
| 50 | + { |
| 51 | + if (maxRequestsPerMinute <= 0) |
| 52 | + { |
| 53 | + throw new ArgumentOutOfRangeException( |
| 54 | + nameof(maxRequestsPerMinute), |
| 55 | + "Must be greater than 0"); |
| 56 | + } |
| 57 | + |
| 58 | + _maxTokens = maxRequestsPerMinute; |
| 59 | + _currentTokens = maxRequestsPerMinute; |
| 60 | + _refillInterval = TimeSpan.FromMinutes(1.0 / maxRequestsPerMinute); |
| 61 | + _lastRefill = DateTime.UtcNow; |
| 62 | + } |
| 63 | + |
| 64 | + /// <summary> |
| 65 | + /// Attempts to acquire a token for making a request. |
| 66 | + /// </summary> |
| 67 | + /// <param name="cancellationToken">Cancellation token</param> |
| 68 | + /// <returns>True if token acquired, false if rate limit exceeded</returns> |
| 69 | + /// <remarks> |
| 70 | + /// This method returns immediately. If no tokens are available, it returns false. |
| 71 | + /// Use <see cref="AcquireAsync"/> if you want to wait for a token to become available. |
| 72 | + /// </remarks> |
| 73 | + /// <exception cref="ObjectDisposedException">Thrown if the rate limiter has been disposed.</exception> |
| 74 | + public async Task<bool> TryAcquireAsync(CancellationToken cancellationToken = default) |
| 75 | + { |
| 76 | + ObjectDisposedException.ThrowIf(_disposed, this); |
| 77 | + |
| 78 | + await _lock.WaitAsync(cancellationToken); |
| 79 | + try |
| 80 | + { |
| 81 | + RefillTokens(); |
| 82 | + |
| 83 | + if (_currentTokens > 0) |
| 84 | + { |
| 85 | + _currentTokens--; |
| 86 | + return true; |
| 87 | + } |
| 88 | + |
| 89 | + return false; |
| 90 | + } |
| 91 | + finally |
| 92 | + { |
| 93 | + _lock.Release(); |
| 94 | + } |
| 95 | + } |
| 96 | + |
| 97 | + /// <summary> |
| 98 | + /// Waits until a token is available, then acquires it. |
| 99 | + /// </summary> |
| 100 | + /// <param name="timeout">Maximum time to wait. If null, waits indefinitely.</param> |
| 101 | + /// <param name="cancellationToken">Cancellation token</param> |
| 102 | + /// <exception cref="RateLimitExceededException">Thrown if timeout expires before a token becomes available.</exception> |
| 103 | + /// <exception cref="ObjectDisposedException">Thrown if the rate limiter has been disposed.</exception> |
| 104 | + /// <remarks> |
| 105 | + /// This method uses adaptive backoff to minimize CPU usage while waiting. |
| 106 | + /// It calculates the next refill time and waits until then before retrying. |
| 107 | + /// </remarks> |
| 108 | + public async Task AcquireAsync( |
| 109 | + TimeSpan? timeout = null, |
| 110 | + CancellationToken cancellationToken = default) |
| 111 | + { |
| 112 | + ObjectDisposedException.ThrowIf(_disposed, this); |
| 113 | + |
| 114 | + var deadline = timeout.HasValue |
| 115 | + ? DateTime.UtcNow + timeout.Value |
| 116 | + : DateTime.MaxValue; |
| 117 | + |
| 118 | + while (DateTime.UtcNow < deadline) |
| 119 | + { |
| 120 | + if (await TryAcquireAsync(cancellationToken)) |
| 121 | + { |
| 122 | + return; |
| 123 | + } |
| 124 | + |
| 125 | + // Wait before retry (adaptive backoff) |
| 126 | + var waitTime = GetNextRefillTime() - DateTime.UtcNow; |
| 127 | + if (waitTime > TimeSpan.Zero) |
| 128 | + { |
| 129 | + await Task.Delay(waitTime, cancellationToken); |
| 130 | + } |
| 131 | + } |
| 132 | + |
| 133 | + throw new RateLimitExceededException( |
| 134 | + $"Rate limit of {_maxTokens} requests per minute exceeded. " + |
| 135 | + $"Timeout of {timeout} expired."); |
| 136 | + } |
| 137 | + |
| 138 | + /// <summary> |
| 139 | + /// Gets the current number of available tokens. |
| 140 | + /// </summary> |
| 141 | + /// <remarks> |
| 142 | + /// This property automatically refills tokens based on elapsed time. |
| 143 | + /// It is thread-safe and can be called concurrently. |
| 144 | + /// </remarks> |
| 145 | + /// <exception cref="ObjectDisposedException">Thrown if the rate limiter has been disposed.</exception> |
| 146 | + public int AvailableTokens |
| 147 | + { |
| 148 | + get |
| 149 | + { |
| 150 | + ObjectDisposedException.ThrowIf(_disposed, this); |
| 151 | + |
| 152 | + _lock.Wait(); |
| 153 | + try |
| 154 | + { |
| 155 | + RefillTokens(); |
| 156 | + return _currentTokens; |
| 157 | + } |
| 158 | + finally |
| 159 | + { |
| 160 | + _lock.Release(); |
| 161 | + } |
| 162 | + } |
| 163 | + } |
| 164 | + |
| 165 | + /// <summary> |
| 166 | + /// Gets the time when the next token will be available. |
| 167 | + /// </summary> |
| 168 | + /// <remarks> |
| 169 | + /// This is useful for displaying wait time to users or for implementing custom retry logic. |
| 170 | + /// </remarks> |
| 171 | + public DateTime GetNextRefillTime() |
| 172 | + { |
| 173 | + return _lastRefill + _refillInterval; |
| 174 | + } |
| 175 | + |
| 176 | + /// <summary> |
| 177 | + /// Refills tokens based on elapsed time since last refill. |
| 178 | + /// </summary> |
| 179 | + /// <remarks> |
| 180 | + /// This method is called internally by TryAcquireAsync and AvailableTokens. |
| 181 | + /// It uses integer division to ensure tokens are added in discrete units. |
| 182 | + /// </remarks> |
| 183 | + private void RefillTokens() |
| 184 | + { |
| 185 | + var now = DateTime.UtcNow; |
| 186 | + var elapsed = now - _lastRefill; |
| 187 | + var tokensToAdd = (int)(elapsed / _refillInterval); |
| 188 | + |
| 189 | + if (tokensToAdd > 0) |
| 190 | + { |
| 191 | + _currentTokens = Math.Min(_maxTokens, _currentTokens + tokensToAdd); |
| 192 | + _lastRefill = now; |
| 193 | + } |
| 194 | + } |
| 195 | + |
| 196 | + /// <summary> |
| 197 | + /// Resets the rate limiter to its initial state (for testing). |
| 198 | + /// </summary> |
| 199 | + /// <remarks> |
| 200 | + /// This method refills the bucket to maximum capacity and resets the refill timer. |
| 201 | + /// It should only be used in testing scenarios. |
| 202 | + /// </remarks> |
| 203 | + /// <exception cref="ObjectDisposedException">Thrown if the rate limiter has been disposed.</exception> |
| 204 | + public void Reset() |
| 205 | + { |
| 206 | + ObjectDisposedException.ThrowIf(_disposed, this); |
| 207 | + |
| 208 | + _lock.Wait(); |
| 209 | + try |
| 210 | + { |
| 211 | + _currentTokens = _maxTokens; |
| 212 | + _lastRefill = DateTime.UtcNow; |
| 213 | + } |
| 214 | + finally |
| 215 | + { |
| 216 | + _lock.Release(); |
| 217 | + } |
| 218 | + } |
| 219 | + |
| 220 | + /// <summary> |
| 221 | + /// Disposes the rate limiter and releases associated resources. |
| 222 | + /// </summary> |
| 223 | + public void Dispose() |
| 224 | + { |
| 225 | + if (_disposed) return; |
| 226 | + |
| 227 | + _disposed = true; |
| 228 | + _lock.Dispose(); |
| 229 | + GC.SuppressFinalize(this); |
| 230 | + } |
| 231 | +} |
0 commit comments