Skip to content

Commit ca67d82

Browse files
BenjaminMichaelisBenjaminMichaelis
committed
Fix rate limiting holes
1 parent cfba677 commit ca67d82

14 files changed

Lines changed: 747 additions & 234 deletions

‎EssentialCSharp.Web.Tests/McpRateLimitingTests.cs‎

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,214 @@
1+
using System.Net;
2+
using System.Text.Json;
3+
using EssentialCSharp.Web.Data;
4+
using EssentialCSharp.Web.Services;
5+
using Microsoft.AspNetCore.Mvc.Testing;
6+
using Microsoft.Extensions.DependencyInjection;
7+
8+
namespace EssentialCSharp.Web.Tests;
9+
10+
/// <summary>
11+
/// Each class gets its own factory so the global limiter starts from a fresh state.
12+
/// </summary>
13+
[NotInParallel("McpTests")]
14+
[ClassDataSource<WebApplicationFactory>(Shared = SharedType.PerClass)]
15+
public class McpDistinctUserRateLimitingTests(WebApplicationFactory factory)
16+
{
17+
[Test]
18+
public async Task DistinctValidMcpUsers_DoNotShareRateLimitBucket()
19+
{
20+
HttpClient client = McpTestHelper.CreateClient(factory);
21+
22+
for (int i = 0; i < 31; i++)
23+
{
24+
(_, string rawToken) = await McpTestHelper.CreateUserAndTokenAsync(
25+
factory,
26+
$"mcp-rate-limit-isolation-{i}",
27+
userPrefix: $"mcp-isolation-{i}");
28+
29+
using var request = McpTestHelper.CreateInitializeRequest("/mcp");
30+
McpTestHelper.AddBearerToken(request, rawToken);
31+
32+
using HttpResponseMessage response = await client.SendAsync(request);
33+
await Assert.That(response.StatusCode)
34+
.IsEqualTo(HttpStatusCode.OK)
35+
.Because($"distinct MCP user request {i + 1} should use its own rate-limit bucket");
36+
}
37+
}
38+
}
39+
40+
[NotInParallel("McpTests")]
41+
[ClassDataSource<WebApplicationFactory>(Shared = SharedType.PerClass)]
42+
public class McpPerUserRateLimitingTests(WebApplicationFactory factory)
43+
{
44+
[Test]
45+
public async Task SingleValidMcpUser_ExceedingTokenBucket_Returns429AndDoesNotCountRejectedRequests()
46+
{
47+
(_, string rawToken) = await McpTestHelper.CreateUserAndTokenAsync(
48+
factory,
49+
"mcp-rate-limit-single-user",
50+
userPrefix: "mcp-single-user");
51+
HttpClient client = McpTestHelper.CreateClient(factory);
52+
List<HttpStatusCode> statuses = [];
53+
string? rateLimitedPayload = null;
54+
string? rateLimitedContentType = null;
55+
TimeSpan? retryAfter = null;
56+
int totalRequests = McpRateLimiterPolicy.AuthenticatedTokenLimit + 15;
57+
58+
for (int i = 0; i < totalRequests; i++)
59+
{
60+
using var request = McpTestHelper.CreateInitializeRequest("/mcp");
61+
McpTestHelper.AddBearerToken(request, rawToken);
62+
63+
using HttpResponseMessage response = await client.SendAsync(request);
64+
statuses.Add(response.StatusCode);
65+
if (response.StatusCode == HttpStatusCode.TooManyRequests && rateLimitedPayload is null)
66+
{
67+
rateLimitedPayload = await response.Content.ReadAsStringAsync();
68+
rateLimitedContentType = response.Content.Headers.ContentType?.MediaType;
69+
retryAfter = response.Headers.RetryAfter?.Delta;
70+
}
71+
}
72+
73+
(long UsageCount, bool HasLastUsedAt) tokenUsage = factory.InServiceScope(services =>
74+
{
75+
var db = services.GetRequiredService<EssentialCSharpWebContext>();
76+
byte[] tokenHash = McpApiTokenService.HashToken(rawToken);
77+
var token = db.McpApiTokens.Single(t => t.TokenHash == tokenHash);
78+
return (token.UsageCount, token.LastUsedAt.HasValue);
79+
});
80+
81+
await Assert.That(statuses.Take(McpRateLimiterPolicy.AuthenticatedTokenLimit)
82+
.All(status => status == HttpStatusCode.OK)).IsTrue();
83+
await Assert.That(statuses.Skip(McpRateLimiterPolicy.AuthenticatedTokenLimit)
84+
.Any(status => status == HttpStatusCode.TooManyRequests)).IsTrue();
85+
86+
int successCount = statuses.Count(status => status == HttpStatusCode.OK);
87+
await Assert.That(successCount).IsLessThan(totalRequests);
88+
await Assert.That(tokenUsage.UsageCount).IsEqualTo((long)successCount);
89+
await Assert.That(tokenUsage.HasLastUsedAt).IsTrue();
90+
91+
string payload = rateLimitedPayload
92+
?? throw new InvalidOperationException("Expected at least one MCP token-bucket rejection.");
93+
await Assert.That(rateLimitedContentType).IsEqualTo("application/json");
94+
95+
TimeSpan retryAfterDelta = retryAfter
96+
?? throw new InvalidOperationException("Expected Retry-After on the MCP token-bucket rejection.");
97+
await Assert.That(retryAfterDelta.TotalSeconds).IsGreaterThan(0d);
98+
99+
using JsonDocument document = JsonDocument.Parse(payload);
100+
JsonElement root = document.RootElement;
101+
await Assert.That(root.GetProperty("jsonrpc").GetString()).IsEqualTo("2.0");
102+
await Assert.That(root.GetProperty("id").ValueKind).IsEqualTo(JsonValueKind.Null);
103+
JsonElement error = root.GetProperty("error");
104+
await Assert.That(error.GetProperty("code").GetInt32()).IsEqualTo(-32000);
105+
await Assert.That(error.GetProperty("message").GetString()).Contains("Rate limit exceeded");
106+
}
107+
}
108+
109+
[NotInParallel("McpTests")]
110+
[ClassDataSource<WebApplicationFactory>(Shared = SharedType.PerClass)]
111+
public class McpAnonymousRateLimitingTests(WebApplicationFactory factory)
112+
{
113+
[Test]
114+
public async Task InvalidMcpBearerRequests_FallBackToAnonymousIpBucket()
115+
{
116+
HttpClient client = McpTestHelper.CreateClient(factory);
117+
118+
for (int i = 0; i < McpRateLimiterPolicy.AnonymousPermitLimit; i++)
119+
{
120+
using var request = McpTestHelper.CreateInitializeRequest("/mcp");
121+
McpTestHelper.AddBearerToken(request, "mcp_invalid_token_that_does_not_exist");
122+
123+
using HttpResponseMessage response = await client.SendAsync(request);
124+
await Assert.That(response.StatusCode)
125+
.IsEqualTo(HttpStatusCode.Unauthorized)
126+
.Because($"invalid MCP bearer request {i + 1} should still challenge before the anonymous bucket is exhausted");
127+
}
128+
129+
using var rateLimitedRequest = McpTestHelper.CreateInitializeRequest("/mcp");
130+
McpTestHelper.AddBearerToken(rateLimitedRequest, "mcp_invalid_token_that_does_not_exist");
131+
132+
using HttpResponseMessage rateLimitedResponse = await client.SendAsync(rateLimitedRequest);
133+
await Assert.That(rateLimitedResponse.StatusCode).IsEqualTo(HttpStatusCode.TooManyRequests);
134+
}
135+
}
136+
137+
[NotInParallel("McpTests")]
138+
[ClassDataSource<WebApplicationFactory>(Shared = SharedType.PerClass)]
139+
public class McpCookieIsolationRateLimitingTests(WebApplicationFactory factory)
140+
{
141+
[Test]
142+
public async Task InvalidMcpBearerRequests_WithDifferentSiteCookies_StillShareAnonymousIpBucket()
143+
{
144+
HttpClient client = McpTestHelper.CreateClient(factory);
145+
146+
for (int i = 0; i < McpRateLimiterPolicy.AnonymousPermitLimit; i++)
147+
{
148+
string cookieUserId = await McpTestHelper.CreateUserAsync(factory, $"mcp-cookie-user-{i}");
149+
(string cookieName, string cookieValue) = await McpTestHelper.CreateIdentityApplicationCookieAsync(factory, cookieUserId);
150+
151+
using var request = McpTestHelper.CreateInitializeRequest("/mcp");
152+
McpTestHelper.AddBearerToken(request, "mcp_invalid_token_that_does_not_exist");
153+
McpTestHelper.AddCookie(request, cookieName, cookieValue);
154+
155+
using HttpResponseMessage response = await client.SendAsync(request);
156+
await Assert.That(response.StatusCode)
157+
.IsEqualTo(HttpStatusCode.Unauthorized)
158+
.Because($"invalid MCP bearer request {i + 1} should ignore the site cookie principal and stay in the anonymous/IP bucket");
159+
}
160+
161+
string finalCookieUserId = await McpTestHelper.CreateUserAsync(factory, "mcp-cookie-user-final");
162+
(string finalCookieName, string finalCookieValue) = await McpTestHelper.CreateIdentityApplicationCookieAsync(factory, finalCookieUserId);
163+
164+
using var rateLimitedRequest = McpTestHelper.CreateInitializeRequest("/mcp");
165+
McpTestHelper.AddBearerToken(rateLimitedRequest, "mcp_invalid_token_that_does_not_exist");
166+
McpTestHelper.AddCookie(rateLimitedRequest, finalCookieName, finalCookieValue);
167+
168+
using HttpResponseMessage rateLimitedResponse = await client.SendAsync(rateLimitedRequest);
169+
await Assert.That(rateLimitedResponse.StatusCode).IsEqualTo(HttpStatusCode.TooManyRequests);
170+
}
171+
}
172+
173+
[NotInParallel("McpTests")]
174+
[ClassDataSource<WebApplicationFactory>(Shared = SharedType.PerClass)]
175+
public class McpGlobalBypassRateLimitingTests(WebApplicationFactory factory)
176+
{
177+
[Test]
178+
public async Task ValidMcpPostRequests_DoNotConsumeGlobalLimiterBudgetForGetShim()
179+
{
180+
(_, string rawToken) = await McpTestHelper.CreateUserAndTokenAsync(
181+
factory,
182+
"mcp-global-bypass",
183+
userPrefix: "mcp-bypass");
184+
HttpClient client = McpTestHelper.CreateClient(factory);
185+
186+
for (int i = 0; i < 10; i++)
187+
{
188+
using var postRequest = McpTestHelper.CreateInitializeRequest("/mcp");
189+
McpTestHelper.AddBearerToken(postRequest, rawToken);
190+
191+
using HttpResponseMessage postResponse = await client.SendAsync(postRequest);
192+
await Assert.That(postResponse.StatusCode).IsEqualTo(HttpStatusCode.OK);
193+
}
194+
195+
for (int i = 0; i < 30; i++)
196+
{
197+
using var getRequest = new HttpRequestMessage(HttpMethod.Get, "/mcp");
198+
McpTestHelper.AddBearerToken(getRequest, rawToken);
199+
getRequest.Headers.Accept.ParseAdd("text/event-stream");
200+
201+
using HttpResponseMessage getResponse = await client.SendAsync(getRequest);
202+
await Assert.That(getResponse.StatusCode)
203+
.IsEqualTo(HttpStatusCode.MethodNotAllowed)
204+
.Because($"global request {i + 1} should still be within the non-MCP GET shim limit");
205+
}
206+
207+
using var rateLimitedGetRequest = new HttpRequestMessage(HttpMethod.Get, "/mcp");
208+
McpTestHelper.AddBearerToken(rateLimitedGetRequest, rawToken);
209+
rateLimitedGetRequest.Headers.Accept.ParseAdd("text/event-stream");
210+
211+
using HttpResponseMessage rateLimitedGetResponse = await client.SendAsync(rateLimitedGetRequest);
212+
await Assert.That(rateLimitedGetResponse.StatusCode).IsEqualTo(HttpStatusCode.TooManyRequests);
213+
}
214+
}

