Skip to content

Commit

Permalink
Combines rate limiting with mocking. Closes #350 (#461)
Browse files Browse the repository at this point in the history
  • Loading branch information
waldekmastykarz authored Dec 23, 2023
1 parent 3eec3d6 commit ad58a21
Show file tree
Hide file tree
Showing 6 changed files with 103 additions and 70 deletions.
1 change: 1 addition & 0 deletions dev-proxy-abstractions/PluginEvents.cs
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ internal ProxyHttpEventArgsBase(SessionEventArgs session, IList<ThrottlerInfo> t

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

public bool HasRequestUrlMatch(ISet<UrlToWatch> watchedUrls)
{
Expand Down
24 changes: 24 additions & 0 deletions dev-proxy-abstractions/ProxyUtils.cs
Original file line number Diff line number Diff line change
Expand Up @@ -319,4 +319,28 @@ public static string ProductVersion
return _productVersion;
}
}

public static void MergeHeaders(IList<HttpHeader> allHeaders, IList<HttpHeader> headersToAdd)
{
foreach (var header in headersToAdd)
{
var existingHeader = allHeaders.FirstOrDefault(h => h.Name.Equals(header.Name, StringComparison.OrdinalIgnoreCase));
if (existingHeader is not null)
{
if (header.Name.Equals("Access-Control-Expose-Headers", StringComparison.OrdinalIgnoreCase))
{
var existingValues = existingHeader.Value.Split(',').Select(v => v.Trim());
var newValues = header.Value.Split(',').Select(v => v.Trim());
var allValues = existingValues.Union(newValues).Distinct();
allHeaders.Remove(existingHeader);
allHeaders.Add(new HttpHeader(header.Name, string.Join(", ", allValues)));
continue;
}

allHeaders.Remove(existingHeader);
}

allHeaders.Add(header);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ public void LoadResponse()
using (StreamReader reader = new StreamReader(stream))
{
var responseString = reader.ReadToEnd();
var response = JsonSerializer.Deserialize<MockResponse>(responseString);
var response = JsonSerializer.Deserialize<MockResponseResponse>(responseString);
if (response is not null)
{
_configuration.CustomResponse = response;
Expand Down
100 changes: 46 additions & 54 deletions dev-proxy-plugins/Behavior/RateLimitingPlugin.cs
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ public class RateLimitConfiguration
public int RetryAfterSeconds { get; set; } = 5;
public RateLimitResponseWhenLimitExceeded WhenLimitExceeded { get; set; } = RateLimitResponseWhenLimitExceeded.Throttle;
public string CustomResponseFile { get; set; } = "rate-limit-response.json";
public MockResponse? CustomResponse { get; set; }
public MockResponseResponse? CustomResponse { get; set; }
}

public class RateLimitingPlugin : BaseProxyPlugin
Expand Down Expand Up @@ -95,55 +95,10 @@ private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorS
return;
}

// add rate limiting headers if reached the threshold percentage
if (_resourcesRemaining <= _configuration.RateLimit - (_configuration.RateLimit * _configuration.WarningThresholdPercent / 100))
if (e.PluginData.TryGetValue(Name, out var pluginData) &&
pluginData is List<HttpHeader> rateLimitingHeaders)
{
var reset = _configuration.ResetFormat == RateLimitResetFormat.SecondsLeft ?
(_resetTime - DateTime.Now).TotalSeconds.ToString("N0") : // drop decimals
new DateTimeOffset(_resetTime).ToUnixTimeSeconds().ToString();
headers.AddRange(new List<HttpHeader> {
new HttpHeader(_configuration.HeaderLimit, _configuration.RateLimit.ToString()),
new HttpHeader(_configuration.HeaderRemaining, _resourcesRemaining.ToString()),
new HttpHeader(_configuration.HeaderReset, reset)
});

// make rate limiting information available for CORS requests
if (request.Headers.FirstOrDefault((h) => h.Name.Equals("Origin", StringComparison.OrdinalIgnoreCase)) is not null)
{
if (!response.Headers.HeaderExists("Access-Control-Allow-Origin"))
{
headers.Add(new HttpHeader("Access-Control-Allow-Origin", "*"));
}
var exposeHeadersHeader = response.Headers.FirstOrDefault((h) => h.Name.Equals("Access-Control-Expose-Headers", StringComparison.OrdinalIgnoreCase));
var headerValue = "";
if (exposeHeadersHeader is null)
{
headerValue = $"{_configuration.HeaderLimit}, {_configuration.HeaderRemaining}, {_configuration.HeaderReset}, {_configuration.HeaderRetryAfter}";
}
else
{
headerValue = exposeHeadersHeader.Value;
if (!headerValue.Contains(_configuration.HeaderLimit))
{
headerValue += $", {_configuration.HeaderLimit}";
}
if (!headerValue.Contains(_configuration.HeaderRemaining))
{
headerValue += $", {_configuration.HeaderRemaining}";
}
if (!headerValue.Contains(_configuration.HeaderReset))
{
headerValue += $", {_configuration.HeaderReset}";
}
if (!headerValue.Contains(_configuration.HeaderRetryAfter))
{
headerValue += $", {_configuration.HeaderRetryAfter}";
}
response.Headers.RemoveHeader("Access-Control-Expose-Headers");
}

headers.Add(new HttpHeader("Access-Control-Expose-Headers", headerValue));
}
ProxyUtils.MergeHeaders(headers, rateLimitingHeaders);
}

// add headers to the original API response, avoiding duplicates
Expand Down Expand Up @@ -244,12 +199,12 @@ _urlsToWatch is null ||
{
if (_configuration.CustomResponse is not null)
{
var headers = _configuration.CustomResponse.Response?.Headers is not null ?
_configuration.CustomResponse.Response.Headers.Select(h => new HttpHeader(h.Key, h.Value)) :
var headers = _configuration.CustomResponse.Headers is not null ?
_configuration.CustomResponse.Headers.Select(h => new HttpHeader(h.Key, h.Value)) :
Array.Empty<HttpHeader>();

// allow custom throttling response
var responseCode = (HttpStatusCode)(_configuration.CustomResponse.Response?.StatusCode ?? 200);
var responseCode = (HttpStatusCode)(_configuration.CustomResponse.StatusCode ?? 200);
if (responseCode == HttpStatusCode.TooManyRequests)
{
e.ThrottledRequests.Add(new ThrottlerInfo(
Expand All @@ -259,8 +214,8 @@ _urlsToWatch is null ||
));
}

string body = _configuration.CustomResponse.Response?.Body is not null ?
JsonSerializer.Serialize(_configuration.CustomResponse.Response.Body, new JsonSerializerOptions { WriteIndented = true }) :
string body = _configuration.CustomResponse.Body is not null ?
JsonSerializer.Serialize(_configuration.CustomResponse.Body, new JsonSerializerOptions { WriteIndented = true }) :
"";
e.Session.GenericResponse(body, responseCode, headers);
state.HasBeenSet = true;
Expand All @@ -272,6 +227,43 @@ _urlsToWatch is null ||
}
}

StoreRateLimitingHeaders(e);
return Task.CompletedTask;
}

private void StoreRateLimitingHeaders(ProxyRequestArgs e)
{
// add rate limiting headers if reached the threshold percentage
if (_resourcesRemaining > _configuration.RateLimit - (_configuration.RateLimit * _configuration.WarningThresholdPercent / 100))
{
return;
}

var headers = new List<HttpHeader>();
var reset = _configuration.ResetFormat == RateLimitResetFormat.SecondsLeft ?
(_resetTime - DateTime.Now).TotalSeconds.ToString("N0") : // drop decimals
new DateTimeOffset(_resetTime).ToUnixTimeSeconds().ToString();
headers.AddRange(new List<HttpHeader>
{
new HttpHeader(_configuration.HeaderLimit, _configuration.RateLimit.ToString()),
new HttpHeader(_configuration.HeaderRemaining, _resourcesRemaining.ToString()),
new HttpHeader(_configuration.HeaderReset, reset)
});

ExposeRateLimitingForCors(headers, e);

e.PluginData.Add(Name, headers);
}

private void ExposeRateLimitingForCors(IList<HttpHeader> headers, ProxyRequestArgs e)
{
var request = e.Session.HttpClient.Request;
if (request.Headers.FirstOrDefault((h) => h.Name.Equals("Origin", StringComparison.OrdinalIgnoreCase)) is null)
{
return;
}

headers.Add(new HttpHeader("Access-Control-Allow-Origin", "*"));
headers.Add(new HttpHeader("Access-Control-Expose-Headers", $"{_configuration.HeaderLimit}, {_configuration.HeaderRemaining}, {_configuration.HeaderReset}, {_configuration.HeaderRetryAfter}"));
}
}
23 changes: 16 additions & 7 deletions dev-proxy-plugins/MockResponses/GraphMockResponsePlugin.cs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
using System.Text.Json;
using System.Text.RegularExpressions;
using Microsoft.DevProxy.Abstractions;
using Microsoft.DevProxy.Plugins.Behavior;
using Titanium.Web.Proxy.Models;

