Skip to content

Commit 1390151

Browse files
waldekmastykarzwaldekmastykarz
authored
Adds support for formatting rate limiting reset time. Closes #331 (#349)
1 parent 1f67e4b commit 1390151

File tree

1 file changed

+80
-44
lines changed

1 file changed

+80
-44
lines changed

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

Lines changed: 80 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,25 @@
1212

1313
namespace Microsoft365.DeveloperProxy.Plugins.Behavior;
1414

15-
public enum RateLimitResponseWhenLimitExceeded {
15+
public enum RateLimitResponseWhenLimitExceeded
16+
{
1617
Throttle,
1718
Custom
1819
}
1920

20-
public class RateLimitConfiguration {
21+
public enum RateLimitResetFormat
22+
{
23+
SecondsLeft,
24+
UtcEpochSeconds
25+
}
26+
27+
public class RateLimitConfiguration
28+
{
2129
public string HeaderLimit { get; set; } = "RateLimit-Limit";
2230
public string HeaderRemaining { get; set; } = "RateLimit-Remaining";
2331
public string HeaderReset { get; set; } = "RateLimit-Reset";
2432
public string HeaderRetryAfter { get; set; } = "Retry-After";
33+
public RateLimitResetFormat ResetFormat { get; set; } = RateLimitResetFormat.SecondsLeft;
2534
public int CostPerRequest { get; set; } = 2;
2635
public int ResetTimeWindowSeconds { get; set; } = 60;
2736
public int WarningThresholdPercent { get; set; } = 80;
@@ -32,7 +41,8 @@ public class RateLimitConfiguration {
3241
public MockResponse? CustomResponse { get; set; }
3342
}
3443

35-
public class RateLimitingPlugin : BaseProxyPlugin {
44+
public class RateLimitingPlugin : BaseProxyPlugin
45+
{
3646
public override string Name => nameof(RateLimitingPlugin);
3747
private readonly RateLimitConfiguration _configuration = new();
3848
// initial values so that we know when we intercept the
@@ -41,31 +51,37 @@ public class RateLimitingPlugin : BaseProxyPlugin {
4151
private DateTime _resetTime = DateTime.MinValue;
4252
private RateLimitingCustomResponseLoader? _loader = null;
4353

44-
private ThrottlingInfo ShouldThrottle(Request request, string throttlingKey) {
54+
private ThrottlingInfo ShouldThrottle(Request request, string throttlingKey)
55+
{
4556
var throttleKeyForRequest = BuildThrottleKey(request);
4657
return new ThrottlingInfo(throttleKeyForRequest == throttlingKey ? _configuration.RetryAfterSeconds : 0, _configuration.HeaderRetryAfter);
4758
}
4859

4960
private void ThrottleResponse(ProxyRequestArgs e) => UpdateProxyResponse(e, HttpStatusCode.TooManyRequests);
5061

51-
private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorStatus) {
62+
private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorStatus)
63+
{
5264
var headers = new List<HttpHeader>();
5365
var body = string.Empty;
5466
var request = e.Session.HttpClient.Request;
5567
var response = e.Session.HttpClient.Response;
5668

5769
// resources exceeded
58-
if (errorStatus == HttpStatusCode.TooManyRequests) {
59-
if (ProxyUtils.IsGraphRequest(request)) {
70+
if (errorStatus == HttpStatusCode.TooManyRequests)
71+
{
72+
if (ProxyUtils.IsGraphRequest(request))
73+
{
6074
string requestId = Guid.NewGuid().ToString();
6175
string requestDate = DateTime.Now.ToString();
6276
headers.AddRange(ProxyUtils.BuildGraphResponseHeaders(request, requestId, requestDate));
6377

6478
body = JsonSerializer.Serialize(new GraphErrorResponseBody(
65-
new GraphErrorResponseError {
79+
new GraphErrorResponseError
80+
{
6681
Code = new Regex("([A-Z])").Replace(errorStatus.ToString(), m => { return $" {m.Groups[1]}"; }).Trim(),
6782
Message = BuildApiErrorMessage(request),
68-
InnerError = new GraphErrorResponseInnerError {
83+
InnerError = new GraphErrorResponseInnerError
84+
{
6985
RequestId = requestId,
7086
Date = requestDate
7187
}
@@ -80,35 +96,47 @@ private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorS
8096
}
8197

8298
// add rate limiting headers if reached the threshold percentage
83-
if (_resourcesRemaining <= _configuration.RateLimit - (_configuration.RateLimit * _configuration.WarningThresholdPercent / 100)) {
99+
if (_resourcesRemaining <= _configuration.RateLimit - (_configuration.RateLimit * _configuration.WarningThresholdPercent / 100))
100+
{
101+
var reset = _configuration.ResetFormat == RateLimitResetFormat.SecondsLeft ?
102+
(_resetTime - DateTime.Now).TotalSeconds.ToString("N0") : // drop decimals
103+
new DateTimeOffset(_resetTime).ToUnixTimeSeconds().ToString();
84104
headers.AddRange(new List<HttpHeader> {
85105
new HttpHeader(_configuration.HeaderLimit, _configuration.RateLimit.ToString()),
86106
new HttpHeader(_configuration.HeaderRemaining, _resourcesRemaining.ToString()),
87-
new HttpHeader(_configuration.HeaderReset, (_resetTime - DateTime.Now).TotalSeconds.ToString("N0")) // drop decimals
107+
new HttpHeader(_configuration.HeaderReset, reset)
88108
});
89109

90110
// make rate limiting information available for CORS requests
91-
if (request.Headers.FirstOrDefault((h) => h.Name.Equals("Origin", StringComparison.OrdinalIgnoreCase)) is not null) {
92-
if (!response.Headers.HeaderExists("Access-Control-Allow-Origin")) {
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+
{
93115
headers.Add(new HttpHeader("Access-Control-Allow-Origin", "*"));
94116
}
95117
var exposeHeadersHeader = response.Headers.FirstOrDefault((h) => h.Name.Equals("Access-Control-Expose-Headers", StringComparison.OrdinalIgnoreCase));
96118
var headerValue = "";
97-
if (exposeHeadersHeader is null) {
119+
if (exposeHeadersHeader is null)
120+
{
98121
headerValue = $"{_configuration.HeaderLimit}, {_configuration.HeaderRemaining}, {_configuration.HeaderReset}, {_configuration.HeaderRetryAfter}";
99122
}
100-
else {
123+
else
124+
{
101125
headerValue = exposeHeadersHeader.Value;
102-
if (!headerValue.Contains(_configuration.HeaderLimit)) {
126+
if (!headerValue.Contains(_configuration.HeaderLimit))
127+
{
103128
headerValue += $", {_configuration.HeaderLimit}";
104129
}
105-
if (!headerValue.Contains(_configuration.HeaderRemaining)) {
130+
if (!headerValue.Contains(_configuration.HeaderRemaining))
131+
{
106132
headerValue += $", {_configuration.HeaderRemaining}";
107133
}
108-
if (!headerValue.Contains(_configuration.HeaderReset)) {
134+
if (!headerValue.Contains(_configuration.HeaderReset))
135+
{
109136
headerValue += $", {_configuration.HeaderReset}";
110137
}
111-
if (!headerValue.Contains(_configuration.HeaderRetryAfter)) {
138+
if (!headerValue.Contains(_configuration.HeaderRetryAfter))
139+
{
112140
headerValue += $", {_configuration.HeaderRetryAfter}";
113141
}
114142
response.Headers.RemoveHeader("Access-Control-Expose-Headers");
@@ -118,32 +146,29 @@ private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorS
118146
}
119147
}
120148

121-
// add headers to the original API response
122-
headers.ForEach(h => {
123-
if (!response.Headers.HeaderExists(h.Name)) {
124-
e.Session.HttpClient.Response.Headers.AddHeader(h);
125-
}
126-
else {
127-
e.Session.HttpClient.Response.Headers.RemoveHeader(h.Name);
128-
e.Session.HttpClient.Response.Headers.AddHeader(h);
129-
}
130-
});
149+
// add headers to the original API response, avoiding duplicates
150+
headers.ForEach(h => e.Session.HttpClient.Response.Headers.RemoveHeader(h.Name));
151+
e.Session.HttpClient.Response.Headers.AddHeaders(headers);
131152
}
132153
private static string BuildApiErrorMessage(Request r) => $"Some error was generated by the proxy. {(ProxyUtils.IsGraphRequest(r) ? ProxyUtils.IsSdkRequest(r) ? "" : String.Join(' ', MessageUtils.BuildUseSdkForErrorsMessage(r)) : "")}";
133154

134-
private string BuildThrottleKey(Request r) {
135-
if (ProxyUtils.IsGraphRequest(r)) {
155+
private string BuildThrottleKey(Request r)
156+
{
157+
if (ProxyUtils.IsGraphRequest(r))
158+
{
136159
return GraphUtils.BuildThrottleKey(r);
137160
}
138-
else {
161+
else
162+
{
139163
return r.RequestUri.Host;
140164
}
141165
}
142166

143167
public override void Register(IPluginEvents pluginEvents,
144168
IProxyContext context,
145169
ISet<UrlToWatch> urlsToWatch,
146-
IConfigurationSection? configSection = null) {
170+
IConfigurationSection? configSection = null)
171+
{
147172
base.Register(pluginEvents, context, urlsToWatch, configSection);
148173

149174
configSection?.Bind(_configuration);
@@ -156,41 +181,49 @@ public override void Register(IPluginEvents pluginEvents,
156181
}
157182

158183
// add rate limiting headers to the response from the API
159-
private async Task OnResponse(object? sender, ProxyResponseArgs e) {
184+
private async Task OnResponse(object? sender, ProxyResponseArgs e)
185+
{
160186
if (_urlsToWatch is null ||
161-
!e.HasRequestUrlMatch(_urlsToWatch)) {
187+
!e.HasRequestUrlMatch(_urlsToWatch))
188+
{
162189
return;
163190
}
164191

165192
UpdateProxyResponse(e, HttpStatusCode.OK);
166193
}
167194

168-
private async Task OnRequest(object? sender, ProxyRequestArgs e) {
195+
private async Task OnRequest(object? sender, ProxyRequestArgs e)
196+
{
169197
var session = e.Session;
170198
var state = e.ResponseState;
171199
if (e.ResponseState.HasBeenSet ||
172200
_urlsToWatch is null ||
173-
!e.ShouldExecute(_urlsToWatch)) {
201+
!e.ShouldExecute(_urlsToWatch))
202+
{
174203
return;
175204
}
176205

177206
// set the initial values for the first request
178-
if (_resetTime == DateTime.MinValue) {
207+
if (_resetTime == DateTime.MinValue)
208+
{
179209
_resetTime = DateTime.Now.AddSeconds(_configuration.ResetTimeWindowSeconds);
180210
}
181-
if (_resourcesRemaining == -1) {
211+
if (_resourcesRemaining == -1)
212+
{
182213
_resourcesRemaining = _configuration.RateLimit;
183214
}
184215

185216
// see if we passed the reset time window
186-
if (DateTime.Now > _resetTime) {
217+
if (DateTime.Now > _resetTime)
218+
{
187219
_resourcesRemaining = _configuration.RateLimit;
188220
_resetTime = DateTime.Now.AddSeconds(_configuration.ResetTimeWindowSeconds);
189221
}
190222

191223
// subtract the cost of the request
192224
_resourcesRemaining -= _configuration.CostPerRequest;
193-
if (_resourcesRemaining < 0) {
225+
if (_resourcesRemaining < 0)
226+
{
194227
_resourcesRemaining = 0;
195228
var request = e.Session.HttpClient.Request;
196229

@@ -204,8 +237,10 @@ _urlsToWatch is null ||
204237
ThrottleResponse(e);
205238
state.HasBeenSet = true;
206239
}
207-
else {
208-
if (_configuration.CustomResponse is not null) {
240+
else
241+
{
242+
if (_configuration.CustomResponse is not null)
243+
{
209244
var headers = _configuration.CustomResponse.ResponseHeaders is not null ?
210245
_configuration.CustomResponse.ResponseHeaders.Select(h => new HttpHeader(h.Key, h.Value)) :
211246
Array.Empty<HttpHeader>();
@@ -215,7 +250,8 @@ _urlsToWatch is null ||
215250
e.Session.GenericResponse(body, (HttpStatusCode)(_configuration.CustomResponse.ResponseCode ?? 200), headers);
216251
state.HasBeenSet = true;
217252
}
218-
else {
253+
else
254+
{
219255
_logger?.LogRequest(new[] { $"Custom behavior not set {_configuration.CustomResponseFile} not found." }, MessageType.Failed, new LoggingContext(e.Session));
220256
}
221257
}

0 commit comments

Comments
 (0)