Centralizes retry-after handling. Closes #239 (#257)

* Centralizes retry-after handling. Closes #239
Improves throttling

* Fixes issues with rate limiting and retry-after

Author: waldekmastykarz

msgraph-developer-proxy-abstractions/PluginEvents.cs

@@ -4,6 +4,7 @@
 using System.CommandLine;
 using System.CommandLine.Invocation;
 using Titanium.Web.Proxy.EventArguments;
+using Titanium.Web.Proxy.Http;
 
 namespace Microsoft.Graph.DeveloperProxy.Abstractions;
 
@@ -12,11 +13,52 @@ public interface IProxyContext {
     ILogger Logger { get; }
 }
 
+public class ThrottlerInfo {
+    /// <summary>
+    /// Throttling key used to identify which requests should be throttled.
+    /// Can be set to a hostname, full URL or a custom string value, that
+    /// represents for example a portion of the API
+    /// </summary>
+    public string ThrottlingKey { get; private set; }
+    /// <summary>
+    /// Function responsible for matching the request to the throttling key.
+    /// Takes as arguments:
+    /// - intercepted request
+    /// - the throttling key
+    /// Returns an instance of ThrottlingInfo that contains information
+    /// whether the request should be throttled or not.
+    /// </summary>
+    public Func<Request, string, ThrottlingInfo> ShouldThrottle { get; private set; }
+    /// <summary>
+    /// Time when the throttling window will be reset
+    /// </summary>
+    public DateTime ResetTime { get; set; }
+
+    public ThrottlerInfo(string throttlingKey, Func<Request, string, ThrottlingInfo> shouldThrottle, DateTime resetTime) {
+        ThrottlingKey = throttlingKey ?? throw new ArgumentNullException(nameof(throttlingKey));
+        ShouldThrottle = shouldThrottle ?? throw new ArgumentNullException(nameof(shouldThrottle));
+        ResetTime = resetTime;
+    }
+}
+
+public class ThrottlingInfo {
+    public int ThrottleForSeconds { get; set; }
+    public string RetryAfterHeaderName { get; set; }
+
+    public ThrottlingInfo(int throttleForSeconds, string retryAfterHeaderName) {
+        ThrottleForSeconds = throttleForSeconds;
+        RetryAfterHeaderName = retryAfterHeaderName ?? throw new ArgumentNullException(nameof(retryAfterHeaderName));
+    }
+}
+
 public class ProxyHttpEventArgsBase {
-    internal ProxyHttpEventArgsBase(SessionEventArgs session) =>
+    internal ProxyHttpEventArgsBase(SessionEventArgs session, IList<ThrottlerInfo> throttledRequests) {
         Session = session ?? throw new ArgumentNullException(nameof(session));
+        ThrottledRequests = throttledRequests ?? throw new ArgumentNullException(nameof(throttledRequests));
+    }
 
     public SessionEventArgs Session { get; }
+    public IList<ThrottlerInfo> ThrottledRequests { get; }
 
     public bool HasRequestUrlMatch(ISet<UrlToWatch> watchedUrls) {
         var match = watchedUrls.FirstOrDefault(r => r.Url.IsMatch(Session.HttpClient.Request.RequestUri.AbsoluteUri));
@@ -25,7 +67,7 @@ public bool HasRequestUrlMatch(ISet<UrlToWatch> watchedUrls) {
 }
 
 public class ProxyRequestArgs : ProxyHttpEventArgsBase {
-    public ProxyRequestArgs(SessionEventArgs session, ResponseState responseState) : base(session) {
+    public ProxyRequestArgs(SessionEventArgs session, IList<ThrottlerInfo> throttledRequests, ResponseState responseState) : base(session, throttledRequests) {
         ResponseState = responseState ?? throw new ArgumentNullException(nameof(responseState));
     }
     public ResponseState ResponseState { get; }
@@ -36,7 +78,7 @@ public bool ShouldExecute(ISet<UrlToWatch> watchedUrls) =>
 }
 
 public class ProxyResponseArgs : ProxyHttpEventArgsBase {
-    public ProxyResponseArgs(SessionEventArgs session, ResponseState responseState) : base(session) {
+    public ProxyResponseArgs(SessionEventArgs session, IList<ThrottlerInfo> throttledRequests, ResponseState responseState) : base(session, throttledRequests) {
         ResponseState = responseState ?? throw new ArgumentNullException(nameof(responseState));
     }
     public ResponseState ResponseState { get; }

msgraph-developer-proxy-plugins/Behavior/RateLimitingPlugin.cs

@@ -5,7 +5,6 @@
 using Microsoft.Graph.DeveloperProxy.Abstractions;
 using System.Net;
 using System.Text.Json;
-using System.Text.Json.Serialization;
 using System.Text.RegularExpressions;
 using Titanium.Web.Proxy.Http;
 using Titanium.Web.Proxy.Models;
@@ -27,46 +26,14 @@ public class RateLimitConfiguration {
 public class RateLimitingPlugin : BaseProxyPlugin {
     public override string Name => nameof(RateLimitingPlugin);
     private readonly RateLimitConfiguration _configuration = new();
-    private readonly Dictionary<string, DateTime> _throttledRequests = new();
     // initial values so that we know when we intercept the
     // first request and can set the initial values
     private int _resourcesRemaining = -1;
     private DateTime _resetTime = DateTime.MinValue;
 
-    private bool ShouldForceThrottle(ProxyRequestArgs e) {
-        var r = e.Session.HttpClient.Request;
-        string key = BuildThrottleKey(r);
-        if (_throttledRequests.TryGetValue(key, out DateTime retryAfterDate)) {
-            if (retryAfterDate > DateTime.Now) {
-                _logger?.LogRequest(new[] { $"Calling {r.Url} again before waiting for the Retry-After period.", "Request will be throttled" }, MessageType.Failed, new LoggingContext(e.Session));
-                // update the retryAfterDate to extend the throttling window to ensure that brute forcing won't succeed.
-                _throttledRequests[key] = retryAfterDate.AddSeconds(_configuration.RetryAfterSeconds);
-                return true;
-            }
-            else {
-                // clean up expired throttled request and ensure that this request is passed through.
-                _throttledRequests.Remove(key);
-                return false;
-            }
-        }
-
-        return false;
-    }
-
-    private void ForceThrottleResponse(ProxyRequestArgs e) => UpdateProxyResponse(e, HttpStatusCode.TooManyRequests);
-
-    private bool ShouldThrottle(ProxyRequestArgs e) {
-        if (_resourcesRemaining > 0) {
-            return false;
-        }
-
-        var r = e.Session.HttpClient.Request;
-        string key = BuildThrottleKey(r);
-
-        _logger?.LogRequest(new[] { $"Exceeded resource limit when calling {r.Url}.", "Request will be throttled" }, MessageType.Failed, new LoggingContext(e.Session));
-        // update the retryAfterDate to extend the throttling window to ensure that brute forcing won't succeed.
-        _throttledRequests[key] = DateTime.Now.AddSeconds(_configuration.RetryAfterSeconds);
-        return true;
+    private ThrottlingInfo ShouldThrottle(Request request, string throttlingKey) {
+        var throttleKeyForRequest = BuildThrottleKey(request);
+        return new ThrottlingInfo(throttleKeyForRequest == throttlingKey ? _configuration.RetryAfterSeconds : 0, _configuration.HeaderRetryAfter);
     }
 
     private void ThrottleResponse(ProxyRequestArgs e) => UpdateProxyResponse(e, HttpStatusCode.TooManyRequests);
@@ -75,24 +42,31 @@ private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorS
         var headers = new List<HttpHeader>();
         var body = string.Empty;
         var request = e.Session.HttpClient.Request;
+        var response = e.Session.HttpClient.Response;
 
-        // override the response body and headers for the error response
-        if (errorStatus != HttpStatusCode.OK &&
-            ProxyUtils.IsGraphRequest(request)) {
-            string requestId = Guid.NewGuid().ToString();
-            string requestDate = DateTime.Now.ToString();
-            headers.AddRange(ProxyUtils.BuildGraphResponseHeaders(request, requestId, requestDate));
-
-            body = JsonSerializer.Serialize(new GraphErrorResponseBody(
-                new GraphErrorResponseError {
-                    Code = new Regex("([A-Z])").Replace(errorStatus.ToString(), m => { return $" {m.Groups[1]}"; }).Trim(),
-                    Message = BuildApiErrorMessage(request),
-                    InnerError = new GraphErrorResponseInnerError {
-                        RequestId = requestId,
-                        Date = requestDate
-                    }
-                })
-            );
+        // resources exceeded
+        if (errorStatus == HttpStatusCode.TooManyRequests) {
+            if (ProxyUtils.IsGraphRequest(request)) {
+                string requestId = Guid.NewGuid().ToString();
+                string requestDate = DateTime.Now.ToString();
+                headers.AddRange(ProxyUtils.BuildGraphResponseHeaders(request, requestId, requestDate));
+
+                body = JsonSerializer.Serialize(new GraphErrorResponseBody(
+                    new GraphErrorResponseError {
+                        Code = new Regex("([A-Z])").Replace(errorStatus.ToString(), m => { return $" {m.Groups[1]}"; }).Trim(),
+                        Message = BuildApiErrorMessage(request),
+                        InnerError = new GraphErrorResponseInnerError {
+                            RequestId = requestId,
+                            Date = requestDate
+                        }
+                    })
+                );
+            }
+
+            headers.Add(new HttpHeader(_configuration.HeaderRetryAfter, _configuration.RetryAfterSeconds.ToString()));
+
+            e.Session.GenericResponse(body ?? string.Empty, errorStatus, headers);
+            return;
         }
 
         // add rate limiting headers if reached the threshold percentage
@@ -102,24 +76,51 @@ private void UpdateProxyResponse(ProxyHttpEventArgsBase e, HttpStatusCode errorS
                 new HttpHeader(_configuration.HeaderRemaining, _resourcesRemaining.ToString()),
                 new HttpHeader(_configuration.HeaderReset, (_resetTime - DateTime.Now).TotalSeconds.ToString("N0")) // drop decimals
             });
-        }
 
-        // send an error response if we are (forced) throttling
-        if (errorStatus == HttpStatusCode.TooManyRequests) {
-            headers.Add(new HttpHeader(_configuration.HeaderRetryAfter, _configuration.RetryAfterSeconds.ToString()));
+            // 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");
+                }
 
-            e.Session.GenericResponse(body ?? string.Empty, errorStatus, headers);
-            return;
+                headers.Add(new HttpHeader("Access-Control-Expose-Headers", headerValue));
+            }
         }
 
-        if (errorStatus == HttpStatusCode.OK) {
-            // add headers to the original API response
-            e.Session.HttpClient.Response.Headers.AddHeaders(headers);
-        }
+        // add headers to the original API response
+        e.Session.HttpClient.Response.Headers.AddHeaders(headers);
     }
     private static string BuildApiErrorMessage(Request r) => $"Some error was generated by the proxy. {(ProxyUtils.IsGraphRequest(r) ? ProxyUtils.IsSdkRequest(r) ? "" : String.Join(' ', MessageUtils.BuildUseSdkForErrorsMessage(r)) : "")}";
 
-    private string BuildThrottleKey(Request r) => $"{r.Method}-{r.Url}";
+    private string BuildThrottleKey(Request r) {
+        if (ProxyUtils.IsGraphRequest(r)) {
+            return GraphUtils.BuildThrottleKey(r);
+        }
+        else {
+            return r.RequestUri.Host;
+        }
+    }
 
     public override void Register(IPluginEvents pluginEvents,
                          IProxyContext context,
@@ -134,8 +135,6 @@ public override void Register(IPluginEvents pluginEvents,
 
     // add rate limiting headers to the response from the API
     private async Task OnResponse(object? sender, ProxyResponseArgs e) {
-        var session = e.Session;
-        var state = e.ResponseState;
         if (_urlsToWatch is null ||
             !e.HasRequestUrlMatch(_urlsToWatch)) {
             return;
@@ -169,44 +168,18 @@ _urlsToWatch is null ||
 
         // subtract the cost of the request
         _resourcesRemaining -= _configuration.CostPerRequest;
-        // avoid communicating negative values
         if (_resourcesRemaining < 0) {
-            _resourcesRemaining = 0;
-        }
+            var request = e.Session.HttpClient.Request;
+
+            _logger?.LogRequest(new[] { $"Exceeded resource limit when calling {request.Url}.", "Request will be throttled" }, MessageType.Failed, new LoggingContext(e.Session));
+            e.ThrottledRequests.Add(new ThrottlerInfo(
+                BuildThrottleKey(request),
+                ShouldThrottle,
+                DateTime.Now.AddSeconds(_configuration.RetryAfterSeconds)
+            ));
 
-        if (ShouldForceThrottle(e)) {
-            ForceThrottleResponse(e);
-            state.HasBeenSet = true;
-        }
-        else if (ShouldThrottle(e)) {
             ThrottleResponse(e);
             state.HasBeenSet = true;
         }
     }
 }
-
-
-internal class GraphErrorResponseBody {
-    [JsonPropertyName("error")]
-    public GraphErrorResponseError Error { get; set; }
-
-    public GraphErrorResponseBody(GraphErrorResponseError error) {
-        Error = error;
-    }
-}
-
-internal class GraphErrorResponseError {
-    [JsonPropertyName("code")]
-    public string Code { get; set; } = string.Empty;
-    [JsonPropertyName("message")]
-    public string Message { get; set; } = string.Empty;
-    [JsonPropertyName("innerError")]
-    public GraphErrorResponseInnerError? InnerError { get; set; }
-}
-
-internal class GraphErrorResponseInnerError {
-    [JsonPropertyName("request-id")]
-    public string RequestId { get; set; } = string.Empty;
-    [JsonPropertyName("date")]
-    public string Date { get; set; } = string.Empty;
-}