Skip to content

Commit ad58a21

Browse files
waldekmastykarzwaldekmastykarz
authored
Combines rate limiting with mocking. Closes #350 (#461)
1 parent 3eec3d6 commit ad58a21

File tree

6 files changed

+103
-70
lines changed

6 files changed

+103
-70
lines changed

‎dev-proxy-abstractions/PluginEvents.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ internal ProxyHttpEventArgsBase(SessionEventArgs session, IList<ThrottlerInfo> t
6666

6767
public SessionEventArgs Session { get; }
6868
public IList<ThrottlerInfo> ThrottledRequests { get; }
69+
public Dictionary<string, object> PluginData { get; set; } = new Dictionary<string, object>();
6970

7071
public bool HasRequestUrlMatch(ISet<UrlToWatch> watchedUrls)
7172
{

‎dev-proxy-abstractions/ProxyUtils.cs

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,4 +319,28 @@ public static string ProductVersion
319319
return _productVersion;
320320
}
321321
}
322+
323+
public static void MergeHeaders(IList<HttpHeader> allHeaders, IList<HttpHeader> headersToAdd)
324+
{
325+
foreach (var header in headersToAdd)
326+
{
327+
var existingHeader = allHeaders.FirstOrDefault(h => h.Name.Equals(header.Name, StringComparison.OrdinalIgnoreCase));
328+
if (existingHeader is not null)
329+
{
330+
if (header.Name.Equals("Access-Control-Expose-Headers", StringComparison.OrdinalIgnoreCase))
331+
{
332+
var existingValues = existingHeader.Value.Split(',').Select(v => v.Trim());
333+
var newValues = header.Value.Split(',').Select(v => v.Trim());
334+
var allValues = existingValues.Union(newValues).Distinct();
335+
allHeaders.Remove(existingHeader);
336+
allHeaders.Add(new HttpHeader(header.Name, string.Join(", ", allValues)));
337+
continue;
338+
}
339+
340+
allHeaders.Remove(existingHeader);
341+
}
342+
343+
allHeaders.Add(header);
344+
}
345+
}
322346
}

‎dev-proxy-plugins/Behavior/RateLimitingCustomResponseLoader.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ public void LoadResponse()
3636
using (StreamReader reader = new StreamReader(stream))
3737
{
3838
var responseString = reader.ReadToEnd();
39-
var response = JsonSerializer.Deserialize<MockResponse>(responseString);
39+
var response = JsonSerializer.Deserialize<MockResponseResponse>(responseString);
4040
if (response is not null)
4141
{
4242
_configuration.CustomResponse = response;

‎dev-proxy-plugins/Behavior/RateLimitingPlugin.cs

Lines changed: 46 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ public class RateLimitConfiguration
3838
public int RetryAfterSeconds { get; set; } = 5;
3939
public RateLimitResponseWhenLimitExceeded WhenLimitExceeded { get; set; } = RateLimitResponseWhenLimitExceeded.Throttle;
4040
public string CustomResponseFile { get; set; } = "rate-limit-response.json";
41-
public MockResponse? CustomResponse { get; set; }
41+
public MockResponseResponse? CustomResponse { get; set; }
4242
}
4343

4444
public class RateLimitingPlugin : BaseProxyPlugin
@@ -95,55 +95,10 @@ private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorS
9595
return;
9696
}
9797

