Skip to content

Commit 093a431

Browse files
vbilopavvbilopav
committed
rate limiter implementation
1 parent edd6fca commit 093a431

8 files changed

Lines changed: 60 additions & 88 deletions

File tree

‎NpgsqlRest/Log.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,4 +238,7 @@ public static partial class Log
238238

239239
[LoggerMessage(Level = LogLevel.Debug, Message = "{description} has set RATE LIMITER POLICY NAME to {name} by the comment annotation.")]
240240
public static partial void RateLimiterPolicySet(this ILogger logger, string description, string name);
241+
242+
[LoggerMessage(Level = LogLevel.Debug, Message = "Endpoint {urlInfo} enabled rate limiter policy: {policyName}")]
243+
public static partial void EndpointEnabledRateLimiterPolicy(this ILogger logger, string urlInfo, string policyName);
241244
}

‎NpgsqlRest/NpgsqlRestBuilder.cs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,15 @@ public static IApplicationBuilder UseNpgsqlRest(this WebApplication builder, Npg
6464
var handler = new NpgsqlRestEndpoint(entry, overloads, logger);
6565
var endpoint = entry.Endpoint;
6666
var methodStr = endpoint.Method.ToString();
67+
var urlInfo = string.Concat(methodStr, " ", endpoint.Path);
6768
var routeBuilder = builder.MapMethods(endpoint.Path, [methodStr], handler.InvokeAsync);
69+
70+
if (endpoint.RateLimiterPolicy is not null || options.DefaultRateLimitingPolicy is not null)
71+
{
72+
var policy = endpoint.RateLimiterPolicy ?? options.DefaultRateLimitingPolicy!;
73+
routeBuilder.RequireRateLimiting(policy);
74+
logger?.EndpointEnabledRateLimiterPolicy(urlInfo, policy);
75+
}
6876

6977
if (options.RouteHandlerCreated is not null)
7078
{
@@ -73,7 +81,6 @@ public static IApplicationBuilder UseNpgsqlRest(this WebApplication builder, Npg
7381

7482
if (options.LogEndpointCreatedInfo)
7583
{
76-
var urlInfo = string.Concat(methodStr, " ", endpoint.Path);
7784
logger?.EndpointCreated(urlInfo);
7885
if (endpoint.InfoEventsStreamingPath is not null)
7986
{

‎NpgsqlRest/Options/NpgsqlRestOptions.cs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,8 @@ public NpgsqlRestOptions(NpgsqlDataSource dataSource)
343343
public Dictionary<string, StringValues> CustomServerSentEventsResponseHeaders { get; set; } = [];
344344

345345
/// <summary>
346-
/// Rate Limiter Options
346+
/// Default rate limiting policy for all requests. Policy must be configured within application rate limiting options.
347+
/// This can be overridden by comment annotations in the database or setting policy for specific endpoints.
347348
/// </summary>
348-
public RateLimiterOptions RateLimiterOptions { get; set; } = new();
349+
public string? DefaultRateLimitingPolicy { get; set; } = null;
349350
}

‎NpgsqlRest/Options/RateLimiterOptions.cs‎

Lines changed: 0 additions & 22 deletions
This file was deleted.

‎NpgsqlRestClient/Builder.cs‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
using Serilog;
1717
using Serilog.Extensions.Logging;
1818
using Serilog.Sinks.OpenTelemetry;
19-
using RateLimiterOptions = NpgsqlRest.RateLimiterOptions;
2019

2120
namespace NpgsqlRestClient;
2221

@@ -933,29 +932,31 @@ public CacheOptions BuildCacheOptions(WebApplication app)
933932

934933
public enum RateLimiterType { FixedWindow, SlidingWindow, TokenBucket, Concurrency }
935934

936-
public RateLimiterOptions? BuildRateLimiter()
935+
public (string? defaultPolicy, bool enabled) BuildRateLimiter()
937936
{
938937
var rateLimiterCfg = _config.Cfg.GetSection("RateLimiterOptions");
939938
if (_config.Exists(rateLimiterCfg) is false || _config.GetConfigBool("Enabled", rateLimiterCfg) is false)
940939
{
941-
return null;
940+
return (null, false);
942941
}
943942
var policiesCfg = rateLimiterCfg.GetSection("Policies");
944943
if (_config.Exists(policiesCfg) is false)
945944
{
946-
return null;
945+
return (null, false);
947946
}
948-
949-
var result = new RateLimiterOptions
950-
{
951-
Enabled = true,
952-
DefaultPolicy = _config.GetConfigStr("DefaultPolicy", rateLimiterCfg),
953-
StatusCode = _config.GetConfigInt("StatusCode", rateLimiterCfg) ?? 429,
954-
Message = _config.GetConfigStr("Message", rateLimiterCfg) ?? "Too many requests. Please try again later."
955-
};
956-
947+
948+
var defaultPolicy = _config.GetConfigStr("DefaultPolicy", rateLimiterCfg);
949+
var message = _config.GetConfigStr("Message", rateLimiterCfg);
957950
Instance.Services.AddRateLimiter(options =>
958951
{
952+
options.RejectionStatusCode = _config.GetConfigInt("StatusCode", rateLimiterCfg) ?? 429;
953+
if (string.IsNullOrEmpty(message) is false)
954+
{
955+
options.OnRejected = async (context, cancellationToken) =>
956+
{
957+
await context.HttpContext.Response.WriteAsync(message, cancellationToken);
958+
};
959+
}
959960
foreach (var sectionCfg in policiesCfg.GetChildren())
960961
{
961962
var type = _config.GetConfigEnum<RateLimiterType?>("Type", sectionCfg);
@@ -1038,6 +1039,6 @@ public enum RateLimiterType { FixedWindow, SlidingWindow, TokenBucket, Concurren
10381039
}
10391040
});
10401041

1041-
return result;
1042+
return (defaultPolicy, true);
10421043
}
10431044
}

‎NpgsqlRestClient/Program.cs‎

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@
116116
builder.Instance.Services.AddRouting();
117117

118118
builder.BuildLogger(cmdRetryStrategy);
119-
var rateLimiterOptions = builder.BuildRateLimiter();
119+
var (rateLimiterDefaultPolicy, rateLimiterEnabled) = builder.BuildRateLimiter();
120120

121121
var (connectionString, retryOpts) = builder.BuildConnectionString();
122122
if (connectionString is null)
@@ -136,10 +136,10 @@
136136

137137
WebApplication app = builder.Build();
138138

139-
// if (rateLimiterOptions is not null && rateLimiterOptions.Enabled)
140-
// {
141-
// app.UseRateLimiter();
142-
// }
139+
if (rateLimiterEnabled)
140+
{
141+
app.UseRateLimiter();
142+
}
143143

144144
// dump encrypted text and exit
145145
if (args.Length >= 1 && string.Equals(args[0], "--encrypt", StringComparison.CurrentCultureIgnoreCase))
@@ -197,18 +197,6 @@
197197

198198
appInstance.ConfigureThreadPool();
199199

200-
/*
201-
//
202-
// The connection name in ConnectionStrings configuration that will be used to execute the metadata query. If this value is null, the default connection string will be used.
203-
//
204-
"MetadataQueryConnectionName": null,
205-
//
206-
// Set the search path to this schema that contains the metadata query function. Default is `public`.
207-
// This is needed when using non superuser connection roles with limited schema access and mapping the metadata function to a specific schema.
208-
// If the connection string contains the same "Search Path=" it will be skipped.
209-
//
210-
"MetadataQuerySchema": null,
211-
*/
212200
NpgsqlRestOptions options = new()
213201
{
214202
DataSource = dataSource,
@@ -261,7 +249,7 @@
261249
UploadOptions = appInstance.CreateUploadOptions(),
262250

263251
CacheOptions = builder.BuildCacheOptions(app),
264-
RateLimiterOptions = rateLimiterOptions ?? new RateLimiterOptions(),
252+
DefaultRateLimitingPolicy = rateLimiterDefaultPolicy
265253
};
266254

267255
app.UseNpgsqlRest(options);

‎NpgsqlRestTests/AuthTests/RateLimiterTests.cs‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ public async Task Test_get_rate_unlimited1()
3939

4040
using var request3 = new HttpRequestMessage(HttpMethod.Get, "/api/get-rate-unlimited1");
4141
using var result3 = await test.Client.SendAsync(request3);
42-
var response3 = await result2.Content.ReadAsStringAsync();
42+
var response3 = await result3.Content.ReadAsStringAsync();
4343

4444
result1.StatusCode.Should().Be(HttpStatusCode.OK);
4545
response1.Should().Be("postgres");
@@ -57,24 +57,23 @@ public async Task Test_get_rate_limited1()
5757
using var request1 = new HttpRequestMessage(HttpMethod.Get, "/api/get-rate-limited1");
5858
using var result1 = await test.Client.SendAsync(request1);
5959
var response1 = await result1.Content.ReadAsStringAsync();
60-
60+
6161
using var request2 = new HttpRequestMessage(HttpMethod.Get, "/api/get-rate-limited1");
6262
using var result2 = await test.Client.SendAsync(request2);
6363
var response2 = await result2.Content.ReadAsStringAsync();
64-
64+
6565
using var request3 = new HttpRequestMessage(HttpMethod.Get, "/api/get-rate-limited1");
6666
using var result3 = await test.Client.SendAsync(request3);
67-
var response3 = await result2.Content.ReadAsStringAsync();
68-
69-
var options = new RateLimiterOptions();
67+
var response3 = await result3.Content.ReadAsStringAsync();
68+
7069
result1.StatusCode.Should().Be(HttpStatusCode.OK);
7170
response1.Should().Be("postgres");
72-
71+
7372
result2.StatusCode.Should().Be(HttpStatusCode.OK);
7473
response2.Should().Be("postgres");
75-
76-
result3.StatusCode.Should().Be(HttpStatusCode.OK);
77-
response3.Should().Be("postgres");
74+
75+
result3.StatusCode.Should().Be(HttpStatusCode.TooManyRequests);
76+
response3.Should().Be("Rate limit exceeded. Please try again later.");
7877
}
7978
}
8079

‎NpgsqlRestTests/Setup/Program.cs‎

Lines changed: 16 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
using Microsoft.AspNetCore.Authentication.Cookies;
33
using Microsoft.AspNetCore.Builder;
44
using Microsoft.AspNetCore.Http;
5+
using Microsoft.AspNetCore.RateLimiting;
56
using Microsoft.Extensions.DependencyInjection;
67
using NpgsqlRest.Auth;
78
using NpgsqlRest.CrudSource;
@@ -68,33 +69,29 @@ public static void Main()
6869

6970
var builder = WebApplication.CreateBuilder([]);
7071

71-
// builder.Services.AddRateLimiter(options =>
72-
// {
73-
// options.AddFixedWindowLimiter("max 2 per second", config =>
74-
// {
75-
// config.PermitLimit = 2;
76-
// config.Window = TimeSpan.FromSeconds(1);
77-
// config.AutoReplenishment = true;
78-
// });
79-
//
80-
// //var fixedWindowRateLimiterOptions = new FixedWindowRateLimiterOptions();
81-
// //var policyName = "max 2 per second (from AddPolicy)";
82-
// // options.AddPolicy(policyName, context =>
83-
// // {
84-
// // return fixedWindowRateLimiterOptions;
85-
// // });
86-
// });
72+
builder.Services.AddRateLimiter(options =>
73+
{
74+
options.RejectionStatusCode = StatusCodes.Status429TooManyRequests;
75+
options.OnRejected = async (context, cancellationToken) =>
76+
{
77+
await context.HttpContext.Response.WriteAsync("Rate limit exceeded. Please try again later.", cancellationToken);
78+
};
79+
options.AddFixedWindowLimiter("max 2 per second", config =>
80+
{
81+
config.PermitLimit = 2;
82+
config.Window = TimeSpan.FromSeconds(1);
83+
config.AutoReplenishment = true;
84+
});
85+
});
8786

88-
//var fixedWindowRateLimiterOptions = new FixedWindowRateLimiterOptions();
89-
9087
builder
9188
.Services
9289
.AddAuthentication()
9390
//.AddBearerToken();
9491
.AddCookie();
9592

9693
var app = builder.Build();
97-
//app.UseRateLimiter();
94+
app.UseRateLimiter();
9895

9996
var authOptions = new NpgsqlRestAuthenticationOptions
10097
{
@@ -196,8 +193,6 @@ public static void Main()
196193
UseDefaultUploadMetadataContextKey = true,
197194
//DefaultUploadMetadataContextKey = "request.upload_metadata",
198195
},
199-
200-
//RateLimiterOptions = new RateLimiterOptions { Enabled = true }
201196
});
202197
app.Run();
203198
}

0 commit comments

Comments
 (0)