Make store proxies lookup dynamic

Author: fatkodima

lib/rack/attack.rb

@@ -6,6 +6,12 @@
 require 'rack/attack/configuration'
 require 'rack/attack/path_normalizer'
 require 'rack/attack/request'
+require 'rack/attack/store_proxy/dalli_proxy'
+require 'rack/attack/store_proxy/mem_cache_store_proxy'
+require 'rack/attack/store_proxy/redis_proxy'
+require 'rack/attack/store_proxy/redis_store_proxy'
+require 'rack/attack/store_proxy/redis_cache_store_proxy'
+require 'rack/attack/store_proxy/active_support_redis_store_proxy'
 
 require 'rack/attack/railtie' if defined?(::Rails)
 
@@ -21,13 +27,6 @@ class IncompatibleStoreError < Error; end
     autoload :Safelist,             'rack/attack/safelist'
     autoload :Blocklist,            'rack/attack/blocklist'
     autoload :Track,                'rack/attack/track'
-    autoload :StoreProxy,           'rack/attack/store_proxy'
-    autoload :DalliProxy,           'rack/attack/store_proxy/dalli_proxy'
-    autoload :MemCacheStoreProxy,   'rack/attack/store_proxy/mem_cache_store_proxy'
-    autoload :RedisProxy,           'rack/attack/store_proxy/redis_proxy'
-    autoload :RedisStoreProxy,      'rack/attack/store_proxy/redis_store_proxy'
-    autoload :RedisCacheStoreProxy, 'rack/attack/store_proxy/redis_cache_store_proxy'
-    autoload :ActiveSupportRedisStoreProxy, 'rack/attack/store_proxy/active_support_redis_store_proxy'
     autoload :Fail2Ban,             'rack/attack/fail2ban'
     autoload :Allow2Ban,            'rack/attack/allow2ban'
 

lib/rack/attack/base_proxy.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require 'delegate'
+
+module Rack
+  class Attack
+    class BaseProxy < SimpleDelegator
+      class << self
+        def proxies
+          @@proxies ||= []
+        end
+
+        def inherited(klass)
+          proxies << klass
+        end
+
+        def lookup(store)
+          proxies.find { |proxy| proxy.handle?(store) }
+        end
+
+        def handle?(_store)
+          raise NotImplementedError
+        end
+      end
+    end
+  end
+end

lib/rack/attack/cache.rb

@@ -13,7 +13,12 @@ def initialize
 
       attr_reader :store
       def store=(store)
-        @store = StoreProxy.build(store)
+        @store =
+          if (proxy = BaseProxy.lookup(store))
+            proxy.new(store)
+          else
+            store
+          end
       end
 
       def count(unprefixed_key, period)

lib/rack/attack/store_proxy.rb

@@ -1,11 +1,11 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/base_proxy'
 
 module Rack
   class Attack
     module StoreProxy
-      class ActiveSupportRedisStoreProxy < SimpleDelegator
+      class ActiveSupportRedisStoreProxy < BaseProxy
         def self.handle?(store)
           defined?(::Redis) &&
             defined?(::ActiveSupport::Cache::RedisStore) &&

lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb

@@ -1,11 +1,11 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/base_proxy'
 
 module Rack
   class Attack
     module StoreProxy
-      class DalliProxy < SimpleDelegator
+      class DalliProxy < BaseProxy
         def self.handle?(store)
           return false unless defined?(::Dalli)
 

lib/rack/attack/store_proxy/dalli_proxy.rb

@@ -1,11 +1,11 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/base_proxy'
 
 module Rack
   class Attack
     module StoreProxy
-      class MemCacheStoreProxy < SimpleDelegator
+      class MemCacheStoreProxy < BaseProxy
         def self.handle?(store)
           defined?(::Dalli) &&
             defined?(::ActiveSupport::Cache::MemCacheStore) &&

lib/rack/attack/store_proxy/mem_cache_store_proxy.rb

@@ -1,11 +1,11 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/base_proxy'
 
 module Rack
   class Attack
     module StoreProxy
-      class RedisCacheStoreProxy < SimpleDelegator
+      class RedisCacheStoreProxy < BaseProxy
         def self.handle?(store)
           store.class.name == "ActiveSupport::Cache::RedisCacheStore"
         end

lib/rack/attack/store_proxy/redis_cache_store_proxy.rb

@@ -1,11 +1,11 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/base_proxy'
 
 module Rack
   class Attack
     module StoreProxy
-      class RedisProxy < SimpleDelegator
+      class RedisProxy < BaseProxy
         def initialize(*args)
           if Gem::Version.new(Redis::VERSION) < Gem::Version.new("3")
             warn 'RackAttack requires Redis gem >= 3.0.0.'
@@ -15,7 +15,7 @@ def initialize(*args)
         end
 
         def self.handle?(store)
-          defined?(::Redis) && store.is_a?(::Redis)
+          defined?(::Redis) && store.class == ::Redis
         end
 
         def read(key)

lib/rack/attack/store_proxy/redis_proxy.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/store_proxy/redis_proxy'
 
 module Rack
   class Attack