Fix RealIpHeader & ClientIdHeader resolving

stefanprodan · Dec 12, 2020 · 5434672 · 5434672
1 parent f6248d3
commit 5434672
Showing 1 changed file with 17 additions and 4 deletions.
diff --git a/src/AspNetCoreRateLimit/Middleware/RateLimitConfiguration.cs b/src/AspNetCoreRateLimit/Middleware/RateLimitConfiguration.cs
@@ -30,18 +30,31 @@ public class RateLimitConfiguration : IRateLimitConfiguration
 
         public virtual void RegisterResolvers()
         {
-            if (!string.IsNullOrEmpty(ClientRateLimitOptions?.ClientIdHeader) || !string.IsNullOrEmpty(IpRateLimitOptions?.ClientIdHeader))
+            string clientIdHeader = GetClientIdHeader();
+            string realIpHeader = GetRealIp();
+
+            if (clientIdHeader != null)
             {
-                ClientResolvers.Add(new ClientHeaderResolveContributor(HttpContextAccessor, ClientRateLimitOptions.ClientIdHeader));
+                ClientResolvers.Add(new ClientHeaderResolveContributor(HttpContextAccessor, clientIdHeader));
             }
 
             // the contributors are resolved in the order of their collection index
-            if (!string.IsNullOrEmpty(ClientRateLimitOptions?.RealIpHeader) || !string.IsNullOrEmpty(IpRateLimitOptions?.RealIpHeader))
+            if (realIpHeader != null)
             {
-                IpResolvers.Add(new IpHeaderResolveContributor(HttpContextAccessor, IpRateLimitOptions.RealIpHeader));
+                IpResolvers.Add(new IpHeaderResolveContributor(HttpContextAccessor, realIpHeader));
             }
 
             IpResolvers.Add(new IpConnectionResolveContributor(HttpContextAccessor));
         }
+
+        protected string GetClientIdHeader()
+        {
+            return ClientRateLimitOptions?.ClientIdHeader ?? IpRateLimitOptions?.ClientIdHeader;
+        }
+
+        protected string GetRealIp()
+        {
+            return IpRateLimitOptions?.RealIpHeader ?? ClientRateLimitOptions?.RealIpHeader;
+        }
     }
 }