98 lines
3.9 KiB
C#
98 lines
3.9 KiB
C#
using System.Text;
|
|
using System.Text.Json;
|
|
using FluentAssertions;
|
|
using Microsoft.AspNetCore.Http;
|
|
using Microsoft.Extensions.Logging.Abstractions;
|
|
using StellaOps.Router.Gateway;
|
|
using StellaOps.Router.Gateway.RateLimit;
|
|
using Xunit;
|
|
|
|
namespace StellaOps.Router.Gateway.Tests;
|
|
|
|
public sealed class RateLimitMiddlewareTests
|
|
{
|
|
[Fact]
|
|
public async Task InvokeAsync_EnforcesEnvironmentLimit_WithRetryAfterAndJsonBody()
|
|
{
|
|
var config = new RateLimitConfig
|
|
{
|
|
ActivationThresholdPer5Min = 0,
|
|
ForEnvironment = new EnvironmentLimitsConfig
|
|
{
|
|
ValkeyConnection = "localhost:6379",
|
|
ValkeyBucket = "bucket",
|
|
Microservices = new Dictionary<string, MicroserviceLimitsConfig>(StringComparer.OrdinalIgnoreCase)
|
|
{
|
|
["scanner"] = new MicroserviceLimitsConfig
|
|
{
|
|
Routes = new Dictionary<string, RouteLimitsConfig>(StringComparer.OrdinalIgnoreCase)
|
|
{
|
|
["scan_submit"] = new RouteLimitsConfig
|
|
{
|
|
Pattern = "/api/scans",
|
|
MatchType = RouteMatchType.Exact,
|
|
Rules = [new RateLimitRule { PerSeconds = 300, MaxRequests = 1 }]
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}.Validate();
|
|
|
|
var store = new InMemoryValkeyRateLimitStore();
|
|
var circuitBreaker = new CircuitBreaker(failureThreshold: 5, timeoutSeconds: 30, halfOpenTimeout: 10);
|
|
var environmentLimiter = new EnvironmentRateLimiter(store, circuitBreaker, NullLogger<EnvironmentRateLimiter>.Instance);
|
|
var service = new RateLimitService(config, instanceLimiter: null, environmentLimiter, NullLogger<RateLimitService>.Instance);
|
|
|
|
var nextCalled = 0;
|
|
var middleware = new RateLimitMiddleware(
|
|
next: async ctx =>
|
|
{
|
|
nextCalled++;
|
|
ctx.Response.StatusCode = StatusCodes.Status200OK;
|
|
await ctx.Response.WriteAsync("ok");
|
|
},
|
|
rateLimitService: service,
|
|
logger: NullLogger<RateLimitMiddleware>.Instance);
|
|
|
|
// First request allowed
|
|
{
|
|
var context = new DefaultHttpContext();
|
|
context.Request.Path = "/api/scans";
|
|
context.Response.Body = new MemoryStream();
|
|
context.Items[RouterHttpContextKeys.TargetMicroservice] = "scanner";
|
|
|
|
await middleware.InvokeAsync(context);
|
|
|
|
context.Response.StatusCode.Should().Be(StatusCodes.Status200OK);
|
|
context.Response.Headers.ContainsKey("Retry-After").Should().BeFalse();
|
|
context.Response.Headers["X-RateLimit-Limit"].ToString().Should().Be("1");
|
|
nextCalled.Should().Be(1);
|
|
}
|
|
|
|
// Second request denied
|
|
{
|
|
var context = new DefaultHttpContext();
|
|
context.Request.Path = "/api/scans";
|
|
context.Response.Body = new MemoryStream();
|
|
context.Items[RouterHttpContextKeys.TargetMicroservice] = "scanner";
|
|
|
|
await middleware.InvokeAsync(context);
|
|
|
|
context.Response.StatusCode.Should().Be(StatusCodes.Status429TooManyRequests);
|
|
context.Response.Headers.ContainsKey("Retry-After").Should().BeTrue();
|
|
|
|
context.Response.Body.Position = 0;
|
|
var body = await new StreamReader(context.Response.Body, Encoding.UTF8).ReadToEndAsync();
|
|
using var json = JsonDocument.Parse(body);
|
|
|
|
json.RootElement.GetProperty("error").GetString().Should().Be("rate_limit_exceeded");
|
|
json.RootElement.GetProperty("scope").GetString().Should().Be("environment");
|
|
json.RootElement.GetProperty("limit").GetInt64().Should().Be(1);
|
|
|
|
nextCalled.Should().Be(1);
|
|
}
|
|
}
|
|
}
|
|
|