98-
// add rate limiting headers if reached the threshold percentage
99-
if (_resourcesRemaining <= _configuration.RateLimit - (_configuration.RateLimit * _configuration.WarningThresholdPercent / 100))
98+
if (e.PluginData.TryGetValue(Name, out var pluginData) &&
99+
pluginData is List<HttpHeader> rateLimitingHeaders)
100100
{
101-
var reset = _configuration.ResetFormat == RateLimitResetFormat.SecondsLeft ?
102-
(_resetTime - DateTime.Now).TotalSeconds.ToString("N0") : // drop decimals
103-
new DateTimeOffset(_resetTime).ToUnixTimeSeconds().ToString();
104-
headers.AddRange(new List<HttpHeader> {
105-
new HttpHeader(_configuration.HeaderLimit, _configuration.RateLimit.ToString()),
106-
new HttpHeader(_configuration.HeaderRemaining, _resourcesRemaining.ToString()),
107-
new HttpHeader(_configuration.HeaderReset, reset)
108-
});
109-
110-
// make rate limiting information available for CORS requests
111-
if (request.Headers.FirstOrDefault((h) => h.Name.Equals("Origin", StringComparison.OrdinalIgnoreCase)) is not null)
112-
{
113-
if (!response.Headers.HeaderExists("Access-Control-Allow-Origin"))
114-
{
115-
headers.Add(new HttpHeader("Access-Control-Allow-Origin", "*"));
116-
}
117-
var exposeHeadersHeader = response.Headers.FirstOrDefault((h) => h.Name.Equals("Access-Control-Expose-Headers", StringComparison.OrdinalIgnoreCase));
118-
var headerValue = "";
119-
if (exposeHeadersHeader is null)
120-
{
121-
headerValue = $"{_configuration.HeaderLimit}, {_configuration.HeaderRemaining}, {_configuration.HeaderReset}, {_configuration.HeaderRetryAfter}";
122-
}
123-
else
124-
{
125-
headerValue = exposeHeadersHeader.Value;
126-
if (!headerValue.Contains(_configuration.HeaderLimit))
127-
{
128-
headerValue += $", {_configuration.HeaderLimit}";
129-
}
130-
if (!headerValue.Contains(_configuration.HeaderRemaining))
131-
{
132-
headerValue += $", {_configuration.HeaderRemaining}";
133-
}
134-
if (!headerValue.Contains(_configuration.HeaderReset))
135-
{
136-
headerValue += $", {_configuration.HeaderReset}";
137-
}
138-
if (!headerValue.Contains(_configuration.HeaderRetryAfter))
139-
{
140-
headerValue += $", {_configuration.HeaderRetryAfter}";
141-
}
142-
response.Headers.RemoveHeader("Access-Control-Expose-Headers");
143-
}
144-
145-
headers.Add(new HttpHeader("Access-Control-Expose-Headers", headerValue));
146-
}
101+
ProxyUtils.MergeHeaders(headers, rateLimitingHeaders);
147102
}
148103

149104
// add headers to the original API response, avoiding duplicates
@@ -244,12 +199,12 @@ _urlsToWatch is null ||
244199
{
245200
if (_configuration.CustomResponse is not null)
246201
{
247-
var headers = _configuration.CustomResponse.Response?.Headers is not null ?
248-
_configuration.CustomResponse.Response.Headers.Select(h => new HttpHeader(h.Key, h.Value)) :
202+
var headers = _configuration.CustomResponse.Headers is not null ?
203+
_configuration.CustomResponse.Headers.Select(h => new HttpHeader(h.Key, h.Value)) :
249204
Array.Empty<HttpHeader>();
250205

251206
// allow custom throttling response
252-
var responseCode = (HttpStatusCode)(_configuration.CustomResponse.Response?.StatusCode ?? 200);
207+
var responseCode = (HttpStatusCode)(_configuration.CustomResponse.StatusCode ?? 200);
253208
if (responseCode == HttpStatusCode.TooManyRequests)
254209
{
255210
e.ThrottledRequests.Add(new ThrottlerInfo(
@@ -259,8 +214,8 @@ _urlsToWatch is null ||
259214
));
260215
}
261216

262-
string body = _configuration.CustomResponse.Response?.Body is not null ?
263-
JsonSerializer.Serialize(_configuration.CustomResponse.Response.Body, new JsonSerializerOptions { WriteIndented = true }) :
217+
string body = _configuration.CustomResponse.Body is not null ?
218+
JsonSerializer.Serialize(_configuration.CustomResponse.Body, new JsonSerializerOptions { WriteIndented = true }) :
264219
"";
265220
e.Session.GenericResponse(body, responseCode, headers);
266221
state.HasBeenSet = true;
@@ -272,6 +227,43 @@ _urlsToWatch is null ||
272227
}
273228
}
274229

