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_proxies/dalli_proxy'
+require 'rack/attack/store_proxies/mem_cache_store_proxy'
+require 'rack/attack/store_proxies/redis_proxy'
+require 'rack/attack/store_proxies/redis_store_proxy'
+require 'rack/attack/store_proxies/redis_cache_store_proxy'
+require 'rack/attack/store_proxies/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/cache.rb

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

…roxy/active_support_redis_store_proxy.rb …xies/active_support_redis_store_proxy.rblib/rack/attack/store_proxy/active_support_redis_store_proxy.rb renamed to lib/rack/attack/store_proxies/active_support_redis_store_proxy.rb lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb renamed to lib/rack/attack/store_proxies/active_support_redis_store_proxy.rb

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

…b/rack/attack/store_proxy/dalli_proxy.rb …rack/attack/store_proxies/dalli_proxy.rblib/rack/attack/store_proxy/dalli_proxy.rb renamed to lib/rack/attack/store_proxies/dalli_proxy.rb lib/rack/attack/store_proxy/dalli_proxy.rb renamed to lib/rack/attack/store_proxies/dalli_proxy.rb

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

…ack/store_proxy/mem_cache_store_proxy.rb …k/store_proxies/mem_cache_store_proxy.rblib/rack/attack/store_proxy/mem_cache_store_proxy.rb renamed to lib/rack/attack/store_proxies/mem_cache_store_proxy.rb lib/rack/attack/store_proxy/mem_cache_store_proxy.rb renamed to lib/rack/attack/store_proxies/mem_cache_store_proxy.rb

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

…k/store_proxy/redis_cache_store_proxy.rb …store_proxies/redis_cache_store_proxy.rblib/rack/attack/store_proxy/redis_cache_store_proxy.rb renamed to lib/rack/attack/store_proxies/redis_cache_store_proxy.rb lib/rack/attack/store_proxy/redis_cache_store_proxy.rb renamed to lib/rack/attack/store_proxies/redis_cache_store_proxy.rb

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

…b/rack/attack/store_proxy/redis_proxy.rb …rack/attack/store_proxies/redis_proxy.rblib/rack/attack/store_proxy/redis_proxy.rb renamed to lib/rack/attack/store_proxies/redis_proxy.rb lib/rack/attack/store_proxy/redis_proxy.rb renamed to lib/rack/attack/store_proxies/redis_proxy.rb

@@ -1,11 +1,11 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/store_proxy'
 
 module Rack
   class Attack
-    module StoreProxy
-      class RedisProxy < SimpleDelegator
+    module StoreProxies
+      class RedisProxy < StoreProxy
         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)

…/attack/store_proxy/redis_store_proxy.rb …ttack/store_proxies/redis_store_proxy.rblib/rack/attack/store_proxy/redis_store_proxy.rb renamed to lib/rack/attack/store_proxies/redis_store_proxy.rb lib/rack/attack/store_proxy/redis_store_proxy.rb renamed to lib/rack/attack/store_proxies/redis_store_proxy.rb

@@ -1,10 +1,10 @@
 # frozen_string_literal: true
 
-require 'delegate'
+require 'rack/attack/store_proxies/redis_proxy'
 
 module Rack
   class Attack
-    module StoreProxy
+    module StoreProxies
       class RedisStoreProxy < RedisProxy
         def self.handle?(store)
           defined?(::Redis::Store) && store.is_a?(::Redis::Store)

lib/rack/attack/store_proxy.rb

@@ -1,20 +1,26 @@
 # frozen_string_literal: true
 
+require 'delegate'
+
 module Rack
   class Attack
-    module StoreProxy
-      PROXIES = [
-        DalliProxy,
-        MemCacheStoreProxy,
-        RedisStoreProxy,
-        RedisProxy,
-        RedisCacheStoreProxy,
-        ActiveSupportRedisStoreProxy
-      ].freeze
+    class StoreProxy < SimpleDelegator
+      class << self
+        def proxies
+          @@proxies ||= []
+        end
+
+        def inherited(klass)
+          proxies << klass
+        end
+
+        def lookup(store)
+          proxies.find { |proxy| proxy.handle?(store) }
+        end
 
-      def self.build(store)
-        klass = PROXIES.find { |proxy| proxy.handle?(store) }
-        klass ? klass.new(store) : store
+        def handle?(_store)
+          raise NotImplementedError
+        end
       end
     end
   end

spec/rack_attack_dalli_proxy_spec.rb

@@ -2,9 +2,9 @@
 
 require_relative 'spec_helper'
 
-describe Rack::Attack::StoreProxy::DalliProxy do
+describe Rack::Attack::StoreProxies::DalliProxy do
   it 'should stub Dalli::Client#with on older clients' do
-    proxy = Rack::Attack::StoreProxy::DalliProxy.new(Class.new)
+    proxy = Rack::Attack::StoreProxies::DalliProxy.new(Class.new)
     proxy.with {} # will not raise an error
   end
 end