From d1e73754f9aa6524d592091f23e4f80c3cbd1b40 Mon Sep 17 00:00:00 2001
From: Cameron Hill <cameron.hill@clickview.com.au>
Date: Thu, 4 Feb 2016 15:37:49 +1100
Subject: [PATCH] Added option to use the controller action as the throttle key

---
 WebApiThrottle/Models/EndpointThrottlingMethod.cs |  8 ++++++++
 WebApiThrottle/Models/RequestIdentity.cs          |  4 ++++
 WebApiThrottle/ThrottlePolicy.cs                  |  4 ++++
 WebApiThrottle/ThrottlingCore.cs                  | 14 +++++++++++++-
 WebApiThrottle/ThrottlingFilter.cs                |  8 ++++++++
 WebApiThrottle/WebApiThrottle.csproj              |  1 +
 6 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 WebApiThrottle/Models/EndpointThrottlingMethod.cs

diff --git a/WebApiThrottle/Models/EndpointThrottlingMethod.cs b/WebApiThrottle/Models/EndpointThrottlingMethod.cs
new file mode 100644
index 0000000..1ce2f14
--- /dev/null
+++ b/WebApiThrottle/Models/EndpointThrottlingMethod.cs
@@ -0,0 +1,8 @@
+namespace WebApiThrottle.Models
+{
+    public enum EndpointThrottlingMethod
+    {
+        Url,
+        Action
+    }
+}
diff --git a/WebApiThrottle/Models/RequestIdentity.cs b/WebApiThrottle/Models/RequestIdentity.cs
index eaa2a32..e86c1cf 100644
--- a/WebApiThrottle/Models/RequestIdentity.cs
+++ b/WebApiThrottle/Models/RequestIdentity.cs
@@ -17,5 +17,9 @@ public class RequestIdentity
         public string ClientKey { get; set; }
 
         public string Endpoint { get; set; }
+
+        public string ActionName { get; set; }
+
+        public string ControllerName { get; set; }
     }
 }
diff --git a/WebApiThrottle/ThrottlePolicy.cs b/WebApiThrottle/ThrottlePolicy.cs
index b892d19..f8a66e8 100644
--- a/WebApiThrottle/ThrottlePolicy.cs
+++ b/WebApiThrottle/ThrottlePolicy.cs
@@ -6,6 +6,8 @@
 
 namespace WebApiThrottle
 {
+    using Models;
+
     /// <summary>
     /// Rate limits policy
     /// </summary>
@@ -80,6 +82,8 @@ public ThrottlePolicy(long? perSecond = null, long? perMinute = null, long? perH
         /// </summary>
         public bool EndpointThrottling { get; set; }
 
+        public EndpointThrottlingMethod EndpointThrottleMethod { get; set; }
+
         public List<string> EndpointWhitelist { get; set; }
 
         public IDictionary<string, RateLimits> EndpointRules { get; set; }
diff --git a/WebApiThrottle/ThrottlingCore.cs b/WebApiThrottle/ThrottlingCore.cs
index 2b3e711..0c2091a 100644
--- a/WebApiThrottle/ThrottlingCore.cs
+++ b/WebApiThrottle/ThrottlingCore.cs
@@ -8,6 +8,8 @@
 
 namespace WebApiThrottle
 {
+    using Models;
+
     /// <summary>
     /// Common code shared between ThrottlingHandler and ThrottlingFilter
     /// </summary>
@@ -130,7 +132,17 @@ internal string ComputeThrottleKey(RequestIdentity requestIdentity, RateLimitPer
 
             if (Policy.EndpointThrottling)
             {
-                keyValues.Add(requestIdentity.Endpoint);
+                if (Policy.EndpointThrottleMethod == EndpointThrottlingMethod.Action &&
+                    !string.IsNullOrWhiteSpace(requestIdentity.ActionName) &&
+                    !string.IsNullOrWhiteSpace(requestIdentity.ControllerName))
+                {
+                    //TODO: think of a better way to do this
+                    keyValues.Add(requestIdentity.ControllerName + "." + requestIdentity.ActionName);
+                }
+                else if (Policy.EndpointThrottleMethod == EndpointThrottlingMethod.Url)
+                {
+                    keyValues.Add(requestIdentity.Endpoint);
+                }
             }
 
             keyValues.Add(period.ToString());
diff --git a/WebApiThrottle/ThrottlingFilter.cs b/WebApiThrottle/ThrottlingFilter.cs
index 648970d..402e30b 100644
--- a/WebApiThrottle/ThrottlingFilter.cs
+++ b/WebApiThrottle/ThrottlingFilter.cs
@@ -228,6 +228,14 @@ protected virtual RequestIdentity SetIndentity(HttpRequestMessage request)
             entry.ClientIp = core.GetClientIp(request).ToString();
             entry.Endpoint = request.RequestUri.AbsolutePath.ToLowerInvariant();
             entry.ClientKey = request.Headers.Contains("Authorization-Token") ? request.Headers.GetValues("Authorization-Token").First() : "anon";
+            
+            var actionDescriptor = request.GetActionDescriptor();
+            if (actionDescriptor == null)
+                return entry;
+
+            entry.ActionName = actionDescriptor.ActionName;
+            if(actionDescriptor.ControllerDescriptor != null)
+                entry.ControllerName = actionDescriptor.ControllerDescriptor.ControllerName;
 
             return entry;
         }
diff --git a/WebApiThrottle/WebApiThrottle.csproj b/WebApiThrottle/WebApiThrottle.csproj
index 5b4959c..b22364e 100644
--- a/WebApiThrottle/WebApiThrottle.csproj
+++ b/WebApiThrottle/WebApiThrottle.csproj
@@ -69,6 +69,7 @@
   </ItemGroup>
   <ItemGroup>
     <Compile Include="Attributes\DisableThrottingAttribute.cs" />
+    <Compile Include="Models\EndpointThrottlingMethod.cs" />
     <Compile Include="Net\DefaultIpAddressParser.cs" />
     <Compile Include="Net\IIpAddressParser.cs" />
     <Compile Include="Net\IpAddressUtil.cs" />