230+
StoreRateLimitingHeaders(e);
275231
return Task.CompletedTask;
276232
}
233+
234+
private void StoreRateLimitingHeaders(ProxyRequestArgs e)
235+
{
236+
// add rate limiting headers if reached the threshold percentage
237+
if (_resourcesRemaining > _configuration.RateLimit - (_configuration.RateLimit * _configuration.WarningThresholdPercent / 100))
238+
{
239+
return;
240+
}
241+
242+
var headers = new List<HttpHeader>();
243+
var reset = _configuration.ResetFormat == RateLimitResetFormat.SecondsLeft ?
244+
(_resetTime - DateTime.Now).TotalSeconds.ToString("N0") : // drop decimals
245+
new DateTimeOffset(_resetTime).ToUnixTimeSeconds().ToString();
246+
headers.AddRange(new List<HttpHeader>
247+
{
248+
new HttpHeader(_configuration.HeaderLimit, _configuration.RateLimit.ToString()),
249+
new HttpHeader(_configuration.HeaderRemaining, _resourcesRemaining.ToString()),
250+
new HttpHeader(_configuration.HeaderReset, reset)
251+
});
252+
253+
ExposeRateLimitingForCors(headers, e);
254+
255+
e.PluginData.Add(Name, headers);
256+
}
257+
258+
private void ExposeRateLimitingForCors(IList<HttpHeader> headers, ProxyRequestArgs e)
259+
{
260+
var request = e.Session.HttpClient.Request;
261+
if (request.Headers.FirstOrDefault((h) => h.Name.Equals("Origin", StringComparison.OrdinalIgnoreCase)) is null)
262+
{
263+
return;
264+
}
265+
266+
headers.Add(new HttpHeader("Access-Control-Allow-Origin", "*"));
267+
headers.Add(new HttpHeader("Access-Control-Expose-Headers", $"{_configuration.HeaderLimit}, {_configuration.HeaderRemaining}, {_configuration.HeaderReset}, {_configuration.HeaderRetryAfter}"));
268+
}
277269
}

‎dev-proxy-plugins/MockResponses/GraphMockResponsePlugin.cs

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
using System.Text.Json;
66
using System.Text.RegularExpressions;
77
using Microsoft.DevProxy.Abstractions;
8+
using Microsoft.DevProxy.Plugins.Behavior;
9+
using Titanium.Web.Proxy.Models;
810

911
namespace Microsoft.DevProxy.Plugins.MockResponses;
1012

@@ -41,8 +43,15 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
4143
var requestId = Guid.NewGuid().ToString();
4244
var requestDate = DateTime.Now.ToString();
4345
var headers = ProxyUtils
44-
.BuildGraphResponseHeaders(e.Session.HttpClient.Request, requestId, requestDate)
45-
.ToDictionary(h => h.Name, h => h.Value);
46+
.BuildGraphResponseHeaders(e.Session.HttpClient.Request, requestId, requestDate);
47+
48+
if (e.PluginData.TryGetValue(nameof(RateLimitingPlugin), out var pluginData) &&
49+
pluginData is List<HttpHeader> rateLimitingHeaders)
50+
{
51+
ProxyUtils.MergeHeaders(headers, rateLimitingHeaders);
52+
}
53+
54+
var headersDictionary = headers.ToDictionary(h => h.Name, h => h.Value);
4655