‎EssentialCSharp.Web.Tests/McpTestHelper.cs‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
using System.Net.Http.Headers;
2+
using System.Text;
3+
using EssentialCSharp.Web.Areas.Identity.Data;
4+
using EssentialCSharp.Web.Data;
5+
using EssentialCSharp.Web.Services;
6+
using Microsoft.AspNetCore.Authentication;
7+
using Microsoft.AspNetCore.Authentication.Cookies;
8+
using Microsoft.AspNetCore.Identity;
9+
using Microsoft.AspNetCore.Mvc.Testing;
10+
using Microsoft.Extensions.DependencyInjection;
11+
using Microsoft.Extensions.Options;
12+
13+
namespace EssentialCSharp.Web.Tests;
14+
15+
internal static class McpTestHelper
16+
{
17+
public static HttpClient CreateClient(WebApplicationFactory factory) => factory.CreateClient(new WebApplicationFactoryClientOptions
18+
{
19+
AllowAutoRedirect = false
20+
});
21+
22+
public static HttpRequestMessage CreateInitializeRequest(string path = "/mcp")
23+
{
24+
var request = new HttpRequestMessage(HttpMethod.Post, path)
25+
{
26+
Content = new StringContent(
27+
"""
28+
{
29+
"jsonrpc": "2.0",
30+
"id": 1,
31+
"method": "initialize",
32+
"params": {
33+
"protocolVersion": "2024-11-05",
34+
"capabilities": {},
35+
"clientInfo": { "name": "test-client", "version": "1.0" }
36+
}
37+
}
38+
""",
39+
Encoding.UTF8, "application/json")
40+
};
41+
42+
request.Headers.Accept.ParseAdd("application/json");
43+
request.Headers.Accept.ParseAdd("text/event-stream");
44+
return request;
45+
}
46+
47+
public static void AddBearerToken(HttpRequestMessage request, string rawToken) =>
48+
request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", rawToken);
49+
50+
public static void AddCookie(HttpRequestMessage request, string cookieName, string cookieValue) =>
51+
request.Headers.Add("Cookie", $"{cookieName}={cookieValue}");
52+
53+
public static async Task<string> CreateUserAsync(WebApplicationFactory factory, string userPrefix)
54+
{
55+
string userId = Guid.NewGuid().ToString();
56+
string suffix = Guid.NewGuid().ToString("N")[..8];
57+
string userName = $"{userPrefix.ToLowerInvariant()}-{suffix}";
58+
59+
using var scope = factory.Services.CreateScope();
60+
var db = scope.ServiceProvider.GetRequiredService<EssentialCSharpWebContext>();
61+
db.Users.Add(new EssentialCSharpWebUser
62+
{
63+
Id = userId,
64+
UserName = userName,
65+
NormalizedUserName = userName.ToUpperInvariant(),
66+
Email = $"{userName}@example.com",
67+
NormalizedEmail = $"{userName.ToUpperInvariant()}@EXAMPLE.COM",
68+
SecurityStamp = Guid.NewGuid().ToString(),
69+
});
70+
await db.SaveChangesAsync();
71+
72+
return userId;
73+
}
74+
75+
public static async Task<(string UserId, string RawToken)> CreateUserAndTokenAsync(
76+
WebApplicationFactory factory,
77+
string tokenName,
78+
string userPrefix = "mcp-test",
79+
DateTime? expiresAt = null)
80+
{
81+
string userId = await CreateUserAsync(factory, userPrefix);
82+
83+
using var scope = factory.Services.CreateScope();
84+
var tokenService = scope.ServiceProvider.GetRequiredService<McpApiTokenService>();
85+
(string rawToken, _) = expiresAt is { } expiry
86+
? await tokenService.CreateTokenAsync(userId, tokenName, expiry)
87+
: await tokenService.CreateTokenAsync(userId, tokenName);
88+
89+
return (userId, rawToken);
90+
}
91+
92+
public static async Task<(string CookieName, string CookieValue)> CreateIdentityApplicationCookieAsync(
93+
WebApplicationFactory factory,
94+
string userId)
95+
{
96+
using var scope = factory.Services.CreateScope();
97+
var signInManager = scope.ServiceProvider.GetRequiredService<SignInManager<EssentialCSharpWebUser>>();
98+
EssentialCSharpWebUser user = await signInManager.UserManager.FindByIdAsync(userId)
99+
?? throw new InvalidOperationException($"Could not find test user '{userId}' to create an identity cookie.");
100+
var principal = await signInManager.CreateUserPrincipalAsync(user);
101+
102+
CookieAuthenticationOptions cookieOptions = scope.ServiceProvider
103+
.GetRequiredService<IOptionsMonitor<CookieAuthenticationOptions>>()
104+
.Get(IdentityConstants.ApplicationScheme);
105+
106+
string cookieName = cookieOptions.Cookie.Name
107+
?? throw new InvalidOperationException("Identity application cookie name is not configured.");
108+
109+
var ticket = new AuthenticationTicket(
110+
principal,
111+
new AuthenticationProperties
112+
{
113+
IssuedUtc = DateTimeOffset.UtcNow,
114+
ExpiresUtc = DateTimeOffset.UtcNow.Add(cookieOptions.ExpireTimeSpan),
115+
},
116+
IdentityConstants.ApplicationScheme);
117+
118+
string cookieValue = cookieOptions.TicketDataFormat.Protect(ticket)
119+
?? throw new InvalidOperationException("Failed to protect the identity application ticket.");
120+
121+
return (cookieName, cookieValue);
122+
}
123+
}

0 commit comments

Comments
 (0)