-
Notifications
You must be signed in to change notification settings - Fork 447
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add StackExchangeRedis atomic Lua script support and refactor to supp…
…ort differing processing strategies
- Loading branch information
1 parent
5b21509
commit 60e6ebc
Showing
27 changed files
with
419 additions
and
112 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
66 changes: 66 additions & 0 deletions
66
src/AspNetCoreRateLimit/Core/ProcessingStrategies/AsyncKeyLockProcessingStrategy.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,66 @@ | ||
| using System; | ||
| using System.Threading; | ||
| using System.Threading.Tasks; | ||
|
|
||
| namespace AspNetCoreRateLimit | ||
| { | ||
| public class AsyncKeyLockProcessingStrategy : ProcessingStrategy | ||
| { | ||
| private readonly RateLimitOptions _options; | ||
| private readonly IRateLimitCounterStore _counterStore; | ||
| private readonly ICounterKeyBuilder _counterKeyBuilder; | ||
| private readonly IRateLimitConfiguration _config; | ||
|
|
||
| public AsyncKeyLockProcessingStrategy(IRateLimitCounterStore counterStore, ICounterKeyBuilder counterKeyBuilder, IRateLimitConfiguration config, RateLimitOptions options) | ||
| : base(counterKeyBuilder, config, options) | ||
| { | ||
| _counterStore = counterStore; | ||
| _counterKeyBuilder = counterKeyBuilder; | ||
| _config = config; | ||
| _options = options; | ||
| } | ||
|
|
||
|
|
||
| /// The key-lock used for limiting requests. | ||
| private static readonly AsyncKeyLock AsyncLock = new AsyncKeyLock(); | ||
|
|
||
| public override async Task<RateLimitCounter> ProcessRequestAsync(ClientRequestIdentity requestIdentity, RateLimitRule rule, CancellationToken cancellationToken = default) | ||
| { | ||
| var counter = new RateLimitCounter | ||
| { | ||
| Timestamp = DateTime.UtcNow, | ||
| Count = 1 | ||
| }; | ||
|
|
||
| var counterId = BuildCounterKey(requestIdentity, rule); | ||
|
|
||
| // serial reads and writes on same key | ||
| using (await AsyncLock.WriterLockAsync(counterId).ConfigureAwait(false)) | ||
| { | ||
| var entry = await _counterStore.GetAsync(counterId, cancellationToken); | ||
|
|
||
| if (entry.HasValue) | ||
| { | ||
| // entry has not expired | ||
| if (entry.Value.Timestamp + rule.PeriodTimespan.Value >= DateTime.UtcNow) | ||
| { | ||
| // increment request count | ||
| var totalCount = entry.Value.Count + _config.RateIncrementer?.Invoke() ?? 1; | ||
|
|
||
| // deep copy | ||
| counter = new RateLimitCounter | ||
| { | ||
| Timestamp = entry.Value.Timestamp, | ||
| Count = totalCount | ||
| }; | ||
| } | ||
| } | ||
|
|
||
| // stores: id (string) - timestamp (datetime) - total_requests (long) | ||
| await _counterStore.SetAsync(counterId, counter, rule.PeriodTimespan.Value, cancellationToken); | ||
| } | ||
|
|
||
| return counter; | ||
| } | ||
| } | ||
| } |
10 changes: 10 additions & 0 deletions
10
src/AspNetCoreRateLimit/Core/ProcessingStrategies/IProcessingStrategy.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,10 @@ | ||
| using System.Threading; | ||
| using System.Threading.Tasks; | ||
|
|
||
| namespace AspNetCoreRateLimit | ||
| { | ||
| public interface IProcessingStrategy | ||
| { | ||
| Task<RateLimitCounter> ProcessRequestAsync(ClientRequestIdentity requestIdentity, RateLimitRule rule, CancellationToken cancellationToken = default); | ||
| } | ||
| } |
7 changes: 7 additions & 0 deletions
7
src/AspNetCoreRateLimit/Core/ProcessingStrategies/IProcessingStrategyFactory.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,7 @@ | ||
| namespace AspNetCoreRateLimit | ||
| { | ||
| public interface IProcessingStrategyFactory | ||
| { | ||
| ProcessingStrategy CreateProcessingStrategy(IRateLimitCounterStore counterStore, ICounterKeyBuilder counterKeyBuilder, IRateLimitConfiguration config, RateLimitOptions options); | ||
| } | ||
| } |
42 changes: 42 additions & 0 deletions
42
src/AspNetCoreRateLimit/Core/ProcessingStrategies/ProcessingStrategy.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| using System; | ||
| using System.Security.Cryptography; | ||
| using System.Text; | ||
| using System.Threading; | ||
| using System.Threading.Tasks; | ||
|
|
||
| namespace AspNetCoreRateLimit | ||
| { | ||
| public abstract class ProcessingStrategy : IProcessingStrategy | ||
| { | ||
| private readonly RateLimitOptions _options; | ||
| private readonly ICounterKeyBuilder _counterKeyBuilder; | ||
| private readonly IRateLimitConfiguration _config; | ||
|
|
||
| public ProcessingStrategy(ICounterKeyBuilder counterKeyBuilder, IRateLimitConfiguration config, RateLimitOptions options) | ||
| : base() | ||
| { | ||
| _counterKeyBuilder = counterKeyBuilder; | ||
| _config = config; | ||
| _options = options; | ||
| } | ||
|
|
||
| protected virtual string BuildCounterKey(ClientRequestIdentity requestIdentity, RateLimitRule rule) | ||
| { | ||
| var key = _counterKeyBuilder.Build(requestIdentity, rule); | ||
|
|
||
| if (_options.EnableEndpointRateLimiting && _config.EndpointCounterKeyBuilder != null) | ||
| { | ||
| key += _config.EndpointCounterKeyBuilder.Build(requestIdentity, rule); | ||
| } | ||
|
|
||
| var bytes = Encoding.UTF8.GetBytes(key); | ||
|
|
||
| using var algorithm = new SHA1Managed(); | ||
| var hash = algorithm.ComputeHash(bytes); | ||
|
|
||
| return Convert.ToBase64String(hash); | ||
| } | ||
|
|
||
| public abstract Task<RateLimitCounter> ProcessRequestAsync(ClientRequestIdentity requestIdentity, RateLimitRule rule, CancellationToken cancellationToken = default); | ||
| } | ||
| } |
27 changes: 27 additions & 0 deletions
27
src/AspNetCoreRateLimit/Core/ProcessingStrategies/ProcessingStrategyFactory.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| using System; | ||
| using StackExchange.Redis; | ||
|
|
||
| namespace AspNetCoreRateLimit | ||
| { | ||
|
|
||
| public class ProcessingStrategyFactory : IProcessingStrategyFactory | ||
| { | ||
| private readonly IConnectionMultiplexer _connectionMultiplexer; | ||
|
|
||
| public ProcessingStrategyFactory(IConnectionMultiplexer connectionMultiplexer = null) | ||
| { | ||
| _connectionMultiplexer = connectionMultiplexer; | ||
| } | ||
|
|
||
| public ProcessingStrategy CreateProcessingStrategy(IRateLimitCounterStore counterStore, ICounterKeyBuilder counterKeyBuilder, IRateLimitConfiguration config, RateLimitOptions options) | ||
| { | ||
| return counterStore switch | ||
| { | ||
| MemoryCacheRateLimitCounterStore => new AsyncKeyLockProcessingStrategy(counterStore, counterKeyBuilder, config, options), | ||
| DistributedCacheRateLimitCounterStore => new AsyncKeyLockProcessingStrategy(counterStore, counterKeyBuilder, config, options), | ||
| StackExchangeRedisRateLimitCounterStore => new StackExchangeRedisProcessingStrategy(_connectionMultiplexer, counterStore, counterKeyBuilder, config, options), | ||
| _ => throw new ArgumentException("Unsupported instance of IRateLimitCounterStore provided") | ||
| }; | ||
| } | ||
| } | ||
| } |
44 changes: 44 additions & 0 deletions
44
src/AspNetCoreRateLimit/Core/ProcessingStrategies/StackExchangeRedisProcessingStrategy.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,44 @@ | ||
| using System; | ||
| using System.Threading; | ||
| using System.Threading.Tasks; | ||
| using StackExchange.Redis; | ||
|
|
||
| namespace AspNetCoreRateLimit | ||
| { | ||
| public class StackExchangeRedisProcessingStrategy : ProcessingStrategy | ||
| { | ||
| private readonly IConnectionMultiplexer _connectionMultiplexer; | ||
| private readonly IRateLimitConfiguration _config; | ||
|
|
||
| public StackExchangeRedisProcessingStrategy(IConnectionMultiplexer connectionMultiplexer, IRateLimitCounterStore counterStore, ICounterKeyBuilder counterKeyBuilder, IRateLimitConfiguration config, RateLimitOptions options) | ||
| : base(counterKeyBuilder, config, options) | ||
| { | ||
| _connectionMultiplexer = connectionMultiplexer ?? throw new ArgumentException("IConnectionMultiplexer was null. Ensure StackExchange.Redis was successfully registered"); | ||
| _config = config; | ||
| } | ||
|
|
||
|
|
||
| static private readonly LuaScript _atomicIncrement = LuaScript.Prepare("local count count = redis.call(\"INCRBYFLOAT\", @key, tonumber(@delta)) if tonumber(count) == @delta then redis.call(\"EXPIRE\", @key, @timeout) end return count"); | ||
|
|
||
| public override async Task<RateLimitCounter> ProcessRequestAsync(ClientRequestIdentity requestIdentity, RateLimitRule rule, CancellationToken cancellationToken = default) | ||
| { | ||
| var counterId = BuildCounterKey(requestIdentity, rule); | ||
| return await IncrementAsync(counterId, rule.PeriodTimespan.Value, _config.RateIncrementer); | ||
| } | ||
|
|
||
| public async Task<RateLimitCounter> IncrementAsync(string counterId, TimeSpan interval, Func<double> RateIncrementer = null) | ||
| { | ||
| var now = DateTime.UtcNow; | ||
| var numberOfIntervals = now.Ticks / interval.Ticks; | ||
| var intervalStart = new DateTime(numberOfIntervals * interval.Ticks, DateTimeKind.Utc); | ||
|
|
||
| // Call the Lua script | ||
| var count = await _connectionMultiplexer.GetDatabase().ScriptEvaluateAsync(_atomicIncrement, new { key = counterId, timeout = interval.TotalSeconds, delta = RateIncrementer?.Invoke() ?? 1D }); | ||
| return new RateLimitCounter | ||
| { | ||
| Count = (double)count, | ||
| Timestamp = intervalStart | ||
| }; | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.