Skip to content

Commit c5c8227

Browse files
waldekmastykarzwaldekmastykarz
authored
Adds custom behavior to rate limiting. Closes #334 (#340)
1 parent e11aa30 commit c5c8227

File tree

4 files changed

+137
-13
lines changed

4 files changed

+137
-13
lines changed

‎m365-developer-proxy-abstractions/PluginEvents.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -194,8 +194,8 @@ public async Task RaiseProxyBeforeResponse(ProxyResponseArgs args) {
194194
await BeforeResponse?.InvokeAsync(this, args, null);
195195
}
196196

197-
public void RaiseProxyAfterResponse(ProxyResponseArgs args) {
198-
AfterResponse?.Invoke(this, args);
197+
public async Task RaiseProxyAfterResponse(ProxyResponseArgs args) {
198+
await AfterResponse?.Invoke(this, args);
199199
}
200200

201201
public void RaiseRequestLogged(RequestLogArgs args) {

‎m365-developer-proxy-plugins/Behavior/RateLimitingCustomResponseLoader.cs‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
using Microsoft365.DeveloperProxy.Abstractions;
5+
using Microsoft365.DeveloperProxy.Plugins.MocksResponses;
6+
using System.Text.Json;
7+
8+
namespace Microsoft365.DeveloperProxy.Plugins.Behavior;
9+
10+
internal class RateLimitingCustomResponseLoader : IDisposable
11+
{
12+
private readonly ILogger _logger;
13+
private readonly RateLimitConfiguration _configuration;
14+
15+
public RateLimitingCustomResponseLoader(ILogger logger, RateLimitConfiguration configuration)
16+
{
17+
_logger = logger ?? throw new ArgumentNullException(nameof(logger));
18+
_configuration = configuration ?? throw new ArgumentNullException(nameof(configuration));
19+
}
20+
21+
private string _customResponseFilePath => Path.Combine(Directory.GetCurrentDirectory(), _configuration.CustomResponseFile);
22+
private FileSystemWatcher? _watcher;
23+
24+
public void LoadResponse()
25+
{
26+
if (!File.Exists(_customResponseFilePath))
27+
{
28+
_logger.LogWarn($"File {_configuration.CustomResponseFile} not found. No response will be provided");
29+
return;
30+
}
31+
32+
try
33+
{
34+
using (FileStream stream = new FileStream(_customResponseFilePath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite))
35+
{
36+
using (StreamReader reader = new StreamReader(stream))
37+
{
38+
var responseString = reader.ReadToEnd();
39+
var response = JsonSerializer.Deserialize<MockResponse>(responseString);
40+
if (response is not null)
41+
{
42+
_configuration.CustomResponse = response;
43+
}
44+
}
45+
}
46+
}
47+
catch (Exception ex)
48+
{
49+
_logger.LogError($"An error has occurred while reading {_configuration.CustomResponseFile}:");
50+
_logger.LogError(ex.Message);
51+
}
52+
}
53+
54+
public void InitResponsesWatcher()
55+
{
56+
if (_watcher is not null)
57+
{
58+
return;
59+
}
60+
61+
string path = Path.GetDirectoryName(_customResponseFilePath) ?? throw new InvalidOperationException($"{_customResponseFilePath} is an invalid path");
62+
if (!File.Exists(_customResponseFilePath))
63+
{
64+
_logger.LogWarn($"File {_configuration.CustomResponseFile} not found. No mocks will be provided");
65+
return;
66+
}
67+
68+
_watcher = new FileSystemWatcher(Path.GetFullPath(path))
69+
{
70+
NotifyFilter = NotifyFilters.CreationTime
71+
| NotifyFilters.FileName
72+
| NotifyFilters.LastWrite
73+
| NotifyFilters.Size,
74+
Filter = Path.GetFileName(_customResponseFilePath)
75+
};
76+
_watcher.Changed += ResponseFile_Changed;
77+
_watcher.Created += ResponseFile_Changed;
78+
_watcher.Deleted += ResponseFile_Changed;
79+
_watcher.Renamed += ResponseFile_Changed;
80+
_watcher.EnableRaisingEvents = true;
81+
82+
LoadResponse();
83+
}
84+
85+
private void ResponseFile_Changed(object sender, FileSystemEventArgs e)
86+
{
87+
LoadResponse();
88+
}
89+
90+
public void Dispose()
91+
{
92+
_watcher?.Dispose();
93+
}
94+
}

‎m365-developer-proxy-plugins/Behavior/RateLimitingPlugin.cs‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
using Microsoft.Extensions.Configuration;
55
using Microsoft365.DeveloperProxy.Abstractions;
6+
using Microsoft365.DeveloperProxy.Plugins.MocksResponses;
67
using System.Net;
78
using System.Text.Json;
89
using System.Text.RegularExpressions;
@@ -11,6 +12,11 @@
1112

1213
namespace Microsoft365.DeveloperProxy.Plugins.Behavior;
1314

15+
public enum RateLimitResponseWhenLimitExceeded {
16+
Throttle,
17+
Custom
18+
}
19+
1420
public class RateLimitConfiguration {
1521
public string HeaderLimit { get; set; } = "RateLimit-Limit";
1622
public string HeaderRemaining { get; set; } = "RateLimit-Remaining";
@@ -21,6 +27,9 @@ public class RateLimitConfiguration {
2127
public int WarningThresholdPercent { get; set; } = 80;
2228
public int RateLimit { get; set; } = 120;
2329
public int RetryAfterSeconds { get; set; } = 5;
30+
public RateLimitResponseWhenLimitExceeded WhenLimitExceeded { get; set; } = RateLimitResponseWhenLimitExceeded.Throttle;
31+
public string CustomResponseFile { get; set; } = "rate-limit-response.json";
32+
public MockResponse? CustomResponse { get; set; }
2433
}
2534

2635
public class RateLimitingPlugin : BaseProxyPlugin {
@@ -30,6 +39,7 @@ public class RateLimitingPlugin : BaseProxyPlugin {
3039
// first request and can set the initial values
3140
private int _resourcesRemaining = -1;
3241
private DateTime _resetTime = DateTime.MinValue;
42+
private RateLimitingCustomResponseLoader? _loader = null;
3343

3444
private ThrottlingInfo ShouldThrottle(Request request, string throttlingKey) {
3545
var throttleKeyForRequest = BuildThrottleKey(request);
@@ -129,6 +139,10 @@ public override void Register(IPluginEvents pluginEvents,
129139
base.Register(pluginEvents, context, urlsToWatch, configSection);
130140

131141
configSection?.Bind(_configuration);
142+
_loader = new RateLimitingCustomResponseLoader(_logger!, _configuration);
143+
// load the responses from the configured mocks file
144+
_loader.InitResponsesWatcher();
145+
132146
pluginEvents.BeforeRequest += OnRequest;
133147
pluginEvents.BeforeResponse += OnResponse;
134148
}
@@ -169,17 +183,34 @@ _urlsToWatch is null ||
169183
// subtract the cost of the request
170184
_resourcesRemaining -= _configuration.CostPerRequest;
171185
if (_resourcesRemaining < 0) {
186+
_resourcesRemaining = 0;
172187
var request = e.Session.HttpClient.Request;
173188

174189
_logger?.LogRequest(new[] { $"Exceeded resource limit when calling {request.Url}.", "Request will be throttled" }, MessageType.Failed, new LoggingContext(e.Session));
175-
e.ThrottledRequests.Add(new ThrottlerInfo(
176-
BuildThrottleKey(request),
177-
ShouldThrottle,
178-
DateTime.Now.AddSeconds(_configuration.RetryAfterSeconds)
179-
));
180-
181-
ThrottleResponse(e);
182-
state.HasBeenSet = true;
190+
if (_configuration.WhenLimitExceeded == RateLimitResponseWhenLimitExceeded.Throttle) {
191+
e.ThrottledRequests.Add(new ThrottlerInfo(
192+
BuildThrottleKey(request),
193+
ShouldThrottle,
194+
DateTime.Now.AddSeconds(_configuration.RetryAfterSeconds)
195+
));
196+
ThrottleResponse(e);
197+
state.HasBeenSet = true;
198+
}
199+
else {
200+
if (_configuration.CustomResponse is not null) {
201+
var headers = _configuration.CustomResponse.ResponseHeaders is not null ?
202+
_configuration.CustomResponse.ResponseHeaders.Select(h => new HttpHeader(h.Key, h.Value)) :
203+
Array.Empty<HttpHeader>();
204+
string body = _configuration.CustomResponse.ResponseBody is not null ?
205+
JsonSerializer.Serialize(_configuration.CustomResponse.ResponseBody, new JsonSerializerOptions { WriteIndented = true }) :
206+
"";
207+
e.Session.GenericResponse(body, (HttpStatusCode)(_configuration.CustomResponse.ResponseCode ?? 200), headers);
208+
state.HasBeenSet = true;
209+
}
210+
else {
211+
_logger?.LogRequest(new[] { $"Custom behavior not set {_configuration.CustomResponseFile} not found." }, MessageType.Failed, new LoggingContext(e.Session));
212+
}
213+
}
183214
}
184215
}
185216
}

‎m365-developer-proxy-plugins/MockResponses/MockResponsesLoader.cs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
using Microsoft365.DeveloperProxy.Abstractions;
55
using System.Text.Json;
6-
using System.IO;
76

87
namespace Microsoft365.DeveloperProxy.Plugins.MocksResponses;
98

@@ -33,8 +32,8 @@ public void LoadResponses() {
3332
var responsesConfig = JsonSerializer.Deserialize<MockResponseConfiguration>(responsesString);
3433
IEnumerable<MockResponse>? configResponses = responsesConfig?.Responses;
3534
if (configResponses is not null) {
36-
_configuration.Responses = configResponses;
37-
_logger.LogInfo($"Mock responses for {configResponses.Count()} url patterns loaded from {_configuration.MocksFile}");
35+
_configuration.Responses = configResponses;
36+
_logger.LogInfo($"Mock responses for {configResponses.Count()} url patterns loaded from {_configuration.MocksFile}");
3837
}
3938
}
4039
}

0 commit comments

Comments
 (0)