4756
var mockResponse = GetMatchingMockResponse(request, e.Session.HttpClient.Request.RequestUri);
4857
if (mockResponse == null)
@@ -51,7 +60,7 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
5160
{
5261
Id = request.Id,
5362
Status = (int)HttpStatusCode.BadGateway,
54-
Headers = headers,
63+
Headers = headersDictionary,
5564
Body = new GraphBatchResponsePayloadResponseBody
5665
{
5766
Error = new GraphBatchResponsePayloadResponseBodyError
@@ -77,13 +86,13 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
7786
{
7887
foreach (var key in mockResponse.Response.Headers.Keys)
7988
{
80-
headers[key] = mockResponse.Response.Headers[key];
89+
headersDictionary[key] = mockResponse.Response.Headers[key];
8190
}
8291
}
8392
// default the content type to application/json unless set in the mock response
84-
if (!headers.Any(h => h.Key.Equals("content-type", StringComparison.OrdinalIgnoreCase)))
93+
if (!headersDictionary.Any(h => h.Key.Equals("content-type", StringComparison.OrdinalIgnoreCase)))
8594
{
86-
headers.Add("content-type", "application/json");
95+
headersDictionary.Add("content-type", "application/json");
8796
}
8897

8998
if (mockResponse.Response?.Body is not null)
@@ -118,7 +127,7 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
118127
{
119128
Id = request.Id,
120129
Status = (int)statusCode,
121-
Headers = headers,
130+
Headers = headersDictionary,
122131
Body = body
123132
};
124133

‎dev-proxy-plugins/MockResponses/MockResponsePlugin.cs

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
using Titanium.Web.Proxy.EventArguments;
1313
using Titanium.Web.Proxy.Http;
1414
using Titanium.Web.Proxy.Models;
15+
using Microsoft.DevProxy.Plugins.Behavior;
1516

1617
namespace Microsoft.DevProxy.Plugins.MockResponses;
1718

@@ -115,12 +116,12 @@ protected virtual Task OnRequest(object? sender, ProxyRequestArgs e)
115116
var matchingResponse = GetMatchingMockResponse(request);
116117
if (matchingResponse is not null)
117118
{
118-
ProcessMockResponse(e.Session, matchingResponse);
119+
ProcessMockResponse(e, matchingResponse);
119120
state.HasBeenSet = true;
120121
}
121122
else if (_configuration.BlockUnmockedRequests)
122123
{
123-
ProcessMockResponse(e.Session, new MockResponse
124+
ProcessMockResponse(e, new MockResponse
124125
{
125126
Request = new()
126127
{
@@ -202,12 +203,12 @@ private bool IsNthRequest(MockResponse mockResponse)
202203
return mockResponse.Request.Nth == nth;
203204
}
204205

205-
private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingResponse)
206+
private void ProcessMockResponse(ProxyRequestArgs e, MockResponse matchingResponse)
206207
{
207208
string? body = null;
208209
string requestId = Guid.NewGuid().ToString();
209210
string requestDate = DateTime.Now.ToString();
210-
var headers = ProxyUtils.BuildGraphResponseHeaders(e.HttpClient.Request, requestId, requestDate);
211+
var headers = ProxyUtils.BuildGraphResponseHeaders(e.Session.HttpClient.Request, requestId, requestDate);
211212
HttpStatusCode statusCode = HttpStatusCode.OK;
212213
if (matchingResponse.Response?.StatusCode is not null)
213214
{
@@ -234,6 +235,12 @@ private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingRespon
234235
headers.Add(new HttpHeader("content-type", "application/json"));
235236
}
236237

238+
if (e.PluginData.TryGetValue(nameof(RateLimitingPlugin), out var pluginData) &&
239+
pluginData is List<HttpHeader> rateLimitingHeaders)
240+
{
241+
ProxyUtils.MergeHeaders(headers, rateLimitingHeaders);
242+
}
243+
237244
if (matchingResponse.Response?.Body is not null)
238245
{
239246
var bodyString = JsonSerializer.Serialize(matchingResponse.Response.Body) as string;
@@ -254,8 +261,8 @@ private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingRespon
254261
else
255262
{
256263
var bodyBytes = File.ReadAllBytes(filePath);
257-
e.GenericResponse(bodyBytes, statusCode, headers);
258-
_logger?.LogRequest([$"{matchingResponse.Response.StatusCode ?? 200} {matchingResponse.Request?.Url}"], MessageType.Mocked, new LoggingContext(e));
264+
e.Session.GenericResponse(bodyBytes, statusCode, headers);
265+
_logger?.LogRequest([$"{matchingResponse.Response.StatusCode ?? 200} {matchingResponse.Request?.Url}"], MessageType.Mocked, new LoggingContext(e.Session));
259266
return;
260267
}
261268
}
@@ -264,8 +271,8 @@ private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingRespon
264271
body = bodyString;
265272
}
266273
}
267-
e.GenericResponse(body ?? string.Empty, statusCode, headers);
274+
e.Session.GenericResponse(body ?? string.Empty, statusCode, headers);
268275

269-
_logger?.LogRequest([$"{matchingResponse.Response?.StatusCode ?? 200} {matchingResponse.Request?.Url}"], MessageType.Mocked, new LoggingContext(e));
276+
_logger?.LogRequest([$"{matchingResponse.Response?.StatusCode ?? 200} {matchingResponse.Request?.Url}"], MessageType.Mocked, new LoggingContext(e.Session));
270277
}
271278
}

0 commit comments

Comments
 (0)