namespace Microsoft.DevProxy.Plugins.MockResponses;

Expand Down Expand Up @@ -41,8 +43,15 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
var requestId = Guid.NewGuid().ToString();
var requestDate = DateTime.Now.ToString();
var headers = ProxyUtils
.BuildGraphResponseHeaders(e.Session.HttpClient.Request, requestId, requestDate)
.ToDictionary(h => h.Name, h => h.Value);
.BuildGraphResponseHeaders(e.Session.HttpClient.Request, requestId, requestDate);

if (e.PluginData.TryGetValue(nameof(RateLimitingPlugin), out var pluginData) &&
pluginData is List<HttpHeader> rateLimitingHeaders)
{
ProxyUtils.MergeHeaders(headers, rateLimitingHeaders);
}

var headersDictionary = headers.ToDictionary(h => h.Name, h => h.Value);

var mockResponse = GetMatchingMockResponse(request, e.Session.HttpClient.Request.RequestUri);
if (mockResponse == null)
Expand All @@ -51,7 +60,7 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
{
Id = request.Id,
Status = (int)HttpStatusCode.BadGateway,
Headers = headers,
Headers = headersDictionary,
Body = new GraphBatchResponsePayloadResponseBody
{
Error = new GraphBatchResponsePayloadResponseBodyError
Expand All @@ -77,13 +86,13 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
{
foreach (var key in mockResponse.Response.Headers.Keys)
{
headers[key] = mockResponse.Response.Headers[key];
headersDictionary[key] = mockResponse.Response.Headers[key];
}
}
// default the content type to application/json unless set in the mock response
if (!headers.Any(h => h.Key.Equals("content-type", StringComparison.OrdinalIgnoreCase)))
if (!headersDictionary.Any(h => h.Key.Equals("content-type", StringComparison.OrdinalIgnoreCase)))
{
headers.Add("content-type", "application/json");
headersDictionary.Add("content-type", "application/json");
}

if (mockResponse.Response?.Body is not null)
Expand Down Expand Up @@ -118,7 +127,7 @@ protected override async Task OnRequest(object? sender, ProxyRequestArgs e)
{
Id = request.Id,
Status = (int)statusCode,
Headers = headers,
Headers = headersDictionary,
Body = body
};

Expand Down
23 changes: 15 additions & 8 deletions dev-proxy-plugins/MockResponses/MockResponsePlugin.cs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
using Titanium.Web.Proxy.EventArguments;
using Titanium.Web.Proxy.Http;
using Titanium.Web.Proxy.Models;
using Microsoft.DevProxy.Plugins.Behavior;

namespace Microsoft.DevProxy.Plugins.MockResponses;

Expand Down Expand Up @@ -115,12 +116,12 @@ protected virtual Task OnRequest(object? sender, ProxyRequestArgs e)
var matchingResponse = GetMatchingMockResponse(request);
if (matchingResponse is not null)
{
ProcessMockResponse(e.Session, matchingResponse);
ProcessMockResponse(e, matchingResponse);
state.HasBeenSet = true;
}
else if (_configuration.BlockUnmockedRequests)
{
ProcessMockResponse(e.Session, new MockResponse
ProcessMockResponse(e, new MockResponse
{
Request = new()
{
Expand Down Expand Up @@ -202,12 +203,12 @@ private bool IsNthRequest(MockResponse mockResponse)
return mockResponse.Request.Nth == nth;
}

private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingResponse)
private void ProcessMockResponse(ProxyRequestArgs e, MockResponse matchingResponse)
{
string? body = null;
string requestId = Guid.NewGuid().ToString();
string requestDate = DateTime.Now.ToString();
var headers = ProxyUtils.BuildGraphResponseHeaders(e.HttpClient.Request, requestId, requestDate);
var headers = ProxyUtils.BuildGraphResponseHeaders(e.Session.HttpClient.Request, requestId, requestDate);
HttpStatusCode statusCode = HttpStatusCode.OK;
if (matchingResponse.Response?.StatusCode is not null)
{
Expand All @@ -234,6 +235,12 @@ private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingRespon
headers.Add(new HttpHeader("content-type", "application/json"));
}

if (e.PluginData.TryGetValue(nameof(RateLimitingPlugin), out var pluginData) &&
pluginData is List<HttpHeader> rateLimitingHeaders)
{
ProxyUtils.MergeHeaders(headers, rateLimitingHeaders);
}

if (matchingResponse.Response?.Body is not null)
{
var bodyString = JsonSerializer.Serialize(matchingResponse.Response.Body) as string;
Expand All @@ -254,8 +261,8 @@ private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingRespon
else
{
var bodyBytes = File.ReadAllBytes(filePath);
e.GenericResponse(bodyBytes, statusCode, headers);
_logger?.LogRequest([$"{matchingResponse.Response.StatusCode ?? 200} {matchingResponse.Request?.Url}"], MessageType.Mocked, new LoggingContext(e));
e.Session.GenericResponse(bodyBytes, statusCode, headers);
_logger?.LogRequest([$"{matchingResponse.Response.StatusCode ?? 200} {matchingResponse.Request?.Url}"], MessageType.Mocked, new LoggingContext(e.Session));
return;
}
}
Expand All @@ -264,8 +271,8 @@ private void ProcessMockResponse(SessionEventArgs e, MockResponse matchingRespon
body = bodyString;
}
}
e.GenericResponse(body ?? string.Empty, statusCode, headers);
e.Session.GenericResponse(body ?? string.Empty, statusCode, headers);

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

0 comments on commit ad58a21

Please sign in to comment.